[ROCm][Kimi-K3][Perf] Fuse MLA decode KV-cache write and Q-prep via AITER

vllm/_aiter_ops.py

@@ -789,6 +789,61 @@ def _rocm_aiter_mla_decode_fwd_lse_fake(
     return torch.empty(q.shape[:-1], dtype=torch.float32, device=q.device)
 
 
+def _rocm_aiter_fused_qk_rope_concat_and_cache_mla_impl(
+    q_nope: torch.Tensor,
+    q_pe: torch.Tensor,
+    kv_c: torch.Tensor,
+    k_pe: torch.Tensor,
+    kv_cache: torch.Tensor,
+    q_out: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    k_scale: torch.Tensor,
+    q_scale: torch.Tensor,
+    positions: torch.Tensor,
+    cos_cache: torch.Tensor,
+    sin_cache: torch.Tensor,
+    is_neox: bool,
+    is_nope_first: bool,
+) -> None:
+    from aiter.ops.cache import fused_qk_rope_concat_and_cache_mla
+
+    fused_qk_rope_concat_and_cache_mla(
+        q_nope,
+        q_pe,
+        kv_c,
+        k_pe,
+        kv_cache,
+        q_out,
+        slot_mapping,
+        k_scale,
+        q_scale,
+        positions,
+        cos_cache,
+        sin_cache,
+        is_neox=is_neox,
+        is_nope_first=is_nope_first,
+    )
+
+
+def _rocm_aiter_fused_qk_rope_concat_and_cache_mla_fake(
+    q_nope: torch.Tensor,
+    q_pe: torch.Tensor,
+    kv_c: torch.Tensor,
+    k_pe: torch.Tensor,
+    kv_cache: torch.Tensor,
+    q_out: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    k_scale: torch.Tensor,
+    q_scale: torch.Tensor,
+    positions: torch.Tensor,
+    cos_cache: torch.Tensor,
+    sin_cache: torch.Tensor,
+    is_neox: bool,
+    is_nope_first: bool,
+) -> None:
+    pass
+
+
 def _rocm_aiter_w8a8_gemm_impl(
     A: torch.Tensor,
     B: torch.Tensor,
@@ -2206,6 +2261,11 @@ def is_triton_unified_attn_enabled(cls) -> bool:
     def is_fp8bmm_enabled(cls) -> bool:
         return cls._AITER_ENABLED and cls._FP8BMM_ENABLED
 
+    @classmethod
+    @if_aiter_supported
+    def is_fused_mla_qkprep_enabled(cls) -> bool:
+        return cls._AITER_ENABLED and cls._MLA_ENABLED
+
     @classmethod
     @if_aiter_supported
     def is_fp4bmm_enabled(cls) -> bool:
@@ -2439,6 +2499,13 @@ def register_ops_once() -> None:
                 fake_impl=_rocm_aiter_mla_decode_fwd_lse_fake,
             )
 
+            direct_register_custom_op(
+                op_name="rocm_aiter_fused_qk_rope_concat_and_cache_mla",
+                op_func=_rocm_aiter_fused_qk_rope_concat_and_cache_mla_impl,
+                mutates_args=["kv_cache", "q_out"],
+                fake_impl=_rocm_aiter_fused_qk_rope_concat_and_cache_mla_fake,
+            )
+
             direct_register_custom_op(
                 op_name="rocm_aiter_w8a8_gemm",
                 op_func=_rocm_aiter_w8a8_gemm_impl,
@@ -3058,6 +3125,41 @@ def mla_decode_fwd_lse(
             kv_scale=kv_scale,
         )
 
+    @staticmethod
+    def fused_qk_rope_concat_and_cache_mla(
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        kv_c: torch.Tensor,
+        k_pe: torch.Tensor,
+        kv_cache: torch.Tensor,
+        q_out: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        k_scale: torch.Tensor,
+        q_scale: torch.Tensor,
+        positions: torch.Tensor,
+        cos_cache: torch.Tensor,
+        sin_cache: torch.Tensor,
+        is_neox: bool,
+        is_nope_first: bool = True,
+    ) -> None:
+        """Fused sparse-MLA decode Q-prep"""
+        torch.ops.vllm.rocm_aiter_fused_qk_rope_concat_and_cache_mla(
+            q_nope,
+            q_pe,
+            kv_c,
+            k_pe,
+            kv_cache,
+            q_out,
+            slot_mapping,
+            k_scale,
+            q_scale,
+            positions,
+            cos_cache,
+            sin_cache,
+            is_neox,
+            is_nope_first,
+        )
+
     @staticmethod
     def per_tensor_quant(
         x: torch.Tensor,

vllm/model_executor/layers/attention/mla_attention.py

@@ -445,6 +445,7 @@ def __init__(
         non_causal_multi_token_decode: bool = False,
         sliding_window: int | None = None,
         prefill_backend_cls: type[MLAPrefillBackend] | None = None,
+        rotary_emb: nn.Module | None = None,
         **extra_impl_args,
     ):
         super().__init__()
@@ -588,6 +589,12 @@ def __init__(
         kv_b_proj._cpu_skip_gemm_dispatch = True
         self.use_direct_call = not current_platform.opaque_attention_op()
 
+        # Disable fused Q-prep path when no rotary_emb is available.
+        self.rotary_emb = rotary_emb
+        self._fused_rope_cos_sin: tuple[torch.Tensor, torch.Tensor] | None = None
+        if getattr(self.impl, "use_fused_qk_rope_cache", False) and rotary_emb is None:
+            self.impl.use_fused_qk_rope_cache = False  # type: ignore[attr-defined]
+
         vllm_config = get_current_vllm_config()
         parallel_config = vllm_config.parallel_config
         self.use_pcp = parallel_config.prefill_context_parallel_size > 1
@@ -759,6 +766,9 @@ def update_kv_cache(
         cache = self.hisparse_cache
         if slot_mapping is None or (cache is not None and cache.dummy_batch):
             return
+        if self._fuses_qk_rope_cache(attn_metadata):
+            # forward_impl writes the cache as part of the fused Q-prep.
+            return
         kv_c_normed, k_pe, slot_mapping = maybe_gather_mla_latent_cache_inputs(
             kv_c_normed,
             k_pe,
@@ -801,13 +811,71 @@ def prepare_kv_cache_update(
         if cache is not None and cache.runtime.is_group_leader:
             cache.prepare_group_for_batch(attn_metadata)
 
+    def can_fuse_qk_rope_cache(self) -> bool:
+        """Batch-independent half of the fused Q-prep decision.
+
+        The MLA wrapper decides RoPE ownership once at load time, while
+        ``_fuses_qk_rope_cache`` decides per batch. Both are built on this, so
+        the wrapper cannot hand RoPE to a kernel that will not run.
+        """
+        if not getattr(self.impl, "use_fused_qk_rope_cache", False):
+            return False
+        if self.rotary_emb is None:
+            return False
+        # An impl that fuses only some batches leaves the wrapper applying
+        # RoPE eagerly on all of them, so the kernel must not apply it again.
+        # Safe only when that rotation is a declared no-op (a NoPE model).
+        if not getattr(self.impl, "always_fuses_qk_rope_cache", False) and not getattr(
+            self.rotary_emb, "is_identity", False
+        ):
+            return False
+        # Skipping the standalone write also skips the PCP latent all-gather,
+        # a collective every rank must reach, and the HiSparse cursor advance.
+        # Neither is reproduced by the fused kernel.
+        return not self.use_pcp and self.hisparse_cache is None
+
+    def _fuses_qk_rope_cache(self, attn_metadata: "MLACommonMetadata | None") -> bool:
+        """Whether this batch folds the KV-cache write into the fused Q-prep.
+
+        ``update_kv_cache`` and ``forward_impl`` must agree: the former skips
+        the standalone write exactly when the latter performs the fused call.
+        Derived only from state both already hold, so no value is carried
+        across the two custom-op boundaries between them.
+        """
+        if not self.can_fuse_qk_rope_cache():
+            return False
+        if getattr(self.impl, "always_fuses_qk_rope_cache", False):
+            # Declared to fuse unconditionally -- sparse routes every token
+            # through MQA, so one call covers the batch, prefill included.
+            # Keyed off the same flag the wrapper uses to hand over RoPE, so
+            # the two agree by construction.
+            return True
+        # Dense MLA sends only decode tokens through MQA, so the fused call
+        # would leave a mixed batch's prefill rows unwritten.
+        if attn_metadata is None:
+            return False
+        if attn_metadata.num_actual_tokens != attn_metadata.num_decode_tokens:
+            return False
+        # Under DCP the kernel skips padded (slot < 0) tokens, but those rows
+        # still join the query all-gather, so their q_out would be garbage.
+        return self.impl.dcp_world_size == 1
+
+    def _get_fused_rope_cos_sin(self) -> tuple[torch.Tensor, torch.Tensor]:
+        """Split rotary_emb.cos_sin_cache into contiguous cos/sin halves."""
+        if self._fused_rope_cos_sin is None:
+            assert self.rotary_emb is not None
+            cos, sin = self.rotary_emb.cos_sin_cache.chunk(2, dim=-1)
+            self._fused_rope_cos_sin = (cos.contiguous(), sin.contiguous())
+        return self._fused_rope_cos_sin
+
     def forward(
         self,
         q: torch.Tensor,
         kv_c_normed: torch.Tensor,
         k_pe: torch.Tensor,
         output_shape: torch.Size | None = None,
         q_dcp_replicated: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
     ) -> torch.Tensor:
         if self.use_direct_call:
             forward_context: ForwardContext = get_forward_context()
@@ -849,6 +917,7 @@ def forward(
                 attn_metadata,
                 output=output,
                 q_dcp_replicated=q_dcp_replicated,
+                positions=positions,
             )
             return output
         else:
@@ -869,6 +938,7 @@ def forward(
                 encoded,
                 kv_cache_dummy_dep=kv_cache_dummy_dep,
                 q_dcp_replicated=q_dcp_replicated,
+                positions=positions,
             )
             return output
 
@@ -887,6 +957,7 @@ def forward_impl(
         quant_col_major: bool | None = None,
         quant_tma_aligned: bool | None = None,
         q_dcp_replicated: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
     ) -> torch.Tensor:
         assert output is not None, "Output tensor must be provided."
 
@@ -1080,7 +1151,35 @@ def forward_impl(
                     mqa_ql_nope = mqa_q_nope.new_empty((B, N, L))
                     torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope.transpose(0, 1))
 
-            if fp8_attention and self.impl.supports_quant_query_input:
+            if self._fuses_qk_rope_cache(attn_metadata):
+                if positions is None:
+                    raise RuntimeError(
+                        "Fused MLA Q-prep is enabled but `positions` was not "
+                        "passed to MLAAttention.forward. The MLA wrapper must "
+                        "forward positions when impl.use_fused_qk_rope_cache is set."
+                    )
+                forward_context = get_forward_context()
+                fc_slot_mapping = forward_context.slot_mapping
+                assert isinstance(fc_slot_mapping, dict), (
+                    f"Expected slot_mapping to be a dict, got {type(fc_slot_mapping)}."
+                )
+                slot_mapping = fc_slot_mapping[self.layer_name].flatten()
+                cos_cache, sin_cache = self._get_fused_rope_cos_sin()
+                assert self.rotary_emb is not None
+                mqa_q = self.impl.fused_qk_rope_concat_and_cache(  # type: ignore[attr-defined]
+                    self,
+                    mqa_ql_nope,
+                    mqa_q_pe,
+                    k_c_normed[:num_mqa_tokens],
+                    k_pe[:num_mqa_tokens],
+                    kv_cache,
+                    slot_mapping[:num_mqa_tokens],
+                    positions[:num_mqa_tokens],
+                    cos_cache,
+                    sin_cache,
+                    self.rotary_emb.is_neox_style,
+                )
+            elif fp8_attention and self.impl.supports_quant_query_input:
                 assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0]
                 assert mqa_ql_nope.shape[1] == mqa_q_pe.shape[1]
                 mqa_q = self._decode_concat_quant_fp8_op(
@@ -1470,6 +1569,7 @@ def unified_mla_attention_with_output(
     quant_col_major: bool | None = None,
     quant_tma_aligned: bool | None = None,
     q_dcp_replicated: torch.Tensor | None = None,
+    positions: torch.Tensor | None = None,
 ) -> None:
     # kv_cache_dummy_dep is not used but accepting it creates a data dependency
     # that ensures torch.compile preserves ordering between KV cache update and
@@ -1493,6 +1593,7 @@ def unified_mla_attention_with_output(
         quant_col_major=quant_col_major,
         quant_tma_aligned=quant_tma_aligned,
         q_dcp_replicated=q_dcp_replicated,
+        positions=positions,
     )
 
 

vllm/model_executor/layers/mla.py

@@ -137,6 +137,14 @@ def __init__(
             topk_indices_buffer=mla_modules.topk_indices_buffer,
             index_group_builder=mla_modules.index_group_builder,
             non_causal_multi_token_decode=non_causal_multi_token_decode,
+            rotary_emb=self.rotary_emb,
+        )
+
+        # Hand RoPE to the kernel only when it can fuse at all *and* fuses
+        # every batch -- otherwise the batches it declines would lose RoPE.
+        self._defer_rope_to_fused_kernel = (
+            self.mla_attn.can_fuse_qk_rope_cache()
+            and getattr(self.mla_attn.impl, "always_fuses_qk_rope_cache", False)
         )
         indexer_op = getattr(self.indexer, "indexer_op", None)
         if indexer_op is not None and hasattr(
@@ -224,7 +232,10 @@ def forward(
             heads *= q_proj_layer.group_size
         q = q.view(-1, heads, self.qk_head_dim)
 
-        if self.rotary_emb is not None:
+        fused_positions: torch.Tensor | None = None
+        if self._defer_rope_to_fused_kernel:
+            fused_positions = positions
+        elif self.rotary_emb is not None:
             q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb(
                 positions, q[..., self.qk_nope_head_dim :], k_pe
             )
@@ -247,6 +258,7 @@ def forward(
             k_pe,
             output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim),
             q_dcp_replicated=q_dcp_replicated,
+            positions=fused_positions,
         )
 
         if self.g_proj is not None:

vllm/models/deepseek_v32/amd/rocm.py

@@ -223,6 +223,7 @@ def _fused_attention(
             mla_kv_cache = self.kv_cache
             mla_k_scale = self._k_scale
 
+        assert self.rotary_emb is not None
         q_c = fused_norm_rope(
             positions,
             q_c,

vllm/models/deepseek_v32/attention.py

@@ -352,6 +352,7 @@ def forward(  # type: ignore[override]
             mla_kv_cache = None if hisparse_cache is not None else self.kv_cache
             mla_k_scale = self._k_scale
 
+        assert self.rotary_emb is not None
         if self.use_pcp or hisparse_cache is None:
             kv_c_out = torch.empty_like(kv_c)
             k_pe_out = torch.empty_like(k_pe)
@@ -531,6 +532,7 @@ def _sparse_indexer_and_attn(
 
         if self._use_sparse_mha(attn_metadata):
             assert kv_c is not None and k_pe is not None
+            assert self.rotary_emb is not None
             mha_q_pe = self.rotary_emb(positions, q_pe)[0] if self._fp8_query else mqa_q
             mha_q = torch.cat((q_nope, mha_q_pe), dim=-1)
             self.forward_impl(