@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 182 lines 5.58 KB
Raw
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
8 from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline
9 from tests.utils import single_gpu_only
10 from vllm.config import CompilationConfig
11
12 from ...utils import compute_acceptance_len
13 from ..utils import run_acceptance_length_eval
14
15
16 @dataclass(frozen=True)
17 class DFlashCorrectnessConfig:
18 model: str
19 draft_model: str
20 expected_accuracy: float
21 expected_acceptance_len: float
22 num_speculative_tokens: int = 16
23 max_model_len: int = 4096
24 max_num_seqs: int = 128
25 num_questions: int = 1319
26 use_chat_completions: bool = False
27 enforce_eager: bool = False
28 disable_flashinfer_sampler: bool = False
29 language_model_only: bool = False
30 chat_template_kwargs: dict[str, object] | None = None
31
32
33 QWEN3_DFLASH = DFlashCorrectnessConfig(
34 model="Qwen/Qwen3-8B",
35 draft_model="z-lab/Qwen3-8B-DFlash-b16",
36 expected_accuracy=0.8,
37 expected_acceptance_len=3.5,
38 )
39
40 QWEN3_8_DFLASH2_NVFP4 = DFlashCorrectnessConfig(
41 model="nvidia/Qwen3.8-27B-NVFP4",
42 draft_model="z-lab/Qwen3.8-27B-DFlash2",
43 expected_accuracy=0.90,
44 expected_acceptance_len=6.14 * 0.95,
45 num_speculative_tokens=7,
46 max_num_seqs=32,
47 use_chat_completions=True,
48 language_model_only=True,
49 chat_template_kwargs={"enable_thinking": False},
50 )
51
52 LAGUNA_DFLASH_NVFP4 = DFlashCorrectnessConfig(
53 model="poolside/Laguna-XS-2.1-NVFP4",
54 draft_model="poolside/Laguna-XS-2.1-DFlash-NVFP4",
55 expected_accuracy=0.7, # Standard GSM8K sanity floor.
56 expected_acceptance_len=3.05 * 0.9,
57 num_speculative_tokens=15,
58 max_model_len=8192,
59 max_num_seqs=32,
60 num_questions=200,
61 use_chat_completions=True,
62 enforce_eager=True,
63 disable_flashinfer_sampler=True,
64 )
65
66
67 @pytest.mark.parametrize("use_mrv2", [False, True])
68 def test_dflash_reference_acceptance_lengths(
69 monkeypatch: pytest.MonkeyPatch,
70 use_mrv2: bool,
71 vllm_runner,
72 ):
73 run_acceptance_length_eval(
74 monkeypatch,
75 vllm_runner,
76 spec_config={
77 "model": QWEN3_DFLASH.model,
78 "trust_remote_code": True,
79 "speculative_config": {
80 "method": "dflash",
81 "model": QWEN3_DFLASH.draft_model,
82 "num_speculative_tokens": QWEN3_DFLASH.num_speculative_tokens,
83 "max_model_len": 32768,
84 },
85 "max_model_len": 32768,
86 "max_num_seqs": 128,
87 "gpu_memory_utilization": 0.85,
88 "enforce_eager": False,
89 "disable_log_stats": False,
90 },
91 # Table 1 in https://arxiv.org/pdf/2602.06036.
92 expected_acceptance_lengths={
93 "mt-bench": 4.24,
94 "humaneval": 6.50,
95 "gsm8k": 6.54 * 0.975,
96 },
97 chat_template_kwargs={"enable_thinking": False},
98 use_mrv2=use_mrv2,
99 )
100
101
102 @single_gpu_only
103 @pytest.mark.parametrize(
104 ("config", "use_mrv2"),
105 [
106 pytest.param(
107 QWEN3_DFLASH,
108 False,
109 id="qwen3-mrv1",
110 ),
111 pytest.param(
112 QWEN3_DFLASH,
113 True,
114 id="qwen3-mrv2",
115 ),
116 pytest.param(
117 LAGUNA_DFLASH_NVFP4,
118 True,
119 id="laguna-nvfp4-mrv2",
120 ),
121 pytest.param(
122 QWEN3_8_DFLASH2_NVFP4,
123 True,
124 id="qwen3.8-dflash2-nvfp4-mrv2",
125 ),
126 ],
127 )
128 def test_dflash_correctness(
129 monkeypatch: pytest.MonkeyPatch,
130 config: DFlashCorrectnessConfig,
131 use_mrv2: bool,
132 vllm_runner,
133 ):
134 """Guard GSM8K accuracy and batched acceptance length for DFlash models."""
135 monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
136 if config.disable_flashinfer_sampler:
137 monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
138
139 with vllm_runner(
140 config.model,
141 block_size=None,
142 trust_remote_code=True,
143 speculative_config={
144 "method": "dflash",
145 "model": config.draft_model,
146 "num_speculative_tokens": config.num_speculative_tokens,
147 "max_model_len": config.max_model_len,
148 },
149 max_model_len=config.max_model_len,
150 max_num_seqs=config.max_num_seqs,
151 gpu_memory_utilization=0.85,
152 enforce_eager=config.enforce_eager,
153 disable_log_stats=False,
154 enable_chunked_prefill=None,
155 compilation_config=CompilationConfig(),
156 language_model_only=config.language_model_only,
157 ) as spec_runner:
158 spec_llm = spec_runner.llm
159 results = evaluate_gsm8k_offline(
160 spec_llm,
161 num_questions=config.num_questions,
162 use_chat_completions=config.use_chat_completions,
163 chat_template_kwargs=config.chat_template_kwargs,
164 )
165 accuracy = results["accuracy"]
166 acceptance_len = compute_acceptance_len(spec_llm.get_metrics())
167 context = (
168 f"DFlash target={config.model}, draft={config.draft_model}, MRV2={use_mrv2}"
169 )
170 print(
171 f"{context}: GSM8K accuracy={accuracy:.3f}, "
172 f"acceptance_len={acceptance_len:.2f}"
173 )
174
175 assert accuracy >= config.expected_accuracy, (
176 f"{context}: GSM8K accuracy {accuracy:.3f} is below "
177 f"{config.expected_accuracy:.3f}; acceptance_len={acceptance_len:.3f}"
178 )
179 assert acceptance_len >= config.expected_acceptance_len, (
180 f"{context}: acceptance_len {acceptance_len:.3f} is below "
181 f"{config.expected_acceptance_len:.3f}; accuracy={accuracy:.3f}"
182 )