kv-landlords / scripts /needle_code_bench.py

Upload kv-quant (INT4/NVFP4 KIVI) work + vLLM fork source

6e668dc verified 5 months ago

13.8 kB

"""Long-range NEEDLE retrieval for code: BF16 vs INT4-KIVI on Laguna-XS.2.
The meaningful long-context test for a KV-cache quantizer: can the model attend
back to a UNIQUE definition placed FAR earlier in the context after the cache
has been quantized to INT4?  We build a long synthetic Python "codebase" of
filler functions padding to a target token length, embed a single
non-guessable needle at a controlled DEPTH near the start, and at the very end
ask the model to recall it.  Grading is by EXACT value / executed equality --
the answer is only right if the model actually attended to the early definition
(it cannot be guessed).
Two needle kinds (both exact-graded):
  * "const"     : ``MAGIC_SEED_<id> = <7-digit prime-ish int>`` -> ask the value.
  * "transform" : ``def secret_transform_<id>(x): return x * A + B`` -> ask the
                  model to compute ``secret_transform_<id>(N)``; correct answer
                  is the exact integer A*N+B.
For each (context_length, depth) cell we run several trials (distinct needle
values / filler) and report retrieval accuracy for BF16 and INT4-KIVI plus
their agreement.  Batch-1 greedy only (the regime Int4KiviCache supports).
Usage
-----
  .venv/bin/python scripts/needle_code_bench.py --lengths 8000,16000 \
       --depths 0.05 --trials 3 --max-new 24
"""
from __future__ import annotations

import argparse
import random
import re
import sys
import time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

ROOT = "/home/alex/poolside-hackathon-kv-quant"
sys.path.insert(0, ROOT)
from int4_kivi.hf_cache import Int4KiviCache

MODEL = "poolside/Laguna-XS.2"


# --------------------------------------------------------------------------- #
# Synthetic codebase / filler
# --------------------------------------------------------------------------- #
_VERBS = ["compute", "process", "build", "merge", "filter", "scan", "encode",
          "decode", "reduce", "expand", "rotate", "shuffle", "balance", "score",
          "align", "cluster", "sample", "weight", "smooth", "sharpen"]
_NOUNS = ["matrix", "vector", "buffer", "token", "graph", "tensor", "window",
          "stream", "record", "payload", "segment", "lattice", "kernel",
          "gradient", "index", "shard", "bucket", "frame", "channel", "node"]


def _filler_fn(rng: random.Random, i: int) -> str:
    """A plausible, in-distribution Python helper that carries no needle info."""
    v = rng.choice(_VERBS)
    n = rng.choice(_NOUNS)
    a = rng.randint(2, 9)
    b = rng.randint(1, 20)
    return (
        f"def {v}_{n}_{i}(items):\n"
        f'    """{v.capitalize()} the {n} {i} with a small affine pass."""\n'
        f"    out = []\n"
        f"    for k, it in enumerate(items):\n"
        f"        out.append(it * {a} + {b} + k)\n"
        f"    return out\n"
    )


def _make_needle(rng: random.Random, kind: str):
    """Return (needle_source, question_text, expected_answer_str, meta)."""
    nid = rng.randint(100000, 999999)
    if kind == "const":
        # 8-digit non-guessable value.
        val = rng.randint(10_000_000, 99_999_999)
        src = (
            f"# Important runtime configuration constant.\n"
            f"MAGIC_SEED_{nid} = {val}\n"
        )
        q = (
            f"In the reference code above there is a constant named "
            f"MAGIC_SEED_{nid}. Reply with ONLY its exact integer value and "
            f"nothing else."
        )
        return src, q, str(val), {"nid": nid, "val": val}
    elif kind == "transform":
        # Use a unique multiplier constant so recall is non-guessable. We ask
        # the model to recall the *added* offset literal, which is a pure
        # attend-back retrieval (no arithmetic), keeping this a clean needle
        # test rather than a reasoning test.
        a = rng.choice([3, 7, 11, 13])
        b = rng.randint(1_000_000, 9_999_999)  # 7-digit, non-guessable
        src = (
            f"# Domain-specific transform used by the pipeline.\n"
            f"def secret_transform_{nid}(x):\n"
            f"    return x * {a} + {b}\n"
        )
        q = (
            f"In the reference code above there is a function named "
            f"secret_transform_{nid} whose body is `return x * {a} + OFFSET`. "
            f"Reply with ONLY the exact integer OFFSET value and nothing else."
        )
        return src, q, str(b), {"nid": nid, "a": a, "b": b}
    else:
        raise ValueError(kind)


