@tduncan / Nvidia-Optimization-Stack-VLLM / commits / dcdd756d75

[CI] GSM8K eval integration test for KV offloading (#46893)

Signed-off-by: Tyler Michael Smith <tyler@tylermsmith.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Signed-off-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Tyler Michael Smith committed Jul 8, 2026 at 17:59 UTC dcdd756d75a9a2cdc9ff257c7d44c67b9b9fee68
5 files changed +330 -10
.buildkite/test_areas/lm_eval.yaml
+41
@@ -298,6 +298,47 @@ steps:
298 - uv pip install --system 'gpt-oss[eval]==0.0.5'
299 - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt
300
301 +- label: LM Eval KV-Offload (1xH200)
302 + key: kv-offload-small
303 + timeout_in_minutes: 30
304 + device: h200_35gb
305 + source_file_dependencies:
306 + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
307 + - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
308 + - vllm/v1/kv_offload/
309 + - vllm/v1/simple_kv_offload/
310 + - tests/evals/gsm8k/test_gsm8k_offloading.py
311 + commands:
312 + - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it"
313 +
314 +- label: LM Eval KV-Offload (2xH100)
315 + key: kv-offload-medium
316 + timeout_in_minutes: 60
317 + device: h100
318 + num_devices: 2
319 + source_file_dependencies:
320 + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
321 + - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
322 + - vllm/v1/kv_offload/
323 + - vllm/v1/simple_kv_offload/
324 + - tests/evals/gsm8k/test_gsm8k_offloading.py
325 + commands:
326 + - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b"
327 +
328 +- label: LM Eval KV-Offload (4xH100)
329 + key: kv-offload-large
330 + timeout_in_minutes: 60
331 + device: h100
332 + num_devices: 4
333 + source_file_dependencies:
334 + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
335 + - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
336 + - vllm/v1/kv_offload/
337 + - vllm/v1/simple_kv_offload/
338 + - tests/evals/gsm8k/test_gsm8k_offloading.py
339 + commands:
340 + - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash"
341 +
342 - label: MRCR Eval Small Models
343 device: h200_35gb
344 timeout_in_minutes: 30
rust/src/server/src/routes/cache.rs
+10 -4
@@ -1,8 +1,9 @@
1 use std::sync::Arc;
2
3 +use axum::Json;
4 use axum::extract::{Query, State};
5 use axum::http::StatusCode;
5 -use serde::Deserialize;
6 +use serde::{Deserialize, Serialize};
7
8 use crate::error::ApiError;
9 use crate::state::AppState;
@@ -16,19 +17,24 @@ pub(crate) struct ResetPrefixCacheParams {
17 reset_external: bool,
18 }
19
20 +#[derive(Debug, Serialize)]
21 +pub(crate) struct ResetPrefixCacheResponse {
22 + success: bool,
23 +}
24 +
25 /// Reset the local prefix cache and optionally the connector-managed external
26 /// cache.
27 pub async fn reset_prefix_cache(
28 State(state): State<Arc<AppState>>,
29 Query(params): Query<ResetPrefixCacheParams>,
24 -) -> Result<StatusCode, ApiError> {
25 - state
30 +) -> Result<Json<ResetPrefixCacheResponse>, ApiError> {
31 + let success = state
32 .engine_core_client()
33 .reset_prefix_cache(params.reset_running_requests, params.reset_external)
34 .await
35 .map_err(|error| utility_call_error("reset_prefix_cache", error))?;
36
31 - Ok(StatusCode::OK)
37 + Ok(Json(ResetPrefixCacheResponse { success }))
38 }
39
40 /// Reset the multi-modal cache.
rust/src/server/src/routes/tests.rs
+4 -1
@@ -4796,7 +4796,10 @@ async fn reset_prefix_cache_route_sends_expected_utility_call() {
4796 let status = response.status();
4797 let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
4798 assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
4799 - assert!(body.is_empty());
4799 + assert_eq!(
4800 + serde_json::from_slice::<serde_json::Value>(&body).expect("json body"),
4801 + json!({"success": true})
4802 + );
4803 engine_task.await.expect("mock engine task");
4804 }
4805
tests/evals/gsm8k/test_gsm8k_offloading.py new
+269
@@ -0,0 +1,269 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +GSM8K correctness test for CPU KV offloading connectors.
5 +
6 +Regression guard for stride computation bugs in the offloading worker
7 +(e.g. https://github.com/vllm-project/vllm/pull/46888) and silent KV
8 +cache data corruption during CPU offloading. Runs GSM8K twice, dropping
9 +the GPU prefix cache (but not the CPU cache) between runs so the second
10 +run reloads offloaded KV data from CPU. The reset only succeeds once
11 +in-flight offload transfers have released their GPU blocks, so retrying
12 +it until success also waits for offloading to complete.
13 +
14 +Correctness is enforced via GSM8K accuracy on both runs: corrupt reloaded
15 +KV drops accuracy below threshold. The reload itself is not asserted, so
16 +a silently skipped reload (e.g. offloading disabled) would not be flagged.
17 +
18 +Covers both KV offloading connectors (OffloadingConnector and
19 +SimpleCPUOffloadConnector) across four architecture families:
20 + - Hybrid Mamba (NemotronH: attention + Mamba)
21 + - Heterogeneous head dim (Gemma 4)
22 + - Hybrid GDN (Qwen 3.5: attention + GatedDeltaNet)
23 + - Compressed attention (DeepSeek-V4-Flash: CSA)
24 +
25 +Usage:
26 + pytest -s -v evals/gsm8k/test_gsm8k_offloading.py
27 +"""
28 +
29 +import json
30 +import time
31 +from dataclasses import dataclass, field
32 +
33 +import pytest
34 +import requests
35 +
36 +from tests.utils import RemoteOpenAIServer
37 +from vllm.platforms import current_platform
38 +
39 +from .gsm8k_eval import evaluate_gsm8k
40 +
41 +if not current_platform.is_cuda_alike():
42 + pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
43 +
44 +NUM_QUESTIONS = 200
45 +NUM_FEWSHOT = 5
46 +
47 +_OFFLOAD_SYNC_TIMEOUT = 60
48 +
49 +
50 +def _kv_transfer_config(connector: str, cpu_gib: int = 4) -> str:
51 + if connector == "OffloadingConnector":
52 + return json.dumps(
53 + {
54 + "kv_connector": "OffloadingConnector",
55 + "kv_role": "kv_both",
56 + "kv_connector_extra_config": {
57 + "spec_name": "CPUOffloadingSpec",
58 + "cpu_bytes_to_use": cpu_gib << 30,
59 + "eviction_policy": "lru",
60 + },
61 + }
62 + )
63 + elif connector == "SimpleCPUOffloadConnector":
64 + return json.dumps(
65 + {
66 + "kv_connector": "SimpleCPUOffloadConnector",
67 + "kv_role": "kv_both",
68 + "kv_connector_extra_config": {
69 + "cpu_bytes_to_use": cpu_gib << 30,
70 + },
71 + }
72 + )
73 + else:
74 + raise ValueError(f"Unknown connector: {connector}")
75 +
76 +
77 +def _force_engine_step(base_url: str) -> None:
78 + """Force an engine step so completed offload transfers get processed."""
79 + requests.post(
80 + f"{base_url}/v1/completions",
81 + json={"prompt": "0", "max_tokens": 1, "temperature": 0},
82 + timeout=60,
83 + ).raise_for_status()
84 +
85 +
86 +def _reset_gpu_prefix_cache(base_url: str) -> None:
87 + """Drop the GPU prefix cache while keeping the CPU (connector) cache, so
88 + the next run must reload KV data through the connector.
89 +
90 + The reset fails while asynchronous offload transfers still hold GPU
91 + blocks, so retry until it succeeds. Requires VLLM_SERVER_DEV_MODE=1.
92 + """
93 + deadline = time.monotonic() + _OFFLOAD_SYNC_TIMEOUT
94 + while True:
95 + resp = requests.post(
96 + f"{base_url}/reset_prefix_cache",
97 + params={"reset_external": "false"},
98 + timeout=30,
99 + )
100 + resp.raise_for_status()
101 + if resp.json().get("success"):
102 + return
103 + assert time.monotonic() < deadline, (
104 + f"prefix cache reset did not succeed within {_OFFLOAD_SYNC_TIMEOUT}s; "
105 + "async offload may be stuck"
106 + )
107 + _force_engine_step(base_url)
108 +
109 +
110 +@dataclass
111 +class OffloadingModelConfig:
112 + id: str
113 + model: str
114 + connector: str
115 + accuracy_threshold: float
116 + tolerance: float = 0.05
117 + extra_server_args: list[str] = field(default_factory=list)
118 + cpu_offload_gib: int = 4
119 + startup_timeout: int = 600
120 +
121 +
122 +MODELS = [
123 + # ── OffloadingConnector ──────────────────────────────────────────
124 + OffloadingModelConfig(
125 + id="offloading-nemotron-h-8b",
126 + model="nvidia/Nemotron-H-8B-Base-8K",
127 + connector="OffloadingConnector",
128 + # Baseline ~0.49 on 200 questions (measured on GB200).
129 + accuracy_threshold=0.39,
130 + ),
131 + OffloadingModelConfig(
132 + id="offloading-gemma-4-e4b-it",
133 + model="google/gemma-4-E4B-it",
134 + connector="OffloadingConnector",
135 + # Baseline ~0.64 on 200 questions (measured on GB200).
136 + accuracy_threshold=0.55,
137 + ),
138 + OffloadingModelConfig(
139 + id="offloading-qwen3.5-35b",
140 + model="Qwen/Qwen3.5-35B-A3B",
141 + connector="OffloadingConnector",
142 + accuracy_threshold=0.75,
143 + extra_server_args=[
144 + "--tensor-parallel-size",
145 + "2",
146 + "--enable-expert-parallel",
147 + ],
148 + startup_timeout=1200,
149 + ),
150 + OffloadingModelConfig(
151 + id="offloading-deepseek-v4-flash",
152 + model="deepseek-ai/DeepSeek-V4-Flash",
153 + connector="OffloadingConnector",
154 + # Baseline ~0.97 on 200 questions (measured on GB200).
155 + accuracy_threshold=0.90,
156 + extra_server_args=[
157 + "--tensor-parallel-size",
158 + "4",
159 + "--enable-expert-parallel",
160 + "--kv-cache-dtype",
161 + "fp8",
162 + "--block-size",
163 + "256",
164 + ],
165 + cpu_offload_gib=16,
166 + startup_timeout=1200,
167 + ),
168 + # ── SimpleCPUOffloadConnector ────────────────────────────────────
169 + OffloadingModelConfig(
170 + id="simple-nemotron-h-8b",
171 + model="nvidia/Nemotron-H-8B-Base-8K",
172 + connector="SimpleCPUOffloadConnector",
173 + accuracy_threshold=0.45,
174 + ),
175 + OffloadingModelConfig(
176 + id="simple-gemma-4-e4b-it",
177 + model="google/gemma-4-E4B-it",
178 + connector="SimpleCPUOffloadConnector",
179 + accuracy_threshold=0.55,
180 + ),
181 + OffloadingModelConfig(
182 + id="simple-qwen3.5-35b",
183 + model="Qwen/Qwen3.5-35B-A3B",
184 + connector="SimpleCPUOffloadConnector",
185 + accuracy_threshold=0.75,
186 + extra_server_args=[
187 + "--tensor-parallel-size",
188 + "2",
189 + "--enable-expert-parallel",
190 + ],
191 + startup_timeout=1200,
192 + ),
193 + OffloadingModelConfig(
194 + id="simple-deepseek-v4-flash",
195 + model="deepseek-ai/DeepSeek-V4-Flash",
196 + connector="SimpleCPUOffloadConnector",
197 + accuracy_threshold=0.90,
198 + extra_server_args=[
199 + "--tensor-parallel-size",
200 + "4",
201 + "--enable-expert-parallel",
202 + "--kv-cache-dtype",
203 + "fp8",
204 + "--block-size",
205 + "256",
206 + ],
207 + cpu_offload_gib=16,
208 + startup_timeout=1200,
209 + ),
210 +]
211 +
212 +
213 +@pytest.mark.parametrize("cfg", MODELS, ids=lambda c: c.id)
214 +def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig):
215 + if "--tensor-parallel-size" in cfg.extra_server_args:
216 + tp_size = int(
217 + cfg.extra_server_args[
218 + cfg.extra_server_args.index("--tensor-parallel-size") + 1
219 + ]
220 + )
221 + if current_platform.device_count() < tp_size:
222 + pytest.skip(f"Requires {tp_size} GPUs")
223 +
224 + # Prefix caching must be explicitly enabled: SimpleCPUOffloadConnector requires it.
225 + server_args = [
226 + "--enforce-eager",
227 + "--max-model-len",
228 + "4096",
229 + "--enable-prefix-caching",
230 + "--no-disable-hybrid-kv-cache-manager",
231 + "--kv-transfer-config",
232 + _kv_transfer_config(cfg.connector, cfg.cpu_offload_gib),
233 + "--trust-remote-code",
234 + "--disable-uvicorn-access-log",
235 + *cfg.extra_server_args,
236 + ]
237 +
238 + with RemoteOpenAIServer(
239 + cfg.model,
240 + server_args,
241 + # /reset_prefix_cache requires dev mode.
242 + env_dict={"VLLM_SERVER_DEV_MODE": "1"},
243 + max_wait_seconds=cfg.startup_timeout,
244 + ) as server:
245 + base_url = f"http://{server.host}:{server.port}"
246 +
247 + for run_idx in range(1, 3):
248 + results = evaluate_gsm8k(
249 + num_questions=NUM_QUESTIONS,
250 + num_shots=NUM_FEWSHOT,
251 + host=f"http://{server.host}",
252 + port=server.port,
253 + )
254 +
255 + print(
256 + f"GSM8K run {run_idx}/2 + {cfg.connector} ({cfg.id}): "
257 + f"accuracy={results['accuracy']:.4f}, "
258 + f"invalid_rate={results['invalid_rate']:.3f}, "
259 + f"latency={results['latency']:.1f}s"
260 + )
261 +
262 + assert results["accuracy"] >= (cfg.accuracy_threshold - cfg.tolerance), (
263 + f"GSM8K run {run_idx}/2 accuracy "
264 + f"{results['accuracy']:.4f} below "
265 + f"{cfg.accuracy_threshold - cfg.tolerance:.4f}"
266 + )
267 +
268 + if run_idx == 1:
269 + _reset_gpu_prefix_cache(base_url)
vllm/entrypoints/serve/dev/cache/api_router.py
+6 -5
@@ -3,7 +3,7 @@
3
4
5 from fastapi import APIRouter, FastAPI, Query, Request
6 -from fastapi.responses import Response
6 +from fastapi.responses import JSONResponse, Response
7
8 from vllm.engine.protocol import EngineClient
9 from vllm.logger import init_logger
@@ -29,18 +29,19 @@ async def reset_prefix_cache(
29 Optionally, if the query parameter `reset_external=true`
30 also resets the external (connector-managed) prefix cache.
31
32 - Note that we currently do not check if the prefix cache
33 - is successfully reset in the API server.
32 + Returns `{"success": bool}`. The reset fails (`success=false`) while
33 + blocks are still held, e.g. by running requests or in-flight async KV
34 + offload transfers; callers may retry.
35
36 Example:
37 POST /reset_prefix_cache?reset_external=true
38 """
39 logger.info("Resetting prefix cache...")
40
40 - await engine_client(raw_request).reset_prefix_cache(
41 + success = await engine_client(raw_request).reset_prefix_cache(
42 reset_running_requests, reset_external
43 )
43 - return Response(status_code=200)
44 + return JSONResponse(content={"success": bool(success)})
45
46
47 @router.post("/reset_mm_cache")