From 0bda0b168ad5bf7311ae639904489a4698a05546 Mon Sep 17 00:00:00 2001 From: hujianmin <76199145+beyondHJM@users.noreply.github.com> Date: Thu, 20 Aug 2026 14:32:31 +0800 Subject: [PATCH] [Fix]: exclude SM120 from attn-res TMA dispatch (#35361) Co-authored-by: 1BIN4 <1741738350@qq.com> Co-authored-by: L-Ark Co-authored-by: Chikati Co-authored-by: mengzili --- python/sglang/srt/layers/attn_residual.py | 15 +++++++++----- .../unit/layers/test_attn_residual.py | 20 +++++++++++++++++++ 2 files changed, 30 insertions(+), 5 deletions(-) create mode 100644 test/registered/unit/layers/test_attn_residual.py diff --git a/python/sglang/srt/layers/attn_residual.py b/python/sglang/srt/layers/attn_residual.py index 2e3cc94e4..d6f00b589 100644 --- a/python/sglang/srt/layers/attn_residual.py +++ b/python/sglang/srt/layers/attn_residual.py @@ -8,7 +8,7 @@ # fast — warp-specialized TMA kernel: cp.async.bulk producer + # online-softmax consumers over a double-buffered chunk ring, out # norm fused, per-nvb tuned launch config, one persistent CTA per -# SM. Taken on SM100+ with H=7168. +# SM. Taken on SM100+ except SM12x with H=7168. # hip — single Triton kernel, everything in one launch; taken on ROCm # within its register budget. # fused — Triton 2-kernel pipeline with full H-parallelism; the fallback @@ -33,15 +33,20 @@ _FAST_SUPPORTED = None _HIP_SHAPE_GATE = None +def _supports_attn_res_tma(capability: tuple[int, int]) -> bool: + """Return whether the device is eligible for the TMA fast path.""" + major, _ = capability + return major >= 10 and major != 12 + + def _use_fast(hidden_size: int) -> bool: - """The TMA kernel needs SM100+ (tcgen05, cp.async.bulk) and its H=7168 - template instantiation; everything else takes the triton pipeline.""" + """The TMA kernel needs SM100+ except SM12x (tcgen05, cp.async.bulk) + and its H=7168 template; everything else takes the triton pipeline.""" global _FAST_SUPPORTED if is_npu(): return False if _FAST_SUPPORTED is None: - major, _ = torch.cuda.get_device_capability() - _FAST_SUPPORTED = major >= 10 + _FAST_SUPPORTED = _supports_attn_res_tma(torch.cuda.get_device_capability()) return _FAST_SUPPORTED and hidden_size == 7168 diff --git a/test/registered/unit/layers/test_attn_residual.py b/test/registered/unit/layers/test_attn_residual.py new file mode 100644 index 000000000..79d85b682 --- /dev/null +++ b/test/registered/unit/layers/test_attn_residual.py @@ -0,0 +1,20 @@ +import unittest + +from sglang.srt.layers.attn_residual import _supports_attn_res_tma +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestAttnResidual(unittest.TestCase): + def test_tma_capability_gate(self): + self.assertFalse(_supports_attn_res_tma((9, 0))) + self.assertTrue(_supports_attn_res_tma((10, 0))) + self.assertTrue(_supports_attn_res_tma((10, 3))) + self.assertTrue(_supports_attn_res_tma((11, 0))) + self.assertFalse(_supports_attn_res_tma((12, 0))) + self.assertTrue(_supports_attn_res_tma((13, 0))) + + +if __name__ == "__main__": + unittest.main()