@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 86aca66191

[Kimi K3 Perf] Avoid KDA mixed-batch gather/scatter, 5.2%~7.7% E2E Throughput Improvement (#56159)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Sep 10, 2026 at 05:57 UTC 86aca6619161b308d675c753f228201e70b59ed7

6 files changed +66 -5

tests/models/kimi_k3/test_kda.py

@@ -330,6 +330,7 @@ def test_chunk_kda_fused_gate_cumsum_matches_unfused(

330 cu_seqlens=cu_seqlens_t,

331 use_qk_l2norm_in_kernel=True,

332 )

333 + output = torch.empty_like(v)

334 new_o, new_ht = chunk_kda_with_fused_gate(

335 q=q.clone(),

336 k=k.clone(),

@@ -343,8 +344,10 @@ def test_chunk_kda_fused_gate_cumsum_matches_unfused(

344 output_final_state=True,

345 cu_seqlens=cu_seqlens_t,

346 use_qk_l2norm_in_kernel=True,

347 + out=output,

348 )

349

350 + assert new_o.data_ptr() == output.data_ptr()

351 assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3)

352 assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3)

353

@@ -442,6 +445,7 @@ def test_packed_kda_decode_correctness(

445 use_qk_l2norm_in_kernel=True,

446 )

447 packed_state = state

448 + packed_output = torch.empty_like(dense_out)

449 packed_out, _ = PACKED_DECODE_IMPLS[impl](

450 mixed_qkv=mixed_qkv,

451 raw_g=raw_g,

@@ -451,8 +455,11 @@ def test_packed_kda_decode_correctness(

455 lower_bound=lower_bound,

456 initial_state=packed_state,

457 state_indices=state_indices,

458 + **({"out": packed_output} if impl == "nvidia" else {}),

459 )

460

461 + if impl == "nvidia":

462 + assert packed_out.data_ptr() == packed_output.data_ptr()

463 assert_close("o", dense_out, packed_out, 1e-3, err_atol=1e-3)

464 assert_close("ht", dense_state, packed_state, 1e-3, err_atol=1e-3)

465

tests/models/kimi_k3/test_kda_metadata.py

@@ -512,6 +512,8 @@ def test_mixed_regular_and_spec_decode_uses_packed_decode_metadata():

512 assert actual.non_spec_query_start_loc is None

513 torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0, 1]))

514 torch.testing.assert_close(actual.spec_token_indx, torch.tensor([2, 3, 4]))

515 + assert actual.non_spec_token_start == 0

516 + assert actual.spec_token_start == 2

517 torch.testing.assert_close(

518 actual.spec_query_start_loc,

519 torch.tensor([0, 3], dtype=torch.int32),

@@ -539,6 +541,8 @@ def test_mixed_regular_and_spec_decode_excludes_request_padding():

541 assert actual.non_spec_state_indices_tensor.shape == (1,)

542 torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0]))

543 torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3]))

544 + assert actual.non_spec_token_start == 0

545 + assert actual.spec_token_start == 1

546

547

548 @pytest.mark.parametrize("mamba_cache_mode", ["none", "align"])

vllm/models/kimi_k3/nvidia/kda.py

+33 -3

@@ -903,6 +903,8 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

903 non_spec_state_indices_tensor = m.non_spec_state_indices_tensor

904 spec_token_indx = m.spec_token_indx

905 non_spec_token_indx = m.non_spec_token_indx

906 + spec_token_start = m.spec_token_start

907 + non_spec_token_start = m.non_spec_token_start

908 spec_state_indices_tensor = m.spec_state_indices_tensor

909 spec_query_start_loc = m.spec_query_start_loc

910 num_accepted_tokens = m.num_accepted_tokens

@@ -981,6 +983,18 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

983 mixed_qkv_spec = mixed_qkv

984 g1_spec, beta_spec = g1, beta

985 mixed_qkv_ns = g1_ns = beta_ns = None

986 + elif spec_token_start is not None:

987 + assert non_spec_token_start is not None

988 + spec_end = spec_token_start + m.num_spec_decode_tokens

989 + non_spec_end = (

990 + non_spec_token_start + m.num_prefill_tokens + m.num_decode_tokens

991 + )

992 + spec_slice = slice(spec_token_start, spec_end)

