wehub-resource-sync
249 lines
8.7 KiB
Bash
Executable File
249 lines
8.7 KiB
Bash
Executable File
#!/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
|