| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import os |
| 4 | from itertools import repeat |
| 5 | from typing import Any |
| 6 | |
| 7 | import pytest |
| 8 | import torch._dynamo.config as dynamo_config |
| 9 | |
| 10 | from tests.utils import ( |
| 11 | large_gpu_mark, |
| 12 | single_gpu_only, |
| 13 | ) |
| 14 | from vllm import SamplingParams |
| 15 | from vllm.logprobs import Logprob |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.sampling_params import StructuredOutputsParams |
| 18 | from vllm.v1.metrics.reader import Metric |
| 19 | |
| 20 | from ....conftest import VllmRunner |
| 21 | from ....models.utils import check_outputs_equal |
| 22 | |
| 23 | MODEL = "Qwen/Qwen3-0.6B" |
| 24 | MTP_MODEL = "meta-llama/Llama-3.2-1B-Instruct" |
| 25 | |
| 26 | # Need to enforce eager for MRV2 while we sort out cudagraph issues. |
| 27 | ENFORCE_EAGER = os.getenv("ENFORCE_EAGER", "0") == "1" |
| 28 | |
| 29 | first_prompt = ( |
| 30 | "The following numbers of the sequence " |
| 31 | + ", ".join(str(i) for i in range(10)) |
| 32 | + " are:" |
| 33 | ) |
| 34 | example_prompts = [first_prompt, "In one word, the capital of France is "] + [ |
| 35 | f"Tell me about the number {i}: " for i in range(32) |
| 36 | ] |
| 37 | |
| 38 | default_params = dict( |
| 39 | temperature=0.0, # greedy |
| 40 | max_tokens=30, |
| 41 | min_tokens=28, |
| 42 | ) |
| 43 | |
| 44 | |
| 45 | @single_gpu_only |
| 46 | def test_without_spec_decoding( |
| 47 | sample_json_schema, |
| 48 | monkeypatch: pytest.MonkeyPatch, |
| 49 | ): |
| 50 | """Test consistency of combos of async scheduling, preemption, |
| 51 | uni/multiproc executor, prefill chunking.""" |
| 52 | struct_outputs = StructuredOutputsParams(json=sample_json_schema) |
| 53 | test_sampling_params: list[dict[str, Any]] = [ |
| 54 | dict(), |
| 55 | # dict(min_tokens=20), |
| 56 | dict(frequency_penalty=-1.0), |
| 57 | dict(bad_words=["the", " the"]), |
| 58 | dict(logprobs=2), |
| 59 | dict(logprobs=2, frequency_penalty=-1.0), |
| 60 | dict(prompt_logprobs=2), |
| 61 | dict(prompt_logprobs=2, logprobs=2), |
| 62 | dict(structured_outputs=struct_outputs), |
| 63 | dict( |
| 64 | structured_outputs=struct_outputs, |
| 65 | logprobs=2, |
| 66 | ), |
| 67 | dict( |
| 68 | structured_outputs=struct_outputs, |
| 69 | frequency_penalty=-1.0, |
| 70 | ), |
| 71 | dict( |
| 72 | structured_outputs=struct_outputs, |
| 73 | logprobs=2, |
| 74 | frequency_penalty=-1.0, |
| 75 | ), |
| 76 | ] |
| 77 | |
| 78 | # test_preemption, executor, async_scheduling, |
| 79 | # spec_config, test_prefill_chunking |
| 80 | test_configs = [ |
| 81 | (False, "mp", False, None, False), |
| 82 | (True, "mp", False, None, True), |
| 83 | (False, "mp", True, None, False), |
| 84 | (False, "uni", True, None, False), |
| 85 | (True, "mp", True, None, False), |
| 86 | (True, "uni", True, None, False), |
| 87 | (False, "mp", True, None, True), |
| 88 | (True, "mp", True, None, True), |
| 89 | (True, "uni", True, None, True), |
| 90 | ] |
| 91 | |
| 92 | if current_platform.is_rocm(): |
| 93 | # On ROCm, Only test with structured_outputs (deterministic) |
| 94 | # and skip chunk_prefill (more variable). |
| 95 | test_configs = [ |
| 96 | cfg |
| 97 | for cfg in test_configs |
| 98 | if not cfg[4] # skip chunk_prefill=True |
| 99 | ] |
| 100 | test_sampling_params = [ |
| 101 | p for p in test_sampling_params if p.get("structured_outputs") is not None |
| 102 | ] |
| 103 | |
| 104 | run_tests(monkeypatch, MODEL, test_configs, test_sampling_params) |
| 105 | |
| 106 | |
| 107 | @single_gpu_only |
| 108 | @large_gpu_mark(min_gb=16) |
| 109 | def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch): |
| 110 | """Test consistency and acceptance rates with some different combos of |
| 111 | preemption, executor, async scheduling, prefill chunking, |
| 112 | spec decoding model length. |
| 113 | """ |
| 114 | spec_config = { |
| 115 | "method": "eagle3", |
| 116 | "num_speculative_tokens": 2, |
| 117 | "model": "nm-testing/Llama3_2_1B_speculator.eagle3", |
| 118 | } |
| 119 | # Set small draft model len to force doesn't-fit-in-drafter case. |
| 120 | spec_config_short = spec_config | {"max_model_len": 50} |
| 121 | |
| 122 | struct_outputs = StructuredOutputsParams(json=sample_json_schema) |
| 123 | |
| 124 | test_sampling_params = [ |
| 125 | dict(), |
| 126 | dict(frequency_penalty=-1.0), |
| 127 | dict(bad_words=["the", " the"]), |
| 128 | dict(logprobs=2), |
| 129 | dict(logprobs=2, frequency_penalty=-1.0), |
| 130 | dict(prompt_logprobs=2), |
| 131 | dict(prompt_logprobs=2, logprobs=2), |
| 132 | dict(structured_outputs=struct_outputs), |
| 133 | dict( |
| 134 | structured_outputs=struct_outputs, |
| 135 | logprobs=2, |
| 136 | frequency_penalty=-1.0, |
| 137 | ), |
| 138 | ] |
| 139 | |
| 140 | # test_preemption, executor, async_scheduling, |
| 141 | # spec_config, test_prefill_chunking |
| 142 | test_configs = [ |
| 143 | (False, "mp", False, None, False), |
| 144 | (False, "mp", False, spec_config, False), |
| 145 | (True, "mp", False, spec_config, True), |
| 146 | (True, "uni", False, spec_config_short, True), |
| 147 | (False, "mp", True, spec_config, False), |
| 148 | (True, "mp", True, spec_config, False), |
| 149 | (False, "mp", True, spec_config_short, True), |
| 150 | (True, "uni", True, spec_config, False), |
| 151 | (True, "uni", True, spec_config_short, False), |
| 152 | (True, "mp", True, spec_config, True), |
| 153 | (True, "uni", True, spec_config_short, True), |
| 154 | ] |
| 155 | |
| 156 | run_tests(monkeypatch, MTP_MODEL, test_configs, test_sampling_params) |
| 157 | |
| 158 | |
| 159 | @pytest.mark.flaky(reruns=2, only_on=current_platform.is_rocm()) |
| 160 | @pytest.mark.skipif( |
| 161 | current_platform.is_xpu(), |
| 162 | reason=("XPU matmul/attention kernels are not batch-invariant"), |
| 163 | ) |
| 164 | def test_with_ngram_gpu_spec_decoding(monkeypatch: pytest.MonkeyPatch): |
| 165 | """Test ngram_gpu speculative decoding with different configurations. |
| 166 | |
| 167 | This test specifically validates ngram_gpu behavior with various: |
| 168 | - Number of speculative tokens (2-6) |
| 169 | - Prompt lookup window sizes (min/max) |
| 170 | - Async scheduling enabled (as in production) |
| 171 | - Different executors and chunking settings |
| 172 | """ |
| 173 | # Variant with larger speculation window |
| 174 | ngram_gpu_config = { |
| 175 | "method": "ngram_gpu", |
| 176 | "num_speculative_tokens": 3, |
| 177 | "prompt_lookup_max": 3, |
| 178 | "prompt_lookup_min": 2, |
| 179 | } |
| 180 | |
| 181 | # Test configurations covering various scenarios |
| 182 | # test_preemption, executor, async_scheduling, |
| 183 | # spec_config, test_prefill_chunking |
| 184 | test_configs = [ |
| 185 | (False, "mp", False, None, False), |
| 186 | (False, "mp", False, ngram_gpu_config, False), |
| 187 | (True, "mp", False, ngram_gpu_config, True), |
| 188 | (False, "mp", True, ngram_gpu_config, False), |
| 189 | (True, "mp", True, ngram_gpu_config, False), |
| 190 | (True, "uni", True, ngram_gpu_config, False), |
| 191 | (True, "mp", True, ngram_gpu_config, True), |
| 192 | ] |
| 193 | |
| 194 | # Use MODEL (Qwen) for ngram_gpu tests as it's lighter weight |
| 195 | # and ngram_gpu doesn't require a specific draft model |
| 196 | run_tests(monkeypatch, MODEL, test_configs, [{}]) |
| 197 | |
| 198 | |
| 199 | @dynamo_config.patch(cache_size_limit=16) |
| 200 | def run_tests( |
| 201 | monkeypatch: pytest.MonkeyPatch, |
| 202 | model: str, |
| 203 | test_configs: list[tuple], |
| 204 | test_sampling_params: list[dict[str, Any]], |
| 205 | ): |
| 206 | """Test consistency of combos of async scheduling, preemption, |
| 207 | uni/multiproc executor with spec decoding.""" |
| 208 | # Flex attention supports float32. |
| 209 | attention_config = {"backend": "FLEX_ATTENTION"} |
| 210 | |
| 211 | with monkeypatch.context() as m: |
| 212 | # lock matmul precision to full FP32 (IEEE) |
| 213 | m.setenv("VLLM_FLOAT32_MATMUL_PRECISION", "highest") |
| 214 | outputs: list[tuple[str, list, list]] = [] |
| 215 | for n, ( |
| 216 | test_preemption, |
| 217 | executor, |
| 218 | async_scheduling, |
| 219 | spec_config, |
| 220 | test_prefill_chunking, |
| 221 | ) in enumerate(test_configs, 1): |
| 222 | test_str = f"{n}/{len(test_configs)}" |
| 223 | test_results = run_test( |
| 224 | model, |
| 225 | test_str, |
| 226 | test_sampling_params, |
| 227 | test_preemption, |
| 228 | executor, |
| 229 | async_scheduling, |
| 230 | spec_config, |
| 231 | test_prefill_chunking=test_prefill_chunking, |
| 232 | attention_config=attention_config, |
| 233 | ) |
| 234 | outputs.append(test_results) |
| 235 | |
| 236 | baseline_config, baseline_tests, _ = outputs[0] |
| 237 | _, _, baseline_acceptances = next( |
| 238 | (o for o in outputs if o[2] is not None), (None, None, None) |
| 239 | ) |
| 240 | |
| 241 | print(f"BASELINE: config=[{baseline_config}], accept_rates={baseline_acceptances}") |
| 242 | |
| 243 | failure = None |
| 244 | for test_config, test_outputs, test_acceptance_rates in outputs[1:]: |
| 245 | for (base_outs, base_logprobs), base_acceptance_rate, ( |
| 246 | test_outs, |
| 247 | test_logprobs, |
| 248 | ), test_acceptance_rate, params in zip( |
| 249 | baseline_tests, |
| 250 | baseline_acceptances or repeat(None), |
| 251 | test_outputs, |
| 252 | test_acceptance_rates or repeat(None), |
| 253 | test_sampling_params, |
| 254 | ): |
| 255 | reason = None |
| 256 | try: |
| 257 | check_outputs_equal( |
| 258 | outputs_0_lst=base_outs, |
| 259 | outputs_1_lst=test_outs, |
| 260 | name_0=f"baseline=[{baseline_config}], params={params}", |
| 261 | name_1=f"config=[{test_config}], params={params}", |
| 262 | ) |
| 263 | except AssertionError as e: |
| 264 | reason = "outputs ", e |
| 265 | |
| 266 | if reason is None: |
| 267 | try: |
| 268 | assert _all_logprobs_match(base_logprobs, test_logprobs) |
| 269 | except AssertionError as e: |
| 270 | reason = "logprobs", e |
| 271 | |
| 272 | if reason is None: |
| 273 | try: |
| 274 | if ( |
| 275 | base_acceptance_rate is not None |
| 276 | and test_acceptance_rate is not None |
| 277 | ): |
| 278 | if "spec_mml=None" in test_config: |
| 279 | # Preemption causes more variance in acceptance rates |
| 280 | if ( |
| 281 | current_platform.is_rocm() |
| 282 | and "preemption=True" in test_config |
| 283 | ): |
| 284 | tolerance = 0.10 |
| 285 | else: |
| 286 | tolerance = 0.05 |
| 287 | assert ( |
| 288 | test_acceptance_rate > base_acceptance_rate |
| 289 | or test_acceptance_rate |
| 290 | == pytest.approx(base_acceptance_rate, rel=tolerance) |
| 291 | ) |
| 292 | else: |
| 293 | # Currently the reported acceptance rate is expected to be |
| 294 | # lower when we sometimes skip drafting altogether. |
| 295 | assert test_acceptance_rate > 0.1 |
| 296 | except AssertionError as e: |
| 297 | reason = "accept ", e |
| 298 | |
| 299 | if reason is None: |
| 300 | print( |
| 301 | f"\033[32mPASSED\033[0m: " |
| 302 | f"config=[{test_config}], params={params}" |
| 303 | f" accept_rate={test_acceptance_rate}" |
| 304 | ) |
| 305 | else: |
| 306 | reason_str, _ = reason |
| 307 | print( |
| 308 | f"\033[31mFAILED\033[0m({reason_str}): " |
| 309 | f"config=[{test_config}], params={params}" |
| 310 | f" accept_rate={test_acceptance_rate}" |
| 311 | ) |
| 312 | if failure is None: |
| 313 | _, failure = reason |
| 314 | |
| 315 | if failure is not None: |
| 316 | raise failure |
| 317 | |
| 318 | |
| 319 | def run_test( |
| 320 | model: str, |
| 321 | test_str: str, |
| 322 | sampling_param_tests: list[dict[str, Any]], |
| 323 | test_preemption: bool, |
| 324 | executor: str, |
| 325 | async_scheduling: bool, |
| 326 | spec_config: dict[str, Any] | None, |
| 327 | test_prefill_chunking: bool, |
| 328 | attention_config: dict[str, Any] | None = None, |
| 329 | ): |
| 330 | spec_decoding = spec_config is not None |
| 331 | cache_arg: dict[str, Any] = ( |
| 332 | # Force preemptions: with 33 blocks (one is the reserved null block) |
| 333 | # the cache holds at most a single max-length request, so the ~34 |
| 334 | # concurrent prompts contend and trigger preemption. (Prompts here are |
| 335 | # << max_model_len, so dropping max_model_len from 4096 to 512 doesn't |
| 336 | # change generation behavior.) |
| 337 | dict(num_gpu_blocks_override=33, max_model_len=512) |
| 338 | if test_preemption |
| 339 | else dict(gpu_memory_utilization=0.9, max_model_len=4096) |
| 340 | ) |
| 341 | spec_mml = (spec_config or {}).get("max_model_len") |
| 342 | spec_method = (spec_config or {}).get("method", "none") |
| 343 | test_config = ( |
| 344 | f"executor={executor}, preemption={test_preemption}, " |
| 345 | f"async_sched={async_scheduling}, " |
| 346 | f"chunk_prefill={test_prefill_chunking}, " |
| 347 | f"spec_decoding={spec_decoding}, spec_method={spec_method}, spec_mml={spec_mml}" |
| 348 | ) |
| 349 | print("-" * 80) |
| 350 | print(f"---- TESTING {test_str}: {test_config}") |
| 351 | print("-" * 80) |
| 352 | |
| 353 | with VllmRunner( |
| 354 | model, |
| 355 | enable_chunked_prefill=test_prefill_chunking, |
| 356 | # Force prefill chunking |
| 357 | max_num_batched_tokens=48 if test_prefill_chunking else None, |
| 358 | enforce_eager=ENFORCE_EAGER, |
| 359 | async_scheduling=async_scheduling, |
| 360 | distributed_executor_backend=executor, |
| 361 | dtype="float32", |
| 362 | speculative_config=spec_config, |
| 363 | disable_log_stats=False, |
| 364 | attention_config=attention_config, |
| 365 | enable_prefix_caching=False if current_platform.is_rocm() else None, |
| 366 | **cache_arg, |
| 367 | ) as vllm_model: |
| 368 | results = [] |
| 369 | acceptance_rates: list[float] | None = [] if spec_decoding else None |
| 370 | for override_params in sampling_param_tests: |
| 371 | metrics_before = vllm_model.llm.get_metrics() |
| 372 | print(f"----------- RUNNING PARAMS: {override_params}") |
| 373 | results.append( |
| 374 | vllm_model.generate( |
| 375 | example_prompts, |
| 376 | sampling_params=SamplingParams(**default_params, **override_params), |
| 377 | return_logprobs=True, |
| 378 | ) |
| 379 | ) |
| 380 | metrics_after = vllm_model.llm.get_metrics() |
| 381 | if acceptance_rates is not None: |
| 382 | acceptance_rate = _get_acceptance_rate(metrics_before, metrics_after) |
| 383 | acceptance_rates.append(acceptance_rate) |
| 384 | print(f"ACCEPTANCE RATE {acceptance_rate}") |
| 385 | |
| 386 | if test_preemption: |
| 387 | preemptions = _get_count( |
| 388 | metrics_before, metrics_after, "vllm:num_preemptions" |
| 389 | ) |
| 390 | assert preemptions > 0, "preemption test had no preemptions" |
| 391 | |
| 392 | if len(results) > 1: |
| 393 | # First check that the different parameter configs |
| 394 | # actually result in different output. |
| 395 | for (other_test_outs, other_test_logprobs), params in zip( |
| 396 | results[1:], sampling_param_tests[1:] |
| 397 | ): |
| 398 | with pytest.raises(AssertionError): |
| 399 | check_outputs_equal( |
| 400 | outputs_0_lst=results[0][0], |
| 401 | outputs_1_lst=other_test_outs, |
| 402 | name_0=f"baseline params={params}", |
| 403 | name_1=f"other params={params}", |
| 404 | ) |
| 405 | assert _all_logprobs_match(results[0][1], other_test_logprobs) |
| 406 | |
| 407 | return test_config, results, acceptance_rates |
| 408 | |
| 409 | |
| 410 | def _all_logprobs_match(req_a, req_b) -> bool: |
| 411 | return ( |
| 412 | req_a == req_b |
| 413 | or len(req_a) == len(req_b) |
| 414 | and all( |
| 415 | len(seq_a) == len(seq_b) |
| 416 | and all(_logprobs_match(a, b) for a, b in zip(seq_a, seq_b)) |
| 417 | for seq_a, seq_b in zip(req_a, req_b) |
| 418 | ) |
| 419 | ) |
| 420 | |
| 421 | |
| 422 | def _logprobs_match( |
| 423 | lps_a: dict[int, Logprob] | None, |
| 424 | lps_b: dict[int, Logprob] | None, |
| 425 | ) -> bool: |
| 426 | if lps_a is None or lps_b is None: |
| 427 | return lps_a is lps_b |
| 428 | rel_tol, abs_tol = 1e-3, 1e-6 |
| 429 | return ( |
| 430 | len(lps_a) == len(lps_b) |
| 431 | and lps_a.keys() == lps_b.keys() |
| 432 | and all( |
| 433 | a.decoded_token == b.decoded_token |
| 434 | and a.rank == pytest.approx(b.rank, rel=0.005) |
| 435 | and a.logprob == pytest.approx(b.logprob, rel=rel_tol, abs=abs_tol) |
| 436 | for a, b in ((lps_a[x], lps_b[x]) for x in lps_a) |
| 437 | ) |
| 438 | ) |
| 439 | |
| 440 | |
| 441 | def _get_acceptance_rate(before: list[Metric], after: list[Metric]) -> float: |
| 442 | draft = _get_count(before, after, "vllm:spec_decode_num_draft_tokens") |
| 443 | accept = _get_count(before, after, "vllm:spec_decode_num_accepted_tokens") |
| 444 | return accept / draft if draft > 0 else 0.0 |
| 445 | |
| 446 | |
| 447 | def _get_count(before: list[Metric], after: list[Metric], name: str) -> int: |
| 448 | before_val = next(m.value for m in before if m.name == name) |
| 449 | after_val = next(m.value for m in after if m.name == name) |
| 450 | return after_val - before_val |