993 + non_spec_slice = slice(non_spec_token_start, non_spec_end)

994 + mixed_qkv_spec = mixed_qkv[spec_slice]

995 + g1_spec, beta_spec = g1[:, spec_slice], beta[:, spec_slice]

996 + mixed_qkv_ns = mixed_qkv[non_spec_slice]

997 + g1_ns, beta_ns = g1[:, non_spec_slice], beta[:, non_spec_slice]

998 else:

999 assert spec_token_indx is not None

1000 assert non_spec_token_indx is not None

@@ -1029,6 +1043,8 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

1043 if m.num_prefills == 0 and m.num_decodes == 0

1044 else None

1045 )

1046 + if spec_token_start is not None:

1047 + spec_out = core_attn_out[:, spec_slice]

1048 if self.use_recoverssm:

1049 from vllm.models.kimi_k3.nvidia.ops.recoverssm import (

1050 kda_recoverssm_verify,

@@ -1076,6 +1092,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

1092 core_attn_out_non_spec = None

1093 if mixed_qkv_ns is not None:

1094 assert g1_ns is not None and beta_ns is not None

1095 + non_spec_out = None

1096 + if non_spec_token_start is not None:

1097 + non_spec_out = core_attn_out[:, non_spec_slice]

1098 if m.num_prefills > 0:

1099 q_ns, k_ns, v_ns = mixed_qkv_ns.split(

1100 self.local_projection_size, dim=-1

@@ -1126,6 +1145,8 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

1145 flashkda_out = (

1146 workspace_out if has_spec_decode else core_attn_out

1147 )[:, : q_ns.shape[1]]

1148 + if non_spec_out is not None:

1149 + flashkda_out = non_spec_out

1150 if checkpoint is not None:

1151 assert non_spec_query_start_loc is not None

1152 num_sequences = initial_state.shape[0]

@@ -1211,7 +1232,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

1232 if q_ns.shape[1] > initial_state.shape[0]:

1233 assert m.flashinfer_prefill_seq_order is not None

1234 flashinfer_out = core_attn_out[:, : q_ns.shape[1]]

1214 - if has_spec_decode:

1235 + if non_spec_out is not None:

1236 + flashinfer_out = non_spec_out

1237 + elif has_spec_decode:

1238 assert self._flashinfer_kda_output_spec is not None

1239 (workspace_out,) = current_workspace_manager().get_simultaneous(

1240 self._flashinfer_kda_output_spec

@@ -1251,6 +1274,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

1274 output_final_state=True,

1275 use_qk_l2norm_in_kernel=True,

1276 cu_seqlens=non_spec_query_start_loc,

1277 + out=non_spec_out,

1278 )

1279 recurrent_state[non_spec_state_indices_tensor] = (

1280 last_recurrent_state.to(recurrent_state.dtype)

@@ -1284,12 +1308,18 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

1308 lower_bound=self.gate_lower_bound,

1309 initial_state=recurrent_state,

1310 state_indices=decode_conv_indices,

1311 + out=non_spec_out,

1312 )

1313

1314 # Restore the scheduler's original token order for mixed batches.

1315 if core_attn_out_spec is not None and core_attn_out_non_spec is not None:

1291 - core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec)

1292 - core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)

1316 + if spec_token_start is None:

1317 + assert spec_token_indx is not None

1318 + assert non_spec_token_indx is not None

1319 + core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec)

1320 + core_attn_out.index_copy_(

1321 + 1, non_spec_token_indx, core_attn_out_non_spec

1322 + )

1323 elif core_attn_out_non_spec is not None:

1324 if (

1325 self.kda_prefill_backend not in ("flashkda", "flashinfer")

vllm/models/kimi_k3/nvidia/kda_metadata.py

+13

@@ -270,6 +270,8 @@ class KDACheckpointMetadata:

270

271 @dataclass

272 class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):

273 + spec_token_start: int | None = None

274 + non_spec_token_start: int | None = None

275 flashinfer_prefill_query_start_loc: torch.Tensor | None = None

276 flashinfer_prefill_seq_order: torch.Tensor | None = None

277 recoverssm_commit: KDARecoverSSMCommitMetadata | None = None

@@ -422,6 +424,8 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

424 spec_sequence_masks_cpu = None

425

426 spec_request_indices = None

427 + spec_token_start = None

