[Bugfix][Kimi-K3] Do not classify a stateless first chunk as a decode (#51483)
Signed-off-by: Oleksandr Zakharchuk <oleksandr.zakharchuk@gmail.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
sashko-zakharchuk committed
Sep 16, 2026 at 15:24 UTC
f8b5c11468f665c75968e3a7c12f16ca074f3a30
2 files changed
+113
-2
tests/models/kimi_k3/test_kda_metadata.py
+88
@@ -934,3 +934,91 @@ def test_aligned_block_table_matches_shared_gdn():
934
)
935
936
torch.testing.assert_close(actual, expected)
937
+
938
+
939
+def _build_non_spec(batch, is_prefilling, full_cuda_graph=False):
940
+ common_attn_metadata = create_common_attn_metadata(
941
+ batch, BLOCK_SIZE, DEVICE
942
+ ).replace(
943
+ is_prefilling=None
944
+ if is_prefilling is None
945
+ else torch.tensor(is_prefilling, dtype=torch.bool)
946
+ )
947
+ builder = _make_builder(
948
+ KimiK3KDAMetadataBuilder,
949
+ num_speculative_tokens=0,
950
+ full_cuda_graph=full_cuda_graph,
951
+ )
952
+ return builder, common_attn_metadata, builder.build(0, common_attn_metadata)
953
+
954
+
955
+def test_one_token_first_chunk_excludes_padding():
956
+ """Neither padding requests nor padding tokens count as prefill work."""
957
+ common = create_common_attn_metadata(
958
+ BatchSpec(seq_lens=[100, 1, 0, 0], query_lens=[1, 1, 0, 0]),
959
+ BLOCK_SIZE,
960
+ DEVICE,
961
+ ).replace(
962
+ is_prefilling=torch.tensor([False, True, False, False], dtype=torch.bool),
963
+ num_actual_tokens=4,
964
+ )
965
+ builder = _make_builder(
966
+ KimiK3KDAMetadataBuilder, num_speculative_tokens=0, full_cuda_graph=False
967
+ )
968
+ actual = builder.build(0, common)
969
+
970
+ assert actual.num_decodes == 1
971
+ assert actual.num_prefills == 1
972
+ assert actual.num_decode_tokens == 1
973
+ assert actual.num_prefill_tokens == 1
974
+
975
+
976
+@pytest.mark.parametrize(
977
+ ("seq_len", "query_len", "is_prefilling", "num_prefills"),
978
+ [
979
+ pytest.param(1, 1, True, 1, id="first-chunk"),
980
+ pytest.param(65, 1, True, 0, id="resumed-chunk"),
981
+ pytest.param(0, 0, True, 0, id="padding"),
982
+ pytest.param(1, 1, None, 0, id="missing-prefill-flag"),
983
+ ],
984
+)
985
+def test_one_token_chunk_classification(
986
+ seq_len, query_len, is_prefilling, num_prefills
987
+):
988
+ """Only a real first chunk with a prefill flag needs state initialization."""
989
+ _, _, actual = _build_non_spec(
990
+ BatchSpec(seq_lens=[100, seq_len], query_lens=[1, query_len]),
991
+ is_prefilling=None if is_prefilling is None else [False, is_prefilling],
992
+ )
993
+
994
+ assert actual.num_prefills == num_prefills
995
+ assert actual.num_decodes == 2 - num_prefills
996
+ assert actual.num_prefill_tokens == num_prefills
997
+ assert actual.num_decode_tokens == 1 + query_len - num_prefills
998
+ if num_prefills:
999
+ assert actual.has_initial_state is not None
1000
+ assert actual.has_initial_state.tolist() == [True, False]
1001
+ else:
1002
+ assert actual.has_initial_state is None
1003
+
1004
+
1005
+def test_cudagraph_capture_batch_stays_decode_only():
1006
+ """Capture rows have no history, but must still select decode kernels."""
1007
+ batch = BatchSpec(seq_lens=[1] * 4, query_lens=[1] * 4)
1008
+ common_attn_metadata = create_common_attn_metadata(
1009
+ batch, BLOCK_SIZE, DEVICE
1010
+ ).replace(is_prefilling=torch.zeros(4, dtype=torch.bool))
1011
+ builder = _make_builder(
1012
+ KimiK3KDAMetadataBuilder,
1013
+ num_speculative_tokens=0,
1014
+ full_cuda_graph=True,
1015
+ )
1016
+ actual = builder.build_for_cudagraph_capture(common_attn_metadata)
1017
+
1018
+ assert actual.num_prefills == 0
1019
+ assert actual.num_decodes == 4
1020
+ assert actual.has_initial_state is None
1021
+ staged = actual.non_spec_state_indices_tensor
1022
+ assert staged is not None
1023
+ assert staged.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr()
1024
+ torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0])
vllm/models/kimi_k3/nvidia/kda_metadata.py
+25
-2
@@ -427,10 +427,33 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
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.
430
+ # V2 already excludes prefills from full decode graphs via has_prefill.
431
+ # Classify first chunks as prefills to mask recycled state;
432
+ # resumed one-token chunks can still use the decode kernels.
433
+ assert m.seq_lens_cpu_upper_bound is not None
434
+ query_lens_cpu = query_start_loc_cpu.diff()
435
+ no_prior_state = (query_lens_cpu > 0) & (
436
+ m.seq_lens_cpu_upper_bound <= query_lens_cpu
437
+ )
438
+ # Capture batches also have seq_len == query_len, but are not prefills.
439
+ if m.is_prefilling is not None:
440
+ no_prior_state &= m.is_prefilling
441
+ else:
442
+ no_prior_state = torch.zeros_like(no_prior_state)
443
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
432
- split_decodes_and_prefills(m, decode_threshold=1)
444
+ split_decodes_and_prefills(
445
+ m.replace(is_prefilling=no_prior_state),
446
+ decode_threshold=1,
447
+ treat_short_extends_as_decodes=False,
448
+ )
449
)
450
+ # Exclude trailing padding from both prefill counts.
451
+ if num_prefills:
452
+ num_prefills -= int((query_lens_cpu[num_decodes:] == 0).sum())
453
+ num_prefill_tokens = (
454
+ int(query_start_loc_cpu[num_decodes + num_prefills])
455
+ - num_decode_tokens
456
+ )
457
num_spec_decode_tokens = 0
458
spec_token_indx = None
459
non_spec_token_indx = None