forked from zhoulin/zl_bench
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:
1 parent
6fa5f5d7a2
commit
ac818291a6
2 files changed
+71
-8
No files matched your search
+39
-7
@@ -173,7 +173,9 @@ class PluginPlanTests(unittest.TestCase):
|
||||
self.assertNotIn("_cudagraph_plugin", names)
|
||||
self.assertNotIn("_mm_cluster_fix_plugin", names)
|
||||
|
||||
def test_musa_same_core_without_graph_or_cluster(self):
|
||||
def test_musa_same_core_without_cluster_fix(self):
|
||||
# _cudagraph_plugin is absent from the plan on every backend: the timing
|
||||
# plugin loads it, so its absence here says nothing about capture support.
|
||||
names = self._plan("musa")
|
||||
self.assertIn("_timing_plugin", names)
|
||||
self.assertNotIn("_cudagraph_plugin", names)
|
||||
@@ -209,11 +211,17 @@ class PluginPlanTests(unittest.TestCase):
|
||||
class CaptureClassificationTests(unittest.TestCase):
|
||||
"""Graph-capture failures split three ways; the wording is vendor text."""
|
||||
|
||||
def test_cuda_and_musa_both_capture(self):
|
||||
self.assertTrue(get_backend("cuda").supports_cuda_graph)
|
||||
# torch_musa supplies Stream/Event/graph/MUSAGraph; measured on S5000.
|
||||
self.assertTrue(get_backend("musa").supports_cuda_graph)
|
||||
|
||||
def test_recoverable_is_the_default(self):
|
||||
# A kernel refusing capture is ordinary and must reach the fallback.
|
||||
v = get_backend("cuda").classify_capture_error(
|
||||
RuntimeError("operation not permitted when stream is capturing"))
|
||||
self.assertEqual(v, "recoverable")
|
||||
for name in ("cuda", "musa"):
|
||||
v = get_backend(name).classify_capture_error(
|
||||
RuntimeError("operation not permitted when stream is capturing"))
|
||||
self.assertEqual(v, "recoverable", name)
|
||||
|
||||
def test_cuda_fatal_markers(self):
|
||||
b = get_backend("cuda")
|
||||
@@ -231,9 +239,33 @@ class CaptureClassificationTests(unittest.TestCase):
|
||||
self.assertEqual(get_backend("cuda").classify_capture_error(exc),
|
||||
"recoverable")
|
||||
|
||||
def test_deny_list_defaults_to_empty(self):
|
||||
self.assertFalse(get_backend("cuda").denies_capture("anything"))
|
||||
self.assertFalse(get_backend("cuda").denies_capture(""))
|
||||
def test_musa_unknown_error_is_poisoning_not_recoverable(self):
|
||||
b = get_backend("musa")
|
||||
exc = RuntimeError("MUSA error: unknown error")
|
||||
self.assertEqual(b.classify_capture_error(exc), "poisoned")
|
||||
# The same text carries no meaning for CUDA, so it must not leak across.
|
||||
self.assertEqual(get_backend("cuda").classify_capture_error(exc),
|
||||
"recoverable")
|
||||
|
||||
def test_musa_rng_seed_failure_stays_recoverable(self):
|
||||
exc = RuntimeError("Cannot call MUSAGeneratorImpl::current_seed during "
|
||||
"MUSA graph capture")
|
||||
self.assertEqual(get_backend("musa").classify_capture_error(exc),
|
||||
"recoverable")
|
||||
|
||||
def test_deny_list_is_per_backend(self):
|
||||
self.assertTrue(get_backend("musa").denies_capture("unique2"))
|
||||
self.assertFalse(get_backend("musa").denies_capture("softmax"))
|
||||
self.assertFalse(get_backend("cuda").denies_capture("unique2"))
|
||||
self.assertFalse(get_backend("musa").denies_capture(""))
|
||||
|
||||
def test_deny_list_is_keyed_by_pytest_function_name(self):
|
||||
# FLAGGEMS_PERF_CURRENT_OP holds the pytest function name minus `test_`,
|
||||
# so a deny entry written as the benchmark FILE name never matches and
|
||||
# fails open. Observed once: "unique" listed, op=unique2 at runtime.
|
||||
musa = get_backend("musa")
|
||||
self.assertTrue(musa.denies_capture("unique2"))
|
||||
self.assertFalse(musa.denies_capture("unique"))
|
||||
|
||||
|
||||
class DeviceArgTests(unittest.TestCase):
|
||||
|
||||
Reference in new issue
Block a user