cache_probe.py

"""Experiment 0 — characterize the KV prefix cache on the running model.

Answers empirically (not from theory) the questions that decide the whole design:
  - Is append-only actually cheap? (prefix reuse)
  - How expensive is editing an EARLY token? (should nuke cache from there)
  - Does a MID-CONTEXT contiguous delete reuse the suffix? (i.e. does
    --cache-reuse work on THIS model — known broken on qwen3-next)

Ground truth from llama-server /completion `timings`:
  prompt_n = tokens freshly PREFILLED (recomputed)
  cache_n  = tokens REUSED from KV cache
So cache_n high + prompt_n low = good reuse; prompt_n high = cache miss.
"""
import json
import sys

import httpx

BASE = "http://127.0.0.1:8080"


def _send(prompt):
    r = httpx.post(f"{BASE}/completion", json={
        "prompt": prompt, "n_predict": 1, "cache_prompt": True,
        "temperature": 0}, timeout=300)
    d = r.json()
    t = d["timings"]
    return {"prompt_n": t["prompt_n"], "cache_n": t.get("cache_n", 0),
            "prompt_ms": t["prompt_ms"]}


def _blocks(n):
    # each block ~deterministic ~40-60 tokens, uniquely labeled so we can
    # append / edit / delete precisely.
    return [f"[Block {i:03d}] The quick brown fox jumps over the lazy dog near "
            f"section {i} of the document, discussing topic number {i} in "
            f"reasonable but unremarkable detail for padding purposes.\n"
            for i in range(n)]


def probe(label, base_prompt, variant_prompt):
    _send(base_prompt)                       # warm the base into the slot
    r = _send(variant_prompt)                # measure the variant vs warm base
    reuse = 100 * r["cache_n"] / max(1, r["cache_n"] + r["prompt_n"])
    print(f"{label:<28} prefilled={r['prompt_n']:>5}  reused={r['cache_n']:>5}  "
          f"reuse={reuse:5.1f}%  prefill_ms={r['prompt_ms']:>7.0f}")
    return r


def main():
    N = 60
    blocks = _blocks(N)
    base = "".join(blocks)
    print(f"base prompt: {N} blocks\n")

    # warm once, measure a pure re-send (should be ~100% reuse)
    _send(base)
    r = _send(base)
    print(f"{'re-send identical':<28} prefilled={r['prompt_n']:>5}  "
          f"reused={r['cache_n']:>5}  (sanity: should be ~all reused)\n")

    # 1) APPEND at the end — the cheap case we want to rely on
    probe("append 3 blocks at END", base, base + "".join(_blocks(3)))

    # 2) EDIT an EARLY block — should invalidate from there on
    early = list(blocks)
    early[2] = "[Block 002] EDITED EARLY TOKENS CHANGE HERE completely different.\n"
    probe("edit block #2 (early)", base, "".join(early))

    # 3) EDIT a LATE block — cache should hold up to it
    late = list(blocks)
    late[N - 3] = "[Block 057] EDITED LATE near the end of the document here.\n"
    probe("edit block #57 (late)", base, "".join(late))

    # 4) DELETE a MID contiguous block — does suffix survive (--cache-reuse)?
    mid_del = blocks[:30] + blocks[31:]      # drop block 30
    probe("delete block #30 (mid)", base, "".join(mid_del))

    # 5) DELETE a bigger mid contiguous span (10 blocks) — cache-reuse needs a
    #    minimum gap; a larger contiguous delete is the friendlier case
    mid_del2 = blocks[:20] + blocks[35:]     # drop blocks 20-34
    probe("delete blocks 20-34 (mid)", base, "".join(mid_del2))

    print("\nreading: if 'delete mid' shows high reuse, --cache-reuse works on "
          "this model and contiguous eviction is viable. If it shows ~reuse only "
          "up to the cut, mid-context editing is OFF the table -> epoch-rebuild "
          "or suffix-only curation only.")


if __name__ == "__main__":
    sys.exit(main())