@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 290 lines 9.63 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """GSM8K correctness test for CPU KV offloading connectors.
4
5 Regression guard for stride computation bugs in the offloading worker
6 (e.g. https://github.com/vllm-project/vllm/pull/46888) and silent KV
7 cache data corruption during CPU offloading. Runs GSM8K twice, dropping
8 the GPU prefix cache (but not the CPU cache) between runs so the second
9 run reloads offloaded KV data from CPU. The reset only succeeds once
10 in-flight offload transfers have released their GPU blocks, so retrying
11 it until success also waits for offloading to complete.
12
13 Correctness is enforced via GSM8K accuracy on both runs: corrupt reloaded
14 KV drops accuracy below threshold. The reload itself is not asserted, so
15 a silently skipped reload (e.g. offloading disabled) would not be flagged.
16
17 Covers both KV offloading connectors (OffloadingConnector and
18 SimpleCPUOffloadConnector) across five architecture families:
19 - Hybrid Mamba (NemotronH: attention + Mamba)
20 - Heterogeneous head dim (Gemma 4)
21 - Hybrid GDN (Qwen 3.5: attention + GatedDeltaNet)
22 - Compressed attention (DeepSeek-V4-Flash: CSA)
23 - Pure MLA (DeepSeek-V2-Lite: TieringOffloadingSpec, TP=2, replicated
24 single-slot host layout)
25
26 Usage:
27 pytest -s -v evals/gsm8k/test_gsm8k_offloading.py
28 """
29
30 import json
31 import time
32 from dataclasses import dataclass, field
33
34 import pytest
35 import requests
36
37 from tests.utils import RemoteOpenAIServer
38 from vllm.platforms import current_platform
39
40 from .gsm8k_eval import evaluate_gsm8k
41
42 if not current_platform.is_cuda_alike():
43 pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
44
45 NUM_QUESTIONS = 200
46 NUM_FEWSHOT = 5
47
48 _OFFLOAD_SYNC_TIMEOUT = 60
49
50
51 def _kv_transfer_config(
52 connector: str, cpu_gib: int = 4, spec_name: str = "CPUOffloadingSpec"
53 ) -> str:
54 if connector == "OffloadingConnector":
55 return json.dumps(
56 {
57 "kv_connector": "OffloadingConnector",
58 "kv_role": "kv_both",
59 "kv_connector_extra_config": {
60 "spec_name": spec_name,
61 "cpu_bytes_to_use": cpu_gib << 30,
62 "eviction_policy": "lru",
63 },
64 }
65 )
66 elif connector == "SimpleCPUOffloadConnector":
67 return json.dumps(
68 {
69 "kv_connector": "SimpleCPUOffloadConnector",
70 "kv_role": "kv_both",
71 "kv_connector_extra_config": {
72 "cpu_bytes_to_use": cpu_gib << 30,
73 },
74 }
75 )
76 else:
77 raise ValueError(f"Unknown connector: {connector}")
78
79
80 def _force_engine_step(base_url: str) -> None:
81 """Force an engine step so completed offload transfers get processed."""
82 requests.post(
83 f"{base_url}/v1/completions",
84 json={"prompt": "0", "max_tokens": 1, "temperature": 0},
85 timeout=60,
86 ).raise_for_status()
87
88
89 def _reset_gpu_prefix_cache(base_url: str) -> None:
90 """Drop the GPU prefix cache while keeping the CPU (connector) cache, so
91 the next run must reload KV data through the connector.
92
93 The reset fails while asynchronous offload transfers still hold GPU
94 blocks, so retry until it succeeds. Requires VLLM_SERVER_DEV_MODE=1.
95 """
96 deadline = time.monotonic() + _OFFLOAD_SYNC_TIMEOUT
97 while True:
98 resp = requests.post(
99 f"{base_url}/reset_prefix_cache",
100 params={"reset_external": "false"},
101 timeout=30,
102 )
103 resp.raise_for_status()
104 if resp.json().get("success"):
105 return
106 assert time.monotonic() < deadline, (
107 f"prefix cache reset did not succeed within {_OFFLOAD_SYNC_TIMEOUT}s; "
108 "async offload may be stuck"
109 )
110 _force_engine_step(base_url)
111
112
113 @dataclass
114 class OffloadingModelConfig:
115 id: str
116 model: str
117 connector: str
118 accuracy_threshold: float
119 tolerance: float = 0.05
120 extra_server_args: list[str] = field(default_factory=list)
121 env_dict: dict[str, str] = field(default_factory=dict)
122 cpu_offload_gib: int = 4
123 startup_timeout: int = 600
124 spec_name: str = "CPUOffloadingSpec"
125
126
127 MODELS = [
128 # ── OffloadingConnector ──────────────────────────────────────────
129 OffloadingModelConfig(
130 id="offloading-nemotron-h-8b",
131 model="nvidia/Nemotron-H-8B-Base-8K",
132 connector="OffloadingConnector",
133 # Baseline ~0.49 on 200 questions (measured on GB200).
134 accuracy_threshold=0.39,
135 cpu_offload_gib=1,
136 ),
137 OffloadingModelConfig(
138 id="offloading-gemma-4-e4b-it",
139 model="google/gemma-4-E4B-it",
140 connector="OffloadingConnector",
141 # Baseline ~0.64 on 200 questions (measured on GB200).
142 accuracy_threshold=0.55,
143 cpu_offload_gib=1,
144 ),
145 OffloadingModelConfig(
146 id="offloading-qwen3.5-35b",
147 model="Qwen/Qwen3.5-35B-A3B",
148 connector="OffloadingConnector",
149 accuracy_threshold=0.75,
150 extra_server_args=[
151 "--tensor-parallel-size",
152 "2",
153 "--enable-expert-parallel",
154 ],
155 startup_timeout=1200,
156 ),
157 OffloadingModelConfig(
158 id="offloading-deepseek-v4-flash",
159 model="deepseek-ai/DeepSeek-V4-Flash",
160 connector="OffloadingConnector",
161 # Baseline ~0.97 on 200 questions (measured on GB200).
162 accuracy_threshold=0.90,
163 extra_server_args=[
164 "--tensor-parallel-size",
165 "4",
166 "--enable-expert-parallel",
167 "--kv-cache-dtype",
168 "fp8",
169 "--block-size",
170 "256",
171 ],
172 cpu_offload_gib=16,
173 startup_timeout=1200,
174 ),
175 OffloadingModelConfig(
176 id="offloading-deepseek-v2-lite-tiering-tp2",
177 model="deepseek-ai/DeepSeek-V2-Lite",
178 connector="OffloadingConnector",
179 # Baseline 0.360/0.345 over two runs (measured on 2xA100).
180 accuracy_threshold=0.35,
181 extra_server_args=[
182 "--tensor-parallel-size",
183 "2",
184 ],
185 cpu_offload_gib=8,
186 startup_timeout=1200,
187 spec_name="TieringOffloadingSpec",
188 ),
189 # ── SimpleCPUOffloadConnector ────────────────────────────────────
190 OffloadingModelConfig(
191 id="simple-nemotron-h-8b",
192 model="nvidia/Nemotron-H-8B-Base-8K",
193 connector="SimpleCPUOffloadConnector",
194 accuracy_threshold=0.43 if current_platform.is_rocm() else 0.45,
195 ),
196 OffloadingModelConfig(
197 id="simple-gemma-4-e4b-it",
198 model="google/gemma-4-E4B-it",
199 connector="SimpleCPUOffloadConnector",
200 accuracy_threshold=0.55,
201 ),
202 OffloadingModelConfig(
203 id="simple-qwen3.5-35b",
204 model="Qwen/Qwen3.5-35B-A3B",
205 connector="SimpleCPUOffloadConnector",
206 accuracy_threshold=0.75,
207 extra_server_args=[
208 "--tensor-parallel-size",
209 "2",
210 "--enable-expert-parallel",
211 ],
212 startup_timeout=1200,
213 ),
214 OffloadingModelConfig(
215 id="simple-deepseek-v4-flash",
216 model="deepseek-ai/DeepSeek-V4-Flash",
217 connector="SimpleCPUOffloadConnector",
218 accuracy_threshold=0.90,
219 extra_server_args=[
220 "--tensor-parallel-size",
221 "4",
222 "--enable-expert-parallel",
223 "--kv-cache-dtype",
224 "fp8",
225 "--block-size",
226 "256",
227 ],
228 cpu_offload_gib=16,
229 startup_timeout=1200,
230 ),
231 ]
232
233
234 @pytest.mark.parametrize("cfg", MODELS, ids=lambda c: c.id)
235 def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig):
236 if "--tensor-parallel-size" in cfg.extra_server_args:
237 tp_size = int(
238 cfg.extra_server_args[
239 cfg.extra_server_args.index("--tensor-parallel-size") + 1
240 ]
241 )
242 if current_platform.device_count() < tp_size:
243 pytest.skip(f"Requires {tp_size} GPUs")
244
245 # Prefix caching must be explicitly enabled: SimpleCPUOffloadConnector requires it.
246 server_args = [
247 "--enforce-eager",
248 "--max-model-len",
249 "4096",
250 "--enable-prefix-caching",
251 "--no-disable-hybrid-kv-cache-manager",
252 "--kv-transfer-config",
253 _kv_transfer_config(cfg.connector, cfg.cpu_offload_gib, cfg.spec_name),
254 "--trust-remote-code",
255 "--disable-uvicorn-access-log",
256 *cfg.extra_server_args,
257 ]
258
259 with RemoteOpenAIServer(
260 cfg.model,
261 server_args,
262 # /reset_prefix_cache requires dev mode.
263 env_dict={"VLLM_SERVER_DEV_MODE": "1", **cfg.env_dict},
264 max_wait_seconds=cfg.startup_timeout,
265 ) as server:
266 base_url = f"http://{server.host}:{server.port}"
267
268 for run_idx in range(1, 3):
269 results = evaluate_gsm8k(
270 num_questions=NUM_QUESTIONS,
271 num_shots=NUM_FEWSHOT,
272 host=f"http://{server.host}",
273 port=server.port,
274 )
275
276 print(
277 f"GSM8K run {run_idx}/2 + {cfg.connector} ({cfg.id}): "
278 f"accuracy={results['accuracy']:.4f}, "
279 f"invalid_rate={results['invalid_rate']:.3f}, "
280 f"latency={results['latency']:.1f}s"
281 )
282
283 assert results["accuracy"] >= (cfg.accuracy_threshold - cfg.tolerance), (
284 f"GSM8K run {run_idx}/2 accuracy "
285 f"{results['accuracy']:.4f} below "
286 f"{cfg.accuracy_threshold - cfg.tolerance:.4f}"
287 )
288
289 if run_idx == 1:
290 _reset_gpu_prefix_cache(base_url)