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.
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