Skip to content

GitLab

Unverified Commit 90f3c01f authored by zhanqiuhu's avatar zhanqiuhu Committed by GitHub
Browse files

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


Signed-off-by: default avatarClaude <noreply@anthropic.com>
Signed-off-by: default avatarZhanqiu Hu <zh338@cornell.edu>
Co-authored-by: default avatarNicolò Lucchesi <nlucches@redhat.com>
parent 807d6803
Show whitespace changes
Inline Side-by-side
Showing with 484 additions and 21 deletions
.buildkite/test_areas/distributed.yaml
View file @ 90f3c01f
...@@ -213,6 +213,19 @@ steps: ...@@ -213,6 +213,19 @@ steps:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
timeout_in_minutes: 30
device: a100
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
- label: Pipeline + Context Parallelism (4 GPUs) - label: Pipeline + Context Parallelism (4 GPUs)
timeout_in_minutes: 60 timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests" working_dir: "/vllm-workspace/tests"
......
tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh 0 → 100755
View file @ 90f3c01f
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# NixlConnector PD + speculative decoding acceptance length test.
# Tests EAGLE3 acceptance length for both RDMA (cuda) and CPU host (cpu)
# KV buffer device paths.
#
# For each kv_buffer_device setting, starts prefill + decode vllm servers
# with NixlConnector, then runs test_spec_decode_acceptance.py to validate
# acceptance length matches the standalone SD baseline.
#
# Usage:
# CUDA_VISIBLE_DEVICES=0,1 bash tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
#
# Environment variables:
# KV_BUFFER_DEVICES - space-separated list of devices to test
# (default: "cuda cpu")
# SD_METHOD - spec decode method (default: eagle3)
# SD_MODEL - drafter model path
# MODEL_NAME - target model (default: meta-llama/Llama-3.1-8B-Instruct)
# NUM_SPEC_TOKENS - number of speculative tokens (default: 3)
# GPU_MEMORY_UTILIZATION - (default: 0.7)
set -x
# ── Model & spec decode config ──────────────────────────────────────────
MODEL_NAME="${MODEL_NAME:-meta-llama/Llama-3.1-8B-Instruct}"
SD_METHOD="${SD_METHOD:-eagle3}"
SD_MODEL="${SD_MODEL:-RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3}"
NUM_SPEC_TOKENS="${NUM_SPEC_TOKENS:-3}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}"
PREFILL_SPEC_CONFIG="{\"method\":\"${SD_METHOD}\",\"model\":\"${SD_MODEL}\",\"num_speculative_tokens\":1,\"max_model_len\":${MAX_MODEL_LEN}}"
DECODE_SPEC_CONFIG="{\"method\":\"${SD_METHOD}\",\"model\":\"${SD_MODEL}\",\"num_speculative_tokens\":${NUM_SPEC_TOKENS},\"max_model_len\":${MAX_MODEL_LEN}}"
# ── Test matrix ──────────────────────────────────────────────────────────
KV_BUFFER_DEVICES="${KV_BUFFER_DEVICES:-cuda cpu}"
# ── Cluster layout ───────────────────────────────────────────────────────
NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1}
NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1}
PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}
DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.7}
BLOCK_SIZE=${BLOCK_SIZE:-16}
GIT_ROOT=$(git rev-parse --show-toplevel)
SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
cleanup_instances() {
echo ""
echo "Cleaning up..."
kill $(jobs -pr) 2>/dev/null || true
sleep 1
kill -9 $(jobs -pr) 2>/dev/null || true
pkill -9 -f "vllm serve.*${MODEL_NAME}" 2>/dev/null || true
pkill -9 -f "toy_proxy_server.*8192" 2>/dev/null || true
sleep 1
echo "Cleanup done."
}
trap cleanup_instances EXIT
trap 'echo " Interrupted."; exit 130' INT TERM
wait_for_server() {
local port=$1
local deadline=600
local elapsed=0
echo "Waiting for server on port ${port}..."
while [ $elapsed -lt $deadline ]; do
if curl -s "localhost:${port}/v1/completions" > /dev/null 2>&1; then
echo "Server on port ${port} ready"
return 0
fi
sleep 2
elapsed=$((elapsed + 2))
done
echo "FAIL: Server on port ${port} did not start within ${deadline}s"
exit 1
}
# ── Resolve GPU list ─────────────────────────────────────────────────────
if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
IFS=',' read -ra ALL_GPUS <<< "$CUDA_VISIBLE_DEVICES"
else
ALL_GPUS=()
if [[ "$SMI_BIN" == *"nvidia"* ]]; then
num=$($SMI_BIN --query-gpu=name --format=csv,noheader | wc -l)
elif [[ "$SMI_BIN" == *"rocm"* ]]; then
num=$($SMI_BIN -l | grep -c GPU)
else
num=1
fi
for (( g=0; g<num; g++ )); do ALL_GPUS+=($g); done
fi
TOTAL_GPUS_NEEDED=$(( (NUM_PREFILL_INSTANCES * PREFILLER_TP_SIZE) + (NUM_DECODE_INSTANCES * DECODER_TP_SIZE) ))
if [[ ${#ALL_GPUS[@]} -lt $TOTAL_GPUS_NEEDED ]]; then
echo "FAIL: Need $TOTAL_GPUS_NEEDED GPUs but only have ${#ALL_GPUS[@]} (CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-not set})"
exit 1
fi
# ── Run one test iteration ───────────────────────────────────────────────
run_test_for_device() {
local kv_device=$1
if [[ "$kv_device" == "cuda" ]]; then
local kv_config='{"kv_connector":"NixlConnector","kv_role":"kv_both"}'
else
local kv_config="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"${kv_device}\"}"
fi
echo ""
echo "================================================================"
echo "NixlConnector PD + Spec Decode Acceptance Test (kv_buffer_device=${kv_device})"
echo "================================================================"
echo "Model: ${MODEL_NAME}"
echo "SD method: ${SD_METHOD}"
echo "SD model: ${SD_MODEL}"
echo "Spec tokens: ${NUM_SPEC_TOKENS}"
echo "KV buffer device: ${kv_device}"
echo "GPUs available: ${ALL_GPUS[*]}"
echo "================================================================"
local PREFILL_HOSTS=()
local PREFILL_PORTS=()
local DECODE_HOSTS=()
local DECODE_PORTS=()
local GPU_IDX=0
# Start prefill instances
for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do
local GPU_ID="${ALL_GPUS[$GPU_IDX]}"
GPU_IDX=$((GPU_IDX + 1))
for (( j=1; j < PREFILLER_TP_SIZE; j++ )); do
GPU_ID="${GPU_ID},${ALL_GPUS[$GPU_IDX]}"
GPU_IDX=$((GPU_IDX + 1))
done
local PORT=$((8100 + i))
local SIDE_CHANNEL_PORT=$((5559 + i))
echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
CUDA_VISIBLE_DEVICES=$GPU_ID \
VLLM_KV_CACHE_LAYOUT='HND' \
UCX_NET_DEVICES=all \
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
vllm serve $MODEL_NAME \
--port $PORT \
--enforce-eager \
--max-model-len $MAX_MODEL_LEN \
--block-size ${BLOCK_SIZE} \
--gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
--tensor-parallel-size $PREFILLER_TP_SIZE \
--kv-transfer-config "$kv_config" \
--speculative-config "$PREFILL_SPEC_CONFIG" \
--attention-backend FLASH_ATTN &
PREFILL_HOSTS+=("localhost")
PREFILL_PORTS+=("$PORT")
done
# Start decode instances
for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do
local GPU_ID="${ALL_GPUS[$GPU_IDX]}"
GPU_IDX=$((GPU_IDX + 1))
for (( j=1; j < DECODER_TP_SIZE; j++ )); do
GPU_ID="${GPU_ID},${ALL_GPUS[$GPU_IDX]}"
GPU_IDX=$((GPU_IDX + 1))
done
local PORT=$((8200 + i))
local SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
CUDA_VISIBLE_DEVICES=$GPU_ID \
VLLM_KV_CACHE_LAYOUT='HND' \
UCX_NET_DEVICES=all \
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
vllm serve $MODEL_NAME \
--port $PORT \
--enforce-eager \
--max-model-len $MAX_MODEL_LEN \
--block-size ${BLOCK_SIZE} \
--gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
--tensor-parallel-size $DECODER_TP_SIZE \
--kv-transfer-config "$kv_config" \
--speculative-config "$DECODE_SPEC_CONFIG" \
--attention-backend FLASH_ATTN &
DECODE_HOSTS+=("localhost")
DECODE_PORTS+=("$PORT")
done
# Wait for servers
for PORT in "${PREFILL_PORTS[@]}"; do
wait_for_server "$PORT"
done
for PORT in "${DECODE_PORTS[@]}"; do
wait_for_server "$PORT"
done
# Start proxy
local PROXY_PORT=8192
echo "Starting proxy server on port $PROXY_PORT..."
python3 "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py" \
--port $PROXY_PORT \
--prefiller-hosts ${PREFILL_HOSTS[*]} \
--prefiller-ports ${PREFILL_PORTS[*]} \
--decoder-hosts ${DECODE_HOSTS[*]} \
--decoder-ports ${DECODE_PORTS[*]} &
sleep 5
# Run test
echo "Running spec decode acceptance test (kv_buffer_device=${kv_device})..."
DECODE_PORT=${DECODE_PORTS[0]} \
TEST_MODEL=$MODEL_NAME \
python3 -m pytest -s -x "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py"
# Tear down before next iteration
cleanup_instances
sleep 3
}
# ── Main: loop over kv_buffer_device values ──────────────────────────────
for device in $KV_BUFFER_DEVICES; do
run_test_for_device "$device"
done
echo "=== All spec decode acceptance tests passed ==="
tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py 0 → 100644
View file @ 90f3c01f
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""NixlConnector PD + EAGLE3 speculative decoding acceptance length test.
- Loads MT-Bench prompts (80 prompts, 256 output tokens)
- Sends through the PD proxy (completions API)
- Scrapes Prometheus metrics from the decode server
- Asserts acceptance length matches standalone EAGLE3 baselines
Baselines from tests/v1/spec_decode/test_acceptance_length.py
(standalone EAGLE3 with same model/drafter on MT-Bench, temp=0).
PD disaggregation via NixlConnector should match within tolerance.
Environment variables (set by spec_decode_acceptance_test.sh):
TEST_MODEL - target model name
DECODE_PORT - port of the decode vLLM server (for /metrics)
"""
import os
from dataclasses import dataclass, field
from types import SimpleNamespace
from urllib.request import urlopen
import openai
import regex as re
from transformers import AutoTokenizer
from vllm.benchmarks.datasets import get_samples
PROXY_BASE_URL = "http://localhost:8192/v1"
DECODE_PORT = os.environ.get("DECODE_PORT", "8200")
MODEL_NAME = os.environ.get("TEST_MODEL", "meta-llama/Llama-3.1-8B-Instruct")
@dataclass
class Eagle3ModelConfig:
verifier: str
drafter: str
expected_acceptance_length: float
expected_acceptance_lengths_per_pos: list[float] = field(default_factory=list)
id: str = ""
rtol: float | None = None
# Standalone EAGLE3 baselines (MT-Bench, 80 prompts, 256 tokens, temp=0).
# Source: tests/v1/spec_decode/test_acceptance_length.py
EAGLE3_MODEL_CONFIGS = [
Eagle3ModelConfig(
verifier="meta-llama/Llama-3.1-8B-Instruct",
drafter="RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3",
expected_acceptance_length=2.60,
expected_acceptance_lengths_per_pos=[0.7296, 0.5208, 0.3545],
id="llama3-8b-eagle3",
),
]
DEFAULT_NUM_PROMPTS = 80
DEFAULT_OUTPUT_LEN = 256
DEFAULT_RTOL = 0.05
def _get_model_config() -> Eagle3ModelConfig:
"""Get the model config matching MODEL_NAME."""
for config in EAGLE3_MODEL_CONFIGS:
if config.verifier == MODEL_NAME:
return config
raise ValueError(
f"No Eagle3ModelConfig found for model {MODEL_NAME}. "
f"Available: {[c.verifier for c in EAGLE3_MODEL_CONFIGS]}"
)
def _get_mt_bench_prompts() -> list[str]:
"""Load MT-Bench prompts via vllm.benchmarks.datasets.get_samples."""
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
args = SimpleNamespace(
dataset_name="hf",
dataset_path="philschmid/mt-bench",
num_prompts=DEFAULT_NUM_PROMPTS,
seed=42,
no_oversample=False,
endpoint_type="openai-chat",
backend="openai-chat",
input_len=None,
output_len=DEFAULT_OUTPUT_LEN,
sharegpt_output_len=DEFAULT_OUTPUT_LEN,
hf_name=None,
hf_split="train",
hf_subset=None,
hf_output_len=DEFAULT_OUTPUT_LEN,
no_stream=True,
disable_shuffle=False,
skip_chat_template=False,
trust_remote_code=False,
enable_multimodal_chat=False,
request_id_prefix="",
)
samples = get_samples(args, tokenizer)
return [sample.prompt for sample in samples]
def _fetch_metric(metric_name: str) -> float:
"""Fetch a single counter metric from the decode server's /metrics."""
url = f"http://localhost:{DECODE_PORT}/metrics"
body = urlopen(url).read().decode()
for line in body.split("\n"):
if line.startswith(metric_name + "{") or line.startswith(metric_name + " "):
return float(line.rsplit(" ", 1)[-1])
raise ValueError(f"Metric {metric_name} not found in decode /metrics")
def _fetch_per_position_acceptance() -> dict[int, float]:
"""Fetch per-position acceptance counts from decode /metrics."""
url = f"http://localhost:{DECODE_PORT}/metrics"
body = urlopen(url).read().decode()
counts: dict[int, float] = {}
for line in body.split("\n"):
if (
"spec_decode_num_accepted_tokens_per_pos_total" in line
and not line.startswith("#")
):
m = re.search(r'position="(\d+)"', line)
if m:
counts[int(m.group(1))] = float(line.rsplit(" ", 1)[-1])
return counts
def test_spec_decode_acceptance_length():
"""Validate PD+SD acceptance length against standalone baseline.
Sends MT-Bench prompts through the PD proxy (completions API),
then checks that the decode server's speculative decoding metrics
match the known standalone baselines.
"""
config = _get_model_config()
rtol = config.rtol if config.rtol is not None else DEFAULT_RTOL
prompts = _get_mt_bench_prompts()
assert len(prompts) == DEFAULT_NUM_PROMPTS, (
f"Expected {DEFAULT_NUM_PROMPTS} prompts, got {len(prompts)}"
)
client = openai.OpenAI(api_key="EMPTY", base_url=PROXY_BASE_URL)
for i, prompt in enumerate(prompts):
resp = client.completions.create(
model=MODEL_NAME,
prompt=prompt,
max_tokens=DEFAULT_OUTPUT_LEN,
temperature=0.0,
top_p=1.0,
)
if i < 3:
text = resp.choices[0].text.strip()[:100]
print(f" [{i}] {prompt[:60]}... -> {text}...")
# ── Extract metrics from decode server ────────────────────────────
n_drafts = _fetch_metric("vllm:spec_decode_num_drafts_total")
n_accepted = _fetch_metric("vllm:spec_decode_num_accepted_tokens_total")
assert n_drafts > 0, "No spec-decode drafts were generated"
acceptance_length = 1 + (n_accepted / n_drafts)
per_pos_counts = _fetch_per_position_acceptance()
per_pos_rates = [
per_pos_counts.get(i, 0) / n_drafts
for i in range(len(config.expected_acceptance_lengths_per_pos))
]
# ── Report ────────────────────────────────────────────────────────
expected = config.expected_acceptance_length
expected_per_pos = config.expected_acceptance_lengths_per_pos
print(
f"\n{config.id}: acceptance_length={acceptance_length:.3f} "
f"(expected={expected:.3f})"
)
print(f" Drafts: {n_drafts:.0f}, Accepted: {n_accepted:.0f}")
for i, (actual, exp) in enumerate(zip(per_pos_rates, expected_per_pos)):
print(f" Position {i}: {actual:.4f} (expected: {exp:.4f})")
# ── Assert overall acceptance length ──────────────────────────────
rel_error = abs(acceptance_length - expected) / expected
assert rel_error <= rtol, (
f"Acceptance length regression for {config.id}! "
f"Expected: {expected:.3f}, "
f"Got: {acceptance_length:.3f}, "
f"Relative error: {rel_error:.2%} (tolerance: {rtol:.0%}). "
f"This may indicate drafter KV was not correctly transferred."
)
# ── Assert per-position acceptance ────────────────────────────────
for i, (actual, exp) in enumerate(zip(per_pos_rates, expected_per_pos)):
if exp > 0:
pos_err = abs(actual - exp) / exp
assert pos_err <= rtol, (
f"Per-position acceptance regression at position {i} "
f"for {config.id}! "
f"Expected: {exp:.4f}, Got: {actual:.4f}, "
f"Relative error: {pos_err:.2%} "
f"(tolerance: {rtol:.0%})"
)
print(
f"\n=== PASS: {config.id} acceptance length {acceptance_length:.3f} "
f"within {rtol:.0%} of {expected:.3f} ==="
)
vllm/v1/worker/gpu_model_runner.py
View file @ 90f3c01f
...@@ -3593,9 +3593,9 @@ class GPUModelRunner( ...@@ -3593,9 +3593,9 @@ class GPUModelRunner(
# Run the model. # Run the model.
# Use persistent buffers for CUDA graphs. # Use persistent buffers for CUDA graphs.
# When spec decode is enabled, delay clearing connector metadata # When spec decode is enabled, defer connector finalization
# until after draft model runs in sample_tokens. # (wait_for_save + clear metadata) until after draft model runs.
clear_kv_metadata = self.speculative_config is None defer_kv_connector_finalize = self.speculative_config is not None
with ( with (
set_forward_context( set_forward_context(
attn_metadata, attn_metadata,
...@@ -3610,7 +3610,8 @@ class GPUModelRunner( ...@@ -3610,7 +3610,8 @@ class GPUModelRunner(
), ),
record_function_or_nullcontext("gpu_model_runner: forward"), record_function_or_nullcontext("gpu_model_runner: forward"),
self.maybe_get_kv_connector_output( self.maybe_get_kv_connector_output(
scheduler_output, clear_metadata=clear_kv_metadata scheduler_output,
defer_finalize=defer_kv_connector_finalize,
) as kv_connector_output, ) as kv_connector_output,
): ):
model_output = self._model_forward( model_output = self._model_forward(
...@@ -3843,11 +3844,11 @@ class GPUModelRunner( ...@@ -3843,11 +3844,11 @@ class GPUModelRunner(
# tokens on the CPU, so they are run after bookkeeping. # tokens on the CPU, so they are run after bookkeeping.
propose_draft_token_ids(valid_sampled_token_ids) propose_draft_token_ids(valid_sampled_token_ids)
# Clear KV connector metadata after draft model runs (if spec decode). # Finalize KV connector (wait_for_save + clear metadata) after
# This was deferred from target model forward to allow draft model # draft model runs. Deferred from target model forward to allow
# to also save its KV cache. # draft model to also save its KV cache.
if self.speculative_config is not None: if spec_config is not None:
self.clear_kv_connector_metadata() self.finalize_kv_connector()
with record_function_or_nullcontext("gpu_model_runner: eplb"): with record_function_or_nullcontext("gpu_model_runner: eplb"):
self.eplb_step() self.eplb_step()
......
vllm/v1/worker/kv_connector_model_runner_mixin.py
View file @ 90f3c01f
...@@ -67,16 +67,27 @@ class KVConnectorModelRunnerMixin: ...@@ -67,16 +67,27 @@ class KVConnectorModelRunnerMixin:
@staticmethod @staticmethod
def maybe_get_kv_connector_output( def maybe_get_kv_connector_output(
scheduler_output: "SchedulerOutput", scheduler_output: "SchedulerOutput",
clear_metadata: bool = True, defer_finalize: bool = False,
) -> AbstractContextManager[KVConnectorOutput | None]: ) -> AbstractContextManager[KVConnectorOutput | None]:
return ( return (
KVConnectorModelRunnerMixin._get_kv_connector_output( KVConnectorModelRunnerMixin._get_kv_connector_output(
scheduler_output, clear_metadata=clear_metadata scheduler_output, defer_finalize=defer_finalize
) )
if has_kv_transfer_group() if has_kv_transfer_group()
else nullcontext() else nullcontext()
) )
@staticmethod
def finalize_kv_connector() -> None:
"""Finalize the KV connector: wait_for_save and clear metadata.
Call after draft model forward when defer_finalize=True was used.
"""
if has_kv_transfer_group():
kv_connector = get_kv_transfer_group()
kv_connector.wait_for_save()
kv_connector.clear_connector_metadata()
# This context manager must be used within an active forward context. # This context manager must be used within an active forward context.
# It encapsulates the entire KV connector lifecycle within execute_model # It encapsulates the entire KV connector lifecycle within execute_model
@staticmethod @staticmethod
...@@ -84,7 +95,7 @@ class KVConnectorModelRunnerMixin: ...@@ -84,7 +95,7 @@ class KVConnectorModelRunnerMixin:
def _get_kv_connector_output( def _get_kv_connector_output(
scheduler_output: "SchedulerOutput", scheduler_output: "SchedulerOutput",
wait_for_save: bool = True, wait_for_save: bool = True,
clear_metadata: bool = True, defer_finalize: bool = False,
) -> Generator[KVConnectorOutput, None, None]: ) -> Generator[KVConnectorOutput, None, None]:
output = KVConnectorOutput() output = KVConnectorOutput()
...@@ -102,7 +113,7 @@ class KVConnectorModelRunnerMixin: ...@@ -102,7 +113,7 @@ class KVConnectorModelRunnerMixin:
try: try:
yield output yield output
finally: finally:
if wait_for_save: if wait_for_save and not defer_finalize:
kv_connector.wait_for_save() kv_connector.wait_for_save()
output.finished_sending, output.finished_recving = ( output.finished_sending, output.finished_recving = (
...@@ -113,14 +124,7 @@ class KVConnectorModelRunnerMixin: ...@@ -113,14 +124,7 @@ class KVConnectorModelRunnerMixin:
output.kv_connector_stats = kv_connector.get_kv_connector_stats() output.kv_connector_stats = kv_connector.get_kv_connector_stats()
output.kv_cache_events = kv_connector.get_kv_connector_kv_cache_events() output.kv_cache_events = kv_connector.get_kv_connector_kv_cache_events()
if clear_metadata: if not defer_finalize:
kv_connector.clear_connector_metadata()
@staticmethod
def clear_kv_connector_metadata() -> None:
"""Clear the KV connector metadata. Call after draft model runs."""
if has_kv_transfer_group():
kv_connector = get_kv_transfer_group()
kv_connector.clear_connector_metadata() kv_connector.clear_connector_metadata()
@staticmethod @staticmethod
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment