[Refactor] Separate ROCm-specific DeepSeek MHA and MLA forward paths (#31531)

test/registered/kernels/test_dcp_lse_combine.py

@@ -267,11 +267,11 @@ class TestCPUReference(CustomTestCase):
 
     def test_flashmla_selects_natural_log_lse(self):
         from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla import (
-            _is_mla_dcp_lse_base_on_e,
+            is_mla_dcp_lse_base_on_e,
         )
 
-        self.assertTrue(_is_mla_dcp_lse_base_on_e("flashmla"))
-        self.assertFalse(_is_mla_dcp_lse_base_on_e("flashinfer_mla"))
+        self.assertTrue(is_mla_dcp_lse_base_on_e("flashmla"))
+        self.assertFalse(is_mla_dcp_lse_base_on_e("flashinfer_mla"))
 
     def test_nan_lse_handled(self):
         from sglang.kernels.ops.attention.dcp_kernels import _lse_weighted_combine_cpu

test/registered/unit/models/test_deepseek_mla_dispatch.py

@@ -62,5 +62,38 @@ class TestDispatchMLASubtype(CustomTestCase):
         self.assertEqual(method, AttnForwardMethod.MLA)
 
 
+class TestResolveRocmForwardMethod(CustomTestCase):
+    """The generic MHA/MLA methods must never reach the CUDA forward paths on
+    ROCm: those were stripped of their AMD branches when the AITER kernels moved
+    into forward_mha_rocm.py / forward_mla_rocm.py."""
+
+    def test_hip_routes_shared_methods_to_rocm(self):
+        with mock.patch.object(abh, "_is_hip", True):
+            self.assertEqual(
+                abh.resolve_rocm_forward_method(AttnForwardMethod.MHA),
+                AttnForwardMethod.MHA_ROCM,
+            )
+            self.assertEqual(
+                abh.resolve_rocm_forward_method(AttnForwardMethod.MHA_ONE_SHOT),
+                AttnForwardMethod.MHA_ONE_SHOT_ROCM,
+            )
+            self.assertEqual(
+                abh.resolve_rocm_forward_method(AttnForwardMethod.MLA),
+                AttnForwardMethod.MLA_ROCM,
+            )
+
+    def test_hip_leaves_platform_specific_methods_alone(self):
+        with mock.patch.object(abh, "_is_hip", True):
+            self.assertEqual(
+                abh.resolve_rocm_forward_method(AttnForwardMethod.MLA_FUSED_ROPE_ROCM),
+                AttnForwardMethod.MLA_FUSED_ROPE_ROCM,
+            )
+
+    def test_non_hip_is_identity(self):
+        with mock.patch.object(abh, "_is_hip", False):
+            for method in AttnForwardMethod:
+                self.assertEqual(abh.resolve_rocm_forward_method(method), method)
+
+
 if __name__ == "__main__":
     unittest.main()