428 + non_spec_token_start = None

429 if num_spec_decodes == 0:

430 # The runner orders ordinary decodes before prefills.

431 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

@@ -518,6 +522,13 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

522 non_spec_token_indx = index[:num_non_spec_tokens]

523 spec_token_indx = index[num_non_spec_tokens:]

524

525 + active_spec_mask = spec_sequence_masks_cpu[query_lens_cpu > 0]

526 + # check if spec / non spec tokens are continuous

527 + if (active_spec_mask[1:] != active_spec_mask[:-1]).sum().item() == 1:

528 + spec_first = active_spec_mask[0].item()

529 + spec_token_start = 0 if spec_first else num_non_spec_tokens

530 + non_spec_token_start = num_spec_decode_tokens if spec_first else 0

531 +

532 # Native spec uses one state slot per step. RecoverSSM keeps

533 # only the current checkpoint slot.

534 spec_state_indices_tensor = block_table_tensor[

@@ -759,6 +770,8 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

770 spec_token_indx=spec_token_indx,

771 non_spec_token_indx=non_spec_token_indx,

772 num_accepted_tokens=num_accepted_tokens,

773 + spec_token_start=spec_token_start,

774 + non_spec_token_start=non_spec_token_start,

775 flashinfer_prefill_query_start_loc=flashinfer_prefill_query_start_loc,

776 flashinfer_prefill_seq_order=flashinfer_prefill_seq_order,

777 recoverssm_commit=recoverssm_commit,

vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py

+6 -1

@@ -598,6 +598,7 @@ def _chunk_kda_fwd_with_cumulative_g(

598 chunk_indices: torch.Tensor | None = None,

599 chunk_size: int = FLA_CHUNK_SIZE,

600 safe_gate: bool = False,

601 + out: torch.Tensor | None = None,

602 ):

603 # `g` must already be chunk-local cumulatively-summed AND scaled by

604 # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)).

@@ -642,7 +643,7 @@ def _chunk_kda_fwd_with_cumulative_g(

643 g=g,

644 A=Aqk,

645 h=h,

645 - o=v,

646 + o=v if out is None else out,

647 scale=scale,

648 cu_seqlens=cu_seqlens,

649 chunk_indices=chunk_indices,

@@ -706,6 +707,7 @@ def chunk_kda_with_fused_gate_fwd(

707 output_final_state: bool,

708 lower_bound: float | None = None,

709 cu_seqlens: torch.Tensor | None = None,

710 + out: torch.Tensor | None = None,

711 ):

712 chunk_size = FLA_CHUNK_SIZE

713 chunk_indices = (

@@ -736,6 +738,7 @@ def chunk_kda_with_fused_gate_fwd(

738 chunk_indices=chunk_indices,

739 chunk_size=chunk_size,

740 safe_gate=lower_bound is not None,

741 + out=out,

742 )

743

744

@@ -787,6 +790,7 @@ def chunk_kda_with_fused_gate(

790 lower_bound: float | None = None,

791 use_qk_l2norm_in_kernel: bool = False,

792 cu_seqlens: torch.Tensor | None = None,

793 + out: torch.Tensor | None = None,

794 **kwargs,

795 ):

796 """Run chunk KDA from raw gate and beta projections."""

@@ -810,6 +814,7 @@ def chunk_kda_with_fused_gate(

814 output_final_state=output_final_state,

815 lower_bound=lower_bound,

816 cu_seqlens=cu_seqlens,

817 + out=out,

818 )

819 return o, final_state

820

vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py

+3 -1

@@ -591,6 +591,7 @@ def fused_recurrent_kda_packed_decode(

591 initial_state: torch.Tensor,

592 state_indices: torch.Tensor,

593 scale: float | None = None,

594 + out: torch.Tensor | None = None,

595 ) -> tuple[torch.Tensor, torch.Tensor]:

596 """Run one-token KDA decode directly from packed post-conv QKV."""

597 if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1:

@@ -639,7 +640,8 @@ def fused_recurrent_kda_packed_decode(

640 if scale is None:

641 scale = K**-0.5

642

642 - out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device)

643 + if out is None:

644 + out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device)

645 grid = (cdiv(V, BV), B * H)

646 fused_recurrent_kda_packed_decode_kernel[grid](

647 mixed_qkv=mixed_qkv,