[CI] Add Kimi-K3-pruned75-DSpark-TP4 gsm8k eval (#54817)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Misha Goin committed
Sep 2, 2026 at 10:18 UTC
872084fb773b4dedf9f72a08724451f72b115f89
5 files changed
+64
-30
.buildkite/test_areas/lm_eval.yaml
+3
-1
@@ -109,9 +109,11 @@ steps:
109
num_devices: 4
110
optional: true
111
source_file_dependencies:
112
- - tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml
112
+ - tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TEP4.yaml
113
+ - tests/evals/gsm8k/configs/Kimi-K3-pruned75-DSpark-TP4.yaml
114
- tests/evals/gsm8k/configs/models-spec-decode.txt
115
- vllm/models/deepseek_v4/
116
+ - vllm/models/kimi_k3/
117
- vllm/model_executor/models/qwen3_dspark.py
118
- vllm/v1/spec_decode/metrics.py
119
- vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TEP4.yaml
renamed
tests/evals/gsm8k/configs/Kimi-K3-pruned75-DSpark-TP4.yaml
new
+26
@@ -0,0 +1,26 @@
1
+model_name: "mgoin/Kimi-K3-pruned75"
2
+accuracy_threshold: 0.33
3
+min_acceptance_length: 4.8
4
+num_questions: 1319
5
+num_fewshot: 5
6
+startup_max_wait_seconds: 1800
7
+server_args: >-
8
+ --trust-remote-code
9
+ --gpu-memory-utilization 0.90
10
+ --tensor-parallel-size 4
11
+ --load-format fastsafetensors
12
+ --no-enable-flashinfer-autotune
13
+ --max-model-len 1048576
14
+ --max-num-batched-tokens 16384
15
+ --max-num-seqs 128
16
+ --kv-cache-dtype fp8
17
+ --attention-config '{"use_prefill_query_quantization":true,"mla_prefill_backend":"flashinfer"}'
18
+ --enable-prefix-caching
19
+ --enable-auto-tool-choice
20
+ --tool-call-parser kimi_k3
21
+ --reasoning-parser kimi_k3
22
+ --speculative-config '{"method":"dspark",
23
+ "model":"RedHatAI/Kimi-K3-speculator.dspark",
24
+ "attention_backend":"FLASH_ATTN","num_speculative_tokens":8,
25
+ "draft_sample_method":"probabilistic",
26
+ "rejection_sample_method":"block"}'
tests/evals/gsm8k/configs/models-spec-decode.txt
+2
-1
@@ -1 +1,2 @@
1
-DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml
1
+DeepSeek-V4-Flash-DSpark-confidence-TEP4.yaml
2
+Kimi-K3-pruned75-DSpark-TP4.yaml
vllm/v1/attention/backends/mla/prefill/flashinfer.py
+33
-28
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, ClassVar
7
import torch
8
9
import vllm.envs as envs
10
+from vllm.utils.gpu_sync_debug import gpu_sync_allowed
11
from vllm.v1.attention.backends.mla.prefill.base import (
12
MLADimensions,
13
MLAPrefillBackend,
@@ -161,39 +162,43 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
162
kv_indptr = qo_indptr.clone()
163
164
assert self._prefill_main is not None
164
- self._prefill_main.plan(
165
- qo_indptr=qo_indptr,
166
- kv_indptr=kv_indptr,
167
- num_qo_heads=num_qo_heads,
168
- num_kv_heads=num_kv_heads,
169
- head_dim_qk=head_dim_qk,
170
- head_dim_vo=head_dim_vo,
171
- causal=True,
172
- sm_scale=global_hyperparameters.sm_scale,
173
- window_left=global_hyperparameters.window_left,
174
- logits_soft_cap=global_hyperparameters.logits_soft_cap,
175
- q_data_type=prefill_metadata.q_data_type,
176
- o_data_type=prefill_metadata.output_dtype,
177
- )
165
+ # FlashInfer's plan() currently copies GPU indptrs to the CPU.
166
+ # This will be fixed by https://github.com/vllm-project/vllm/pull/52657.
167
+ with gpu_sync_allowed():
168
+ self._prefill_main.plan(
169
+ qo_indptr=qo_indptr,
170
+ kv_indptr=kv_indptr,
171
+ num_qo_heads=num_qo_heads,
172
+ num_kv_heads=num_kv_heads,
173
+ head_dim_qk=head_dim_qk,
174
+ head_dim_vo=head_dim_vo,
175
+ causal=True,
176
+ sm_scale=global_hyperparameters.sm_scale,
177
+ window_left=global_hyperparameters.window_left,
178
+ logits_soft_cap=global_hyperparameters.logits_soft_cap,
179
+ q_data_type=prefill_metadata.q_data_type,
180
+ o_data_type=prefill_metadata.output_dtype,
181
+ )
182
183
if has_context:
184
chunked_context = prefill_metadata.chunked_context
185
assert chunked_context is not None
186
for chunk in chunked_context.chunks:
183
- self._prefill_chunks[chunk.index].plan(
184
- qo_indptr=chunk.query_start_loc,
185
- kv_indptr=chunk.cu_seq_lens,
186
- num_qo_heads=num_qo_heads,
187
- num_kv_heads=num_kv_heads,
188
- head_dim_qk=head_dim_qk,
189
- head_dim_vo=head_dim_vo,
190
- causal=False,
191
- sm_scale=global_hyperparameters.sm_scale,
192
- window_left=global_hyperparameters.window_left,
193
- logits_soft_cap=global_hyperparameters.logits_soft_cap,
194
- q_data_type=prefill_metadata.q_data_type,
195
- o_data_type=prefill_metadata.output_dtype,
196
- )
187
+ with gpu_sync_allowed():
188
+ self._prefill_chunks[chunk.index].plan(
189
+ qo_indptr=chunk.query_start_loc,
190
+ kv_indptr=chunk.cu_seq_lens,
191
+ num_qo_heads=num_qo_heads,
192
+ num_kv_heads=num_kv_heads,
193
+ head_dim_qk=head_dim_qk,
194
+ head_dim_vo=head_dim_vo,
195
+ causal=False,
196
+ sm_scale=global_hyperparameters.sm_scale,
197
+ window_left=global_hyperparameters.window_left,
198
+ logits_soft_cap=global_hyperparameters.logits_soft_cap,
199
+ q_data_type=prefill_metadata.q_data_type,
200
+ o_data_type=prefill_metadata.output_dtype,
201
+ )
202
203
def supports_out(self) -> bool:
204
# Planned with head_dim_vo == v_head_dim, so the output is unpadded.