# --------------------------------------------------------------------------- #
# Build a context of ~target_tokens with the needle at a given depth fraction
# --------------------------------------------------------------------------- #
def build_context(tok, target_tokens: int, depth: float, kind: str,
                  rng: random.Random):
    """Return (reference_code_text, question, expected_str, meta, n_pre_tokens).
    ``depth`` is the fraction of the *filler* placed BEFORE the needle (0.0 =
    needle right at the start, 0.5 = needle in the middle).  We grow filler
    functions until the whole reference block reaches target_tokens.
    """
    needle_src, question, expected, meta = _make_needle(rng, kind)

    # Generate a large pool of filler, then split around the needle by depth.
    fillers: list[str] = []
    i = 0
    # Build until the joined text (incl. needle) is at/above target.
    # Estimate ~ a few tokens per char; just append and recount periodically.
    pre, post = [], []
    # First, mass-produce filler text and measure.
    block_parts: list[str] = []
    while True:
        fn = _filler_fn(rng, i)
        block_parts.append(fn)
        i += 1
        if i % 40 == 0:
            joined = "\n".join(block_parts)
            ntok = tok(joined, return_tensors="pt").input_ids.shape[1]
            if ntok >= target_tokens:
                break
        if i > 200000:  # safety
            break

    # Now place needle at the requested depth among the filler functions.
    n_fn = len(block_parts)
    split = int(n_fn * depth)
    parts = block_parts[:split] + [needle_src] + block_parts[split:]
    reference = "\n".join(parts)

    # Trim to roughly target_tokens (keep the needle!). We trim from the END so
    # the needle (near the front for small depth) is preserved.
    ids = tok(reference, return_tensors="pt").input_ids[0]
    if ids.shape[0] > target_tokens * 1.15:
        ids = ids[: int(target_tokens * 1.15)]
        reference = tok.decode(ids, skip_special_tokens=True)
        # Re-confirm needle survived; if trimmed off (shouldn't for depth<0.7),
        # re-insert near the front.
        if meta_token(meta, kind) not in reference:
            reference = needle_src + "\n" + reference

    n_pre = tok(reference, return_tensors="pt").input_ids.shape[1]
    return reference, question, expected, meta, n_pre


def meta_token(meta, kind: str) -> str:
    if kind == "const":
        return f"MAGIC_SEED_{meta['nid']}"
    return f"secret_transform_{meta['nid']}"


# --------------------------------------------------------------------------- #
# Prompt + generation
# --------------------------------------------------------------------------- #
def build_input_ids(tok, reference: str, question: str, device):
    sys_msg = (
        "You are a careful code assistant. You will be shown a long Python "
        "reference codebase, then asked one question about a specific symbol "
        "defined in it. Answer using ONLY the information in the reference."
    )
    user_msg = (
        "Reference codebase:\n\n"
        f"```python\n{reference}\n```\n\n"
        f"{question}"
    )
    msgs = [
        {"role": "system", "content": sys_msg},
        {"role": "user", "content": user_msg},
    ]
    return tok.apply_chat_template(
        msgs, add_generation_prompt=True, return_tensors="pt", return_dict=False,
    ).to(device)


