[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>

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()

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()

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()

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(

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(

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.

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,