[Fix]: exclude SM120 from attn-res TMA dispatch (#35361)

Co-authored-by: 1BIN4 <1741738350@qq.com>
Co-authored-by: L-Ark <fliangae@connect.ust.hk>
Co-authored-by: Chikati <jxudn@connect.ust.hk>
Co-authored-by: mengzili <zilim@ust.hk>

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
+    """The TMA kernel needs SM100+ except SM12x (tcgen05, cp.async.bulk)
-    template instantiation; everything else takes the triton pipeline."""
+    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 = _supports_attn_res_tma(torch.cuda.get_device_capability())
-        _FAST_SUPPORTED = major >= 10
     return _FAST_SUPPORTED and hidden_size == 7168
 
 

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()