From ac818291a634946d0acb7de23eeee590d515f8ee Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Wed, 9 Sep 2026 17:56:56 +0800 Subject: [PATCH] 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 --- backends/musa.py | 33 ++++++++++++++++++++++++++++++++- test_backends.py | 46 +++++++++++++++++++++++++++++++++++++++------- 2 files changed, 71 insertions(+), 8 deletions(-) diff --git a/backends/musa.py b/backends/musa.py index 21ea4f5..18f0add 100644 --- a/backends/musa.py +++ b/backends/musa.py @@ -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 diff --git a/test_backends.py b/test_backends.py index e105d2d..c7b600a 100644 --- a/test_backends.py +++ b/test_backends.py @@ -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):