From 382970ee6ca490aeaaaf4e32c53695b581ff61ba Mon Sep 17 00:00:00 2001 From: yinfengLiu Date: Tue, 22 Sep 2026 04:44:41 +0800 Subject: [PATCH] [Kimi-K3][AMD] Return KDA and MLA projection outputs directly (#50592) Signed-off-by: Liuyinfeng01 Co-authored-by: Liuyinfeng01 <199041580+LiuYinfeng01@users.noreply.github.com> Co-authored-by: Andreas Karatzas --- .../kimi_k3/test_amd_kda_direct_return.py | 72 +++++++++++++++++++ .../kimi_k3/test_amd_mla_direct_return.py | 38 ++++++++++ .../kimi_k3/test_nvidia_kda_direct_return.py | 31 ++++++++ .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 5 +- vllm/models/kimi_k3/amd/kda.py | 5 +- vllm/models/kimi_k3/amd/linear.py | 41 ++++++----- vllm/models/kimi_k3/nvidia/model.py | 11 --- 7 files changed, 170 insertions(+), 33 deletions(-) create mode 100644 tests/models/kimi_k3/test_amd_kda_direct_return.py create mode 100644 tests/models/kimi_k3/test_amd_mla_direct_return.py create mode 100644 tests/models/kimi_k3/test_nvidia_kda_direct_return.py diff --git a/tests/models/kimi_k3/test_amd_kda_direct_return.py b/tests/models/kimi_k3/test_amd_kda_direct_return.py new file mode 100644 index 0000000000..e46efe7681 --- /dev/null +++ b/tests/models/kimi_k3/test_amd_kda_direct_return.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch +from torch import nn + +from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention +from vllm.models.kimi_k3.amd.linear import KimiDecoderLayer + + +class _ReturningAttention(nn.Module): + def __init__(self, result: torch.Tensor) -> None: + super().__init__() + self.result = result + + def forward( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + return self.result + + +class _TupleProjection(nn.Module): + def __init__(self, result: torch.Tensor) -> None: + super().__init__() + self.result = result + + def forward(self, _: torch.Tensor) -> tuple[torch.Tensor, None]: + return self.result, None + + +def _make_layer(self_attn: nn.Module) -> KimiDecoderLayer: + layer = object.__new__(KimiDecoderLayer) + nn.Module.__init__(layer) + layer.self_attn = self_attn + return layer + + +def test_kda_self_attention_returns_projection_storage_directly() -> None: + hidden_states = torch.randn(4, 8) + projected = torch.randn_like(hidden_states) + layer = _make_layer(_ReturningAttention(projected)) + + output = layer._run_self_attn(torch.arange(4), hidden_states) + + assert output is projected + assert output.data_ptr() == projected.data_ptr() + + +def test_k3_kda_forward_returns_output_projection_storage() -> None: + hidden_states = torch.randn(4, 8) + projected = torch.randn_like(hidden_states) + + def fill_core_output(**kwargs: torch.Tensor) -> None: + kwargs["core_attn_out"].fill_(1) + + attention = SimpleNamespace( + local_projection_size=4, + head_dim=2, + local_num_heads=2, + in_proj_padding=0, + in_proj_qkvgfab=_TupleProjection(torch.randn(4, 20)), + f_b_proj=_TupleProjection(torch.randn(4, 4)), + o_proj=_TupleProjection(projected), + _forward=fill_core_output, + ) + + output = KimiK3DeltaAttention.forward(attention, hidden_states, torch.arange(4)) + + assert output is projected + assert output.data_ptr() == projected.data_ptr() diff --git a/tests/models/kimi_k3/test_amd_mla_direct_return.py b/tests/models/kimi_k3/test_amd_mla_direct_return.py new file mode 100644 index 0000000000..a1c120737d --- /dev/null +++ b/tests/models/kimi_k3/test_amd_mla_direct_return.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + +from vllm.models.kimi_k3.amd.linear import KimiDecoderLayer, KimiMLAAttention + + +class _DirectMLA(KimiMLAAttention): + def __init__(self, result: torch.Tensor) -> None: + nn.Module.__init__(self) + self.result = result + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + return self.result + + +def _make_layer(self_attn: nn.Module) -> KimiDecoderLayer: + layer = object.__new__(KimiDecoderLayer) + nn.Module.__init__(layer) + layer.self_attn = self_attn + return layer + + +def test_mla_self_attention_returns_projection_storage_directly() -> None: + hidden_states = torch.randn(4, 8) + projected = torch.randn_like(hidden_states) + layer = _make_layer(_DirectMLA(projected)) + + output = layer._run_self_attn(torch.arange(4), hidden_states) + + assert output is projected + assert output.data_ptr() == projected.data_ptr() diff --git a/tests/models/kimi_k3/test_nvidia_kda_direct_return.py b/tests/models/kimi_k3/test_nvidia_kda_direct_return.py new file mode 100644 index 0000000000..5b40d7ef6f --- /dev/null +++ b/tests/models/kimi_k3/test_nvidia_kda_direct_return.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + +from vllm.models.kimi_k3.nvidia.model import KimiDecoderLayer + + +class _ReturningSharedKDA(nn.Module): + def __init__(self, result: torch.Tensor) -> None: + super().__init__() + self.result = result + + def forward( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + return self.result + + +def test_low_rank_kda_caller_returns_projection_storage_directly() -> None: + hidden_states = torch.randn(4, 8) + projected = torch.randn_like(hidden_states) + layer = object.__new__(KimiDecoderLayer) + nn.Module.__init__(layer) + layer.self_attn = _ReturningSharedKDA(projected) + + output = layer._run_self_attn(torch.arange(4), hidden_states) + + assert output is projected + assert output.data_ptr() == projected.data_ptr() diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index bf81cdd75f..3892d1a27a 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -329,8 +329,7 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention): self, hidden_states: torch.Tensor, positions: torch.Tensor, - output: torch.Tensor, - ) -> None: + ) -> torch.Tensor: num_tokens = hidden_states.size(0) projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] if self.use_full_rank_gate: @@ -375,7 +374,7 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention): core_attn_out=core_attn_out, ) core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") - output[:] = self.o_proj(core_attn_out)[0] + return self.o_proj(core_attn_out)[0] @eager_break_during_capture def _forward( diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index a02bf90d38..ed631aea7e 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -271,8 +271,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): self, hidden_states: torch.Tensor, positions: torch.Tensor, - output: torch.Tensor, - ) -> None: + ) -> torch.Tensor: num_tokens = hidden_states.size(0) projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] @@ -306,7 +305,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): core_attn_out=core_attn_out, ) core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") - output[:] = self.o_proj(core_attn_out)[0] + return self.o_proj(core_attn_out)[0] @eager_break_during_capture def _forward( diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py index 2a013ed0fc..65dad661ab 100644 --- a/vllm/models/kimi_k3/amd/linear.py +++ b/vllm/models/kimi_k3/amd/linear.py @@ -456,9 +456,8 @@ class KimiMLAAttention(nn.Module): self, positions: torch.Tensor, hidden_states: torch.Tensor, - output: torch.Tensor, - ) -> None: - output[:] = self.mla_attn(positions, hidden_states) + ) -> torch.Tensor: + return self.mla_attn(positions, hidden_states) class KimiDecoderLayer(nn.Module): @@ -582,25 +581,29 @@ class KimiDecoderLayer(nn.Module): positions: torch.Tensor, hidden_states: torch.Tensor, ) -> torch.Tensor: - attn_output = torch.empty_like(hidden_states) - self.self_attn( + return self.self_attn( hidden_states=hidden_states, positions=positions, - output=attn_output, ) - return attn_output def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None, + prefix_delta: torch.Tensor | None = None, **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> ( + tuple[torch.Tensor, torch.Tensor] + | tuple[torch.Tensor, torch.Tensor, torch.Tensor] + ): if self.use_attn_residuals: assert residual is not None - return self.forward_attn_residual(positions, hidden_states, residual) + return self.forward_attn_residual( + positions, hidden_states, residual, prefix_delta + ) + assert prefix_delta is None # Self Attention if residual is None: residual = hidden_states @@ -620,7 +623,8 @@ class KimiDecoderLayer(nn.Module): positions: torch.Tensor, hidden_states: torch.Tensor, block_residual: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: + prefix_delta: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: prefix_sum = hidden_states hidden_states = _apply_attn_res( prefix_sum, @@ -628,6 +632,7 @@ class KimiDecoderLayer(nn.Module): self.self_attention_res_proj, self.self_attention_res_norm, self.prev_valid_blocks, + delta=prefix_delta, output_norm=self.input_layernorm, block_write_idx=(self.block_write_idx if self.is_block_write_layer else -1), ) @@ -657,8 +662,7 @@ class KimiDecoderLayer(nn.Module): ) hidden_states = self.mlp(hidden_states) - prefix_sum = prefix_sum + hidden_states - return prefix_sum, block_residual + return prefix_sum, block_residual, hidden_states class KimiLinearModel(nn.Module, EagleModelMixin): @@ -813,24 +817,28 @@ class KimiLinearModel(nn.Module, EagleModelMixin): if residual is not None: block_residual[:, : residual.size(1), :].copy_(residual) residual = block_residual + prefix_delta = None for layer_idx, layer in enumerate( self.layers[self.start_layer : self.end_layer], start=self.start_layer, ): - hidden_states, residual = layer( + hidden_states, residual, prefix_delta = layer( positions=positions, hidden_states=hidden_states, residual=residual, + prefix_delta=prefix_delta, ) if (layer_idx + 1) in self.aux_hidden_state_layers: - # AMD attn-res layer already returns prefix_sum + MLP delta as - # hidden_states; the override drops the block bank in residual. self._maybe_add_hidden_state( - aux_hidden_states, layer_idx + 1, hidden_states, residual + aux_hidden_states, + layer_idx + 1, + hidden_states + prefix_delta, + residual, ) if not get_pp_group().is_last_rank: + hidden_states = hidden_states + prefix_delta return IntermediateTensors( {"hidden_states": hidden_states, "residual": residual} ) @@ -841,6 +849,7 @@ class KimiLinearModel(nn.Module, EagleModelMixin): self.output_attn_res_proj, self.output_attn_res_norm, attn_res_block_num, + delta=prefix_delta, ) # NOTE: the final norm is applied in compute_logits instead of here, so # the MTP draft model receives the pre-norm hidden states. diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 737184b9d8..5dedf44b7c 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -922,14 +922,12 @@ class KimiDecoderLayer(nn.Module): aux_stream=aux_stream, run_gemm_rs_ar=run_gemm_rs_ar, ) - self._self_attn_writes_output = False else: self.self_attn = KimiLinearGatedDeltaNetAttention( config, vllm_config, prefix=f"{prefix}.self_attn", ) - self._self_attn_writes_output = True else: qk_nope_head_dim = config.qk_nope_head_dim qk_rope_head_dim = config.qk_rope_head_dim @@ -959,7 +957,6 @@ class KimiDecoderLayer(nn.Module): aux_stream=aux_stream, run_gemm_rs_ar=run_gemm_rs_ar, ) - self._self_attn_writes_output = False if self.use_sequence_parallel: self.self_attn.o_proj.reduce_results = False @@ -1025,14 +1022,6 @@ class KimiDecoderLayer(nn.Module): positions: torch.Tensor, hidden_states: torch.Tensor, ) -> torch.Tensor: - if self._self_attn_writes_output: - output = torch.empty_like(hidden_states) - self.self_attn( - hidden_states=hidden_states, - positions=positions, - output=output, - ) - return output return self.self_attn( hidden_states=hidden_states, positions=positions,