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>
Two failure modes that cost this harness real results, both of them silent.
A wrong path does not fail as a wrong path. pytest cannot collect a benchmark
directory that is not there, so benchmark/conftest.py never registers --level
and --mode, and the run dies with "unrecognized arguments: --level core --mode
kernel" -- an error about the flags, several layers away from the missing
directory. run_pytest.sh now checks the FlagGems directory, the benchmark file,
SHAPE_FILE and REPLAY_FROM before starting pytest, prints
"PREFLIGHT_FAILED check=<the named check>" and exits 3, distinct from pytest's
own 1 and 4 so the batch layer can tell a broken environment from a broken
operator. Only the file half of a node is checked; a missing ::test_<fn> already
produces a clear "not found: ..." from pytest and needs no help.
FlagGems keeps libtuner winners in a sqlite DB under $FLAGGEMS_CACHE_DIR
(default ~/.flaggems), shared by every run on the machine. Whichever side of an
A/B runs first therefore hands its winners to the second, which then compiles
only those: measured on router_gemm, the second of two shared-cache runs
compiled 10 variants instead of 234 and finished in 17s instead of 52s. That is
how an mm A/B in this repo produced "945 variants -> 33, 136 converts -> 10" and
read as a compiler result when it was cache inheritance; isolated, both runs
compile 234.
ZL_ISOLATE_TUNER_CACHE=1 gives a run its own cache directory, dropped afterwards
with the Triton one since the chosen configs are already in autotune_records/.
It defaults to 0, because cold tuning is expensive and the standard REPLAY_FROM
A/B pins configs anyway; only record-against-record is exposed. The default is
unchanged, but the caliber is no longer invisible: either way it goes into the
run header and run_identity.json as tuner_cache, which is what was missing when
the mm numbers could not be explained after the fact.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The sharded autotune sweep could publish a config set with holes, and a run
using it reports itself as replaying while every uncovered key silently
self-tunes. That is the exact failure the record/replay caliber exists to
prevent, and it was reachable two ways.
Shape lookup was op_name-only, but BLAS operators keep their shapes under the
class name: core_shapes.yaml has BlasBenchmark, not mm. Every BLAS operator
therefore resolved to zero shapes and the plugin disabled itself with a message
that read like a normal skip. It now walks the benchmark class chain the way
FlagGems' set_shapes does. The chain is recovered statically -- the class comes
from the op_name= keyword of the constructor call in the test module, its bases
from that module plus benchmark/base.py -- because the plugin runs at
pytest_configure, before any Benchmark instance exists. Resolved keys are named
in the log, and a total miss lists the candidates it tried.
Two cases can still produce partial coverage, and both now refuse rather than
shard a subset:
- --level comprehensive, where set_shapes() merges set_more_shapes() on top of
the yaml. That addition is arbitrary per-subclass Python and invisible here.
Measured on mm: 5 of 29 shapes swept, 90 AUTOTUNE_REPLAY_FALLBACK
reason=key_missing in the run that consumed the result.
- a shard exiting non-zero, whose shapes were never swept. This used to append a
WARNING to a success message and publish the merge anyway; observed with 3 of
5 shards dead. The whole sweep is now discarded, worker logs kept.
Both fall back to the normal serial sweep, which costs wall time rather than the
same-config premise.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
run_pytest.sh hardcoded --level core, so the only shapes ever reachable were
each operator's core_shapes entry. FlagGems merges a subclass's
set_more_shapes() on top of that at --level comprehensive, and for BLAS that
addition is where the large GEMMs live: mm goes from 5 shapes to 29, taking in
four large-K tuples plus models_shapes/*.yaml, with N/K up to 152064. The
comprehensive path also yields each shape a second time (transposed, for BLAS),
so the case count doubles. None of that was reachable without editing the driver.
BENCH_LEVEL defaults to core, so existing runs are unchanged, and rejects
anything other than core/comprehensive rather than passing it through to pytest.
The chosen level is echoed in the run header, written into the archived
shapes.yaml, and recorded in run_identity.json, because an A/B whose two sides
used different levels is not comparable and is otherwise invisible afterwards.
Sweep cost scales with the shape count, so comprehensive on a BLAS operator
wants PARALLEL_WARMUP_GPUS.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Everything about which accelerator a run targets was previously spelled CUDA:
device discovery went through nvidia-smi, pinning through CUDA_VISIBLE_DEVICES,
the vendor guard only recognised NVIDIA, and the -p plugin list was hardcoded in
run_pytest.sh with no way to drop the CUDA-only members. Running the same
operator corpus on another vendor's device therefore meant editing the drivers.
The selection now hangs off one ZL_BACKEND variable (cuda/musa/rocm/ascend, with
the usual aliases):
- backends/ owns per-backend device enumeration, visibility env keys and
capability flags, so nothing above it names a vendor;
- plugin_plan.py derives the -p list from those capabilities instead of a
hardcoded array, which is what lets a non-CUDA run omit the graph-capture and
cluster-fix plugins rather than watch them fail;
- _timing_plugin selects the timer from the same capabilities: CUDA keeps
graph capture with its per-shape fallback, other backends stay on upstream
do_bench and record timer=direct, because capture-then-fallback on a device
without graph support only produces noise;
- _device_setup_plugin replaces the NVIDIA-only vendor guard;
- providers/flaggems.py resolves the pytest node, so a caller may pass TEST_NODE
directly instead of the OP/OP_FILE pair, and SHAPE_MODE=upstream can defer to
FlagGems' own core_shapes;
- _run_identity.py writes run_identity.json next to each run, recording backend,
device, timer, autotune mode and the resolved test node. Two runs whose
calibers differ are otherwise indistinguishable after the fact.
run_batch gains --backend and --devices (--gpus kept as an alias) and carries
backend/device/timer into ops_status.csv. test_backends.py covers the registry,
alias resolution, visibility env and plugin planning.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
run_batch.py drives run_pytest.sh per operator across GPUs (stable-hash
sharding, per-op subprocess isolation, process-group timeouts, retry with
deterministic-failure cutoff, two-level dtype fallback, .complete resume,
per-op REPLAY_FROM). batch_summary.py aggregates run.log tables into
summary.csv. ops/ holds the curated assets: dual-repo inventories rebuilt
via AST scan + pytest collect verification, shape sets migrated from the
old regression harness and merged with upstream core_shapes class-name
keys (upstream's set_shapes falls back op_name -> MRO class name ->
1-D DEFAULT_SHAPES, so replacing the shape file without class keys
crashes the BLAS family), and a dismiss list where all 76 entries carry
verified reasons. Validated end to end: 1036-op full screen with zero
failures.
Also fix two cudagraph plugin regressions: newer torch appends "enable
device-side assertions" to every CUDA error, so the loose fatal-error
marker disabled the documented do_bench fallback entirely; and an aborted
graph capture can leave the default CUDA RNG generator stuck in capturing
state, poisoning every later torch.randn - captures now run under a
throwaway RNG state. run_pytest.sh gains an optional DTYPES passthrough.
Upstream moved from single /workspace/FlagGems-dev to /workspace/dev/FlagGems
(flag_gems) plus /workspace/dev/FlagGems-vllm (flaggems_vllm). Patch LibTuner
per actually-imported package at collection finish instead of hardcoding
flag_gems, and update the default FLAGGEMS_DIR.
Also stop dropping pre_hook configs for libtuner kernels: hopper mm's TMA
configs carry a pre_hook even with USE_FLAGTUNE=0, so record silently skipped
them and replay fell back. Upstream LibTuner.run now re-attaches the hook by
kwargs match on cache read (same path as its own ConfigCache round-trip), so
pre_hook-less injection is safe there; plain triton Autotuner keeps the drop.
For shape-heavy ops the autotune sweep dominates wall time, not the
measurement. _parallel_warmup_plugin shards the sweep across N GPUs and
hands the merged configs to the normal single-GPU serial measurement, so
timing stays comparable. Activated by PARALLEL_WARMUP_GPUS=N (silent
no-op when unset, so it can stay in the plugin list); output layout,
REPLAY_FROM usage and the latency table are unchanged.
Measurement is deliberately not parallelized: N processes saturating one
box couple through the power/thermal budget, so per-card latency gets
dragged by its neighbours by an amount that does not reproduce. Config
keys whose winner differs across shards are counted and reported as a
WARNING -- that count is how much to trust the run.
Fixes found while auditing:
- Shards no longer overwrite CUDA_VISIBLE_DEVICES with a bare shard
index, which a caller who had selected idle cards (e.g. 6,7) would see
re-interpreted as absolute ids 0,1 -- silently benchmarking on the busy
cards they were avoiding.
- Interrupting the sweep no longer leaves N subprocesses holding GPUs;
children are terminated and reaped before the exception propagates
(BaseException, since KeyboardInterrupt is not an Exception).
- run_pytest.sh claims its output dir with a bare mkdir and retreats to
a -2/-3 suffix on collision. The second-resolution timestamp meant two
concurrent runs shared one directory and overwrote each other.
- Replay eviction failures now emit a distinct compile_*_no_evict marker.
LibTuner.cache is a sqlite-backed ConfigCache with no __delitem__, so
the bad config could not be dropped and the retry re-read it, while the
log still claimed a clean fallback to live autotune.
Also trims comment density in both shell scripts and reworks the README:
promotes the parallel and cudagraph sections out from under the A/B flow,
groups the env table by purpose, documents that two record-mode runs are
not comparable, and marks ab_fold_test.sh as a caliber example whose
switch upstream has already removed.
LibTuner.cache is a view onto flag_gems' persistent sqlite config DB
(~/.flaggems/config_cache/TunedConfig_*.db), which survives across runs.
Both record and replay assumed a cold, in-process cache:
- record captured the chosen config by diffing cache keys before/after
run(). On a warm DB the key is already present, LibTuner.run takes the
cached branch without writing it again, so the diff was always empty
and nothing was recorded for any libtuner kernel. Only @triton.autotune
kernels (in-process cache) made it into the json -- e.g. a
fused_marlin_moe_mxfp4 run recorded moe_sum_kernel alone, missing both
MXFP4 GEMMs.
- replay only injected when the key was absent from the cache, so a warm
DB skipped injection entirely: the run reported "replaying N entries"
while actually self-tuning.
Record now reads back this call's own self.cache[key] after run();
replay overwrites unconditionally. Verified on fused_marlin_moe_mxfp4:
recorded entries 1 -> 6 (both GEMMs present), replay injects all 6 with
zero AUTOTUNE_REPLAY_FALLBACK and reproduces latency.
README: note that "fresh tune per side" requires dropping the sqlite DB
(not merely omitting REPLAY_FROM), and how to verify record coverage.
Plain autotune is the common case (fast, same-config A/B); FlagTune's
exhaustive first-run search is now opt-in via USE_FLAGTUNE=1.
ab_fold_test.sh automates the FLAGGEMS_MXFP4_FOLDSCALE=0/1 record-replay
comparison over the four DeepSeek-V4-Flash trace shape sets, which are
now tracked under runs/shape/.
- _cudagraph_plugin: warm up autotune/JIT explicitly before graph capture
(internal 5-iter warmup is too short), tag fallback markers with the
failing phase; first-run latency no longer jitters run-to-run.
- _device_guard_plugin (new): set GEMS_VENDOR via torch probe before
importing flag_gems, avoiding its timeout-less nvidia-smi subprocess
probe that can hang import in fork-broken environments.
- _pretty_report_plugin (new) + _term_style (new): fold inputs identical
across all result rows into a legend line, color status/plugin
tags/markers on the live terminal; run.log is ANSI-stripped and keeps
upstream SUCCESS/column wording for grep compatibility.
- run_pytest.sh: make USE_FLAGTUNE overridable, group all knobs into a
config section with Chinese comments, add start/end banners.
- README: document the warmup semantics, new plugins/env vars, and the
A/B rule that both sides must use the same USE_FLAGTUNE.
AharrypotterandClaude Opus 5
zhoulin