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

python/sglang/srt/batch_overlap/operations.py

@@ -75,7 +75,7 @@ def _resolve_tbo_child_contexts():
     """Return (child_ctx_a, child_ctx_b) derived from the active TboAttnBackend,
     or (None, None) if the active backend is not a TBO dispatcher (e.g. a
     backend that handles TBO splitting internally like DeepSeek MHA's
-    _resolve_attn_backend path)."""
+    resolve_attn_backend path)."""
     # Lazy import to avoid circular dependency at module load time.
     from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
 

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

@@ -16,6 +16,23 @@ from sglang.srt.utils import is_sm100_or_sm110_supported, use_intel_amx_backend
 
 MHA_ONE_SHOT_SUPPORTED_BACKENDS = ["fa3", "flashinfer", "flashmla"]
 
+# ROCm runs dedicated MHA/MLA implementations (forward_mha_rocm.py /
+# forward_mla_rocm.py) so the shared CUDA paths carry no AMD branches. Backend
+# handlers keep returning the generic method; the platform swap happens here.
+# MHA_CHUNKED_KV has no ROCm entry because its accumulation step needs the
+# CUDA-only merge_state_v2 kernel.
+_ROCM_FORWARD_METHODS = {
+    AttnForwardMethod.MHA: AttnForwardMethod.MHA_ROCM,
+    AttnForwardMethod.MHA_ONE_SHOT: AttnForwardMethod.MHA_ONE_SHOT_ROCM,
+    AttnForwardMethod.MLA: AttnForwardMethod.MLA_ROCM,
+}
+
+
+def resolve_rocm_forward_method(method: AttnForwardMethod) -> AttnForwardMethod:
+    if not _is_hip:
+        return method
+    return _ROCM_FORWARD_METHODS.get(method, method)
+
 
 class AttentionBackendRegistry:
     _handlers = {}

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

@@ -1,13 +1,17 @@
 from .forward_methods import AttnForwardMethod
 from .forward_mha import DeepseekMHAForwardMixin
+from .forward_mha_rocm import DeepseekMHARocmForwardMixin
 from .forward_mla import DeepseekMLAForwardMixin
 from .forward_mla_fused_rope_cpu import DeepseekMLACpuForwardMixin
-from .forward_mla_fused_rope_rocm import DeepseekMLARocmForwardMixin
+from .forward_mla_fused_rope_rocm import DeepseekMLAFusedRopeRocmForwardMixin
+from .forward_mla_rocm import DeepseekMLARocmForwardMixin
 
 __all__ = [
     "AttnForwardMethod",
     "DeepseekMHAForwardMixin",
+    "DeepseekMHARocmForwardMixin",
     "DeepseekMLACpuForwardMixin",
     "DeepseekMLAForwardMixin",
+    "DeepseekMLAFusedRopeRocmForwardMixin",
     "DeepseekMLARocmForwardMixin",
 ]

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

@@ -30,3 +30,12 @@ class AttnForwardMethod(IntEnum):
 
     # Use Deepseek V3.2 sparse multi-latent attention for NPU
     DSA_NPU = auto()
+
+    # Use multi-head attention for ROCm
+    MHA_ROCM = auto()
+
+    # Use one-shot multi-head attention for ROCm
+    MHA_ONE_SHOT_ROCM = auto()
+
+    # Use absorbed multi-latent attention for ROCm
+    MLA_ROCM = auto()

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

@@ -14,9 +14,6 @@ from sglang.srt.layers.dcp import (
     all_gather_kv_cache_for_mha_extend,
     filter_dcp_local_kv_indices,
 )
-from sglang.srt.layers.quantization.fp8_utils import (
-    materialize_bpreshuffle_fp8_scale_tuple,
-)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.forward_context import (
     get_attn_backend,
@@ -24,22 +21,15 @@ from sglang.srt.model_executor.forward_context import (
 )
 from sglang.srt.models.deepseek_common.utils import (
     _is_cuda,
-    _is_hip,
     _is_musa,
     _is_npu,
-    _use_aiter_bpreshuffle_gfx95,
-    _use_aiter_gfx95,
 )
 from sglang.srt.runtime_context import (
     get_exec,
     get_parallel,
     get_schedule,
 )
-from sglang.srt.utils import BumpAllocator, get_bool_env_var, next_power_of_2
+from sglang.srt.utils import BumpAllocator, next_power_of_2
-
-_use_fp8_prefill_attn = (
-    get_bool_env_var("SGLANG_AITER_FP8_PREFILL_ATTN", "True") and _use_aiter_gfx95
-)
 
 if TYPE_CHECKING:
     from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
@@ -51,21 +41,15 @@ if _is_cuda:
 elif _is_musa:
     from sgl_kernel import concat_mla_k
 
-if _use_aiter_gfx95:
-    from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
 
-    from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype
+def resolve_attn_backend(forward_batch: ForwardBatch):
-    from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant
-
-
-def _resolve_attn_backend(forward_batch: ForwardBatch):
     backend = get_attn_backend()
     if isinstance(backend, TboAttnBackend):
         backend = backend.primary
     return backend
 
 
-def _forward_dsa_indexer_for_mha(
+def forward_dsa_indexer_for_mha(
     indexer,
     *,
     hidden_states: torch.Tensor,
@@ -173,39 +157,12 @@ class DeepseekMHAForwardMixin:
             # DSA Indexer: cache quantized keys, auto-skip topk for sequences <= dsa_index_topk
 
             if self.use_dsa:
-                # DSA requires unquantized q_lora for the indexer. When q_b_proj is FP8
-                # on gfx95, we can still use fused RMSNorm+FP8 quant, but MUST request
-                # the unquantized output for q_lora; otherwise q_lora becomes the (fp8,scale)
-                # tuple.
-                if (
-                    _use_aiter_gfx95
-                    and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
-                ):
-                    q_quanted, q_lora, _, _ = fused_rms_fp8_group_quant(
-                        q,
-                        self.q_a_layernorm.weight,
-                        self.q_a_layernorm.variance_epsilon,
-                        None,
-                        None,
-                        None,
-                        group_size=128,
-                        dtype_quant=torch.float8_e4m3fn,
-                        res1=None,
-                        output_unquantized_inp1=True,
-                        transpose_scale=False,
-                    )
-                    if _use_aiter_bpreshuffle_gfx95:
-                        q_quanted = materialize_bpreshuffle_fp8_scale_tuple(q_quanted)
-                    q = self.q_b_proj(q_quanted)[0].view(
-                        -1, self.num_local_heads, self.qk_head_dim
-                    )
-                else:
                 q_lora = self.q_a_layernorm(q)
                 q = self.q_b_proj(q_lora)[0].view(
                     -1, self.num_local_heads, self.qk_head_dim
                 )
                 if self.should_run_indexer():
-                    _forward_dsa_indexer_for_mha(
+                    forward_dsa_indexer_for_mha(
                         self.indexer,
                         hidden_states=hidden_states,
                         q_lora=q_lora,
@@ -213,34 +170,6 @@ class DeepseekMHAForwardMixin:
                         forward_batch=forward_batch,
                         layer_id=self.layer_id,
                     )
-            elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
-                # MXFP4: fused RMSNorm + quant
-                q, _, _, _ = fused_rms_mxfp4_quant(
-                    q,
-                    self.q_a_layernorm.weight,
-                    self.q_a_layernorm.variance_epsilon,
-                    None,
-                    None,
-                    None,
-                )
-                q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
-            elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
-                q, _, _, _ = fused_rms_fp8_group_quant(
-                    q,
-                    self.q_a_layernorm.weight,
-                    self.q_a_layernorm.variance_epsilon,
-                    None,
-                    None,
-                    None,
-                    group_size=128,
-                    dtype_quant=torch.float8_e4m3fn,
-                    res1=None,
-                    output_unquantized_inp1=False,
-                    transpose_scale=False,
-                )
-                if _use_aiter_bpreshuffle_gfx95:
-                    q = materialize_bpreshuffle_fp8_scale_tuple(q)
-                q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
             else:
                 q = self.q_a_layernorm(q)
                 q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
@@ -255,24 +184,6 @@ class DeepseekMHAForwardMixin:
         kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
         latent_cache = latent_cache.unsqueeze(1)
 
-        if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
-            kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant(
-                kv_a,
-                self.kv_a_layernorm.weight,
-                self.kv_a_layernorm.variance_epsilon,
-                None,
-                None,
-                None,
-                group_size=128,
-                dtype_quant=torch.float8_e4m3fn,
-                res1=None,
-                output_unquantized_inp1=True,  # return unqaunt kv_a
-                transpose_scale=False,
-            )
-            if _use_aiter_bpreshuffle_gfx95:
-                kv_a_quanted = materialize_bpreshuffle_fp8_scale_tuple(kv_a_quanted)
-
-        else:
         kv_a = self.kv_a_layernorm(kv_a)
 
         k_pe = latent_cache[:, :, self.kv_lora_rank :]
@@ -281,7 +192,7 @@ class DeepseekMHAForwardMixin:
         # (fused RoPE + quantize for Q/K, direct FP8 KV-cache write) and
         # returns FP8 tensors ready for its kernel. Backends without the
         # hook fall through to the BF16 path below.
-        backend = _resolve_attn_backend(forward_batch)
+        backend = resolve_attn_backend(forward_batch)
         if hasattr(backend, "prepare_prefill_qkv"):
             q_out, k_out, v_out = backend.prepare_prefill_qkv(
                 q=q,
@@ -333,27 +244,8 @@ class DeepseekMHAForwardMixin:
                         q.dtype,
                         forward_batch,
                     )
-        if _use_fp8_prefill_attn and self.kv_b_proj.weight.dtype == torch.uint8:
-            # MXFP4 weights + FP8 prefill: fuse GEMM, nope/v split, and k_pe cat
-            # into a single kernel (fused_gemm_afp4wfp4_split_cat) that writes k and v
-            # directly in FP8, avoiding a separate elementwise cast
-            k, v = self.kv_b_proj(
-                (
-                    kv_a,
-                    k_pe.expand(-1, self.num_local_heads, -1),
-                    self.qk_nope_head_dim,
-                    self.v_head_dim,
-                    fp8_dtype,
-                )
-            )[0]
-        else:
-            if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
-                kv = self.kv_b_proj(kv_a_quanted)[0]
-            else:
         kv = self.kv_b_proj(kv_a)[0]
-            kv = kv.view(
+        kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
-                -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
-            )
         k_nope = kv[..., : self.qk_nope_head_dim]
         v = kv[..., self.qk_nope_head_dim :]
 
@@ -464,7 +356,7 @@ class DeepseekMHAForwardMixin:
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
         # kv_b_proj needs BF16 input, but legacy q.dtype was BF16 by accident.
-        backend = _resolve_attn_backend(forward_batch)
+        backend = resolve_attn_backend(forward_batch)
         pack_fn = getattr(backend, "pack_prefix_chunk_kv", None)
         kv_a_dtype = torch.bfloat16 if pack_fn is not None else q.dtype
 
@@ -531,7 +423,7 @@ class DeepseekMHAForwardMixin:
         k_pe: torch.Tensor,
         forward_batch: ForwardBatch,
     ):
-        if _is_cuda or _use_aiter_gfx95:
+        if _is_cuda:
             # Save latent cache
             get_token_to_kv_pool().set_mla_kv_buffer(
                 self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe
@@ -556,7 +448,7 @@ class DeepseekMHAForwardMixin:
         dst_dtype: torch.dtype,
         forward_batch: ForwardBatch,
     ):
-        if _is_cuda or _use_aiter_gfx95:
+        if _is_cuda:
             kv_indices = filter_dcp_local_kv_indices(kv_indices=kv_indices)
             kv_a, k_pe = get_token_to_kv_pool().get_mla_kv_buffer(
                 self.attn_mha, kv_indices, dst_dtype
@@ -632,9 +524,6 @@ class DeepseekMHAForwardMixin:
                 attn_dtype = k_nope.dtype
             k = k_nope.new_empty(*k_shape, dtype=attn_dtype)
             concat_and_cast_mha_k_triton(k, k_nope, k_pe)
-        elif _is_hip and self.current_attention_backend == "aiter":
-            k = k_nope.new_empty(*k_shape)
-            concat_and_cast_mha_k_triton(k, k_nope, k_pe)
         else:
             k = k_nope.new_empty(*k_shape)
             k[..., : self.qk_nope_head_dim] = k_nope

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

@@ -0,0 +1,324 @@
+"""AMD/ROCm multi-head attention forward path for DeepSeek models.
+
+`AttnForwardMethod.MHA_ROCM` and `AttnForwardMethod.MHA_ONE_SHOT_ROCM` route
+here, which keeps every AITER/gfx95 kernel choice out of the shared
+`forward_mha.py`. Only the prepare step is platform specific; the attention
+cores in `DeepseekMHAForwardMixin` are reused as-is.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.ops.attention.utils import concat_and_cast_mha_k_triton
+from sglang.srt.layers.communicator import get_attn_tp_context
+from sglang.srt.layers.dcp import (
+    all_gather_kv_cache_for_mha_extend,
+    filter_dcp_local_kv_indices,
+)
+from sglang.srt.layers.quantization.fp8_utils import (
+    materialize_bpreshuffle_fp8_scale_tuple,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mha import (
+    forward_dsa_indexer_for_mha,
+    resolve_attn_backend,
+)
+from sglang.srt.models.deepseek_common.utils import (
+    _use_aiter_bpreshuffle_gfx95,
+    _use_aiter_gfx95,
+)
+from sglang.srt.runtime_context import get_exec, get_parallel
+from sglang.srt.utils import BumpAllocator, get_bool_env_var
+
+if TYPE_CHECKING:
+    from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
+
+_use_fp8_prefill_attn = (
+    get_bool_env_var("SGLANG_AITER_FP8_PREFILL_ATTN", "True") and _use_aiter_gfx95
+)
+
+if _use_aiter_gfx95:
+    from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
+
+    from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype
+    from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant
+
+
+class DeepseekMHARocmForwardMixin:
+
+    def forward_normal_rocm_prepare(
+        self: DeepseekV2AttentionMLA,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+        zero_allocator: BumpAllocator,
+    ):
+        if self.q_lora_rank is not None:
+            q, latent_cache = (
+                get_attn_tp_context()
+                .fetch_qkv_latent()
+                .split(
+                    [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
+                    dim=-1,
+                )
+            )
+
+            # DSA Indexer: cache quantized keys, auto-skip topk for sequences <= dsa_index_topk
+
+            if self.use_dsa:
+                # DSA requires unquantized q_lora for the indexer. When q_b_proj is FP8
+                # on gfx95, we can still use fused RMSNorm+FP8 quant, but MUST request
+                # the unquantized output for q_lora; otherwise q_lora becomes the (fp8,scale)
+                # tuple.
+                if (
+                    _use_aiter_gfx95
+                    and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
+                ):
+                    q_quanted, q_lora, _, _ = fused_rms_fp8_group_quant(
+                        q,
+                        self.q_a_layernorm.weight,
+                        self.q_a_layernorm.variance_epsilon,
+                        None,
+                        None,
+                        None,
+                        group_size=128,
+                        dtype_quant=torch.float8_e4m3fn,
+                        res1=None,
+                        output_unquantized_inp1=True,
+                        transpose_scale=False,
+                    )
+                    if _use_aiter_bpreshuffle_gfx95:
+                        q_quanted = materialize_bpreshuffle_fp8_scale_tuple(q_quanted)
+                    q = self.q_b_proj(q_quanted)[0].view(
+                        -1, self.num_local_heads, self.qk_head_dim
+                    )
+                else:
+                    q_lora = self.q_a_layernorm(q)
+                    q = self.q_b_proj(q_lora)[0].view(
+                        -1, self.num_local_heads, self.qk_head_dim
+                    )
+                if self.should_run_indexer():
+                    forward_dsa_indexer_for_mha(
+                        self.indexer,
+                        hidden_states=hidden_states,
+                        q_lora=q_lora,
+                        positions=positions,
+                        forward_batch=forward_batch,
+                        layer_id=self.layer_id,
+                    )
+            elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
+                # MXFP4: fused RMSNorm + quant
+                q, _, _, _ = fused_rms_mxfp4_quant(
+                    q,
+                    self.q_a_layernorm.weight,
+                    self.q_a_layernorm.variance_epsilon,
+                    None,
+                    None,
+                    None,
+                )
+                q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+            elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
+                q, _, _, _ = fused_rms_fp8_group_quant(
+                    q,
+                    self.q_a_layernorm.weight,
+                    self.q_a_layernorm.variance_epsilon,
+                    None,
+                    None,
+                    None,
+                    group_size=128,
+                    dtype_quant=torch.float8_e4m3fn,
+                    res1=None,
+                    output_unquantized_inp1=False,
+                    transpose_scale=False,
+                )
+                if _use_aiter_bpreshuffle_gfx95:
+                    q = materialize_bpreshuffle_fp8_scale_tuple(q)
+                q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+            else:
+                q = self.q_a_layernorm(q)
+                q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+
+        else:
+            q = self.q_proj(hidden_states)[0].view(
+                -1, self.num_local_heads, self.qk_head_dim
+            )
+            latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
+
+        _, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+        latent_cache = latent_cache.unsqueeze(1)
+
+        if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
+            kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant(
+                kv_a,
+                self.kv_a_layernorm.weight,
+                self.kv_a_layernorm.variance_epsilon,
+                None,
+                None,
+                None,
+                group_size=128,
+                dtype_quant=torch.float8_e4m3fn,
+                res1=None,
+                output_unquantized_inp1=True,  # return unqaunt kv_a
+                transpose_scale=False,
+            )
+            if _use_aiter_bpreshuffle_gfx95:
+                kv_a_quanted = materialize_bpreshuffle_fp8_scale_tuple(kv_a_quanted)
+        else:
+            kv_a = self.kv_a_layernorm(kv_a)
+
+        k_pe = latent_cache[:, :, self.kv_lora_rank :]
+
+        # Backend prefill hook: the backend owns the BF16->FP8 transition
+        # (fused RoPE + quantize for Q/K, direct FP8 KV-cache write) and
+        # returns FP8 tensors ready for its kernel. Backends without the
+        # hook fall through to the BF16 path below.
+        backend = resolve_attn_backend(forward_batch)
+        if hasattr(backend, "prepare_prefill_qkv"):
+            q_out, k_out, v_out = backend.prepare_prefill_qkv(
+                q=q,
+                q_pe=q_pe,
+                kv_a=kv_a,
+                k_pe=k_pe,
+                positions=positions,
+                layer=self,
+                forward_batch=forward_batch,
+            )
+            return q_out, k_out, v_out, forward_batch
+
+        if self.rotary_emb is not None:
+            q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
+        q[..., self.qk_nope_head_dim :] = q_pe
+
+        self._set_mla_kv_buffer_rocm(latent_cache, kv_a, k_pe, forward_batch)
+        if (
+            forward_batch.mha_one_shot
+            and sum(forward_batch.extend_prefix_lens_cpu) != 0
+        ):
+            if (
+                self.use_dsa
+                and self.kv_cache_dtype == "fp8_e4m3"
+                and (
+                    not get_exec().kernel.dsa_decode_backend == "trtllm"
+                    or not get_exec().kernel.dsa_prefill_backend == "trtllm"
+                )
+            ):
+                # FP8 path: dequantize DSA-specific FP8 format to BF16
+                kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_dsa(forward_batch)
+            else:
+                # BF16/FP16 path: directly fetch from cache
+                if get_parallel().dcp_enabled:
+                    kv_a, k_pe = all_gather_kv_cache_for_mha_extend(
+                        get_token_to_kv_pool(),
+                        self.attn_mha,
+                        forward_batch.attn_dcp_metadata.dcp_local_prefix_kv_indices,
+                        forward_batch.seq_lens,
+                        forward_batch.extend_prefix_lens,
+                        forward_batch.extend_prefix_lens_cpu,
+                        forward_batch.extend_seq_lens,
+                        kv_a,
+                        k_pe,
+                    )
+                else:
+                    kv_a, k_pe = self._get_mla_kv_buffer_rocm(
+                        forward_batch.fetch_mha_one_shot_kv_indices(),
+                        q.dtype,
+                        forward_batch,
+                    )
+        if _use_fp8_prefill_attn and self.kv_b_proj.weight.dtype == torch.uint8:
+            # MXFP4 weights + FP8 prefill: fuse GEMM, nope/v split, and k_pe cat
+            # into a single kernel (fused_gemm_afp4wfp4_split_cat) that writes k and v
+            # directly in FP8, avoiding a separate elementwise cast
+            k, v = self.kv_b_proj(
+                (
+                    kv_a,
+                    k_pe.expand(-1, self.num_local_heads, -1),
+                    self.qk_nope_head_dim,
+                    self.v_head_dim,
+                    fp8_dtype,
+                )
+            )[0]
+        else:
+            if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
+                kv = self.kv_b_proj(kv_a_quanted)[0]
+            else:
+                kv = self.kv_b_proj(kv_a)[0]
+            kv = kv.view(
+                -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
+            )
+            k_nope = kv[..., : self.qk_nope_head_dim]
+            v = kv[..., self.qk_nope_head_dim :]
+
+            k = self._concat_and_cast_mha_k_rocm(k_nope, k_pe)
+        return q, k, v, forward_batch
+
+    def forward_normal_one_shot_rocm_prepare(
+        self: DeepseekV2AttentionMLA,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+        zero_allocator: BumpAllocator,
+    ):
+        forward_batch.mha_one_shot = True
+        return self.forward_normal_rocm_prepare(
+            positions, hidden_states, forward_batch, zero_allocator
+        )
+
+    def _concat_and_cast_mha_k_rocm(
+        self: DeepseekV2AttentionMLA,
+        k_nope: torch.Tensor,
+        k_pe: torch.Tensor,
+    ):
+        k_shape = (k_nope.shape[0], self.num_local_heads, self.qk_head_dim)
+        k = k_nope.new_empty(*k_shape)
+        if self.current_attention_backend == "aiter":
+            concat_and_cast_mha_k_triton(k, k_nope, k_pe)
+        else:
+            k[..., : self.qk_nope_head_dim] = k_nope
+            k[..., self.qk_nope_head_dim :] = k_pe
+        return k
+
+    def _set_mla_kv_buffer_rocm(
+        self: DeepseekV2AttentionMLA,
+        latent_cache: torch.Tensor,
+        kv_a: torch.Tensor,
+        k_pe: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ):
+        if _use_aiter_gfx95:
+            get_token_to_kv_pool().set_mla_kv_buffer(
+                self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe
+            )
+        else:
+            latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1)
+            latent_cache[:, :, self.kv_lora_rank :] = k_pe.clone()
+            get_token_to_kv_pool().set_kv_buffer(
+                self.attn_mha, forward_batch.out_cache_loc, latent_cache, None
+            )
+
+    def _get_mla_kv_buffer_rocm(
+        self: DeepseekV2AttentionMLA,
+        kv_indices: torch.Tensor,
+        dst_dtype: torch.dtype,
+        forward_batch: ForwardBatch,
+    ):
+        if _use_aiter_gfx95:
+            kv_indices = filter_dcp_local_kv_indices(kv_indices=kv_indices)
+            kv_a, k_pe = get_token_to_kv_pool().get_mla_kv_buffer(
+                self.attn_mha, kv_indices, dst_dtype
+            )
+            kv_a = kv_a.squeeze(1)
+        else:
+            latent_cache_buf = get_token_to_kv_pool().get_key_buffer(
+                self.attn_mha.layer_id
+            )
+            latent_cache = latent_cache_buf[kv_indices].contiguous().to(dst_dtype)
+            kv_a, k_pe = latent_cache.split(
+                [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+            )
+            kv_a = kv_a.squeeze(1).contiguous()
+        return kv_a, k_pe

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

@@ -8,7 +8,6 @@ import torch
 
 from sglang.kernels.ops.kvcache.cache_ops import absorbed_bmm_concat_cast_q_fp8
 from sglang.kernels.ops.quantization.fp8_kernel import (
-    fp8_dtype,
     per_tensor_quant_mla_fp8,
     per_token_group_quant_mla_deep_gemm_masked_fp8,
 )
@@ -28,9 +27,6 @@ from sglang.srt.layers.dcp import (
     dcp_a2a_lse_reduce,
 )
 from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run
-from sglang.srt.layers.quantization.fp8_utils import (
-    materialize_bpreshuffle_fp8_scale_tuple,
-)
 from sglang.srt.layers.radix_attention import unified_attention_with_output
 from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
 from sglang.srt.lora.deepseek_mla_correction import (
@@ -62,12 +58,8 @@ from sglang.srt.models.deepseek_common.utils import (
     _is_cpu,
     _is_cublas_ge_129,
     _is_cuda,
-    _is_gfx95_supported,
     _is_hip,
     _is_musa,
-    _use_aiter,
-    _use_aiter_bpreshuffle_gfx95,
-    _use_aiter_gfx95,
 )
 from sglang.srt.runtime_context import get_exec, get_parallel
 from sglang.srt.state_capturer.indexer_topk import (
@@ -101,7 +93,16 @@ def _select_local_dcp_heads_for_autotune(
     return attn_output.narrow(1, rank * num_local_heads, num_local_heads)
 
 
-def _is_mla_dcp_lse_base_on_e(attention_backend: Optional[str]) -> bool:
+def is_dcp_mla_decode_phase(forward_batch: ForwardBatch) -> bool:
+    if not get_parallel().dcp_enabled:
+        return False
+    return (
+        forward_batch.forward_mode.is_decode()
+        or forward_batch.forward_mode.is_target_verify()
+    )
+
+
+def is_mla_dcp_lse_base_on_e(attention_backend: Optional[str]) -> bool:
     # FlashMLA exposes natural-log softmax LSE. FlashInfer MLA and the other
     # currently supported MLA DCP decode backends expose base-2 LSE.
     return attention_backend == "flashmla"
@@ -111,58 +112,7 @@ if _is_cuda:
     from sglang.kernels.ops.gemm import bmm_fp8
 
 
-if _use_aiter:
+def should_defer_dsa_cp_kv_gather(
-    # aiter ROCm/aiter#2958 renamed the public `fused_qk_rmsnorm` in
-    # `aiter.ops.fused_qk_norm_rope_cache_quant` to a private `_fused_qk_rmsnorm`
-    # and introduced a unified entry point in `aiter.ops.fused_qk_rmsnorm_group_quant`
-    # with a different (in-place, kwarg-only, no-return) signature. Probe for the
-    # new symbol first so SGLang works with both pre- and post-#2958 aiter without
-    # requiring the docker pin to be bumped atomically.
-    try:
-        from aiter.ops.enum import QuantType as _AiterQuantType
-        from aiter.ops.fused_qk_rmsnorm_group_quant import (
-            fused_qk_rmsnorm as _aiter_fused_qk_rmsnorm_unified,
-        )
-
-        def fused_qk_rmsnorm_bf16(q, q_weight, q_eps, k, k_weight, k_eps):
-            q_out = torch.empty_like(q)
-            k_out = torch.empty_like(k)
-            _aiter_fused_qk_rmsnorm_unified(
-                q_out_quantized=q_out,
-                k_out=k_out,
-                q=q,
-                q_weight=q_weight,
-                q_epsilon=q_eps,
-                k=k,
-                k_weight=k_weight,
-                k_epsilon=k_eps,
-                quant_type=_AiterQuantType.No,
-            )
-            return q_out, k_out
-
-    except ImportError:
-        from aiter.ops.fused_qk_norm_rope_cache_quant import (
-            fused_qk_rmsnorm as fused_qk_rmsnorm_bf16,
-        )
-
-    from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (
-        batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant,
-    )
-if _use_aiter_gfx95:
-    from aiter.ops.triton.fused_fp8_quant import (
-        fused_flatten_fp8_group_quant,
-        fused_rms_fp8_group_quant,
-    )
-
-    from sglang.srt.layers.quantization.rocm_mxfp4_utils import (
-        batched_gemm_afp4wfp4_pre_quant,
-        fused_flatten_mxfp4_quant,
-        fused_rms_mxfp4_quant,
-    )
-    from sglang.srt.layers.rocm_linear_utils import fused_qk_rope_cat_and_cache_mla
-
-
-def _should_defer_dsa_cp_kv_gather(
     *,
     dsa_prefill_cp: bool,
     fuse_rope_for_trtllm_mla: bool,
@@ -207,7 +157,7 @@ class DeepseekMLAForwardMixin:
             return False
         if not self.use_dsa:
             return False
-        if self.use_deep_gemm_bmm or _is_hip:
+        if self.use_deep_gemm_bmm:
             return False
         if is_kv_b_lora_active(self):
             return False
@@ -290,12 +240,6 @@ class DeepseekMLAForwardMixin:
             return None
         if self._fuse_rope_for_trtllm_mla(forward_batch):
             return None
-        if self._skip_rope_for_dsa_tilelang_fused():
-            return None
-        if self._skip_rope_for_aiter_fused_mla():
-            return None
-        if _use_aiter and _is_gfx95_supported and not self.use_dsa:
-            return None
         # Graph/compile surfaces run their own dispatch; the python-side
         # stash handshake is eager-only.
         if is_graph_dsa_split_op_surface(forward_batch):
@@ -356,11 +300,7 @@ class DeepseekMLAForwardMixin:
         # weights and skip the per-layer Q all-gather (bf16 decode absorb only).
         q_replicate_active = (
             get_parallel().dcp_replicate_q_proj
-            and get_parallel().dcp_enabled
+            and is_dcp_mla_decode_phase(forward_batch)
-            and (
-                forward_batch.forward_mode.is_decode()
-                or forward_batch.forward_mode.is_target_verify()
-            )
             and not self.use_deep_gemm_bmm
             and self.w_kc_qrep is not None
             and self.q_b_proj_qrep_weight is not None
@@ -393,67 +333,6 @@ class DeepseekMLAForwardMixin:
                 with torch.cuda.stream(self.alt_stream):
                     k_nope = self.kv_a_layernorm(k_nope)
                 current_stream.wait_stream(self.alt_stream)
-            else:
-                if _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
-                    q, _, k_nope, *_ = fused_rms_mxfp4_quant(
-                        q,
-                        self.q_a_layernorm.weight,
-                        self.q_a_layernorm.variance_epsilon,
-                        k_nope,
-                        self.kv_a_layernorm.weight,
-                        self.kv_a_layernorm.variance_epsilon,
-                    )
-                else:
-                    q_lora = None
-                    if (
-                        _use_aiter_gfx95
-                        and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
-                    ):
-                        if self.use_dsa:
-                            q_quanted, q_lora, k_nope, _ = fused_rms_fp8_group_quant(
-                                q,
-                                self.q_a_layernorm.weight,
-                                self.q_a_layernorm.variance_epsilon,
-                                k_nope,
-                                self.kv_a_layernorm.weight,
-                                self.kv_a_layernorm.variance_epsilon,
-                                group_size=128,
-                                dtype_quant=torch.float8_e4m3fn,
-                                res1=None,
-                                output_unquantized_inp1=True,
-                                transpose_scale=False,
-                            )
-                            if _use_aiter_bpreshuffle_gfx95:
-                                q_quanted = materialize_bpreshuffle_fp8_scale_tuple(
-                                    q_quanted
-                                )
-                            q = q_quanted
-                        else:
-                            q, _, k_nope, _ = fused_rms_fp8_group_quant(
-                                q,
-                                self.q_a_layernorm.weight,
-                                self.q_a_layernorm.variance_epsilon,
-                                k_nope,
-                                self.kv_a_layernorm.weight,
-                                self.kv_a_layernorm.variance_epsilon,
-                                group_size=128,
-                                dtype_quant=torch.float8_e4m3fn,
-                                res1=None,
-                                output_unquantized_inp1=False,
-                                transpose_scale=False,
-                            )
-                            if _use_aiter_bpreshuffle_gfx95:
-                                q = materialize_bpreshuffle_fp8_scale_tuple(q)
-
-                    elif _use_aiter:
-                        q, k_nope = fused_qk_rmsnorm_bf16(
-                            q,
-                            self.q_a_layernorm.weight,
-                            self.q_a_layernorm.variance_epsilon,
-                            k_nope,
-                            self.kv_a_layernorm.weight,
-                            self.kv_a_layernorm.variance_epsilon,
-                        )
             else:
                 q = self.q_a_layernorm(q)
                 k_nope = self.kv_a_layernorm(k_nope)
@@ -639,50 +518,6 @@ class DeepseekMLAForwardMixin:
                     expected_m,
                 )
                 q_nope_out = q_nope_out[:, :expected_m, :]
-            elif _is_hip:
-                # TODO(haishaw): add bmm_fp8 to ROCm
-                if _use_aiter_gfx95 and self.w_kc.dtype == torch.uint8:
-                    x = q_nope.transpose(0, 1)
-                    q_nope_out = torch.empty(
-                        x.shape[0],
-                        x.shape[1],
-                        self.w_kc.shape[2],
-                        device=x.device,
-                        dtype=torch.bfloat16,
-                    )
-                    batched_gemm_afp4wfp4_pre_quant(
-                        x,
-                        self.w_kc.transpose(-2, -1),
-                        self.w_scale_k.transpose(-2, -1),
-                        torch.bfloat16,
-                        q_nope_out,
-                    )
-                else:
-                    if (
-                        _use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn
-                    ) or (
-                        get_is_capture_mode()
-                        and self.w_kc.dtype == torch.float8_e4m3fnuz
-                    ):
-                        # fp8 Triton kernel: always on gfx950,
-                        # cudagraph-only on gfx942 (hides launch overhead)
-                        q_nope_out = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
-                            X=q_nope,
-                            WQ=self.w_kc.transpose(-1, -2),
-                            w_scale=self.w_scale,
-                            group_size=128,
-                            YQ=None,  # allocate (B, M, N)
-                            transpose_bm=False,  # (B, M, N)
-                            transpose_bm_in=True,  # (M, B, K)
-                            dtype=torch.bfloat16,
-                        )
-
-                    else:
-                        q_nope_out = torch.bmm(
-                            q_nope.to(torch.bfloat16).transpose(0, 1),
-                            self.w_kc.to(torch.bfloat16) * self.w_scale,
-                        )
-
             elif self.w_kc.dtype == torch.float8_e4m3fn:
                 if _is_cpu:
                     q_nope_out = torch.bmm(
@@ -720,19 +555,9 @@ class DeepseekMLAForwardMixin:
                 q_nope_out = apply_kv_b_lora_q_correction(self, q_nope, q_nope_out)
 
         fuse_rope_for_trtllm_mla = self._fuse_rope_for_trtllm_mla(forward_batch)
-        skip_rope_for_dsa_tilelang_fused = self._skip_rope_for_dsa_tilelang_fused()
-        skip_rope_for_aiter_fused_mla = self._skip_rope_for_aiter_fused_mla()
         if (
             self.rotary_emb is not None
-            and (not fuse_rope_for_trtllm_mla)
+            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
-                or self.current_attention_backend == "triton"
-            )
             # Already applied at the q-prep/indexer overlap fork.
             and not self._q8kv8_qprep_overlap_pending
         ):
@@ -770,7 +595,7 @@ class DeepseekMLAForwardMixin:
 
         dsa_prefill_cp = dsa_use_prefill_cp(forward_batch)
         mla_prefill_cp = mla_use_prefill_cp(forward_batch)
-        defer_kv_gather_until_after_rope = _should_defer_dsa_cp_kv_gather(
+        defer_kv_gather_until_after_rope = should_defer_dsa_cp_kv_gather(
             dsa_prefill_cp=dsa_prefill_cp,
             fuse_rope_for_trtllm_mla=fuse_rope_for_trtllm_mla,
         )
@@ -796,10 +621,7 @@ class DeepseekMLAForwardMixin:
 
         # all_gather q_pe, q_nope_out,take tp8 as an example， q_pe [B, H, ROPE_DIM], q_nope_out [B, H, NOPE_DIM] gathered to [B, H * dcp_world_size, ROPE_DIM] [B, H * dcp_world_size, NOPE_DIM] for decode batch, and all gather k_pe, k_nope for extend batch.
         if get_parallel().dcp_enabled:
-            if (
+            if is_dcp_mla_decode_phase(forward_batch):
-                forward_batch.forward_mode.is_decode()
-                or forward_batch.forward_mode.is_target_verify()
-            ):
                 if not q_replicate_active:
                     q_nope_out, q_pe = all_gather_q_for_mla_decode(
                         q_nope_out=q_nope_out,
@@ -852,76 +674,6 @@ class DeepseekMLAForwardMixin:
         save_kv_cache = True
 
         if self.current_attention_backend in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS:
-            if self._skip_rope_for_dsa_tilelang_fused() and self.rotary_emb is not None:
-                cos = self.rotary_emb.cos_cache
-                sin = self.rotary_emb.sin_cache
-                kv_cache_dtype = (
-                    fp8_dtype if self.kv_cache_dtype == "fp8_e4m3" else q_nope_out.dtype
-                )
-                q_cat, _, k_pe_fused, _ = fused_qk_rope_cat_and_cache_mla(
-                    q_nope_out,
-                    q_pe,
-                    k_nope,
-                    k_pe,
-                    get_token_to_kv_pool().get_key_buffer(self.attn_mqa.layer_id),
-                    forward_batch.out_cache_loc,
-                    positions,
-                    cos,
-                    sin,
-                    self.attn_mqa.k_scale,
-                    self.rotary_emb.is_neox_style,
-                    q_out_dtype=kv_cache_dtype,
-                )
-                save_kv_cache = False
-                # On decode, pass q_cat directly to attn_mqa with q_rope=None so
-                # dsa_backend.forward_decode reuses q_cat as a zero-copy view
-                # (`q.contiguous().view(...)` fast-path) instead of running the
-                # redundant `concat_mla_absorb_q_general(q_nope_fused, q_pe_fused)`
-                # that would otherwise rebuild a tensor byte-identical to q_cat.
-                # On ROCm tilelang decode, this eliminates the
-                # `CatArrayBatchedCopy<OpaqueType<1u>, ...>` kernel that used to
-                # fire once per layer per decode step (~2.6 us / layer saved).
-                # Prefill keeps the split form because dsa_backend.forward_extend
-                # asserts `q_rope is not None`.
-                if forward_batch.forward_mode.is_decode_or_idle():
-                    if llama_4_scaling is not None:
-                        # llama_4_scaling applies only to the q_nope portion;
-                        # mutate in place via the slice view of q_cat.
-                        q_cat[..., : self.kv_lora_rank] *= llama_4_scaling
-                    attn_output = self.attn_mqa(
-                        q_cat,
-                        None,
-                        None,
-                        forward_batch,
-                        q_rope=None,
-                        k_rope=k_pe_fused,
-                        save_kv_cache=save_kv_cache,
-                        **(
-                            dict(topk_indices=topk_indices)
-                            if topk_indices is not None
-                            else {}
-                        ),
-                    )
-                else:
-                    q_nope_fused = q_cat[..., : self.kv_lora_rank]
-                    q_pe_fused = q_cat[..., self.kv_lora_rank :]
-                    if llama_4_scaling is not None:
-                        q_nope_fused *= llama_4_scaling
-                    attn_output = self.attn_mqa(
-                        q_nope_fused,
-                        None,
-                        None,
-                        forward_batch,
-                        q_rope=q_pe_fused,
-                        k_rope=k_pe_fused,
-                        save_kv_cache=save_kv_cache,
-                        **(
-                            dict(topk_indices=topk_indices)
-                            if topk_indices is not None
-                            else {}
-                        ),
-                    )
-            else:
             extra_args = {}
             if self._fuse_rope_for_trtllm_mla(forward_batch):
                 extra_args = {
@@ -952,10 +704,7 @@ class DeepseekMLAForwardMixin:
                     topk_indices=topk_indices,
                 )
                 attn_output = fusion_plan.attn_output_buf
-                elif (
+            elif is_dcp_mla_decode_phase(forward_batch):
-                    forward_batch.forward_mode.is_decode()
-                    or forward_batch.forward_mode.is_target_verify()
-                ) and get_parallel().dcp_enabled:
                 # set return_lse=True to correct attn_output
                 attn_output, lse = self.attn_mqa_for_dcp_decode(
                     q_nope_out,
@@ -986,31 +735,6 @@ class DeepseekMLAForwardMixin:
                         else {}
                     ),
                 )
-        else:
-            if _use_aiter_gfx95 and self.current_attention_backend == "aiter":
-                cos = self.rotary_emb.cos_cache
-                sin = self.rotary_emb.sin_cache
-
-                kv_cache_dtype = (
-                    fp8_dtype if self.kv_cache_dtype == "fp8_e4m3" else q_nope_out.dtype
-                )
-
-                q, _, _, k = fused_qk_rope_cat_and_cache_mla(
-                    q_nope_out,
-                    q_pe,
-                    k_nope,
-                    k_pe,
-                    get_token_to_kv_pool().get_key_buffer(self.attn_mqa.layer_id),
-                    forward_batch.out_cache_loc,
-                    positions,
-                    cos,
-                    sin,
-                    self.attn_mqa.k_scale,
-                    self.rotary_emb.is_neox_style,
-                    q_out_dtype=kv_cache_dtype,
-                )
-
-                save_kv_cache = False
         else:
             q = torch.cat([q_nope_out, q_pe], dim=-1)
             k = torch.cat([k_nope, k_pe], dim=-1)
@@ -1029,10 +753,7 @@ class DeepseekMLAForwardMixin:
             )
 
         # correct attn_output with respect to lse from other ranks
-        if (
+        if is_dcp_mla_decode_phase(forward_batch):
-            forward_batch.forward_mode.is_decode()
-            or forward_batch.forward_mode.is_target_verify()
-        ) and get_parallel().dcp_enabled:
             attn_output = attn_output.view(
                 -1,
                 self.num_local_heads * get_parallel().attn_dcp_size,
@@ -1047,7 +768,7 @@ class DeepseekMLAForwardMixin:
                 )
             else:
                 dcp_comm_backend = get_parallel().dcp_comm_backend
-                is_lse_base_on_e = _is_mla_dcp_lse_base_on_e(
+                is_lse_base_on_e = is_mla_dcp_lse_base_on_e(
                     self.current_attention_backend
                 )
                 if dcp_comm_backend in ("a2a", "fi_a2a"):
@@ -1101,83 +822,6 @@ class DeepseekMLAForwardMixin:
             attn_bmm_output = (
                 attn_bmm_output[:, :expected_m, :].transpose(0, 1).flatten(1, 2)
             )
-        elif _is_hip:
-            # TODO(haishaw): add bmm_fp8 to ROCm
-            if _use_aiter_gfx95 and self.w_vc.dtype == torch.uint8:
-                x = attn_output.transpose(0, 1)
-                B_heads, M_batch = x.shape[0], x.shape[1]
-                N_vdim = self.w_vc.shape[2]
-                # Allocate in (batch, heads, dim) so the post-GEMM
-                # transpose+flatten is a free view instead of a copy.
-                _bmm_buf = torch.empty(
-                    M_batch,
-                    B_heads,
-                    N_vdim,
-                    device=x.device,
-                    dtype=torch.bfloat16,
-                )
-                attn_bmm_output = _bmm_buf.transpose(0, 1)
-                batched_gemm_afp4wfp4_pre_quant(
-                    x,
-                    self.w_vc.transpose(-2, -1),
-                    self.w_scale_v.transpose(-2, -1),
-                    torch.bfloat16,
-                    attn_bmm_output,
-                )
-            else:
-                _bmm_buf = None
-                if _use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn:
-                    attn_bmm_output = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
-                        X=attn_output,
-                        WQ=self.w_vc.transpose(-1, -2),
-                        w_scale=self.w_scale,
-                        group_size=128,
-                        YQ=None,
-                        transpose_bm=False,
-                        transpose_bm_in=True,
-                        dtype=torch.bfloat16,
-                    )
-                else:
-                    attn_bmm_output = torch.bmm(
-                        attn_output.to(torch.bfloat16).transpose(0, 1),
-                        self.w_vc.to(torch.bfloat16) * self.w_scale,
-                    )
-
-            if _bmm_buf is not None:
-                # _bmm_buf is already (batch, heads, dim) contiguous
-                if self.o_proj.weight.dtype == torch.uint8:
-                    attn_bmm_output = fused_flatten_mxfp4_quant(_bmm_buf)
-                elif self.o_proj.weight.dtype == torch.float8_e4m3fn:
-                    attn_bmm_output = fused_flatten_fp8_group_quant(
-                        _bmm_buf,
-                        group_size=128,
-                        dtype_quant=torch.float8_e4m3fn,
-                        transpose_scale=False,
-                    )
-                    if _use_aiter_bpreshuffle_gfx95:
-                        attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple(
-                            attn_bmm_output
-                        )
-                else:
-                    attn_bmm_output = _bmm_buf.flatten(1, 2)
-            elif self.o_proj.weight.dtype == torch.uint8:
-                attn_bmm_output = attn_bmm_output.transpose(0, 1)
-                attn_bmm_output = fused_flatten_mxfp4_quant(attn_bmm_output)
-            elif self.o_proj.weight.dtype == torch.float8_e4m3fn:
-                attn_bmm_output = attn_bmm_output.transpose(0, 1)
-                attn_bmm_output = fused_flatten_fp8_group_quant(
-                    attn_bmm_output,
-                    group_size=128,
-                    dtype_quant=torch.float8_e4m3fn,
-                    transpose_scale=False,
-                )
-                if _use_aiter_bpreshuffle_gfx95:
-                    attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple(
-                        attn_bmm_output
-                    )
-            else:
-                attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
-
         elif self.w_vc.dtype == torch.float8_e4m3fn:
             if _is_cpu:
                 attn_bmm_output = torch.bmm(
@@ -1276,27 +920,6 @@ class DeepseekMLAForwardMixin:
             and get_attn_backend().data_type == torch.float8_e4m3fn
         )
 
-    def _skip_rope_for_dsa_tilelang_fused(self: DeepseekV2AttentionMLA) -> bool:
-        """
-        Check if we should skip rope and use fused rope+cache path for TileLang DSA on gfx95.
-        """
-        return (
-            _use_aiter_gfx95
-            and self.current_attention_backend in ("dsa", "nsa")
-            and (
-                get_exec().kernel.dsa_decode_backend == "tilelang"
-                or get_exec().kernel.dsa_prefill_backend == "tilelang"
-            )
-        )
-
-    def _skip_rope_for_aiter_fused_mla(self: DeepseekV2AttentionMLA) -> bool:
-        """
-        Skip rope in prepare and let the fused kernel in forward_absorb_core handle it,
-        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 == "aiter"
-
 
 # Fuses the absorb BMM (`q_nope @ w_kc`) with `unified_attention_with_output`
 # into one eager split op under both PCG and BCG. Without this, the bf16

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

@@ -29,7 +29,7 @@ if _is_hip:
     )
 
 
-class DeepseekMLARocmForwardMixin:
+class DeepseekMLAFusedRopeRocmForwardMixin:
 
     def init_mla_fused_rope_rocm_forward(self: DeepseekV2AttentionMLA):
         self.rocm_fused_decode_mla = get_bool_env_var(

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

@@ -0,0 +1,849 @@
+"""AMD/ROCm absorbed multi-latent attention forward path for DeepSeek models.
+
+`AttnForwardMethod.MLA_ROCM` routes here, which keeps every AITER/gfx95 kernel
+choice out of the shared `forward_mla.py` and lets non-AMD builds avoid
+importing `aiter` altogether.
+
+The BMM absorb steps stay module level to keep ROCm kernel selection localized.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+
+from sglang.kernels.ops.quantization.fp8_kernel import (
+    fp8_dtype,
+    per_token_group_quant_mla_deep_gemm_masked_fp8,
+)
+from sglang.srt.environ import envs
+from sglang.srt.layers import deep_gemm_wrapper
+from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
+from sglang.srt.layers.communicator import get_attn_tp_context
+from sglang.srt.layers.cp.utils import is_cp_v2_active
+from sglang.srt.layers.dcp import (
+    all_gather_kv_cache_for_mla_extend,
+    all_gather_q_for_mla_decode,
+    cp_lse_ag_out_rs_mla,
+    dcp_a2a_lse_reduce,
+)
+from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run
+from sglang.srt.layers.quantization.fp8_utils import (
+    materialize_bpreshuffle_fp8_scale_tuple,
+)
+from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
+from sglang.srt.lora.deepseek_mla_correction import (
+    apply_q_correction as apply_kv_b_lora_q_correction,
+)
+from sglang.srt.lora.deepseek_mla_correction import (
+    apply_v_correction as apply_kv_b_lora_v_correction,
+)
+from sglang.srt.lora.deepseek_mla_correction import (
+    is_kv_b_lora_active,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla import (
+    _select_local_dcp_heads_for_autotune,
+    is_dcp_mla_decode_phase,
+    is_mla_dcp_lse_base_on_e,
+    should_defer_dsa_cp_kv_gather,
+)
+from sglang.srt.models.deepseek_common.utils import (
+    FORWARD_ABSORB_CORE_ATTENTION_BACKENDS,
+    _is_gfx95_supported,
+    _use_aiter,
+    _use_aiter_bpreshuffle_gfx95,
+    _use_aiter_gfx95,
+)
+from sglang.srt.runtime_context import get_exec, get_parallel
+from sglang.srt.state_capturer.indexer_topk import (
+    maybe_capture_indexer_topk,
+)
+from sglang.srt.utils import BumpAllocator
+
+logger = logging.getLogger(__name__)
+_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
+
+if TYPE_CHECKING:
+    from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
+
+if _use_aiter:
+    # aiter ROCm/aiter#2958 renamed the public `fused_qk_rmsnorm` in
+    # `aiter.ops.fused_qk_norm_rope_cache_quant` to a private `_fused_qk_rmsnorm`
+    # and introduced a unified entry point in `aiter.ops.fused_qk_rmsnorm_group_quant`
+    # with a different (in-place, kwarg-only, no-return) signature. Probe for the
+    # new symbol first so SGLang works with both pre- and post-#2958 aiter without
+    # requiring the docker pin to be bumped atomically.
+    try:
+        from aiter.ops.enum import QuantType as _AiterQuantType
+        from aiter.ops.fused_qk_rmsnorm_group_quant import (
+            fused_qk_rmsnorm as _aiter_fused_qk_rmsnorm_unified,
+        )
+
+        def fused_qk_rmsnorm_bf16(q, q_weight, q_eps, k, k_weight, k_eps):
+            q_out = torch.empty_like(q)
+            k_out = torch.empty_like(k)
+            _aiter_fused_qk_rmsnorm_unified(
+                q_out_quantized=q_out,
+                k_out=k_out,
+                q=q,
+                q_weight=q_weight,
+                q_epsilon=q_eps,
+                k=k,
+                k_weight=k_weight,
+                k_epsilon=k_eps,
+                quant_type=_AiterQuantType.No,
+            )
+            return q_out, k_out
+
+    except ImportError:
+        from aiter.ops.fused_qk_norm_rope_cache_quant import (
+            fused_qk_rmsnorm as fused_qk_rmsnorm_bf16,
+        )
+
+    from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (
+        batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant,
+    )
+
+if _use_aiter_gfx95:
+    from aiter.ops.triton.fused_fp8_quant import (
+        fused_flatten_fp8_group_quant,
+        fused_rms_fp8_group_quant,
+    )
+
+    from sglang.srt.layers.quantization.rocm_mxfp4_utils import (
+        batched_gemm_afp4wfp4_pre_quant,
+        fused_flatten_mxfp4_quant,
+        fused_rms_mxfp4_quant,
+    )
+    from sglang.srt.layers.rocm_linear_utils import fused_qk_rope_cat_and_cache_mla
+
+
+def rocm_absorb_q_bmm(
+    attn: DeepseekV2AttentionMLA,
+    q_nope: torch.Tensor,
+    *,
+    is_capture_mode: bool,
+) -> torch.Tensor:
+    """Absorb ``q_nope @ w_kc`` on HIP/AITER (pre-transpose layout)."""
+    # TODO(haishaw): add bmm_fp8 to ROCm
+    if _use_aiter_gfx95 and attn.w_kc.dtype == torch.uint8:
+        x = q_nope.transpose(0, 1)
+        q_nope_out = torch.empty(
+            x.shape[0],
+            x.shape[1],
+            attn.w_kc.shape[2],
+            device=x.device,
+            dtype=torch.bfloat16,
+        )
+        batched_gemm_afp4wfp4_pre_quant(
+            x,
+            attn.w_kc.transpose(-2, -1),
+            attn.w_scale_k.transpose(-2, -1),
+            torch.bfloat16,
+            q_nope_out,
+        )
+    else:
+        if (_use_aiter_gfx95 and attn.w_kc.dtype == torch.float8_e4m3fn) or (
+            is_capture_mode and attn.w_kc.dtype == torch.float8_e4m3fnuz
+        ):
+            # fp8 Triton kernel: always on gfx950,
+            # cudagraph-only on gfx942 (hides launch overhead)
+            q_nope_out = (
+                batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
+                    X=q_nope,
+                    WQ=attn.w_kc.transpose(-1, -2),
+                    w_scale=attn.w_scale,
+                    group_size=128,
+                    YQ=None,  # allocate (B, M, N)
+                    transpose_bm=False,  # (B, M, N)
+                    transpose_bm_in=True,  # (M, B, K)
+                    dtype=torch.bfloat16,
+                )
+            )
+        else:
+            q_nope_out = torch.bmm(
+                q_nope.to(torch.bfloat16).transpose(0, 1),
+                attn.w_kc.to(torch.bfloat16) * attn.w_scale,
+            )
+    return q_nope_out
+
+
+def rocm_absorb_v_bmm(
+    attn: DeepseekV2AttentionMLA,
+    attn_output: torch.Tensor,
+) -> torch.Tensor:
+    """Absorb ``attn_output @ w_vc`` (+ optional fused flatten quant) on HIP."""
+    # TODO(haishaw): add bmm_fp8 to ROCm
+    if _use_aiter_gfx95 and attn.w_vc.dtype == torch.uint8:
+        x = attn_output.transpose(0, 1)
+        B_heads, M_batch = x.shape[0], x.shape[1]
+        N_vdim = attn.w_vc.shape[2]
+        # Allocate in (batch, heads, dim) so the post-GEMM
+        # transpose+flatten is a free view instead of a copy.
+        _bmm_buf = torch.empty(
+            M_batch,
+            B_heads,
+            N_vdim,
+            device=x.device,
+            dtype=torch.bfloat16,
+        )
+        attn_bmm_output = _bmm_buf.transpose(0, 1)
+        batched_gemm_afp4wfp4_pre_quant(
+            x,
+            attn.w_vc.transpose(-2, -1),
+            attn.w_scale_v.transpose(-2, -1),
+            torch.bfloat16,
+            attn_bmm_output,
+        )
+    else:
+        _bmm_buf = None
+        if _use_aiter_gfx95 and attn.w_kc.dtype == torch.float8_e4m3fn:
+            attn_bmm_output = (
+                batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
+                    X=attn_output,
+                    WQ=attn.w_vc.transpose(-1, -2),
+                    w_scale=attn.w_scale,
+                    group_size=128,
+                    YQ=None,
+                    transpose_bm=False,
+                    transpose_bm_in=True,
+                    dtype=torch.bfloat16,
+                )
+            )
+        else:
+            attn_bmm_output = torch.bmm(
+                attn_output.to(torch.bfloat16).transpose(0, 1),
+                attn.w_vc.to(torch.bfloat16) * attn.w_scale,
+            )
+
+    if _bmm_buf is not None:
+        # _bmm_buf is already (batch, heads, dim) contiguous
+        if attn.o_proj.weight.dtype == torch.uint8:
+            attn_bmm_output = fused_flatten_mxfp4_quant(_bmm_buf)
+        elif attn.o_proj.weight.dtype == torch.float8_e4m3fn:
+            attn_bmm_output = fused_flatten_fp8_group_quant(
+                _bmm_buf,
+                group_size=128,
+                dtype_quant=torch.float8_e4m3fn,
+                transpose_scale=False,
+            )
+            if _use_aiter_bpreshuffle_gfx95:
+                attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple(
+                    attn_bmm_output
+                )
+        else:
+            attn_bmm_output = _bmm_buf.flatten(1, 2)
+    elif attn.o_proj.weight.dtype == torch.uint8:
+        attn_bmm_output = attn_bmm_output.transpose(0, 1)
+        attn_bmm_output = fused_flatten_mxfp4_quant(attn_bmm_output)
+    elif attn.o_proj.weight.dtype == torch.float8_e4m3fn:
+        attn_bmm_output = attn_bmm_output.transpose(0, 1)
+        attn_bmm_output = fused_flatten_fp8_group_quant(
+            attn_bmm_output,
+            group_size=128,
+            dtype_quant=torch.float8_e4m3fn,
+            transpose_scale=False,
+        )
+        if _use_aiter_bpreshuffle_gfx95:
+            attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple(attn_bmm_output)
+    else:
+        attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
+
+    return attn_bmm_output
+
+
+def _fused_rope_cat_and_cache(
+    attn: DeepseekV2AttentionMLA,
+    q_nope_out: torch.Tensor,
+    q_pe: torch.Tensor,
+    k_nope: torch.Tensor,
+    k_pe: torch.Tensor,
+    positions: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """RoPE + concat + KV-cache write via the AITER fused kernel on gfx95."""
+    kv_cache_dtype = (
+        fp8_dtype if attn.kv_cache_dtype == "fp8_e4m3" else q_nope_out.dtype
+    )
+    return fused_qk_rope_cat_and_cache_mla(
+        q_nope_out,
+        q_pe,
+        k_nope,
+        k_pe,
+        get_token_to_kv_pool().get_key_buffer(attn.attn_mqa.layer_id),
+        out_cache_loc,
+        positions,
+        attn.rotary_emb.cos_cache,
+        attn.rotary_emb.sin_cache,
+        attn.attn_mqa.k_scale,
+        attn.rotary_emb.is_neox_style,
+        q_out_dtype=kv_cache_dtype,
+    )
+
+
+class DeepseekMLARocmForwardMixin:
+
+    def forward_absorb_rocm_prepare(
+        self: DeepseekV2AttentionMLA,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+        zero_allocator: BumpAllocator,
+        llama_4_scaling: Optional[torch.Tensor] = None,
+        prev_topk_indices: Optional[torch.Tensor] = None,
+    ):
+        from sglang.srt.model_executor.runner import get_is_capture_mode
+
+        q_replicate_active = (
+            get_parallel().dcp_replicate_q_proj
+            and is_dcp_mla_decode_phase(forward_batch)
+            and not self.use_deep_gemm_bmm
+            and self.w_kc_qrep is not None
+            and self.q_b_proj_qrep_weight is not None
+        )
+        q_lora = None
+        topk_indices = None
+        if self.q_lora_rank is not None:
+            q, latent_cache = (
+                get_attn_tp_context()
+                .fetch_qkv_latent()
+                .split(
+                    [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
+                    dim=-1,
+                )
+            )
+            k_nope = latent_cache[..., : self.kv_lora_rank]
+
+            # overlap qk norm
+            if self.alt_stream is not None and get_is_capture_mode():
+                current_stream = torch.cuda.current_stream()
+                self.alt_stream.wait_stream(current_stream)
+                q = self.q_a_layernorm(q)
+                with torch.cuda.stream(self.alt_stream):
+                    k_nope = self.kv_a_layernorm(k_nope)
+                current_stream.wait_stream(self.alt_stream)
+            elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
+                q, _, k_nope, *_ = fused_rms_mxfp4_quant(
+                    q,
+                    self.q_a_layernorm.weight,
+                    self.q_a_layernorm.variance_epsilon,
+                    k_nope,
+                    self.kv_a_layernorm.weight,
+                    self.kv_a_layernorm.variance_epsilon,
+                )
+            elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
+                if self.use_dsa:
+                    q_quanted, q_lora, k_nope, _ = fused_rms_fp8_group_quant(
+                        q,
+                        self.q_a_layernorm.weight,
+                        self.q_a_layernorm.variance_epsilon,
+                        k_nope,
+                        self.kv_a_layernorm.weight,
+                        self.kv_a_layernorm.variance_epsilon,
+                        group_size=128,
+                        dtype_quant=torch.float8_e4m3fn,
+                        res1=None,
+                        output_unquantized_inp1=True,
+                        transpose_scale=False,
+                    )
+                    if _use_aiter_bpreshuffle_gfx95:
+                        q_quanted = materialize_bpreshuffle_fp8_scale_tuple(q_quanted)
+                    q = q_quanted
+                else:
+                    q, _, k_nope, _ = fused_rms_fp8_group_quant(
+                        q,
+                        self.q_a_layernorm.weight,
+                        self.q_a_layernorm.variance_epsilon,
+                        k_nope,
+                        self.kv_a_layernorm.weight,
+                        self.kv_a_layernorm.variance_epsilon,
+                        group_size=128,
+                        dtype_quant=torch.float8_e4m3fn,
+                        res1=None,
+                        output_unquantized_inp1=False,
+                        transpose_scale=False,
+                    )
+                    if _use_aiter_bpreshuffle_gfx95:
+                        q = materialize_bpreshuffle_fp8_scale_tuple(q)
+            elif _use_aiter:
+                q, k_nope = fused_qk_rmsnorm_bf16(
+                    q,
+                    self.q_a_layernorm.weight,
+                    self.q_a_layernorm.variance_epsilon,
+                    k_nope,
+                    self.kv_a_layernorm.weight,
+                    self.kv_a_layernorm.variance_epsilon,
+                )
+            else:
+                q = self.q_a_layernorm(q)
+                k_nope = self.kv_a_layernorm(k_nope)
+
+            # q_lora needed by indexer
+            if self.use_dsa:
+                if q_lora is None:
+                    q_lora = q
+
+            # overlap q_b_proj and indexer during decode
+            if (
+                self.alt_stream is not None
+                and get_is_capture_mode()
+                and forward_batch.forward_mode.is_decode_or_idle()
+                and q_lora is not None
+                and not q_replicate_active
+            ):
+                current_stream = torch.cuda.current_stream()
+                self.alt_stream.wait_stream(current_stream)
+                with torch.cuda.stream(self.alt_stream):
+                    k_nope = k_nope.unsqueeze(1)
+                    q = self.q_b_proj_forward(q)
+                if self.should_run_indexer(prev_topk_indices):
+                    topk_indices = self.indexer(
+                        x=hidden_states,
+                        q_lora=q_lora,
+                        positions=positions,
+                        forward_batch=forward_batch,
+                        layer_id=self.layer_id,
+                    )
+                else:
+                    # skip_topk reuses prev layer's indices; mirror into this
+                    # layer's slot so the captured buffer matches what's used.
+                    topk_indices = maybe_capture_indexer_topk(
+                        self.layer_id, prev_topk_indices
+                    )
+                current_stream.wait_stream(self.alt_stream)
+            else:
+                k_nope = k_nope.unsqueeze(1)
+                if q_replicate_active:
+                    q = torch.nn.functional.linear(q, self.q_b_proj_qrep_weight).view(
+                        -1,
+                        self.num_local_heads * get_parallel().attn_dcp_size,
+                        self.qk_head_dim,
+                    )
+                else:
+                    q = self.q_b_proj_forward(q)
+
+                if q_lora is not None:
+                    if self.should_run_indexer(prev_topk_indices):
+                        topk_indices = self.indexer(
+                            x=hidden_states,
+                            q_lora=q_lora,
+                            positions=positions,
+                            forward_batch=forward_batch,
+                            layer_id=self.layer_id,
+                        )
+                    else:
+                        topk_indices = maybe_capture_indexer_topk(
+                            self.layer_id, prev_topk_indices
+                        )
+        else:
+            if q_replicate_active:
+                q = torch.nn.functional.linear(
+                    hidden_states, self.q_b_proj_qrep_weight
+                ).view(
+                    -1,
+                    self.num_local_heads * get_parallel().attn_dcp_size,
+                    self.qk_head_dim,
+                )
+            else:
+                q = self.q_proj(hidden_states)[0].view(
+                    -1, self.num_local_heads, self.qk_head_dim
+                )
+            latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
+            k_nope = latent_cache[..., : self.kv_lora_rank]
+            k_nope = self.kv_a_layernorm(k_nope).unsqueeze(1)
+
+        q_nope, q_pe, k_pe = self._split_q_nope_pe(q, latent_cache)
+
+        if q_replicate_active:
+            q_nope_out = (
+                torch.bmm(q_nope.transpose(0, 1), self.w_kc_qrep)
+                .transpose(0, 1)
+                .contiguous()
+            )
+        else:
+            _kvb_q = None
+            if _SGLANG_EXPERIMENTAL_LORA_OPTI:
+                # Fork the kv_b q-correction A-step onto the LoRA side stream to overlap the bmm.
+                from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
+                    kv_b_lora_q_prepare,
+                )
+
+                _kvb_q = kv_b_lora_q_prepare(self, q_nope)
+
+            if self.use_deep_gemm_bmm:
+                (
+                    q_nope_val,
+                    q_nope_scale,
+                    masked_m,
+                    expected_m,
+                    aligned_m,
+                ) = per_token_group_quant_mla_deep_gemm_masked_fp8(
+                    q_nope.transpose(0, 1)
+                )
+                q_nope_out = q_nope.new_empty(
+                    (self.num_local_heads, aligned_m, self.kv_lora_rank)
+                )
+                deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
+                    (q_nope_val, q_nope_scale),
+                    (self.w_kc, self.w_scale_k),
+                    q_nope_out,
+                    masked_m,
+                    expected_m,
+                )
+                q_nope_out = q_nope_out[:, :expected_m, :]
+            else:
+                q_nope_out = rocm_absorb_q_bmm(
+                    self, q_nope, is_capture_mode=get_is_capture_mode()
+                )
+
+            q_nope_out = q_nope_out.transpose(0, 1)
+            if _SGLANG_EXPERIMENTAL_LORA_OPTI:
+                from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
+                    kv_b_lora_q_apply,
+                )
+
+                q_nope_out = kv_b_lora_q_apply(self, q_nope, q_nope_out, _kvb_q)
+            elif is_kv_b_lora_active(self):
+                q_nope_out = apply_kv_b_lora_q_correction(self, q_nope, q_nope_out)
+
+        fuse_rope_for_trtllm_mla = self._fuse_rope_for_trtllm_mla(forward_batch)
+        if (
+            self.rotary_emb is not None
+            and (not fuse_rope_for_trtllm_mla)
+            and (not self._skip_rope_for_dsa_tilelang_fused())
+            and (not self._skip_rope_for_aiter_fused_mla())
+            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)
+
+        dsa_prefill_cp = dsa_use_prefill_cp(forward_batch)
+        mla_prefill_cp = mla_use_prefill_cp(forward_batch)
+        defer_kv_gather_until_after_rope = should_defer_dsa_cp_kv_gather(
+            dsa_prefill_cp=dsa_prefill_cp,
+            fuse_rope_for_trtllm_mla=fuse_rope_for_trtllm_mla,
+        )
+        if dsa_prefill_cp and not defer_kv_gather_until_after_rope:
+            from sglang.srt.layers.attention.dsa_backend import materialize_full_kv_cp
+
+            k_nope, k_pe = materialize_full_kv_cp(
+                self,
+                forward_batch,
+                latent_cache,
+                k_nope,
+                k_pe,
+            )
+        elif mla_prefill_cp and not is_cp_v2_active(forward_batch):
+            # CP-v1 gathers the latent here; CP-v2 gathers it in the attention
+            # backend via the strategy (materialize_full_mla_kv).
+            k_nope, k_pe = self.rebuild_cp_kv_cache(
+                latent_cache,
+                forward_batch,
+                k_nope,
+                k_pe,
+            )
+
+        # all_gather q_pe, q_nope_out,take tp8 as an example， q_pe [B, H, ROPE_DIM], q_nope_out [B, H, NOPE_DIM] gathered to [B, H * dcp_world_size, ROPE_DIM] [B, H * dcp_world_size, NOPE_DIM] for decode batch, and all gather k_pe, k_nope for extend batch.
+        if get_parallel().dcp_enabled:
+            if is_dcp_mla_decode_phase(forward_batch):
+                if not q_replicate_active:
+                    q_nope_out, q_pe = all_gather_q_for_mla_decode(
+                        q_nope_out=q_nope_out,
+                        q_pe=q_pe,
+                    )
+            elif forward_batch.forward_mode.is_extend():
+                # for extend, gather kv
+                all_gather_kv_cache_for_mla_extend(
+                    get_token_to_kv_pool(),
+                    self.attn_mqa,
+                    forward_batch.extend_prefix_lens_cpu,
+                    forward_batch.attn_dcp_metadata.dcp_local_prefix_kv_indices,
+                    forward_batch.attn_dcp_metadata.dcp_extend_prefix_lens_sum,
+                    forward_batch.attn_dcp_metadata.dcp_kv_buffer,
+                    self.kv_lora_rank,
+                    k_nope,
+                    k_pe,
+                )
+            else:
+                logger.warning(
+                    f"not supported forward_mode {forward_batch.forward_mode}"
+                )
+
+        return (
+            q_pe,
+            k_pe,
+            q_nope_out,
+            k_nope,
+            forward_batch,
+            zero_allocator,
+            positions,
+            topk_indices,
+            llama_4_scaling,
+        )
+
+    def forward_absorb_rocm_core(
+        self: DeepseekV2AttentionMLA,
+        q_pe,
+        k_pe,
+        q_nope_out,
+        k_nope,
+        forward_batch,
+        zero_allocator,
+        positions,
+        topk_indices,
+        llama_4_scaling,
+    ):
+        save_kv_cache = True
+
+        if self.current_attention_backend in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS:
+            if self._skip_rope_for_dsa_tilelang_fused() and self.rotary_emb is not None:
+                q_cat, _, k_pe_fused, _ = _fused_rope_cat_and_cache(
+                    self,
+                    q_nope_out,
+                    q_pe,
+                    k_nope,
+                    k_pe,
+                    positions,
+                    forward_batch.out_cache_loc,
+                )
+                save_kv_cache = False
+                # On decode, pass q_cat directly to attn_mqa with q_rope=None so
+                # dsa_backend.forward_decode reuses q_cat as a zero-copy view
+                # (`q.contiguous().view(...)` fast-path) instead of running the
+                # redundant `concat_mla_absorb_q_general(q_nope_fused, q_pe_fused)`
+                # that would otherwise rebuild a tensor byte-identical to q_cat.
+                # On ROCm tilelang decode, this eliminates the
+                # `CatArrayBatchedCopy<OpaqueType<1u>, ...>` kernel that used to
+                # fire once per layer per decode step (~2.6 us / layer saved).
+                # Prefill keeps the split form because dsa_backend.forward_extend
+                # asserts `q_rope is not None`.
+                if forward_batch.forward_mode.is_decode_or_idle():
+                    if llama_4_scaling is not None:
+                        # llama_4_scaling applies only to the q_nope portion;
+                        # mutate in place via the slice view of q_cat.
+                        q_cat[..., : self.kv_lora_rank] *= llama_4_scaling
+                    attn_output = self.attn_mqa(
+                        q_cat,
+                        None,
+                        None,
+                        forward_batch,
+                        q_rope=None,
+                        k_rope=k_pe_fused,
+                        save_kv_cache=save_kv_cache,
+                        **(
+                            dict(topk_indices=topk_indices)
+                            if topk_indices is not None
+                            else {}
+                        ),
+                    )
+                else:
+                    q_nope_fused = q_cat[..., : self.kv_lora_rank]
+                    q_pe_fused = q_cat[..., self.kv_lora_rank :]
+                    if llama_4_scaling is not None:
+                        q_nope_fused *= llama_4_scaling
+                    attn_output = self.attn_mqa(
+                        q_nope_fused,
+                        None,
+                        None,
+                        forward_batch,
+                        q_rope=q_pe_fused,
+                        k_rope=k_pe_fused,
+                        save_kv_cache=save_kv_cache,
+                        **(
+                            dict(topk_indices=topk_indices)
+                            if topk_indices is not None
+                            else {}
+                        ),
+                    )
+            else:
+                extra_args = {}
+                if self._fuse_rope_for_trtllm_mla(forward_batch):
+                    extra_args = {
+                        "cos_sin_cache": self.rotary_emb.cos_sin_cache,
+                        "is_neox": self.rotary_emb.is_neox_style,
+                        "llama_4_scaling": llama_4_scaling,
+                    }
+                if is_dcp_mla_decode_phase(forward_batch):
+                    # set return_lse=True to correct attn_output
+                    attn_output, lse = self.attn_mqa_for_dcp_decode(
+                        q_nope_out,
+                        k_nope,
+                        k_nope,
+                        forward_batch,
+                        q_rope=q_pe,
+                        k_rope=k_pe,
+                        **extra_args,
+                        **(
+                            dict(topk_indices=topk_indices)
+                            if topk_indices is not None
+                            else {}
+                        ),
+                    )
+                else:
+                    attn_output = self.attn_mqa(
+                        q_nope_out,
+                        k_nope,
+                        k_nope,
+                        forward_batch,
+                        q_rope=q_pe,
+                        k_rope=k_pe,
+                        **extra_args,
+                        **(
+                            dict(topk_indices=topk_indices)
+                            if topk_indices is not None
+                            else {}
+                        ),
+                    )
+        else:
+            if self._skip_rope_for_aiter_fused_mla():
+                q, _, _, k = _fused_rope_cat_and_cache(
+                    self,
+                    q_nope_out,
+                    q_pe,
+                    k_nope,
+                    k_pe,
+                    positions,
+                    forward_batch.out_cache_loc,
+                )
+                save_kv_cache = False
+            else:
+                q = torch.cat([q_nope_out, q_pe], dim=-1)
+                k = torch.cat([k_nope, k_pe], dim=-1)
+
+            # Apply llama 4 scaling if provided
+            if llama_4_scaling is not None:
+                q *= llama_4_scaling
+
+            attn_output = self.attn_mqa(
+                q,
+                k,
+                k_nope,
+                forward_batch,
+                save_kv_cache=save_kv_cache,
+                **(dict(topk_indices=topk_indices) if topk_indices is not None else {}),
+            )
+
+        # correct attn_output with respect to lse from other ranks
+        if is_dcp_mla_decode_phase(forward_batch):
+            attn_output = attn_output.view(
+                -1,
+                self.num_local_heads * get_parallel().attn_dcp_size,
+                self.kv_lora_rank,
+            )
+            if get_in_autotune_dummy_run():
+                # The synthetic FlashInfer MoE autotune pass discards model
+                # outputs. Avoid an unnecessary cross-node MNNVL exchange of
+                # zero attention partials.
+                attn_output = _select_local_dcp_heads_for_autotune(
+                    attn_output, self.num_local_heads
+                )
+            else:
+                dcp_comm_backend = get_parallel().dcp_comm_backend
+                is_lse_base_on_e = is_mla_dcp_lse_base_on_e(
+                    self.current_attention_backend
+                )
+                if dcp_comm_backend in ("a2a", "fi_a2a"):
+                    # A2A exchange of head partials + LSE, then local Triton combine.
+                    attn_output = dcp_a2a_lse_reduce(
+                        attn_output.contiguous(),
+                        lse.contiguous(),
+                        get_parallel().dcp_group,
+                        is_lse_base_on_e=is_lse_base_on_e,
+                        comm_backend=dcp_comm_backend,
+                    )
+                else:
+                    attn_output = cp_lse_ag_out_rs_mla(
+                        attn_output,
+                        lse,
+                        get_parallel().dcp_group,
+                        is_lse_base_on_e=is_lse_base_on_e,
+                    )
+                    attn_output = attn_output.transpose(0, 1)
+        attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)
+
+        _kvb_v = None
+        if _SGLANG_EXPERIMENTAL_LORA_OPTI:
+            # Fork the kv_b v-correction A-step onto the LoRA side stream to overlap the bmm.
+            from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
+                kv_b_lora_v_prepare,
+            )
+
+            _kvb_v = kv_b_lora_v_prepare(self, attn_output)
+
+        if self.use_deep_gemm_bmm:
+            (
+                attn_output_val,
+                attn_output_scale,
+                masked_m,
+                expected_m,
+                aligned_m,
+            ) = per_token_group_quant_mla_deep_gemm_masked_fp8(
+                attn_output.transpose(0, 1)
+            )
+            attn_bmm_output = attn_output.new_empty(
+                (self.num_local_heads, aligned_m, self.v_head_dim)
+            )
+            deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
+                (attn_output_val, attn_output_scale),
+                (self.w_vc, self.w_scale_v),
+                attn_bmm_output,
+                masked_m,
+                expected_m,
+            )
+            attn_bmm_output = (
+                attn_bmm_output[:, :expected_m, :].transpose(0, 1).flatten(1, 2)
+            )
+        else:
+            attn_bmm_output = rocm_absorb_v_bmm(self, attn_output)
+
+        if _SGLANG_EXPERIMENTAL_LORA_OPTI:
+            from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
+                kv_b_lora_v_apply,
+            )
+
+            attn_bmm_output = kv_b_lora_v_apply(
+                self, attn_output, attn_bmm_output, _kvb_v
+            )
+        elif is_kv_b_lora_active(self):
+            attn_bmm_output = apply_kv_b_lora_v_correction(
+                self, attn_output, attn_bmm_output
+            )
+        output, _ = self.o_proj(attn_bmm_output)
+
+        if self.next_skip_topk is None:
+            return output
+
+        # Return topk_indices for the next layer when enabling index cache
+        if not self.next_skip_topk:
+            return output, None
+        else:
+            return output, topk_indices
+
+    def _skip_rope_for_dsa_tilelang_fused(self: DeepseekV2AttentionMLA) -> bool:
+        """
+        Check if we should skip rope and use fused rope+cache path for TileLang DSA on gfx95.
+        """
+        return (
+            _use_aiter_gfx95
+            and self.current_attention_backend in ("dsa", "nsa")
+            and (
+                get_exec().kernel.dsa_decode_backend == "tilelang"
+                or get_exec().kernel.dsa_prefill_backend == "tilelang"
+            )
+        )
+
+    def _skip_rope_for_aiter_fused_mla(self: DeepseekV2AttentionMLA) -> bool:
+        """
+        Skip rope in prepare and let the fused kernel in forward_absorb_rocm_core handle it,
+        when running aiter-backend MLA on gfx95 (i.e., the `else` branch in
+        forward_absorb_rocm_core that calls fused_qk_rope_cat_and_cache_mla).
+        """
+        return _use_aiter_gfx95 and self.current_attention_backend == "aiter"

