diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py index 3c36909a1..6e24068a1 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 8cfcec83e..da0af746f 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/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 ( - _use_aiter_gfx95 - and self.current_attention_backend - not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS - ) + return _use_aiter_gfx95 and self.current_attention_backend == "aiter" # Fuses the absorb BMM (`q_nope @ w_kc`) with `unified_attention_with_output` diff --git a/test/registered/unit/models/test_deepseek_mla_dispatch.py b/test/registered/unit/models/test_deepseek_mla_dispatch.py new file mode 100644 index 000000000..3296ae870 --- /dev/null +++ b/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()