@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 267 lines 9.32 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Engine-level parity: ReplaySSM standard decode vs the baseline SSM kernel."""
4
5 import pytest
6
7 import vllm.envs as envs
8 from vllm.v1.metrics.reader import Counter
9
10 from ...models.utils import check_logprobs_close
11 from ...utils import large_gpu_mark, multi_gpu_test
12
13 # Mamba2 (Nemotron-3) hybrids.
14 MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16"
15 # FlashInfer 0.7.0's standard SSU cannot compile Nano's DIM=80 configuration.
16 MAMBA2_FLASHINFER_MODEL = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4"
17 MODELS = [
18 pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)),
19 ]
20 FLASHINFER_MODELS = [
21 pytest.param(MAMBA2_FLASHINFER_MODEL, marks=large_gpu_mark(min_gb=40)),
22 ]
23
24 PROMPTS = [
25 "The capital of France is",
26 "Once upon a time, in a small village,",
27 ]
28
29
30 @pytest.fixture(autouse=True)
31 def _use_v1_model_runner_by_default(monkeypatch):
32 # Triton ReplaySSM is V1-only. FlashInfer V2 tests override this locally.
33 with monkeypatch.context() as patch:
34 patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
35 envs.disable_envs_cache()
36 yield
37 envs.disable_envs_cache()
38
39
40 def _check_replayssm_parity(
41 vllm_runner,
42 model_name,
43 *,
44 tensor_parallel_size=1,
45 mamba_backend: str = "triton",
46 name_1: str = "replayssm",
47 expected_v2: bool | None = None,
48 ):
49 # Compare logprobs, not greedy ids: ReplaySSM's fp arithmetic can flip a
50 # near-tie. Baseline and ReplaySSM use the same backend so only ReplaySSM
51 # behavior varies.
52 common = dict(
53 max_model_len=1024,
54 trust_remote_code=True,
55 enable_prefix_caching=False,
56 mamba_cache_mode="none",
57 tensor_parallel_size=tensor_parallel_size,
58 mamba_backend=mamba_backend,
59 )
60 with vllm_runner(model_name, **common) as llm:
61 if expected_v2 is not None:
62 assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2
63 baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
64 with vllm_runner(
65 model_name,
66 use_replayssm=True,
67 replayssm_buffer_len=16,
68 **common,
69 ) as llm:
70 if expected_v2 is not None:
71 assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2
72 replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
73
74 check_logprobs_close(
75 outputs_0_lst=baseline,
76 outputs_1_lst=replay,
77 name_0="baseline",
78 name_1=name_1,
79 )
80
81
82 @pytest.mark.parametrize("model_name", MODELS)
83 def test_replayssm_decode_matches_baseline(vllm_runner, model_name):
84 _check_replayssm_parity(vllm_runner, model_name)
85
86
87 @multi_gpu_test(num_gpus=2)
88 @pytest.mark.parametrize("model_name", [MAMBA2_MODEL])
89 def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name):
90 # Tensor-parallel correctness: ReplaySSM's caches and checkpoint state are
91 # sharded per rank, so TP2 decode must still match the baseline at TP2.
92 _check_replayssm_parity(vllm_runner, model_name, tensor_parallel_size=2)
93
94
95 @pytest.mark.parametrize("model_name", FLASHINFER_MODELS)
96 @pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"])
97 def test_replayssm_flashinfer_decode_matches_baseline(
98 vllm_runner, model_name, monkeypatch, use_v2_model_runner
99 ):
100 try:
101 with monkeypatch.context() as patch:
102 patch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner)))
103 envs.disable_envs_cache()
104 _check_replayssm_parity(
105 vllm_runner,
106 model_name,
107 mamba_backend="flashinfer",
108 name_1="replayssm_flashinfer",
109 expected_v2=use_v2_model_runner,
110 )
111 finally:
112 # The context restores the environment before the final cache reset.
113 envs.disable_envs_cache()
114
115
116 @pytest.mark.parametrize("model_name", FLASHINFER_MODELS)
117 def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_name):
118 common = dict(
119 max_model_len=1024,
120 trust_remote_code=True,
121 enable_prefix_caching=False,
122 mamba_cache_mode="none",
123 speculative_config={
124 "method": "ngram",
125 "num_speculative_tokens": 3,
126 "prompt_lookup_max": 3,
127 },
128 )
129 with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm:
130 baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
131 with vllm_runner(
132 model_name,
133 use_replayssm=True,
134 replayssm_buffer_len=16,
135 mamba_backend="flashinfer",
136 **common,
137 ) as llm:
138 replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
139
140 check_logprobs_close(
141 outputs_0_lst=baseline,
142 outputs_1_lst=replay,
143 name_0="baseline_spec",
144 name_1="replayssm_flashinfer_spec",
145 )
146
147
148 @multi_gpu_test(num_gpus=2)
149 @large_gpu_mark(min_gb=40)
150 @pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"])
151 def test_replayssm_flashinfer_mtp(vllm_runner, monkeypatch, use_v2_model_runner):
152 common = dict(
153 max_model_len=1024,
154 trust_remote_code=True,
155 enable_prefix_caching=False,
156 mamba_cache_mode="none",
157 mamba_backend="flashinfer",
158 tensor_parallel_size=2,
159 disable_log_stats=False,
160 speculative_config={"method": "mtp", "num_speculative_tokens": 3},
161 )
162 try:
163 with monkeypatch.context() as patch:
164 patch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner)))
165 envs.disable_envs_cache()
166 with vllm_runner(
167 MAMBA2_FLASHINFER_MODEL,
168 use_replayssm=True,
169 replayssm_buffer_len=16,
170 **common,
171 ) as llm:
172 assert (
173 llm.llm.llm_engine.vllm_config.use_v2_model_runner
174 is use_v2_model_runner
175 )
176 outputs = llm.generate_greedy(PROMPTS, max_tokens=32)
177 draft_count = sum(
178 metric.value
179 for metric in llm.llm.get_metrics()
180 if isinstance(metric, Counter)
181 and metric.name == "vllm:spec_decode_num_drafts"
182 )
183 finally:
184 envs.disable_envs_cache()
185
186 # At least one request must run past the 16-token replay window; another
187 # may legitimately stop early on EOS.
188 assert any(len(token_ids) > 16 for token_ids, _ in outputs)
189 assert draft_count > 0
190
191
192 # Prefix spans several mamba blocks; prefix caching only reuses full blocks.
193 _PC_SENTENCE = (
194 "In a detailed survey of state space models, the authors compared many "
195 "architectures across a wide range of long-context language tasks and "
196 "measured their throughput, memory use, and accuracy in careful detail. "
197 )
198 _PC_PREFIX = _PC_SENTENCE * 120
199 PREFIX_CACHING_PROMPTS = [
200 _PC_PREFIX + "The most important conclusion was that",
201 _PC_PREFIX + "Surprisingly, the experiments showed that",
202 _PC_PREFIX + "The most important conclusion was that",
203 ]
204
205
206 def _prefix_cache_hits(llm) -> int:
207 return sum(
208 m.value
209 for m in llm.llm.get_metrics()
210 if isinstance(m, Counter) and m.name == "vllm:prefix_cache_hits"
211 )
212
213
214 def _check_replayssm_prefix_caching_parity(
215 vllm_runner, model_name, *, tensor_parallel_size=1
216 ):
217 # align mode materializes the exact SSM state at each block boundary, so
218 # ReplaySSM's cached prefixes must match the always-materialized baseline.
219 common = dict(
220 max_model_len=8192,
221 trust_remote_code=True,
222 enable_prefix_caching=True,
223 enable_chunked_prefill=True,
224 mamba_cache_mode="align",
225 disable_log_stats=False, # required for llm.get_metrics()
226 tensor_parallel_size=tensor_parallel_size,
227 )
228 with vllm_runner(model_name, **common) as llm:
229 baseline = llm.generate_greedy_logprobs(
230 PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5
231 )
232 with vllm_runner(
233 model_name, use_replayssm=True, replayssm_buffer_len=16, **common
234 ) as llm:
235 # Prime the cache, then measure, so cache hits are deterministic.
236 llm.generate_greedy_logprobs(
237 PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5
238 )
239 replay = llm.generate_greedy_logprobs(
240 PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5
241 )
242 replay_hits = _prefix_cache_hits(llm)
243
244 # Without real cache hits the cached path is never exercised.
245 assert replay_hits > 0, (
246 "ReplaySSM align-mode run produced no prefix-cache hits; the shared "
247 "prefix may be shorter than one mamba block, so prefix caching is inert"
248 )
249 check_logprobs_close(
250 outputs_0_lst=baseline,
251 outputs_1_lst=replay,
252 name_0="baseline_align_pc",
253 name_1="replayssm_align_pc",
254 )
255
256
257 @pytest.mark.parametrize("model_name", MODELS)
258 def test_replayssm_prefix_caching_matches_baseline(vllm_runner, model_name):
259 _check_replayssm_prefix_caching_parity(vllm_runner, model_name)
260
261
262 @multi_gpu_test(num_gpus=2)
263 @pytest.mark.parametrize("model_name", [MAMBA2_MODEL])
264 def test_replayssm_prefix_caching_matches_baseline_tp2(vllm_runner, model_name):
265 _check_replayssm_prefix_caching_parity(
266 vllm_runner, model_name, tensor_parallel_size=2
267 )