Lumbridge Bench

kbench/specdec.py

@@ -0,0 +1,106 @@
+"""Speculative-decoding metrics, read from the engine rather than inferred.
+
+Throughput alone cannot tell you whether speculative decoding is working. A draft model that
+is rejected almost every time still produces tokens — just slowly, having paid for the draft
+pass as well as the verify pass — and on a synthetic benchmark that can even look like a
+modest win. The number that says whether it is working is the **acceptance rate**: of the
+tokens the draft proposed, how many survived verification.
+
+vLLM exposes this on its Prometheus endpoint as counters, so a rate for one sweep point is a
+delta across that point rather than the process-lifetime total. Reading the total instead would
+mix in every request since the server started, including warm-up.
+
+Parsing is pure standard library and `httpx` is only needed to fetch — the import is deferred
+so the parser can be tested, and reasoned about, without the HTTP client installed. Same rule as
+`kbench/schema.py`: the logic that decides what a number means must not depend on the machinery
+that fetches it.
+
+This is deliberately best-effort. An engine without these counters, or with metrics disabled,
+gets `None` and the sweep continues — an absent metric must never fail a benchmark run.
+
+Why it matters for this bench specifically: `perf.py` sends a random nonce prefix and
+`ignore_eos`, both correct for defeating prefix caching and pinning output length, and both
+adversarial to speculative decoding. Forced continuation past a natural stop is degenerate
+text, and a draft model's acceptance on degenerate text is not its acceptance on real work. So
+a speculative target measured with the synthetic prompt needs its acceptance rate reported
+beside the throughput, or the throughput will be read as if it transfers.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:  # pragma: no cover
+    import httpx
+
+# vLLM V1 names. The older V0 engine used vllm:spec_decode_* with the same meaning; both are
+# accepted so a mixed fleet does not silently report nothing.
+ACCEPTED = ("vllm:spec_decode_num_accepted_tokens_total",)
+DRAFT = ("vllm:spec_decode_num_draft_tokens_total",)
+
+_SAMPLE = re.compile(r"^(?P<name>[a-zA-Z_:][\w:]*)(?P<labels>\{[^}]*\})?\s+(?P<value>[-+0-9.eE]+)$")
+
+
+@dataclass(frozen=True)
+class SpecCounters:
+    accepted: float
+    draft: float
+
+    def __sub__(self, other: "SpecCounters") -> "SpecCounters":
+        return SpecCounters(self.accepted - other.accepted, self.draft - other.draft)
+
+    @property
+    def acceptance_rate(self) -> float | None:
+        """Accepted / proposed over this window. None when nothing was proposed."""
+        if self.draft <= 0:
+            return None
+        return round(self.accepted / self.draft, 4)
+
+
+def parse_metrics(text: str) -> SpecCounters | None:
+    """Pull the speculative counters out of a Prometheus exposition payload."""
+    totals: dict[str, float] = {}
+    for line in text.splitlines():
+        line = line.strip()
+        if not line or line.startswith("#"):
+            continue
+        m = _SAMPLE.match(line)
+        if not m:
+            continue
+        name = m.group("name")
+        if name in ACCEPTED or name in DRAFT:
+            try:
+                # Counters are per-label-set (one per model); sum them. A single-model server
+                # has one series, and summing is still correct there.
+                totals[name] = totals.get(name, 0.0) + float(m.group("value"))
+            except ValueError:
+                continue
+
+    accepted = next((totals[n] for n in ACCEPTED if n in totals), None)
+    draft = next((totals[n] for n in DRAFT if n in totals), None)
+    if accepted is None or draft is None:
+        return None
+    return SpecCounters(accepted=accepted, draft=draft)
+
+
+def metrics_url(base_url: str) -> str:
+    """/metrics sits at the server root, not under the OpenAI /v1 prefix."""
+    root = base_url.rstrip("/")
+    for suffix in ("/v1", "/openai/v1"):
+        if root.endswith(suffix):
+            root = root[: -len(suffix)]
+            break
+    return f"{root}/metrics"
+
+
+async def read_counters(client: "httpx.AsyncClient", base_url: str) -> SpecCounters | None:
+    """Best-effort read. Any failure means "no data", never an exception into the sweep."""
+    try:
+        resp = await client.get(metrics_url(base_url), timeout=5.0)
+        if resp.status_code != 200:
+            return None
+        return parse_metrics(resp.text)
+    except Exception:  # noqa: BLE001 - a missing metrics endpoint is an ordinary outcome
+        return None