[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>
Wentao Ye committed
Aug 12, 2026 at 14:51 UTC
fe889ac9255491e82e58012d7787e920c6c7fbc2
1 file changed
+36
-15
vllm/models/kimi_k3/nvidia/kda.py
+36
-15
@@ -46,6 +46,7 @@ from vllm.platforms import current_platform
46
from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
47
from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
48
from vllm.v1.attention.backend import AttentionBackend
49
+from vllm.v1.worker.workspace import current_workspace_manager
50
51
logger = init_logger(__name__)
52
@@ -187,20 +188,12 @@ def _flashkda_prefill(
188
lower_bound: float,
189
initial_state: torch.Tensor,
190
cu_seqlens: torch.Tensor,
191
+ out: torch.Tensor,
192
+ final_state: torch.Tensor,
193
+ workspace: torch.Tensor,
194
) -> tuple[torch.Tensor, torch.Tensor]:
195
import vllm._flashkda_C # noqa: F401
196
193
- out = torch.empty(v.shape, dtype=v.dtype, device=v.device)
194
- final_state = torch.empty_like(initial_state)
195
- workspace = torch.empty(
196
- torch.ops._flashkda_C.get_workspace_size(
197
- q.shape[0] * q.shape[1],
198
- q.shape[2],
199
- cu_seqlens.numel() - 1,
200
- ),
201
- dtype=torch.uint8,
202
- device=q.device,
203
- )
197
# FlashKDA hardcodes dense Q/K/V/G strides. Beta may be row-strided because
198
# FlashKDA materializes its transposed [H, T] layout internally.
199
# TODO: Teach FlashKDA to consume beta in [T, H] layout directly instead
@@ -436,6 +429,21 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
429
vllm_config.model_config.dtype,
430
self.gate_lower_bound,
431
)
432
+ self._flashkda_buffer_specs: (
433
+ tuple[tuple[tuple[int, ...], torch.dtype], ...] | None
434
+ ) = None
435
+ if self.kda_prefill_backend == "flashkda":
436
+ T = vllm_config.scheduler_config.max_num_batched_tokens
437
+ N = vllm_config.scheduler_config.max_num_seqs
438
+ H, D = self.local_num_heads, self.head_dim
439
+ import vllm._flashkda_C # noqa: F401
440
+
441
+ workspace_size = torch.ops._flashkda_C.get_workspace_size(T, H, N)
442
+ self._flashkda_buffer_specs = (
443
+ ((1, T, H, D), self.model_config.dtype),
444
+ ((N, H, D, D), self.get_state_dtype()[1]),
445
+ ((workspace_size,), torch.uint8),
446
+ )
447
448
self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
449
decode_norm_weight = None
@@ -696,6 +704,15 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
704
)
705
if self.kda_prefill_backend == "flashkda":
706
assert self.gate_lower_bound is not None
707
+ assert self._flashkda_buffer_specs is not None
708
+ workspace_out, final_state, workspace = (
709
+ current_workspace_manager().get_simultaneous(
710
+ *self._flashkda_buffer_specs
711
+ )
712
+ )
713
+ flashkda_out = (
714
+ workspace_out if has_spec_decode else core_attn_out
715
+ )[:, : q_ns.shape[1]]
716
(
717
core_attn_out_non_spec,
718
last_recurrent_state,
@@ -710,6 +727,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
727
lower_bound=self.gate_lower_bound,
728
initial_state=initial_state,
729
cu_seqlens=non_spec_query_start_loc,
730
+ out=flashkda_out,
731
+ final_state=final_state[: initial_state.shape[0]],
732
+ workspace=workspace,
733
)
734
else:
735
(
@@ -766,10 +786,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
786
core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
787
core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
788
elif core_attn_out_non_spec is not None:
769
- # TODO: prefill and decode kernels write directly to core_attn_out
770
- core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
771
- 0, :num_actual_tokens
772
- ]
789
+ if self.kda_prefill_backend != "flashkda" or m.num_prefills == 0:
790
+ # TODO: decode kernels write directly to core_attn_out
791
+ core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
792
+ 0, :num_actual_tokens
793
+ ]
794
else:
795
assert core_attn_out_spec is not None
796
# Triton normalizes in place, so this is a self-copy with no device