scripts/quant_ab.py

"""Downstream evidence: INT4/INT3 + per-channel-K (KIVI) vs vLLM's NVFP4 KV baseline.

Across several prompts on Laguna-XS.2, reports for each scheme:
  * K-RMSE / V-RMSE   — reconstruction error (cheap proxy; tracks KL)
  * top-1 agreement   — teacher-forced vs BF16 (identical context, no drift)
  * mean KL(bf16||scheme) in nats — output-distribution distortion

Protocol is production-faithful: each 16-token page is quantized once when it
fills (frozen thereafter), and the partial hot page stays BF16. Teacher forcing
replays BF16's own tokens so every scheme sees identical context.

Schemes:
  nvfp4-baseline  K,V = nvfp4 / headdim / absmax        (what vLLM ships, 4-bit)
  int4-kivi       K = int4 / per-channel / mse,  V = int4 / per-token / mse
  int3-kivi       K = int3 / per-channel / mse,  V = int3 / per-token / mse
  int3-naive      K,V = int3 / headdim / absmax         (3-bit done the vLLM way)

Usage:
    python -m scripts.quant_ab [--max-new 384] [--n-prompts 3]
"""
from __future__ import annotations

import argparse
import sys
import time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

sys.path.insert(0, "/home/alex/poolside-hackathon-kv-quant")
from kv_quant import BLOCK, PAGE, roundtrip, rmse_cell

MODEL = "poolside/Laguna-XS.2"

PROMPTS = [
    "Solve step by step. A train leaves city A at 60 km/h. Two hours later a second "
    "train leaves the same station on the same track at 90 km/h. How many hours after "
    "the second train departs will it catch up to the first train? Show your reasoning.",
    "Explain step by step how the quicksort algorithm works, including how partitioning "
    "works and its time complexity in the best, average, and worst cases. Give a small "
    "worked example.",
    "Write a Python function that merges two sorted linked lists into one sorted list, "
    "then explain step by step how it works and analyze its time and space complexity.",
]

# scheme -> per-K / per-V (format, layout, calib) cells + data bit-width.
SCHEMES = {
    "nvfp4-baseline": {"k": ("nvfp4", "headdim", "absmax"), "v": ("nvfp4", "headdim", "absmax"), "bits": 4},
    "int4-kivi":      {"k": ("int4", "channel", "mse"),     "v": ("int4", "headdim", "mse"),     "bits": 4},
    "int3-kivi":      {"k": ("int3", "channel", "mse"),     "v": ("int3", "headdim", "mse"),     "bits": 3},
    "int3-naive":     {"k": ("int3", "headdim", "absmax"),  "v": ("int3", "headdim", "absmax"),  "bits": 3},
}
BASELINE = "nvfp4-baseline"


def mem_ratio(bits: int) -> float:
    """vs BF16: data bits/8 + one 1-byte scale per 16-elem block."""
    return 2.0 / (bits / 8.0 + 1.0 / BLOCK)


class PageSim:
    """Freeze-at-fill quantization on a live DynamicCache: completed pages are
    quantized once and kept; the partial hot page stays BF16."""

    def __init__(self, scheme):
        self.scheme = scheme
        self.n_frozen = 0

    def update(self, cache):
        if self.scheme is None:
            return
        n_pages = cache.layers[0].keys.shape[2] // PAGE
        if n_pages <= self.n_frozen:
            return
        lo, hi = self.n_frozen * PAGE, n_pages * PAGE
        for layer in cache.layers:
            k, v = layer.keys[0], layer.values[0]
            qk = roundtrip(k[:, lo:hi], *self.scheme["k"])
            qv = roundtrip(v[:, lo:hi], *self.scheme["v"])
            layer.keys = torch.cat([k[:, :lo], qk, k[:, hi:]], dim=1).unsqueeze(0)
            layer.values = torch.cat([v[:, :lo], qv, v[:, hi:]], dim=1).unsqueeze(0)
        self.n_frozen = n_pages


def _eos_set(model):
    eos = getattr(model.config, "eos_token_id", None)
    if isinstance(eos, (list, tuple)):
        return set(eos)
    return {eos} if eos is not None else set()


def _prefill(model, input_ids, cache, device):
    pos = torch.arange(input_ids.shape[1], device=device)
    return model(input_ids=input_ids, past_key_values=cache, use_cache=True,
                 cache_position=pos, position_ids=pos.unsqueeze(0))


def _step(model, tok_id, cache, abs_pos, device):
    cp = torch.tensor([abs_pos], device=device)
    return model(input_ids=torch.tensor([[tok_id]], device=device), past_key_values=cache,
                 use_cache=True, cache_position=cp, position_ids=cp.unsqueeze(0))


