forked from zhoulin/zl_bench
feat: make the harness backend-agnostic
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>
This commit is contained in:
1 parent
35a928a516
commit
c8ea586d0f
24 files changed
+1448
-159
No files matched your search
+15
-2
@@ -1,7 +1,9 @@
|
||||
"""pytest plugin: swap `triton.testing.do_bench` for cudagraph-based timing.
|
||||
|
||||
Loaded via `-p _cudagraph_plugin` in run_pytest.sh's PLUGINS list (comment that
|
||||
line out to time with plain do_bench). Monkey-patches do_bench →
|
||||
Invoked by `_timing_plugin` when the active backend declares
|
||||
`supports_cuda_graph` (CUDA only). Do not put this module on the default `-p`
|
||||
list for MUSA/ROCm/Ascend: those backends must start on plain do_bench, not
|
||||
attempt capture-then-fallback. Monkey-patches do_bench →
|
||||
do_bench_cudagraph (kernel-only latency) at session configure, after Triton is
|
||||
loaded but before any benchmark calls it. Kernels that can't be graph-captured
|
||||
fall back to plain do_bench and print a BENCHMARK_DIRECT_NO_CUDAGRAPH marker
|
||||
@@ -167,6 +169,17 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
|
||||
|
||||
def pytest_configure(config):
|
||||
from _term_style import tag
|
||||
try:
|
||||
from backends import get_backend
|
||||
if not get_backend().supports_cuda_graph:
|
||||
print(f"{tag('[cudagraph-plugin]')} skipped: backend "
|
||||
f"{get_backend().name!r} does not support CUDA Graph",
|
||||
file=sys.stderr, flush=True)
|
||||
return
|
||||
except Exception as exc:
|
||||
print(f"{tag('[cudagraph-plugin]')} backend probe failed ({exc}); "
|
||||
"not patching do_bench", file=sys.stderr, flush=True)
|
||||
return
|
||||
_tt.do_bench = _patched_do_bench
|
||||
print(f"{tag('[cudagraph-plugin]')} triton.testing.do_bench → do_bench_cudagraph",
|
||||
file=sys.stderr, flush=True)
|
||||
Reference in new issue
Block a user