From dcdd756d75a9a2cdc9ff257c7d44c67b9b9fee68 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Wed, 8 Jul 2026 17:59:48 -0400 Subject: [PATCH] [CI] GSM8K eval integration test for KV offloading (#46893) Signed-off-by: Tyler Michael Smith Signed-off-by: Tyler Michael Smith Signed-off-by: Itay Etelis Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Signed-off-by: Or Ozeri Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Itay Etelis Co-authored-by: Or Ozeri Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test_areas/lm_eval.yaml | 41 +++ rust/src/server/src/routes/cache.rs | 14 +- rust/src/server/src/routes/tests.rs | 5 +- tests/evals/gsm8k/test_gsm8k_offloading.py | 269 ++++++++++++++++++ .../entrypoints/serve/dev/cache/api_router.py | 11 +- 5 files changed, 330 insertions(+), 10 deletions(-) create mode 100644 tests/evals/gsm8k/test_gsm8k_offloading.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 4cf7774ebd..feba8f26eb 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -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 diff --git a/rust/src/server/src/routes/cache.rs b/rust/src/server/src/routes/cache.rs index 580b91d413..0626d2ffff 100644 --- a/rust/src/server/src/routes/cache.rs +++ b/rust/src/server/src/routes/cache.rs @@ -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>, Query(params): Query, -) -> Result { - state +) -> Result, 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. diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 520d504383..adf7103d05 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -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::(&body).expect("json body"), + json!({"success": true}) + ); engine_task.await.expect("mock engine task"); } diff --git a/tests/evals/gsm8k/test_gsm8k_offloading.py b/tests/evals/gsm8k/test_gsm8k_offloading.py new file mode 100644 index 0000000000..f652dcaf1b --- /dev/null +++ b/tests/evals/gsm8k/test_gsm8k_offloading.py @@ -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) diff --git a/vllm/entrypoints/serve/dev/cache/api_router.py b/vllm/entrypoints/serve/dev/cache/api_router.py index c274717c0a..63737393f2 100644 --- a/vllm/entrypoints/serve/dev/cache/api_router.py +++ b/vllm/entrypoints/serve/dev/cache/api_router.py @@ -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")