The cudagraph plugin swapped do_bench for do_bench_cudagraph, but upstream's
do_bench_cudagraph names torch.cuda directly: torch.cuda.stream, Event,
synchronize, CUDAGraph, graph. FlagTree's do_bench already resolves the device
through driver.active.get_device_interface(), its cudagraph twin does not. On a
non-CUDA backend every capture therefore raised and every measurement silently
fell back to plain do_bench, so the plugin looked enabled while nothing was
being captured.
The plugin now resolves the interface itself and takes the stream, event,
synchronize and graph calls from it. CUDA still calls upstream's
do_bench_cudagraph verbatim, so its numbers do not move; other backends get an
otherwise step-for-step port. The scratch-RNG guard around capture is skipped
when the interface exposes no default_generators, which is the case on MThreads
and where there is no generator registration to protect against.
The replay timings are handed to _summarize_statistics as a plain list of
floats, matching do_bench and upstream's cudagraph path. This fork's helper runs
Python's statistics module over them, so wrapping them in a tensor makes "mean"
raise and makes a quantile request return tensors that later fail to serialize.
Verified on a Moore Threads S5000 through torch_musa, which supplies Stream,
Event, synchronize, graph and MUSAGraph. softmax and gelu_backward capture and
replay with no fallback marker. RNG operators cannot be captured there yet:
dropout aborts with "Cannot call MUSAGeneratorImpl::current_seed during MUSA
graph capture" and takes the documented BENCHMARK_DIRECT_NO_CUDAGRAPH fallback,
so a corpus mixing RNG and non-RNG operators mixes two timing regimes and the
marker is the thing to grep before comparing them.
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.
- _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.
Aharrypotter
zhoulin