def gen_bf16(model, input_ids, max_new, device, eos):
    """BF16 greedy; returns (gold_tokens, ref_logits [N,V] cpu, bf16 cache)."""
    cache = DynamicCache()
    logits, toks = [], []
    with torch.no_grad():
        out = _prefill(model, input_ids, cache, device)
        logits.append(out.logits[0, -1].float().cpu())
        toks.append(out.logits[0, -1].argmax().item())
        abs_pos = input_ids.shape[1]
        for _ in range(max_new - 1):
            out = _step(model, toks[-1], cache, abs_pos, device)
            logits.append(out.logits[0, -1].float().cpu())
            toks.append(out.logits[0, -1].argmax().item())
            abs_pos += 1
            if toks[-1] in eos:
                break
    return toks, torch.stack(logits), cache


def teacher_forced(model, input_ids, gold, scheme, device):
    """Replay gold through a frozen-page scheme cache; logits [len(gold), V] cpu."""
    cache = DynamicCache()
    sim = PageSim(scheme)
    logits = []
    with torch.no_grad():
        out = _prefill(model, input_ids, cache, device)
        sim.update(cache)
        logits.append(out.logits[0, -1].float().cpu())
        abs_pos = input_ids.shape[1]
        for t in gold[:-1]:
            out = _step(model, t, cache, abs_pos, device)
            sim.update(cache)
            logits.append(out.logits[0, -1].float().cpu())
            abs_pos += 1
    return torch.stack(logits)


def fidelity(ref, scheme_logits):
    top1 = (scheme_logits.argmax(-1) == ref.argmax(-1)).float().mean().item()
    logp = torch.log_softmax(ref, dim=-1)
    logq = torch.log_softmax(scheme_logits, dim=-1)
    kl = (logp.exp() * (logp - logq)).sum(-1).mean().item()
    return top1, kl


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--max-new", type=int, default=384)
    ap.add_argument("--n-prompts", type=int, default=3)
    args = ap.parse_args()

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

    agg = {n: {"top1": [], "kl": [], "krmse": [], "vrmse": []} for n in SCHEMES}
    for pi, prompt in enumerate(PROMPTS[:args.n_prompts]):
        input_ids = tok.apply_chat_template(
            [{"role": "user", "content": prompt}],
            add_generation_prompt=True, return_tensors="pt", return_dict=False).to(device)
        t0 = time.time()
        gold, ref_logits, cache = gen_bf16(model, input_ids, args.max_new, device, eos)
        ctx = input_ids.shape[1] + len(gold)

        for layer in cache.layers:                       # RMSE on the BF16 cache
            S = layer.keys.shape[2]
            nf = (S // BLOCK) * BLOCK
            K, V = layer.keys[0, :, :nf], layer.values[0, :, :nf]
            for n, s in SCHEMES.items():
                agg[n]["krmse"].append(rmse_cell(K, *s["k"]))
                agg[n]["vrmse"].append(rmse_cell(V, *s["v"]))

        for n, s in SCHEMES.items():                     # teacher-forced KL
            top1, kl = fidelity(ref_logits, teacher_forced(model, input_ids, gold, s, device))
            agg[n]["top1"].append(top1)
            agg[n]["kl"].append(kl)
        print(f"[prompt {pi}] ctx={ctx} tokens, {time.time()-t0:.0f}s", flush=True)

    avg = lambda xs: sum(xs) / max(len(xs), 1)
    base_kl = avg(agg[BASELINE]["kl"])
    print("\n" + "=" * 78)
    print(f"AGGREGATE over {args.n_prompts} prompts  (production-faithful frozen-page protocol)")
    print(f"  {'scheme':<15} {'bits':>4} {'mem×':>5} {'K-RMSE':>8} {'V-RMSE':>8} "
          f"{'top-1':>7} {'KL':>8} {'KL vs base':>11}")
    print(f"  {'-'*15} {'-'*4} {'-'*5} {'-'*8} {'-'*8} {'-'*7} {'-'*8} {'-'*11}")
    for n, s in SCHEMES.items():
        kl = avg(agg[n]["kl"])
        kld = "baseline" if n == BASELINE else f"{100*(base_kl-kl)/max(base_kl,1e-12):+.0f}%"
        print(f"  {n:<15} {s['bits']:>4} {mem_ratio(s['bits']):>4.2f}x "
              f"{avg(agg[n]['krmse']):>8.5f} {avg(agg[n]['vrmse']):>8.5f} "
              f"{100*avg(agg[n]['top1']):>6.1f}% {kl:>8.5f} {kld:>11}")


if __name__ == "__main__":
    main()