26S30n1
#!/usr/bin/env bash
# ============================================================================
# STAGE 3a-8
#
# vLLM Prefix Caching A/B Benchmark
#
# Environment
# - vLLM : 0.29.0
# - GPU : NVIDIA B300 x 8
# - Model : Meta-Llama-3.1-70B-Instruct
# - TP : 8
#
# Goal
# Compare Prefix Caching ON vs OFF under the SAME workload.
#
# Workload
# Fixed prefix : 32,000 tokens
# Random suffix: 768 tokens
# Total input : 32,768 tokens
# Output : 256 tokens
# Requests : 100
# Concurrency : 8
#
# Important
# vLLM 0.29.0 does NOT support --enable-metrics.
# Therefore this script intentionally does NOT use --enable-metrics.
#
# vLLM 0.29.0 supports:
# --enable-prefix-caching
# --no-enable-prefix-caching
# --kv-cache-metrics
# --kv-cache-metrics-sample
#
# ============================================================================
set -Eeuo pipefail
# ----------------------------------------------------------------------------
# Configuration
# ----------------------------------------------------------------------------
MODEL="${MODEL:-/mnt/local-nvme-cache/models/Meta-Llama-3.1-70B-Instruct}"
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-8000}"
TP="${TP:-8}"
GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.90}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-65536}"
# Prefix caching workload
PREFIX_LEN="${PREFIX_LEN:-32000}"
RANDOM_INPUT_LEN="${RANDOM_INPUT_LEN:-768}"
OUTPUT_LEN="${OUTPUT_LEN:-256}"
NUM_PROMPTS="${NUM_PROMPTS:-100}"
MAX_CONCURRENCY="${MAX_CONCURRENCY:-8}"
REQUEST_RATE="${REQUEST_RATE:-inf}"
SEED="${SEED:-42}"
# Number of independent runs per condition.
REPEATS="${REPEATS:-2}"
# KV cache metrics sampling.
#
# 0.01 = 1% of KV blocks
# This minimizes benchmark overhead.
KV_CACHE_METRICS_SAMPLE="${KV_CACHE_METRICS_SAMPLE:-0.01}"
# ----------------------------------------------------------------------------
# Paths
# ----------------------------------------------------------------------------
ROOT_DIR="${ROOT_DIR:-$PWD/stage3a-8-prefix-cache}"
RESULT_DIR="$ROOT_DIR/results"
LOG_DIR="$ROOT_DIR/logs"
METRIC_DIR="$ROOT_DIR/metrics"
SERVER_INFO_DIR="$ROOT_DIR/server-info"
mkdir -p \
"$RESULT_DIR" \
"$LOG_DIR" \
"$METRIC_DIR" \
"$SERVER_INFO_DIR"
SERVER_LOG="$LOG_DIR/vllm-server.log"
PID_FILE="$ROOT_DIR/vllm-server.pid"
# ----------------------------------------------------------------------------
# Cleanup
# ----------------------------------------------------------------------------
SERVER_PID=""
cleanup() {
echo
echo "======================================================================"
echo "[CLEANUP]"
echo "======================================================================"
if [[ -n "${SERVER_PID:-}" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then
echo "[INFO] Stopping vLLM server PID=$SERVER_PID"
kill "$SERVER_PID" 2>/dev/null || true
for _ in {1..30}; do
if ! kill -0 "$SERVER_PID" 2>/dev/null; then
break
fi
sleep 1
done
if kill -0 "$SERVER_PID" 2>/dev/null; then
echo "[WARN] Server did not exit gracefully. Sending SIGKILL."
kill -9 "$SERVER_PID" 2>/dev/null || true
fi
fi
SERVER_PID=""
rm -f "$PID_FILE"
# Make sure no stale vLLM process remains.
pkill -f "vllm serve.*$PORT" 2>/dev/null || true
sleep 3
}
trap cleanup EXIT INT TERM
# ----------------------------------------------------------------------------
# Utility functions
# ----------------------------------------------------------------------------
timestamp() {
date '+%Y-%m-%d %H:%M:%S'
}
die() {
echo
echo "[ERROR] $*" >&2
exit 1
}
log() {
echo "[$(timestamp)] $*"
}
# ----------------------------------------------------------------------------
# Environment validation
# ----------------------------------------------------------------------------
echo
echo "======================================================================"
echo " STAGE 3a-8 : Prefix Caching A/B"
echo "======================================================================"
echo
echo "MODEL : $MODEL"
echo "vLLM expected : 0.29.0"
echo "Tensor Parallel : $TP"
echo "GPU Memory Utilization : $GPU_MEMORY_UTILIZATION"
echo "Max Model Length : $MAX_MODEL_LEN"
echo
echo "Fixed Prefix : $PREFIX_LEN"
echo "Random Suffix : $RANDOM_INPUT_LEN"
echo "Total Input : $((PREFIX_LEN + RANDOM_INPUT_LEN))"
echo "Output : $OUTPUT_LEN"
echo
echo "Requests : $NUM_PROMPTS"
echo "Concurrency : $MAX_CONCURRENCY"
echo "Request Rate : $REQUEST_RATE"
echo "Seed : $SEED"
echo "Repeats / condition : $REPEATS"
echo
echo "Results : $RESULT_DIR"
echo "Logs : $LOG_DIR"
echo "Metrics : $METRIC_DIR"
echo
command -v python >/dev/null 2>&1 \
|| die "python not found"
command -v vllm >/dev/null 2>&1 \
|| die "vllm command not found"
command -v curl >/dev/null 2>&1 \
|| die "curl not found"
[[ -d "$MODEL" ]] \
|| die "Model directory does not exist: $MODEL"
# ----------------------------------------------------------------------------
# Version check
# ----------------------------------------------------------------------------
VLLM_VERSION="$(
python - <<'PY'
import vllm
print(vllm.__version__)
PY
)"
echo "[INFO] Installed vLLM version: $VLLM_VERSION"
if [[ "$VLLM_VERSION" != "0.29.0" ]]; then
echo
echo "[WARN] This script was designed for vLLM 0.29.0."
echo "[WARN] Detected version: $VLLM_VERSION"
echo
fi
# ----------------------------------------------------------------------------
# GPU validation
# ----------------------------------------------------------------------------
if command -v nvidia-smi >/dev/null 2>&1; then
echo
echo "======================================================================"
echo "[GPU]"
echo "======================================================================"
nvidia-smi \
--query-gpu=index,name,memory.total,driver_version \
--format=csv
GPU_COUNT="$(nvidia-smi --query-gpu=count --format=csv,noheader | head -1)"
echo
echo "[INFO] GPU count: $GPU_COUNT"
[[ "$GPU_COUNT" -ge 8 ]] \
|| die "At least 8 GPUs are required."
else
echo "[WARN] nvidia-smi not found."
fi
# ----------------------------------------------------------------------------
# Save environment information
# ----------------------------------------------------------------------------
ENV_INFO="$SERVER_INFO_DIR/environment.txt"
{
echo "timestamp=$(timestamp)"
echo "vllm_version=$VLLM_VERSION"
echo "python=$(python --version 2>&1)"
echo
echo "[vllm]"
vllm --version 2>&1 || true
echo
echo "[pip]"
python -m pip show vllm 2>&1 || true
echo
echo "[nvidia-smi]"
nvidia-smi 2>&1 || true
} > "$ENV_INFO"
# ----------------------------------------------------------------------------
# Verify required CLI arguments
# ----------------------------------------------------------------------------
echo
echo "======================================================================"
echo "[CLI VALIDATION]"
echo "======================================================================"
if vllm serve --help 2>&1 | grep -q -- '--enable-prefix-caching'; then
echo "[OK] --enable-prefix-caching"
else
die "This vLLM installation does not expose --enable-prefix-caching"
fi
if vllm serve --help 2>&1 | grep -q -- '--kv-cache-metrics'; then
echo "[OK] --kv-cache-metrics"
else
echo "[WARN] --kv-cache-metrics is not available."
echo "[WARN] Continuing without KV cache metrics."
KV_METRICS_SUPPORTED="false"
fi
if vllm serve --help 2>&1 | grep -q -- '--enable-metrics'; then
echo "[INFO] --enable-metrics exists in this installation."
else
echo "[OK] --enable-metrics is not used."
fi
# Determine if kv-cache metrics are supported.
KV_METRICS_SUPPORTED="${KV_METRICS_SUPPORTED:-true}"
# ----------------------------------------------------------------------------
# Start vLLM server
# ----------------------------------------------------------------------------
start_server() {
local CACHE_MODE="$1"
local RUN_ID="$2"
local CACHE_ARG=""
local MODE_TAG=""
if [[ "$CACHE_MODE" == "ON" ]]; then
CACHE_ARG="--enable-prefix-caching"
MODE_TAG="prefix-cache-ON"
elif [[ "$CACHE_MODE" == "OFF" ]]; then
CACHE_ARG="--no-enable-prefix-caching"
MODE_TAG="prefix-cache-OFF"
else
die "Unknown cache mode: $CACHE_MODE"
fi
SERVER_LOG="$LOG_DIR/${MODE_TAG}-${RUN_ID}.log"
echo
echo "======================================================================"
echo "[SERVER START]"
echo "======================================================================"
echo "Mode : $CACHE_MODE"
echo "Run : $RUN_ID"
echo "Log : $SERVER_LOG"
echo
rm -f "$PID_FILE"
# ------------------------------------------------------------------------
# Base arguments
# ------------------------------------------------------------------------
local SERVER_ARGS=(
serve "$MODEL"
--tensor-parallel-size "$TP"
--gpu-memory-utilization "$GPU_MEMORY_UTILIZATION"
--max-model-len "$MAX_MODEL_LEN"
--host "$HOST"
--port "$PORT"
"$CACHE_ARG"
# Keep prefix hashing identical between ON/OFF experiments.
--prefix-caching-hash-algo sha256
)
# ------------------------------------------------------------------------
# KV cache residency metrics
#
# IMPORTANT:
# vLLM 0.29.0 does not have --enable-metrics.
#
# --kv-cache-metrics is the relevant option for KV cache residency
# metrics and requires normal log stats to remain enabled.
# ------------------------------------------------------------------------
if [[ "$KV_METRICS_SUPPORTED" == "true" ]]; then
SERVER_ARGS+=(
--kv-cache-metrics
--kv-cache-metrics-sample "$KV_CACHE_METRICS_SAMPLE"
)
fi
# ------------------------------------------------------------------------
# Start server
# ------------------------------------------------------------------------
echo "[COMMAND]"
printf 'vllm'
printf ' %q' "${SERVER_ARGS[@]}"
echo
echo
nohup vllm "${SERVER_ARGS[@]}" \
> "$SERVER_LOG" 2>&1 &
SERVER_PID=$!
echo "$SERVER_PID" > "$PID_FILE"
echo "[INFO] vLLM PID=$SERVER_PID"
# ------------------------------------------------------------------------
# Wait for server
# ------------------------------------------------------------------------
echo "[INFO] Waiting for vLLM server..."
local READY=0
for i in {1..180}; do
if ! kill -0 "$SERVER_PID" 2>/dev/null; then
echo
echo "[ERROR] vLLM server exited unexpectedly."
echo
tail -100 "$SERVER_LOG" || true
return 1
fi
if curl -sf \
"http://${HOST}:${PORT}/health" \
>/dev/null 2>&1; then
READY=1
break
fi
if (( i % 10 == 0 )); then
echo "[INFO] Waiting... ${i}s"
fi
sleep 1
done
if [[ "$READY" != "1" ]]; then
echo
echo "[ERROR] vLLM server did not become ready."
echo
tail -150 "$SERVER_LOG" || true
return 1
fi
echo
echo "[OK] vLLM server is READY."
echo
# ------------------------------------------------------------------------
# Verify OpenAI endpoint
# ------------------------------------------------------------------------
curl -sf \
"http://${HOST}:${PORT}/v1/models" \
> "$SERVER_INFO_DIR/${MODE_TAG}-${RUN_ID}-models.json" \
|| true
# ------------------------------------------------------------------------
# Capture metrics if endpoint exists
# ------------------------------------------------------------------------
capture_metrics \
"$MODE_TAG" \
"$RUN_ID" \
"server-start"
return 0
}
# ----------------------------------------------------------------------------
# Metrics collector
# ----------------------------------------------------------------------------
capture_metrics() {
local MODE="$1"
local RUN_ID="$2"
local LABEL="$3"
local OUT="$METRIC_DIR/${MODE}-${RUN_ID}-${LABEL}.txt"
if curl -sf \
"http://${HOST}:${PORT}/metrics" \
> "$OUT" 2>/dev/null; then
echo "[OK] Prometheus metrics captured: $OUT"
# Extract potentially useful cache-related metrics.
grep -Ei \
'prefix|cache|prompt_tokens|kv' \
"$OUT" \
> "${OUT}.filtered" \
2>/dev/null || true
else
echo "[INFO] /metrics endpoint unavailable."
echo "[INFO] Continuing without Prometheus metrics."
rm -f "$OUT"
fi
}
# ----------------------------------------------------------------------------
# Stop vLLM
# ----------------------------------------------------------------------------
stop_server() {
echo
echo "======================================================================"
echo "[SERVER STOP]"
echo "======================================================================"
if [[ -n "${SERVER_PID:-}" ]] &&
kill -0 "$SERVER_PID" 2>/dev/null; then
echo "[INFO] Stopping PID=$SERVER_PID"
kill "$SERVER_PID" 2>/dev/null || true
for _ in {1..30}; do
if ! kill -0 "$SERVER_PID" 2>/dev/null; then
break
fi
sleep 1
done
if kill -0 "$SERVER_PID" 2>/dev/null; then
echo "[WARN] Sending SIGKILL."
kill -9 "$SERVER_PID" 2>/dev/null || true
fi
fi
SERVER_PID=""
rm -f "$PID_FILE"
# Make absolutely sure no server remains.
pkill -f "vllm serve.*$PORT" 2>/dev/null || true
echo "[INFO] Waiting for GPU memory cleanup..."
sleep 8
# Show GPU status after shutdown.
if command -v nvidia-smi >/dev/null 2>&1; then
nvidia-smi \
--query-compute-apps=pid,name,used_memory \
--format=csv,noheader 2>/dev/null || true
fi
}
# ----------------------------------------------------------------------------
# Run benchmark
# ----------------------------------------------------------------------------
run_benchmark() {
local CACHE_MODE="$1"
local RUN_ID="$2"
local MODE_TAG="prefix-cache-${CACHE_MODE}"
local RESULT_FILE="$RESULT_DIR/${MODE_TAG}-${RUN_ID}.json"
local DETAILED_FILE="$RESULT_DIR/${MODE_TAG}-${RUN_ID}-detailed.json"
echo
echo "======================================================================"
echo "[BENCHMARK]"
echo "======================================================================"
echo "Prefix Cache : $CACHE_MODE"
echo "Run : $RUN_ID"
echo
echo "Input:"
echo " Prefix : ${PREFIX_LEN}"
echo " Random : ${RANDOM_INPUT_LEN}"
echo " Total : $((PREFIX_LEN + RANDOM_INPUT_LEN))"
echo
echo "Output : ${OUTPUT_LEN}"
echo "Requests : ${NUM_PROMPTS}"
echo "Concurrency : ${MAX_CONCURRENCY}"
echo
# ------------------------------------------------------------------------
# Benchmark command
#
# random-prefix-len = 32000
# random-input-len = 768
#
# => total input = 32768
# ------------------------------------------------------------------------
local BENCH_ARGS=(
bench serve
--backend openai
--model "$MODEL"
--dataset-name random
--random-input-len "$RANDOM_INPUT_LEN"
--random-output-len "$OUTPUT_LEN"
--random-prefix-len "$PREFIX_LEN"
--random-range-ratio 0
--num-prompts "$NUM_PROMPTS"
--max-concurrency "$MAX_CONCURRENCY"
--request-rate "$REQUEST_RATE"
--seed "$SEED"
--percentile-metrics ttft,tpot,itl
--metric-percentiles 50,95,99
--save-result
--result-filename "$RESULT_FILE"
--save-detailed
--detailed-result-filename "$DETAILED_FILE"
)
echo "[COMMAND]"
printf 'vllm'
printf ' %q' "${BENCH_ARGS[@]}"
echo
echo
# ------------------------------------------------------------------------
# Run
# ------------------------------------------------------------------------
vllm "${BENCH_ARGS[@]}"
local RC=$?
echo
echo "[INFO] Benchmark exit code: $RC"
# Capture metrics after benchmark.
capture_metrics \
"$MODE_TAG" \
"$RUN_ID" \
"after-benchmark"
return "$RC"
}
# ----------------------------------------------------------------------------
# Extract benchmark summary
# ----------------------------------------------------------------------------
print_result_summary() {
local FILE="$1"
echo
echo "------------------------------------------------------------------"
echo "Result: $FILE"
echo "------------------------------------------------------------------"
if [[ ! -f "$FILE" ]]; then
echo "[WARN] Result file not found."
return
fi
python - "$FILE" <<'PY'
import json
import sys
path = sys.argv[1]
try:
with open(path) as f:
d = json.load(f)
print(json.dumps(d, indent=2))
except Exception as e:
print(f"[WARN] Could not parse result: {e}")
PY
}
# ----------------------------------------------------------------------------
# Execute A/B tests
# ----------------------------------------------------------------------------
run_condition() {
local MODE="$1"
echo
echo
echo "######################################################################"
echo "# CONDITION: Prefix Cache $MODE"
echo "######################################################################"
for ((RUN=1; RUN<=REPEATS; RUN++)); do
echo
echo "##################################################################"
echo "# $MODE RUN $RUN / $REPEATS"
echo "##################################################################"
# ------------------------------------------------------------
# Start fresh server for EVERY run.
#
# This is important:
# A/B must not inherit KV cache state from another run.
# ------------------------------------------------------------
start_server "$MODE" "$RUN" \
|| die "Failed to start vLLM server for $MODE run $RUN"
# ------------------------------------------------------------
# Give server a short stabilization period.
# ------------------------------------------------------------
sleep 5
# ------------------------------------------------------------
# Benchmark
# ------------------------------------------------------------
run_benchmark "$MODE" "$RUN" \
|| {
echo
echo "[ERROR] Benchmark failed."
echo "[ERROR] Server log:"
tail -100 "$SERVER_LOG" || true
stop_server
return 1
}
# ------------------------------------------------------------
# Print result
# ------------------------------------------------------------
print_result_summary \
"$RESULT_DIR/prefix-cache-${MODE}-${RUN}.json"
# ------------------------------------------------------------
# Stop server
# ------------------------------------------------------------
stop_server
# ------------------------------------------------------------
# Cooling / GPU cleanup
# ------------------------------------------------------------
echo "[INFO] Waiting before next condition..."
sleep 10
done
}
# ----------------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------------
echo
echo "======================================================================"
echo " START A/B TEST"
echo "======================================================================"
echo
START_TIME="$(date +%s)"
# ------------------------------------------------------------------------
# Prefix Cache OFF
# ------------------------------------------------------------------------
run_condition "OFF"
# ------------------------------------------------------------------------
# Prefix Cache ON
# ------------------------------------------------------------------------
run_condition "ON"
END_TIME="$(date +%s)"
ELAPSED=$((END_TIME - START_TIME))
# ----------------------------------------------------------------------------
# Final summary
# ----------------------------------------------------------------------------
echo
echo
echo "======================================================================"
echo " A/B TEST COMPLETE"
echo "======================================================================"
echo
echo "Elapsed: ${ELAPSED}s"
echo
echo "Results:"
find "$RESULT_DIR" \
-maxdepth 1 \
-type f \
-name '*.json' \
-printf ' %f\n' \
2>/dev/null | sort || true
echo
echo "Logs:"
find "$LOG_DIR" \
-maxdepth 1 \
-type f \
-printf ' %f\n' \
2>/dev/null | sort || true
echo
echo "Metrics:"
find "$METRIC_DIR" \
-maxdepth 1 \
-type f \
-printf ' %f\n' \
2>/dev/null | sort || true
echo
echo "======================================================================"
echo " IMPORTANT INTERPRETATION"
echo "======================================================================"
echo
echo "This test compares:"
echo
echo " Prefix Cache OFF"
echo " vs"
echo " Prefix Cache ON"
echo
echo "with:"
echo
echo " Fixed prefix = ${PREFIX_LEN} tokens"
echo " Random suffix = ${RANDOM_INPUT_LEN} tokens"
echo " Total input = $((PREFIX_LEN + RANDOM_INPUT_LEN)) tokens"
echo " Output = ${OUTPUT_LEN} tokens"
echo " Requests = ${NUM_PROMPTS}"
echo " Concurrency = ${MAX_CONCURRENCY}"
echo
echo "Primary metric:"
echo " TTFT P50"
echo
echo "Secondary:"
echo " TTFT P95"
echo " TTFT P99"
echo " TPOT"
echo " ITL"
echo
echo "======================================================================"
echo " DONE"
echo "======================================================================"
engineer