From 3fa936cff409f754cc1ab1485f7733ded57f4aca Mon Sep 17 00:00:00 2001 From: cch-zuzuche <332574719+cch-zuzuche@users.noreply.github.com> Date: Wed, 23 Sep 2026 10:08:35 +0800 Subject: [PATCH 1/2] bench: prefix-reuse reproducer, interleaved-traffic dose-response, and a reuse-correctness needle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three controls for prefix-reuse behaviour, which this repo has no reproducer for. They came out of #174 and each one exists because its absence produced a wrong answer first. bench/prefix_alternation.py — two independent long conversations advanced in strict alternation, reading prompt_tokens_details.cached_tokens per request. The pass criterion is "this turn reused the whole prefix the previous turn established" (cached >= 0.9 * prev_prompt), not "a high hit rate": the failure mode is a clean 0%, so an average hides it. It reads the engine's own environment via /proc so an arm label cannot be whatever the caller exported. bench/interleave_dose.py — the two controls that make the above legible. --mode single advances one conversation with nothing else on the box; if that is clean the cache works and the defect is about interleaving. --mode dose injects one independent request of varying size between two turns of a long conversation. The knee this finds is the quantity to compare across arms: it moves with the KV pool, which is how capacity is separated from a fixed structural trigger. Both size their documents after calibrating tokens/char against the running tokenizer — a hardcoded ratio is a silent arm change, and the first version of this asked for 55,000 tokens, produced 48,000, and landed below the knee it was looking for. bench/needle_reuse.py — needle_test.py sends cold prompts and so never exercises the reuse path. Anything that changes which KV gets restored (retention intervals, checkpoint ordering, dtype, block promotion) is invisible to it, and the failure it misses is the worst kind: a high hit rate with a stale restored state. This one puts the passcode inside the reused prefix and requires both a high cache percentage and a correct answer. An optional interleaved second conversation can be advanced between the turns, and a needle at depth 0.95 sits at the edge of the reused region. All three take --tokens/--rounds/--doses and follow bench/ conventions (VLLM_API, VLLM_API_KEY or ../api_key.txt, optional VLLM_MODEL). Verified end to end against a running server. --- bench/interleave_dose.py | 168 ++++++++++++++++++++++++++ bench/needle_reuse.py | 132 +++++++++++++++++++++ bench/prefix_alternation.py | 227 ++++++++++++++++++++++++++++++++++++ 3 files changed, 527 insertions(+) create mode 100755 bench/interleave_dose.py create mode 100755 bench/needle_reuse.py create mode 100755 bench/prefix_alternation.py diff --git a/bench/interleave_dose.py b/bench/interleave_dose.py new file mode 100755 index 00000000..245d1049 --- /dev/null +++ b/bench/interleave_dose.py @@ -0,0 +1,168 @@ +#!/usr/bin/env python3 +"""Interleaved-traffic dose-response, and the single-conversation control. + +Both controls exist because the obvious experiment is misleading. + +1. Single conversation, advancing N turns with nothing else on the box. If this + is clean, the cache is working and the defect is about interleaving, not about + caching being broken. Without this arm, "conversation X lost its prefix" is + uninterpretable. + +2. One conversation with a second, independent request injected between two of + its turns, sweeping the injected size. The knee this finds is the quantity to + compare across arms: it moves with pool size roughly linearly, which is how + capacity is separated from a fixed structural trigger. + +The injected document is sized AFTER calibrating tokens/char against the running +tokenizer. A hardcoded ratio is a silent arm change — the first version of this +script asked for 55,000 tokens, produced 48,000, and landed below the knee it was +trying to find. + +Usage: + python bench/interleave_dose.py --mode single --tokens 130000 --turns 3 + python bench/interleave_dose.py --mode dose --tokens 60000 --doses 0,5000,30000,60000 +""" +import argparse +import json +import os +import sys +import time +import urllib.request + +HERE = os.path.dirname(os.path.abspath(__file__)) + + +def _key(path): # same convention as quality_battery.py + try: + return open(path).read().strip() + except OSError: + return "" + + +KEY = os.environ.get("VLLM_API_KEY") or _key(os.path.join(HERE, "..", "api_key.txt")) +API = os.environ.get("VLLM_API", "http://127.0.0.1:18020/v1") +MODEL = os.environ.get("VLLM_MODEL", "qwen3.8-27b") + +WORDS = ("the scheduler interleaves prefill chunks with decode steps while a mamba " + "state snapshot is materialised at the last prefill chunk boundary so the " + "attention prefix cache is the intersection of per-group hit sets").split() + +RATIO = None # calibrated at runtime + + +def make_doc_chars(n_chars, salt): + out, cur, i = [], 0, 0 + while cur < n_chars: + w = WORDS[i % len(WORDS)] + out.append(w) + cur += len(w) + 1 + i += 1 + return f"[{salt}] " + " ".join(out) + + +def make_doc(tokens, salt): + return make_doc_chars(int(tokens / RATIO), salt) + + +def ask(messages, max_tokens=24): + body = json.dumps({ + "model": MODEL, "messages": messages, "max_tokens": max_tokens, + "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}, + }).encode() + req = urllib.request.Request( + API + "/chat/completions", data=body, + headers={"Authorization": "Bearer " + KEY, "Content-Type": "application/json"}) + t0 = time.time() + d = json.load(urllib.request.urlopen(req, timeout=1800)) + el = time.time() - t0 + u = d["usage"] + det = u.get("prompt_tokens_details") or {} + return {"prompt": u["prompt_tokens"], "cached": det.get("cached_tokens", 0), + "secs": el, "text": (d["choices"][0]["message"].get("content") or "").strip()} + + +def calibrate(): + global RATIO + n_chars = 20000 + doc = make_doc_chars(n_chars, "calib") + r = ask([{"role": "user", "content": "CALIB\n" + doc}], max_tokens=1) + RATIO = r["prompt"] / float(len(doc)) + print(f"calibration: {len(doc):,} chars -> {r['prompt']:,} tok ({RATIO:.4f} tok/char)") + + +def mode_single(args): + doc = make_doc(args.tokens, f"single-{int(time.time())}") + conv = [{"role": "user", "content": doc}] + print(f"\nsingle conversation, {args.turns} turns, no other traffic") + prev = 0 + for t in range(1, args.turns + 1): + if t > 1: + conv.append({"role": "user", "content": f"round {t}: one short sentence."}) + r = ask(conv) + if t == 1: + verdict = "COLD" + else: + verdict = "ok" if r["cached"] > 0.5 * prev else "PREFIX-LOST" + print(f" turn {t} {verdict:<12} prompt={r['prompt']:>8,} " + f"cached={r['cached']:>8,} ({r['cached'] / max(r['prompt'], 1) * 100:5.1f}%) " + f"{r['secs']:6.2f}s") + sys.stdout.flush() + conv.append({"role": "assistant", "content": r["text"] or "ok"}) + prev = r["prompt"] + print("\nclean here means the cache works; interleaving is what breaks it.") + + +def mode_dose(args): + doses = [int(x) for x in args.doses.split(",")] + print(f"\nbig conversation {args.tokens:,} tok; injected sizes {doses}") + print(f" {'injected':>10} | {'turn2 prompt':>13} {'cached':>10} {'reused':>7} {'secs':>7} | verdict") + print(" " + "-" * 68) + results = [] + for dose in doses: + salt = f"dose{dose}-{int(time.time())}" + conv = [{"role": "user", "content": make_doc(args.tokens, salt + "-A")}] + r1 = ask(conv) + conv.append({"role": "assistant", "content": r1["text"] or "ok"}) + if dose > 0: + ask([{"role": "user", "content": make_doc(dose, salt + "-B")}]) + conv.append({"role": "user", "content": "round 2: one short sentence."}) + r2 = ask(conv) + pct = r2["cached"] / max(r2["prompt"], 1) * 100 + verdict = "ok" if pct > 50 else ("partial" if pct > 1 else "WIPED") + print(f" {dose:>10,} | {r2['prompt']:>13,} {r2['cached']:>10,} {pct:>6.1f}% " + f"{r2['secs']:>6.2f}s | {verdict}") + sys.stdout.flush() + results.append({"dose": dose, "prompt": r2["prompt"], + "cached": r2["cached"], "pct": round(pct, 1), + "secs": round(r2["secs"], 2), "verdict": verdict}) + print("\ncompare the knee across arms; it should move with the KV pool size.") + return results + + +def main(): + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + ap.add_argument("--mode", choices=("single", "dose"), default="dose") + ap.add_argument("--tokens", type=int, default=60000, + help="the long conversation's size, in tokens") + ap.add_argument("--turns", type=int, default=3, help="--mode single only") + ap.add_argument("--doses", default="0,5000,30000,60000", + help="--mode dose only; injected sizes in tokens") + ap.add_argument("--out", default=None) + args = ap.parse_args() + + calibrate() + if args.mode == "single": + mode_single(args) + return 0 + results = mode_dose(args) + if args.out: + with open(args.out, "w") as f: + json.dump({"args": vars(args), "ratio": RATIO, "results": results}, + f, ensure_ascii=False, indent=1) + print(f"raw records -> {args.out}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/bench/needle_reuse.py b/bench/needle_reuse.py new file mode 100755 index 00000000..fd6b15ec --- /dev/null +++ b/bench/needle_reuse.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +"""Reuse correctness: the answer lives inside the prefix that was reused. + +Why this is separate from needle_test.py. That probe sends cold prompts, so it +never exercises the prefix-reuse path. Anything that changes *which* KV is +restored — retention intervals, checkpoint ordering, KV dtype, block promotion — +is invisible to it, and the failure it misses is the worst kind: a high hit rate +with a stale or wrong restored state. A hit-rate-only test cannot see that. + +So this test does both, and requires both: + + turn 1 long document with a passcode buried at a given depth, stored in the + conversation, establishing the cache + turn 2 same conversation, asking only for the passcode -> must reuse turn 1's + prefix + pass cached/prompt over 50% AND the passcode is correct + +An optional interleaved second conversation can be advanced between the turns, +which is the case where reuse actually matters. A needle at depth 0.95 sits right +at the edge of the reused region and is the most sensitive position. + +Usage: + python bench/needle_reuse.py --tokens 100000 --depths 0.05,0.5,0.95 + python bench/needle_reuse.py --tokens 100000 --depths 0.5 --interleave 60000 +""" +import argparse +import json +import os +import sys +import time +import urllib.request + +HERE = os.path.dirname(os.path.abspath(__file__)) + + +def _key(path): # same convention as quality_battery.py + try: + return open(path).read().strip() + except OSError: + return "" + + +KEY = os.environ.get("VLLM_API_KEY") or _key(os.path.join(HERE, "..", "api_key.txt")) +API = os.environ.get("VLLM_API", "http://127.0.0.1:18020/v1") +MODEL = os.environ.get("VLLM_MODEL", "qwen3.8-27b") + +# ~46 chars, ~11 tokens of filler per unit +UNIT = "All work and no play makes Jack a dull boy. " + + +def ask(messages, max_tokens=48): + body = json.dumps({ + "model": MODEL, "messages": messages, "max_tokens": max_tokens, + "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}, + }).encode() + req = urllib.request.Request( + API + "/chat/completions", data=body, + headers={"Authorization": "Bearer " + KEY, "Content-Type": "application/json"}) + t0 = time.time() + d = json.load(urllib.request.urlopen(req, timeout=1800)) + u = d["usage"] + det = u.get("prompt_tokens_details") or {} + return {"prompt": u["prompt_tokens"], "cached": det.get("cached_tokens", 0), + "secs": time.time() - t0, + "ans": (d["choices"][0]["message"].get("content") or "").strip()} + + +def main(): + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + ap.add_argument("--tokens", type=int, default=100000) + ap.add_argument("--depths", default="0.05,0.5,0.95") + ap.add_argument("--interleave", type=int, default=0, + help="advance a second conversation of this size between turns") + ap.add_argument("--out", default=None) + args = ap.parse_args() + + filler = UNIT * int(args.tokens / 11) + print(f" ~{args.tokens:,} tok per conversation; needle depths {args.depths}") + results, allok = [], True + + for ds in args.depths.split(","): + dep = float(ds) + # a distinct passcode per depth, so one case cannot pass on another's cache + needle = f"PASS{dep:.2f}".replace(".", "") + "XYZ" + cut = int(len(filler) * dep) + ctx = (filler[:cut] + f"\n\nThe secret passcode is {needle}. Remember it exactly.\n\n" + + filler[cut:]) + + conv = [{"role": "user", "content": ctx + + "\n\nPlease acknowledge with the single word OK."}] + r1 = ask(conv, max_tokens=8) + conv.append({"role": "assistant", "content": r1["ans"] or "OK"}) + + if args.interleave > 0: + ask([{"role": "user", "content": + f"UNRELATED-{time.time():.6f}\n" + UNIT * int(args.interleave / 11)}], + max_tokens=4) + + conv.append({"role": "user", + "content": "What is the secret passcode? Reply with the passcode only."}) + r2 = ask(conv, max_tokens=48) + + pct = r2["cached"] / max(r2["prompt"], 1) * 100 + correct = needle in r2["ans"] + # BOTH conditions. A hit without the right answer is the silent failure. + if pct > 50 and correct: + verdict = "OK" + elif pct > 50: + verdict = "WRONG-ANSWER" + else: + verdict = "NO-REUSE" + allok = allok and verdict == "OK" + print(f" depth {dep:>5.0%} cache {pct:5.1f}% {r2['secs']:5.2f}s " + f"{verdict:<13} answer={r2['ans'][:32]!r}") + sys.stdout.flush() + results.append({"depth": dep, "needle": needle, "cached": r2["cached"], + "prompt": r2["prompt"], "pct": round(pct, 1), + "secs": round(r2["secs"], 2), "correct": correct, + "verdict": verdict}) + + print(f" -> {'all passed' if allok else 'FAILURES PRESENT'}") + if args.out: + with open(args.out, "w") as f: + json.dump({"args": vars(args), "results": results}, f, + ensure_ascii=False, indent=1) + print(f" raw records -> {args.out}") + return 0 if allok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/bench/prefix_alternation.py b/bench/prefix_alternation.py new file mode 100755 index 00000000..93509614 --- /dev/null +++ b/bench/prefix_alternation.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python3 +"""Prefix reuse under alternating conversations — reproducer and A/B arm. + +Two independent long conversations (A and B), advanced in strict alternation. +For A, every one of B's turns is interleaved traffic, and vice versa. Each turn +is its own request, which is what an agentic loop does. + +The pass criterion is not "a high hit rate" but "this turn reused the whole +prefix the previous turn established": ``cached_tokens >= 0.9 * prev_prompt``. +Failure is reported as PREFIX-LOST and means the turn re-prefilled in full. + +Regimes are sharp, so read the size table in the issue rather than one run: + + one conversation, no second one ~99.7% reused + two conversations, sequential ~96.6-99.7% + two conversations, alternating 0% every turn (above a size knee) + +Usage: + python bench/prefix_alternation.py --target-tokens 60000 --rounds 3 + python bench/prefix_alternation.py --target-tokens 96000 --noise 2 --rounds 4 + HQ_UNIT=qwen38-hq-vllm python bench/prefix_alternation.py # engine-env label +""" +import argparse +import datetime as dt +import json +import os +import subprocess +import sys +import time +import urllib.request + +HERE = os.path.dirname(os.path.abspath(__file__)) + + +def _key(path): # same convention as quality_battery.py + try: + return open(path).read().strip() + except OSError: + return "" + + +KEY = os.environ.get("VLLM_API_KEY") or _key(os.path.join(HERE, "..", "api_key.txt")) +API = os.environ.get("VLLM_API", "http://127.0.0.1:18020/v1") +MODEL = os.environ.get("VLLM_MODEL", "qwen3.8-27b") +UNIT = os.environ.get("HQ_UNIT", "qwen38-hq-vllm") + +LOG = [] + + +def call(messages, max_tokens=24): + body = json.dumps({ + "model": MODEL, "messages": messages, "max_tokens": max_tokens, + "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}, + }).encode() + req = urllib.request.Request( + API + "/chat/completions", data=body, + headers={"Authorization": "Bearer " + KEY, "Content-Type": "application/json"}) + t0 = time.time() + d = json.load(urllib.request.urlopen(req, timeout=1800)) + el = time.time() - t0 + u = d["usage"] + det = u.get("prompt_tokens_details") or {} + return {"prompt": u["prompt_tokens"], "cached": det.get("cached_tokens", 0), + "gen": u["completion_tokens"], "secs": el, + "text": (d["choices"][0]["message"].get("content") or "").strip()} + + +FILLERS = { + "A": ("The scheduler interleaves prefill chunks with decode steps. " + "A Mamba state snapshot is materialised at the last prefill chunk boundary. " + "Attention prefix caching is the intersection of per-group hit sets. "), + "B": ("Quantisation geometry decides how many tokens fit in the pool. " + "Speculative decoding acceptance is what turns steps into throughput. " + "Chunked prefill bounds how long a co-tenant can be starved. "), + "N": ("Unrelated interleaved traffic allocates blocks between turns. " + "This text exists only to consume KV blocks and force eviction order. "), +} + + +def make_doc(side, n_chars): + filler = FILLERS[side] + body = filler * (n_chars // len(filler) + 1) + lines = [body[i:i + 900] for i in range(0, n_chars, 900)] + return "\n".join(f"[{side}{i:05d}] {ln}" for i, ln in enumerate(lines))[:n_chars] + + +def log(rec): + LOG.append(rec) + print(json.dumps(rec, ensure_ascii=False), flush=True) + + +def engine_env(): + """Read the ENGINE process's environment, not this shell's. + + Otherwise an arm label is whatever the caller happened to export, which is + how a run gets reported against the wrong arm. Returns the retention setting + and max_model_len, or placeholders if the unit cannot be read. + """ + try: + pid = subprocess.check_output( + ["systemctl", "show", UNIT, "-p", "MainPID", "--value"], + text=True, timeout=10).strip() + raw = open(f"/proc/{pid}/environ", "rb").read().decode(errors="replace") + d = dict(x.split("=", 1) for x in raw.split("\0") if "=" in x) + extra = d.get("EXTRA_ARGS", "") + ret = "" + for tok, nxt in zip(extra.split(), extra.split()[1:]): + if tok == "--prefix-cache-retention-interval": + ret = nxt + if ret == "": + ret = d.get("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "") + return ret, d.get("MAX_LEN", "?") + except Exception as exc: # noqa: BLE001 + return f"", "?" + + +def main(): + ap = argparse.ArgumentParser( + description="Prefix reuse under alternating conversations (A/B arm).") + ap.add_argument("--target-tokens", type=int, default=100000, + help="base document size per conversation, in tokens") + ap.add_argument("--rounds", type=int, default=8, help="turns per conversation") + ap.add_argument("--noise", type=int, default=2, + help="unrelated requests injected between turns") + ap.add_argument("--noise-tokens", type=int, default=4000) + ap.add_argument("--budget-min", type=float, default=25.0) + ap.add_argument("--out", default=None) + args = ap.parse_args() + + retention, maxlen = engine_env() + tag = f"RET{retention}_ML{maxlen}" + t_start = time.time() + deadline = t_start + args.budget_min * 60 + out = args.out or os.path.join( + HERE, f"prefix-alternation-{tag}-{dt.datetime.now():%Y%m%d-%H%M%S}.json") + print(f"══ alternating-conversation prefix reuse ══ engine: retention={retention} " + f"MAX_LEN={maxlen} (read from /proc/, not this shell)") + print(f" {args.target_tokens:,} tok/conversation rounds={args.rounds} " + f"noise={args.noise}x{args.noise_tokens} tok budget={args.budget_min:.0f} min") + print(f" output -> {out}\n") + + # Calibrate tokens/char against THIS stack's tokenizer; a hardcoded ratio + # silently changes the arm you think you ran by tens of percent. + cal = call([{"role": "user", "content": "CALIB\n" + make_doc("A", 20000)}], max_tokens=1) + ratio = cal["prompt"] / 20000.0 + n_chars = int(args.target_tokens / ratio) + print(f"calibration: 20,000 chars -> {cal['prompt']:,} tok ({ratio:.4f} tok/char)") + print(f"per-conversation document: {n_chars:,} chars ≈ {args.target_tokens:,} tok\n") + log({"event": "calib", "ratio": ratio, "n_chars": n_chars, "arm": tag}) + + docs = {s: make_doc(s, n_chars) for s in ("A", "B")} + conv = {s: [{"role": "user", "content": docs[s]}] for s in ("A", "B")} + prev_prompt = {s: 0 for s in ("A", "B")} + stats = {"reqs": 0, "lost": 0, "healthy": 0, "cold": 0} + fail_rows = [] + + for rnd in range(1, args.rounds + 1): + for side in ("A", "B"): + if time.time() > deadline: + print(f"\nbudget exhausted before round {rnd} {side}") + rnd = args.rounds + 1 + break + q = (f"This is {side}'s question for round {rnd}. Answer in one sentence, " + f"do not restate the document.") + conv[side].append({"role": "user", "content": q}) + res = call(conv[side], max_tokens=24) + conv[side].append({"role": "assistant", "content": res["text"] or "OK"}) + + stats["reqs"] += 1 + if rnd == 1: + stats["cold"] += 1 + verdict = "COLD" + else: + need = 0.9 * prev_prompt[side] + if res["cached"] < need: + stats["lost"] += 1 + verdict = "PREFIX-LOST" + fail_rows.append((rnd, side, prev_prompt[side], res["cached"], res["secs"])) + else: + stats["healthy"] += 1 + verdict = "ok" + pct = (res["cached"] / res["prompt"] * 100) if res["prompt"] else 0 + print(f" r{rnd:>2} {side} {verdict:<11} " + f"prompt={res['prompt']:>7,} cached={res['cached']:>7,} ({pct:5.1f}%) " + f"{res['secs']:6.2f}s") + log({"event": "turn", "round": rnd, "side": side, "verdict": verdict, + "prompt": res["prompt"], "cached": res["cached"], "pct": round(pct, 1), + "secs": round(res["secs"], 2), "prev_prompt": prev_prompt[side]}) + prev_prompt[side] = res["prompt"] + + for k in range(args.noise): + if time.time() > deadline: + break + nd = f"NOISE-{side}-{rnd}-{k}-{time.time():.6f}\n" + make_doc( + "N", int(args.noise_tokens / ratio)) + nr = call([{"role": "user", "content": nd}], max_tokens=4) + log({"event": "noise", "round": rnd, "side": side, "k": k, + "prompt": nr["prompt"], "cached": nr["cached"], + "secs": round(nr["secs"], 2)}) + if time.time() > deadline: + break + + el = time.time() - t_start + print("\n" + "=" * 84) + print(f"arm={tag} {el:.0f}s conversation requests {stats['reqs']} " + f"(cold {stats['cold']} / healthy {stats['healthy']} / **lost {stats['lost']}**)") + if fail_rows: + print("\nprefix reuse lost:") + for r, s, pp, c, sec in fail_rows: + print(f" r{r} {s} previous prefix {pp:,} tok, reused {c:,} tok " + f"(short by {pp - c:,}, {sec:.1f}s)") + print("\nthis arm REPRODUCES the defect.") + else: + print("\n✅ no prefix loss: every turn reused the previous turn's whole prefix.") + print(" ⚠️ This only says so for this arm. To claim a fix, run the dense") + print(" arm too — if that is also clean, the test failed to reproduce") + print(" rather than the fix working.") + with open(out, "w") as f: + json.dump({"arm": tag, "args": vars(args), "stats": stats, + "elapsed_s": el, "log": LOG}, f, ensure_ascii=False, indent=1) + print(f"\nraw records -> {out}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From c475b66ac2e4bb54e5bbe7372cbf1d9efbf3f41d Mon Sep 17 00:00:00 2001 From: mhenrichsen Date: Wed, 23 Sep 2026 09:33:19 +0200 Subject: [PATCH 2/2] bench/prefix_alternation: read the arm from the engine's command line, judge turns against #102's floor, salt each run - The engine is found by its API port in /proc (or HQ_UNIT's MainPID, user unit first), and the retention interval and max_model_len are read from its command line: this repo's launcher passes both as flags, so the environment alone labelled a 13056 arm ''. - A healthy turn reuses max(0, prev//B - 1)*B (#102: the block holding the previous request's end is never served), B from vllm:cache_config_info, or --block when the engine was launched with an explicit --block-size, where the metric reports that value instead. The flat 90% rule marked every healthy turn under ~20 blocks PREFIX-LOST. - Documents carry a per-run salt, so round 1 is cold even when a previous run's prefix is still cached. --- bench/prefix_alternation.py | 121 ++++++++++++++++++++++++++++++------ 1 file changed, 101 insertions(+), 20 deletions(-) diff --git a/bench/prefix_alternation.py b/bench/prefix_alternation.py index 93509614..9c82832e 100755 --- a/bench/prefix_alternation.py +++ b/bench/prefix_alternation.py @@ -6,8 +6,13 @@ is its own request, which is what an agentic loop does. The pass criterion is not "a high hit rate" but "this turn reused the whole -prefix the previous turn established": ``cached_tokens >= 0.9 * prev_prompt``. -Failure is reported as PREFIX-LOST and means the turn re-prefilled in full. +prefix the previous turn established". A hybrid model never serves back the block +holding the previous request's end (#102), so a healthy turn reuses +``max(0, prev_prompt // B - 1) * B`` tokens, B being the engine's hybrid attention +block, read from ``vllm:cache_config_info`` on /metrics. A turn below that is +PREFIX-LOST. If the block cannot be read, the fallback is 90% of the previous +prompt, which misreads short conversations: with B = 2176 a healthy turn under +~44K tokens is below 90%. Regimes are sharp, so read the size table in the issue rather than one run: @@ -21,6 +26,7 @@ HQ_UNIT=qwen38-hq-vllm python bench/prefix_alternation.py # engine-env label """ import argparse +import re import datetime as dt import json import os @@ -42,11 +48,42 @@ def _key(path): # same convention as quality_battery.py KEY = os.environ.get("VLLM_API_KEY") or _key(os.path.join(HERE, "..", "api_key.txt")) API = os.environ.get("VLLM_API", "http://127.0.0.1:18020/v1") MODEL = os.environ.get("VLLM_MODEL", "qwen3.8-27b") -UNIT = os.environ.get("HQ_UNIT", "qwen38-hq-vllm") +UNIT = os.environ.get("HQ_UNIT", "") LOG = [] +def engine_block(): + """The engine's resolved hybrid attention block (what its boot line calls the + attention block size), from /metrics; None if it cannot be read. When the engine + was launched with an explicit --block-size (CTX=huge passes 128, KVarN's tile), + cache_config_info reports that value rather than the resolved hybrid block, so it + is not trusted: pass --block from the boot line instead.""" + try: + argv = [x.decode(errors="replace") + for x in open(f"/proc/{_engine_pid()}/cmdline", "rb").read().split(b"\0")] + if any(x == "--block-size" or x.startswith("--block-size=") for x in argv): + return None + except Exception: # noqa: BLE001 + pass + try: + base = API[: -len("/v1")] if API.endswith("/v1") else API + req = urllib.request.Request(base + "/metrics", + headers={"Authorization": "Bearer " + KEY}) + text = urllib.request.urlopen(req, timeout=30).read().decode() + m = re.search(r'^vllm:cache_config_info\{[^}]*\bblock_size="(\d+)"', text, re.M) + return int(m.group(1)) if m else None + except Exception: # noqa: BLE001 + return None + + +def healthy_floor(prev_prompt, block): + """Least cached_tokens a turn that reused the whole previous prefix can show.""" + if not block: + return 0.9 * prev_prompt + return max(0, prev_prompt // block - 1) * block + + def call(messages, max_tokens=24): body = json.dumps({ "model": MODEL, "messages": messages, "max_tokens": max_tokens, @@ -78,11 +115,15 @@ def call(messages, max_tokens=24): } +RUN = f"{time.time_ns():x}" # per-run salt: round 1 must be cold even if a previous run's prefix is cached + + def make_doc(side, n_chars): filler = FILLERS[side] body = filler * (n_chars // len(filler) + 1) lines = [body[i:i + 900] for i in range(0, n_chars, 900)] - return "\n".join(f"[{side}{i:05d}] {ln}" for i, ln in enumerate(lines))[:n_chars] + return f"[run {RUN}]\n" + "\n".join( + f"[{side}{i:05d}] {ln}" for i, ln in enumerate(lines))[:n_chars] def log(rec): @@ -90,27 +131,61 @@ def log(rec): print(json.dumps(rec, ensure_ascii=False), flush=True) +def _engine_pid(): + """PID of the serving engine: HQ_UNIT's MainPID if set (user unit first, then + system), else the `vllm serve` process on this API's port, found from /proc. + This repo's launcher execs `vllm serve`, so the port is on its command line; the + port scan also works inside a container, where there is no unit.""" + if UNIT: + for scope in (["--user"], []): + try: + pid = subprocess.check_output( + ["systemctl", *scope, "show", UNIT, "-p", "MainPID", "--value"], + text=True, timeout=10, stderr=subprocess.DEVNULL).strip() + if pid and pid != "0": + return pid + except Exception: # noqa: BLE001 + pass + port = API.split("://", 1)[-1].split("/", 1)[0].rsplit(":", 1)[-1] + for pid in filter(str.isdigit, os.listdir("/proc")): + try: + argv = open(f"/proc/{pid}/cmdline", "rb").read().split(b"\0") + except OSError: + continue + args = [x.decode(errors="replace") for x in argv] + if "serve" in args and any("vllm" in x for x in args[:3]) and \ + any(x == "--port" and nxt == port for x, nxt in zip(args, args[1:])): + return pid + raise RuntimeError(f"no vllm serve on port {port} (set HQ_UNIT to name its unit)") + + def engine_env(): - """Read the ENGINE process's environment, not this shell's. + """Read the ENGINE process's settings, not this shell's. Otherwise an arm label is whatever the caller happened to export, which is - how a run gets reported against the wrong arm. Returns the retention setting - and max_model_len, or placeholders if the unit cannot be read. + how a run gets reported against the wrong arm. The launcher computes both + values and passes them as flags, so the command line is read first and the + environment is only the fallback (the deprecated env spelling of the + retention interval). Returns placeholders if the engine cannot be read. """ try: - pid = subprocess.check_output( - ["systemctl", "show", UNIT, "-p", "MainPID", "--value"], - text=True, timeout=10).strip() + pid = _engine_pid() + argv = [x.decode(errors="replace") + for x in open(f"/proc/{pid}/cmdline", "rb").read().split(b"\0")] raw = open(f"/proc/{pid}/environ", "rb").read().decode(errors="replace") - d = dict(x.split("=", 1) for x in raw.split("\0") if "=" in x) - extra = d.get("EXTRA_ARGS", "") - ret = "" - for tok, nxt in zip(extra.split(), extra.split()[1:]): - if tok == "--prefix-cache-retention-interval": - ret = nxt - if ret == "": - ret = d.get("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "") - return ret, d.get("MAX_LEN", "?") + env = dict(x.split("=", 1) for x in raw.split("\0") if "=" in x) + + def flag(name): + for i, tok in enumerate(argv): + if tok == name and i + 1 < len(argv): + return argv[i + 1] + if tok.startswith(name + "="): + return tok.split("=", 1)[1] + return None + + ret = flag("--prefix-cache-retention-interval") \ + or env.get("VLLM_PREFIX_CACHE_RETENTION_INTERVAL") or "dense" + return ret, flag("--max-model-len") or env.get("MAX_LEN", "?") except Exception as exc: # noqa: BLE001 return f"", "?" @@ -126,9 +201,14 @@ def main(): ap.add_argument("--noise-tokens", type=int, default=4000) ap.add_argument("--budget-min", type=float, default=25.0) ap.add_argument("--out", default=None) + ap.add_argument("--block", type=int, default=0, + help="the engine's hybrid attention block, from its boot line " + "('Setting attention block size to N tokens'); default: read " + "from /metrics, which is only right without an explicit --block-size") args = ap.parse_args() retention, maxlen = engine_env() + block = args.block or engine_block() tag = f"RET{retention}_ML{maxlen}" t_start = time.time() deadline = t_start + args.budget_min * 60 @@ -136,6 +216,7 @@ def main(): HERE, f"prefix-alternation-{tag}-{dt.datetime.now():%Y%m%d-%H%M%S}.json") print(f"══ alternating-conversation prefix reuse ══ engine: retention={retention} " f"MAX_LEN={maxlen} (read from /proc/, not this shell)") + print(f" engine attention block: {block if block else 'unreadable, falling back to 90% (pass --block )'}") print(f" {args.target_tokens:,} tok/conversation rounds={args.rounds} " f"noise={args.noise}x{args.noise_tokens} tok budget={args.budget_min:.0f} min") print(f" output -> {out}\n") @@ -172,7 +253,7 @@ def main(): stats["cold"] += 1 verdict = "COLD" else: - need = 0.9 * prev_prompt[side] + need = healthy_floor(prev_prompt[side], block) if res["cached"] < need: stats["lost"] += 1 verdict = "PREFIX-LOST"