diff --git a/benchmarks/kv_cache_watermark.sh b/benchmarks/kv_cache_watermark.sh new file mode 100755 index 0000000000..258afa9fce --- /dev/null +++ b/benchmarks/kv_cache_watermark.sh @@ -0,0 +1,248 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Reproducible demonstration of the KV cache watermark (`--watermark`) for +# reducing preemption thrashing. +# +# The watermark is the fraction of total KV cache blocks the scheduler keeps +# free when admitting a waiting/preempted request into the running queue. +# +# Why this workload triggers thrashing: +# Requests are admitted based on the KV cache they need *at admission time*. +# With `--scheduler-reserve-full-isl` (default) the input length is reserved up +# front, but the *output* length is unknown and unreserved. A decode-heavy +# workload (output >> input) at high concurrency therefore over-admits while +# requests are short, then runs out of KV cache as they all grow during decode +# -> the scheduler preempts (recompute) recently-admitted requests, re-prefills +# them later, and repeats. The watermark keeps a block of KV cache free so +# running requests can grow into it instead of triggering this churn. +# +# This script launches `vllm serve` under a deliberately KV-constrained config +# and a decode-heavy workload, sweeping the watermark across several values, and +# reports the preemption count (scraped from /metrics), throughput, and latency +# percentiles for each. It then plots the results. +# +# Default workload: concurrency 200, input ~300 tokens, output ~4000 tokens +# (+/- 20% variance), sized to run each config for ~5 minutes. +# +# Usage: +# benchmarks/kv_cache_watermark.sh +# MODEL=Qwen/Qwen2.5-14B-Instruct TP=2 benchmarks/kv_cache_watermark.sh +# +# Run inside the vLLM virtualenv (so `vllm` and `python` resolve to it). +set -euo pipefail + +# ---- Config (override via environment) ------------------------------------- +MODEL=${MODEL:-Qwen/Qwen2.5-7B-Instruct} +TP=${TP:-1} +PORT=${PORT:-8000} +URL="http://127.0.0.1:${PORT}" +# Constrain the KV cache to a *near-critical* size: large enough that the engine +# can run stably, but small enough that greedy over-admission tips it into +# preemption thrashing. (Independent of GPU size, so the demo is reproducible.) +# At the default workload this fits ~1.5x the mean concurrent KV demand. +KV_CACHE_MEMORY_GB=${KV_CACHE_MEMORY_GB:-16} +MAX_MODEL_LEN=${MAX_MODEL_LEN:-8192} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-256} +# Optional weight loader (e.g. fastsafetensors on the GCP cluster). +LOAD_FORMAT=${LOAD_FORMAT:-auto} +# Decode-heavy workload: moderate input, long output, with length variance. The +# long output means preempted requests have generated a lot before eviction, so +# resuming them re-prefills a long sequence (high recomputation cost). +INPUT_LEN=${INPUT_LEN:-1000} +OUTPUT_LEN=${OUTPUT_LEN:-5000} +RANGE_RATIO=${RANGE_RATIO:-0.2} +CONCURRENCY=${CONCURRENCY:-128} +# Enough prompts to keep each config saturated for ~5+ minutes. +NUM_PROMPTS=${NUM_PROMPTS:-450} +OUTDIR=${OUTDIR:-./watermark_bench_results} +# Watermark fractions compared. "label value" per line; value=0 disables it. +CONFIGS=${CONFIGS:-"off 0 +w0.02 0.02 +w0.05 0.05 +w0.10 0.10 +w0.15 0.15"} + +KV_CACHE_MEMORY_BYTES=$((KV_CACHE_MEMORY_GB * 1024 * 1024 * 1024)) +mkdir -p "$OUTDIR" + +SERVER_PID="" +cleanup() { [[ -n "$SERVER_PID" ]] && kill "$SERVER_PID" 2>/dev/null || true; } +trap cleanup EXIT + +scrape_preemptions() { + # Sum the vllm:num_preemptions_total counter across engines. + python - "${URL}/metrics" <<'PY' +import sys, urllib.request +total = 0.0 +try: + body = urllib.request.urlopen(sys.argv[1], timeout=10).read().decode("utf-8", "replace") + for line in body.splitlines(): + if line.startswith("vllm:num_preemptions_total"): + total += float(line.rsplit(" ", 1)[-1]) +except Exception as e: # noqa: BLE001 + print(f"scrape error: {e}", file=sys.stderr) +print(int(total)) +PY +} + +wait_for_server() { + for _ in $(seq 1 300); do + if curl -s "${URL}/health" >/dev/null 2>&1; then return 0; fi + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "ERROR: server process exited during startup" >&2; return 1 + fi + sleep 5 + done + echo "ERROR: server did not become ready" >&2; return 1 +} + +run_one() { + local label=$1 watermark=$2 + echo + echo "==================== watermark: ${label} (${watermark}) ====================" + vllm serve "$MODEL" \ + --tensor-parallel-size "$TP" \ + --load-format "$LOAD_FORMAT" \ + --kv-cache-memory-bytes "$KV_CACHE_MEMORY_BYTES" \ + --max-model-len "$MAX_MODEL_LEN" \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --no-enable-prefix-caching \ + --watermark "$watermark" \ + --port "$PORT" >"${OUTDIR}/serve_${label}.log" 2>&1 & + SERVER_PID=$! + wait_for_server + sleep 5 + + local pre post + pre=$(scrape_preemptions) + vllm bench serve \ + --backend vllm \ + --base-url "$URL" \ + --model "$MODEL" \ + --dataset-name random \ + --random-input-len "$INPUT_LEN" \ + --random-output-len "$OUTPUT_LEN" \ + --random-range-ratio "$RANGE_RATIO" \ + --ignore-eos \ + --num-prompts "$NUM_PROMPTS" \ + --max-concurrency "$CONCURRENCY" \ + --percentile-metrics "ttft,tpot,itl,e2el" \ + --metric-percentiles "50,90,99" \ + --save-result \ + --result-dir "$OUTDIR" \ + --result-filename "bench_${label}.json" + post=$(scrape_preemptions) + echo "${label} ${watermark} $((post - pre))" >>"${OUTDIR}/preemptions.txt" + + kill "$SERVER_PID" 2>/dev/null || true + for _ in $(seq 1 60); do curl -s "${URL}/health" >/dev/null 2>&1 || break; sleep 2; done + SERVER_PID="" + sleep 10 +} + +: >"${OUTDIR}/preemptions.txt" +while read -r label watermark; do + [[ -z "${label:-}" ]] && continue + run_one "$label" "$watermark" +done <<<"$CONFIGS" + +echo +echo "==================== summary ====================" +python - "$OUTDIR" <<'PY' +import json, os, sys +outdir = sys.argv[1] +pre = {} +order = [] +for line in open(os.path.join(outdir, "preemptions.txt")): + label, watermark, n = line.split() + pre[label] = (float(watermark), int(n)) + order.append(label) + +def g(d, *names): + for n in names: + if d.get(n) is not None: + return d[n] + return float("nan") + +cols = ["watermark", "frac", "preempt", "out_tok/s", "req/s", + "TTFT_p50", "TTFT_p99", "ITL_p99", "E2EL_p50"] +print(" ".join(f"{c:>10}" for c in cols)) +rows = [] +for label in order: + watermark, n = pre[label] + d = json.load(open(os.path.join(outdir, f"bench_{label}.json"))) + rows.append(dict( + label=label, watermark=watermark, preempt=n, + out_tok_s=g(d, "output_throughput"), + req_s=g(d, "request_throughput"), + ttft_p50=g(d, "p50_ttft_ms", "median_ttft_ms"), + ttft_p99=g(d, "p99_ttft_ms"), + itl_p99=g(d, "p99_itl_ms"), + e2el_p50=g(d, "p50_e2el_ms", "median_e2el_ms"), + )) + print(" ".join(f"{str(v):>10}" for v in [ + label, watermark, n, + f"{rows[-1]['out_tok_s']:.0f}", + f"{rows[-1]['req_s']:.3f}", + f"{rows[-1]['ttft_p50']/1000:.2f}", + f"{rows[-1]['ttft_p99']/1000:.2f}", + f"{rows[-1]['itl_p99']:.2f}", + f"{rows[-1]['e2el_p50']/1000:.1f}", + ])) +print("\n(TTFT/E2EL in seconds; ITL in ms. Lower preempt is better.)") + +# ---- Plot ------------------------------------------------------------------- +try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt +except Exception as e: # noqa: BLE001 + print(f"\n(skip plot: matplotlib unavailable: {e})") + sys.exit(0) + +x = [r["watermark"] for r in rows] +xt = [f"{r['watermark']:g}\n({r['label']})" for r in rows] +idx = list(range(len(rows))) + +fig, axes = plt.subplots(2, 2, figsize=(12, 8)) +fig.suptitle( + f"KV cache watermark sweep — {os.path.basename(os.path.abspath(outdir))}", + fontsize=12, +) + +ax = axes[0][0] +ax.bar(idx, [r["preempt"] for r in rows], color="tab:red") +ax.set_title("Preemptions (lower is better)") +ax.set_ylabel("preemptions") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[0][1] +ax.plot(idx, [r["out_tok_s"] for r in rows], "o-", color="tab:green") +ax.set_title("Output throughput (higher is better)") +ax.set_ylabel("tokens/s") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[1][0] +ax.plot(idx, [r["itl_p99"] for r in rows], "o-", color="tab:blue") +ax.set_title("Inter-token latency p99 (lower is better)") +ax.set_ylabel("ITL p99 (ms)") +ax.set_xlabel("watermark fraction") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[1][1] +ax.plot(idx, [r["ttft_p50"] / 1000 for r in rows], "o-", label="TTFT p50") +ax.plot(idx, [r["ttft_p99"] / 1000 for r in rows], "o-", label="TTFT p99") +ax.plot(idx, [r["e2el_p50"] / 1000 for r in rows], "o-", label="E2EL p50") +ax.set_title("Latency (lower is better)") +ax.set_ylabel("seconds") +ax.set_xlabel("watermark fraction") +ax.set_xticks(idx); ax.set_xticklabels(xt) +ax.legend() + +fig.tight_layout(rect=(0, 0, 1, 0.95)) +out_png = os.path.join(outdir, "watermark_results.png") +fig.savefig(out_png, dpi=120) +print(f"\nWrote plot: {out_png}") +PY diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 4d652beec8..1b789152e9 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -1849,6 +1849,8 @@ def create_scheduler_with_priority( enable_chunked_prefill=True, is_encoder_decoder=model_config.is_encoder_decoder, policy="priority", # Enable priority scheduling + # Ensure admission/preemption mechanics are deterministic + watermark=0.0, ) # Cache config, optionally force APC cache_config = CacheConfig( diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 7213a669c5..7f34250cb2 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -90,6 +90,8 @@ def create_scheduler( enable_chunked_prefill=enable_chunked_prefill, async_scheduling=async_scheduling, is_encoder_decoder=model_config.is_encoder_decoder, + # Ensure admission/preemption mechanics are deterministic + watermark=0.0, ) # Cache config, optionally force APC cache_config = CacheConfig( diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 9669bd1cc4..95f3ed48d4 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -143,6 +143,13 @@ class SchedulerConfig: checking the first chunk. Prevents over-admission and KV cache thrashing with chunked prefill.""" + watermark: float = Field(default=0.0, ge=0.0, lt=1.0) + """Fraction of total KV cache blocks to keep free (the watermark) when + admitting waiting or preempted requests into the running queue. This headroom + helps avoid frequent KV cache eviction and the resulting repeated preemption + of requests when GPU memory is scarce. Must be in the range [0.0, 1.0); 0.0 + (the default) disables the watermark.""" + async_scheduling: bool | None = None """If set to False, disable async scheduling. Async scheduling helps to avoid gaps in GPU utilization, leading to better latency and throughput. diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 0490cbc3e4..f0dade8371 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -600,6 +600,8 @@ class EngineArgs: scheduler_reserve_full_isl: bool = SchedulerConfig.scheduler_reserve_full_isl + watermark: float = SchedulerConfig.watermark + disable_hybrid_kv_cache_manager: bool | None = ( SchedulerConfig.disable_hybrid_kv_cache_manager ) @@ -1408,6 +1410,7 @@ class EngineArgs: "--scheduler-reserve-full-isl", **scheduler_kwargs["scheduler_reserve_full_isl"], ) + scheduler_group.add_argument("--watermark", **scheduler_kwargs["watermark"]) scheduler_group.add_argument( "--disable-hybrid-kv-cache-manager", **scheduler_kwargs["disable_hybrid_kv_cache_manager"], @@ -2045,6 +2048,7 @@ class EngineArgs: max_long_partial_prefills=self.max_long_partial_prefills, long_prefill_token_threshold=self.long_prefill_token_threshold, scheduler_reserve_full_isl=self.scheduler_reserve_full_isl, + watermark=self.watermark, disable_hybrid_kv_cache_manager=self.disable_hybrid_kv_cache_manager, async_scheduling=self.async_scheduling, stream_interval=self.stream_interval, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 9f0bfc5880..9af54e0a24 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -17,7 +17,7 @@ from vllm.v1.kv_cache_interface import ( get_kv_cache_spec_sliding_window, ) from vllm.v1.metrics.stats import PrefixCacheStats -from vllm.v1.request import Request +from vllm.v1.request import Request, RequestStatus logger = init_logger(__name__) @@ -122,6 +122,7 @@ class KVCacheManager: dcp_world_size: int = 1, pcp_world_size: int = 1, metrics_collector: KVCacheMetricsCollector | None = None, + watermark: float = 0.0, ) -> None: self.max_model_len = max_model_len # When unset, fall back to `max_model_len` so the recycling-aware cap @@ -155,6 +156,11 @@ class KVCacheManager: self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) self.block_pool = self.coordinator.block_pool self.kv_cache_config = kv_cache_config + + # Watermark: minimum number of KV cache blocks to keep free when + # admitting waiting/preempted requests, to avoid frequent preemptions. + assert watermark >= 0.0, "watermark must be non-negative" + self.watermark_blocks = int(watermark * kv_cache_config.num_blocks) self.kv_cache_event_metadata = tuple( ( get_kv_cache_spec_kind(group.kv_cache_spec).value, @@ -247,6 +253,7 @@ class KVCacheManager: num_encoder_tokens: int = 0, full_sequence_must_fit: bool = False, reserved_blocks: int = 0, + has_scheduled_reqs: bool = True, ) -> KVCacheBlocks | None: """Add slots for a request with new tokens to append. @@ -277,6 +284,8 @@ class KVCacheManager: made if it fits within (free blocks - reserved_blocks). Used to gate async KV-connector loads so their initial allocation cannot consume blocks an already in-flight (prefilling) sequence is relying on. + has_scheduled_reqs: Whether any requests are already scheduled to run + this step, controls whether watermark is applied. Blocks layout: ``` @@ -351,6 +360,15 @@ class KVCacheManager: self.max_model_len, ) + watermark_blocks = 0 + # The watermark is applied to waiting/preempted requests only, and only + # when there's at least one request already scheduled. + if has_scheduled_reqs and request.status in ( + RequestStatus.WAITING, + RequestStatus.PREEMPTED, + ): + watermark_blocks = self.watermark_blocks + if full_sequence_must_fit: # First check and fail if the full request sequence won't fit. full_num_tokens = min(request.num_tokens, self.max_model_len) @@ -364,7 +382,8 @@ class KVCacheManager: num_tokens_main_model=full_num_tokens, apply_admission_cap=True, ) - if num_blocks_to_allocate > self.block_pool.get_num_free_blocks(): + required_blocks = num_blocks_to_allocate + watermark_blocks + if required_blocks > self.block_pool.get_num_free_blocks(): return None num_tokens_main_model = total_computed_tokens + num_new_tokens @@ -392,8 +411,11 @@ class KVCacheManager: num_tokens_main_model=num_tokens_main_model, ) + # Keep `reserved_blocks` free for other in-flight sequences, and an + # additional watermark of headroom for waiting/preempted admissions. available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks - if num_blocks_to_allocate > available_blocks: + required_blocks = num_blocks_to_allocate + watermark_blocks + if required_blocks > available_blocks: # Cannot allocate new blocks return None diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 160cdb74f5..9a3a9ffa7d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -242,6 +242,7 @@ class Scheduler(SchedulerInterface): scheduler_block_size=self.block_size, hash_block_size=hash_block_size, metrics_collector=self.kv_metrics_collector, + watermark=self.scheduler_config.watermark, ) # Bind GPU block pool to the KV connector. This must happen after # kv_cache_manager is constructed so block_pool is available. @@ -826,6 +827,7 @@ class Scheduler(SchedulerInterface): num_encoder_tokens=num_encoder_tokens, full_sequence_must_fit=self.scheduler_reserve_full_isl, reserved_blocks=reserved_blocks, + has_scheduled_reqs=bool(self.running), ) if new_blocks is None: @@ -2198,12 +2200,8 @@ class Scheduler(SchedulerInterface): ) def _inflight_prefill_reserved_blocks(self) -> int: - """Blocks in-flight prefills still need to finish (their reservation). + """Num blocks in-flight prefills still need to finish (their reservation).""" - Sums remaining full-ISL blocks over `self._inflight_prefills` (running - prefills + in-progress async loads). The candidate async load isn't yet - in the set, so it's naturally excluded. - """ return sum( self._request_remaining_blocks(req) for req in self._inflight_prefills )