Karti Tripathi
107 lines
4.5 KiB
Python
107 lines
4.5 KiB
Python
"""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
|