Karti Tripathi
320 lines
12 KiB
Python
320 lines
12 KiB
Python
"""Serving performance measurement.
|
|
|
|
A self-contained async load generator rather than a wrapper around
|
|
`vllm bench serve`. Reasons:
|
|
|
|
* It works against ANY OpenAI-compatible endpoint -- vLLM, SGLang,
|
|
llama.cpp, ollama, or a submitted model served however we choose to serve
|
|
it. Submissions are the point of the site; we cannot assume vLLM.
|
|
* It needs nothing installed on the target box. The runner talks HTTP.
|
|
* We control the metric definitions, so numbers stay comparable across
|
|
engines rather than inheriting each engine's benchmarking conventions.
|
|
|
|
TWO MEASUREMENT TRAPS THIS CODE AVOIDS -- both silently produce numbers that
|
|
look great and mean nothing:
|
|
|
|
1. PREFIX CACHING. vLLM caches shared prompt prefixes. If every concurrent
|
|
request sends the same prompt, prefill is nearly free after the first and
|
|
throughput is wildly overstated. Every request here gets a unique nonce
|
|
prefix so each one actually does its own prefill.
|
|
|
|
2. VARIABLE OUTPUT LENGTH. If the model decides when to stop, concurrency
|
|
levels finish at different token counts and tok/s is not comparable
|
|
between runs. We send `ignore_eos` so every request emits exactly
|
|
`output_tokens` tokens.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import statistics
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
|
|
import httpx
|
|
|
|
from .results import PerfPoint, PerfResult
|
|
from .specdec import read_counters
|
|
|
|
# Filler vocabulary for synthetic prompts. Ordinary words rather than repeated
|
|
# junk so tokenization stays close to what real traffic looks like.
|
|
_FILLER = (
|
|
"the quick brown fox jumps over a lazy dog while distant thunder rolls "
|
|
"across open water and the harbour lights flicker against low cloud cover "
|
|
"as fishing boats return with the morning tide carrying nets and crates "
|
|
).split()
|
|
|
|
|
|
def _synthetic_prompt(approx_tokens: int, nonce: str) -> str:
|
|
"""Build a prompt of roughly `approx_tokens` tokens, unique per request.
|
|
|
|
The nonce leads so it lands in the first block and defeats prefix-cache
|
|
reuse across concurrent requests.
|
|
"""
|
|
# ~0.75 tokens per word for ordinary English text.
|
|
n_words = max(1, int(approx_tokens / 0.75))
|
|
body = " ".join(_FILLER[i % len(_FILLER)] for i in range(n_words))
|
|
return f"[{nonce}] {body}"
|
|
|
|
|
|
@dataclass
|
|
class RequestResult:
|
|
ok: bool
|
|
ttft_ms: float | None = None
|
|
total_s: float | None = None
|
|
output_tokens: int = 0
|
|
input_tokens: int = 0
|
|
error: str | None = None
|
|
|
|
@property
|
|
def tpot_ms(self) -> float | None:
|
|
"""Inter-token latency: decode time divided by tokens after the first."""
|
|
if self.ttft_ms is None or self.total_s is None or self.output_tokens < 2:
|
|
return None
|
|
decode_ms = self.total_s * 1000 - self.ttft_ms
|
|
return decode_ms / (self.output_tokens - 1)
|
|
|
|
|
|
async def _one_request(
|
|
client: httpx.AsyncClient,
|
|
base_url: str,
|
|
model: str,
|
|
input_tokens: int,
|
|
output_tokens: int,
|
|
prompt: str | None = None,
|
|
pin_output: bool = True,
|
|
) -> RequestResult:
|
|
nonce = uuid.uuid4().hex[:12]
|
|
# A supplied prompt still gets the nonce, because prefix-cache defeat is orthogonal to
|
|
# whether the text is representative — reusing one real prompt across a concurrency level
|
|
# would let the engine serve most of it from cache and report an inflated number.
|
|
content = f"[{nonce}] {prompt}" if prompt else _synthetic_prompt(input_tokens, nonce)
|
|
payload = {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": content}],
|
|
"max_tokens": output_tokens,
|
|
"temperature": 0.0,
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
}
|
|
if pin_output:
|
|
# vLLM/SGLang extension: force exactly max_tokens of output so runs are
|
|
# comparable. Harmless (ignored) on engines that do not support it.
|
|
#
|
|
# Turned OFF for speculative-decoding measurement. Forcing generation past a natural
|
|
# stop produces degenerate continuation, and a draft model's acceptance rate on
|
|
# degenerate text is not its acceptance rate on real work — so pinning the length here
|
|
# buys comparability at the cost of measuring the wrong thing.
|
|
payload["ignore_eos"] = True
|
|
|
|
start = time.perf_counter()
|
|
ttft: float | None = None
|
|
counted = 0
|
|
usage_in = 0
|
|
usage_out = 0
|
|
|
|
try:
|
|
async with client.stream(
|
|
"POST", f"{base_url}/chat/completions", json=payload
|
|
) as resp:
|
|
if resp.status_code != 200:
|
|
body = (await resp.aread()).decode()[:200]
|
|
return RequestResult(ok=False, error=f"HTTP {resp.status_code}: {body}")
|
|
|
|
async for line in resp.aiter_lines():
|
|
if not line.startswith("data: "):
|
|
continue
|
|
data = line[6:]
|
|
if data == "[DONE]":
|
|
break
|
|
try:
|
|
chunk = json.loads(data)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
if usage := chunk.get("usage"):
|
|
usage_in = usage.get("prompt_tokens", 0)
|
|
usage_out = usage.get("completion_tokens", 0)
|
|
|
|
for choice in chunk.get("choices") or []:
|
|
delta = choice.get("delta") or {}
|
|
# Engines disagree on where reasoning text lands: vLLM with
|
|
# --reasoning-parser emits `reasoning`, others use
|
|
# `reasoning_content`. Count all of them -- reasoning tokens
|
|
# cost the same decode time and the user waits for them
|
|
# either way. Checking only `content` would report TTFT as
|
|
# the time to the *answer*, which on a reasoning model is
|
|
# thousands of tokens late.
|
|
if any(
|
|
isinstance(delta.get(k), str) and delta[k]
|
|
for k in ("content", "reasoning", "reasoning_content")
|
|
):
|
|
# The opening chunk carries role + empty content; it is
|
|
# the stream opening, not a token, and is excluded by
|
|
# the emptiness check above.
|
|
if ttft is None:
|
|
ttft = (time.perf_counter() - start) * 1000
|
|
counted += 1
|
|
except Exception as exc: # noqa: BLE001 - report, do not abort the sweep
|
|
return RequestResult(ok=False, error=f"{type(exc).__name__}: {exc}")
|
|
|
|
total = time.perf_counter() - start
|
|
return RequestResult(
|
|
ok=True,
|
|
ttft_ms=ttft,
|
|
total_s=total,
|
|
# Prefer server-reported usage; fall back to counted chunks.
|
|
output_tokens=usage_out or counted,
|
|
input_tokens=usage_in,
|
|
)
|
|
|
|
|
|
async def _sweep_point(
|
|
base_url: str,
|
|
model: str,
|
|
concurrency: int,
|
|
input_tokens: int,
|
|
output_tokens: int,
|
|
timeout_s: float,
|
|
prompts: list[str] | None = None,
|
|
pin_output: bool = True,
|
|
) -> PerfPoint:
|
|
limits = httpx.Limits(max_connections=concurrency + 8)
|
|
async with httpx.AsyncClient(timeout=timeout_s, limits=limits) as client:
|
|
# Warm up so the first measured request does not absorb graph capture,
|
|
# weight paging, or connection setup.
|
|
await _one_request(client, base_url, model, 32, 8)
|
|
|
|
# Bracket the measured window. These are lifetime counters, so the rate for THIS point
|
|
# is the delta — reading the total would fold in warm-up and every earlier point.
|
|
before = await read_counters(client, base_url)
|
|
|
|
start = time.perf_counter()
|
|
results = await asyncio.gather(
|
|
*(
|
|
_one_request(
|
|
client,
|
|
base_url,
|
|
model,
|
|
input_tokens,
|
|
output_tokens,
|
|
# Cycle rather than repeat: every concurrent request in a point gets a
|
|
# different prompt, so one unusually easy or hard example cannot set the
|
|
# whole level's number.
|
|
prompt=prompts[i % len(prompts)] if prompts else None,
|
|
pin_output=pin_output,
|
|
)
|
|
for i in range(concurrency)
|
|
)
|
|
)
|
|
wall = time.perf_counter() - start
|
|
after = await read_counters(client, base_url)
|
|
|
|
ok = [r for r in results if r.ok]
|
|
failed = [r for r in results if not r.ok]
|
|
|
|
total_out = sum(r.output_tokens for r in ok)
|
|
ttfts = [r.ttft_ms for r in ok if r.ttft_ms is not None]
|
|
tpots = [r.tpot_ms for r in ok if r.tpot_ms is not None]
|
|
per_stream = [
|
|
r.output_tokens / r.total_s for r in ok if r.total_s and r.output_tokens
|
|
]
|
|
|
|
def pct(values: list[float], p: float) -> float | None:
|
|
if not values:
|
|
return None
|
|
s = sorted(values)
|
|
idx = min(len(s) - 1, round(p * (len(s) - 1)))
|
|
return round(s[idx], 2)
|
|
|
|
# Prefill rate: input tokens processed per second, inferred from TTFT.
|
|
# At concurrency > 1 this is per-stream and includes queueing, so it reads
|
|
# low; treat the concurrency-1 value as the true prefill capability.
|
|
prefill = None
|
|
if ttfts and ok:
|
|
reported = [r.input_tokens for r in ok if r.input_tokens]
|
|
mean_in = statistics.mean(reported) if reported else input_tokens
|
|
mean_ttft_s = statistics.mean(ttfts) / 1000
|
|
if mean_ttft_s > 0:
|
|
prefill = round(mean_in / mean_ttft_s, 1)
|
|
|
|
# Delta across the measured window only. `None` when the engine exposes no counters, which
|
|
# is an ordinary outcome (no speculative decoding, or metrics disabled) and not an error.
|
|
spec = (after - before) if (before and after) else None
|
|
acceptance = spec.acceptance_rate if spec else None
|
|
drafted = int(spec.draft) if spec and spec.draft > 0 else None
|
|
|
|
# With a natural stop, requests no longer emit identical token counts, so the spread is
|
|
# itself a caveat on comparability and has to travel with the number.
|
|
out_counts = [r.output_tokens for r in ok if r.output_tokens]
|
|
stdev = (
|
|
round(statistics.stdev(out_counts), 2)
|
|
if not pin_output and len(out_counts) > 1
|
|
else None
|
|
)
|
|
|
|
return PerfPoint(
|
|
concurrency=concurrency,
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
n_requests=concurrency,
|
|
completed=len(ok),
|
|
failed=len(failed),
|
|
duration_s=round(wall, 3),
|
|
output_tps_total=round(total_out / wall, 2) if wall else 0.0,
|
|
output_tps_per_stream=round(statistics.mean(per_stream), 2) if per_stream else 0.0,
|
|
ttft_p50_ms=pct(ttfts, 0.50),
|
|
ttft_p95_ms=pct(ttfts, 0.95),
|
|
tpot_p50_ms=pct(tpots, 0.50),
|
|
prefill_tps=prefill,
|
|
spec_acceptance_rate=acceptance,
|
|
spec_draft_tokens=drafted,
|
|
natural_stop=not pin_output,
|
|
output_token_stdev=stdev,
|
|
error=failed[0].error if failed else None,
|
|
)
|
|
|
|
|
|
async def run_sweep(
|
|
base_url: str,
|
|
model: str,
|
|
concurrencies: tuple[int, ...] = (1, 8, 32),
|
|
input_tokens: int = 1024,
|
|
output_tokens: int = 256,
|
|
timeout_s: float = 600.0,
|
|
engine: str = "unknown",
|
|
progress: bool = True,
|
|
prompts: list[str] | None = None,
|
|
pin_output: bool = True,
|
|
) -> PerfResult:
|
|
"""Run a concurrency sweep and return a PerfResult.
|
|
|
|
`prompts` and `pin_output=False` together give the mode a speculative-decoding target needs:
|
|
representative text, stopping naturally. Both defaults stay as they were, so an ordinary
|
|
throughput sweep is unchanged and remains directly comparable to every result already
|
|
committed.
|
|
"""
|
|
points: list[PerfPoint] = []
|
|
for c in concurrencies:
|
|
if progress:
|
|
print(f" concurrency {c:>3} ...", end="", flush=True)
|
|
point = await _sweep_point(
|
|
base_url, model, c, input_tokens, output_tokens, timeout_s,
|
|
prompts=prompts, pin_output=pin_output,
|
|
)
|
|
points.append(point)
|
|
if progress:
|
|
status = f"{point.output_tps_total:>8.1f} tok/s total"
|
|
if point.spec_acceptance_rate is not None:
|
|
status += f" accept {point.spec_acceptance_rate:.0%}"
|
|
status += f" | {point.output_tps_per_stream:>6.1f} /stream"
|
|
if point.ttft_p50_ms is not None:
|
|
status += f" | TTFT p50 {point.ttft_p50_ms:>7.0f}ms"
|
|
if point.failed:
|
|
status += f" | {point.failed} FAILED"
|
|
print(status, flush=True)
|
|
|
|
return PerfResult(engine=engine, points=points)
|