From 4e333f070b54e24e91021dc1d3bf070b9d494c7e Mon Sep 17 00:00:00 2001 From: Rita Brugarolas Brufau Date: Fri, 18 Sep 2026 16:32:04 +0000 Subject: [PATCH 1/2] [ROCm][Kimi-K3][Perf] Fuse MLA decode KV-cache write and Q-prep via AITER On a stock build, Kimi-K3 MLA decode runs two kernels per layer per step: the fp8 KV-cache write (`vllm::concat_and_cache_mla_kernel`) and the `[ql_nope | q_pe]` concat plus that query's static fp8 quant, which the `mla_decode_concat_quant_fp8` custom op emits as a single `triton_poi_fused__to_copy_cat_clamp_mul_reciprocal_view` kernel. Collapse both into one AITER `fused_qk_rope_concat_and_cache_mla` launch. Measured on MI355X at TP=8, plain `vllm serve` on the stock upstream nightly with no speculative decoding and no compilation-config override: 288 -> 144 launches and 1351.8 -> 682.8 us per rank over 6 decode steps, i.e. 9.39 -> 4.74 us per MLA layer per decode step. End-to-end serving throughput is unchanged within measurement noise and this change does not claim a throughput win. K3 MLA is NoPE, while the kernel always applies RoPE and takes the cos/sin caches as required arguments. Feed it an identity rotation (cos=1, sin=0) so the rotation is a no-op -- the approach ROCm/ATOM uses for K3. The cache is a single row: the `_opt` decode kernel clamps `pos` into `[0, cos_cache.size(0))`, so real positions pass through untouched and K3's 1M max_position_embeddings does not cost ~268 MB to express a no-op. AITER's general decode kernel does not clamp, so the head counts that would select it are excluded. Kept entirely in `vllm/models/kimi_k3/amd/mla.py`, mirroring the q/kv RMSNorm fusion this wrapper already owns (#52494); no shared MLA code is touched, so this carries no dependency on #47757 and cannot affect any other model or backend. Everything that is fixed at load time -- platform, KV dtype, shapes, DCP/PCP/HiSparse, head padding and the W_UK variant -- is resolved once in `_fused_qk_prep_supported()`, so the only per-batch condition is that the batch is decode-only; anything else falls through to the stock `MLAAttention` path. Signed-off-by: Rita Brugarolas Brufau Co-Authored-By: Claude Signed-off-by: Rita Brugarolas Brufau --- vllm/models/kimi_k3/amd/mla.py | 196 +++++++++++++++++++++++++++++++-- 1 file changed, 188 insertions(+), 8 deletions(-) diff --git a/vllm/models/kimi_k3/amd/mla.py b/vllm/models/kimi_k3/amd/mla.py index e76c0e80ad45..4db85523b12e 100644 --- a/vllm/models/kimi_k3/amd/mla.py +++ b/vllm/models/kimi_k3/amd/mla.py @@ -7,16 +7,65 @@ import torch from vllm._aiter_ops import rocm_aiter_ops +from vllm.model_executor.layers.attention.attention import get_attention_context from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.mla import MultiHeadLatentAttentionWrapper +from vllm.platforms import current_platform + +# AITER dispatch constants, mirrored from csrc/kernels/cache_kernels.cu. With all +# three satisfied the `_opt` decode kernel is selected, which clamps `pos` into +# the cos/sin cache; below MIN_SIZE the *general* kernel is chosen, which does +# not clamp and would read past the single-row identity cache below. At TP=8 K3 +# has 12 heads/rank so 512 * 12 = 6144 clears it; at TP=32 it would be 3 heads. +_OPT_KV_LORA_RANK = 512 +_OPT_ROT_DIM = 64 +_OPT_MIN_SIZE = 2048 class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper): - """Kimi-K3 MLA wrapper with eager AITER q/kv RMSNorm fusion.""" + """Kimi-K3 MLA wrapper with eager AITER q/kv RMSNorm fusion and a fused + decode Q-prep path. + + Both fusions are kept here rather than in the shared MLA layer so only AMD + Kimi-K3 changes, mirroring the q/kv RMSNorm fusion this class already owns. + """ def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self._use_eager_qk_rmsnorm_fusion = bool(rocm_aiter_ops.is_enabled()) + self._fused_qk_prep = self._fused_qk_prep_supported() + self._identity_rope: tuple[torch.Tensor, torch.Tensor] | None = None + + def _fused_qk_prep_supported(self) -> bool: + """Whether this layer can ever take the fused decode path. + + Everything checked here is fixed at load time, so the per-batch test in + ``forward`` is only about the shape of the batch itself. + """ + attn = self.mla_attn + if not ( + current_platform.is_rocm() + and rocm_aiter_ops.is_mla_enabled() + and attn.kv_cache_dtype.startswith("fp8") + and attn.kv_cache_dtype != "fp8_ds_mla" + and self.kv_lora_rank == _OPT_KV_LORA_RANK + and self.qk_rope_head_dim == _OPT_ROT_DIM + and self.kv_lora_rank * self.num_heads >= _OPT_MIN_SIZE + ): + return False + # The fused kernel writes q_out only for slot >= 0, so padded rows would + # be garbage after a DCP query all-gather. PCP and HiSparse rewrite the + # KV inputs/slots before the standalone write, which this path bypasses. + if attn.impl.dcp_world_size > 1 or attn.use_pcp: + return False + if attn.hisparse_cache is not None or attn.q_pad_num_heads is not None: + return False + # Only the W_UK variants reproduced in _fused_decode are supported. + return attn.is_aiter_triton_fp4_bmm_enabled or ( + not attn.is_aiter_triton_fp8_bmm_enabled + and not attn.is_amx_bmm_enabled + and attn.W_UK_T is not None + ) def _normalize_q_kv( self, @@ -38,6 +87,107 @@ def _normalize_q_kv( return q_layernorm(q_c), kv_layernorm(kv_c) + def _fused_decode( + self, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor, + output_shape: tuple[int, int], + attn_metadata, + layer, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> torch.Tensor: + """Decode-only MLA through one fused AITER launch. + + Replaces the fp8 KV-cache write plus the ``[ql_nope | q_pe]`` concat and + that query's static fp8 quant with a single + ``fused_qk_rope_concat_and_cache_mla``. On an untouched upstream build + the concat and quant already arrive as one inductor-fused kernel, so + this collapses two launches into one. + """ + from aiter.ops.cache import fused_qk_rope_concat_and_cache_mla + + attn = self.mla_attn + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + # (B, N, P) -> (N, B, P) + q_nope_t = q_nope.transpose(0, 1) + + # Mirror MLAAttention.forward_impl's W_UK dispatch, *including its + # precedence* -- fp4 is checked before fp8 there, and both flags can be + # set at once. ql_nope must come out bit-identical to what the unfused + # path feeds its concat+quant, because this only replaces what happens + # after the bmm. Both variants return bf16, which the kernel consumes. + if attn.is_aiter_triton_fp4_bmm_enabled: + from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4 + + ql_nope = batched_gemm_a16wfp4( + q_nope_t, + attn.W_K, + attn.W_K_scale, + transpose_bm=True, + prequant=True, + y_scale=layer._q_scale, + ) + else: + assert attn.W_UK_T is not None # guaranteed by _fused_qk_prep_supported + B, N = q_nope.shape[0], q_nope.shape[1] + L = attn.W_UK_T.shape[-1] + ql_nope = q_nope_t.new_empty((B, N, L)) + torch.bmm(q_nope_t, attn.W_UK_T, out=ql_nope.transpose(0, 1)) + + if self._identity_rope is None: + # cos = 1, sin = 0 makes the kernel's RoPE the identity, which is + # what a NoPE model needs. One row suffices because the `_opt` + # kernel clamps `pos` into [0, cos_cache.size(0)); a full-length + # cache would cost ~268 MB at K3's 1M max_position_embeddings. + half = self.qk_rope_head_dim // 2 + opts = {"dtype": kv_c_normed.dtype, "device": kv_c_normed.device} + self._identity_rope = ( + torch.ones(1, half, **opts), + torch.zeros(1, half, **opts), + ) + cos_cache, sin_cache = self._identity_rope + + # An fp8 KV cache is allocated as uint8 and re-viewed as fp8 before use; + # the AITER kernel rejects the raw uint8 dtype outright. + fp8_dtype = current_platform.fp8_dtype() + if kv_cache.dtype != fp8_dtype: + kv_cache = kv_cache.view(fp8_dtype) + + head_size = attn.kv_lora_rank + self.qk_rope_head_dim + num_tokens, num_heads = ql_nope.shape[:2] + q_out = torch.empty( + (num_tokens, num_heads, head_size), + dtype=fp8_dtype, + device=ql_nope.device, + ) + fused_qk_rope_concat_and_cache_mla( + # the W_UK bmm hands back a transposed view; the kernel reads dense + ql_nope.contiguous(), + q_pe, + kv_c_normed, + k_pe.squeeze(1), + kv_cache.view(kv_cache.shape[0], -1, head_size), + q_out, + slot_mapping.flatten(), + layer._k_scale, + layer._q_scale, + positions, + cos_cache, + sin_cache, + is_neox=True, + is_nope_first=True, + ) + + # forward_mqa is MLA-specific and not on AttentionImpl; MLAAttention + # annotates its own call to it the same way. + attn_out, _ = attn.impl.forward_mqa(q_out, kv_cache, attn_metadata, attn) # type: ignore[attr-defined] + output = q.new_empty(output_shape) + attn._v_up_proj(attn_out, out=output) + return output + def forward( self, positions: torch.Tensor, @@ -106,13 +256,43 @@ def forward( if self.dcp_q_replicate: q_dcp_replicated, q = q, q_proj_layer._local_view(q) - attn_out = self.mla_attn( - q, - kv_c_normed, - k_pe, - output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim), - q_dcp_replicated=q_dcp_replicated, - ) + output_shape = (hidden_states.shape[0], self.num_heads * self.v_head_dim) + + # The fused decode path covers only the MQA slice, so it is taken for + # decode-only batches; anything else falls through to MLAAttention. + attn_metadata = layer = kv_cache = slot_mapping = None + fuse = self._fused_qk_prep and q_dcp_replicated is None + if fuse: + attn_metadata, layer, kv_cache, slot_mapping = get_attention_context( + self.mla_attn.layer_name + ) + fuse = ( + attn_metadata is not None + and slot_mapping is not None + and kv_cache.numel() > 0 + and attn_metadata.num_actual_tokens == attn_metadata.num_decode_tokens + ) + + if fuse: + attn_out = self._fused_decode( + q, + kv_c_normed, + k_pe, + positions, + output_shape, + attn_metadata, + layer, + kv_cache, + slot_mapping, + ) + else: + attn_out = self.mla_attn( + q, + kv_c_normed, + k_pe, + output_shape=output_shape, + q_dcp_replicated=q_dcp_replicated, + ) if self.g_proj is not None: attn_out = attn_out * self.g_proj(hidden_states)[0].sigmoid() From 4b6ee759f2e648134100a1e2e3f919b3ee55c2e6 Mon Sep 17 00:00:00 2001 From: Rita Brugarolas Brufau Date: Wed, 23 Sep 2026 19:11:56 +0000 Subject: [PATCH 2/2] [Kimi-K3][AMD] Address review: drop unreachable gates and comments Per review on #57640: - remove the three explanatory comment blocks - drop `current_platform.is_rocm()`: kimi_k3/__init__.py only imports the amd/ wrapper under is_rocm(), so it cannot be instantiated elsewhere - drop the `fp8_ds_mla` check: only FLASHMLA_SPARSE and FLASHINFER_MLA_SPARSE_SM120 promote to that dtype, both NVIDIA-only - drop `is_amx_bmm_enabled`: set only by AMXMLAImpl, which is CPU-only All four gates are unreachable on this path, so the fused path's behavior is unchanged. Verified by kernel census on 8x MI355X TP=8: the fused kernel still fires 144x (24 MLA layers x 6 decode steps) with zero concat_and_cache_mla and zero triton_poi_fused, matching the pre-change build. Co-authored-by: Claude Signed-off-by: Rita Brugarolas Brufau --- vllm/models/kimi_k3/amd/mla.py | 24 +++--------------------- 1 file changed, 3 insertions(+), 21 deletions(-) diff --git a/vllm/models/kimi_k3/amd/mla.py b/vllm/models/kimi_k3/amd/mla.py index 4db85523b12e..11df7d035f31 100644 --- a/vllm/models/kimi_k3/amd/mla.py +++ b/vllm/models/kimi_k3/amd/mla.py @@ -12,11 +12,6 @@ from vllm.model_executor.layers.mla import MultiHeadLatentAttentionWrapper from vllm.platforms import current_platform -# AITER dispatch constants, mirrored from csrc/kernels/cache_kernels.cu. With all -# three satisfied the `_opt` decode kernel is selected, which clamps `pos` into -# the cos/sin cache; below MIN_SIZE the *general* kernel is chosen, which does -# not clamp and would read past the single-row identity cache below. At TP=8 K3 -# has 12 heads/rank so 512 * 12 = 6144 clears it; at TP=32 it would be 3 heads. _OPT_KV_LORA_RANK = 512 _OPT_ROT_DIM = 64 _OPT_MIN_SIZE = 2048 @@ -24,11 +19,7 @@ class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper): """Kimi-K3 MLA wrapper with eager AITER q/kv RMSNorm fusion and a fused - decode Q-prep path. - - Both fusions are kept here rather than in the shared MLA layer so only AMD - Kimi-K3 changes, mirroring the q/kv RMSNorm fusion this class already owns. - """ + decode Q-prep path.""" def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) @@ -37,17 +28,10 @@ def __init__(self, *args, **kwargs) -> None: self._identity_rope: tuple[torch.Tensor, torch.Tensor] | None = None def _fused_qk_prep_supported(self) -> bool: - """Whether this layer can ever take the fused decode path. - - Everything checked here is fixed at load time, so the per-batch test in - ``forward`` is only about the shape of the batch itself. - """ attn = self.mla_attn if not ( - current_platform.is_rocm() - and rocm_aiter_ops.is_mla_enabled() + rocm_aiter_ops.is_mla_enabled() and attn.kv_cache_dtype.startswith("fp8") - and attn.kv_cache_dtype != "fp8_ds_mla" and self.kv_lora_rank == _OPT_KV_LORA_RANK and self.qk_rope_head_dim == _OPT_ROT_DIM and self.kv_lora_rank * self.num_heads >= _OPT_MIN_SIZE @@ -62,9 +46,7 @@ def _fused_qk_prep_supported(self) -> bool: return False # Only the W_UK variants reproduced in _fused_decode are supported. return attn.is_aiter_triton_fp4_bmm_enabled or ( - not attn.is_aiter_triton_fp8_bmm_enabled - and not attn.is_amx_bmm_enabled - and attn.W_UK_T is not None + not attn.is_aiter_triton_fp8_bmm_enabled and attn.W_UK_T is not None ) def _normalize_q_kv(