@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e281cb721c
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 07:15 UTC

[CI] Add MultiConnector (Nixl+Offloading) e2e edge case tests (#39343)

Signed-off-by: ZhanqiuHu <zhu@redhat.com>

zhanqiuhu committed Apr 10, 2026 at 13:35 UTC e281cb721ca21f39afff53b02f59ddedb3977172
3 files changed +665
.buildkite/test_areas/distributed.yaml
+14
@@ -295,6 +295,20 @@ steps:
295 - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
296 - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
297
298 +- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
299 + timeout_in_minutes: 30
300 + working_dir: "/vllm-workspace/tests"
301 + num_devices: 2
302 + source_file_dependencies:
303 + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
304 + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
305 + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
306 + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
307 + - tests/v1/kv_connector/nixl_integration/
308 + commands:
309 + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
310 + - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
311 +
312 - label: Pipeline + Context Parallelism (4 GPUs)
313 timeout_in_minutes: 60
314 working_dir: "/vllm-workspace/tests"
tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh new
+174
@@ -0,0 +1,174 @@
1 +#!/bin/bash
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +#
5 +# Integration edge-case tests for MultiConnector (NixlConnector + OffloadingConnector).
6 +#
7 +# Launches a P/D setup where both prefill and decode instances use MultiConnector
8 +# wrapping NixlConnector and OffloadingConnector, then runs scenario-based edge
9 +# case tests including Prometheus metrics validation.
10 +#
11 +# Tests cover: block-size boundaries, decode-side cache-hit scenarios
12 +# (cold / full / partial), direct decode (control), and prefill-side CPU
13 +# offload recovery after GPU eviction.
14 +#
15 +# Usage:
16 +# bash tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
17 +#
18 +# Environment variables:
19 +# MODEL_NAMES - model to test (default: Qwen/Qwen3-0.6B)
20 +# KV_CACHE_MEMORY_BYTES - GPU KV cache size in bytes (default: 268435456 = 256 MiB)
21 +# BLOCK_SIZE - KV cache block size (default: 128)
22 +# VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve
23 +set -xe
24 +
25 +# ── Configuration ────────────────────────────────────────────────────────
26 +
27 +MODEL_NAMES=${MODEL_NAMES:-}
28 +if [[ -n "$MODEL_NAMES" ]]; then
29 + MODELS=("$MODEL_NAMES")
30 +else
31 + MODELS=("Qwen/Qwen3-0.6B")
32 +fi
33 +
34 +KV_CACHE_MEMORY_BYTES=${KV_CACHE_MEMORY_BYTES:-268435456} # 256 MiB
35 +MAX_MODEL_LEN=${MAX_MODEL_LEN:-2048}
36 +BLOCK_SIZE=${BLOCK_SIZE:-128}
37 +VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
38 +
39 +GIT_ROOT=$(git rev-parse --show-toplevel)
40 +
41 +# ── KV transfer config ──────────────────────────────────────────────────
42 +
43 +KV_CONFIG='{
44 + "kv_connector":"MultiConnector",
45 + "kv_role":"kv_both",
46 + "kv_connector_extra_config":{
47 + "connectors":[
48 + {"kv_connector":"NixlConnector","kv_role":"kv_both"},
49 + {"kv_connector":"OffloadingConnector","kv_role":"kv_both",
50 + "kv_connector_extra_config":{"cpu_bytes_to_use":2147483648}}
51 + ]
52 + }
53 +}'
54 +KV_CONFIG=$(echo "$KV_CONFIG" | tr -d '[:space:]')
55 +
56 +# ── Helpers ──────────────────────────────────────────────────────────────
57 +
58 +trap 'kill $(jobs -pr) 2>/dev/null || true' SIGINT SIGTERM EXIT
59 +
60 +wait_for_server() {
61 + local port=$1
62 + timeout 1200 bash -c "
63 + until curl -s localhost:${port}/v1/completions > /dev/null; do
64 + sleep 1
65 + done" && return 0 || return 1
66 +}
67 +
68 +cleanup_instances() {
69 + echo "Cleaning up any running vLLM instances and proxy..."
70 + pkill -f "vllm serve" || true
71 + pkill -f "toy_proxy_server.py" || true
72 + sleep 2
73 +}
74 +
75 +# ── Run tests for one model ──────────────────────────────────────────────
76 +
77 +run_tests_for_model() {
78 + local model_name=$1
79 +
80 + echo "================================================================"
81 + echo "Testing model: $model_name (MultiConnector edge cases)"
82 + echo "================================================================"
83 +
84 + local PREFILL_PORT=8100
85 + local DECODE_PORT=8200
86 + local PROXY_PORT=8192
87 + local PREFILL_GPU=0
88 + local DECODE_GPU=1
89 + local PREFILL_SIDE_CHANNEL_PORT=5559
90 + local DECODE_SIDE_CHANNEL_PORT=5659
91 +
92 + # ── Start prefill instance ──
93 + echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT"
94 + BASE_CMD="CUDA_VISIBLE_DEVICES=$PREFILL_GPU \
95 + VLLM_KV_CACHE_LAYOUT='HND' \
96 + UCX_NET_DEVICES=all \
97 + VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \
98 + vllm serve \"$model_name\" \
99 + --port $PREFILL_PORT \
100 + --enforce-eager \
101 + --block-size ${BLOCK_SIZE} \
102 + --max-model-len $MAX_MODEL_LEN \
103 + --kv-cache-memory-bytes $KV_CACHE_MEMORY_BYTES \
104 + --tensor-parallel-size 1 \
105 + --kv-transfer-config '$KV_CONFIG'"
106 +
107 + if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
108 + IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
109 + for arg in "${extra_args[@]}"; do
110 + BASE_CMD="${BASE_CMD} $arg"
111 + done
112 + fi
113 + eval "$BASE_CMD &"
114 +
115 + # ── Start decode instance ──
116 + echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT"
117 + BASE_CMD="CUDA_VISIBLE_DEVICES=$DECODE_GPU \
118 + VLLM_KV_CACHE_LAYOUT='HND' \
119 + UCX_NET_DEVICES=all \
120 + VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \
121 + vllm serve \"$model_name\" \
122 + --port $DECODE_PORT \
123 + --enforce-eager \
124 + --block-size ${BLOCK_SIZE} \
125 + --max-model-len $MAX_MODEL_LEN \
126 + --kv-cache-memory-bytes $KV_CACHE_MEMORY_BYTES \
127 + --tensor-parallel-size 1 \
128 + --kv-transfer-config '$KV_CONFIG'"
129 +
130 + if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
131 + IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
132 + for arg in "${extra_args[@]}"; do
133 + BASE_CMD="${BASE_CMD} $arg"
134 + done
135 + fi
136 + eval "$BASE_CMD &"
137 +
138 + # ── Wait for servers ──
139 + echo "Waiting for prefill instance on port $PREFILL_PORT to start..."
140 + wait_for_server "$PREFILL_PORT"
141 + echo "Waiting for decode instance on port $DECODE_PORT to start..."
142 + wait_for_server "$DECODE_PORT"
143 +
144 + # ── Start proxy ──
145 + echo "Starting proxy server on port $PROXY_PORT"
146 + python3 "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py" \
147 + --port "$PROXY_PORT" \
148 + --prefiller-hosts localhost \
149 + --prefiller-ports "$PREFILL_PORT" \
150 + --decoder-hosts localhost \
151 + --decoder-ports "$DECODE_PORT" &
152 + sleep 5
153 +
154 + # ── Run edge case tests ──
155 + echo "Running MultiConnector edge case tests for $model_name"
156 + PREFILL_PORT=$PREFILL_PORT \
157 + DECODE_PORT=$DECODE_PORT \
158 + PROXY_PORT=$PROXY_PORT \
159 + BLOCK_SIZE=$BLOCK_SIZE \
160 + python3 -m pytest -s -x \
161 + "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py"
162 +
163 + # ── Cleanup ──
164 + cleanup_instances
165 + sleep 3
166 +}
167 +
168 +# ── Main ─────────────────────────────────────────────────────────────────
169 +
170 +for model in "${MODELS[@]}"; do
171 + run_tests_for_model "$model"
172 +done
173 +
174 +echo "All MultiConnector edge case tests passed!"
tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py new
+477
@@ -0,0 +1,477 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Integration edge-case tests for MultiConnector (NixlConnector + OffloadingConnector).
5 +
6 +Tests cover:
7 + - Output correctness across block-size boundaries (proxy vs direct prefill).
8 + - Decode-side Prometheus metrics validation (local_cache_hit,
9 + external_kv_transfer, local_compute) for cold/warm/partial cache scenarios.
10 + - Prefill-side CPU offload recovery after GPU cache eviction.
11 +
12 +Requires running servers started by run_multi_connector_edge_case_test.sh.
13 +"""
14 +
15 +import os
16 +import time
17 +import urllib.request
18 +
19 +import openai
20 +import regex as re
21 +
22 +# ── Server configuration from environment ─────────────────────────────────
23 +
24 +PREFILL_HOST = os.getenv("PREFILL_HOST", "localhost")
25 +PREFILL_PORT = os.environ["PREFILL_PORT"]
26 +DECODE_HOST = os.getenv("DECODE_HOST", "localhost")
27 +DECODE_PORT = os.environ["DECODE_PORT"]
28 +PROXY_HOST = os.getenv("PROXY_HOST", "localhost")
29 +PROXY_PORT = os.environ["PROXY_PORT"]
30 +BLOCK_SIZE = int(os.getenv("BLOCK_SIZE", "128"))
31 +
32 +# ── OpenAI clients ────────────────────────────────────────────────────────
33 +
34 +decode_client = openai.OpenAI(
35 + api_key="EMPTY",
36 + base_url=f"http://{DECODE_HOST}:{DECODE_PORT}/v1",
37 +)
38 +prefill_client = openai.OpenAI(
39 + api_key="EMPTY",
40 + base_url=f"http://{PREFILL_HOST}:{PREFILL_PORT}/v1",
41 +)
42 +proxy_client = openai.OpenAI(
43 + api_key="EMPTY",
44 + base_url=f"http://{PROXY_HOST}:{PROXY_PORT}/v1",
45 +)
46 +
47 +_MODEL = None
48 +
49 +
50 +def _get_model() -> str:
51 + global _MODEL
52 + if _MODEL is None:
53 + models = decode_client.models.list()
54 + _MODEL = models.data[0].id
55 + return _MODEL
56 +
57 +
58 +def _complete(client: openai.OpenAI, prompt: str, max_tokens: int = 20):
59 + """Send a completion request and return (text, prompt_tokens)."""
60 + resp = client.completions.create(
61 + model=_get_model(),
62 + prompt=prompt,
63 + max_tokens=max_tokens,
64 + temperature=0,
65 + )
66 + return resp.choices[0].text, resp.usage.prompt_tokens
67 +
68 +
69 +# ── Prometheus metrics helpers ────────────────────────────────────────────
70 +
71 +_METRIC_RE = re.compile(
72 + r'vllm:prompt_tokens_by_source_total\{.*?source="([^"]+)".*?\}\s+'
73 + r"([\d.eE+\-]+)"
74 +)
75 +
76 +
77 +def _fetch_metrics(host: str, port: str) -> dict[str, float]:
78 + """Scrape prompt_tokens_by_source counters from a vLLM server."""
79 + body = urllib.request.urlopen(f"http://{host}:{port}/metrics").read().decode()
80 + result = {
81 + "local_compute": 0.0,
82 + "local_cache_hit": 0.0,
83 + "external_kv_transfer": 0.0,
84 + }
85 + for m in _METRIC_RE.finditer(body):
86 + source, val = m.group(1), float(m.group(2))
87 + if source in result:
88 + result[source] += val
89 + return result
90 +
91 +
92 +def _fetch_decode_metrics() -> dict[str, float]:
93 + return _fetch_metrics(DECODE_HOST, DECODE_PORT)
94 +
95 +
96 +def _fetch_prefill_metrics() -> dict[str, float]:
97 + return _fetch_metrics(PREFILL_HOST, PREFILL_PORT)
98 +
99 +
100 +_NIXL_BYTES_RE = re.compile(r"vllm:nixl_bytes_transferred_sum\b.*?\s+([\d.eE+\-]+)")
101 +
102 +
103 +def _fetch_nixl_bytes(host: str, port: str) -> float:
104 + """Scrape total NIXL bytes transferred from a vLLM server."""
105 + body = urllib.request.urlopen(f"http://{host}:{port}/metrics").read().decode()
106 + total = 0.0
107 + for m in _NIXL_BYTES_RE.finditer(body):
108 + total += float(m.group(1))
109 + return total
110 +
111 +
112 +_OFFLOAD_BYTES_RE = re.compile(
113 + r'vllm:kv_offload_total_bytes_total\{.*?transfer_type="([^"]+)".*?\}\s+'
114 + r"([\d.eE+\-]+)"
115 +)
116 +
117 +
118 +def _fetch_offload_bytes(host: str, port: str) -> dict[str, float]:
119 + """Scrape kv_offload_total_bytes counters (CPU_to_GPU / GPU_to_CPU)."""
120 + body = urllib.request.urlopen(f"http://{host}:{port}/metrics").read().decode()
121 + result = {"CPU_to_GPU": 0.0, "GPU_to_CPU": 0.0}
122 + for m in _OFFLOAD_BYTES_RE.finditer(body):
123 + transfer_type, val = m.group(1), float(m.group(2))
124 + if transfer_type in result:
125 + result[transfer_type] += val
126 + return result
127 +
128 +
129 +def _metrics_delta(before: dict, after: dict) -> dict[str, float]:
130 + return {k: after.get(k, 0) - before.get(k, 0) for k in before}
131 +
132 +
133 +# ── Prompts (unique per test to avoid cross-test cache interference) ──────
134 +
135 +SHORT_PROMPT = "Red Hat is "
136 +
137 +MEDIUM_PROMPT = (
138 + "Red Hat is the best company in the world to work for because it works "
139 + "on open source software, which means that all the contributions are "
140 + "delivered to the community. As a result,"
141 +)
142 +
143 +
144 +def _make_prompt(n_tokens: int) -> str:
145 + """Build a prompt of ~n_tokens tokens (1 word ~ 1 token)."""
146 + return "word " * n_tokens
147 +
148 +
149 +BLOCK_BOUNDARY_PROMPT = _make_prompt(BLOCK_SIZE)
150 +ABOVE_BOUNDARY_PROMPT = _make_prompt(BLOCK_SIZE + 2)
151 +MULTI_BLOCK_PROMPT = _make_prompt(BLOCK_SIZE * 4)
152 +
153 +FULL_CACHE_HIT_PROMPT = ( # noqa: E501
154 + "The history of computing begins with Charles Babbage who designed the "
155 + "Analytical Engine in the 1830s which is considered the first general "
156 + "purpose computer design in history. Ada Lovelace is widely regarded as "
157 + "the first computer programmer for her work on the Analytical Engine. "
158 + "The modern era of computing began with Alan Turing who formalized the "
159 + "concept of computation with his Turing machine in 1936. During World "
160 + "War Two Turing worked at Bletchley Park to break the Enigma cipher. "
161 + "After the war the first electronic computers were built including ENIAC "
162 + "at the University of Pennsylvania and Colossus at Bletchley Park. "
163 + "These early machines filled entire rooms and used vacuum tubes for logic. "
164 + "The invention of the transistor at Bell Labs in 1947 revolutionized "
165 + "computing by making smaller and more reliable machines possible. "
166 + "The integrated circuit followed in the late 1950s combining multiple "
167 + "transistors on a single chip. This led to the microprocessor in the 1970s "
168 + "and eventually to the personal computer revolution of the 1980s."
169 +)
170 +
171 +PARTIAL_CACHE_PREFIX = ( # noqa: E501
172 + "Machine learning has transformed the field of artificial intelligence "
173 + "by enabling computers to learn patterns from data without being "
174 + "explicitly programmed for every task. The field has evolved dramatically "
175 + "since its inception in the 1950s when Arthur Samuel coined the term while "
176 + "working at IBM. Early approaches focused on symbolic reasoning and expert "
177 + "systems that encoded human knowledge as rules. The statistical revolution "
178 + "of the 1990s shifted the paradigm toward data driven methods. Support "
179 + "vector machines and random forests became popular for classification tasks. "
180 + "The breakthrough of deep learning in 2012 with AlexNet winning ImageNet "
181 + "changed everything. Neural networks with many layers could automatically "
182 + "learn hierarchical feature representations from raw data."
183 +)
184 +PARTIAL_CACHE_EXTENDED = PARTIAL_CACHE_PREFIX + (
185 + " Transformers have become the dominant architecture for natural language "
186 + "processing tasks including translation, summarization, and generation. "
187 + "The attention mechanism allows models to weigh the importance of different "
188 + "parts of the input sequence. Large language models like GPT and BERT "
189 + "demonstrated that pre-training on massive text corpora followed by fine "
190 + "tuning on specific tasks could achieve state of the art results across "
191 + "a wide range of benchmarks. Scaling laws suggest that larger models "
192 + "trained on more data continue to improve in capability."
193 +)
194 +
195 +# ═══════════════════════════════════════════════════════════════════════════
196 +# Output correctness across block-size boundaries (decode-side metrics)
197 +#
198 +# Each test sends via proxy, verifies output matches prefill_direct at
199 +# temperature=0, and checks decode-side metrics for NIXL transfer.
200 +# ═══════════════════════════════════════════════════════════════════════════
201 +
202 +
203 +def test_short_prompt_correctness():
204 + """Short prompt (< block_size): output matches prefill, NIXL used."""
205 + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
206 + m0 = _fetch_decode_metrics()
207 + proxy_text, _ = _complete(proxy_client, SHORT_PROMPT)
208 + time.sleep(1)
209 + m1 = _fetch_decode_metrics()
210 + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
211 + d = _metrics_delta(m0, m1)
212 +
213 + prefill_text, _ = _complete(prefill_client, SHORT_PROMPT)
214 + print(f"SHORT PROMPT: {proxy_text=}, nixl_bytes_delta={n1 - n0}")
215 + assert proxy_text == prefill_text
216 + assert d["external_kv_transfer"] > 0, (
217 + "NIXL transfer did not occur — decode may have silently fallen back "
218 + "to local compute"
219 + )
220 + assert n1 - n0 > 0, (
221 + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}"
222 + )
223 +
224 +
225 +def test_block_boundary_correctness():
226 + """Exactly block_size tokens: output matches prefill, NIXL used."""
227 + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
228 + m0 = _fetch_decode_metrics()
229 + proxy_text, pt = _complete(proxy_client, BLOCK_BOUNDARY_PROMPT)
230 + time.sleep(1)
231 + m1 = _fetch_decode_metrics()
232 + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
233 + d = _metrics_delta(m0, m1)
234 +
235 + prefill_text, _ = _complete(prefill_client, BLOCK_BOUNDARY_PROMPT)
236 + print(f"BLOCK BOUNDARY: {pt} prompt tokens, nixl_bytes_delta={n1 - n0}")
237 + assert proxy_text == prefill_text
238 + assert d["external_kv_transfer"] > 0, (
239 + "NIXL transfer did not occur — decode may have silently fallen back "
240 + "to local compute"
241 + )
242 + assert n1 - n0 > 0, (
243 + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}"
244 + )
245 +
246 +
247 +def test_above_block_boundary_correctness():
248 + """Just above block_size (partial second block): output matches prefill."""
249 + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
250 + m0 = _fetch_decode_metrics()
251 + proxy_text, pt = _complete(proxy_client, ABOVE_BOUNDARY_PROMPT)
252 + time.sleep(1)
253 + m1 = _fetch_decode_metrics()
254 + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
255 + d = _metrics_delta(m0, m1)
256 +
257 + prefill_text, _ = _complete(prefill_client, ABOVE_BOUNDARY_PROMPT)
258 + print(f"ABOVE BOUNDARY: {pt} prompt tokens, nixl_bytes_delta={n1 - n0}")
259 + assert proxy_text == prefill_text
260 + assert d["external_kv_transfer"] > 0, (
261 + "NIXL transfer did not occur — decode may have silently fallen back "
262 + "to local compute"
263 + )
264 + assert n1 - n0 > 0, (
265 + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}"
266 + )
267 +
268 +
269 +def test_multi_block_correctness():
270 + """Multi-block prompt (~4x block_size): output matches prefill."""
271 + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
272 + m0 = _fetch_decode_metrics()
273 + proxy_text, pt = _complete(proxy_client, MULTI_BLOCK_PROMPT)
274 + time.sleep(1)
275 + m1 = _fetch_decode_metrics()
276 + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
277 + d = _metrics_delta(m0, m1)
278 +
279 + prefill_text, _ = _complete(prefill_client, MULTI_BLOCK_PROMPT)
280 + print(f"MULTI BLOCK: {pt} prompt tokens, nixl_bytes_delta={n1 - n0}")
281 + assert proxy_text == prefill_text
282 + assert d["external_kv_transfer"] > 0, (
283 + "NIXL transfer did not occur — decode may have silently fallen back "
284 + "to local compute"
285 + )
286 + assert n1 - n0 > 0, (
287 + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}"
288 + )
289 +
290 +
291 +# ═══════════════════════════════════════════════════════════════════════════
292 +# Decode-side KV source validation via Prometheus metrics
293 +#
294 +# Scrape vllm:prompt_tokens_by_source_total from the DECODE server to
295 +# verify which code path (GPU prefix, NIXL, local compute) was exercised.
296 +# ═══════════════════════════════════════════════════════════════════════════
297 +
298 +
299 +def test_cold_decode_no_cache_hit_metrics():
300 + """Cold decode: external_kv_transfer==P, local_cache_hit==0."""
301 + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
302 + m0 = _fetch_decode_metrics()
303 + proxy_text, P = _complete(proxy_client, MEDIUM_PROMPT)
304 + time.sleep(1)
305 + m1 = _fetch_decode_metrics()
306 + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
307 + d = _metrics_delta(m0, m1)
308 +
309 + print(f"COLD DECODE: {P} prompt tokens, metrics delta: {d}")
310 + print(f" nixl_bytes_delta={n1 - n0}")
311 + assert len(proxy_text) > 0, "proxy returned empty response"
312 + assert d["external_kv_transfer"] == P, (
313 + f"expected external_kv_transfer={P}, got {d['external_kv_transfer']}"
314 + )
315 + assert d["local_compute"] == 1, (
316 + f"expected local_compute=1, got {d['local_compute']}"
317 + )
318 + assert d["local_cache_hit"] == 0, (
319 + f"expected local_cache_hit=0, got {d['local_cache_hit']}"
320 + )
321 + assert n1 - n0 > 0, (
322 + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}"
323 + )
324 +
325 +
326 +def test_full_decode_gpu_cache_hit_metrics():
327 + """Prime decode, resend via proxy: local_cache_hit==cached blocks."""
328 + decode_text, _ = _complete(decode_client, FULL_CACHE_HIT_PROMPT)
329 +
330 + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
331 + m0 = _fetch_decode_metrics()
332 + proxy_text, P = _complete(proxy_client, FULL_CACHE_HIT_PROMPT)
333 + time.sleep(1)
334 + m1 = _fetch_decode_metrics()
335 + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
336 + d = _metrics_delta(m0, m1)
337 +
338 + cached = (P // BLOCK_SIZE) * BLOCK_SIZE
339 + expected_nixl = P - cached
340 +
341 + print(f"FULL CACHE HIT: {P} tokens, cached={cached}, nixl={expected_nixl}")
342 + print(f" metrics delta: {d}, nixl_bytes_delta={n1 - n0}")
343 + assert len(proxy_text) > 0, "proxy returned empty response"
344 + assert d["local_cache_hit"] == cached, (
345 + f"expected local_cache_hit={cached}, got {d['local_cache_hit']}"
346 + )
347 + assert d["external_kv_transfer"] == expected_nixl, (
348 + f"expected external_kv_transfer={expected_nixl}, "
349 + f"got {d['external_kv_transfer']}"
350 + )
351 + assert d["local_compute"] == 1, (
352 + f"expected local_compute=1 (recomputed last token), got {d['local_compute']}"
353 + )
354 + assert n1 - n0 > 0, (
355 + f"expected nixl_bytes_transferred to increase (partial NIXL for "
356 + f"uncached tail), got delta={n1 - n0}"
357 + )
358 +
359 +
360 +def test_partial_decode_gpu_cache_hit_metrics():
361 + """Prime with prefix, extend via proxy: partial local_cache_hit."""
362 + _, prefix_tokens = _complete(decode_client, PARTIAL_CACHE_PREFIX)
363 + cached = (prefix_tokens // BLOCK_SIZE) * BLOCK_SIZE
364 + assert cached >= BLOCK_SIZE, (
365 + f"PARTIAL_CACHE_PREFIX too short ({prefix_tokens} tokens) for partial "
366 + f"cache hit test with block_size={BLOCK_SIZE}"
367 + )
368 +
369 + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
370 + m0 = _fetch_decode_metrics()
371 + proxy_text, P = _complete(proxy_client, PARTIAL_CACHE_EXTENDED)
372 + time.sleep(1)
373 + m1 = _fetch_decode_metrics()
374 + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
375 + d = _metrics_delta(m0, m1)
376 +
377 + expected_nixl = P - cached
378 +
379 + print(f"PARTIAL CACHE HIT: {P} tokens, cached={cached}, nixl={expected_nixl}")
380 + print(f" metrics delta: {d}, nixl_bytes_delta={n1 - n0}")
381 + assert len(proxy_text) > 0, "proxy returned empty response"
382 + assert d["external_kv_transfer"] == expected_nixl, (
383 + f"expected external_kv_transfer={expected_nixl}, "
384 + f"got {d['external_kv_transfer']}"
385 + )
386 + assert d["local_cache_hit"] == cached, (
387 + f"expected local_cache_hit={cached}, got {d['local_cache_hit']}"
388 + )
389 + assert d["local_compute"] == 1, (
390 + f"expected local_compute=1 (recomputed last token), got {d['local_compute']}"
391 + )
392 + assert n1 - n0 > 0, (
393 + f"expected nixl_bytes_transferred to increase (NIXL for uncached "
394 + f"tail), got delta={n1 - n0}"
395 + )
396 +
397 +
398 +def test_decode_direct_all_local_compute():
399 + """Direct decode (no proxy): local_compute==P, no transfers."""
400 + prompt = "The speed of light is approximately"
401 + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
402 + m0 = _fetch_decode_metrics()
403 + text, P = _complete(decode_client, prompt)
404 + time.sleep(1)
405 + m1 = _fetch_decode_metrics()
406 + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
407 + d = _metrics_delta(m0, m1)
408 +
409 + print(f"DIRECT DECODE: {text!r} ({P} tokens), metrics delta: {d}")
410 + print(f" nixl_bytes_delta={n1 - n0}")
411 + assert len(text.strip()) > 0, "empty output from direct decode"
412 + assert d["local_compute"] == P, (
413 + f"expected local_compute={P}, got {d['local_compute']}"
414 + )
415 + assert d["external_kv_transfer"] == 0, (
416 + f"expected external_kv_transfer=0, got {d['external_kv_transfer']}"
417 + )
418 + assert n1 - n0 == 0, (
419 + f"expected no nixl_bytes_transferred for direct decode, got delta={n1 - n0}"
420 + )
421 +
422 +
423 +# ═══════════════════════════════════════════════════════════════════════════
424 +# Prefill-side CPU offload validation via Prometheus metrics
425 +#
426 +# Scrape vllm:prompt_tokens_by_source_total from the PREFILL server.
427 +# Exercises the OffloadingConnector read path: after GPU cache eviction,
428 +# the OffloadingConnector restores KV from CPU (NixlConnector cannot help
429 +# for direct requests without kv_transfer_params).
430 +# ═══════════════════════════════════════════════════════════════════════════
431 +
432 +EVICTION_PROMPT = ( # noqa: E501
433 + "Quantum computing leverages quantum mechanical phenomena like "
434 + "superposition and entanglement to perform computations that would be "
435 + "intractable for classical computers. This has implications for "
436 + "cryptography, drug discovery, and optimization problems. Richard Feynman "
437 + "first proposed the idea of quantum computing in 1982 when he observed "
438 + "that simulating quantum systems on classical computers was exponentially "
439 + "hard. Peter Shor developed a quantum algorithm for factoring large "
440 + "numbers in polynomial time which threatens RSA encryption. Grover search "
441 + "algorithm provides a quadratic speedup for unstructured search problems. "
442 + "Companies like IBM Google and Rigetti are building quantum processors "
443 + "with increasing numbers of qubits. Error correction remains a major "
444 + "challenge as quantum states are extremely fragile and prone to decoherence."
445 +)
446 +
447 +
448 +def test_prefill_cpu_offload_after_gpu_eviction():
449 + """Prefill-side: evict GPU, re-request directly, CPU offload restores KV."""
450 + text1, P = _complete(prefill_client, EVICTION_PROMPT, max_tokens=30)
451 +
452 + for i in range(100):
453 + _complete(prefill_client, f"Eviction prompt number {i}: " + _make_prompt(200))
454 +
455 + ob0 = _fetch_offload_bytes(PREFILL_HOST, PREFILL_PORT)
456 + m0 = _fetch_prefill_metrics()
457 + text2, _ = _complete(prefill_client, EVICTION_PROMPT, max_tokens=30)
458 +
459 + cpu_to_gpu_delta = 0.0
460 + for _ in range(10):
461 + time.sleep(1)
462 + ob1 = _fetch_offload_bytes(PREFILL_HOST, PREFILL_PORT)
463 + cpu_to_gpu_delta = ob1["CPU_to_GPU"] - ob0["CPU_to_GPU"]
464 + if cpu_to_gpu_delta > 0:
465 + break
466 +
467 + m1 = _fetch_prefill_metrics()
468 + d = _metrics_delta(m0, m1)
469 +
470 + print(f"PREFILL CPU OFFLOAD: run1={text1[:60]!r}, run2={text2[:60]!r}")
471 + print(f" prefill metrics delta: {d}")
472 + print(f" cpu_to_gpu bytes delta: {cpu_to_gpu_delta}")
473 + assert text1 == text2, f"inconsistent after eviction: {text1=!r}, {text2=!r}"
474 + assert cpu_to_gpu_delta > 0, (
475 + f"expected cpu_to_gpu bytes > 0 (OffloadingConnector should restore "
476 + f"KV from CPU to GPU), got {cpu_to_gpu_delta}"
477 + )