Aharrypotter/zl_bench
1
0 Fork 0
forked from zhoulin/zl_bench
Code Pull Requests Activity
Files
d45806eacd30004cfa81e7937c80ba84c8f355e5
zl_bench/test_backends.py
T
AharrypotterandClaude Opus 5 c95474a432 bench: refuse to report a comparison whose two sides have different calibers 
Two runs on different calibers produce a report that looks exactly like a
valid one — no step errors, the numbers just mean something else. This repo
has published three figures that were an artifact of such a difference:
router_gemm -8.2% (record vs record), mm "945 variants -> 33" (shared vs
inherited libtuner cache), replication_pad3d +43% vs +124.5% (timer).

batch_summary --ab / --ab-rounds now reads every participating run's
run_identity.json and compares backend/device/timer/bench_level/shape_mode/
tuner_cache/direct_repeat before parsing a single latency. Mismatch, or a
missing identity, prints !!! [caliber] and exits 4. autotune_mode is judged
separately: record pinning a replay is the normal shape, both sides record
is the exact trap above and is refused by name. --allow-caliber-mismatch
releases it, still printing what it found.

Recording on a *shared* tuner cache warns rather than refuses: the pinned
config depends on what else ran on the machine, which leaves relative
conclusions intact and absolute latencies unquotable. Refusing it would
reject the default every A/B here currently uses.

ab_repeat.sh closes the one axis that had no check: after round 1 it
fingerprints both sides' ttgir (relative path + content, so a changed
variant *set* counts too). Byte-identical means the switch changed no
generated code for this op, and the remaining N-1 rounds would measure
noise and report it as "no difference" — so the op is skipped with
SWITCH_NO_IR_DIFF unless SWITCH_ALLOW_NO_IR_DIFF=1.

README leads with a caliber table: each axis, its measured magnitude, what
pins it, and which code checks it. The fourth column is code because
documentation did not hold — the README already said "the B side must use
REPLAY_FROM" when the record-vs-record number was published.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-10 14:32:49 +08:00

546 lines
23 KiB
Python
Raw Blame History

#!/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))
class CaliberGuardTests(unittest.TestCase):
"""两次运行口径不同产出的报告，和有效报告长得一模一样。
这些用例钉住的是本仓已经踩过的三次：record 对 record、共享 tuner cache、
timer 不同。它们都不会让任何一步报错，只会让数字变成另一个含义。
"""
BASE = {
"backend": "musa", "device": "S5000", "timer": "cuda_graph",
"bench_level": "core", "shape_mode": "default",
"tuner_cache": "isolated", "direct_repeat": 5,
"autotune_mode": "replay",
}
def _run(self, root, name, **over):
d = Path(root, name)
d.mkdir(parents=True)
if over.get("_no_identity"):
return d
ident = dict(self.BASE)
ident.update({k: v for k, v in over.items() if not k.startswith("_")})
Path(d, "run_identity.json").write_text(
json.dumps(ident), encoding="utf-8")
return d
def test_same_caliber_has_no_complaints(self):
from batch_summary import caliber_complaints
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a")), ("B", self._run(tmp, "b"))]
self.assertEqual(caliber_complaints(runs), [])
def test_record_side_pinning_the_other_is_normal(self):
"""A/B 的正常形态：一侧 record 定 config，另一侧 replay 它。"""
from batch_summary import caliber_complaints
with tempfile.TemporaryDirectory() as tmp:
runs = [("r1_A", self._run(tmp, "a", autotune_mode="record")),
("r1_B", self._run(tmp, "b"))]
self.assertEqual(caliber_complaints(runs), [])
def test_record_against_record_is_flagged(self):
from batch_summary import caliber_complaints
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a", autotune_mode="record")),
("B", self._run(tmp, "b", autotune_mode="record"))]
self.assertTrue(any("record" in c for c in caliber_complaints(runs)))
def test_each_caliber_field_is_checked(self):
from batch_summary import CALIBER_FIELDS, caliber_complaints
others = {"backend": "cuda", "device": "H20", "timer": "direct",
"bench_level": "comprehensive", "shape_mode": "yaml",
"tuner_cache": "shared", "direct_repeat": 1}
self.assertEqual(set(others), set(CALIBER_FIELDS))
for field, other in others.items():
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a")),
("B", self._run(tmp, "b", **{field: other}))]
got = caliber_complaints(runs)
self.assertTrue(any(c.startswith(field) for c in got),
f"{field} 不一致时没被拦下: {got}")
def test_missing_identity_is_flagged_not_assumed_equal(self):
from batch_summary import caliber_complaints
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a")),
("B", self._run(tmp, "b", _no_identity=True))]
self.assertTrue(any("run_identity" in c
for c in caliber_complaints(runs)))
def test_shared_cache_record_warns_but_does_not_block(self):
from batch_summary import caliber_complaints, caliber_hazards
with tempfile.TemporaryDirectory() as tmp:
runs = [("r1_A", self._run(tmp, "a", autotune_mode="record",
tuner_cache="shared(~/.flaggems)")),
("r1_B", self._run(tmp, "b",
tuner_cache="shared(~/.flaggems)"))]
self.assertEqual(caliber_complaints(runs), [])
self.assertTrue(caliber_hazards(runs))
def test_isolated_cache_has_no_hazard(self):
from batch_summary import caliber_hazards
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a", autotune_mode="record")),
("B", self._run(tmp, "b"))]
self.assertEqual(caliber_hazards(runs), [])
def test_strict_refuses_and_permissive_continues(self):
from batch_summary import _report_caliber
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a")),
("B", self._run(tmp, "b", timer="direct"))]
self.assertFalse(_report_caliber(runs, "op", strict=True))
self.assertTrue(_report_caliber(runs, "op", strict=False))
LOG = ("Operator: op (dtype=torch.float16, mode=kernel, level=core)\n"
"SUCCESS 1.000 2.000 0.500 shape\n")
def test_report_ab_refuses_mismatched_pair(self):
from batch_summary import CaliberMismatch, report_ab
with tempfile.TemporaryDirectory() as tmp:
a = self._run(tmp, "a")
b = self._run(tmp, "b", timer="direct")
for d in (a, b):
Path(d, "run.log").write_text(self.LOG, encoding="utf-8")
with self.assertRaises(CaliberMismatch):
report_ab(a / "run.log", b / "run.log", "op")
# 放行开关在，但必须是显式的
self.assertEqual(
len(report_ab(a / "run.log", b / "run.log", "op", strict=False)), 1)
if __name__ == "__main__":
unittest.main()
View Git Blame Copy Permalink