S124 · LLM Serving — KV Cache, Batching, Speculative Decoding

#!/usr/bin/env bash
# serve.sh — launch vLLM as an OpenAI-compatible server.
set -euo pipefail
 
pip install "vllm==0.5.4" "openai>=1.30"
 
python -m vllm.entrypoints.openai.api_server \
    --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
    --max-model-len 2048 \
    --gpu-memory-utilization 0.9 \
    --enable-prefix-caching \
    --port 8000 &
 
# wait for the server to be reachable
for i in {1..60}; do
    if curl -sf http://localhost:8000/v1/models >/dev/null; then
        echo "vLLM ready after ${i}s"; break
    fi
    sleep 1
done

#!/usr/bin/env python3
# bench.py — hit the vLLM server at increasing concurrency and print throughput.
import asyncio
import time
from statistics import mean, quantiles
 
from openai import AsyncOpenAI
 
MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
client = AsyncOpenAI(base_url="http://localhost:8000/v1", api_key="dummy")
 
SYSTEM = "You are a poet. Answer in exactly four short rhyming lines."
 
async def one_request(i: int) -> tuple[float, int, float]:
    t0 = time.time()
    resp = await client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": f"Write about the number {i}."},
        ],
        max_tokens=120,
        temperature=0.7,
    )
    latency = time.time() - t0
    tokens = resp.usage.completion_tokens
    return latency, tokens, tokens / latency
 
async def sweep(concurrency: int) -> None:
    t0 = time.time()
    results = await asyncio.gather(*[one_request(i) for i in range(concurrency)])
    wall = time.time() - t0
    lats = [r[0] for r in results]
    total_tokens = sum(r[1] for r in results)
    p50, p95 = quantiles(lats, n=20)[9], quantiles(lats, n=20)[18]
    print(
        f"conc={concurrency:3d}  wall={wall:5.1f}s  "
        f"throughput={total_tokens/wall:7.0f} tok/s  "
        f"lat p50={p50:.2f}s  p95={p95:.2f}s"
    )
 
async def main() -> None:
    for c in [1, 2, 4, 8, 16, 32]:
        await sweep(c)
 
if __name__ == "__main__":
    asyncio.run(main())