Lumbridge Bench

kbench/perf.py

@@ -0,0 +1,319 @@
+"""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)