@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 90f3c01fa4
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 7, 02:45 UTC

[Spec Decode][KV Connector] Fix KV transfer in PD + speculative decoding (#35158)

Signed-off-by: Claude <noreply@anthropic.com> Signed-off-by: Zhanqiu Hu <zh338@cornell.edu> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>

zhanqiuhu committed Mar 6, 2026 at 02:50 UTC 90f3c01fa4dfc00d13beb8ae758d43365f7ba91f
5 files changed +484 -21
.buildkite/test_areas/distributed.yaml
+13
@@ -213,6 +213,19 @@ steps:
213 - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
214 - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
215
216 +- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
217 + timeout_in_minutes: 30
218 + device: a100
219 + working_dir: "/vllm-workspace/tests"
220 + num_devices: 2
221 + source_file_dependencies:
222 + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
223 + - vllm/v1/worker/kv_connector_model_runner_mixin.py
224 + - tests/v1/kv_connector/nixl_integration/
225 + commands:
226 + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
227 + - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
228 +
229 - label: Pipeline + Context Parallelism (4 GPUs)
230 timeout_in_minutes: 60
231 working_dir: "/vllm-workspace/tests"
tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh new
+237
@@ -0,0 +1,237 @@
1 +#!/bin/bash
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +#
5 +# NixlConnector PD + speculative decoding acceptance length test.
6 +# Tests EAGLE3 acceptance length for both RDMA (cuda) and CPU host (cpu)
7 +# KV buffer device paths.
8 +#
9 +# For each kv_buffer_device setting, starts prefill + decode vllm servers
10 +# with NixlConnector, then runs test_spec_decode_acceptance.py to validate
11 +# acceptance length matches the standalone SD baseline.
12 +#
13 +# Usage:
14 +# CUDA_VISIBLE_DEVICES=0,1 bash tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
15 +#
16 +# Environment variables:
17 +# KV_BUFFER_DEVICES - space-separated list of devices to test
18 +# (default: "cuda cpu")
19 +# SD_METHOD - spec decode method (default: eagle3)
20 +# SD_MODEL - drafter model path
21 +# MODEL_NAME - target model (default: meta-llama/Llama-3.1-8B-Instruct)
22 +# NUM_SPEC_TOKENS - number of speculative tokens (default: 3)
23 +# GPU_MEMORY_UTILIZATION - (default: 0.7)
24 +set -x
25 +
26 +# ── Model & spec decode config ──────────────────────────────────────────
27 +
28 +MODEL_NAME="${MODEL_NAME:-meta-llama/Llama-3.1-8B-Instruct}"
29 +SD_METHOD="${SD_METHOD:-eagle3}"
30 +SD_MODEL="${SD_MODEL:-RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3}"
31 +NUM_SPEC_TOKENS="${NUM_SPEC_TOKENS:-3}"
32 +MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}"
33 +
34 +PREFILL_SPEC_CONFIG="{\"method\":\"${SD_METHOD}\",\"model\":\"${SD_MODEL}\",\"num_speculative_tokens\":1,\"max_model_len\":${MAX_MODEL_LEN}}"
35 +DECODE_SPEC_CONFIG="{\"method\":\"${SD_METHOD}\",\"model\":\"${SD_MODEL}\",\"num_speculative_tokens\":${NUM_SPEC_TOKENS},\"max_model_len\":${MAX_MODEL_LEN}}"
36 +
37 +# ── Test matrix ──────────────────────────────────────────────────────────
38 +
39 +KV_BUFFER_DEVICES="${KV_BUFFER_DEVICES:-cuda cpu}"
40 +
41 +# ── Cluster layout ───────────────────────────────────────────────────────
42 +
43 +NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1}
44 +NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1}
45 +PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}
46 +DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
47 +GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.7}
48 +BLOCK_SIZE=${BLOCK_SIZE:-16}
49 +
50 +GIT_ROOT=$(git rev-parse --show-toplevel)
51 +
52 +SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
53 +
54 +cleanup_instances() {
55 + echo ""
56 + echo "Cleaning up..."
57 + kill $(jobs -pr) 2>/dev/null || true
58 + sleep 1
59 + kill -9 $(jobs -pr) 2>/dev/null || true
60 + pkill -9 -f "vllm serve.*${MODEL_NAME}" 2>/dev/null || true
61 + pkill -9 -f "toy_proxy_server.*8192" 2>/dev/null || true
62 + sleep 1
63 + echo "Cleanup done."
64 +}
65 +trap cleanup_instances EXIT
66 +trap 'echo " Interrupted."; exit 130' INT TERM
67 +
68 +wait_for_server() {
69 + local port=$1
70 + local deadline=600
71 + local elapsed=0
72 + echo "Waiting for server on port ${port}..."
73 + while [ $elapsed -lt $deadline ]; do
74 + if curl -s "localhost:${port}/v1/completions" > /dev/null 2>&1; then
75 + echo "Server on port ${port} ready"
76 + return 0
77 + fi
78 + sleep 2
79 + elapsed=$((elapsed + 2))
80 + done
81 + echo "FAIL: Server on port ${port} did not start within ${deadline}s"
82 + exit 1
83 +}
84 +
85 +# ── Resolve GPU list ─────────────────────────────────────────────────────
86 +
87 +if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
88 + IFS=',' read -ra ALL_GPUS <<< "$CUDA_VISIBLE_DEVICES"
89 +else
90 + ALL_GPUS=()
91 + if [[ "$SMI_BIN" == *"nvidia"* ]]; then
92 + num=$($SMI_BIN --query-gpu=name --format=csv,noheader | wc -l)
93 + elif [[ "$SMI_BIN" == *"rocm"* ]]; then
94 + num=$($SMI_BIN -l | grep -c GPU)
95 + else
96 + num=1
97 + fi
98 + for (( g=0; g<num; g++ )); do ALL_GPUS+=($g); done
99 +fi
100 +
101 +TOTAL_GPUS_NEEDED=$(( (NUM_PREFILL_INSTANCES * PREFILLER_TP_SIZE) + (NUM_DECODE_INSTANCES * DECODER_TP_SIZE) ))
102 +if [[ ${#ALL_GPUS[@]} -lt $TOTAL_GPUS_NEEDED ]]; then
103 + echo "FAIL: Need $TOTAL_GPUS_NEEDED GPUs but only have ${#ALL_GPUS[@]} (CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-not set})"
104 + exit 1
105 +fi
106 +
107 +# ── Run one test iteration ───────────────────────────────────────────────
108 +
109 +run_test_for_device() {
110 + local kv_device=$1
111 +
112 + if [[ "$kv_device" == "cuda" ]]; then
113 + local kv_config='{"kv_connector":"NixlConnector","kv_role":"kv_both"}'
114 + else
115 + local kv_config="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"${kv_device}\"}"
116 + fi
117 +
118 + echo ""
119 + echo "================================================================"
120 + echo "NixlConnector PD + Spec Decode Acceptance Test (kv_buffer_device=${kv_device})"
121 + echo "================================================================"
122 + echo "Model: ${MODEL_NAME}"
123 + echo "SD method: ${SD_METHOD}"
124 + echo "SD model: ${SD_MODEL}"
125 + echo "Spec tokens: ${NUM_SPEC_TOKENS}"
126 + echo "KV buffer device: ${kv_device}"
127 + echo "GPUs available: ${ALL_GPUS[*]}"
128 + echo "================================================================"
129 +
130 + local PREFILL_HOSTS=()
131 + local PREFILL_PORTS=()
132 + local DECODE_HOSTS=()
133 + local DECODE_PORTS=()
134 + local GPU_IDX=0
135 +
136 + # Start prefill instances
137 + for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do
138 + local GPU_ID="${ALL_GPUS[$GPU_IDX]}"
139 + GPU_IDX=$((GPU_IDX + 1))
140 + for (( j=1; j < PREFILLER_TP_SIZE; j++ )); do
141 + GPU_ID="${GPU_ID},${ALL_GPUS[$GPU_IDX]}"
142 + GPU_IDX=$((GPU_IDX + 1))
143 + done
144 +
145 + local PORT=$((8100 + i))
146 + local SIDE_CHANNEL_PORT=$((5559 + i))
147 +
148 + echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
149 + CUDA_VISIBLE_DEVICES=$GPU_ID \
150 + VLLM_KV_CACHE_LAYOUT='HND' \
151 + UCX_NET_DEVICES=all \
152 + VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
153 + vllm serve $MODEL_NAME \
154 + --port $PORT \
155 + --enforce-eager \
156 + --max-model-len $MAX_MODEL_LEN \
157 + --block-size ${BLOCK_SIZE} \
158 + --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
159 + --tensor-parallel-size $PREFILLER_TP_SIZE \
160 + --kv-transfer-config "$kv_config" \
161 + --speculative-config "$PREFILL_SPEC_CONFIG" \
162 + --attention-backend FLASH_ATTN &
163 +
164 + PREFILL_HOSTS+=("localhost")
165 + PREFILL_PORTS+=("$PORT")
166 + done
167 +
168 + # Start decode instances
169 + for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do
170 + local GPU_ID="${ALL_GPUS[$GPU_IDX]}"
171 + GPU_IDX=$((GPU_IDX + 1))
172 + for (( j=1; j < DECODER_TP_SIZE; j++ )); do
173 + GPU_ID="${GPU_ID},${ALL_GPUS[$GPU_IDX]}"
174 + GPU_IDX=$((GPU_IDX + 1))
175 + done
176 +
177 + local PORT=$((8200 + i))
178 + local SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
179 +
180 + echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
181 + CUDA_VISIBLE_DEVICES=$GPU_ID \
182 + VLLM_KV_CACHE_LAYOUT='HND' \
183 + UCX_NET_DEVICES=all \
184 + VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
185 + vllm serve $MODEL_NAME \
186 + --port $PORT \
187 + --enforce-eager \
188 + --max-model-len $MAX_MODEL_LEN \
189 + --block-size ${BLOCK_SIZE} \
190 + --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
191 + --tensor-parallel-size $DECODER_TP_SIZE \
192 + --kv-transfer-config "$kv_config" \
193 + --speculative-config "$DECODE_SPEC_CONFIG" \
194 + --attention-backend FLASH_ATTN &
195 +
196 + DECODE_HOSTS+=("localhost")
197 + DECODE_PORTS+=("$PORT")
198 + done
199 +
200 + # Wait for servers
201 + for PORT in "${PREFILL_PORTS[@]}"; do
202 + wait_for_server "$PORT"
203 + done
204 + for PORT in "${DECODE_PORTS[@]}"; do
205 + wait_for_server "$PORT"
206 + done
207 +
208 + # Start proxy
209 + local PROXY_PORT=8192
210 + echo "Starting proxy server on port $PROXY_PORT..."
211 + python3 "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py" \
212 + --port $PROXY_PORT \
213 + --prefiller-hosts ${PREFILL_HOSTS[*]} \
214 + --prefiller-ports ${PREFILL_PORTS[*]} \
215 + --decoder-hosts ${DECODE_HOSTS[*]} \
216 + --decoder-ports ${DECODE_PORTS[*]} &
217 +
218 + sleep 5
219 +
220 + # Run test
221 + echo "Running spec decode acceptance test (kv_buffer_device=${kv_device})..."
222 + DECODE_PORT=${DECODE_PORTS[0]} \
223 + TEST_MODEL=$MODEL_NAME \
224 + python3 -m pytest -s -x "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py"
225 +
226 + # Tear down before next iteration
227 + cleanup_instances
228 + sleep 3
229 +}
230 +
231 +# ── Main: loop over kv_buffer_device values ──────────────────────────────
232 +
233 +for device in $KV_BUFFER_DEVICES; do
234 + run_test_for_device "$device"
235 +done
236 +
237 +echo "=== All spec decode acceptance tests passed ==="
tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py new
+208
@@ -0,0 +1,208 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""NixlConnector PD + EAGLE3 speculative decoding acceptance length test.
4 +
5 + - Loads MT-Bench prompts (80 prompts, 256 output tokens)
6 + - Sends through the PD proxy (completions API)
7 + - Scrapes Prometheus metrics from the decode server
8 + - Asserts acceptance length matches standalone EAGLE3 baselines
9 +
10 +Baselines from tests/v1/spec_decode/test_acceptance_length.py
11 +(standalone EAGLE3 with same model/drafter on MT-Bench, temp=0).
12 +PD disaggregation via NixlConnector should match within tolerance.
13 +
14 +Environment variables (set by spec_decode_acceptance_test.sh):
15 + TEST_MODEL - target model name
16 + DECODE_PORT - port of the decode vLLM server (for /metrics)
17 +"""
18 +
19 +import os
20 +from dataclasses import dataclass, field
21 +from types import SimpleNamespace
22 +from urllib.request import urlopen
23 +
24 +import openai
25 +import regex as re
26 +from transformers import AutoTokenizer
27 +
28 +from vllm.benchmarks.datasets import get_samples
29 +
30 +PROXY_BASE_URL = "http://localhost:8192/v1"
31 +DECODE_PORT = os.environ.get("DECODE_PORT", "8200")
32 +MODEL_NAME = os.environ.get("TEST_MODEL", "meta-llama/Llama-3.1-8B-Instruct")
33 +
34 +
35 +@dataclass
36 +class Eagle3ModelConfig:
37 + verifier: str
38 + drafter: str
39 + expected_acceptance_length: float
40 + expected_acceptance_lengths_per_pos: list[float] = field(default_factory=list)
41 + id: str = ""
42 + rtol: float | None = None
43 +
44 +
45 +# Standalone EAGLE3 baselines (MT-Bench, 80 prompts, 256 tokens, temp=0).
46 +# Source: tests/v1/spec_decode/test_acceptance_length.py
47 +EAGLE3_MODEL_CONFIGS = [
48 + Eagle3ModelConfig(
49 + verifier="meta-llama/Llama-3.1-8B-Instruct",
50 + drafter="RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3",
51 + expected_acceptance_length=2.60,
52 + expected_acceptance_lengths_per_pos=[0.7296, 0.5208, 0.3545],
53 + id="llama3-8b-eagle3",
54 + ),
55 +]
56 +
57 +DEFAULT_NUM_PROMPTS = 80
58 +DEFAULT_OUTPUT_LEN = 256
59 +DEFAULT_RTOL = 0.05
60 +
61 +
62 +def _get_model_config() -> Eagle3ModelConfig:
63 + """Get the model config matching MODEL_NAME."""
64 + for config in EAGLE3_MODEL_CONFIGS:
65 + if config.verifier == MODEL_NAME:
66 + return config
67 + raise ValueError(
68 + f"No Eagle3ModelConfig found for model {MODEL_NAME}. "
69 + f"Available: {[c.verifier for c in EAGLE3_MODEL_CONFIGS]}"
70 + )
71 +
72 +
73 +def _get_mt_bench_prompts() -> list[str]:
74 + """Load MT-Bench prompts via vllm.benchmarks.datasets.get_samples."""
75 + tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
76 + args = SimpleNamespace(
77 + dataset_name="hf",
78 + dataset_path="philschmid/mt-bench",
79 + num_prompts=DEFAULT_NUM_PROMPTS,
80 + seed=42,
81 + no_oversample=False,
82 + endpoint_type="openai-chat",
83 + backend="openai-chat",
84 + input_len=None,
85 + output_len=DEFAULT_OUTPUT_LEN,
86 + sharegpt_output_len=DEFAULT_OUTPUT_LEN,
87 + hf_name=None,
88 + hf_split="train",
89 + hf_subset=None,
90 + hf_output_len=DEFAULT_OUTPUT_LEN,
91 + no_stream=True,
92 + disable_shuffle=False,
93 + skip_chat_template=False,
94 + trust_remote_code=False,
95 + enable_multimodal_chat=False,
96 + request_id_prefix="",
97 + )
98 + samples = get_samples(args, tokenizer)
99 + return [sample.prompt for sample in samples]
100 +
101 +
102 +def _fetch_metric(metric_name: str) -> float:
103 + """Fetch a single counter metric from the decode server's /metrics."""
104 + url = f"http://localhost:{DECODE_PORT}/metrics"
105 + body = urlopen(url).read().decode()
106 + for line in body.split("\n"):
107 + if line.startswith(metric_name + "{") or line.startswith(metric_name + " "):
108 + return float(line.rsplit(" ", 1)[-1])
109 + raise ValueError(f"Metric {metric_name} not found in decode /metrics")
110 +
111 +
112 +def _fetch_per_position_acceptance() -> dict[int, float]:
113 + """Fetch per-position acceptance counts from decode /metrics."""
114 + url = f"http://localhost:{DECODE_PORT}/metrics"
115 + body = urlopen(url).read().decode()
116 + counts: dict[int, float] = {}
117 + for line in body.split("\n"):
118 + if (
119 + "spec_decode_num_accepted_tokens_per_pos_total" in line
120 + and not line.startswith("#")
121 + ):
122 + m = re.search(r'position="(\d+)"', line)
123 + if m:
124 + counts[int(m.group(1))] = float(line.rsplit(" ", 1)[-1])
125 + return counts
126 +
127 +
128 +def test_spec_decode_acceptance_length():
129 + """Validate PD+SD acceptance length against standalone baseline.
130 +
131 + Sends MT-Bench prompts through the PD proxy (completions API),
132 + then checks that the decode server's speculative decoding metrics
133 + match the known standalone baselines.
134 + """
135 + config = _get_model_config()
136 + rtol = config.rtol if config.rtol is not None else DEFAULT_RTOL
137 +
138 + prompts = _get_mt_bench_prompts()
139 + assert len(prompts) == DEFAULT_NUM_PROMPTS, (
140 + f"Expected {DEFAULT_NUM_PROMPTS} prompts, got {len(prompts)}"
141 + )
142 +
143 + client = openai.OpenAI(api_key="EMPTY", base_url=PROXY_BASE_URL)
144 + for i, prompt in enumerate(prompts):
145 + resp = client.completions.create(
146 + model=MODEL_NAME,
147 + prompt=prompt,
148 + max_tokens=DEFAULT_OUTPUT_LEN,
149 + temperature=0.0,
150 + top_p=1.0,
151 + )
152 + if i < 3:
153 + text = resp.choices[0].text.strip()[:100]
154 + print(f" [{i}] {prompt[:60]}... -> {text}...")
155 +
156 + # ── Extract metrics from decode server ────────────────────────────
157 + n_drafts = _fetch_metric("vllm:spec_decode_num_drafts_total")
158 + n_accepted = _fetch_metric("vllm:spec_decode_num_accepted_tokens_total")
159 +
160 + assert n_drafts > 0, "No spec-decode drafts were generated"
161 +
162 + acceptance_length = 1 + (n_accepted / n_drafts)
163 +
164 + per_pos_counts = _fetch_per_position_acceptance()
165 + per_pos_rates = [
166 + per_pos_counts.get(i, 0) / n_drafts
167 + for i in range(len(config.expected_acceptance_lengths_per_pos))
168 + ]
169 +
170 + # ── Report ────────────────────────────────────────────────────────
171 + expected = config.expected_acceptance_length
172 + expected_per_pos = config.expected_acceptance_lengths_per_pos
173 +
174 + print(
175 + f"\n{config.id}: acceptance_length={acceptance_length:.3f} "
176 + f"(expected={expected:.3f})"
177 + )
178 + print(f" Drafts: {n_drafts:.0f}, Accepted: {n_accepted:.0f}")
179 + for i, (actual, exp) in enumerate(zip(per_pos_rates, expected_per_pos)):
180 + print(f" Position {i}: {actual:.4f} (expected: {exp:.4f})")
181 +
182 + # ── Assert overall acceptance length ──────────────────────────────
183 + rel_error = abs(acceptance_length - expected) / expected
184 +
185 + assert rel_error <= rtol, (
186 + f"Acceptance length regression for {config.id}! "
187 + f"Expected: {expected:.3f}, "
188 + f"Got: {acceptance_length:.3f}, "
189 + f"Relative error: {rel_error:.2%} (tolerance: {rtol:.0%}). "
190 + f"This may indicate drafter KV was not correctly transferred."
191 + )
192 +
193 + # ── Assert per-position acceptance ────────────────────────────────
194 + for i, (actual, exp) in enumerate(zip(per_pos_rates, expected_per_pos)):
195 + if exp > 0:
196 + pos_err = abs(actual - exp) / exp
197 + assert pos_err <= rtol, (
198 + f"Per-position acceptance regression at position {i} "
199 + f"for {config.id}! "
200 + f"Expected: {exp:.4f}, Got: {actual:.4f}, "
201 + f"Relative error: {pos_err:.2%} "
202 + f"(tolerance: {rtol:.0%})"
203 + )
204 +
205 + print(
206 + f"\n=== PASS: {config.id} acceptance length {acceptance_length:.3f} "
207 + f"within {rtol:.0%} of {expected:.3f} ==="
208 + )
vllm/v1/worker/gpu_model_runner.py
+10 -9
@@ -3593,9 +3593,9 @@ class GPUModelRunner(
3593
3594 # Run the model.
3595 # Use persistent buffers for CUDA graphs.
3596 - # When spec decode is enabled, delay clearing connector metadata
3597 - # until after draft model runs in sample_tokens.
3598 - clear_kv_metadata = self.speculative_config is None
3596 + # When spec decode is enabled, defer connector finalization
3597 + # (wait_for_save + clear metadata) until after draft model runs.
3598 + defer_kv_connector_finalize = self.speculative_config is not None
3599 with (
3600 set_forward_context(
3601 attn_metadata,
@@ -3610,7 +3610,8 @@ class GPUModelRunner(
3610 ),
3611 record_function_or_nullcontext("gpu_model_runner: forward"),
3612 self.maybe_get_kv_connector_output(
3613 - scheduler_output, clear_metadata=clear_kv_metadata
3613 + scheduler_output,
3614 + defer_finalize=defer_kv_connector_finalize,
3615 ) as kv_connector_output,
3616 ):
3617 model_output = self._model_forward(
@@ -3843,11 +3844,11 @@ class GPUModelRunner(
3844 # tokens on the CPU, so they are run after bookkeeping.
3845 propose_draft_token_ids(valid_sampled_token_ids)
3846
3846 - # Clear KV connector metadata after draft model runs (if spec decode).
3847 - # This was deferred from target model forward to allow draft model
3848 - # to also save its KV cache.
3849 - if self.speculative_config is not None:
3850 - self.clear_kv_connector_metadata()
3847 + # Finalize KV connector (wait_for_save + clear metadata) after
3848 + # draft model runs. Deferred from target model forward to allow
3849 + # draft model to also save its KV cache.
3850 + if spec_config is not None:
3851 + self.finalize_kv_connector()
3852
3853 with record_function_or_nullcontext("gpu_model_runner: eplb"):
3854 self.eplb_step()
vllm/v1/worker/kv_connector_model_runner_mixin.py
+16 -12
@@ -67,16 +67,27 @@ class KVConnectorModelRunnerMixin:
67 @staticmethod
68 def maybe_get_kv_connector_output(
69 scheduler_output: "SchedulerOutput",
70 - clear_metadata: bool = True,
70 + defer_finalize: bool = False,
71 ) -> AbstractContextManager[KVConnectorOutput | None]:
72 return (
73 KVConnectorModelRunnerMixin._get_kv_connector_output(
74 - scheduler_output, clear_metadata=clear_metadata
74 + scheduler_output, defer_finalize=defer_finalize
75 )
76 if has_kv_transfer_group()
77 else nullcontext()
78 )
79
80 + @staticmethod
81 + def finalize_kv_connector() -> None:
82 + """Finalize the KV connector: wait_for_save and clear metadata.
83 +
84 + Call after draft model forward when defer_finalize=True was used.
85 + """
86 + if has_kv_transfer_group():
87 + kv_connector = get_kv_transfer_group()
88 + kv_connector.wait_for_save()
89 + kv_connector.clear_connector_metadata()
90 +
91 # This context manager must be used within an active forward context.
92 # It encapsulates the entire KV connector lifecycle within execute_model
93 @staticmethod
@@ -84,7 +95,7 @@ class KVConnectorModelRunnerMixin:
95 def _get_kv_connector_output(
96 scheduler_output: "SchedulerOutput",
97 wait_for_save: bool = True,
87 - clear_metadata: bool = True,
98 + defer_finalize: bool = False,
99 ) -> Generator[KVConnectorOutput, None, None]:
100 output = KVConnectorOutput()
101
@@ -102,7 +113,7 @@ class KVConnectorModelRunnerMixin:
113 try:
114 yield output
115 finally:
105 - if wait_for_save:
116 + if wait_for_save and not defer_finalize:
117 kv_connector.wait_for_save()
118
119 output.finished_sending, output.finished_recving = (
@@ -113,16 +124,9 @@ class KVConnectorModelRunnerMixin:
124 output.kv_connector_stats = kv_connector.get_kv_connector_stats()
125 output.kv_cache_events = kv_connector.get_kv_connector_kv_cache_events()
126
116 - if clear_metadata:
127 + if not defer_finalize:
128 kv_connector.clear_connector_metadata()
129
119 - @staticmethod
120 - def clear_kv_connector_metadata() -> None:
121 - """Clear the KV connector metadata. Call after draft model runs."""
122 - if has_kv_transfer_group():
123 - kv_connector = get_kv_transfer_group()
124 - kv_connector.clear_connector_metadata()
125 -
130 @staticmethod
131 def use_uniform_kv_cache(
132 attn_groups: list[list[AttentionGroup]],