python/sglang/srt/models/deepseek_v2.py

@@ -162,12 +162,15 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
 )
 from sglang.srt.models.deepseek_common.attention_backend_handler import (
     AttentionBackendRegistry,
+    resolve_rocm_forward_method,
 )
 from sglang.srt.models.deepseek_common.attention_forward_methods import (
     AttnForwardMethod,
     DeepseekMHAForwardMixin,
+    DeepseekMHARocmForwardMixin,
     DeepseekMLACpuForwardMixin,
     DeepseekMLAForwardMixin,
+    DeepseekMLAFusedRopeRocmForwardMixin,
     DeepseekMLARocmForwardMixin,
 )
 from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
@@ -1711,8 +1714,10 @@ class DeepseekV2MoE(nn.Module):
 class DeepseekV2AttentionMLA(
     nn.Module,
     DeepseekMHAForwardMixin,
+    DeepseekMHARocmForwardMixin,
     DeepseekMLAForwardMixin,
     DeepseekMLARocmForwardMixin,
+    DeepseekMLAFusedRopeRocmForwardMixin,
     DeepseekMLACpuForwardMixin,
 ):
 
@@ -2021,7 +2026,7 @@ class DeepseekV2AttentionMLA(
         self.current_attention_backend = attention_backend
 
         handler = AttentionBackendRegistry.get_handler(attention_backend)
-        return handler(self, forward_batch)
+        return resolve_rocm_forward_method(handler(self, forward_batch))
 
     def op_prepare(self, state):
         state.attn_intermediate_state = self.forward_prepare(
@@ -2117,6 +2122,23 @@ class DeepseekV2AttentionMLA(
                 llama_4_scaling,
                 prev_topk_indices,
             )
+        elif attn_forward_method == AttnForwardMethod.MHA_ROCM:
+            inner_state = self.forward_normal_rocm_prepare(
+                positions, hidden_states, forward_batch, zero_allocator
+            )
+        elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT_ROCM:
+            inner_state = self.forward_normal_one_shot_rocm_prepare(
+                positions, hidden_states, forward_batch, zero_allocator
+            )
+        elif attn_forward_method == AttnForwardMethod.MLA_ROCM:
+            inner_state = self.forward_absorb_rocm_prepare(
+                positions,
+                hidden_states,
+                forward_batch,
+                zero_allocator,
+                llama_4_scaling,
+                prev_topk_indices,
+            )
         elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM:
             inner_state = self.forward_absorb_fused_mla_rope_prepare(
                 positions, hidden_states, forward_batch, zero_allocator
@@ -2172,6 +2194,12 @@ class DeepseekV2AttentionMLA(
             return self.forward_normal_one_shot_core(*inner_state)
         elif attn_forward_method == AttnForwardMethod.MLA:
             return self.forward_absorb_core(*inner_state)
+        elif attn_forward_method == AttnForwardMethod.MHA_ROCM:
+            return self.forward_normal_core(*inner_state)
+        elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT_ROCM:
+            return self.forward_normal_one_shot_core(*inner_state)
+        elif attn_forward_method == AttnForwardMethod.MLA_ROCM:
+            return self.forward_absorb_rocm_core(*inner_state)
         elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM:
             return self.forward_absorb_fused_mla_rope_core(*inner_state)
         elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_CPU:

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.assertTrue(is_mla_dcp_lse_base_on_e("flashmla"))
-        self.assertFalse(_is_mla_dcp_lse_base_on_e("flashinfer_mla"))
+        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()