[K3 Perf] Flash kda out kernel for prefill, 1.1~1.4x kernel performance improvement (#51311)

Signed-off-by: yewentao256 <zhyanwentao@126.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

vllm/models/kimi_k3/nvidia/kda.py

@@ -46,6 +46,7 @@ from vllm.platforms import current_platform
 from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
 from vllm.v1.attention.backend import AttentionBackend
+from vllm.v1.worker.workspace import current_workspace_manager
 
 logger = init_logger(__name__)
 
@@ -187,20 +188,12 @@ def _flashkda_prefill(
     lower_bound: float,
     initial_state: torch.Tensor,
     cu_seqlens: torch.Tensor,
+    out: torch.Tensor,
+    final_state: torch.Tensor,
+    workspace: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     import vllm._flashkda_C  # noqa: F401
 
-    out = torch.empty(v.shape, dtype=v.dtype, device=v.device)
-    final_state = torch.empty_like(initial_state)
-    workspace = torch.empty(
-        torch.ops._flashkda_C.get_workspace_size(
-            q.shape[0] * q.shape[1],
-            q.shape[2],
-            cu_seqlens.numel() - 1,
-        ),
-        dtype=torch.uint8,
-        device=q.device,
-    )
     # FlashKDA hardcodes dense Q/K/V/G strides. Beta may be row-strided because
     # FlashKDA materializes its transposed [H, T] layout internally.
     # TODO: Teach FlashKDA to consume beta in [T, H] layout directly instead
@@ -436,6 +429,21 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
             vllm_config.model_config.dtype,
             self.gate_lower_bound,
         )
+        self._flashkda_buffer_specs: (
+            tuple[tuple[tuple[int, ...], torch.dtype], ...] | None
+        ) = None
+        if self.kda_prefill_backend == "flashkda":
+            T = vllm_config.scheduler_config.max_num_batched_tokens
+            N = vllm_config.scheduler_config.max_num_seqs
+            H, D = self.local_num_heads, self.head_dim
+            import vllm._flashkda_C  # noqa: F401
+
+            workspace_size = torch.ops._flashkda_C.get_workspace_size(T, H, N)
+            self._flashkda_buffer_specs = (
+                ((1, T, H, D), self.model_config.dtype),
+                ((N, H, D, D), self.get_state_dtype()[1]),
+                ((workspace_size,), torch.uint8),
+            )
 
         self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
         decode_norm_weight = None
@@ -696,6 +704,15 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
                 )
                 if self.kda_prefill_backend == "flashkda":
                     assert self.gate_lower_bound is not None
+                    assert self._flashkda_buffer_specs is not None
+                    workspace_out, final_state, workspace = (
+                        current_workspace_manager().get_simultaneous(
+                            *self._flashkda_buffer_specs
+                        )
+                    )
+                    flashkda_out = (
+                        workspace_out if has_spec_decode else core_attn_out
+                    )[:, : q_ns.shape[1]]
                     (
                         core_attn_out_non_spec,
                         last_recurrent_state,
@@ -710,6 +727,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
                         lower_bound=self.gate_lower_bound,
                         initial_state=initial_state,
                         cu_seqlens=non_spec_query_start_loc,
+                        out=flashkda_out,
+                        final_state=final_state[: initial_state.shape[0]],
+                        workspace=workspace,
                     )
                 else:
                     (
@@ -766,10 +786,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
             core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
             core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
         elif core_attn_out_non_spec is not None:
-            # TODO: prefill and decode kernels write directly to core_attn_out
-            core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
-                0, :num_actual_tokens
-            ]
+            if self.kda_prefill_backend != "flashkda" or m.num_prefills == 0:
+                # TODO: decode kernels write directly to core_attn_out
+                core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
+                    0, :num_actual_tokens
+                ]
         else:
             assert core_attn_out_spec is not None
         # Triton normalizes in place, so this is a self-copy with no device