forked from zhoulin/zl_bench
AharrypotterandClaude Opus 5
29af59be9b
Three pieces from the fork's feature/musa-benchmark-support branch that never reached main, rebased onto the backend abstraction rather than copied verbatim. _musa_benchmark_compat_plugin drops FlagGems benchmark skip guards written as `torch.cuda.is_available()` where the code under test dispatches through flag_gems.device. On MThreads those turn a working operator into "not measured": test_fft.py skipped entirely, and unskipped on an S5000 it produces 25 rows. The plugin pins each override to one file and one exact skip reason and only acts when flag_gems has dispatched to the mthreads vendor; two guards were probed on device and deliberately left in place, which is why they are named in its docstring rather than silently absent. It is wired through a new DeviceBackend.compat_plugins rather than a name test in the planner, so the policy stays where every other backend difference lives. Compat plugins are planned first because they rewrite collection-time markers and nothing may read the item list before them; a test pins that position and another pins that they do not leak to other backends. batch_summary gains --ab: the speedup column compares gems against torch, not one build against another, so an A/B needs the two run.logs paired by operator, dtype and size detail. Only rows that succeeded on both sides with a positive latency count, so an incomplete pairing cannot look like a complete range, and the range is annotated when either side carries a BENCHMARK_DIRECT_NO_CUDAGRAPH marker, since a fallback row was timed under a different regime. ab_rlc_test.sh is the same caliber as ab_fold_test.sh with the comparison axis moved to FLAGTREE_<VENDOR>_RLC_ENHANCE. Both are caliber examples, not regression scripts: if the switch under test no longer exists the script still runs every round and prints a full table showing no difference, so confirm the axis is live before reusing either. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
359 lines
15 KiB
Python
359 lines
15 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 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 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()
|