[Kimi-K3][AMD] Return KDA and MLA projection outputs directly (#50592)
Signed-off-by: Liuyinfeng01 <yinfeliu@amd.com> Co-authored-by: Liuyinfeng01 <199041580+LiuYinfeng01@users.noreply.github.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
yinfengLiu committed
Sep 22, 2026 at 04:44 UTC
382970ee6ca490aeaaaf4e32c53695b581ff61ba
7 files changed
+170
-33
tests/models/kimi_k3/test_amd_kda_direct_return.py
new
+72
@@ -0,0 +1,72 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import torch
7
+from torch import nn
8
+
9
+from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention
10
+from vllm.models.kimi_k3.amd.linear import KimiDecoderLayer
11
+
12
+
13
+class _ReturningAttention(nn.Module):
14
+ def __init__(self, result: torch.Tensor) -> None:
15
+ super().__init__()
16
+ self.result = result
17
+
18
+ def forward(
19
+ self, hidden_states: torch.Tensor, positions: torch.Tensor
20
+ ) -> torch.Tensor:
21
+ return self.result
22
+
23
+
24
+class _TupleProjection(nn.Module):
25
+ def __init__(self, result: torch.Tensor) -> None:
26
+ super().__init__()
27
+ self.result = result
28
+
29
+ def forward(self, _: torch.Tensor) -> tuple[torch.Tensor, None]:
30
+ return self.result, None
31
+
32
+
33
+def _make_layer(self_attn: nn.Module) -> KimiDecoderLayer:
34
+ layer = object.__new__(KimiDecoderLayer)
35
+ nn.Module.__init__(layer)
36
+ layer.self_attn = self_attn
37
+ return layer
38
+
39
+
40
+def test_kda_self_attention_returns_projection_storage_directly() -> None:
41
+ hidden_states = torch.randn(4, 8)
42
+ projected = torch.randn_like(hidden_states)
43
+ layer = _make_layer(_ReturningAttention(projected))
44
+
45
+ output = layer._run_self_attn(torch.arange(4), hidden_states)
46
+
47
+ assert output is projected
48
+ assert output.data_ptr() == projected.data_ptr()
49
+
50
+
51
+def test_k3_kda_forward_returns_output_projection_storage() -> None:
52
+ hidden_states = torch.randn(4, 8)
53
+ projected = torch.randn_like(hidden_states)
54
+
55
+ def fill_core_output(**kwargs: torch.Tensor) -> None:
56
+ kwargs["core_attn_out"].fill_(1)
57
+
58
+ attention = SimpleNamespace(
59
+ local_projection_size=4,
60
+ head_dim=2,
61
+ local_num_heads=2,
62
+ in_proj_padding=0,
63
+ in_proj_qkvgfab=_TupleProjection(torch.randn(4, 20)),
64
+ f_b_proj=_TupleProjection(torch.randn(4, 4)),
65
+ o_proj=_TupleProjection(projected),
66
+ _forward=fill_core_output,
67
+ )
68
+
69
+ output = KimiK3DeltaAttention.forward(attention, hidden_states, torch.arange(4))
70
+
71
+ assert output is projected
72
+ assert output.data_ptr() == projected.data_ptr()
tests/models/kimi_k3/test_amd_mla_direct_return.py
new
+38
@@ -0,0 +1,38 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+from torch import nn
6
+
7
+from vllm.models.kimi_k3.amd.linear import KimiDecoderLayer, KimiMLAAttention
8
+
9
+
10
+class _DirectMLA(KimiMLAAttention):
11
+ def __init__(self, result: torch.Tensor) -> None:
12
+ nn.Module.__init__(self)
13
+ self.result = result
14
+
15
+ def forward(
16
+ self,
17
+ positions: torch.Tensor,
18
+ hidden_states: torch.Tensor,
19
+ ) -> torch.Tensor:
20
+ return self.result
21
+
22
+
23
+def _make_layer(self_attn: nn.Module) -> KimiDecoderLayer:
24
+ layer = object.__new__(KimiDecoderLayer)
25
+ nn.Module.__init__(layer)
26
+ layer.self_attn = self_attn
27
+ return layer
28
+
29
+
30
+def test_mla_self_attention_returns_projection_storage_directly() -> None:
31
+ hidden_states = torch.randn(4, 8)
32
+ projected = torch.randn_like(hidden_states)
33
+ layer = _make_layer(_DirectMLA(projected))
34
+
35
+ output = layer._run_self_attn(torch.arange(4), hidden_states)
36
+
37
+ assert output is projected
38
+ assert output.data_ptr() == projected.data_ptr()
tests/models/kimi_k3/test_nvidia_kda_direct_return.py
new
+31
@@ -0,0 +1,31 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+from torch import nn
6
+
7
+from vllm.models.kimi_k3.nvidia.model import KimiDecoderLayer
8
+
9
+
10
+class _ReturningSharedKDA(nn.Module):
11
+ def __init__(self, result: torch.Tensor) -> None:
12
+ super().__init__()
13
+ self.result = result
14
+
15
+ def forward(
16
+ self, hidden_states: torch.Tensor, positions: torch.Tensor
17
+ ) -> torch.Tensor:
18
+ return self.result
19
+
20
+
21
+def test_low_rank_kda_caller_returns_projection_storage_directly() -> None:
22
+ hidden_states = torch.randn(4, 8)
23
+ projected = torch.randn_like(hidden_states)
24
+ layer = object.__new__(KimiDecoderLayer)
25
+ nn.Module.__init__(layer)
26
+ layer.self_attn = _ReturningSharedKDA(projected)
27
+
28
+ output = layer._run_self_attn(torch.arange(4), hidden_states)
29
+
30
+ assert output is projected
31
+ assert output.data_ptr() == projected.data_ptr()
vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
+2
-3
@@ -329,8 +329,7 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
329
self,
330
hidden_states: torch.Tensor,
331
positions: torch.Tensor,
332
- output: torch.Tensor,
333
- ) -> None:
332
+ ) -> torch.Tensor:
333
num_tokens = hidden_states.size(0)
334
projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
335
if self.use_full_rank_gate:
@@ -375,7 +374,7 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
374
core_attn_out=core_attn_out,
375
)
376
core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
378
- output[:] = self.o_proj(core_attn_out)[0]
377
+ return self.o_proj(core_attn_out)[0]
378
379
@eager_break_during_capture
380
def _forward(
vllm/models/kimi_k3/amd/kda.py
+2
-3
@@ -271,8 +271,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
271
self,
272
hidden_states: torch.Tensor,
273
positions: torch.Tensor,
274
- output: torch.Tensor,
275
- ) -> None:
274
+ ) -> torch.Tensor:
275
num_tokens = hidden_states.size(0)
276
projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
277
@@ -306,7 +305,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
305
core_attn_out=core_attn_out,
306
)
307
core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
309
- output[:] = self.o_proj(core_attn_out)[0]
308
+ return self.o_proj(core_attn_out)[0]
309
310
@eager_break_during_capture
311
def _forward(
vllm/models/kimi_k3/amd/linear.py
+25
-16
@@ -456,9 +456,8 @@ class KimiMLAAttention(nn.Module):
456
self,
457
positions: torch.Tensor,
458
hidden_states: torch.Tensor,
459
- output: torch.Tensor,
460
- ) -> None:
461
- output[:] = self.mla_attn(positions, hidden_states)
459
+ ) -> torch.Tensor:
460
+ return self.mla_attn(positions, hidden_states)
461
462
463
class KimiDecoderLayer(nn.Module):
@@ -582,25 +581,29 @@ class KimiDecoderLayer(nn.Module):
581
positions: torch.Tensor,
582
hidden_states: torch.Tensor,
583
) -> torch.Tensor:
585
- attn_output = torch.empty_like(hidden_states)
586
- self.self_attn(
584
+ return self.self_attn(
585
hidden_states=hidden_states,
586
positions=positions,
589
- output=attn_output,
587
)
591
- return attn_output
588
589
def forward(
590
self,
591
positions: torch.Tensor,
592
hidden_states: torch.Tensor,
593
residual: torch.Tensor | None,
594
+ prefix_delta: torch.Tensor | None = None,
595
**kwargs,
599
- ) -> tuple[torch.Tensor, torch.Tensor]:
596
+ ) -> (
597
+ tuple[torch.Tensor, torch.Tensor]
598
+ | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
599
+ ):
600
if self.use_attn_residuals:
601
assert residual is not None
602
- return self.forward_attn_residual(positions, hidden_states, residual)
602
+ return self.forward_attn_residual(
603
+ positions, hidden_states, residual, prefix_delta
604
+ )
605
606
+ assert prefix_delta is None
607
# Self Attention
608
if residual is None:
609
residual = hidden_states
@@ -620,7 +623,8 @@ class KimiDecoderLayer(nn.Module):
623
positions: torch.Tensor,
624
hidden_states: torch.Tensor,
625
block_residual: torch.Tensor,
623
- ) -> tuple[torch.Tensor, torch.Tensor]:
626
+ prefix_delta: torch.Tensor | None,
627
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
628
prefix_sum = hidden_states
629
hidden_states = _apply_attn_res(
630
prefix_sum,
@@ -628,6 +632,7 @@ class KimiDecoderLayer(nn.Module):
632
self.self_attention_res_proj,
633
self.self_attention_res_norm,
634
self.prev_valid_blocks,
635
+ delta=prefix_delta,
636
output_norm=self.input_layernorm,
637
block_write_idx=(self.block_write_idx if self.is_block_write_layer else -1),
638
)
@@ -657,8 +662,7 @@ class KimiDecoderLayer(nn.Module):
662
)
663
664
hidden_states = self.mlp(hidden_states)
660
- prefix_sum = prefix_sum + hidden_states
661
- return prefix_sum, block_residual
665
+ return prefix_sum, block_residual, hidden_states
666
667
668
class KimiLinearModel(nn.Module, EagleModelMixin):
@@ -813,24 +817,28 @@ class KimiLinearModel(nn.Module, EagleModelMixin):
817
if residual is not None:
818
block_residual[:, : residual.size(1), :].copy_(residual)
819
residual = block_residual
820
+ prefix_delta = None
821
822
for layer_idx, layer in enumerate(
823
self.layers[self.start_layer : self.end_layer],
824
start=self.start_layer,
825
):
821
- hidden_states, residual = layer(
826
+ hidden_states, residual, prefix_delta = layer(
827
positions=positions,
828
hidden_states=hidden_states,
829
residual=residual,
830
+ prefix_delta=prefix_delta,
831
)
832
if (layer_idx + 1) in self.aux_hidden_state_layers:
827
- # AMD attn-res layer already returns prefix_sum + MLP delta as
828
- # hidden_states; the override drops the block bank in residual.
833
self._maybe_add_hidden_state(
830
- aux_hidden_states, layer_idx + 1, hidden_states, residual
834
+ aux_hidden_states,
835
+ layer_idx + 1,
836
+ hidden_states + prefix_delta,
837
+ residual,
838
)
839
840
if not get_pp_group().is_last_rank:
841
+ hidden_states = hidden_states + prefix_delta
842
return IntermediateTensors(
843
{"hidden_states": hidden_states, "residual": residual}
844
)
@@ -841,6 +849,7 @@ class KimiLinearModel(nn.Module, EagleModelMixin):
849
self.output_attn_res_proj,
850
self.output_attn_res_norm,
851
attn_res_block_num,
852
+ delta=prefix_delta,
853
)
854
# NOTE: the final norm is applied in compute_logits instead of here, so
855
# the MTP draft model receives the pre-norm hidden states.
vllm/models/kimi_k3/nvidia/model.py
-11
@@ -922,14 +922,12 @@ class KimiDecoderLayer(nn.Module):
922
aux_stream=aux_stream,
923
run_gemm_rs_ar=run_gemm_rs_ar,
924
)
925
- self._self_attn_writes_output = False
925
else:
926
self.self_attn = KimiLinearGatedDeltaNetAttention(
927
config,
928
vllm_config,
929
prefix=f"{prefix}.self_attn",
930
)
932
- self._self_attn_writes_output = True
931
else:
932
qk_nope_head_dim = config.qk_nope_head_dim
933
qk_rope_head_dim = config.qk_rope_head_dim
@@ -959,7 +957,6 @@ class KimiDecoderLayer(nn.Module):
957
aux_stream=aux_stream,
958
run_gemm_rs_ar=run_gemm_rs_ar,
959
)
962
- self._self_attn_writes_output = False
960
961
if self.use_sequence_parallel:
962
self.self_attn.o_proj.reduce_results = False
@@ -1025,14 +1022,6 @@ class KimiDecoderLayer(nn.Module):
1022
positions: torch.Tensor,
1023
hidden_states: torch.Tensor,
1024
) -> torch.Tensor:
1028
- if self._self_attn_writes_output:
1029
- output = torch.empty_like(hidden_states)
1030
- self.self_attn(
1031
- hidden_states=hidden_states,
1032
- positions=positions,
1033
- output=output,
1034
- )
1035
- return output
1025
return self.self_attn(
1026
hidden_states=hidden_states,
1027
positions=positions,