[AMD] [Fix] Fix --attention-backend triton work for DeepSeek MLA on MI355 (null-K + decode dispatch + RoPE) (#30355)

python/sglang/srt/models/deepseek_common/attention_backend_handler.py

@@ -31,7 +31,11 @@ class AttentionBackendRegistry:
 
 def _dispatch_mla_subtype(attn, forward_batch):
     if _is_hip:
-        if attn.rocm_fused_decode_mla and forward_batch.forward_mode.is_decode():
+        if (
+            attn.rocm_fused_decode_mla
+            and forward_batch.forward_mode.is_decode()
+            and attn.current_attention_backend == "aiter"
+        ):
             return AttnForwardMethod.MLA_FUSED_ROPE_ROCM
         else:
             return AttnForwardMethod.MLA

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -552,7 +552,12 @@ class DeepseekMLAForwardMixin:
             and (not fuse_rope_for_trtllm_mla)
             and (not skip_rope_for_dsa_tilelang_fused)
             and (not skip_rope_for_aiter_fused_mla)
-            and (not _use_aiter or not _is_gfx95_supported or self.use_dsa)
+            and (
+                not _use_aiter
+                or not _is_gfx95_supported
+                or self.use_dsa
+                or self.current_attention_backend == "triton"
+            )
         ):
             q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
 
@@ -758,7 +763,7 @@ class DeepseekMLAForwardMixin:
                         ),
                     )
         else:
-            if _use_aiter_gfx95:
+            if _use_aiter_gfx95 and self.current_attention_backend == "aiter":
                 cos = self.rotary_emb.cos_cache
                 sin = self.rotary_emb.sin_cache
 
@@ -1038,11 +1043,7 @@ class DeepseekMLAForwardMixin:
         when running aiter-backend MLA on gfx95 (i.e., the `else` branch in forward_absorb_core
         that calls fused_qk_rope_cat_and_cache_mla).
         """
-        return (
+        return _use_aiter_gfx95 and self.current_attention_backend == "aiter"
-            _use_aiter_gfx95
-            and self.current_attention_backend
-            not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS
-        )
 
 
 # Fuses the absorb BMM (`q_nope @ w_kc`) with `unified_attention_with_output`

test/registered/unit/models/test_deepseek_mla_dispatch.py

@@ -0,0 +1,67 @@
+"""Hermetic unit tests for DeepSeek MLA attention-method dispatch on ROCm.
+
+`_dispatch_mla_subtype` picks the forward method for MLA attention. On ROCm the
+fused-decode-MLA + fused-RoPE fast path (`MLA_FUSED_ROPE_ROCM`) is only correct
+for the aiter attention backend; taking it under the triton backend GPU-faults
+on gfx95 (MI355). This test pins the dispatch table so the triton MLA path stays
+on the plain `MLA` method.
+
+Pure Python (no GPU, no model weights): `_is_hip` is patched and `attn` /
+`forward_batch` are lightweight fakes. Runs on any PR-CI lane.
+"""
+
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+from sglang.srt.models.deepseek_common import attention_backend_handler as abh
+from sglang.srt.models.deepseek_common.attention_forward_methods.forward_methods import (
+    AttnForwardMethod,
+)
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
+register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+
+def _fake_forward_batch(is_decode: bool):
+    return SimpleNamespace(forward_mode=SimpleNamespace(is_decode=lambda: is_decode))
+
+
+def _fake_attn(backend: str, rocm_fused_decode_mla: bool = True):
+    return SimpleNamespace(
+        current_attention_backend=backend,
+        rocm_fused_decode_mla=rocm_fused_decode_mla,
+    )
+
+
+class TestDispatchMLASubtype(CustomTestCase):
+    def test_hip_aiter_decode_takes_fused_rope(self):
+        # aiter + fused-decode + decode -> fused ROPE fast path (unchanged).
+        with mock.patch.object(abh, "_is_hip", True):
+            method = abh._dispatch_mla_subtype(
+                _fake_attn("aiter"), _fake_forward_batch(is_decode=True)
+            )
+        self.assertEqual(method, AttnForwardMethod.MLA_FUSED_ROPE_ROCM)
+
+    def test_hip_triton_decode_stays_plain_mla(self):
+        # The fix: triton backend must NOT take the aiter-only fused path even
+        # with rocm_fused_decode_mla set -- that path GPU-faults on gfx95.
+        with mock.patch.object(abh, "_is_hip", True):
+            method = abh._dispatch_mla_subtype(
+                _fake_attn("triton"), _fake_forward_batch(is_decode=True)
+            )
+        self.assertEqual(method, AttnForwardMethod.MLA)
+
+    def test_hip_aiter_extend_stays_plain_mla(self):
+        # Fused path is decode-only; extend/prefill uses plain MLA.
+        with mock.patch.object(abh, "_is_hip", True):
+            method = abh._dispatch_mla_subtype(
+                _fake_attn("aiter"), _fake_forward_batch(is_decode=False)
+            )
+        self.assertEqual(method, AttnForwardMethod.MLA)
+
+
+if __name__ == "__main__":
+    unittest.main()