A high-throughput and memory-efficient inference and serving engine for LLMs
[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>
Nick Hill committed
Jun 11, 2026 at 08:27 UTC
4085ff7cb43d03bbfd05707238ea58a1561f87c2
7 files changed
+291
-8
benchmarks/kv_cache_watermark.sh
new
+248
@@ -0,0 +1,248 @@
1
+#!/bin/bash
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+#
5
+# Reproducible demonstration of the KV cache watermark (`--watermark`) for
6
+# reducing preemption thrashing.
7
+#
8
+# The watermark is the fraction of total KV cache blocks the scheduler keeps
9
+# free when admitting a waiting/preempted request into the running queue.
10
+#
11
+# Why this workload triggers thrashing:
12
+# Requests are admitted based on the KV cache they need *at admission time*.
13
+# With `--scheduler-reserve-full-isl` (default) the input length is reserved up
14
+# front, but the *output* length is unknown and unreserved. A decode-heavy
15
+# workload (output >> input) at high concurrency therefore over-admits while
16
+# requests are short, then runs out of KV cache as they all grow during decode
17
+# -> the scheduler preempts (recompute) recently-admitted requests, re-prefills
18
+# them later, and repeats. The watermark keeps a block of KV cache free so
19
+# running requests can grow into it instead of triggering this churn.
20
+#
21
+# This script launches `vllm serve` under a deliberately KV-constrained config
22
+# and a decode-heavy workload, sweeping the watermark across several values, and
23
+# reports the preemption count (scraped from /metrics), throughput, and latency
24
+# percentiles for each. It then plots the results.
25
+#
26
+# Default workload: concurrency 200, input ~300 tokens, output ~4000 tokens
27
+# (+/- 20% variance), sized to run each config for ~5 minutes.
28
+#
29
+# Usage:
30
+# benchmarks/kv_cache_watermark.sh
31
+# MODEL=Qwen/Qwen2.5-14B-Instruct TP=2 benchmarks/kv_cache_watermark.sh
32
+#
33
+# Run inside the vLLM virtualenv (so `vllm` and `python` resolve to it).
34
+set -euo pipefail
35
+
36
+# ---- Config (override via environment) -------------------------------------
37
+MODEL=${MODEL:-Qwen/Qwen2.5-7B-Instruct}
38
+TP=${TP:-1}
39
+PORT=${PORT:-8000}
40
+URL="http://127.0.0.1:${PORT}"
41
+# Constrain the KV cache to a *near-critical* size: large enough that the engine
42
+# can run stably, but small enough that greedy over-admission tips it into
43
+# preemption thrashing. (Independent of GPU size, so the demo is reproducible.)
44
+# At the default workload this fits ~1.5x the mean concurrent KV demand.
45
+KV_CACHE_MEMORY_GB=${KV_CACHE_MEMORY_GB:-16}
46
+MAX_MODEL_LEN=${MAX_MODEL_LEN:-8192}
47
+MAX_NUM_SEQS=${MAX_NUM_SEQS:-256}
48
+# Optional weight loader (e.g. fastsafetensors on the GCP cluster).
49
+LOAD_FORMAT=${LOAD_FORMAT:-auto}
50
+# Decode-heavy workload: moderate input, long output, with length variance. The
51
+# long output means preempted requests have generated a lot before eviction, so
52
+# resuming them re-prefills a long sequence (high recomputation cost).
53
+INPUT_LEN=${INPUT_LEN:-1000}
54
+OUTPUT_LEN=${OUTPUT_LEN:-5000}
55
+RANGE_RATIO=${RANGE_RATIO:-0.2}
56
+CONCURRENCY=${CONCURRENCY:-128}
57
+# Enough prompts to keep each config saturated for ~5+ minutes.
58
+NUM_PROMPTS=${NUM_PROMPTS:-450}
59
+OUTDIR=${OUTDIR:-./watermark_bench_results}
60
+# Watermark fractions compared. "label value" per line; value=0 disables it.
61
+CONFIGS=${CONFIGS:-"off 0
62
+w0.02 0.02
63
+w0.05 0.05
64
+w0.10 0.10
65
+w0.15 0.15"}
66
+
67
+KV_CACHE_MEMORY_BYTES=$((KV_CACHE_MEMORY_GB * 1024 * 1024 * 1024))
68
+mkdir -p "$OUTDIR"
69
+
70
+SERVER_PID=""
71
+cleanup() { [[ -n "$SERVER_PID" ]] && kill "$SERVER_PID" 2>/dev/null || true; }
72
+trap cleanup EXIT
73
+
74
+scrape_preemptions() {
75
+ # Sum the vllm:num_preemptions_total counter across engines.
76
+ python - "${URL}/metrics" <<'PY'
77
+import sys, urllib.request
78
+total = 0.0
79
+try:
80
+ body = urllib.request.urlopen(sys.argv[1], timeout=10).read().decode("utf-8", "replace")
81
+ for line in body.splitlines():
82
+ if line.startswith("vllm:num_preemptions_total"):
83
+ total += float(line.rsplit(" ", 1)[-1])
84
+except Exception as e: # noqa: BLE001
85
+ print(f"scrape error: {e}", file=sys.stderr)
86
+print(int(total))
87
+PY
88
+}
89
+
90
+wait_for_server() {
91
+ for _ in $(seq 1 300); do
92
+ if curl -s "${URL}/health" >/dev/null 2>&1; then return 0; fi
93
+ if ! kill -0 "$SERVER_PID" 2>/dev/null; then
94
+ echo "ERROR: server process exited during startup" >&2; return 1
95
+ fi
96
+ sleep 5
97
+ done
98
+ echo "ERROR: server did not become ready" >&2; return 1
99
+}
100
+
101
+run_one() {
102
+ local label=$1 watermark=$2
103
+ echo
104
+ echo "==================== watermark: ${label} (${watermark}) ===================="
105
+ vllm serve "$MODEL" \
106
+ --tensor-parallel-size "$TP" \
107
+ --load-format "$LOAD_FORMAT" \
108
+ --kv-cache-memory-bytes "$KV_CACHE_MEMORY_BYTES" \
109
+ --max-model-len "$MAX_MODEL_LEN" \
110
+ --max-num-seqs "$MAX_NUM_SEQS" \
111
+ --no-enable-prefix-caching \
112
+ --watermark "$watermark" \
113
+ --port "$PORT" >"${OUTDIR}/serve_${label}.log" 2>&1 &
114
+ SERVER_PID=$!
115
+ wait_for_server
116
+ sleep 5
117
+
118
+ local pre post
119
+ pre=$(scrape_preemptions)
120
+ vllm bench serve \
121
+ --backend vllm \
122
+ --base-url "$URL" \
123
+ --model "$MODEL" \
124
+ --dataset-name random \
125
+ --random-input-len "$INPUT_LEN" \
126
+ --random-output-len "$OUTPUT_LEN" \
127
+ --random-range-ratio "$RANGE_RATIO" \
128
+ --ignore-eos \
129
+ --num-prompts "$NUM_PROMPTS" \
130
+ --max-concurrency "$CONCURRENCY" \
131
+ --percentile-metrics "ttft,tpot,itl,e2el" \
132
+ --metric-percentiles "50,90,99" \
133
+ --save-result \
134
+ --result-dir "$OUTDIR" \
135
+ --result-filename "bench_${label}.json"
136
+ post=$(scrape_preemptions)
137
+ echo "${label} ${watermark} $((post - pre))" >>"${OUTDIR}/preemptions.txt"
138
+
139
+ kill "$SERVER_PID" 2>/dev/null || true
140
+ for _ in $(seq 1 60); do curl -s "${URL}/health" >/dev/null 2>&1 || break; sleep 2; done
141
+ SERVER_PID=""
142
+ sleep 10
143
+}
144
+
145
+: >"${OUTDIR}/preemptions.txt"
146
+while read -r label watermark; do
147
+ [[ -z "${label:-}" ]] && continue
148
+ run_one "$label" "$watermark"
149
+done <<<"$CONFIGS"
150
+
151
+echo
152
+echo "==================== summary ===================="
153
+python - "$OUTDIR" <<'PY'
154
+import json, os, sys
155
+outdir = sys.argv[1]
156
+pre = {}
157
+order = []
158
+for line in open(os.path.join(outdir, "preemptions.txt")):
159
+ label, watermark, n = line.split()
160
+ pre[label] = (float(watermark), int(n))
161
+ order.append(label)
162
+
163
+def g(d, *names):
164
+ for n in names:
165
+ if d.get(n) is not None:
166
+ return d[n]
167
+ return float("nan")
168
+
169
+cols = ["watermark", "frac", "preempt", "out_tok/s", "req/s",
170
+ "TTFT_p50", "TTFT_p99", "ITL_p99", "E2EL_p50"]
171
+print(" ".join(f"{c:>10}" for c in cols))
172
+rows = []
173
+for label in order:
174
+ watermark, n = pre[label]
175
+ d = json.load(open(os.path.join(outdir, f"bench_{label}.json")))
176
+ rows.append(dict(
177
+ label=label, watermark=watermark, preempt=n,
178
+ out_tok_s=g(d, "output_throughput"),
179
+ req_s=g(d, "request_throughput"),
180
+ ttft_p50=g(d, "p50_ttft_ms", "median_ttft_ms"),
181
+ ttft_p99=g(d, "p99_ttft_ms"),
182
+ itl_p99=g(d, "p99_itl_ms"),
183
+ e2el_p50=g(d, "p50_e2el_ms", "median_e2el_ms"),
184
+ ))
185
+ print(" ".join(f"{str(v):>10}" for v in [
186
+ label, watermark, n,
187
+ f"{rows[-1]['out_tok_s']:.0f}",
188
+ f"{rows[-1]['req_s']:.3f}",
189
+ f"{rows[-1]['ttft_p50']/1000:.2f}",
190
+ f"{rows[-1]['ttft_p99']/1000:.2f}",
191
+ f"{rows[-1]['itl_p99']:.2f}",
192
+ f"{rows[-1]['e2el_p50']/1000:.1f}",
193
+ ]))
194
+print("\n(TTFT/E2EL in seconds; ITL in ms. Lower preempt is better.)")
195
+
196
+# ---- Plot -------------------------------------------------------------------
197
+try:
198
+ import matplotlib
199
+ matplotlib.use("Agg")
200
+ import matplotlib.pyplot as plt
201
+except Exception as e: # noqa: BLE001
202
+ print(f"\n(skip plot: matplotlib unavailable: {e})")
203
+ sys.exit(0)
204
+
205
+x = [r["watermark"] for r in rows]
206
+xt = [f"{r['watermark']:g}\n({r['label']})" for r in rows]
207
+idx = list(range(len(rows)))
208
+
209
+fig, axes = plt.subplots(2, 2, figsize=(12, 8))
210
+fig.suptitle(
211
+ f"KV cache watermark sweep — {os.path.basename(os.path.abspath(outdir))}",
212
+ fontsize=12,
213
+)
214
+
215
+ax = axes[0][0]
216
+ax.bar(idx, [r["preempt"] for r in rows], color="tab:red")
217
+ax.set_title("Preemptions (lower is better)")
218
+ax.set_ylabel("preemptions")
219
+ax.set_xticks(idx); ax.set_xticklabels(xt)
220
+
221
+ax = axes[0][1]
222
+ax.plot(idx, [r["out_tok_s"] for r in rows], "o-", color="tab:green")
223
+ax.set_title("Output throughput (higher is better)")
224
+ax.set_ylabel("tokens/s")
225
+ax.set_xticks(idx); ax.set_xticklabels(xt)
226
+
227
+ax = axes[1][0]
228
+ax.plot(idx, [r["itl_p99"] for r in rows], "o-", color="tab:blue")
229
+ax.set_title("Inter-token latency p99 (lower is better)")
230
+ax.set_ylabel("ITL p99 (ms)")
231
+ax.set_xlabel("watermark fraction")
232
+ax.set_xticks(idx); ax.set_xticklabels(xt)
233
+
234
+ax = axes[1][1]
235
+ax.plot(idx, [r["ttft_p50"] / 1000 for r in rows], "o-", label="TTFT p50")
236
+ax.plot(idx, [r["ttft_p99"] / 1000 for r in rows], "o-", label="TTFT p99")
237
+ax.plot(idx, [r["e2el_p50"] / 1000 for r in rows], "o-", label="E2EL p50")
238
+ax.set_title("Latency (lower is better)")
239
+ax.set_ylabel("seconds")
240
+ax.set_xlabel("watermark fraction")
241
+ax.set_xticks(idx); ax.set_xticklabels(xt)
242
+ax.legend()
243
+
244
+fig.tight_layout(rect=(0, 0, 1, 0.95))
245
+out_png = os.path.join(outdir, "watermark_results.png")
246
+fig.savefig(out_png, dpi=120)
247
+print(f"\nWrote plot: {out_png}")
248
+PY
tests/v1/core/test_scheduler.py
+2
@@ -1849,6 +1849,8 @@ def create_scheduler_with_priority(
1849
enable_chunked_prefill=True,
1850
is_encoder_decoder=model_config.is_encoder_decoder,
1851
policy="priority", # Enable priority scheduling
1852
+ # Ensure admission/preemption mechanics are deterministic
1853
+ watermark=0.0,
1854
)
1855
# Cache config, optionally force APC
1856
cache_config = CacheConfig(
tests/v1/core/utils.py
+2
@@ -90,6 +90,8 @@ def create_scheduler(
90
enable_chunked_prefill=enable_chunked_prefill,
91
async_scheduling=async_scheduling,
92
is_encoder_decoder=model_config.is_encoder_decoder,
93
+ # Ensure admission/preemption mechanics are deterministic
94
+ watermark=0.0,
95
)
96
# Cache config, optionally force APC
97
cache_config = CacheConfig(
vllm/config/scheduler.py
+7
@@ -143,6 +143,13 @@ class SchedulerConfig:
143
checking the first chunk. Prevents over-admission and KV cache thrashing
144
with chunked prefill."""
145
146
+ watermark: float = Field(default=0.0, ge=0.0, lt=1.0)
147
+ """Fraction of total KV cache blocks to keep free (the watermark) when
148
+ admitting waiting or preempted requests into the running queue. This headroom
149
+ helps avoid frequent KV cache eviction and the resulting repeated preemption
150
+ of requests when GPU memory is scarce. Must be in the range [0.0, 1.0); 0.0
151
+ (the default) disables the watermark."""
152
+
153
async_scheduling: bool | None = None
154
"""If set to False, disable async scheduling. Async scheduling helps to
155
avoid gaps in GPU utilization, leading to better latency and throughput.
vllm/engine/arg_utils.py
+4
@@ -600,6 +600,8 @@ class EngineArgs:
600
601
scheduler_reserve_full_isl: bool = SchedulerConfig.scheduler_reserve_full_isl
602
603
+ watermark: float = SchedulerConfig.watermark
604
+
605
disable_hybrid_kv_cache_manager: bool | None = (
606
SchedulerConfig.disable_hybrid_kv_cache_manager
607
)
@@ -1408,6 +1410,7 @@ class EngineArgs:
1410
"--scheduler-reserve-full-isl",
1411
**scheduler_kwargs["scheduler_reserve_full_isl"],
1412
)
1413
+ scheduler_group.add_argument("--watermark", **scheduler_kwargs["watermark"])
1414
scheduler_group.add_argument(
1415
"--disable-hybrid-kv-cache-manager",
1416
**scheduler_kwargs["disable_hybrid_kv_cache_manager"],
@@ -2045,6 +2048,7 @@ class EngineArgs:
2048
max_long_partial_prefills=self.max_long_partial_prefills,
2049
long_prefill_token_threshold=self.long_prefill_token_threshold,
2050
scheduler_reserve_full_isl=self.scheduler_reserve_full_isl,
2051
+ watermark=self.watermark,
2052
disable_hybrid_kv_cache_manager=self.disable_hybrid_kv_cache_manager,
2053
async_scheduling=self.async_scheduling,
2054
stream_interval=self.stream_interval,
vllm/v1/core/kv_cache_manager.py
+25
-3
@@ -17,7 +17,7 @@ from vllm.v1.kv_cache_interface import (
17
get_kv_cache_spec_sliding_window,
18
)
19
from vllm.v1.metrics.stats import PrefixCacheStats
20
-from vllm.v1.request import Request
20
+from vllm.v1.request import Request, RequestStatus
21
22
logger = init_logger(__name__)
23
@@ -122,6 +122,7 @@ class KVCacheManager:
122
dcp_world_size: int = 1,
123
pcp_world_size: int = 1,
124
metrics_collector: KVCacheMetricsCollector | None = None,
125
+ watermark: float = 0.0,
126
) -> None:
127
self.max_model_len = max_model_len
128
# When unset, fall back to `max_model_len` so the recycling-aware cap
@@ -155,6 +156,11 @@ class KVCacheManager:
156
self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
157
self.block_pool = self.coordinator.block_pool
158
self.kv_cache_config = kv_cache_config
159
+
160
+ # Watermark: minimum number of KV cache blocks to keep free when
161
+ # admitting waiting/preempted requests, to avoid frequent preemptions.
162
+ assert watermark >= 0.0, "watermark must be non-negative"
163
+ self.watermark_blocks = int(watermark * kv_cache_config.num_blocks)
164
self.kv_cache_event_metadata = tuple(
165
(
166
get_kv_cache_spec_kind(group.kv_cache_spec).value,
@@ -247,6 +253,7 @@ class KVCacheManager:
253
num_encoder_tokens: int = 0,
254
full_sequence_must_fit: bool = False,
255
reserved_blocks: int = 0,
256
+ has_scheduled_reqs: bool = True,
257
) -> KVCacheBlocks | None:
258
"""Add slots for a request with new tokens to append.
259
@@ -277,6 +284,8 @@ class KVCacheManager:
284
made if it fits within (free blocks - reserved_blocks). Used to gate
285
async KV-connector loads so their initial allocation cannot consume
286
blocks an already in-flight (prefilling) sequence is relying on.
287
+ has_scheduled_reqs: Whether any requests are already scheduled to run
288
+ this step, controls whether watermark is applied.
289
290
Blocks layout:
291
```
@@ -351,6 +360,15 @@ class KVCacheManager:
360
self.max_model_len,
361
)
362
363
+ watermark_blocks = 0
364
+ # The watermark is applied to waiting/preempted requests only, and only
365
+ # when there's at least one request already scheduled.
366
+ if has_scheduled_reqs and request.status in (
367
+ RequestStatus.WAITING,
368
+ RequestStatus.PREEMPTED,
369
+ ):
370
+ watermark_blocks = self.watermark_blocks
371
+
372
if full_sequence_must_fit:
373
# First check and fail if the full request sequence won't fit.
374
full_num_tokens = min(request.num_tokens, self.max_model_len)
@@ -364,7 +382,8 @@ class KVCacheManager:
382
num_tokens_main_model=full_num_tokens,
383
apply_admission_cap=True,
384
)
367
- if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():
385
+ required_blocks = num_blocks_to_allocate + watermark_blocks
386
+ if required_blocks > self.block_pool.get_num_free_blocks():
387
return None
388
389
num_tokens_main_model = total_computed_tokens + num_new_tokens
@@ -392,8 +411,11 @@ class KVCacheManager:
411
num_tokens_main_model=num_tokens_main_model,
412
)
413
414
+ # Keep `reserved_blocks` free for other in-flight sequences, and an
415
+ # additional watermark of headroom for waiting/preempted admissions.
416
available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks
396
- if num_blocks_to_allocate > available_blocks:
417
+ required_blocks = num_blocks_to_allocate + watermark_blocks
418
+ if required_blocks > available_blocks:
419
# Cannot allocate new blocks
420
return None
421
vllm/v1/core/sched/scheduler.py
+3
-5
@@ -242,6 +242,7 @@ class Scheduler(SchedulerInterface):
242
scheduler_block_size=self.block_size,
243
hash_block_size=hash_block_size,
244
metrics_collector=self.kv_metrics_collector,
245
+ watermark=self.scheduler_config.watermark,
246
)
247
# Bind GPU block pool to the KV connector. This must happen after
248
# kv_cache_manager is constructed so block_pool is available.
@@ -826,6 +827,7 @@ class Scheduler(SchedulerInterface):
827
num_encoder_tokens=num_encoder_tokens,
828
full_sequence_must_fit=self.scheduler_reserve_full_isl,
829
reserved_blocks=reserved_blocks,
830
+ has_scheduled_reqs=bool(self.running),
831
)
832
833
if new_blocks is None:
@@ -2198,12 +2200,8 @@ class Scheduler(SchedulerInterface):
2200
)
2201
2202
def _inflight_prefill_reserved_blocks(self) -> int:
2201
- """Blocks in-flight prefills still need to finish (their reservation).
2203
+ """Num blocks in-flight prefills still need to finish (their reservation)."""
2204
2203
- Sums remaining full-ISL blocks over `self._inflight_prefills` (running
2204
- prefills + in-progress async loads). The candidate async load isn't yet
2205
- in the set, so it's naturally excluded.
2206
- """
2205
return sum(
2206
self._request_remaining_blocks(req) for req in self._inflight_prefills
2207
)