@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 248 lines 8.66 KB
Raw
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