feat: add the S5000 MUSA RLC product campaign

A frozen 32-workload campaign that measures how FlagTree's RemoveLayoutConversions
pass behaves on a Moore Threads S5000, through FlagGems public dispatch and this
repository's existing plugins.

The manifest is the contract: source commits, the single physical device, the
three screen cells off/mask0/candidate, the native-event timer, and per-workload
dtypes, timeouts and canary membership. prepare validates it against
ops/inventory_main.csv and freezes an immutable task matrix; the runner executes
the matrix serially, and the worker drives one operator cell through pytest with
_ir_meta_plugin, _autotune_record_plugin and the rest.

Three properties are worth calling out because they decide whether a verdict
means anything.

Convert topology is scoped to the measured operator. A benchmark call builds its
inputs once and then runs the operator many times under the timer, so counting
ttg.convert_layout over every executed kernel credits the candidate with
converts removed from input construction that no timer can observe. Each dumped
TTGIR file is joined to its launch count from index.tsv and only variants
launched more than measured_launch_floor contribute; the unscoped total and the
excluded files stay in the receipt, and a cell that cannot be joined to its
index is classified X rather than trusted.

Correctness is scoped to the variant a workload measures. Running a whole test
file lets an unrelated variant decide the verdict, so a workload may pin node
ids or a -k selector, never both, and any narrowing must state its reason.

Targets are resolved before the campaign starts. The inventory can lag the
pinned FlagGems commit, and such a task previously failed only inside the
worker's preflight, which exited silently and left an empty log. The runner now
resolves every benchmark and correctness target against the real tree, a
workload may name a target the inventory does not know about if it says why,
and each preflight check names itself on failure.
This commit is contained in:
Aharrypotter committed 2026-09-02 11:35:30 +08:00
1 parent 22139e4b87
commit 8bbea472bb
5 files changed
+2541

No files matched your search

