[Perf] [AMD] mla: Fuse RoPE, concat and KV-cache write for MLA decode on gfx1250

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

@@ -34,6 +34,10 @@
     view_aiter_fused_rms_transposed_fp8_scale_tuple,
 )
 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+from sglang.srt.layers.rotary_embedding.base import RotaryEmbedding
+from sglang.srt.layers.rotary_embedding.rope_variant import (
+    DeepseekScalingRotaryEmbedding,
+)
 from sglang.srt.lora.deepseek_mla_correction import (
     apply_q_correction as apply_kv_b_lora_q_correction,
 )
@@ -153,6 +157,15 @@ def fused_qk_rmsnorm_bf16(q, q_weight, q_eps, k, k_weight, k_eps):
         fused_qk_rope_cat_and_cache_mla,
     )
 
+# The gfx1250 launch of this kernel needs aiter with ROCm/aiter#5501.
+_use_gfx1250_fused_mla_rope = _use_aiter and _is_gfx1250_supported
+if _use_gfx1250_fused_mla_rope:
+    from aiter.ops.triton.fusions.fused_kv_cache import (
+        fused_qk_rope_cat_and_cache_mla,
+    )
+
+_GFX1250_FUSED_ROPE_BACKENDS = ("aiter", "triton")
+
 
 _ABSORB_BMM_DECODE_MAX_M = 256
 _absorb_bmm_num_cu = 0
