iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye
2026-09-10 17:57:01 +08:00
committed by GitHub
parent c9355e25e8
commit 86aca66191
6 changed files with 66 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/kimi_k3/test_kda.py
+7
View File
@@ -330,6 +330,7 @@ def test_chunk_kda_fused_gate_cumsum_matches_unfused(
cu_seqlens=cu_seqlens_t,
use_qk_l2norm_in_kernel=True,
)
output = torch.empty_like(v)
new_o, new_ht = chunk_kda_with_fused_gate(
q=q.clone(),
k=k.clone(),
@@ -343,8 +344,10 @@ def test_chunk_kda_fused_gate_cumsum_matches_unfused(
output_final_state=True,
cu_seqlens=cu_seqlens_t,
use_qk_l2norm_in_kernel=True,
out=output,
)
assert new_o.data_ptr() == output.data_ptr()
assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3)
assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3)
@@ -442,6 +445,7 @@ def test_packed_kda_decode_correctness(
use_qk_l2norm_in_kernel=True,
)
packed_state = state
packed_output = torch.empty_like(dense_out)
packed_out, _ = PACKED_DECODE_IMPLS[impl](
mixed_qkv=mixed_qkv,
raw_g=raw_g,
@@ -451,8 +455,11 @@ def test_packed_kda_decode_correctness(
lower_bound=lower_bound,
initial_state=packed_state,
state_indices=state_indices,
**({"out": packed_output} if impl == "nvidia" else {}),
)
if impl == "nvidia":
assert packed_out.data_ptr() == packed_output.data_ptr()
assert_close("o", dense_out, packed_out, 1e-3, err_atol=1e-3)
assert_close("ht", dense_state, packed_state, 1e-3, err_atol=1e-3)
tests/models/kimi_k3/test_kda_metadata.py
+4
View File
@@ -512,6 +512,8 @@ def test_mixed_regular_and_spec_decode_uses_packed_decode_metadata():
assert actual.non_spec_query_start_loc is None
torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0, 1]))
torch.testing.assert_close(actual.spec_token_indx, torch.tensor([2, 3, 4]))
assert actual.non_spec_token_start == 0
assert actual.spec_token_start == 2
torch.testing.assert_close(
actual.spec_query_start_loc,
torch.tensor([0, 3], dtype=torch.int32),
@@ -539,6 +541,8 @@ def test_mixed_regular_and_spec_decode_excludes_request_padding():
assert actual.non_spec_state_indices_tensor.shape == (1,)
torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0]))
torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3]))
assert actual.non_spec_token_start == 0
assert actual.spec_token_start == 1
@pytest.mark.parametrize("mamba_cache_mode", ["none", "align"])
vllm/models/kimi_k3/nvidia/kda.py
+33 -3
View File
@@ -903,6 +903,8 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
non_spec_state_indices_tensor = m.non_spec_state_indices_tensor
spec_token_indx = m.spec_token_indx
non_spec_token_indx = m.non_spec_token_indx
spec_token_start = m.spec_token_start
non_spec_token_start = m.non_spec_token_start
spec_state_indices_tensor = m.spec_state_indices_tensor
spec_query_start_loc = m.spec_query_start_loc
num_accepted_tokens = m.num_accepted_tokens
@@ -981,6 +983,18 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
mixed_qkv_spec = mixed_qkv
g1_spec, beta_spec = g1, beta
mixed_qkv_ns = g1_ns = beta_ns = None
elif spec_token_start is not None:
assert non_spec_token_start is not None
spec_end = spec_token_start + m.num_spec_decode_tokens
non_spec_end = (
non_spec_token_start + m.num_prefill_tokens + m.num_decode_tokens
)
spec_slice = slice(spec_token_start, spec_end)
non_spec_slice = slice(non_spec_token_start, non_spec_end)
mixed_qkv_spec = mixed_qkv[spec_slice]
g1_spec, beta_spec = g1[:, spec_slice], beta[:, spec_slice]
mixed_qkv_ns = mixed_qkv[non_spec_slice]
g1_ns, beta_ns = g1[:, non_spec_slice], beta[:, non_spec_slice]
else:
assert spec_token_indx is not None
assert non_spec_token_indx is not None
@@ -1029,6 +1043,8 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
if m.num_prefills == 0 and m.num_decodes == 0
else None
)
if spec_token_start is not None:
spec_out = core_attn_out[:, spec_slice]
if self.use_recoverssm:
from vllm.models.kimi_k3.nvidia.ops.recoverssm import (
kda_recoverssm_verify,
@@ -1076,6 +1092,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
core_attn_out_non_spec = None
if mixed_qkv_ns is not None:
assert g1_ns is not None and beta_ns is not None
non_spec_out = None
if non_spec_token_start is not None:
non_spec_out = core_attn_out[:, non_spec_slice]
if m.num_prefills > 0:
q_ns, k_ns, v_ns = mixed_qkv_ns.split(
self.local_projection_size, dim=-1
@@ -1126,6 +1145,8 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
flashkda_out = (
workspace_out if has_spec_decode else core_attn_out
)[:, : q_ns.shape[1]]
if non_spec_out is not None:
flashkda_out = non_spec_out
if checkpoint is not None:
assert non_spec_query_start_loc is not None
num_sequences = initial_state.shape[0]
@@ -1211,7 +1232,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
if q_ns.shape[1] > initial_state.shape[0]:
assert m.flashinfer_prefill_seq_order is not None
flashinfer_out = core_attn_out[:, : q_ns.shape[1]]
if has_spec_decode:
if non_spec_out is not None:
flashinfer_out = non_spec_out
elif has_spec_decode:
assert self._flashinfer_kda_output_spec is not None
(workspace_out,) = current_workspace_manager().get_simultaneous(
self._flashinfer_kda_output_spec
@@ -1251,6 +1274,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
output_final_state=True,
use_qk_l2norm_in_kernel=True,
cu_seqlens=non_spec_query_start_loc,
out=non_spec_out,
)
recurrent_state[non_spec_state_indices_tensor] = (
last_recurrent_state.to(recurrent_state.dtype)
@@ -1284,12 +1308,18 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
lower_bound=self.gate_lower_bound,
initial_state=recurrent_state,
state_indices=decode_conv_indices,
out=non_spec_out,
)
# Restore the scheduler's original token order for mixed batches.
if core_attn_out_spec is not None and core_attn_out_non_spec is not None:
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)
if spec_token_start is None:
assert spec_token_indx is not None
assert non_spec_token_indx is not None
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:
if (
self.kda_prefill_backend not in ("flashkda", "flashinfer")
vllm/models/kimi_k3/nvidia/kda_metadata.py
+13
View File
@@ -270,6 +270,8 @@ class KDACheckpointMetadata:
@dataclass
class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
spec_token_start: int | None = None
non_spec_token_start: int | None = None
flashinfer_prefill_query_start_loc: torch.Tensor | None = None
flashinfer_prefill_seq_order: torch.Tensor | None = None
recoverssm_commit: KDARecoverSSMCommitMetadata | None = None
@@ -422,6 +424,8 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
spec_sequence_masks_cpu = None
spec_request_indices = None
spec_token_start = None
non_spec_token_start = None
if num_spec_decodes == 0:
# The runner orders ordinary decodes before prefills.
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
@@ -518,6 +522,13 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
non_spec_token_indx = index[:num_non_spec_tokens]
spec_token_indx = index[num_non_spec_tokens:]
active_spec_mask = spec_sequence_masks_cpu[query_lens_cpu > 0]
# check if spec / non spec tokens are continuous
if (active_spec_mask[1:] != active_spec_mask[:-1]).sum().item() == 1:
spec_first = active_spec_mask[0].item()
spec_token_start = 0 if spec_first else num_non_spec_tokens
non_spec_token_start = num_spec_decode_tokens if spec_first else 0
# Native spec uses one state slot per step. RecoverSSM keeps
# only the current checkpoint slot.
spec_state_indices_tensor = block_table_tensor[
@@ -759,6 +770,8 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
spec_token_indx=spec_token_indx,
non_spec_token_indx=non_spec_token_indx,
num_accepted_tokens=num_accepted_tokens,
spec_token_start=spec_token_start,
non_spec_token_start=non_spec_token_start,
flashinfer_prefill_query_start_loc=flashinfer_prefill_query_start_loc,
flashinfer_prefill_seq_order=flashinfer_prefill_seq_order,
recoverssm_commit=recoverssm_commit,
vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py
Vendored
+6 -1
View File
@@ -598,6 +598,7 @@ def _chunk_kda_fwd_with_cumulative_g(
chunk_indices: torch.Tensor | None = None,
chunk_size: int = FLA_CHUNK_SIZE,
safe_gate: bool = False,
out: torch.Tensor | None = None,
):
# `g` must already be chunk-local cumulatively-summed AND scaled by
# RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)).
@@ -642,7 +643,7 @@ def _chunk_kda_fwd_with_cumulative_g(
g=g,
A=Aqk,
h=h,
o=v,
o=v if out is None else out,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
@@ -706,6 +707,7 @@ def chunk_kda_with_fused_gate_fwd(
output_final_state: bool,
lower_bound: float | None = None,
cu_seqlens: torch.Tensor | None = None,
out: torch.Tensor | None = None,
):
chunk_size = FLA_CHUNK_SIZE
chunk_indices = (
@@ -736,6 +738,7 @@ def chunk_kda_with_fused_gate_fwd(
chunk_indices=chunk_indices,
chunk_size=chunk_size,
safe_gate=lower_bound is not None,
out=out,
)
@@ -787,6 +790,7 @@ def chunk_kda_with_fused_gate(
lower_bound: float | None = None,
use_qk_l2norm_in_kernel: bool = False,
cu_seqlens: torch.Tensor | None = None,
out: torch.Tensor | None = None,
**kwargs,
):
"""Run chunk KDA from raw gate and beta projections."""
@@ -810,6 +814,7 @@ def chunk_kda_with_fused_gate(
output_final_state=output_final_state,
lower_bound=lower_bound,
cu_seqlens=cu_seqlens,
out=out,
)
return o, final_state
vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py
Vendored
+3 -1
View File
@@ -591,6 +591,7 @@ def fused_recurrent_kda_packed_decode(
initial_state: torch.Tensor,
state_indices: torch.Tensor,
scale: float | None = None,
out: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Run one-token KDA decode directly from packed post-conv QKV."""
if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1:
@@ -639,7 +640,8 @@ def fused_recurrent_kda_packed_decode(
if scale is None:
scale = K**-0.5
out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device)
if out is None:
out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device)
grid = (cdiv(V, BV), B * H)
fused_recurrent_kda_packed_decode_kernel[grid](
mixed_qkv=mixed_qkv,