@@ -0,0 +1,567 @@
from __future__ import annotations
import copy
import importlib.util
import json
import sys
import tempfile
import unittest
from types import SimpleNamespace
from pathlib import Path
HERE = Path(__file__).resolve().parent
CAMPAIGN = HERE.parent
REPO_ROOT = CAMPAIGN.parents[1]
REFERENCE = CAMPAIGN / "manifest" / "musa-rlc-32-product-v1.json"
SPEC = importlib.util.spec_from_file_location(
"prepare_musa_rlc_product", CAMPAIGN / "prepare_musa_rlc_product.py"
)
assert SPEC is not None and SPEC.loader is not None
MODULE = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(MODULE)
RUN_SPEC = importlib.util.spec_from_file_location(
"run_musa_rlc_product", CAMPAIGN / "run_musa_rlc_product.py"
)
assert RUN_SPEC is not None and RUN_SPEC.loader is not None
RUNNER = importlib.util.module_from_spec(RUN_SPEC)
RUN_SPEC.loader.exec_module(RUNNER)
COMPAT_SPEC = importlib.util.spec_from_file_location(
"_musa_benchmark_compat_plugin", REPO_ROOT / "_musa_benchmark_compat_plugin.py"
)
assert COMPAT_SPEC is not None and COMPAT_SPEC.loader is not None
COMPAT = importlib.util.module_from_spec(COMPAT_SPEC)
COMPAT_SPEC.loader.exec_module(COMPAT)
def _manifest():
return json.loads(REFERENCE.read_text(encoding="utf-8"))
class MusaRlcProductTests(unittest.TestCase):
def test_musa_fft_compatibility_override_is_narrow(self):
stale = SimpleNamespace(
name="skipif", args=(True,), kwargs={"reason": "CUDA is not available"}
)
unrelated = SimpleNamespace(
name="skipif", args=(True,), kwargs={"reason": "another reason"}
)
fft = SimpleNamespace(path=Path("benchmark/test_fft.py"), name="test_perf_fft")
other = SimpleNamespace(path=Path("benchmark/test_fft.py"), name="test_other")
self.assertTrue(COMPAT._is_stale_fft_cuda_guard(fft, stale))
self.assertFalse(COMPAT._is_stale_fft_cuda_guard(fft, unrelated))
self.assertFalse(COMPAT._is_stale_fft_cuda_guard(other, stale))
def test_runner_preserves_preflight_failure_receipts(self):
runner = (CAMPAIGN / "run_musa_rlc_product.py").read_text(encoding="utf-8")
ensure_parent = "out.mkdir(parents=True, exist_ok=True)"
write_receipt = '_atomic_json(out / "cell-receipt.json", receipt)'
self.assertIn(ensure_parent, runner)
self.assertLess(runner.index(ensure_parent), runner.index(write_receipt))
def test_worker_dumps_only_executed_ttgir_variants(self):
worker = (CAMPAIGN / "run_pytest_musa_rlc_product.sh").read_text(
encoding="utf-8"
)
self.assertIn(
'FLAGGEMS_PERF_TTGIR_DUMP_DIR="$OUT_DIR/ttgir"', worker
)
self.assertIn("-p _musa_benchmark_compat_plugin", worker)
self.assertIn(
'PYTHONPATH="$worker_dir:$ZL_BENCH_DIR:$FLAGGEMS_DIR/src:$FLAGTREE_SITE"',
worker,
)
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
dumped = root / "ttgir" / "shape" / "kernel" / "used.ttgir"
cached = root / "triton-cache" / "hash" / "loser.ttgir"
dumped.parent.mkdir(parents=True)
cached.parent.mkdir(parents=True)
dumped.write_text("ttg.convert_layout\nttg.convert_layout\n", encoding="utf-8")
cached.write_text("ttg.convert_layout\n" * 99, encoding="utf-8")
scope = RUNNER._convert_count(root)
self.assertEqual(scope["total_convert_count"], 2)
self.assertEqual(scope["ttgir_files"], 1)
# Without an index the measured scope cannot be established, so the
# count stays unscoped and the cell is marked incomplete.
self.assertEqual(scope["measured_convert_count"], 2)
self.assertFalse(scope["scope_complete"])
@staticmethod
def _write_cell(root, rows):
"""Materialise a cell's ttgir dump plus the index the dumper writes."""
ttgir = root / "ttgir"
ttgir.mkdir(parents=True)
lines = [
"\t".join(
[
"shape", "kernel", "file", "launches", "sweep_launches",
"num_warps", "num_stages", "constexprs", "cache_hash",
]
)
]
for shape, kernel, name, launches, converts in rows:
if name.startswith("("):
lines.append(
"\t".join([shape, kernel, name, str(launches), "0", "4", "3", "{}", "H"])
)
continue
path = ttgir / shape / kernel / name
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text("ttg.convert_layout\n" * converts, encoding="utf-8")
lines.append(
"\t".join([shape, kernel, name, str(launches), "0", "4", "3", "{}", "H"])
)
(ttgir / "index.tsv").write_text("\n".join(lines) + "\n", encoding="utf-8")
def test_convert_count_excludes_input_construction_kernels(self):
# Reproduces the sealed gelu_backward shape: every convert sits in the
# once-launched input builder, while the timed kernel has none.
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
self._write_cell(
root,
[
("1048576", "randn_kernel", "w16s3.ttgir", 1, 4),
("1048576", "gelu_none_kernel_rank_1", "w4s3.ttgir", 1, 0),
("1048576", "gelu_backward_none_kernel_rank_1", "w4s3.ttgir", 3795, 0),
],
)
scope = RUNNER._convert_count(root)
self.assertEqual(scope["total_convert_count"], 4)
self.assertEqual(scope["measured_convert_count"], 0)
self.assertTrue(scope["scope_complete"])
self.assertEqual(
[row["file"] for row in scope["excluded_unmeasured"]],
["1048576/randn_kernel/w16s3.ttgir"],
)
def test_convert_count_keeps_repeatedly_launched_kernels(self):
# Reproduces the sealed sort shape: converts in the dominant kernel and
# in an autotune loser that was indexed but never dumped.
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
self._write_cell(
root,
[
("1024x1024", "sweep", "w4s3.ttgir", 280, 3),
("1024x1024", "broadcast_to_kernel", "w4s3.ttgir", 140, 1),
("-", "arange_kernel", "(sweep loser, not dumped)", 0, 0),
],
)
scope = RUNNER._convert_count(root)
self.assertEqual(scope["measured_convert_count"], 4)
self.assertEqual(scope["total_convert_count"], 4)
self.assertTrue(scope["scope_complete"])
self.assertEqual(scope["excluded_unmeasured"], [])
def test_convert_count_reports_dumped_but_unindexed_files(self):
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
self._write_cell(root, [("s", "measured", "w4s3.ttgir", 140, 2)])
stray = root / "ttgir" / "s" / "stray" / "w4s3.ttgir"
stray.parent.mkdir(parents=True)
stray.write_text("ttg.convert_layout\n", encoding="utf-8")
scope = RUNNER._convert_count(root)
# An unindexed dump is counted rather than dropped, and flagged.
self.assertEqual(scope["measured_convert_count"], 3)
self.assertFalse(scope["scope_complete"])
self.assertEqual(scope["unindexed_files"], ["s/stray/w4s3.ttgir"])
def test_incomplete_convert_scope_fails_closed(self):
def cell(convert_count, latency, complete=True):
return {
"status": "PASS",
"convert_count": convert_count,
"convert_scope_complete": complete,
"median_gems_latency_ms": latency,
"replay_fallback_count": 0,
}
cells = {
"a_off": cell(8, 0.14),
"a_mask0": cell(8, 0.14),
"a_candidate": cell(0, 0.12, complete=False),
"b_candidate": cell(0, 0.121),
"b_off": cell(8, 0.141),
}
classification, detail = RUNNER.classify_task({"status": "PASS"}, cells)
self.assertEqual(classification, "X")
self.assertEqual(
detail["reason"], "measured convert topology scope is incomplete"
)
def test_correctness_command_narrows_to_the_measured_variant(self):
python = Path("/usr/bin/python3")
plain = RUNNER.correctness_command(
{"correctness_file": "tests/test_sort.py"}, python
)
self.assertEqual(plain[-3:], ["tests/test_sort.py", "-rs", "--quick"])
self.assertNotIn("-k", plain)
selected = RUNNER.correctness_command(
{
"correctness_file": "tests/test_upsample_bicubic2d_aa.py",
"correctness_selector": "not backward",
},
python,
)
self.assertIn("tests/test_upsample_bicubic2d_aa.py", selected)
self.assertEqual(selected[selected.index("-k") + 1], "not backward")
pinned = RUNNER.correctness_command(
{
"correctness_file": "tests/test_mm.py",
"correctness_node_ids": ["tests/test_mm.py::test_mm"],
},
python,
)
self.assertIn("tests/test_mm.py::test_mm", pinned)
self.assertNotIn("tests/test_mm.py", pinned)
def test_narrowed_correctness_requires_a_stated_reason(self):
manifest = copy.deepcopy(_manifest())
target = next(
row for row in manifest["workloads"] if row["id"] == "upsample_bicubic2d_aa"
)
self.assertEqual(target["correctness_selector"], "not backward")
target.pop("correctness_scope_reason")
errors = MODULE.validate_manifest(manifest)
self.assertIn(
"workload upsample_bicubic2d_aa narrows correctness and must state "
"correctness_scope_reason",
errors,
)
both = copy.deepcopy(_manifest())
row = next(item for item in both["workloads"] if item["id"] == "sort")
row["correctness_selector"] = "sort"
row["correctness_node_ids"] = ["tests/test_sort.py::test_sort"]
row["correctness_scope_reason"] = "reason"
self.assertIn(
"workload sort must not set both correctness_selector and "
"correctness_node_ids",
MODULE.validate_manifest(both),
)
def test_measured_launch_floor_is_a_frozen_protocol_field(self):
manifest = copy.deepcopy(_manifest())
self.assertEqual(manifest["protocol"]["measured_launch_floor"], 1)
plan, _ = MODULE.build_plan(manifest, "all")
self.assertTrue(
all(task["measured_launch_floor"] == 1 for task in plan["tasks"])
)
scoped = next(
task for task in plan["tasks"] if task["task_id"] == "fused_marlin_moe_mxfp4"
)
self.assertEqual(scoped["correctness_selector"], "mxfp4")
self.assertTrue(scoped["correctness_scope_reason"])
manifest["protocol"]["measured_launch_floor"] = -1
self.assertIn(
"protocol.measured_launch_floor must be a non-negative integer",
MODULE.validate_manifest(manifest),
)
def test_compat_plugin_only_overrides_verified_stale_guards(self):
def mark(reason, name="skipif"):
return SimpleNamespace(name=name, args=(True,), kwargs={"reason": reason})
def item(file_name, test_name):
return SimpleNamespace(path=Path("tests") / file_name, name=test_name)
# Verified on S5000: unskipped, these pass.
self.assertTrue(
COMPAT._is_stale_cuda_guard(
item("test_top_k_per_row_decode.py", "test_accuracy_top_k"),
mark("CUDA device required"),
)
)
self.assertTrue(
COMPAT._is_stale_fft_cuda_guard(
item("test_fft.py", "test_perf_fft"), mark("CUDA is not available")
)
)
# The fft guard stays pinned to the one test it was verified against.
self.assertFalse(
COMPAT._is_stale_cuda_guard(
item("test_fft.py", "test_other"), mark("CUDA is not available")
)
)
# Probed and rejected: grid_sample hardcodes device="cuda" in its own
# helper, and cutlass_scaled_mm gates on vLLM as well as CUDA.
self.assertFalse(
COMPAT._is_stale_cuda_guard(
item("test_grid_sample.py", "test_nearest_zeros_4d_small"),
mark("CUDA is required."),
)
)
self.assertFalse(
COMPAT._is_stale_cuda_guard(
item("test_cutlass_scaled_mm.py", "test_cutlass_scaled_mm"),
mark("Feature requires vLLM and NVIDIA Hopper architecture"),
)
)
# Only skipif marks, and only the exact reason string.
self.assertFalse(
COMPAT._is_stale_cuda_guard(
item("test_top_k_per_row_decode.py", "test_accuracy_top_k"),
mark("CUDA device required", name="xfail"),
)
)
self.assertFalse(
COMPAT._is_stale_cuda_guard(
item("test_top_k_per_row_decode.py", "test_accuracy_top_k"),
mark("some other reason"),
)
)
def test_compat_plugin_strips_module_level_guards(self):
# top_k_per_row_decode carries its guard as a module-level pytestmark,
# which resolves through the node chain rather than the item's own marks.
stale = SimpleNamespace(
name="skipif", args=(True,), kwargs={"reason": "CUDA device required"}
)
keep = SimpleNamespace(
name="skipif", args=(True,), kwargs={"reason": "vLLM not installed"}
)
module = SimpleNamespace(own_markers=[stale, keep])
function = SimpleNamespace(
path=Path("tests/test_top_k_per_row_decode.py"),
name="test_accuracy_top_k",
own_markers=[],
)
function.listchain = lambda: [module, function]
sys.modules.setdefault(
"flag_gems", SimpleNamespace(vendor_name="mthreads")
)
COMPAT.pytest_collection_modifyitems([function])
self.assertEqual(module.own_markers, [keep])
def test_benchmark_selector_requires_a_stated_reason(self):
manifest = copy.deepcopy(_manifest())
target = next(
row for row in manifest["workloads"] if row["id"] == "grid_sample"
)
self.assertEqual(target["benchmark_selector"], "dtype0")
target.pop("benchmark_scope_reason")
self.assertIn(
"workload grid_sample narrows the benchmark and must state "
"benchmark_scope_reason",
MODULE.validate_manifest(manifest),
)
def test_declared_benchmark_dtypes_must_exist_in_the_cli_vocabulary(self):
manifest = copy.deepcopy(_manifest())
self.assertEqual(
next(row for row in manifest["workloads"] if row["id"] == "randperm")[
"benchmark_dtypes"
],
["int32"],
)
# int64 is not a --dtypes choice; the sealed run only found that out on
# the device, as an argparse error inside a screen cell.
next(row for row in manifest["workloads"] if row["id"] == "randperm")[
"benchmark_dtypes"
] = ["int32", "int64"]
self.assertIn(
"workload randperm declares benchmark dtype(s) the benchmark CLI "
"does not accept: int64",
MODULE.validate_manifest(manifest),
)
def test_stale_inventory_needs_an_override_reason(self):
manifest = copy.deepcopy(_manifest())
marlin = next(
row
for row in manifest["workloads"]
if row["id"] == "fused_marlin_moe_mxfp4"
)
self.assertEqual(marlin["op_file"], "fused_marlin_moe_w4a16_mxfp4")
self.assertTrue(marlin["inventory_override_reason"])
with tempfile.TemporaryDirectory() as temporary:
inventory = Path(temporary)
(inventory / "inventory_main.csv").write_text(
"repo,op,op_file,file,marked_skip,collected,op_names\n"
+ "".join(
f"main,{row['op']},{row['op_file']},f.py,,yes,{row['id']}\n"
for row in manifest["workloads"]
if row["id"] != "fused_marlin_moe_mxfp4"
),
encoding="utf-8",
)
# The stated override lets the missing row through.
self.assertEqual(
MODULE.validate_inventory(manifest["workloads"], inventory), []
)
marlin.pop("inventory_override_reason")
self.assertEqual(
MODULE.validate_inventory(manifest["workloads"], inventory),
[
"inventory missing ('main', 'fused_marlin_moe_w4a16_mxfp4', "
"'fused_marlin_moe_w4a16_mxfp4')"
],
)
def test_missing_benchmark_target_fails_before_the_campaign_starts(self):
matrix = {
"tasks": [
{
"task_id": "present",
"op_file": "sort",
"correctness_file": "tests/test_sort.py",
},
{
# The shape of the sealed marlin defect: the inventory named
# a merged benchmark file that upstream had already split.
"task_id": "fused_marlin_moe_mxfp4",
"op_file": "fused_marlin_moe",
"correctness_file": "tests/test_fused_marlin_moe.py",
},
]
}
with tempfile.TemporaryDirectory() as temporary:
flaggems = Path(temporary)
(flaggems / "benchmark").mkdir()
(flaggems / "tests").mkdir()
(flaggems / "benchmark" / "test_sort.py").write_text("", encoding="utf-8")
(flaggems / "tests" / "test_sort.py").write_text("", encoding="utf-8")
(flaggems / "tests" / "test_fused_marlin_moe.py").write_text(
"", encoding="utf-8"
)
missing = RUNNER.missing_targets(matrix, flaggems)
self.assertEqual(len(missing), 1)
self.assertIn("fused_marlin_moe_mxfp4 benchmark", missing[0])
self.assertIn("test_fused_marlin_moe.py", missing[0])
def test_worker_names_the_failed_preflight_check(self):
worker = (CAMPAIGN / "run_pytest_musa_rlc_product.sh").read_text(encoding="utf-8")
self.assertIn("MUSA_RLC_PRODUCT_PREFLIGHT_FAILED", worker)
self.assertIn(
'preflight "benchmark file exists: '
'$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py"',
worker,
)
# The benchmark selector must reach pytest.
self.assertIn('selector_args+=(-k "$BENCH_SELECTOR")', worker)
self.assertIn('"${selector_args[@]}"', worker)
def test_correctness_runs_with_the_vendor_compat_plugin(self):
command = RUNNER.correctness_command(
{"correctness_file": "tests/test_top_k_per_row_decode.py"},
Path("/usr/bin/python3"),
)
self.assertEqual(
command[command.index("-p") + 1], "_musa_benchmark_compat_plugin"
)
def test_valid_manifest_and_plan_sizes(self):
manifest = _manifest()
self.assertEqual(MODULE.validate_manifest(manifest), [])
canary, canary_budget = MODULE.build_plan(manifest, "canary")
full, full_budget = MODULE.build_plan(manifest, "all")
self.assertEqual(len(canary["tasks"]), 8)
self.assertEqual(len(full["tasks"]), 32)
self.assertEqual(canary["tasks"][0]["task_id"], "dropout")
self.assertEqual(canary["tasks"][1]["task_id"], "softmax")
self.assertEqual(canary["tasks"][0]["benchmark_dtypes"], ["float16"])
self.assertEqual(
next(task for task in canary["tasks"] if task["task_id"] == "fft")[
"benchmark_dtypes"
],
["complex64"],
)
self.assertEqual(canary_budget["gpu_count"], 1)
self.assertEqual(full_budget["gpu_count"], 1)
self.assertTrue(canary_budget["serial_measurement"])
def test_invalid_device_and_canary_count_fail_closed(self):
manifest = copy.deepcopy(_manifest())
manifest["device"]["physical_device"] = "0"
manifest["workloads"][0]["canary"] = False
errors = MODULE.validate_manifest(manifest)
self.assertIn("device.physical_device must equal '7'", errors)
self.assertIn("exactly eight workloads must be canaries", errors)
def test_invalid_timer_and_default_policy_fail_closed(self):
manifest = copy.deepcopy(_manifest())
manifest["protocol"]["timer"] = "host wall clock"
manifest["production_default_policy"] = "enabled"
errors = MODULE.validate_manifest(manifest)
self.assertIn("protocol.timer must bind the MUSA native-event path", errors)
self.assertIn("production_default_policy must be unchanged-off", errors)
def test_runtime_classifier_positive_and_fail_closed_paths(self):
def cell(convert_count, latency, status="PASS", replay=0):
return {
"status": status,
"convert_count": convert_count,
"median_gems_latency_ms": latency,
"replay_fallback_count": replay,
}
correct = {"status": "PASS"}
cells = {
"a_off": cell(8, 0.14),
"a_mask0": cell(8, 0.14),
"a_candidate": cell(0, 0.12),
"b_candidate": cell(0, 0.121),
"b_off": cell(8, 0.141),
}
classification, _ = RUNNER.classify_task(correct, cells)
self.assertEqual(classification, "G")
invalid = copy.deepcopy(cells)
invalid["a_mask0"]["convert_count"] = 7
classification, detail = RUNNER.classify_task(correct, invalid)
self.assertEqual(classification, "X")
self.assertEqual(detail["reason"], "mask0 topology differs from off")
classification, _ = RUNNER.classify_task({"status": "FAIL"}, cells)
self.assertEqual(classification, "K")
classification, detail = RUNNER.classify_task(
{"status": "FAIL"},
cells,
{
"candidate": {"status": "FAIL"},
"mask0": {"status": "FAIL"},
"off": {"status": "FAIL"},
},
)
self.assertEqual(classification, "X")
self.assertEqual(
detail["reason"],
"baseline correctness fails under off/mask0/candidate",
)
def test_full_preparation_requires_passing_canary_receipt(self):
manifest = _manifest()
self.assertEqual(
MODULE.validate_canary_receipt(manifest, "all", None),
["full campaign requires --canary-receipt"],
)
with tempfile.TemporaryDirectory() as temporary:
receipt = Path(temporary) / "campaign-receipt.json"
receipt.write_text(
json.dumps(
{
"campaign_id": manifest["campaign_id"],
"selection": "canary",
"status": "TERMINAL",
"canary_gate_pass": True,
}
),
encoding="utf-8",
)
self.assertEqual(
MODULE.validate_canary_receipt(manifest, "all", receipt), []
)
if __name__ == "__main__":
unittest.main()