@torch.no_grad()
def generate(model, input_ids, max_new, mode, config):
    cache = Int4KiviCache(config=config) if mode == "int4" else DynamicCache()
    out = model.generate(
        input_ids, max_new_tokens=max_new, past_key_values=cache,
        use_cache=True, do_sample=False, num_beams=1,
    )
    return out[0, input_ids.shape[1]:].tolist(), cache


def grade(response: str, expected: str) -> bool:
    """Exact-value match: the expected integer must appear as a standalone
    number in the response (and be the first number, to avoid echoing the
    function's coefficients)."""
    nums = re.findall(r"-?\d+", response.replace(",", ""))
    if not nums:
        return False
    # Correct if the expected value is the first integer emitted, OR appears
    # and no other distinct integer precedes it. We accept "first integer".
    return nums[0] == expected


# --------------------------------------------------------------------------- #
# Main sweep
# --------------------------------------------------------------------------- #
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--lengths", type=str, default="8000,16000,32000")
    ap.add_argument("--depths", type=str, default="0.05")
    ap.add_argument("--trials", type=int, default=3)
    ap.add_argument("--kinds", type=str, default="const,transform")
    ap.add_argument("--max-new", type=int, default=24)
    ap.add_argument("--seed", type=int, default=1234)
    args = ap.parse_args()

    lengths = [int(x) for x in args.lengths.split(",") if x]
    depths = [float(x) for x in args.depths.split(",") if x]
    kinds = [k for k in args.kinds.split(",") if k]

    print(f"[load] {MODEL} ...", flush=True)
    tok = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL, dtype=torch.bfloat16, device_map="auto"
    )
    model.eval()
    device = next(model.parameters()).device

    print(f"\n{'':#<92}")
    print("  NEEDLE-IN-CODEBASE RETRIEVAL  (exact-value graded, greedy, batch-1)")
    print(f"{'':#<92}")
    header = (f"  {'len':>7} {'depth':>6} {'kind':>10} {'ctx':>8} "
              f"{'BF16':>10} {'INT4':>10} {'agree':>7} {'ratio':>6} {'t/gen':>6}")
    print(header)
    print(f"{'':-<92}")

    rows = []
    example = None  # (length, kind, question, expected, bf16_resp, int4_resp)

    for length in lengths:
        for depth in depths:
            for kind in kinds:
                rng = random.Random(args.seed + length + int(depth * 1000)
                                    + hash(kind) % 1000)
                bf_ok = iq_ok = agree = 0
                ctxs, ratios, times = [], [], []
                for t in range(args.trials):
                    reference, question, expected, meta, n_pre = build_context(
                        tok, length, depth, kind, rng
                    )
                    input_ids = build_input_ids(tok, reference, question, device)
                    ctx = int(input_ids.shape[1])
                    ctxs.append(ctx)

                    res = {}
                    for mode in ("bf16", "int4"):
                        t0 = time.time()
                        ids, cache = generate(model, input_ids, args.max_new,
                                              mode, model.config)
                        dt = time.time() - t0
                        resp = tok.decode(ids, skip_special_tokens=True)
                        ok = grade(resp, expected)
                        res[mode] = (ok, resp, dt)
                        if mode == "int4":
                            try:
                                ratios.append(cache.compression_ratio_vs_bf16())
                            except Exception:
                                pass
                            times.append(dt)
                    bf_ok += res["bf16"][0]
                    iq_ok += res["int4"][0]
                    agree += (res["bf16"][0] == res["int4"][0])
                    if example is None and res["int4"][0] and res["bf16"][0]:
                        example = (ctx, kind, question, expected,
                                   res["bf16"][1].strip()[:120],
                                   res["int4"][1].strip()[:120])

                n = args.trials
                ctx_m = sum(ctxs) // len(ctxs)
                ratio_m = (sum(ratios) / len(ratios)) if ratios else 0.0
                t_m = (sum(times) / len(times)) if times else 0.0
                print(f"  {length:>7} {depth:>6.2f} {kind:>10} {ctx_m:>8} "
                      f"{bf_ok:>4}/{n:<5} {iq_ok:>4}/{n:<5} {agree:>3}/{n:<3} "
                      f"{ratio_m:>5.2f}x {t_m:>5.1f}s", flush=True)
                rows.append((length, depth, kind, ctx_m, bf_ok, iq_ok, agree, n,
                             ratio_m))

    print(f"{'':-<92}")
    # Aggregate by length.
    print("\n  RETRIEVAL ACCURACY BY CONTEXT LENGTH (summed over depths/kinds):")
    by_len: dict[int, list[int]] = {}
    for (length, depth, kind, ctx, bf, iq, ag, n, ratio) in rows:
        acc = by_len.setdefault(length, [0, 0, 0, 0])
        acc[0] += bf; acc[1] += iq; acc[2] += ag; acc[3] += n
    print(f"  {'target_len':>11} {'BF16':>14} {'INT4-KIVI':>14} {'agreement':>12}")
    for length in sorted(by_len):
        bf, iq, ag, n = by_len[length]
        print(f"  {length:>11} {bf:>6}/{n:<4} ({100*bf/n:>3.0f}%) "
              f"{iq:>6}/{n:<4} ({100*iq/n:>3.0f}%)  {ag:>4}/{n:<4} "
              f"({100*ag/n:>3.0f}%)")

    if example:
        ctx, kind, q, exp, bf_r, iq_r = example
        print(f"\n{'':-<92}")
        print(f"  REPRESENTATIVE NEEDLE (ctx={ctx} tok, kind={kind})")
        print(f"  Q: {q}")
        print(f"  expected: {exp}")
        print(f"  BF16  ->: {bf_r!r}")
        print(f"  INT4  ->: {iq_r!r}")
        print(f"{'':-<92}")

    print("\n[done]", flush=True)


