[AMD][kimik3][ROCm][Perf] Fuse MLA q/kv RMSNorm in AMD Kimi-K3 MLA wrapper (#52494)
Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
Rita Brugarolas committed
Sep 4, 2026 at 03:14 UTC
3ff4f02dfe69abc1a0375d1ea8d8d5cb25609fcc
2 files changed
+123
-2
vllm/models/kimi_k3/amd/linear.py
+3
-2
@@ -36,7 +36,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
36
MambaStateDtypeCalculator,
37
MambaStateShapeCalculator,
38
)
39
-from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
39
+from vllm.model_executor.layers.mla import MLAModules
40
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
41
from vllm.model_executor.layers.vocab_parallel_embedding import (
42
ParallelLMHead,
@@ -63,6 +63,7 @@ from vllm.model_executor.models.utils import (
63
)
64
from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention
65
from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner
66
+from vllm.models.kimi_k3.amd.mla import KimiK3MultiHeadLatentAttentionWrapper
67
from vllm.models.kimi_k3.amd.ops.attn_res import attn_res
68
from vllm.sequence import IntermediateTensors
69
from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
@@ -438,7 +439,7 @@ class KimiMLAAttention(nn.Module):
439
topk_indices_buffer=None,
440
g_proj=getattr(self, "g_proj", None),
441
)
441
- self.mla_attn = MultiHeadLatentAttentionWrapper(
442
+ self.mla_attn = KimiK3MultiHeadLatentAttentionWrapper(
443
self.hidden_size,
444
self.num_local_heads,
445
self.scaling,
vllm/models/kimi_k3/amd/mla.py
new
+120
@@ -0,0 +1,120 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""AMD-specific MLA wrapper for Kimi-K3."""
4
+
5
+from typing import cast
6
+
7
+import torch
8
+
9
+from vllm._aiter_ops import rocm_aiter_ops
10
+from vllm.model_executor.layers.layernorm import RMSNorm
11
+from vllm.model_executor.layers.mla import MultiHeadLatentAttentionWrapper
12
+
13
+
14
+class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper):
15
+ """Kimi-K3 MLA wrapper with eager AITER q/kv RMSNorm fusion."""
16
+
17
+ def __init__(self, *args, **kwargs) -> None:
18
+ super().__init__(*args, **kwargs)
19
+ self._use_eager_qk_rmsnorm_fusion = bool(rocm_aiter_ops.is_enabled())
20
+
21
+ def _normalize_q_kv(
22
+ self,
23
+ q_c: torch.Tensor,
24
+ kv_c: torch.Tensor,
25
+ ) -> tuple[torch.Tensor, torch.Tensor]:
26
+ q_layernorm = cast(RMSNorm, self.q_a_layernorm)
27
+ kv_layernorm = cast(RMSNorm, self.kv_a_layernorm)
28
+
29
+ if self._use_eager_qk_rmsnorm_fusion and not torch.compiler.is_compiling():
30
+ return torch.ops.vllm.fused_mla_dual_rms_norm(
31
+ q_c,
32
+ q_layernorm.weight,
33
+ kv_c,
34
+ kv_layernorm.weight,
35
+ q_layernorm.variance_epsilon,
36
+ kv_layernorm.variance_epsilon,
37
+ )
38
+
39
+ return q_layernorm(q_c), kv_layernorm(kv_c)
40
+
41
+ def forward(
42
+ self,
43
+ positions: torch.Tensor,
44
+ hidden_states: torch.Tensor,
45
+ llama_4_scaling: torch.Tensor | None = None,
46
+ ) -> torch.Tensor:
47
+ q_c = None
48
+
49
+ if self.q_lora_rank is not None:
50
+ assert self.fused_qkv_a_proj is not None, (
51
+ "fused_qkv_a_proj is required when q_lora_rank is not None"
52
+ )
53
+ assert self.q_a_layernorm is not None, (
54
+ "q_a_layernorm is required when q_lora_rank is not None"
55
+ )
56
+ assert self.q_b_proj is not None, (
57
+ "q_b_proj is required when q_lora_rank is not None"
58
+ )
59
+
60
+ qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
61
+ q_c, kv_lora = qkv_lora.split(
62
+ [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
63
+ dim=-1,
64
+ )
65
+ kv_c, k_pe = kv_lora.split(
66
+ [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
67
+ )
68
+ q_proj_input, kv_c_normed = self._normalize_q_kv(q_c, kv_c)
69
+ q_proj_layer = self.q_b_proj
70
+ else:
71
+ assert self.kv_a_proj_with_mqa is not None, (
72
+ "kv_a_proj_with_mqa is required when q_lora_rank is None"
73
+ )
74
+ assert self.q_proj is not None, (
75
+ "q_proj is required when q_lora_rank is None"
76
+ )
77
+ kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]
78
+ kv_c, k_pe = kv_lora.split(
79
+ [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
80
+ )
81
+ kv_c_normed = self.kv_a_layernorm(kv_c)
82
+ q_proj_layer = self.q_proj
83
+ q_proj_input = hidden_states
84
+
85
+ # Add head dim of 1 to k_pe.
86
+ k_pe = k_pe.unsqueeze(1)
87
+
88
+ q = q_proj_layer(q_proj_input)[0]
89
+ heads = self.num_heads
90
+ if self.dcp_q_replicate:
91
+ heads *= q_proj_layer.group_size
92
+ q = q.view(-1, heads, self.qk_head_dim)
93
+
94
+ if self.rotary_emb is not None:
95
+ q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb(
96
+ positions, q[..., self.qk_nope_head_dim :], k_pe
97
+ )
98
+
99
+ if self.indexer and self.is_sparse and not self.skip_topk:
100
+ self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb)
101
+
102
+ if llama_4_scaling is not None:
103
+ q *= llama_4_scaling
104
+
105
+ q_dcp_replicated = None
106
+ if self.dcp_q_replicate:
107
+ q_dcp_replicated, q = q, q_proj_layer._local_view(q)
108
+
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
+ )
116
+
117
+ if self.g_proj is not None:
118
+ attn_out = attn_out * self.g_proj(hidden_states)[0].sigmoid()
119
+
120
+ return self.o_proj(attn_out)[0]