@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ee59a7c615
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 19:00 UTC

[Tests] Add GSM8k check to SpecDec E2E tests (#34772)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>

Benjamin Chislett committed Feb 25, 2026 at 07:51 UTC ee59a7c61574485cf4ddbc6037ba557941be5c56
2 files changed +237 -143
tests/evals/gsm8k/gsm8k_eval.py
+95 -43
@@ -110,29 +110,16 @@ async def call_vllm_api(
110 return "", 0
111
112
113 -def evaluate_gsm8k(
113 +def _build_gsm8k_prompts(
114 num_questions: int = 1319,
115 num_shots: int = 5,
116 - max_tokens: int = 256,
117 - host: str = "http://127.0.0.1",
118 - port: int = 8000,
119 - temperature: float = 0.0,
120 - seed: int | None = 42,
121 -) -> dict[str, float | int]:
122 - """
123 - Evaluate GSM8K accuracy using vLLM serve endpoint.
124 -
125 - Returns dict with accuracy, invalid_rate, latency, etc.
126 - """
127 - base_url = f"{host}:{port}"
128 -
129 - # Load GSM8K train and test data
116 +) -> tuple[list[str], list[int]]:
117 + """Build few-shot GSM8K completion prompts and ground-truth labels."""
118 + if num_questions == 0:
119 + return [], []
120 train_data, test_data = load_gsm8k_data()
131 -
132 - # Limit to available test questions
121 num_questions = min(num_questions, len(test_data))
122
135 - # Build few-shot examples from train split (like lm-eval does)
123 few_shot_examples = ""
124 for i in range(num_shots):
125 few_shot_examples += (
@@ -140,25 +127,74 @@ def evaluate_gsm8k(
127 f"Answer: {train_data[i]['answer']}\n\n"
128 )
129
143 - # Prepare test questions and labels from test split
144 - questions = []
130 + prompts = []
131 labels = []
132 for i in range(num_questions):
147 - questions.append(f"Question: {test_data[i]['question']}\nAnswer:")
133 + prompts.append(
134 + few_shot_examples + f"Question: {test_data[i]['question']}\nAnswer:"
135 + )
136 labels.append(get_answer_value(test_data[i]["answer"]))
137
138 assert all(label != INVALID for label in labels), "Some labels are invalid"
139 + return prompts, labels
140 +
141 +
142 +def _score_gsm8k(
143 + states: list[str],
144 + output_tokens: list[int],
145 + labels: list[int],
146 + num_shots: int,
147 + max_tokens: int,
148 + latency: float,
149 +) -> dict[str, float | int]:
150 + """Score GSM8K responses and return a results dict."""
151 + num_questions = len(labels)
152 + preds = [get_answer_value(state) for state in states]
153 + accuracy = np.mean(np.array(preds) == np.array(labels))
154 + invalid_rate = np.mean(np.array(preds) == INVALID)
155 + total_output_tokens = sum(output_tokens)
156 + tokens_per_second = total_output_tokens / latency if latency > 0 else 0.0
157 +
158 + return {
159 + "accuracy": accuracy,
160 + "invalid_rate": invalid_rate,
161 + "latency": latency,
162 + "questions_per_second": num_questions / latency if latency > 0 else 0.0,
163 + "total_output_tokens": total_output_tokens,
164 + "tokens_per_second": tokens_per_second,
165 + "num_questions": num_questions,
166 + "num_shots": num_shots,
167 + "max_tokens": max_tokens,
168 + "timestamp": time.time(),
169 + }
170 +
171 +
172 +def evaluate_gsm8k(
173 + num_questions: int = 1319,
174 + num_shots: int = 5,
175 + max_tokens: int = 256,
176 + host: str = "http://127.0.0.1",
177 + port: int = 8000,
178 + temperature: float = 0.0,
179 + seed: int | None = 42,
180 +) -> dict[str, float | int]:
181 + """
182 + Evaluate GSM8K accuracy using vLLM serve endpoint.
183 +
184 + Returns dict with accuracy, invalid_rate, latency, etc.
185 + """
186 + base_url = f"{host}:{port}"
187 + prompts, labels = _build_gsm8k_prompts(num_questions, num_shots)
188 + num_questions = len(prompts)
189
152 - # Run evaluation
190 async def run_async_evaluation():
191 states: list[str] = [""] * num_questions
192 output_tokens: list[int] = [0] * num_questions
193
194 async def get_answer(session: aiohttp.ClientSession, i: int) -> tuple[str, int]:
158 - prompt = few_shot_examples + questions[i]
195 answer, tokens = await call_vllm_api(
196 session=session,
161 - prompt=prompt,
197 + prompt=prompts[i],
198 temperature=temperature,
199 max_tokens=max_tokens,
200 stop=["Question", "Assistant:", "<|separator|>"],
@@ -183,27 +219,43 @@ def evaluate_gsm8k(
219 states, output_tokens = asyncio.run(run_async_evaluation())
220 latency = time.perf_counter() - tic
221
186 - # Compute metrics
187 - preds = [get_answer_value(state) for state in states]
188 - accuracy = np.mean(np.array(preds) == np.array(labels))
189 - invalid_rate = np.mean(np.array(preds) == INVALID)
190 - total_output_tokens = sum(output_tokens)
191 - tokens_per_second = total_output_tokens / latency if latency > 0 else 0.0
222 + return _score_gsm8k(states, output_tokens, labels, num_shots, max_tokens, latency)
223
193 - result = {
194 - "accuracy": accuracy,
195 - "invalid_rate": invalid_rate,
196 - "latency": latency,
197 - "questions_per_second": num_questions / latency,
198 - "total_output_tokens": total_output_tokens,
199 - "tokens_per_second": tokens_per_second,
200 - "num_questions": num_questions,
201 - "num_shots": num_shots,
202 - "max_tokens": max_tokens,
203 - "timestamp": time.time(),
204 - }
224
206 - return result
225 +def evaluate_gsm8k_offline(
226 + llm,
227 + num_questions: int = 1319,
228 + num_shots: int = 5,
229 + max_tokens: int = 256,
230 + temperature: float = 0.0,
231 +) -> dict[str, float | int]:
232 + """Evaluate GSM8K accuracy using an offline vllm.LLM object.
233 +
234 + Same prompts and scoring as evaluate_gsm8k(), but runs generation
235 + directly via llm.generate() instead of calling a server over HTTP.
236 + """
237 + from vllm import SamplingParams
238 +
239 + prompts, labels = _build_gsm8k_prompts(num_questions, num_shots)
240 +
241 + sampling_params = SamplingParams(
242 + temperature=temperature,
243 + max_tokens=max_tokens,
244 + stop=["Question", "Assistant:", "<|separator|>"],
245 + )
246 +
247 + print(
248 + f"Running offline GSM8K evaluation: {len(prompts)} questions, {num_shots}-shot"
249 + )
250 +
251 + tic = time.perf_counter()
252 + outputs = llm.generate(prompts, sampling_params)
253 + latency = time.perf_counter() - tic
254 +
255 + states = [o.outputs[0].text for o in outputs]
256 + output_tokens = [len(o.outputs[0].token_ids) for o in outputs]
257 +
258 + return _score_gsm8k(states, output_tokens, labels, num_shots, max_tokens, latency)
259
260
261 def main() -> None:
tests/v1/e2e/test_spec_decode.py
+142 -100
@@ -8,6 +8,7 @@ from typing import Any
8 import pytest
9 import torch
10
11 +from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline
12 from tests.utils import get_attn_backend_list_based_on_platform, large_gpu_mark
13 from vllm import LLM, SamplingParams
14 from vllm.assets.base import VLLM_S3_BUCKET_URL
@@ -35,53 +36,57 @@ def _skip_if_insufficient_gpus_for_tp(tp_size: int):
36 Messages = list[dict[str, Any]]
37
38
38 -def get_test_prompts(
39 - mm_enabled: bool, quiet: bool = False, num_prompts: int = 100
40 -) -> list[Messages]:
41 - prompt_types = ["repeat", "sentence"]
39 +def get_test_prompts(mm_enabled: bool, num_prompts: int = 100) -> list[Messages]:
40 + prompt_types = ["repeat", "gsm8k"]
41 if mm_enabled:
42 prompt_types.append("mm")
44 - prompts = []
43 + prompts: list[Messages] = []
44
46 - random.seed(0)
47 - random_prompt_type_choices = random.choices(prompt_types, k=num_prompts)
48 -
49 - if not quiet:
50 - print(f"Prompt types: {random_prompt_type_choices}")
45 + num_repeat_prompts = num_prompts // len(prompt_types)
46 + if mm_enabled:
47 + num_gsm8k_prompts = num_prompts // len(prompt_types)
48 + num_mm_prompts = num_prompts - num_repeat_prompts - num_gsm8k_prompts
49 + else:
50 + num_mm_prompts = 0
51 + num_gsm8k_prompts = num_prompts - num_repeat_prompts
52
53 # Generate a mixed batch of prompts, some of which can be easily
54 # predicted by n-gram matching and some which likely cannot.
54 - for kind in random_prompt_type_choices:
55 + random.seed(0)
56 + for _ in range(num_repeat_prompts):
57 word_choices = ["test", "temp", "hello", "where"]
58 word = random.choice(word_choices)
57 - prompt: str | list[dict[str, Any]] = ""
58 - if kind == "repeat":
59 - prompt = f"""
60 - please repeat the word '{word}' 10 times.
61 - give no other output than the word at least ten times in a row,
62 - in lowercase with spaces between each word and without quotes.
63 - """
64 - elif kind == "sentence":
65 - prompt = f"""
66 - please give a ten-word sentence that
67 - uses the word {word} at least once.
68 - give no other output than that simple sentence without quotes.
69 - """
70 - elif kind == "mm":
71 - placeholders = [
59 + prompts.append(
60 + [
61 {
73 - "type": "image_url",
74 - "image_url": {
75 - "url": f"{VLLM_S3_BUCKET_URL}/{VLM_IMAGES_DIR}/stop_sign.jpg"
76 - },
62 + "role": "user",
63 + "content": f"""
64 + please repeat the word '{word}' 10 times.
65 + give no other output than the word at least ten times in a row,
66 + in lowercase with spaces between each word and without quotes.
67 + """,
68 }
69 ]
79 - prompt = [
80 - *placeholders,
81 - {"type": "text", "text": "The meaning of the image is"},
82 - ]
83 - else:
84 - raise ValueError(f"Unknown prompt type: {kind}")
70 + )
71 + prompts.extend(
72 + [{"role": "user", "content": prompt}]
73 + for prompt in _build_gsm8k_prompts(
74 + num_questions=num_gsm8k_prompts, num_shots=5
75 + )[0]
76 + )
77 + for _ in range(num_mm_prompts):
78 + placeholders = [
79 + {
80 + "type": "image_url",
81 + "image_url": {
82 + "url": f"{VLLM_S3_BUCKET_URL}/{VLM_IMAGES_DIR}/stop_sign.jpg"
83 + },
84 + }
85 + ]
86 + prompt = [
87 + *placeholders,
88 + {"type": "text", "text": "The meaning of the image is"},
89 + ]
90 prompts.append([{"role": "user", "content": prompt}])
91
92 return prompts
@@ -113,6 +118,25 @@ def model_name():
118 return "meta-llama/Llama-3.1-8B-Instruct"
119
120
121 +def evaluate_llm_for_gsm8k(llm: LLM, expected_accuracy_threshold: float = 0.70) -> None:
122 + """Evaluate the LLM on GSM8K and check that accuracy is above a sanity threshold.
123 +
124 + The default threshold assumes the LLM uses the same target model as the "model_name"
125 + fixture, with max model len == 4096. Precomputed reference value is 75% to 80%
126 + on GSM8K with greedy decoding, so we check that it's above a sanity threshold of 70%
127 + to verify that the model is correct.
128 + """
129 + if expected_accuracy_threshold <= 0.0:
130 + print("Skipping GSM8K evaluation")
131 + return
132 + results = evaluate_gsm8k_offline(llm)
133 + accuracy = results["accuracy"]
134 + print(f"GSM8K accuracy: {accuracy:.3f}")
135 + assert accuracy >= expected_accuracy_threshold, (
136 + f"Expected GSM8K accuracy >= {expected_accuracy_threshold}, got {accuracy:.3f}"
137 + )
138 +
139 +
140 @pytest.fixture(autouse=True)
141 def reset_torch_dynamo():
142 """Reset torch dynamo cache before each test"""
@@ -138,41 +162,14 @@ def reset_torch_dynamo():
162 )
163 def test_ngram_and_suffix_correctness(
164 speculative_config: dict,
141 - monkeypatch: pytest.MonkeyPatch,
142 - sampling_config: SamplingParams,
165 model_name: str,
166 ):
145 - """
146 - Compare the outputs of an original LLM and a speculative LLM
147 - should be the same when using ngram speculative decoding.
148 - """
149 - test_prompts = get_test_prompts(mm_enabled=False)
150 -
151 - ref_llm = LLM(model=model_name, max_model_len=1024)
152 - ref_outputs = ref_llm.chat(test_prompts, sampling_config)
153 - del ref_llm
154 - torch.cuda.empty_cache()
155 - cleanup_dist_env_and_memory()
156 -
167 spec_llm = LLM(
168 model=model_name,
169 speculative_config=speculative_config,
160 - max_model_len=1024,
170 + max_model_len=4096,
171 )
162 - spec_outputs = spec_llm.chat(test_prompts, sampling_config)
163 - matches = 0
164 - misses = 0
165 - for ref_output, spec_output in zip(ref_outputs, spec_outputs):
166 - if ref_output.outputs[0].text == spec_output.outputs[0].text:
167 - matches += 1
168 - else:
169 - misses += 1
170 - print(f"ref_output: {ref_output.outputs[0].text}")
171 - print(f"spec_output: {spec_output.outputs[0].text}")
172 -
173 - # Heuristic: expect at least 66% of the prompts to match exactly
174 - # Upon failure, inspect the outputs to check for inaccuracy.
175 - assert matches >= int(0.66 * len(ref_outputs))
172 + evaluate_llm_for_gsm8k(spec_llm)
173 del spec_llm
174 torch.cuda.empty_cache()
175 cleanup_dist_env_and_memory()
@@ -238,10 +235,10 @@ def test_suffix_decoding_acceptance(
235
236
237 @pytest.mark.parametrize(
241 - "model_path",
238 + ["model_path", "expected_accuracy_threshold"],
239 [
243 - "RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3",
244 - "RedHatAI/Qwen3-8B-speculator.eagle3",
240 + ("RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3", 0.7), # ref: 75%-80%
241 + ("RedHatAI/Qwen3-8B-speculator.eagle3", 0.8), # ref: 87%-92%
242 ],
243 ids=["llama3_eagle3_speculator", "qwen3_eagle3_speculator"],
244 )
@@ -249,6 +246,7 @@ def test_speculators_model_integration(
246 monkeypatch: pytest.MonkeyPatch,
247 sampling_config: SamplingParams,
248 model_path: str,
249 + expected_accuracy_threshold: float,
250 ):
251 """
252 Test that speculators models work with the simplified integration.
@@ -262,7 +260,8 @@ def test_speculators_model_integration(
260 2. Verifier model is extracted from speculator config
261 3. Speculative decoding is automatically enabled
262 4. Text generation works correctly
265 - 5. Output matches reference (non-speculative) generation
263 + 5. GSM8k accuracy of the model passes a sanity check when speculative decoding on
264 + 6. Output matches reference (non-speculative) generation
265 """
266 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
267
@@ -270,7 +269,10 @@ def test_speculators_model_integration(
269 test_prompts = get_test_prompts(mm_enabled=False)
270
271 # First run: Direct speculator model (simplified integration)
273 - spec_llm = LLM(model=model_path, max_model_len=1024)
272 + spec_llm = LLM(model=model_path, max_model_len=4096)
273 + evaluate_llm_for_gsm8k(
274 + spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
275 + )
276 spec_outputs = spec_llm.chat(test_prompts, sampling_config)
277
278 # Verify speculative config was auto-detected
@@ -297,7 +299,7 @@ def test_speculators_model_integration(
299 cleanup_dist_env_and_memory()
300
301 # Second run: Reference without speculative decoding
300 - ref_llm = LLM(model=verifier_model, max_model_len=1024)
302 + ref_llm = LLM(model=verifier_model, max_model_len=4096)
303 ref_outputs = ref_llm.chat(test_prompts, sampling_config)
304 del ref_llm
305 torch.cuda.empty_cache()
@@ -318,19 +320,27 @@ def test_speculators_model_integration(
320
321
322 @pytest.mark.parametrize(
321 - ["model_setup", "mm_enabled", "enable_chunked_prefill", "model_impl"],
323 + [
324 + "model_setup",
325 + "mm_enabled",
326 + "enable_chunked_prefill",
327 + "model_impl",
328 + "expected_accuracy_threshold",
329 + ],
330 [
331 (
332 ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
333 False,
334 False,
335 "auto",
336 + 0.8, # ref: 90%
337 ),
338 (
339 ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
340 False,
341 False,
342 "transformers",
343 + 0.8, # ref: 90%
344 ),
345 pytest.param(
346 (
@@ -342,6 +352,7 @@ def test_speculators_model_integration(
352 False,
353 False,
354 "auto",
355 + 0.8, # ref: 90%
356 marks=pytest.mark.skip(
357 reason="architecture of its eagle3 is LlamaForCausalLMEagle3"
358 ),
@@ -356,6 +367,7 @@ def test_speculators_model_integration(
367 False,
368 False,
369 "auto",
370 + 0.7, # TODO, update this with a reference value when re-enabling this case
371 marks=pytest.mark.skip(
372 reason="Skipping due to its head_dim not being a a multiple of 32"
373 ),
@@ -370,6 +382,7 @@ def test_speculators_model_integration(
382 False,
383 True,
384 "auto",
385 + 0.7, # ref: 75%-80%
386 marks=large_gpu_mark(min_gb=40),
387 ), # works on 4x H100
388 (
@@ -382,6 +395,7 @@ def test_speculators_model_integration(
395 False,
396 False,
397 "auto",
398 + 0.7, # ref: 75%-80%
399 ),
400 pytest.param(
401 (
@@ -393,7 +407,8 @@ def test_speculators_model_integration(
407 False,
408 False,
409 "auto",
396 - marks=large_gpu_mark(min_gb=80),
410 + 0.8, # ref: 90%
411 + # marks=large_gpu_mark(min_gb=80),
412 ), # works on 4x H100
413 pytest.param(
414 (
@@ -405,6 +420,7 @@ def test_speculators_model_integration(
420 True,
421 True,
422 "auto",
423 + 0.8, # ref: 90%
424 marks=large_gpu_mark(min_gb=80),
425 ), # works on 4x H100
426 (
@@ -417,6 +433,7 @@ def test_speculators_model_integration(
433 False,
434 False,
435 "auto",
436 + 0.0, # dummy model, skip gsm8k check
437 ),
438 ],
439 ids=[
@@ -437,10 +454,18 @@ def test_eagle_correctness(
454 sampling_config: SamplingParams,
455 model_setup: tuple[str, str, str, int],
456 mm_enabled: bool,
457 + expected_accuracy_threshold: float,
458 enable_chunked_prefill: bool,
459 model_impl: str,
460 attn_backend: str,
461 ):
462 + """
463 + Compare the outputs of a original LLM and a speculative LLM
464 + which should be the same when using eagle speculative decoding. Due to some variance
465 + in the engine, it is possible for some outputs to differ, so we expect that at least
466 + 6/10 output tokens match exactly, and that the GSM8k accuracy is above
467 + a precomputed reference threshold for each model.
468 + """
469 if attn_backend == "TREE_ATTN":
470 # TODO: Fix this flaky test
471 pytest.skip(
@@ -461,11 +486,6 @@ def test_eagle_correctness(
486
487 # Generate test prompts inside the function instead of using fixture
488 test_prompts = get_test_prompts(mm_enabled)
464 - """
465 - Compare the outputs of a original LLM and a speculative LLM
466 - should be the same when using eagle speculative decoding.
467 - model_setup: (method, model_name, eagle_model_name, tp_size)
468 - """
489 # Determine attention config
490 # Scout requires default backend selection because vision encoder has
491 # head_dim 88 being incompatible with FLASH_ATTN and needs to fall back
@@ -505,6 +525,9 @@ def test_eagle_correctness(
525 tensor_parallel_size=tp_size,
526 attention_config=attention_config,
527 )
528 + evaluate_llm_for_gsm8k(
529 + ref_llm, expected_accuracy_threshold=expected_accuracy_threshold
530 + )
531 ref_outputs = ref_llm.chat(test_prompts, sampling_config)
532 del ref_llm
533 torch.cuda.empty_cache()
@@ -526,6 +549,9 @@ def test_eagle_correctness(
549 model_impl=model_impl,
550 attention_config=attention_config,
551 )
552 + evaluate_llm_for_gsm8k(
553 + spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
554 + )
555 spec_outputs = spec_llm.chat(test_prompts, sampling_config)
556 matches = 0
557 misses = 0
@@ -546,10 +572,10 @@ def test_eagle_correctness(
572
573
574 @pytest.mark.parametrize(
549 - ["model_setup", "mm_enabled"],
575 + ["model_setup", "mm_enabled", "expected_accuracy_threshold"],
576 [
551 - (("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False),
552 - (("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1), False),
577 + (("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False, 0.5), # ref: 65%-70%
578 + (("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1), False, 0.0), # dummy model
579 ],
580 ids=["mimo", "deepseek"],
581 )
@@ -558,14 +584,17 @@ def test_mtp_correctness(
584 sampling_config: SamplingParams,
585 model_setup: tuple[str, str, int],
586 mm_enabled: bool,
587 + expected_accuracy_threshold: float,
588 ):
562 - # Generate test prompts inside the function instead of using fixture
563 - test_prompts = get_test_prompts(mm_enabled)
589 """
590 Compare the outputs of a original LLM and a speculative LLM
566 - should be the same when using MTP speculative decoding.
567 - model_setup: (method, model_name, tp_size)
591 + which should be the same when using MTP speculative decoding. Due to some variance
592 + in the engine, it is possible for some outputs to differ, so we expect that at least
593 + 6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed
594 + reference threshold for each model.
595 """
596 + # Generate test prompts inside the function instead of using fixture
597 + test_prompts = get_test_prompts(mm_enabled)
598 with monkeypatch.context() as m:
599 m.setenv("VLLM_MLA_DISABLE", "1")
600
@@ -579,6 +608,9 @@ def test_mtp_correctness(
608 trust_remote_code=True,
609 )
610 ref_outputs = ref_llm.chat(test_prompts, sampling_config)
611 + evaluate_llm_for_gsm8k(
612 + ref_llm, expected_accuracy_threshold=expected_accuracy_threshold
613 + )
614 del ref_llm
615 torch.cuda.empty_cache()
616 cleanup_dist_env_and_memory()
@@ -594,6 +626,9 @@ def test_mtp_correctness(
626 },
627 max_model_len=2048,
628 )
629 + evaluate_llm_for_gsm8k(
630 + spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
631 + )
632 spec_outputs = spec_llm.chat(test_prompts, sampling_config)
633 matches = 0
634 misses = 0
@@ -621,12 +656,13 @@ class ArgsTest:
656 num_speculative_tokens: int
657 expected_acceptance_rate: float
658 expected_acceptance_len: float
659 + expected_gsm8k_accuracy: float = 0.0 # skip by default
660 # Defaults
661 enforce_eager: bool = True
662 parallel_drafting: bool = False
663 target_tensor_parallel_size: int = 1
664 draft_tensor_parallel_size: int = 1
629 - max_model_len: int = 1024
665 + max_model_len: int = 2048
666 gpu_memory_utilization: float = 0.5
667 dataset: str = "test_prompts"
668 num_prompts: int = 100
@@ -639,8 +675,9 @@ cases = [
675 draft_model="Qwen/Qwen3-0.6B",
676 sampling_config=greedy_sampling(),
677 num_speculative_tokens=3, # K
642 - expected_acceptance_len=3 + 1, # K + 1
643 - expected_acceptance_rate=1.0,
678 + expected_acceptance_len=0.98 * (3 + 1), # epsilon discount of K + 1
679 + expected_acceptance_rate=0.98, # slight epsilon
680 + expected_gsm8k_accuracy=0.25, # ref: 35-40%
681 ),
682 # Smaller draft model, stochastic sampling.
683 ArgsTest(
@@ -648,8 +685,9 @@ cases = [
685 draft_model="Qwen/Qwen3-0.6B",
686 sampling_config=stochastic_sampling(),
687 num_speculative_tokens=3,
651 - expected_acceptance_len=2.8 + 1,
652 - expected_acceptance_rate=0.9,
688 + expected_acceptance_len=3.4, # ref: 3.7
689 + expected_acceptance_rate=0.80, # ref: 0.90
690 + expected_gsm8k_accuracy=0.5, # ref: 60%. Note gsm8k always runs greedy sampling
691 ),
692 ]
693
@@ -669,9 +707,8 @@ def test_draft_model_realistic_example():
707 num_speculative_tokens=3,
708 sampling_config=greedy_sampling(),
709 enforce_eager=False,
672 - # values below are not derived, but just prevent a regression
673 - expected_acceptance_len=2.8,
674 - expected_acceptance_rate=0.55,
710 + expected_acceptance_len=2.6, # ref: 2.86
711 + expected_acceptance_rate=0.5, # ref: 0.62
712 )
713 assert_draft_model_correctness(args)
714
@@ -685,9 +722,8 @@ def test_draft_model_parallel_drafting():
722 sampling_config=greedy_sampling(),
723 parallel_drafting=True,
724 enforce_eager=False,
688 - # values below are collected from a stable run, with ~5% tolerance
689 - expected_acceptance_len=2.375,
690 - expected_acceptance_rate=0.45,
725 + expected_acceptance_len=2.3, # ref: 2.52
726 + expected_acceptance_rate=0.4, # ref: 0.51
727 )
728 assert_draft_model_correctness(args)
729
@@ -723,6 +759,7 @@ def test_draft_model_tensor_parallelism():
759 draft_tensor_parallel_size=2,
760 **some_high_acceptance_metrics(),
761 enforce_eager=False,
762 + expected_gsm8k_accuracy=0.5,
763 )
764 assert_draft_model_correctness(sd_case)
765
@@ -797,9 +834,14 @@ def assert_draft_model_correctness(args: ArgsTest):
834 # we don't check the outputs, only check the metrics
835 spec_llm.chat(test_prompts, args.sampling_config)
836 metrics = spec_llm.get_metrics()
800 -
837 acceptance_rate: float = compute_acceptance_rate(metrics)
838 acceptance_len: float = compute_acceptance_len(metrics)
839 +
840 + # Need to evaluate after getting metrics to avoid polluting the AR
841 + evaluate_llm_for_gsm8k(
842 + spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy
843 + )
844 +
845 del spec_llm # CLEANUP
846 torch.cuda.empty_cache()
847 cleanup_dist_env_and_memory()
@@ -817,7 +859,7 @@ def assert_draft_model_correctness(args: ArgsTest):
859
860 def get_messages(dataset: str, n: int) -> list[Messages]:
861 if dataset == "test_prompts":
820 - return get_test_prompts(mm_enabled=False, quiet=True, num_prompts=n)
862 + return get_test_prompts(mm_enabled=False, num_prompts=n)
863 elif dataset == "likaixin/InstructCoder":
864 return get_instruct_coder_messages(n=n)
865 else:
@@ -828,8 +870,8 @@ def some_high_acceptance_metrics() -> dict:
870 return {
871 "sampling_config": greedy_sampling(),
872 "num_speculative_tokens": 3,
831 - "expected_acceptance_len": 2.8 + 1,
832 - "expected_acceptance_rate": 0.90,
873 + "expected_acceptance_len": 3.4, # ref: 3.75
874 + "expected_acceptance_rate": 0.8, # ref: 0.9
875 }
876
877