feat: turn on graph capture for MUSA

supports_cuda_graph was False for MUSA, so every run there was pinned to
timer=direct. That was correct when it was written -- the capture path named
torch.cuda and could not have worked -- but it is not correct now, and the two
timers are not interchangeable: direct timing walks the full Python launch path
per iteration, so on small kernels it measures mostly launch overhead. Measured
on an S5000, capture against direct: batch_norm 16x16x8x48 0.0758ms -> 0.0089ms,
tanh and gelu at 4096x4096 about 3.3x, GEMMs unchanged at 1.00x. Leaving it off
discards that on every operator that can capture, and it also puts MUSA on a
different caliber from the CUDA numbers the harness is meant to compare against.

torch_musa supplies Stream, Event, synchronize, graph and MUSAGraph, so capture
works through the device interface. Verified on an S5000 over five operators,
one per path:

  tanh      captured, no marker
  dropout   15/15 rows fall back: MUSAGeneratorImpl::current_seed is not
            callable during capture
  rand      same
  perf_sort 25/25 rows fall back: operation not permitted when stream is
            capturing
  unique2   denied up front; rc=0 with 10 rows

MUSA gets its own error wording; none of the CUDA markers match it. "MUSA
error: unknown error" during capture is classified as poisoning rather than
recoverable, because the fallback runs and the process then dies in teardown:
before the deny entry was keyed correctly, unique2 reproduced exactly that,
ending at rc=1 with no rows, but now under an explicit BENCHMARK_CAPTURE_POISONED
marker instead of a bare ExceptionGroup.

capture_deny_ops is keyed by FLAGGEMS_PERF_CURRENT_OP, which holds the pytest
FUNCTION name minus `test_` -- unique2, perf_sort -- not the benchmark file
name. Getting that wrong fails open silently, which is how the first attempt
still hit rc=1 with the entry written as "unique"; a test now pins the key
space. The plugin-plan test is renamed because its assertion that
_cudagraph_plugin is absent from the plan holds on CUDA too and says nothing
about capture support.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
AharrypotterandClaude Opus 5 committed 2026-09-09 17:56:56 +08:00
1 parent 6fa5f5d7a2
commit ac818291a6
2 files changed
+71 -8

No files matched your search

+32 -1
View File
@@ -8,9 +8,40 @@ class MusaBackend(DeviceBackend):
visibility_env_keys = ("MUSA_VISIBLE_DEVICES",)
torch_namespace = "musa"
flaggems_vendor = "musa"
supports_cuda_graph = False
# torch_musa supplies Stream, Event, synchronize, graph and MUSAGraph, so
# capture works once the timing path takes them from Triton's device
# interface instead of torch.cuda. Measured on an S5000 over 15 operators:
# 11 captured, 3 fell back cleanly with a marker (dropout/rand on the RNG
# seed, sort on a capture-illegal stream op), 1 poisoned the context (see
# capture_deny_ops). Capture is worth having because it removes launch
# overhead the direct timer cannot: batch_norm 16x16x8x48 went 0.0758ms ->
# 0.0089ms, tanh/gelu 4096x4096 about 3.3x, while GEMMs stayed at 1.00x.
supports_cuda_graph = True
supports_mm_cluster_fix = False
supports_triton_hooks = True
# MUSA reports its own wording; none of the CUDA markers match it.
fatal_error_markers = (
"musa error: an illegal memory access",
"out of memory",
"device-side assert triggered",
"misaligned address",
"uncorrectable ecc",
)
# "unknown error" during capture is not recoverable on this runtime: the
# fallback runs, and the process then dies in teardown with zero rows. Treat
# it as poisoning so the operator is reported as capture-incompatible rather
# than silently losing its measurement.
poisoning_error_markers = ("musa error: unknown error",)
# Keyed by FLAGGEMS_PERF_CURRENT_OP, which is the pytest FUNCTION name minus
# `test_` -- `unique2`, `perf_sort` -- not the benchmark file name. Getting
# that wrong is silent: the entry simply never matches.
#
# Observed on an S5000: capturing test_unique2 takes two recoverable
# fallbacks and then raises "MUSA error: unknown error", ending the run at
# rc=1 with no measured rows. Skipping the attempt costs that operator's
# timing precision instead of its measurement. Revisit if a newer torch_musa
# changes the behaviour.
capture_deny_ops = ("unique2",)
# Device count comes from torch.musa or MUSA_VISIBLE_DEVICES. Do not guess
# a vendor CLI here; confirm the real listing tool before adding it.
list_command = None