if __name__ == "__main__":
    main()

	"""Long-range NEEDLE retrieval for code: BF16 vs INT4-KIVI on Laguna-XS.2.

	The meaningful long-context test for a KV-cache quantizer: can the model attend
	back to a UNIQUE definition placed FAR earlier in the context after the cache
	has been quantized to INT4? We build a long synthetic Python "codebase" of
	filler functions padding to a target token length, embed a single
	non-guessable needle at a controlled DEPTH near the start, and at the very end
	ask the model to recall it. Grading is by EXACT value / executed equality --
	the answer is only right if the model actually attended to the early definition
	(it cannot be guessed).

	Two needle kinds (both exact-graded):
	* "const" : ``MAGIC_SEED_<id> = <7-digit prime-ish int>`` -> ask the value.
	* "transform" : ``def secret_transform_<id>(x): return x * A + B`` -> ask the
	model to compute ``secret_transform_<id>(N)``; correct answer
	is the exact integer A*N+B.

	For each (context_length, depth) cell we run several trials (distinct needle
	values / filler) and report retrieval accuracy for BF16 and INT4-KIVI plus
	their agreement. Batch-1 greedy only (the regime Int4KiviCache supports).

	Usage
	-----
	.venv/bin/python scripts/needle_code_bench.py --lengths 8000,16000 \
	--depths 0.05 --trials 3 --max-new 24
	"""
	from __future__ import annotations

	import argparse
	import random
	import re
	import sys
	import time

	import torch
	from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

	ROOT = "/home/alex/poolside-hackathon-kv-quant"
	sys.path.insert(0, ROOT)
	from int4_kivi.hf_cache import Int4KiviCache

	MODEL = "poolside/Laguna-XS.2"


	# --------------------------------------------------------------------------- #
	# Synthetic codebase / filler
	# --------------------------------------------------------------------------- #
	_VERBS = ["compute", "process", "build", "merge", "filter", "scan", "encode",
	"decode", "reduce", "expand", "rotate", "shuffle", "balance", "score",
	"align", "cluster", "sample", "weight", "smooth", "sharpen"]
	_NOUNS = ["matrix", "vector", "buffer", "token", "graph", "tensor", "window",
	"stream", "record", "payload", "segment", "lattice", "kernel",
	"gradient", "index", "shard", "bucket", "frame", "channel", "node"]


	def _filler_fn(rng: random.Random, i: int) -> str:
	"""A plausible, in-distribution Python helper that carries no needle info."""
	v = rng.choice(_VERBS)
	n = rng.choice(_NOUNS)
	a = rng.randint(2, 9)
	b = rng.randint(1, 20)
	return (
	f"def {v}_{n}_{i}(items):\n"
	f' """{v.capitalize()} the {n} {i} with a small affine pass."""\n'
	f" out = []\n"
	f" for k, it in enumerate(items):\n"
	f" out.append(it * {a} + {b} + k)\n"
	f" return out\n"
	)


	def _make_needle(rng: random.Random, kind: str):
	"""Return (needle_source, question_text, expected_answer_str, meta)."""
	nid = rng.randint(100000, 999999)
	if kind == "const":
	# 8-digit non-guessable value.
	val = rng.randint(10_000_000, 99_999_999)
	src = (
	f"# Important runtime configuration constant.\n"
	f"MAGIC_SEED_{nid} = {val}\n"
	)
	q = (
	f"In the reference code above there is a constant named "
	f"MAGIC_SEED_{nid}. Reply with ONLY its exact integer value and "
	f"nothing else."
	)
	return src, q, str(val), {"nid": nid, "val": val}
	elif kind == "transform":
	# Use a unique multiplier constant so recall is non-guessable. We ask
	# the model to recall the added offset literal, which is a pure
	# attend-back retrieval (no arithmetic), keeping this a clean needle
	# test rather than a reasoning test.
	a = rng.choice([3, 7, 11, 13])
	b = rng.randint(1_000_000, 9_999_999) # 7-digit, non-guessable
	src = (
	f"# Domain-specific transform used by the pipeline.\n"
	f"def secret_transform_{nid}(x):\n"
	f" return x * {a} + {b}\n"
	)
	q = (
	f"In the reference code above there is a function named "
	f"secret_transform_{nid} whose body is `return x * {a} + OFFSET`. "
	f"Reply with ONLY the exact integer OFFSET value and nothing else."
	)
	return src, q, str(b), {"nid": nid, "a": a, "b": b}
	else:
	raise ValueError(kind)


	# --------------------------------------------------------------------------- #
	# Build a context of ~target_tokens with the needle at a given depth fraction
	# --------------------------------------------------------------------------- #
	def build_context(tok, target_tokens: int, depth: float, kind: str,
	rng: random.Random):
	"""Return (reference_code_text, question, expected_str, meta, n_pre_tokens).

	``depth`` is the fraction of the filler placed BEFORE the needle (0.0 =
	needle right at the start, 0.5 = needle in the middle). We grow filler
	functions until the whole reference block reaches target_tokens.
	"""
	needle_src, question, expected, meta = _make_needle(rng, kind)

	# Generate a large pool of filler, then split around the needle by depth.
	fillers: list[str] = []
	i = 0
	# Build until the joined text (incl. needle) is at/above target.
	# Estimate ~ a few tokens per char; just append and recount periodically.
	pre, post = [], []
	# First, mass-produce filler text and measure.
	block_parts: list[str] = []
	while True:
	fn = _filler_fn(rng, i)
	block_parts.append(fn)
	i += 1
	if i % 40 == 0:
	joined = "\n".join(block_parts)
	ntok = tok(joined, return_tensors="pt").input_ids.shape[1]
	if ntok >= target_tokens:
	break
	if i > 200000: # safety
	break

	# Now place needle at the requested depth among the filler functions.
	n_fn = len(block_parts)
	split = int(n_fn * depth)
	parts = block_parts[:split] + [needle_src] + block_parts[split:]
	reference = "\n".join(parts)

	# Trim to roughly target_tokens (keep the needle!). We trim from the END so
	# the needle (near the front for small depth) is preserved.
	ids = tok(reference, return_tensors="pt").input_ids[0]
	if ids.shape[0] > target_tokens * 1.15:
	ids = ids[: int(target_tokens * 1.15)]
	reference = tok.decode(ids, skip_special_tokens=True)
	# Re-confirm needle survived; if trimmed off (shouldn't for depth<0.7),
	# re-insert near the front.
	if meta_token(meta, kind) not in reference:
	reference = needle_src + "\n" + reference

	n_pre = tok(reference, return_tensors="pt").input_ids.shape[1]
	return reference, question, expected, meta, n_pre


	def meta_token(meta, kind: str) -> str:
	if kind == "const":
	return f"MAGIC_SEED_{meta['nid']}"
	return f"secret_transform_{meta['nid']}"


	# --------------------------------------------------------------------------- #
	# Prompt + generation
	# --------------------------------------------------------------------------- #
	def build_input_ids(tok, reference: str, question: str, device):
	sys_msg = (
	"You are a careful code assistant. You will be shown a long Python "
	"reference codebase, then asked one question about a specific symbol "
	"defined in it. Answer using ONLY the information in the reference."
	)
	user_msg = (
	"Reference codebase:\n\n"
	f"```python\n{reference}\n```\n\n"
	f"{question}"
	)
	msgs = [
	{"role": "system", "content": sys_msg},
	{"role": "user", "content": user_msg},
	]
	return tok.apply_chat_template(
	msgs, add_generation_prompt=True, return_tensors="pt", return_dict=False,
	).to(device)


	@torch.no_grad()
	def generate(model, input_ids, max_new, mode, config):
	cache = Int4KiviCache(config=config) if mode == "int4" else DynamicCache()
	out = model.generate(
	input_ids, max_new_tokens=max_new, past_key_values=cache,
	use_cache=True, do_sample=False, num_beams=1,
	)
	return out[0, input_ids.shape[1]:].tolist(), cache


	def grade(response: str, expected: str) -> bool:
	"""Exact-value match: the expected integer must appear as a standalone
	number in the response (and be the first number, to avoid echoing the
	function's coefficients)."""
	nums = re.findall(r"-?\d+", response.replace(",", ""))
	if not nums:
	return False
	# Correct if the expected value is the first integer emitted, OR appears
	# and no other distinct integer precedes it. We accept "first integer".
	return nums[0] == expected


	# --------------------------------------------------------------------------- #
	# Main sweep
	# --------------------------------------------------------------------------- #
	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--lengths", type=str, default="8000,16000,32000")
	ap.add_argument("--depths", type=str, default="0.05")
	ap.add_argument("--trials", type=int, default=3)
	ap.add_argument("--kinds", type=str, default="const,transform")
	ap.add_argument("--max-new", type=int, default=24)
	ap.add_argument("--seed", type=int, default=1234)
	args = ap.parse_args()

	lengths = [int(x) for x in args.lengths.split(",") if x]
	depths = [float(x) for x in args.depths.split(",") if x]
	kinds = [k for k in args.kinds.split(",") if k]

	print(f"[load] {MODEL} ...", flush=True)
	tok = AutoTokenizer.from_pretrained(MODEL)
	model = AutoModelForCausalLM.from_pretrained(
	MODEL, dtype=torch.bfloat16, device_map="auto"
	)
	model.eval()
	device = next(model.parameters()).device

	print(f"\n{'':#<92}")
	print(" NEEDLE-IN-CODEBASE RETRIEVAL (exact-value graded, greedy, batch-1)")
	print(f"{'':#<92}")
	header = (f" {'len':>7} {'depth':>6} {'kind':>10} {'ctx':>8} "
	f"{'BF16':>10} {'INT4':>10} {'agree':>7} {'ratio':>6} {'t/gen':>6}")
	print(header)
	print(f"{'':-<92}")

	rows = []
	example = None # (length, kind, question, expected, bf16_resp, int4_resp)

	for length in lengths:
	for depth in depths:
	for kind in kinds:
	rng = random.Random(args.seed + length + int(depth * 1000)
	+ hash(kind) % 1000)
	bf_ok = iq_ok = agree = 0
	ctxs, ratios, times = [], [], []
	for t in range(args.trials):
	reference, question, expected, meta, n_pre = build_context(
	tok, length, depth, kind, rng
	)
	input_ids = build_input_ids(tok, reference, question, device)
	ctx = int(input_ids.shape[1])
	ctxs.append(ctx)

	res = {}
	for mode in ("bf16", "int4"):
	t0 = time.time()
	ids, cache = generate(model, input_ids, args.max_new,
	mode, model.config)
	dt = time.time() - t0
	resp = tok.decode(ids, skip_special_tokens=True)
	ok = grade(resp, expected)
	res[mode] = (ok, resp, dt)
	if mode == "int4":
	try:
	ratios.append(cache.compression_ratio_vs_bf16())
	except Exception:
	pass
	times.append(dt)
	bf_ok += res["bf16"][0]
	iq_ok += res["int4"][0]
	agree += (res["bf16"][0] == res["int4"][0])
	if example is None and res["int4"][0] and res["bf16"][0]:
	example = (ctx, kind, question, expected,
	res["bf16"][1].strip()[:120],
	res["int4"][1].strip()[:120])

	n = args.trials
	ctx_m = sum(ctxs) // len(ctxs)
	ratio_m = (sum(ratios) / len(ratios)) if ratios else 0.0
	t_m = (sum(times) / len(times)) if times else 0.0
	print(f" {length:>7} {depth:>6.2f} {kind:>10} {ctx_m:>8} "
	f"{bf_ok:>4}/{n:<5} {iq_ok:>4}/{n:<5} {agree:>3}/{n:<3} "
	f"{ratio_m:>5.2f}x {t_m:>5.1f}s", flush=True)
	rows.append((length, depth, kind, ctx_m, bf_ok, iq_ok, agree, n,
	ratio_m))

	print(f"{'':-<92}")
	# Aggregate by length.
	print("\n RETRIEVAL ACCURACY BY CONTEXT LENGTH (summed over depths/kinds):")
	by_len: dict[int, list[int]] = {}
	for (length, depth, kind, ctx, bf, iq, ag, n, ratio) in rows:
	acc = by_len.setdefault(length, [0, 0, 0, 0])
	acc[0] += bf; acc[1] += iq; acc[2] += ag; acc[3] += n
	print(f" {'target_len':>11} {'BF16':>14} {'INT4-KIVI':>14} {'agreement':>12}")
	for length in sorted(by_len):
	bf, iq, ag, n = by_len[length]
	print(f" {length:>11} {bf:>6}/{n:<4} ({100*bf/n:>3.0f}%) "
	f"{iq:>6}/{n:<4} ({100*iq/n:>3.0f}%) {ag:>4}/{n:<4} "
	f"({100*ag/n:>3.0f}%)")

	if example:
	ctx, kind, q, exp, bf_r, iq_r = example
	print(f"\n{'':-<92}")
	print(f" REPRESENTATIVE NEEDLE (ctx={ctx} tok, kind={kind})")
	print(f" Q: {q}")
	print(f" expected: {exp}")
	print(f" BF16 ->: {bf_r!r}")
	print(f" INT4 ->: {iq_r!r}")
	print(f"{'':-<92}")

	print("\n[done]", flush=True)


	if __name__ == "__main__":
	main()