| 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.utils import multi_gpu_only, single_gpu_only |
| 9 | from vllm import SamplingParams |
| 10 | from vllm.config import CompilationConfig, VllmConfig, replace |
| 11 | from vllm.config.kernel import MoEBackend |
| 12 | from vllm.engine.arg_utils import EngineArgs |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.v1.spec_decode.draft_model import DraftModelProposer |
| 15 | |
| 16 | from ..utils import ( |
| 17 | Messages, |
| 18 | compute_acceptance_len, |
| 19 | compute_acceptance_rate, |
| 20 | evaluate_llm_for_gsm8k, |
| 21 | get_instruct_coder_messages, |
| 22 | get_test_prompts, |
| 23 | greedy_sampling, |
| 24 | stochastic_sampling, |
| 25 | ) |
| 26 | |
| 27 | |
| 28 | @dataclass |
| 29 | class ArgsTest: |
| 30 | target_model: str |
| 31 | draft_model: str |
| 32 | sampling_config: SamplingParams |
| 33 | num_speculative_tokens: int |
| 34 | expected_acceptance_rate: float |
| 35 | expected_acceptance_len: float |
| 36 | expected_gsm8k_accuracy: float = 0.0 # skip by default |
| 37 | # Defaults |
| 38 | enforce_eager: bool = True |
| 39 | parallel_drafting: bool = False |
| 40 | target_tensor_parallel_size: int = 1 |
| 41 | draft_tensor_parallel_size: int = 1 |
| 42 | max_model_len: int = 2048 |
| 43 | gpu_memory_utilization: float = 0.5 |
| 44 | dataset: str = "test_prompts" |
| 45 | # Doubling 100 to 200 reduces sampling standard error by about 29%. |
| 46 | num_prompts: int = 200 |
| 47 | |
| 48 | |
| 49 | def get_messages(dataset: str, n: int) -> list[Messages]: |
| 50 | if dataset == "test_prompts": |
| 51 | return get_test_prompts(mm_enabled=False, num_prompts=n) |
| 52 | if dataset == "likaixin/InstructCoder": |
| 53 | return get_instruct_coder_messages(n=n) |
| 54 | raise NotImplementedError(f"Dataset '{dataset}' not implemented") |
| 55 | |
| 56 | |
| 57 | def some_high_acceptance_metrics() -> dict: |
| 58 | return { |
| 59 | "sampling_config": greedy_sampling(), |
| 60 | "num_speculative_tokens": 3, |
| 61 | "expected_acceptance_len": 3.4, # ref: 3.75 |
| 62 | "expected_acceptance_rate": 0.8, # ref: 0.9 |
| 63 | } |
| 64 | |
| 65 | |
| 66 | cases = [ |
| 67 | # Same model for draft and target, greedy sampling. |
| 68 | ArgsTest( |
| 69 | target_model="Qwen/Qwen3-0.6B", |
| 70 | draft_model="Qwen/Qwen3-0.6B", |
| 71 | sampling_config=greedy_sampling(), |
| 72 | num_speculative_tokens=3, # K |
| 73 | expected_acceptance_len=0.98 * (3 + 1), # epsilon discount of K + 1 |
| 74 | expected_acceptance_rate=0.98, # slight epsilon |
| 75 | expected_gsm8k_accuracy=0.25, # ref: 35-40% |
| 76 | ), |
| 77 | # Smaller draft model, stochastic sampling. |
| 78 | ArgsTest( |
| 79 | target_model="Qwen/Qwen3-1.7B", |
| 80 | draft_model="Qwen/Qwen3-0.6B", |
| 81 | sampling_config=stochastic_sampling(), |
| 82 | num_speculative_tokens=3, |
| 83 | expected_acceptance_len=3.4, # ref: 3.7 |
| 84 | expected_acceptance_rate=0.80, # ref: 0.90 |
| 85 | expected_gsm8k_accuracy=0.5, # ref: 60%. Note gsm8k always runs greedy sampling |
| 86 | ), |
| 87 | ] |
| 88 | |
| 89 | |
| 90 | @pytest.mark.parametrize("args", cases) |
| 91 | @pytest.mark.parametrize("enforce_eager", [True, False]) |
| 92 | @single_gpu_only |
| 93 | def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool, vllm_runner): |
| 94 | args.enforce_eager = enforce_eager |
| 95 | assert_draft_model_correctness(args, vllm_runner) |
| 96 | |
| 97 | |
| 98 | @single_gpu_only |
| 99 | def test_draft_model_realistic_example(vllm_runner): |
| 100 | args = ArgsTest( |
| 101 | target_model="Qwen/Qwen3-1.7B", |
| 102 | draft_model="Qwen/Qwen3-0.6B", |
| 103 | dataset="likaixin/InstructCoder", |
| 104 | num_speculative_tokens=3, |
| 105 | sampling_config=greedy_sampling(), |
| 106 | enforce_eager=False, |
| 107 | expected_acceptance_len=2.6, # ref: 2.86 |
| 108 | expected_acceptance_rate=0.5, # ref: 0.62 |
| 109 | ) |
| 110 | assert_draft_model_correctness(args, vllm_runner) |
| 111 | |
| 112 | |
| 113 | @single_gpu_only |
| 114 | def test_draft_model_parallel_drafting(vllm_runner): |
| 115 | args = ArgsTest( |
| 116 | target_model="Qwen/Qwen3-1.7B", |
| 117 | draft_model="amd/PARD-Qwen3-0.6B", |
| 118 | dataset="likaixin/InstructCoder", |
| 119 | num_speculative_tokens=3, |
| 120 | sampling_config=greedy_sampling(), |
| 121 | parallel_drafting=True, |
| 122 | enforce_eager=False, |
| 123 | expected_acceptance_len=2.3, # ref: 2.52 |
| 124 | expected_acceptance_rate=0.4, # ref: 0.51 |
| 125 | ) |
| 126 | assert_draft_model_correctness(args, vllm_runner) |
| 127 | |
| 128 | |
| 129 | @pytest.mark.parametrize( |
| 130 | "models", |
| 131 | [ |
| 132 | # target_model, draft_model |
| 133 | ("Qwen/Qwen3-1.7B-FP8", "Qwen/Qwen3-0.6B"), # target quantized |
| 134 | ("Qwen/Qwen3-1.7B", "Qwen/Qwen3-0.6B-FP8"), # draft quantized |
| 135 | ], |
| 136 | ids=["target_quantized", "draft_quantized"], |
| 137 | ) |
| 138 | @pytest.mark.parametrize("enforce_eager", [True, False]) |
| 139 | @single_gpu_only |
| 140 | def test_draft_model_quantization( |
| 141 | models: tuple[str, str], enforce_eager: bool, vllm_runner |
| 142 | ): |
| 143 | tgt_model, draft_model = models |
| 144 | sd_case = ArgsTest( |
| 145 | target_model=tgt_model, |
| 146 | draft_model=draft_model, |
| 147 | **some_high_acceptance_metrics(), |
| 148 | enforce_eager=enforce_eager, |
| 149 | ) |
| 150 | assert_draft_model_correctness(sd_case, vllm_runner) |
| 151 | |
| 152 | |
| 153 | @single_gpu_only |
| 154 | def test_draft_model_transformers_backend(vllm_runner): |
| 155 | """Ensure that Attention prefixes don't collide between the target and |
| 156 | draft models when both models use the Transformers modelling backend.""" |
| 157 | sd_case = ArgsTest( |
| 158 | target_model="allenai/OLMo-2-0425-1B-Instruct", |
| 159 | draft_model="allenai/OLMo-2-0425-1B-Instruct", |
| 160 | sampling_config=greedy_sampling(), |
| 161 | num_speculative_tokens=3, |
| 162 | expected_acceptance_len=0.98 * (3 + 1), # epsilon discount of K + 1 |
| 163 | expected_acceptance_rate=0.98, # slight epsilon |
| 164 | ) |
| 165 | assert_draft_model_correctness(sd_case, vllm_runner) |
| 166 | |
| 167 | |
| 168 | @multi_gpu_only(num_gpus=2) |
| 169 | def test_draft_model_tensor_parallelism(vllm_runner): |
| 170 | """Ensure spec decode works when running with TP > 1.""" |
| 171 | sd_case = ArgsTest( |
| 172 | target_model="Qwen/Qwen3-1.7B", |
| 173 | target_tensor_parallel_size=2, |
| 174 | draft_model="Qwen/Qwen3-0.6B", |
| 175 | draft_tensor_parallel_size=2, |
| 176 | **some_high_acceptance_metrics(), |
| 177 | enforce_eager=False, |
| 178 | expected_gsm8k_accuracy=0.5, |
| 179 | ) |
| 180 | assert_draft_model_correctness(sd_case, vllm_runner) |
| 181 | |
| 182 | |
| 183 | @multi_gpu_only(num_gpus=2) |
| 184 | def test_draft_model_engine_args_tensor_parallelism(): |
| 185 | """Ensure the vllm_config for the draft model is created correctly, |
| 186 | and independently of the target model (quantization, TP, etc.)""" |
| 187 | engine_args = EngineArgs( |
| 188 | model="Qwen/Qwen3-1.7B-FP8", # <<< tgt quantized |
| 189 | tensor_parallel_size=2, |
| 190 | speculative_config={ |
| 191 | "model": "Qwen/Qwen3-0.6B", # <<< draft not quantized |
| 192 | "method": "draft_model", |
| 193 | "num_speculative_tokens": 3, |
| 194 | "draft_tensor_parallel_size": 1, # <<< valid arg name |
| 195 | }, |
| 196 | ) |
| 197 | target_config: VllmConfig = engine_args.create_engine_config() |
| 198 | assert target_config.parallel_config.tensor_parallel_size == 2 |
| 199 | assert target_config.quant_config.get_name() == "fp8" |
| 200 | |
| 201 | speculative_config = target_config.speculative_config |
| 202 | draft_config: VllmConfig = replace( |
| 203 | target_config, |
| 204 | quant_config=None, |
| 205 | parallel_config=replace( |
| 206 | speculative_config.draft_parallel_config, |
| 207 | rank=target_config.parallel_config.rank, |
| 208 | ), |
| 209 | model_config=speculative_config.draft_model_config, |
| 210 | ) |
| 211 | assert draft_config.parallel_config.tensor_parallel_size == 1 |
| 212 | assert draft_config.quant_config is None |
| 213 | |
| 214 | |
| 215 | @pytest.mark.cpu_test |
| 216 | @pytest.mark.parametrize("dcp_size", [1, 2]) |
| 217 | def test_draft_model_inherits_target_dcp_config(monkeypatch, dcp_size): |
| 218 | monkeypatch.setattr(current_platform, "device_count", lambda: 16) |
| 219 | target_config = EngineArgs( |
| 220 | model="Qwen/Qwen3-1.7B", |
| 221 | tensor_parallel_size=16, |
| 222 | decode_context_parallel_size=dcp_size, |
| 223 | cp_kv_cache_interleave_size=16, |
| 224 | dcp_comm_backend="ag_rs", |
| 225 | distributed_executor_backend="mp", |
| 226 | speculative_config={ |
| 227 | "model": "Qwen/Qwen3-0.6B", |
| 228 | "method": "draft_model", |
| 229 | "num_speculative_tokens": 3, |
| 230 | }, |
| 231 | ).create_engine_config() |
| 232 | proposer = object.__new__(DraftModelProposer) |
| 233 | proposer.vllm_config = target_config |
| 234 | proposer.speculative_config = target_config.speculative_config |
| 235 | draft_config = proposer._create_draft_vllm_config() |
| 236 | draft_parallel = draft_config.parallel_config |
| 237 | assert draft_parallel.decode_context_parallel_size == dcp_size |
| 238 | assert draft_parallel.cp_kv_cache_interleave_size == 16 |
| 239 | assert draft_parallel.dcp_comm_backend == "ag_rs" |
| 240 | assert ( |
| 241 | draft_parallel.dcp_q_replicate == target_config.parallel_config.dcp_q_replicate |
| 242 | ) |
| 243 | assert draft_parallel is not target_config.parallel_config |
| 244 | assert target_config.parallel_config.decode_context_parallel_size == dcp_size |
| 245 | |
| 246 | |
| 247 | def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig: |
| 248 | """Replicate SpecDecodeBaseProposer._create_draft_vllm_config logic |
| 249 | so we can test it without instantiating a full proposer.""" |
| 250 | spec_cfg = vllm_config.speculative_config |
| 251 | if spec_cfg.moe_backend is not None: |
| 252 | return replace( |
| 253 | vllm_config, |
| 254 | kernel_config=replace( |
| 255 | vllm_config.kernel_config, |
| 256 | moe_backend=spec_cfg.moe_backend, |
| 257 | ), |
| 258 | ) |
| 259 | return vllm_config |
| 260 | |
| 261 | |
| 262 | def _platform_moe_backend(monkeypatch: pytest.MonkeyPatch) -> MoEBackend: |
| 263 | if current_platform.is_rocm(): |
| 264 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 265 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") |
| 266 | return "aiter" |
| 267 | if current_platform.is_cuda(): |
| 268 | return "flashinfer_trtllm" |
| 269 | return "auto" |
| 270 | |
| 271 | |
| 272 | def test_draft_model_moe_backend_override(monkeypatch: pytest.MonkeyPatch): |
| 273 | """When moe_backend is set in speculative_config, the draft VllmConfig |
| 274 | should use it while the target keeps its own setting.""" |
| 275 | target_moe_backend = _platform_moe_backend(monkeypatch) |
| 276 | engine_args = EngineArgs( |
| 277 | model="Qwen/Qwen3-1.7B", |
| 278 | tensor_parallel_size=1, |
| 279 | moe_backend=target_moe_backend, |
| 280 | speculative_config={ |
| 281 | "model": "Qwen/Qwen3-0.6B", |
| 282 | "method": "draft_model", |
| 283 | "num_speculative_tokens": 3, |
| 284 | "moe_backend": "triton", |
| 285 | }, |
| 286 | ) |
| 287 | tgt_config: VllmConfig = engine_args.create_engine_config() |
| 288 | assert tgt_config.kernel_config.moe_backend == target_moe_backend |
| 289 | assert tgt_config.speculative_config.moe_backend == "triton" |
| 290 | |
| 291 | draft_config = _apply_draft_moe_backend(tgt_config) |
| 292 | assert draft_config.kernel_config.moe_backend == "triton" |
| 293 | # Target config must be unaffected. |
| 294 | assert tgt_config.kernel_config.moe_backend == target_moe_backend |
| 295 | |
| 296 | |
| 297 | def test_draft_model_moe_backend_inherits_target(monkeypatch: pytest.MonkeyPatch): |
| 298 | """When moe_backend is not set in speculative_config, the draft should |
| 299 | inherit the target's moe_backend.""" |
| 300 | target_moe_backend = _platform_moe_backend(monkeypatch) |
| 301 | engine_args = EngineArgs( |
| 302 | model="Qwen/Qwen3-1.7B", |
| 303 | tensor_parallel_size=1, |
| 304 | moe_backend=target_moe_backend, |
| 305 | speculative_config={ |
| 306 | "model": "Qwen/Qwen3-0.6B", |
| 307 | "method": "draft_model", |
| 308 | "num_speculative_tokens": 3, |
| 309 | }, |
| 310 | ) |
| 311 | tgt_config: VllmConfig = engine_args.create_engine_config() |
| 312 | assert tgt_config.kernel_config.moe_backend == target_moe_backend |
| 313 | assert tgt_config.speculative_config.moe_backend is None |
| 314 | |
| 315 | draft_config = _apply_draft_moe_backend(tgt_config) |
| 316 | assert draft_config.kernel_config.moe_backend == target_moe_backend |
| 317 | assert draft_config is tgt_config |
| 318 | |
| 319 | |
| 320 | def test_draft_model_moe_backend_default_auto(): |
| 321 | """When neither target nor draft set moe_backend explicitly, both should |
| 322 | default to 'auto'.""" |
| 323 | engine_args = EngineArgs( |
| 324 | model="Qwen/Qwen3-1.7B", |
| 325 | tensor_parallel_size=1, |
| 326 | speculative_config={ |
| 327 | "model": "Qwen/Qwen3-0.6B", |
| 328 | "method": "draft_model", |
| 329 | "num_speculative_tokens": 3, |
| 330 | }, |
| 331 | ) |
| 332 | tgt_config: VllmConfig = engine_args.create_engine_config() |
| 333 | assert tgt_config.kernel_config.moe_backend == "auto" |
| 334 | assert tgt_config.speculative_config.moe_backend is None |
| 335 | |
| 336 | draft_config = _apply_draft_moe_backend(tgt_config) |
| 337 | assert draft_config.kernel_config.moe_backend == "auto" |
| 338 | assert draft_config is tgt_config |
| 339 | |
| 340 | |
| 341 | def test_draft_model_engine_args_rejects_invalid_tp_argname(): |
| 342 | """The user should pass "draft_tensor_parallel_size" rather than |
| 343 | "tensor_parallel_size". We enforce this with validation.""" |
| 344 | engine_args = EngineArgs( |
| 345 | model="Qwen/Qwen3-1.7B", |
| 346 | tensor_parallel_size=1, |
| 347 | speculative_config={ |
| 348 | "model": "Qwen/Qwen3-0.6B", |
| 349 | "method": "draft_model", |
| 350 | "num_speculative_tokens": 3, |
| 351 | "tensor_parallel_size": 1, # <<< invalid arg name |
| 352 | }, |
| 353 | ) |
| 354 | with pytest.raises(ValueError, match="draft_tensor_parallel_size"): |
| 355 | engine_args.create_engine_config() |
| 356 | |
| 357 | |
| 358 | def assert_draft_model_correctness(args: ArgsTest, vllm_runner): |
| 359 | """Compare the outputs using and not using speculative decoding. |
| 360 | In the greedy decoding case, the outputs must match EXACTLY.""" |
| 361 | test_prompts: list[Messages] = get_messages( |
| 362 | dataset=args.dataset, n=args.num_prompts |
| 363 | ) |
| 364 | |
| 365 | with vllm_runner( |
| 366 | args.target_model, |
| 367 | block_size=None, |
| 368 | trust_remote_code=False, |
| 369 | enable_chunked_prefill=None, |
| 370 | compilation_config=CompilationConfig(), |
| 371 | speculative_config={ |
| 372 | "model": args.draft_model, |
| 373 | "method": "draft_model", |
| 374 | "num_speculative_tokens": args.num_speculative_tokens, |
| 375 | "max_model_len": args.max_model_len, |
| 376 | "enforce_eager": args.enforce_eager, |
| 377 | "draft_tensor_parallel_size": args.draft_tensor_parallel_size, |
| 378 | "parallel_drafting": args.parallel_drafting, |
| 379 | }, |
| 380 | max_num_seqs=100, # limit cudagraph capture runtime |
| 381 | max_model_len=args.max_model_len, |
| 382 | gpu_memory_utilization=args.gpu_memory_utilization, |
| 383 | tensor_parallel_size=args.target_tensor_parallel_size, |
| 384 | enforce_eager=args.enforce_eager, |
| 385 | disable_log_stats=False, # enables get_metrics() |
| 386 | ) as spec_runner: |
| 387 | spec_llm = spec_runner.llm |
| 388 | |
| 389 | # we don't check the outputs, only check the metrics |
| 390 | spec_llm.chat(test_prompts, args.sampling_config) |
| 391 | metrics = spec_llm.get_metrics() |
| 392 | acceptance_rate: float = compute_acceptance_rate(metrics) |
| 393 | acceptance_len: float = compute_acceptance_len(metrics) |
| 394 | |
| 395 | # Need to evaluate after getting metrics to avoid polluting the AR |
| 396 | evaluate_llm_for_gsm8k( |
| 397 | spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy |
| 398 | ) |
| 399 | |
| 400 | print( |
| 401 | f"spec-decode: target={args.target_model}, draft={args.draft_model}, " |
| 402 | f"temperature={args.sampling_config.temperature:.2f}, " |
| 403 | f"acceptance_rate={acceptance_rate:.2f}, " |
| 404 | f"acceptance_len={acceptance_len:.2f}, " |
| 405 | ) |
| 406 | |
| 407 | context = ( |
| 408 | f"draft_model target={args.target_model}, draft={args.draft_model}, " |
| 409 | f"eager={args.enforce_eager}, " |
| 410 | f"acceptance_rate={acceptance_rate:.3f}, " |
| 411 | f"acceptance_len={acceptance_len:.3f}" |
| 412 | ) |
| 413 | assert acceptance_rate >= args.expected_acceptance_rate, ( |
| 414 | f"{context}; expected acceptance_rate >= " |
| 415 | f"{args.expected_acceptance_rate:.3f}" |
| 416 | ) |
| 417 | assert acceptance_len >= args.expected_acceptance_len, ( |
| 418 | f"{context}; expected acceptance_len >= {args.expected_acceptance_len:.3f}" |
| 419 | ) |
| 420 | # draft_model supports async scheduling; assert it is active by default. |
| 421 | has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling |
| 422 | |
| 423 | assert has_async, "Expected async_scheduling=True for draft_model spec decode" |