[Spec Decode] Register Qwen3.6 dSpark acceptance coverage (#51310)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Michael Goin committed
Aug 7, 2026 at 17:07 UTC
c8fe1d5715a75a0e71712ec6717dd7ff71d1a8d0
2 files changed
+125
-103
tests/evals/gsm8k/gsm8k_eval.py
+7
-1
@@ -292,6 +292,7 @@ def evaluate_gsm8k_offline(
292
temperature: float = 0.0,
293
gen_prefix: str = "",
294
use_chat_completions: bool = False,
295
+ chat_template_kwargs: dict[str, object] | None = None,
296
) -> dict[str, float | int]:
297
"""Evaluate GSM8K accuracy using an offline vllm.LLM object.
298
@@ -300,6 +301,7 @@ def evaluate_gsm8k_offline(
301
302
When ``use_chat_completions=True``, prompts go through the chat template via
303
``llm.chat()`` instead of raw completion (for instruction-tuned models).
304
+ ``chat_template_kwargs`` are forwarded to ``llm.chat()`` when provided.
305
"""
306
from vllm import SamplingParams
307
@@ -318,7 +320,11 @@ def evaluate_gsm8k_offline(
320
tic = time.perf_counter()
321
if use_chat_completions:
322
conversations = [[{"role": "user", "content": p}] for p in prompts]
321
- outputs = llm.chat(conversations, sampling_params)
323
+ outputs = llm.chat(
324
+ conversations,
325
+ sampling_params,
326
+ chat_template_kwargs=chat_template_kwargs,
327
+ )
328
else:
329
outputs = llm.generate(prompts, sampling_params)
330
latency = time.perf_counter() - tic
tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py
+118
-102
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from dataclasses import dataclass
5
+
6
import pytest
7
import torch
8
@@ -10,120 +12,134 @@ from vllm.distributed import cleanup_dist_env_and_memory
12
13
from ...utils import compute_acceptance_len, compute_acceptance_rate
14
13
-
14
-def test_gemma4_dspark_correctness_and_acceptance_rate(
15
+REGRESSION_TOLERANCE = 0.95
16
+
17
+
18
+@dataclass(frozen=True)
19
+class DSparkCorrectnessConfig:
20
+ model: str
21
+ draft_model: str
22
+ reference_accuracy: float
23
+ reference_acceptance_rate: float
24
+ reference_acceptance_len: float
25
+ num_speculative_tokens: int = 7
26
+ max_model_len: int = 4096
27
+ max_num_seqs: int | None = None
28
+ num_questions: int = 1319
29
+ max_tokens: int = 256
30
+ use_chat_completions: bool = False
31
+ chat_template_kwargs: dict[str, object] | None = None
32
+ attention_backend: str | None = None
33
+ gpu_memory_utilization: float = 0.92
34
+ enforce_eager: bool = False
35
+ disable_flashinfer_sampler: bool = False
36
+ language_model_only: bool = False
37
+
38
+
39
+# References from 12 full GSM8K runs at temperature 1.0: accuracy 0.782-0.814,
40
+# acceptance rate 0.418-0.434, acceptance length 3.928-4.037.
41
+QWEN3_DSPARK_DEEPSPEC = DSparkCorrectnessConfig(
42
+ model="Qwen/Qwen3-4B-FP8",
43
+ draft_model="deepseek-ai/dspark_qwen3_4b_block7",
44
+ reference_accuracy=0.801,
45
+ reference_acceptance_rate=0.428,
46
+ reference_acceptance_len=3.994,
47
+)
48
+
49
+# References from five 200-question runs at temperature 1.0: accuracy
50
+# 0.900-0.955, acceptance rate 0.578-0.595, acceptance length 5.044-5.167.
51
+GEMMA4_DSPARK_DEEPSPEC = DSparkCorrectnessConfig(
52
+ model="RedHatAI/gemma-4-12B-it-NVFP4",
53
+ draft_model="deepseek-ai/dspark_gemma4_12b_block7",
54
+ reference_accuracy=0.92,
55
+ reference_acceptance_rate=0.58,
56
+ reference_acceptance_len=5.0,
57
+ max_num_seqs=32,
58
+ num_questions=200,
59
+ use_chat_completions=True,
60
+ gpu_memory_utilization=0.85,
61
+)
62
+
63
+# Reference from 200 GSM8K questions at temperature 1.0.
64
+QWEN3_6_DSPARK_SPECULATORS = DSparkCorrectnessConfig(
65
+ model="RedHatAI/Qwen3.6-35B-A3B-NVFP4",
66
+ draft_model="RedHatAI/Qwen3.6-35B-A3B-speculator.dspark",
67
+ reference_accuracy=0.85,
68
+ reference_acceptance_rate=0.41,
69
+ reference_acceptance_len=4.25,
70
+ num_speculative_tokens=8,
71
+ max_num_seqs=32,
72
+ num_questions=200,
73
+ use_chat_completions=True,
74
+ chat_template_kwargs={"enable_thinking": False},
75
+ gpu_memory_utilization=0.85,
76
+ language_model_only=True,
77
+)
78
+
79
+
80
+@pytest.mark.parametrize(
81
+ "config",
82
+ [
83
+ pytest.param(QWEN3_DSPARK_DEEPSPEC, id="qwen3-deepspec"),
84
+ pytest.param(GEMMA4_DSPARK_DEEPSPEC, id="gemma4-deepspec"),
85
+ pytest.param(QWEN3_6_DSPARK_SPECULATORS, id="qwen3.6-speculators"),
86
+ ],
87
+)
88
+def test_dspark_correctness_and_acceptance_rate(
89
monkeypatch: pytest.MonkeyPatch,
90
+ config: DSparkCorrectnessConfig,
91
):
17
- """
18
- E2E test for Gemma4 DSpark speculative decoding: acceptance rate/length
19
- regression coverage plus GSM8K correctness, at temperature=1.0 to exercise
20
- the probabilistic draft-sampling/rejection-sampling path (not just greedy).
21
-
22
- gemma-4-12B is instruct-tuned, so GSM8K is run through the chat template
23
- (use_chat_completions=True; raw few-shot completion collapses to a few
24
- percent). Reference: measured over 5 runs of 200 GSM8K questions at
25
- temperature=1.0 (prefix caching disabled):
26
- accuracy: min=0.900 max=0.955 mean=0.937
27
- acceptance_rate: min=0.578 max=0.595 mean=0.588
28
- acceptance_len: min=5.044 max=5.167 mean=5.116
29
- Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
30
- """
31
- monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
32
-
33
- spec_llm = LLM(
34
- model="RedHatAI/gemma-4-12B-it-NVFP4",
35
- trust_remote_code=True,
36
- speculative_config={
37
- "method": "dspark",
38
- "model": "deepseek-ai/dspark_gemma4_12b_block7",
39
- "num_speculative_tokens": 7,
40
- "draft_sample_method": "probabilistic",
41
- },
42
- max_model_len=4096,
43
- max_num_seqs=32,
44
- gpu_memory_utilization=0.85,
45
- enforce_eager=True,
46
- enable_prefix_caching=False,
47
- disable_log_stats=False,
48
- )
92
+ """Guard GSM8K accuracy and acceptance metrics for DSpark models."""
93
+ speculative_config = {
94
+ "method": "dspark",
95
+ "model": config.draft_model,
96
+ "num_speculative_tokens": config.num_speculative_tokens,
97
+ "draft_sample_method": "probabilistic",
98
+ }
99
+ if config.attention_backend is not None:
100
+ speculative_config["attention_backend"] = config.attention_backend
101
+
102
+ llm_kwargs = {
103
+ "model": config.model,
104
+ "trust_remote_code": True,
105
+ "speculative_config": speculative_config,
106
+ "max_model_len": config.max_model_len,
107
+ "gpu_memory_utilization": config.gpu_memory_utilization,
108
+ "enforce_eager": config.enforce_eager,
109
+ "enable_prefix_caching": False,
110
+ "disable_log_stats": False,
111
+ }
112
+ if config.max_num_seqs is not None:
113
+ llm_kwargs["max_num_seqs"] = config.max_num_seqs
114
+ if config.language_model_only:
115
+ llm_kwargs["language_model_only"] = True
116
+
117
+ spec_llm = LLM(**llm_kwargs)
118
try:
119
results = evaluate_gsm8k_offline(
51
- spec_llm, num_questions=200, temperature=1.0, use_chat_completions=True
120
+ spec_llm,
121
+ num_questions=config.num_questions,
122
+ max_tokens=config.max_tokens,
123
+ temperature=1.0,
124
+ use_chat_completions=config.use_chat_completions,
125
+ chat_template_kwargs=config.chat_template_kwargs,
126
)
53
- gsm8k_accuracy = results["accuracy"]
54
-
127
+ accuracy = results["accuracy"]
128
metrics = spec_llm.get_metrics()
129
acceptance_rate = compute_acceptance_rate(metrics)
130
acceptance_len = compute_acceptance_len(metrics)
131
print(
59
- f"Gemma4 DSpark acceptance_rate={acceptance_rate:.2f}, "
60
- f"acceptance_len={acceptance_len:.2f}, "
61
- f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
132
+ f"{config.model}: GSM8K accuracy={accuracy:.3f}, "
133
+ f"acceptance_rate={acceptance_rate:.2f}, "
134
+ f"acceptance_len={acceptance_len:.2f}"
135
)
136
64
- assert acceptance_rate >= 0.58 * 0.9
65
- assert acceptance_len >= 5.0 * 0.9
66
- assert gsm8k_accuracy >= 0.92 * 0.9
137
+ assert accuracy >= config.reference_accuracy * REGRESSION_TOLERANCE
138
+ assert (
139
+ acceptance_rate >= config.reference_acceptance_rate * REGRESSION_TOLERANCE
140
+ )
141
+ assert acceptance_len >= config.reference_acceptance_len * REGRESSION_TOLERANCE
142
finally:
143
del spec_llm
144
torch.accelerator.empty_cache()
145
cleanup_dist_env_and_memory()
71
-
72
-
73
-@pytest.fixture
74
-def dspark_config():
75
- target_model = "Qwen/Qwen3-4B-FP8"
76
- draft_model = "deepseek-ai/dspark_qwen3_4b_block7"
77
-
78
- return dict(
79
- model=target_model,
80
- trust_remote_code=True,
81
- speculative_config={
82
- "method": "dspark",
83
- "model": draft_model,
84
- "num_speculative_tokens": 7,
85
- "attention_backend": "FLASH_ATTN",
86
- "draft_sample_method": "probabilistic",
87
- },
88
- max_model_len=4096,
89
- disable_log_stats=False,
90
- )
91
-
92
-
93
-def test_dspark_correctness_and_acceptance_rate(dspark_config):
94
- """
95
- E2E test for DSpark speculative decoding: acceptance rate/length
96
- regression coverage plus GSM8K correctness, at temperature=1.0 to
97
- exercise the probabilistic draft-sampling/rejection-sampling path
98
- (not just greedy).
99
-
100
- Uses Qwen/Qwen3-4B-FP8 as target with the dspark_qwen3_4b_block7 draft
101
- model. Reference: measured over 12 runs of the full GSM8K set at
102
- temperature=1.0 (prefix caching disabled to avoid cross-run reuse):
103
- accuracy: min=0.782 max=0.814 mean=0.801
104
- acceptance_rate: min=0.418 max=0.434 mean=0.428
105
- acceptance_len: min=3.928 max=4.037 mean=3.994
106
- Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
107
- """
108
- spec_llm = LLM(**dspark_config)
109
-
110
- results = evaluate_gsm8k_offline(spec_llm, temperature=1.0)
111
- gsm8k_accuracy = results["accuracy"]
112
-
113
- metrics = spec_llm.get_metrics()
114
- acceptance_rate = compute_acceptance_rate(metrics)
115
- acceptance_len = compute_acceptance_len(metrics)
116
-
117
- print(
118
- f"DSpark acceptance_rate={acceptance_rate:.2f}, "
119
- f"acceptance_len={acceptance_len:.2f}, "
120
- f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
121
- )
122
-
123
- assert acceptance_rate >= 0.428 * 0.9
124
- assert acceptance_len >= 3.994 * 0.9
125
- assert gsm8k_accuracy >= 0.801 * 0.9
126
-
127
- del spec_llm
128
- torch.accelerator.empty_cache()
129
- cleanup_dist_env_and_memory()