iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Tyler Michael Smith
2026-07-08 17:59:48 -04:00
committed by GitHub
co-authored by Claude Opus 4.6 Itay Etelis Itay Etelis Or Ozeri mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 0d2f4e7c9c
commit dcdd756d75
5 changed files with 330 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/lm_eval.yaml
+41
View File
@@ -298,6 +298,47 @@ steps:
- uv pip install --system 'gpt-oss[eval]==0.0.5'
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt
- label: LM Eval KV-Offload (1xH200)
key: kv-offload-small
timeout_in_minutes: 30
device: h200_35gb
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
- vllm/v1/kv_offload/
- vllm/v1/simple_kv_offload/
- tests/evals/gsm8k/test_gsm8k_offloading.py
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it"
- label: LM Eval KV-Offload (2xH100)
key: kv-offload-medium
timeout_in_minutes: 60
device: h100
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
- vllm/v1/kv_offload/
- vllm/v1/simple_kv_offload/
- tests/evals/gsm8k/test_gsm8k_offloading.py
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b"
- label: LM Eval KV-Offload (4xH100)
key: kv-offload-large
timeout_in_minutes: 60
device: h100
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
- vllm/v1/kv_offload/
- vllm/v1/simple_kv_offload/
- tests/evals/gsm8k/test_gsm8k_offloading.py
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash"
- label: MRCR Eval Small Models
device: h200_35gb
timeout_in_minutes: 30
rust/src/server/src/routes/cache.rs
+10 -4
View File
@@ -1,8 +1,9 @@
use std::sync::Arc;
use axum::Json;
use axum::extract::{Query, State};
use axum::http::StatusCode;
use serde::Deserialize;
use serde::{Deserialize, Serialize};
use crate::error::ApiError;
use crate::state::AppState;
@@ -16,19 +17,24 @@ pub(crate) struct ResetPrefixCacheParams {
reset_external: bool,
}
#[derive(Debug, Serialize)]
pub(crate) struct ResetPrefixCacheResponse {
success: bool,
}
/// Reset the local prefix cache and optionally the connector-managed external
/// cache.
pub async fn reset_prefix_cache(
State(state): State<Arc<AppState>>,
Query(params): Query<ResetPrefixCacheParams>,
) -> Result<StatusCode, ApiError> {
state
) -> Result<Json<ResetPrefixCacheResponse>, ApiError> {
let success = state
.engine_core_client()
.reset_prefix_cache(params.reset_running_requests, params.reset_external)
.await
.map_err(|error| utility_call_error("reset_prefix_cache", error))?;
Ok(StatusCode::OK)
Ok(Json(ResetPrefixCacheResponse { success }))
}
/// Reset the multi-modal cache.
rust/src/server/src/routes/tests.rs
+4 -1
View File
@@ -4796,7 +4796,10 @@ async fn reset_prefix_cache_route_sends_expected_utility_call() {
let status = response.status();
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
assert!(body.is_empty());
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&body).expect("json body"),
json!({"success": true})
);
engine_task.await.expect("mock engine task");
}
tests/evals/gsm8k/test_gsm8k_offloading.py
+269
View File
@@ -0,0 +1,269 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
GSM8K correctness test for CPU KV offloading connectors.
Regression guard for stride computation bugs in the offloading worker
(e.g. https://github.com/vllm-project/vllm/pull/46888) and silent KV
cache data corruption during CPU offloading. Runs GSM8K twice, dropping
the GPU prefix cache (but not the CPU cache) between runs so the second
run reloads offloaded KV data from CPU. The reset only succeeds once
in-flight offload transfers have released their GPU blocks, so retrying
it until success also waits for offloading to complete.
Correctness is enforced via GSM8K accuracy on both runs: corrupt reloaded
KV drops accuracy below threshold. The reload itself is not asserted, so
a silently skipped reload (e.g. offloading disabled) would not be flagged.
Covers both KV offloading connectors (OffloadingConnector and
SimpleCPUOffloadConnector) across four architecture families:
- Hybrid Mamba (NemotronH: attention + Mamba)
- Heterogeneous head dim (Gemma 4)
- Hybrid GDN (Qwen 3.5: attention + GatedDeltaNet)
- Compressed attention (DeepSeek-V4-Flash: CSA)
Usage:
pytest -s -v evals/gsm8k/test_gsm8k_offloading.py
"""
import json
import time
from dataclasses import dataclass, field
import pytest
import requests
from tests.utils import RemoteOpenAIServer
from vllm.platforms import current_platform
from .gsm8k_eval import evaluate_gsm8k
if not current_platform.is_cuda_alike():
pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
NUM_QUESTIONS = 200
NUM_FEWSHOT = 5
_OFFLOAD_SYNC_TIMEOUT = 60
def _kv_transfer_config(connector: str, cpu_gib: int = 4) -> str:
if connector == "OffloadingConnector":
return json.dumps(
{
"kv_connector": "OffloadingConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"spec_name": "CPUOffloadingSpec",
"cpu_bytes_to_use": cpu_gib << 30,
"eviction_policy": "lru",
},
}
)
elif connector == "SimpleCPUOffloadConnector":
return json.dumps(
{
"kv_connector": "SimpleCPUOffloadConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"cpu_bytes_to_use": cpu_gib << 30,
},
}
)
else:
raise ValueError(f"Unknown connector: {connector}")
def _force_engine_step(base_url: str) -> None:
"""Force an engine step so completed offload transfers get processed."""
requests.post(
f"{base_url}/v1/completions",
json={"prompt": "0", "max_tokens": 1, "temperature": 0},
timeout=60,
).raise_for_status()
def _reset_gpu_prefix_cache(base_url: str) -> None:
"""Drop the GPU prefix cache while keeping the CPU (connector) cache, so
the next run must reload KV data through the connector.
The reset fails while asynchronous offload transfers still hold GPU
blocks, so retry until it succeeds. Requires VLLM_SERVER_DEV_MODE=1.
"""
deadline = time.monotonic() + _OFFLOAD_SYNC_TIMEOUT
while True:
resp = requests.post(
f"{base_url}/reset_prefix_cache",
params={"reset_external": "false"},
timeout=30,
)
resp.raise_for_status()
if resp.json().get("success"):
return
assert time.monotonic() < deadline, (
f"prefix cache reset did not succeed within {_OFFLOAD_SYNC_TIMEOUT}s; "
"async offload may be stuck"
)
_force_engine_step(base_url)
@dataclass
class OffloadingModelConfig:
id: str
model: str
connector: str
accuracy_threshold: float
tolerance: float = 0.05
extra_server_args: list[str] = field(default_factory=list)
cpu_offload_gib: int = 4
startup_timeout: int = 600
MODELS = [
# ── OffloadingConnector ──────────────────────────────────────────
OffloadingModelConfig(
id="offloading-nemotron-h-8b",
model="nvidia/Nemotron-H-8B-Base-8K",
connector="OffloadingConnector",
# Baseline ~0.49 on 200 questions (measured on GB200).
accuracy_threshold=0.39,
),
OffloadingModelConfig(
id="offloading-gemma-4-e4b-it",
model="google/gemma-4-E4B-it",
connector="OffloadingConnector",
# Baseline ~0.64 on 200 questions (measured on GB200).
accuracy_threshold=0.55,
),
OffloadingModelConfig(
id="offloading-qwen3.5-35b",
model="Qwen/Qwen3.5-35B-A3B",
connector="OffloadingConnector",
accuracy_threshold=0.75,
extra_server_args=[
"--tensor-parallel-size",
"2",
"--enable-expert-parallel",
],
startup_timeout=1200,
),
OffloadingModelConfig(
id="offloading-deepseek-v4-flash",
model="deepseek-ai/DeepSeek-V4-Flash",
connector="OffloadingConnector",
# Baseline ~0.97 on 200 questions (measured on GB200).
accuracy_threshold=0.90,
extra_server_args=[
"--tensor-parallel-size",
"4",
"--enable-expert-parallel",
"--kv-cache-dtype",
"fp8",
"--block-size",
"256",
],
cpu_offload_gib=16,
startup_timeout=1200,
),
# ── SimpleCPUOffloadConnector ────────────────────────────────────
OffloadingModelConfig(
id="simple-nemotron-h-8b",
model="nvidia/Nemotron-H-8B-Base-8K",
connector="SimpleCPUOffloadConnector",
accuracy_threshold=0.45,
),
OffloadingModelConfig(
id="simple-gemma-4-e4b-it",
model="google/gemma-4-E4B-it",
connector="SimpleCPUOffloadConnector",
accuracy_threshold=0.55,
),
OffloadingModelConfig(
id="simple-qwen3.5-35b",
model="Qwen/Qwen3.5-35B-A3B",
connector="SimpleCPUOffloadConnector",
accuracy_threshold=0.75,
extra_server_args=[
"--tensor-parallel-size",
"2",
"--enable-expert-parallel",
],
startup_timeout=1200,
),
OffloadingModelConfig(
id="simple-deepseek-v4-flash",
model="deepseek-ai/DeepSeek-V4-Flash",
connector="SimpleCPUOffloadConnector",
accuracy_threshold=0.90,
extra_server_args=[
"--tensor-parallel-size",
"4",
"--enable-expert-parallel",
"--kv-cache-dtype",
"fp8",
"--block-size",
"256",
],
cpu_offload_gib=16,
startup_timeout=1200,
),
]
@pytest.mark.parametrize("cfg", MODELS, ids=lambda c: c.id)
def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig):
if "--tensor-parallel-size" in cfg.extra_server_args:
tp_size = int(
cfg.extra_server_args[
cfg.extra_server_args.index("--tensor-parallel-size") + 1
]
)
if current_platform.device_count() < tp_size:
pytest.skip(f"Requires {tp_size} GPUs")
# Prefix caching must be explicitly enabled: SimpleCPUOffloadConnector requires it.
server_args = [
"--enforce-eager",
"--max-model-len",
"4096",
"--enable-prefix-caching",
"--no-disable-hybrid-kv-cache-manager",
"--kv-transfer-config",
_kv_transfer_config(cfg.connector, cfg.cpu_offload_gib),
"--trust-remote-code",
"--disable-uvicorn-access-log",
*cfg.extra_server_args,
]
with RemoteOpenAIServer(
cfg.model,
server_args,
# /reset_prefix_cache requires dev mode.
env_dict={"VLLM_SERVER_DEV_MODE": "1"},
max_wait_seconds=cfg.startup_timeout,
) as server:
base_url = f"http://{server.host}:{server.port}"
for run_idx in range(1, 3):
results = evaluate_gsm8k(
num_questions=NUM_QUESTIONS,
num_shots=NUM_FEWSHOT,
host=f"http://{server.host}",
port=server.port,
)
print(
f"GSM8K run {run_idx}/2 + {cfg.connector} ({cfg.id}): "
f"accuracy={results['accuracy']:.4f}, "
f"invalid_rate={results['invalid_rate']:.3f}, "
f"latency={results['latency']:.1f}s"
)
assert results["accuracy"] >= (cfg.accuracy_threshold - cfg.tolerance), (
f"GSM8K run {run_idx}/2 accuracy "
f"{results['accuracy']:.4f} below "
f"{cfg.accuracy_threshold - cfg.tolerance:.4f}"
)
if run_idx == 1:
_reset_gpu_prefix_cache(base_url)
vllm/entrypoints/serve/dev/cache/api_router.py
Vendored
+6 -5
View File
@@ -3,7 +3,7 @@
from fastapi import APIRouter, FastAPI, Query, Request
from fastapi.responses import Response
from fastapi.responses import JSONResponse, Response
from vllm.engine.protocol import EngineClient
from vllm.logger import init_logger
@@ -29,18 +29,19 @@ async def reset_prefix_cache(
Optionally, if the query parameter `reset_external=true`
also resets the external (connector-managed) prefix cache.
Note that we currently do not check if the prefix cache
is successfully reset in the API server.
Returns `{"success": bool}`. The reset fails (`success=false`) while
blocks are still held, e.g. by running requests or in-flight async KV
offload transfers; callers may retry.
Example:
POST /reset_prefix_cache?reset_external=true
"""
logger.info("Resetting prefix cache...")
await engine_client(raw_request).reset_prefix_cache(
success = await engine_client(raw_request).reset_prefix_cache(
reset_running_requests, reset_external
)
return Response(status_code=200)
return JSONResponse(content={"success": bool(success)})
@router.post("/reset_mm_cache")