forked from zhoulin/zl_bench
Aharrypotter
4329b0c835
Both the _timing_plugin docstring and the README env table still described the behaviour from before MUSA declared supports_cuda_graph: "other backends do not attempt CUDA Graph at all", "非 CUDA 即使设 1 也拒绝". Capture selection has been driven by the backend capability, not by "is this CUDA", since the MUSA fallback policy landed. Reading the old text is most likely why the S5000 campaigns were launched with ZL_ENABLE_CUDAGRAPH=0 without anyone checking what auto would do. Also states what timer=cuda_graph does and does not mean, since the two are easy to conflate: it says capture was attempted. Measured on S5000 with auto and configs pinned, dropout/rand/masked_select select the graph path and then fall back on every single shape — BENCHMARK_DIRECT_NO_CUDAGRAPH on all 15 rows, zero poisoned, zero rows lost — so the whole table is direct-timed under a cuda_graph label. ZL_ENABLE_CUDAGRAPH=0 selects the same path and records it plainly, which is why the existing campaign data stands as measured.
71 lines
3.2 KiB
Python
71 lines
3.2 KiB
Python
"""Select the timing path from backend capabilities, before any shape runs.
|
|
|
|
Whether capture is attempted comes from the backend's `supports_cuda_graph`,
|
|
not from "is this CUDA". CUDA and MUSA both declare it, so both get graph
|
|
capture with per-shape fallback to plain do_bench; a backend that declares
|
|
nothing leaves `triton.testing.do_bench` in place and records `timer=direct`.
|
|
|
|
`timer=cuda_graph` therefore means capture was *attempted*, not that any row
|
|
was captured. On S5000 every shape of dropout/rand/masked_select falls back
|
|
(BENCHMARK_DIRECT_NO_CUDAGRAPH on every row, nothing poisoned, no rows lost),
|
|
so the whole table is direct-timed under a cuda_graph label. Setting
|
|
ZL_ENABLE_CUDAGRAPH=0 there selects the same timing path and says so plainly.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
|
|
from backends import get_backend
|
|
from _run_identity import update_identity
|
|
|
|
|
|
def _graph_requested(backend) -> bool:
|
|
raw = os.environ.get("ZL_ENABLE_CUDAGRAPH", "auto").strip().lower()
|
|
if raw in ("0", "false", "no", "off", "never", "direct"):
|
|
return False
|
|
if raw in ("1", "true", "yes", "on", "always"):
|
|
return bool(backend.supports_cuda_graph)
|
|
return bool(backend.supports_cuda_graph)
|
|
|
|
|
|
def pytest_configure(config):
|
|
from _term_style import tag
|
|
backend = get_backend()
|
|
raw = os.environ.get("ZL_ENABLE_CUDAGRAPH", "auto").strip().lower()
|
|
force_on = raw in ("1", "true", "yes", "on", "always")
|
|
if force_on and not backend.supports_cuda_graph:
|
|
print(f"{tag('[timing-plugin]')} ZL_ENABLE_CUDAGRAPH=1 ignored on "
|
|
f"backend={backend.name} (no CUDA Graph); using timer=direct",
|
|
file=sys.stderr, flush=True)
|
|
if _graph_requested(backend):
|
|
import _cudagraph_plugin as _cg
|
|
_cg.pytest_configure(config)
|
|
# Recorded on this branch too: the fallback rows a captured run produces
|
|
# are timed by the direct path and take the same repeat count.
|
|
update_identity(timer="cuda_graph", graph_capture="attempted",
|
|
direct_repeat=_cg.direct_repeat())
|
|
print(f"{tag('[timing-plugin]')} timer=cuda_graph backend={backend.name}",
|
|
file=sys.stderr, flush=True)
|
|
return
|
|
# A wholly-direct run gets the same repetition the fallback rows get, so the
|
|
# "put everything on one timer" advice does not silently cost the samples.
|
|
# Patching only when repeat > 1 keeps the default path untouched: upstream's
|
|
# do_bench stays in place exactly as before.
|
|
import _cudagraph_plugin as _cg
|
|
repeat = _cg.direct_repeat()
|
|
if repeat > 1:
|
|
def _repeated(fn, warmup=25, rep=100, grad_to_none=None,
|
|
quantiles=None, return_mode="mean"):
|
|
return _cg.run_direct(fn, warmup, rep, grad_to_none,
|
|
quantiles, return_mode)
|
|
import triton.testing as _tt
|
|
_tt.do_bench = _repeated
|
|
update_identity(timer="direct", graph_capture="disabled",
|
|
direct_repeat=repeat)
|
|
print(f"{tag('[timing-plugin]')} timer=direct backend={backend.name} "
|
|
f"(upstream do_bench; graph capture not selected)"
|
|
+ (f"; median of {repeat} calls" if repeat > 1 else ""),
|
|
file=sys.stderr, flush=True)
|