Gahow Wang
837df6bc9e
Measures TTFT to serve a reused prefix of length L from each KV tier on a single H20 (Qwen3-Coder-30B-A3B, vLLM 0.18.1): miss (recompute), CPU-tier hit (native DRAM offload), GPU-tier hit (HBM prefix cache). Each measured request is bracketed by /metrics scrapes so the tier is verified (vllm:prefix_cache_hits vs external_prefix_cache_hits), not assumed. Result: GPU hit is ~flat (42->111 ms over 1k->64k tokens); CPU hit is transfer-bound (PCIe H2D ~54 GB/s, 57->272 ms); miss grows superlinearly (78 ms -> 15.2 s). GPU beats CPU 1.4-2.5x (gap grows with context); miss/CPU up to 56x, miss/GPU up to 137x. pcie_transfer.py is the independent CPU-hit floor backstop. Evidence for the GPU-hit-first principle (paper section 2.2). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
"""Shared helpers for v2 GPU-hit-first experiments."""
|
|
from __future__ import annotations
|
|
|
|
import random
|
|
import time
|
|
import requests
|
|
|
|
# Qwen3-Coder geometry (from config.json): 48 layers, 4 KV heads, head_dim 128, bf16
|
|
KV_BYTES_PER_TOKEN = 98304 # 96 KiB
|
|
VOCAB = 151936
|
|
# Safe token-id range: avoid low special-ish ids and the high special tokens (>=151643)
|
|
TOK_LO, TOK_HI = 1000, 151000
|
|
|
|
|
|
def make_token_prompt(length: int, seed: int) -> list[int]:
|
|
"""Deterministic, content-addressed token-id prompt of exact `length`.
|
|
|
|
Same (length, seed) -> same ids -> prefix-cache hit.
|
|
Different seed -> fresh ids -> miss.
|
|
"""
|
|
rng = random.Random(seed)
|
|
return [rng.randint(TOK_LO, TOK_HI) for _ in range(length)]
|
|
|
|
|
|
def scrape_prefix_cache(endpoint: str) -> dict:
|
|
"""Return cumulative prefix-cache counters from vLLM /metrics.
|
|
|
|
Keys: gpu_hits, gpu_queries, ext_hits, ext_queries (floats, cumulative).
|
|
"""
|
|
out = {"gpu_hits": 0.0, "gpu_queries": 0.0, "ext_hits": 0.0, "ext_queries": 0.0}
|
|
try:
|
|
txt = requests.get(f"{endpoint}/metrics", timeout=10).text
|
|
except Exception:
|
|
return out
|
|
for line in txt.splitlines():
|
|
if line.startswith("#") or not line:
|
|
continue
|
|
try:
|
|
name, val = line.rsplit(" ", 1)
|
|
v = float(val)
|
|
except ValueError:
|
|
continue
|
|
# strip prometheus labels and match only the cumulative _total counters
|
|
# (exclude _created epoch-timestamp series, which would dominate the sum)
|
|
metric = name.split("{", 1)[0]
|
|
if metric == "vllm:external_prefix_cache_hits_total":
|
|
out["ext_hits"] += v
|
|
elif metric == "vllm:external_prefix_cache_queries_total":
|
|
out["ext_queries"] += v
|
|
elif metric == "vllm:prefix_cache_hits_total":
|
|
out["gpu_hits"] += v
|
|
elif metric == "vllm:prefix_cache_queries_total":
|
|
out["gpu_queries"] += v
|
|
return out
|
|
|
|
|
|
def measure_ttft(endpoint: str, model: str, prompt_ids: list[int],
|
|
max_tokens: int = 1, timeout: float = 600.0) -> dict:
|
|
"""Send one streaming /v1/completions request; return TTFT and e2e seconds.
|
|
|
|
TTFT = time from send to first streamed token chunk (== prefill wall time).
|
|
"""
|
|
url = f"{endpoint}/v1/completions"
|
|
payload = {
|
|
"model": model,
|
|
"prompt": prompt_ids,
|
|
"max_tokens": max_tokens,
|
|
"temperature": 0.0,
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
}
|
|
t0 = time.perf_counter()
|
|
ttft = None
|
|
usage = None
|
|
with requests.post(url, json=payload, stream=True, timeout=timeout) as r:
|
|
r.raise_for_status()
|
|
for raw in r.iter_lines():
|
|
if not raw:
|
|
continue
|
|
line = raw.decode("utf-8") if isinstance(raw, bytes) else raw
|
|
if not line.startswith("data: "):
|
|
continue
|
|
data = line[6:]
|
|
if data.strip() == "[DONE]":
|
|
break
|
|
import json as _json
|
|
obj = _json.loads(data)
|
|
if obj.get("usage"):
|
|
usage = obj["usage"]
|
|
choices = obj.get("choices") or []
|
|
if ttft is None and choices and choices[0].get("text"):
|
|
ttft = time.perf_counter() - t0
|
|
e2e = time.perf_counter() - t0
|
|
return {"ttft_s": ttft if ttft is not None else e2e, "e2e_s": e2e, "usage": usage}
|
|
|
|
|
|
def wait_healthy(endpoint: str, timeout: float = 900.0) -> bool:
|
|
deadline = time.time() + timeout
|
|
while time.time() < deadline:
|
|
try:
|
|
if requests.get(f"{endpoint}/health", timeout=5).status_code == 200:
|
|
return True
|
|
except Exception:
|
|
pass
|
|
time.sleep(3)
|
|
return False
|