[Core] Add kvcache watermark to reduce preemptions (#44594)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: njhill

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

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(

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(

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.

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 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             "--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 @@ def create_engine_config(
             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,

vllm/v1/core/kv_cache_manager.py

@@ -17,7 +17,7 @@
     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 @@ def __init__(
         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 @@ def __init__(
         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 @@ def allocate_slots(
         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 @@ def allocate_slots(
                 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 @@ def allocate_slots(
             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 @@ def allocate_slots(
                 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 @@ def allocate_slots(
             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