[ROCm][Kimi-K3][Perf] Fuse MLA decode KV-cache write and Q-prep via AITER (#57640)
Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Rita Brugarolas committed
Sep 30, 2026 at 10:20 UTC
8f24ab30a36d62386bb0d30d2aa1c689806604af
1 file changed
+170
-8
vllm/models/kimi_k3/amd/mla.py
+170
-8
@@ -7,16 +7,47 @@ from typing import cast
7
import torch
8
9
from vllm._aiter_ops import rocm_aiter_ops
10
+from vllm.model_executor.layers.attention.attention import get_attention_context
11
from vllm.model_executor.layers.layernorm import RMSNorm
12
from vllm.model_executor.layers.mla import MultiHeadLatentAttentionWrapper
13
+from vllm.platforms import current_platform
14
+
15
+_OPT_KV_LORA_RANK = 512
16
+_OPT_ROT_DIM = 64
17
+_OPT_MIN_SIZE = 2048
18
19
20
class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper):
15
- """Kimi-K3 MLA wrapper with eager AITER q/kv RMSNorm fusion."""
21
+ """Kimi-K3 MLA wrapper with eager AITER q/kv RMSNorm fusion and a fused
22
+ decode Q-prep path."""
23
24
def __init__(self, *args, **kwargs) -> None:
25
super().__init__(*args, **kwargs)
26
self._use_eager_qk_rmsnorm_fusion = bool(rocm_aiter_ops.is_enabled())
27
+ self._fused_qk_prep = self._fused_qk_prep_supported()
28
+ self._identity_rope: tuple[torch.Tensor, torch.Tensor] | None = None
29
+
30
+ def _fused_qk_prep_supported(self) -> bool:
31
+ attn = self.mla_attn
32
+ if not (
33
+ rocm_aiter_ops.is_mla_enabled()
34
+ and attn.kv_cache_dtype.startswith("fp8")
35
+ and self.kv_lora_rank == _OPT_KV_LORA_RANK
36
+ and self.qk_rope_head_dim == _OPT_ROT_DIM
37
+ and self.kv_lora_rank * self.num_heads >= _OPT_MIN_SIZE
38
+ ):
39
+ return False
40
+ # The fused kernel writes q_out only for slot >= 0, so padded rows would
41
+ # be garbage after a DCP query all-gather. PCP and HiSparse rewrite the
42
+ # KV inputs/slots before the standalone write, which this path bypasses.
43
+ if attn.impl.dcp_world_size > 1 or attn.use_pcp:
44
+ return False
45
+ if attn.hisparse_cache is not None or attn.q_pad_num_heads is not None:
46
+ return False
47
+ # Only the W_UK variants reproduced in _fused_decode are supported.
48
+ return attn.is_aiter_triton_fp4_bmm_enabled or (
49
+ not attn.is_aiter_triton_fp8_bmm_enabled and attn.W_UK_T is not None
50
+ )
51
52
def _normalize_q_kv(
53
self,
@@ -38,6 +69,107 @@ class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper):
69
70
return q_layernorm(q_c), kv_layernorm(kv_c)
71
72
+ def _fused_decode(
73
+ self,
74
+ q: torch.Tensor,
75
+ kv_c_normed: torch.Tensor,
76
+ k_pe: torch.Tensor,
77
+ positions: torch.Tensor,
78
+ output_shape: tuple[int, int],
79
+ attn_metadata,
80
+ layer,
81
+ kv_cache: torch.Tensor,
82
+ slot_mapping: torch.Tensor,
83
+ ) -> torch.Tensor:
84
+ """Decode-only MLA through one fused AITER launch.
85
+
86
+ Replaces the fp8 KV-cache write plus the ``[ql_nope | q_pe]`` concat and
87
+ that query's static fp8 quant with a single
88
+ ``fused_qk_rope_concat_and_cache_mla``. On an untouched upstream build
89
+ the concat and quant already arrive as one inductor-fused kernel, so
90
+ this collapses two launches into one.
91
+ """
92
+ from aiter.ops.cache import fused_qk_rope_concat_and_cache_mla
93
+
94
+ attn = self.mla_attn
95
+ q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
96
+ # (B, N, P) -> (N, B, P)
97
+ q_nope_t = q_nope.transpose(0, 1)
98
+
99
+ # Mirror MLAAttention.forward_impl's W_UK dispatch, *including its
100
+ # precedence* -- fp4 is checked before fp8 there, and both flags can be
101
+ # set at once. ql_nope must come out bit-identical to what the unfused
102
+ # path feeds its concat+quant, because this only replaces what happens
103
+ # after the bmm. Both variants return bf16, which the kernel consumes.
104
+ if attn.is_aiter_triton_fp4_bmm_enabled:
105
+ from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
106
+
107
+ ql_nope = batched_gemm_a16wfp4(
108
+ q_nope_t,
109
+ attn.W_K,
110
+ attn.W_K_scale,
111
+ transpose_bm=True,
112
+ prequant=True,
113
+ y_scale=layer._q_scale,
114
+ )
115
+ else:
116
+ assert attn.W_UK_T is not None # guaranteed by _fused_qk_prep_supported
117
+ B, N = q_nope.shape[0], q_nope.shape[1]
118
+ L = attn.W_UK_T.shape[-1]
119
+ ql_nope = q_nope_t.new_empty((B, N, L))
120
+ torch.bmm(q_nope_t, attn.W_UK_T, out=ql_nope.transpose(0, 1))
121
+
122
+ if self._identity_rope is None:
123
+ # cos = 1, sin = 0 makes the kernel's RoPE the identity, which is
124
+ # what a NoPE model needs. One row suffices because the `_opt`
125
+ # kernel clamps `pos` into [0, cos_cache.size(0)); a full-length
126
+ # cache would cost ~268 MB at K3's 1M max_position_embeddings.
127
+ half = self.qk_rope_head_dim // 2
128
+ opts = {"dtype": kv_c_normed.dtype, "device": kv_c_normed.device}
129
+ self._identity_rope = (
130
+ torch.ones(1, half, **opts),
131
+ torch.zeros(1, half, **opts),
132
+ )
133
+ cos_cache, sin_cache = self._identity_rope
134
+
135
+ # An fp8 KV cache is allocated as uint8 and re-viewed as fp8 before use;
136
+ # the AITER kernel rejects the raw uint8 dtype outright.
137
+ fp8_dtype = current_platform.fp8_dtype()
138
+ if kv_cache.dtype != fp8_dtype:
139
+ kv_cache = kv_cache.view(fp8_dtype)
140
+
141
+ head_size = attn.kv_lora_rank + self.qk_rope_head_dim
142
+ num_tokens, num_heads = ql_nope.shape[:2]
143
+ q_out = torch.empty(
144
+ (num_tokens, num_heads, head_size),
145
+ dtype=fp8_dtype,
146
+ device=ql_nope.device,
147
+ )
148
+ fused_qk_rope_concat_and_cache_mla(
149
+ # the W_UK bmm hands back a transposed view; the kernel reads dense
150
+ ql_nope.contiguous(),
151
+ q_pe,
152
+ kv_c_normed,
153
+ k_pe.squeeze(1),
154
+ kv_cache.view(kv_cache.shape[0], -1, head_size),
155
+ q_out,
156
+ slot_mapping.flatten(),
157
+ layer._k_scale,
158
+ layer._q_scale,
159
+ positions,
160
+ cos_cache,
161
+ sin_cache,
162
+ is_neox=True,
163
+ is_nope_first=True,
164
+ )
165
+
166
+ # forward_mqa is MLA-specific and not on AttentionImpl; MLAAttention
167
+ # annotates its own call to it the same way.
168
+ attn_out, _ = attn.impl.forward_mqa(q_out, kv_cache, attn_metadata, attn) # type: ignore[attr-defined]
169
+ output = q.new_empty(output_shape)
170
+ attn._v_up_proj(attn_out, out=output)
171
+ return output
172
+
173
def forward(
174
self,
175
positions: torch.Tensor,
@@ -106,13 +238,43 @@ class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper):
238
if self.dcp_q_replicate:
239
q_dcp_replicated, q = q, q_proj_layer._local_view(q)
240
109
- attn_out = self.mla_attn(
110
- q,
111
- kv_c_normed,
112
- k_pe,
113
- output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim),
114
- q_dcp_replicated=q_dcp_replicated,
115
- )
241
+ output_shape = (hidden_states.shape[0], self.num_heads * self.v_head_dim)
242
+
243
+ # The fused decode path covers only the MQA slice, so it is taken for
244
+ # decode-only batches; anything else falls through to MLAAttention.
245
+ attn_metadata = layer = kv_cache = slot_mapping = None
246
+ fuse = self._fused_qk_prep and q_dcp_replicated is None
247
+ if fuse:
248
+ attn_metadata, layer, kv_cache, slot_mapping = get_attention_context(
249
+ self.mla_attn.layer_name
250
+ )
251
+ fuse = (
252
+ attn_metadata is not None
253
+ and slot_mapping is not None
254
+ and kv_cache.numel() > 0
255
+ and attn_metadata.num_actual_tokens == attn_metadata.num_decode_tokens
256
+ )
257
+
258
+ if fuse:
259
+ attn_out = self._fused_decode(
260
+ q,
261
+ kv_c_normed,
262
+ k_pe,
263
+ positions,
264
+ output_shape,
265
+ attn_metadata,
266
+ layer,
267
+ kv_cache,
268
+ slot_mapping,
269
+ )
270
+ else:
271
+ attn_out = self.mla_attn(
272
+ q,
273
+ kv_c_normed,
274
+ k_pe,
275
+ output_shape=output_shape,
276
+ q_dcp_replicated=q_dcp_replicated,
277
+ )
278
279
if self.g_proj is not None:
280
attn_out = attn_out * self.g_proj(hidden_states)[0].sigmoid()