@@ -435,18 +448,29 @@ def _fused_rope_cat_and_cache(
     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."""
+    """RoPE + concat + KV-cache write via the AITER fused kernel on gfx95 or gfx1250."""
     kv_cache_dtype = (
         fp8_dtype if attn.kv_cache_dtype == "fp8_e4m3" else q_nope_out.dtype
     )
     # Gluon MLA decode (bh16bn128) requires bf16 Q; vLLM #50563.
+    # gfx1250 MLA decode (aiter ASM, paged Gluon, triton) takes an unquantized Q;
+    # q_nope_out is bf16 there, since rocm_absorb_q_bmm computes in bf16.
     q_out_dtype = (
         q_nope_out.dtype
-        if attn.kv_cache_dtype == "fp8_e4m3"
-        and attn.current_attention_backend == "aiter"
-        and q_nope_out.shape[-2] == 12
+        if (
+            attn.kv_cache_dtype == "fp8_e4m3"
+            and attn.current_attention_backend == "aiter"
+            and q_nope_out.shape[-2] == 12
+        )
+        or _use_gfx1250_fused_mla_rope
         else kv_cache_dtype
     )
+    if _use_gfx1250_fused_mla_rope:
+        # sglang's rope keeps [cos | sin] in one cache; the kernel takes strided
+        # cos/sin, so the two halves go in as views.
+        cos, sin = attn.rotary_emb.cos_sin_cache.chunk(2, dim=-1)
+    else:
+        cos, sin = attn.rotary_emb.cos_cache, attn.rotary_emb.sin_cache
     kv_pool = get_token_to_kv_pool()
     if isinstance(kv_pool, HiSparseDSATokenToKVPool):
         # The fused write bypasses set_mla_kv_buffer()'s logical-to-device mapping.
@@ -461,8 +485,8 @@ def _fused_rope_cat_and_cache(
         kv_pool.layer_router.read(attn.attn_mqa.layer_id).key(),
         out_cache_loc,
         positions,
-        attn.rotary_emb.cos_cache,
-        attn.rotary_emb.sin_cache,
+        cos,
+        sin,
         attn.attn_mqa.k_scale,
         attn.rotary_emb.is_neox_style,
         q_out_dtype=q_out_dtype,
@@ -794,37 +818,7 @@ def forward_absorb_rocm_prepare(
                 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)
-
-        force_rope_for_aiter_dcp_decode = (
-            self.dcp_size > 1
-            and (
-                forward_batch.forward_mode.is_decode()
-                or forward_batch.forward_mode.is_target_verify()
-                or forward_batch.forward_mode.is_draft_extend_v2()
-            )
-            and _use_aiter_gfx95
-            and self.current_attention_backend
-            not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS
-        )
-        if self.rotary_emb is not None and (
-            force_rope_for_aiter_dcp_decode
-            or (
-                (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
-                        not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS
-                        and self.current_attention_backend != "aiter"
-                    )
-                )
-            )
-        ):
+        if self._applies_rope_in_prepare(forward_batch):
             q_pe, k_pe = self.rotary_emb(positions, q_pe, 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.
@@ -872,6 +866,41 @@ def forward_absorb_rocm_prepare(
             q_nope if fuse_bmm_rope_cache else None,
         )
 
+    def _applies_rope_in_prepare(
+        self: DeepseekV2AttentionMLA, forward_batch: ForwardBatch
+    ) -> bool:
+        """Whether prepare applies rope; otherwise core or the backend fuses it."""
+        force_rope_for_aiter_dcp_decode = (
+            self.dcp_size > 1
+            and (
+                forward_batch.forward_mode.is_decode()
+                or forward_batch.forward_mode.is_target_verify()
+                or forward_batch.forward_mode.is_draft_extend_v2()
+            )
+            and _use_aiter_gfx95
+            and self.current_attention_backend
+            not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS
+        )
+        return self.rotary_emb is not None and (
+            force_rope_for_aiter_dcp_decode
+            or (
+                (not self._fuse_rope_for_trtllm_mla(forward_batch))
+                and (not self._skip_rope_for_dsa_tilelang_fused())
+                and (not self._skip_rope_for_aiter_fused_mla())
+                and (not self._fuse_rope_for_gfx1250_mla(forward_batch))
+                and (
+                    not _use_aiter
+                    or not _is_gfx95_supported
+                    or self.use_dsa
+                    or (
+                        self.current_attention_backend
+                        not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS
+                        and self.current_attention_backend != "aiter"
+                    )
+                )
+            )
+        )
+
     def forward_absorb_rocm_core(
         self: DeepseekV2AttentionMLA,
         q_pe,
@@ -1032,7 +1061,9 @@ def forward_absorb_rocm_core(
                 **(dict(topk_indices=topk_indices) if topk_indices is not None else {}),
             )
         else:
-            if self._skip_rope_for_aiter_fused_mla():
+            if self._skip_rope_for_aiter_fused_mla() or self._fuse_rope_for_gfx1250_mla(
+                forward_batch
+            ):
                 q, _, _, k = _fused_rope_cat_and_cache(
                     self,
                     q_nope_out,
@@ -1228,6 +1259,25 @@ def _skip_rope_for_dsa_tilelang_fused(self: DeepseekV2AttentionMLA) -> bool:
             )
         )
 
+    def _fuse_rope_for_gfx1250_mla(
+        self: DeepseekV2AttentionMLA, forward_batch: ForwardBatch
+    ) -> bool:
+        """Skip rope in prepare and fuse it with the KV-cache write in core on gfx1250."""
+        # The platform term goes first: other platforms never look up the KV pool.
+        return (
+            _use_gfx1250_fused_mla_rope
+            and forward_batch.forward_mode.is_decode()
+            and self.current_attention_backend in _GFX1250_FUSED_ROPE_BACKENDS
+            and not self.use_dsa
+            and self.dcp_size == 1
+            # The kernel applies a plain rope from its cos_sin_cache; subclasses such
+            # as Bailing's video mrope compute theirs differently.
+            and type(self.rotary_emb)
+            in (RotaryEmbedding, DeepseekScalingRotaryEmbedding)
+            # The kernel accepts only a bf16 or fp8 pool.
+            and get_token_to_kv_pool().dtype in (torch.bfloat16, fp8_dtype)
+        )
+
     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,

test/registered/unit/models/test_rocm_hisparse_fused_kv.py

@@ -54,6 +54,8 @@ def invoke(self, locations):
                 lambda *args, **kwargs: args,
                 create=True,
             ),
+            # The gfx1250 fused branch excludes DSA, so these DSA pools use the gfx95 one.
+            patch.object(forward_mla_rocm, "_use_gfx1250_fused_mla_rope", False),
         ):
             args = forward_mla_rocm._fused_rope_cat_and_cache(
                 self.attn, torch.empty(0), None, None, None, None, locations