forked from zhoulin/zl_bench
AharrypotterandClaude Opus 5
9665d36a7c
do_bench already returns a median over the iterations inside one call, which suppresses iteration-level jitter. It cannot see drift across a call: clocks ramping, another process arriving. Direct timing is where that shows, because it walks the full Python launch path once per iteration, so on small kernels the number carries launch overhead and CPU jitter that graph capture removes. ZL_DIRECT_REPEAT=N repeats the call and takes the median over calls, which is the layer the inner median misses. It defaults to 1, so nothing changes unless asked. The knob covers both ways a row ends up on the direct timer: a run that selected timer=direct outright, and the rows a captured run falls back on, whose marker now carries repeat=N since those rows were already a different caliber from the captured ones and can now also be a different sample count. Repetition is skipped when the caller asks for quantiles or return_mode "all", because the result is then a sequence and a median across calls is undefined for it. On MThreads this is not a refinement. A convert-reduction screen over 21 operators puts the largest reductions at sort (65->15), dropout (48->0) and rand (24->0), and all three are exactly the operators that cannot be graph-captured there -- RNG seeding and a capture-illegal stream op. The headline numbers on that device are produced by the noisiest timer, which until now had no repetition at all. _cudagraph_plugin imports triton at module scope, so the tests load it against a stub: the repeat logic is arithmetic and should stay testable without a GPU. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
436 lines
18 KiB
Python
436 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""No-GPU unit tests for zl_bench backend adapters and plugin policy.
|
|
|
|
FlagGems pytest invocations pass an explicit node, so this file is not
|
|
collected during operator runs. Run it with `python test_backends.py`.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
if str(ROOT) not in sys.path:
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from backends import ( # noqa: E402
|
|
UnknownBackendError,
|
|
get_backend,
|
|
normalize_backend_name,
|
|
)
|
|
from backends.base import ALL_VISIBILITY_KEYS # noqa: E402
|
|
from batch_summary import SUMMARY_FIELDS # noqa: E402
|
|
from plugin_plan import plan_plugins # noqa: E402
|
|
from providers.flaggems import ( # noqa: E402
|
|
VENDOR_BY_BACKEND,
|
|
emit_run_env,
|
|
flaggems_vendor,
|
|
parse_test_node,
|
|
)
|
|
from run_batch import STATUS_FIELDS, resolve_device_arg # noqa: E402
|
|
from _run_identity import init_identity, load_identity, update_identity # noqa: E402
|
|
|
|
|
|
class BackendResolveTests(unittest.TestCase):
|
|
def test_empty_and_alias_resolve_to_cuda(self):
|
|
self.assertEqual(normalize_backend_name(""), "cuda")
|
|
self.assertEqual(get_backend("nvidia").name, "cuda")
|
|
self.assertEqual(get_backend("nv").name, "cuda")
|
|
self.assertEqual(get_backend("").name, "cuda")
|
|
|
|
def test_other_aliases(self):
|
|
self.assertEqual(get_backend("amd").name, "rocm")
|
|
self.assertEqual(get_backend("hip").name, "rocm")
|
|
self.assertEqual(get_backend("npu").name, "ascend")
|
|
self.assertEqual(get_backend("huawei").name, "ascend")
|
|
|
|
def test_unknown_backend_fails_closed(self):
|
|
with self.assertRaises(UnknownBackendError):
|
|
get_backend("tpu")
|
|
with self.assertRaises(UnknownBackendError):
|
|
get_backend("not-a-backend")
|
|
|
|
|
|
class VisibilityTests(unittest.TestCase):
|
|
def test_musa_drops_leftover_cuda(self):
|
|
env = {"CUDA_VISIBLE_DEVICES": "7", "PATH": "/bin"}
|
|
out = get_backend("musa").apply_visibility(env, "1")
|
|
self.assertNotIn("CUDA_VISIBLE_DEVICES", out)
|
|
self.assertEqual(out["MUSA_VISIBLE_DEVICES"], "1")
|
|
self.assertEqual(out["PATH"], "/bin")
|
|
for key in ALL_VISIBILITY_KEYS:
|
|
if key != "MUSA_VISIBLE_DEVICES":
|
|
self.assertNotIn(key, out)
|
|
|
|
def test_rocm_writes_hip_and_rocr(self):
|
|
out = get_backend("rocm").apply_visibility(
|
|
{"CUDA_VISIBLE_DEVICES": "0"}, "3")
|
|
self.assertEqual(out["HIP_VISIBLE_DEVICES"], "3")
|
|
self.assertEqual(out["ROCR_VISIBLE_DEVICES"], "3")
|
|
self.assertNotIn("CUDA_VISIBLE_DEVICES", out)
|
|
|
|
def test_ascend_writes_both_keys(self):
|
|
out = get_backend("ascend").apply_visibility({}, "2")
|
|
self.assertEqual(out["ASCEND_RT_VISIBLE_DEVICES"], "2")
|
|
self.assertEqual(out["ASCEND_VISIBLE_DEVICES"], "2")
|
|
self.assertNotIn("CUDA_VISIBLE_DEVICES", out)
|
|
|
|
def test_cuda_only_sets_cuda(self):
|
|
out = get_backend("cuda").apply_visibility(
|
|
{"MUSA_VISIBLE_DEVICES": "9"}, "0")
|
|
self.assertEqual(out["CUDA_VISIBLE_DEVICES"], "0")
|
|
self.assertNotIn("MUSA_VISIBLE_DEVICES", out)
|
|
|
|
|
|
class FlagGemsVendorTests(unittest.TestCase):
|
|
def test_vendor_map(self):
|
|
self.assertEqual(VENDOR_BY_BACKEND["cuda"], "nvidia")
|
|
self.assertEqual(VENDOR_BY_BACKEND["rocm"], "amd")
|
|
self.assertEqual(flaggems_vendor(get_backend("cuda")), "nvidia")
|
|
self.assertEqual(flaggems_vendor(get_backend("rocm")), "amd")
|
|
self.assertEqual(flaggems_vendor(get_backend("musa")), "musa")
|
|
self.assertEqual(flaggems_vendor(get_backend("ascend")), "ascend")
|
|
|
|
|
|
class FlagGemsNodeTests(unittest.TestCase):
|
|
_NODE = (
|
|
"/workspace/FlagGems/benchmark/test_fused_marlin_moe.py"
|
|
"::test_fused_marlin_moe_mxfp4"
|
|
)
|
|
|
|
def test_parse_user_command_node(self):
|
|
parsed = parse_test_node(self._NODE)
|
|
self.assertEqual(parsed["op"], "fused_marlin_moe_mxfp4")
|
|
self.assertEqual(parsed["op_file"], "fused_marlin_moe")
|
|
self.assertEqual(parsed["flaggems_dir"], "/workspace/FlagGems")
|
|
self.assertTrue(parsed["test_file"].endswith(
|
|
"test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4"))
|
|
|
|
def test_test_node_defaults_to_upstream_shapes(self):
|
|
env = emit_run_env(
|
|
test_node=self._NODE,
|
|
op="softmax",
|
|
op_file="softmax",
|
|
flaggems_dir="/workspace/dev/FlagGems",
|
|
)
|
|
self.assertEqual(env["OP"], "fused_marlin_moe_mxfp4")
|
|
self.assertEqual(env["OP_FILE"], "fused_marlin_moe")
|
|
self.assertEqual(env["FLAGGEMS_DIR"], "/workspace/FlagGems")
|
|
self.assertEqual(env["SHAPE_MODE"], "upstream")
|
|
self.assertEqual(env["PASS_SHAPE_FILE"], "0")
|
|
self.assertEqual(env["TEST_FILE"], self._NODE)
|
|
|
|
def test_plain_run_keeps_inline_shapes(self):
|
|
env = emit_run_env(
|
|
op="fused_marlin_moe_mxfp4",
|
|
op_file="fused_marlin_moe",
|
|
flaggems_dir="/workspace/dev/FlagGems",
|
|
)
|
|
self.assertEqual(env["SHAPE_MODE"], "inline")
|
|
self.assertEqual(env["PASS_SHAPE_FILE"], "1")
|
|
self.assertIn("test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4",
|
|
env["TEST_FILE"])
|
|
|
|
def test_explicit_shape_file_wins(self):
|
|
env = emit_run_env(
|
|
test_node=self._NODE,
|
|
shape_file="/tmp/my_shapes.yaml",
|
|
)
|
|
self.assertEqual(env["SHAPE_MODE"], "file")
|
|
self.assertEqual(env["PASS_SHAPE_FILE"], "1")
|
|
|
|
def test_resolve_uses_existing_local_root(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp) / "FlagGems"
|
|
py = root / "benchmark" / "test_fused_marlin_moe.py"
|
|
py.parent.mkdir(parents=True)
|
|
py.write_text("# stub\n", encoding="utf-8")
|
|
env = emit_run_env(
|
|
op="fused_marlin_moe_mxfp4",
|
|
op_file="fused_marlin_moe",
|
|
flaggems_dir=str(root),
|
|
)
|
|
self.assertEqual(env["FLAGGEMS_DIR"], str(root))
|
|
self.assertTrue(Path(env["TEST_FILE"].split("::", 1)[0]).is_file())
|
|
|
|
|
|
class PluginPlanTests(unittest.TestCase):
|
|
def _plan(self, backend, **kwargs):
|
|
kwargs.setdefault("extra", [])
|
|
kwargs.setdefault("enable_autotune", True)
|
|
kwargs.setdefault("enable_ir", True)
|
|
kwargs.setdefault("enable_mm_cluster", False)
|
|
return plan_plugins(backend, **kwargs)
|
|
|
|
def test_cuda_does_not_put_cudagraph_on_the_list(self):
|
|
names = self._plan("cuda")
|
|
self.assertIn("_timing_plugin", names)
|
|
self.assertIn("_device_setup_plugin", names)
|
|
self.assertNotIn("_cudagraph_plugin", names)
|
|
self.assertNotIn("_mm_cluster_fix_plugin", names)
|
|
|
|
def test_musa_loads_its_compat_plugin_first(self):
|
|
names = self._plan("musa")
|
|
self.assertIn("_musa_benchmark_compat_plugin", names)
|
|
# It rewrites collection-time markers, so nothing may inspect the item
|
|
# list before it runs.
|
|
self.assertEqual(names[0], "_musa_benchmark_compat_plugin")
|
|
|
|
def test_compat_plugins_do_not_leak_to_other_backends(self):
|
|
for name in ("cuda", "rocm", "ascend"):
|
|
self.assertNotIn("_musa_benchmark_compat_plugin", self._plan(name), name)
|
|
|
|
def test_musa_same_core_without_cluster_fix(self):
|
|
# _cudagraph_plugin is absent from the plan on every backend: the timing
|
|
# plugin loads it, so its absence here says nothing about capture support.
|
|
names = self._plan("musa")
|
|
self.assertIn("_timing_plugin", names)
|
|
self.assertNotIn("_cudagraph_plugin", names)
|
|
self.assertNotIn("_mm_cluster_fix_plugin", names)
|
|
|
|
def test_mm_cluster_ignored_on_non_cuda(self):
|
|
names = self._plan("musa", enable_mm_cluster=True)
|
|
self.assertNotIn("_mm_cluster_fix_plugin", names)
|
|
names = self._plan("cuda", enable_mm_cluster=True)
|
|
self.assertIn("_mm_cluster_fix_plugin", names)
|
|
|
|
def test_extra_plugins_isolated(self):
|
|
names = self._plan("cuda", extra=["_cudagraph_plugin"])
|
|
self.assertIn("_cudagraph_plugin", names)
|
|
names = self._plan("cuda", extra=[])
|
|
self.assertNotIn("_cudagraph_plugin", names)
|
|
|
|
def test_extra_none_reads_env(self):
|
|
old = os.environ.get("ZL_EXTRA_PLUGINS")
|
|
os.environ["ZL_EXTRA_PLUGINS"] = "_cudagraph_plugin"
|
|
try:
|
|
names = plan_plugins(
|
|
"cuda", extra=None, enable_autotune=True,
|
|
enable_ir=True, enable_mm_cluster=False)
|
|
self.assertIn("_cudagraph_plugin", names)
|
|
finally:
|
|
if old is None:
|
|
os.environ.pop("ZL_EXTRA_PLUGINS", None)
|
|
else:
|
|
os.environ["ZL_EXTRA_PLUGINS"] = old
|
|
|
|
|
|
class CaptureClassificationTests(unittest.TestCase):
|
|
"""Graph-capture failures split three ways; the wording is vendor text."""
|
|
|
|
def test_cuda_and_musa_both_capture(self):
|
|
self.assertTrue(get_backend("cuda").supports_cuda_graph)
|
|
# torch_musa supplies Stream/Event/graph/MUSAGraph; measured on S5000.
|
|
self.assertTrue(get_backend("musa").supports_cuda_graph)
|
|
|
|
def test_recoverable_is_the_default(self):
|
|
# A kernel refusing capture is ordinary and must reach the fallback.
|
|
for name in ("cuda", "musa"):
|
|
v = get_backend(name).classify_capture_error(
|
|
RuntimeError("operation not permitted when stream is capturing"))
|
|
self.assertEqual(v, "recoverable", name)
|
|
|
|
def test_cuda_fatal_markers(self):
|
|
b = get_backend("cuda")
|
|
self.assertEqual(
|
|
b.classify_capture_error(RuntimeError(
|
|
"CUDA error: an illegal memory access was encountered")), "fatal")
|
|
self.assertEqual(
|
|
b.classify_capture_error(RuntimeError("CUDA out of memory")), "fatal")
|
|
|
|
def test_torch_dsa_footer_does_not_make_everything_fatal(self):
|
|
# torch appends this to every AcceleratorError; matching it loosely
|
|
# would classify each capture refusal as fatal and kill the fallback.
|
|
exc = RuntimeError("Cannot call current_seed during capture. Compile with "
|
|
"`TORCH_USE_CUDA_DSA` to enable device-side assertions.")
|
|
self.assertEqual(get_backend("cuda").classify_capture_error(exc),
|
|
"recoverable")
|
|
|
|
def test_musa_unknown_error_is_poisoning_not_recoverable(self):
|
|
b = get_backend("musa")
|
|
exc = RuntimeError("MUSA error: unknown error")
|
|
self.assertEqual(b.classify_capture_error(exc), "poisoned")
|
|
# The same text carries no meaning for CUDA, so it must not leak across.
|
|
self.assertEqual(get_backend("cuda").classify_capture_error(exc),
|
|
"recoverable")
|
|
|
|
def test_musa_rng_seed_failure_stays_recoverable(self):
|
|
exc = RuntimeError("Cannot call MUSAGeneratorImpl::current_seed during "
|
|
"MUSA graph capture")
|
|
self.assertEqual(get_backend("musa").classify_capture_error(exc),
|
|
"recoverable")
|
|
|
|
def test_deny_list_is_per_backend(self):
|
|
self.assertTrue(get_backend("musa").denies_capture("unique2"))
|
|
self.assertFalse(get_backend("musa").denies_capture("softmax"))
|
|
self.assertFalse(get_backend("cuda").denies_capture("unique2"))
|
|
self.assertFalse(get_backend("musa").denies_capture(""))
|
|
|
|
def test_deny_list_is_keyed_by_pytest_function_name(self):
|
|
# FLAGGEMS_PERF_CURRENT_OP holds the pytest function name minus `test_`,
|
|
# so a deny entry written as the benchmark FILE name never matches and
|
|
# fails open. Observed once: "unique" listed, op=unique2 at runtime.
|
|
musa = get_backend("musa")
|
|
self.assertTrue(musa.denies_capture("unique2"))
|
|
self.assertFalse(musa.denies_capture("unique"))
|
|
|
|
|
|
class DirectRepeatTests(unittest.TestCase):
|
|
"""Repetition on the direct timer.
|
|
|
|
_cudagraph_plugin imports triton at module scope, so it is loaded here
|
|
against a stub: the repeat logic is plain arithmetic and should stay
|
|
testable on a machine with no GPU and no Triton.
|
|
"""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls._saved = {k: sys.modules.get(k) for k in ("triton", "triton.testing")}
|
|
triton = types.ModuleType("triton")
|
|
testing = types.ModuleType("triton.testing")
|
|
testing.do_bench = lambda *a, **k: 0.0
|
|
testing.do_bench_cudagraph = lambda *a, **k: 0.0
|
|
testing._summarize_statistics = lambda *a, **k: 0.0
|
|
triton.testing = testing
|
|
sys.modules["triton"] = triton
|
|
sys.modules["triton.testing"] = testing
|
|
import _cudagraph_plugin # noqa: E402
|
|
cls.cg = _cudagraph_plugin
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
for name, mod in cls._saved.items():
|
|
if mod is None:
|
|
sys.modules.pop(name, None)
|
|
else:
|
|
sys.modules[name] = mod
|
|
sys.modules.pop("_cudagraph_plugin", None)
|
|
|
|
def _run(self, value, **kwargs):
|
|
"""-> (result, how many times do_bench was called)"""
|
|
seq, calls = [3.0, 1.0, 2.0, 9.0, 5.0], []
|
|
self.cg._ORIGINAL_DO_BENCH = (
|
|
lambda fn, **kw: (calls.append(1), seq[(len(calls) - 1) % len(seq)])[1])
|
|
old = os.environ.get("ZL_DIRECT_REPEAT")
|
|
os.environ["ZL_DIRECT_REPEAT"] = value
|
|
try:
|
|
out = self.cg.run_direct(lambda: None, 25, 100, None,
|
|
kwargs.get("quantiles"),
|
|
kwargs.get("return_mode", "mean"))
|
|
finally:
|
|
if old is None:
|
|
os.environ.pop("ZL_DIRECT_REPEAT", None)
|
|
else:
|
|
os.environ["ZL_DIRECT_REPEAT"] = old
|
|
return out, len(calls)
|
|
|
|
def test_default_is_one_call(self):
|
|
self.assertEqual(self._run("1"), (3.0, 1))
|
|
|
|
def test_unset_is_one_call(self):
|
|
old = os.environ.pop("ZL_DIRECT_REPEAT", None)
|
|
try:
|
|
self.assertEqual(self.cg.direct_repeat(), 1)
|
|
finally:
|
|
if old is not None:
|
|
os.environ["ZL_DIRECT_REPEAT"] = old
|
|
|
|
def test_median_over_calls(self):
|
|
# 3, 1, 2, 9, 5 -> 3
|
|
self.assertEqual(self._run("5"), (3.0, 5))
|
|
|
|
def test_garbage_and_zero_fall_back_to_one(self):
|
|
self.assertEqual(self._run("abc"), (3.0, 1))
|
|
self.assertEqual(self._run("0"), (3.0, 1))
|
|
self.assertEqual(self._run("-4"), (3.0, 1))
|
|
|
|
def test_sequence_results_are_not_repeated(self):
|
|
# quantiles / return_mode="all" make do_bench return a sequence, and a
|
|
# median across calls is undefined for it; take the single call.
|
|
self.assertEqual(self._run("5", quantiles=[0.5])[1], 1)
|
|
self.assertEqual(self._run("5", return_mode="all")[1], 1)
|
|
|
|
|
|
class DeviceArgTests(unittest.TestCase):
|
|
def test_devices_preferred(self):
|
|
self.assertEqual(resolve_device_arg("0,1", None), "0,1")
|
|
self.assertEqual(resolve_device_arg(None, "2"), "2")
|
|
self.assertIsNone(resolve_device_arg(None, None))
|
|
|
|
def test_both_set_is_error(self):
|
|
with self.assertRaises(ValueError):
|
|
resolve_device_arg("0", "1")
|
|
|
|
|
|
class CsvFieldTests(unittest.TestCase):
|
|
def test_status_fields(self):
|
|
for name in ("backend", "device", "gpu", "timer"):
|
|
self.assertIn(name, STATUS_FIELDS)
|
|
|
|
def test_summary_fields(self):
|
|
for name in ("backend", "device", "timer"):
|
|
self.assertIn(name, SUMMARY_FIELDS)
|
|
self.assertLess(SUMMARY_FIELDS.index("backend"),
|
|
SUMMARY_FIELDS.index("fg_op_name"))
|
|
|
|
|
|
class IdentityMergeTests(unittest.TestCase):
|
|
def test_init_then_update(self):
|
|
old_out = os.environ.get("OUT_DIR")
|
|
try:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
init_identity(tmp, backend="cuda", device="0", op="softmax")
|
|
update_identity(tmp, timer="direct")
|
|
ident = load_identity(tmp)
|
|
self.assertEqual(ident["schema_version"], 1)
|
|
self.assertEqual(ident["provider"], "flaggems")
|
|
self.assertEqual(ident["backend"], "cuda")
|
|
self.assertEqual(ident["device"], "0")
|
|
self.assertEqual(ident["timer"], "direct")
|
|
self.assertEqual(ident["op"], "softmax")
|
|
finally:
|
|
if old_out is None:
|
|
os.environ.pop("OUT_DIR", None)
|
|
else:
|
|
os.environ["OUT_DIR"] = old_out
|
|
|
|
|
|
class AutotuneReplayBackendTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self._old = os.environ.get("ZL_BACKEND")
|
|
os.environ["ZL_BACKEND"] = "cuda"
|
|
|
|
def tearDown(self):
|
|
if self._old is None:
|
|
os.environ.pop("ZL_BACKEND", None)
|
|
else:
|
|
os.environ["ZL_BACKEND"] = self._old
|
|
|
|
def test_missing_meta_still_replays(self):
|
|
from _autotune_record_plugin import _replay_backend_ok
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
self.assertTrue(_replay_backend_ok(tmp))
|
|
|
|
def test_matching_backend_ok(self):
|
|
from _autotune_record_plugin import _replay_backend_ok
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
Path(tmp, "meta.json").write_text(
|
|
json.dumps({"backend": "cuda"}), encoding="utf-8")
|
|
self.assertTrue(_replay_backend_ok(tmp))
|
|
|
|
def test_mismatch_refuses_replay(self):
|
|
from _autotune_record_plugin import _replay_backend_ok
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
Path(tmp, "meta.json").write_text(
|
|
json.dumps({"backend": "musa"}), encoding="utf-8")
|
|
self.assertFalse(_replay_backend_ok(tmp))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|