@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 167 lines 5.33 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """CPU speculative-decoding end-to-end correctness."""
4
5 import os
6
7 import pytest
8
9 from tests.models.utils import check_logprobs_close
10 from tests.v1.e2e.spec_decode.utils import compute_acceptance_len
11 from vllm import LLM, SamplingParams
12 from vllm.platforms import current_platform
13 from vllm.triton_utils import HAS_TRITON
14
15 pytestmark = pytest.mark.cpu_model
16
17 if not current_platform.is_cpu():
18 pytest.skip("skipping CPU-only tests", allow_module_level=True)
19
20 # Bound the KV cache so the run does not scale with host memory.
21 os.environ.setdefault("VLLM_CPU_KVCACHE_SPACE", "1")
22
23 MODEL = "Qwen/Qwen3-0.6B"
24 NUM_LOGPROBS = 5
25 SP = SamplingParams(max_tokens=48, temperature=0, logprobs=NUM_LOGPROBS)
26
27 # Repetitive prompts so the ngram proposer finds matches and actually drafts;
28 # without hits the run is indistinguishable from plain decoding.
29 PROMPTS = [
30 "The capital of France is Paris. The capital of Italy is Rome. "
31 "The capital of Spain is Madrid. The capital of France is",
32 "def add(a, b):\n return a + b\n\ndef sub(a, b):\n return a - b\n\n"
33 "def mul(a, b):\n return",
34 "one two three four five one two three four five one two three four",
35 ]
36
37
38 def _make_llm(**overrides) -> LLM:
39 base = dict(
40 model=MODEL,
41 dtype="bfloat16",
42 max_model_len=1024,
43 disable_log_stats=False,
44 )
45 base.update(overrides)
46 return LLM(**base)
47
48
49 def _tuples(outputs) -> list[tuple[list[int], str, object]]:
50 return [
51 (list(o.outputs[0].token_ids), o.outputs[0].text, o.outputs[0].logprobs)
52 for o in outputs
53 ]
54
55
56 @pytest.fixture(scope="module")
57 def baseline_refs():
58 llm = _make_llm()
59 refs = _tuples(llm.generate(PROMPTS, SP))
60 del llm
61 return refs
62
63
64 def _spec_metric(llm: LLM, name: str) -> float:
65 for metric in llm.get_metrics():
66 if metric.name == name:
67 return float(getattr(metric, "value", 0.0))
68 return 0.0
69
70
71 def test_ngram_spec_decode_matches_baseline(baseline_refs):
72 """Greedy output must be unchanged by speculative decoding.
73
74 The drafted-token check keeps the test from passing vacuously when the
75 proposer silently produces nothing.
76 """
77 llm = _make_llm(
78 speculative_config={
79 "method": "ngram",
80 "prompt_lookup_max": 5,
81 "prompt_lookup_min": 3,
82 "num_speculative_tokens": 3,
83 },
84 )
85 got = _tuples(llm.generate(PROMPTS, SP))
86 drafted = _spec_metric(llm, "vllm:spec_decode_num_draft_tokens")
87 accepted = _spec_metric(llm, "vllm:spec_decode_num_accepted_tokens")
88 del llm
89
90 assert drafted > 0, "ngram proposer drafted no tokens; the test would be vacuous"
91 assert accepted > 0, f"no draft tokens accepted out of {drafted} drafted"
92
93 check_logprobs_close(
94 outputs_0_lst=baseline_refs,
95 outputs_1_lst=got,
96 name_0="no_spec_decode",
97 name_1="ngram_spec_decode",
98 )
99
100
101 SMOKE_PROMPTS = [
102 "The capital of France is",
103 "2 + 2 equals",
104 "In one word, the color of the sky is",
105 "Q: If a train travels 60 miles in 1.5 hours, what is its average speed?\nA:",
106 ]
107 SMOKE_SP = SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True)
108
109 # Methods routed to the V2 runner, using the same target/draft pairs as the GPU
110 # suite. MTP needs no separate draft checkpoint: its speculator layer is
111 # trained jointly with, and shipped inside, the target model's own weights.
112 SMOKE_CONFIGS = [
113 pytest.param(
114 "meta-llama/Llama-3.2-1B-Instruct",
115 {
116 "method": "eagle3",
117 "model": "nm-testing/Llama3_2_1B_speculator.eagle3",
118 "num_speculative_tokens": 3,
119 },
120 id="eagle3",
121 ),
122 pytest.param(
123 "Qwen/Qwen3-4B",
124 {
125 "method": "dflash",
126 "model": "z-lab/Qwen3-4B-DFlash-b16",
127 # The checkpoint allows 16, but accepted tokens plateau around 64
128 # regardless, so a larger draft only burns CPU on rejected tokens.
129 "num_speculative_tokens": 4,
130 },
131 id="dflash",
132 ),
133 pytest.param(
134 "Qwen/Qwen3.5-0.8B",
135 {
136 "method": "mtp",
137 "num_speculative_tokens": 3,
138 },
139 id="mtp",
140 ),
141 ]
142
143 # Measured 1.85 (eagle3), 1.99 (dflash), 2.29 (mtp) on CPU; this floor leaves
144 # room for numeric drift while still failing if speculation stops paying off.
145 MIN_ACCEPTANCE_LEN = 1.5
146
147
148 @pytest.mark.skipif(
149 not HAS_TRITON, reason="the V2 runner needs triton-cpu, which CI builds"
150 )
151 @pytest.mark.parametrize("model,speculative_config", SMOKE_CONFIGS)
152 def test_v2_speculator_smoke(model: str, speculative_config: dict):
153 """Gates that drafts are accepted often enough for speculation to pay off."""
154 llm = _make_llm(model=model, speculative_config=speculative_config)
155 try:
156 assert llm.llm_engine.vllm_config.use_v2_model_runner, (
157 f"{speculative_config['method']} fell back to the V1 runner; "
158 "this test is meant to cover the V2 path"
159 )
160 llm.generate(SMOKE_PROMPTS, SMOKE_SP)
161 acceptance = compute_acceptance_len(llm.get_metrics())
162 assert acceptance >= MIN_ACCEPTANCE_LEN, (
163 f"acceptance length {acceptance:.3f} below {MIN_ACCEPTANCE_LEN}; "
164 "drafts are being rejected far more than expected"
165 )
166 finally:
167 del llm