| """Long-context test: does the INT4-KIVI vs NVFP4-baseline gap grow with context? |
| |
| Scores long in-distribution sequences (real source code, in-distribution for a |
| code model) in a SINGLE forward per scheme. KV quantization is injected by |
| patching scaled_dot_product_attention so position t attends to a quantized |
| prefix. Per-position KL(bf16||scheme) and top-1 agreement are binned by sequence |
| position, revealing whether quantization error compounds as context lengthens. |
| |
| All cached K/V are quantized (the <16-token bf16 hot page is omitted — negligible |
| at long context, and a slightly conservative choice). |
| |
| Usage: |
| python -m scripts.quant_longctx [--ctx 8000] [--windows 3] |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import glob |
| import sys |
|
|
| import torch |
| import torch.nn.functional as F |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| sys.path.insert(0, "/home/alex/poolside-hackathon-kv-quant") |
| from kv_quant import roundtrip |
|
|
| MODEL = "poolside/Laguna-XS.2" |
| N_ALPHAS = 32 |
|
|
| SCHEMES = { |
| "nvfp4-baseline": {"k": ("nvfp4", "headdim", "absmax"), "v": ("nvfp4", "headdim", "absmax")}, |
| "int4-kivi": {"k": ("int4", "channel", "mse"), "v": ("int4", "headdim", "mse")}, |
| "int3-kivi": {"k": ("int3", "channel", "mse"), "v": ("int3", "headdim", "mse")}, |
| "int3-naive": {"k": ("int3", "headdim", "absmax"), "v": ("int3", "headdim", "absmax")}, |
| } |
| BASELINE = "nvfp4-baseline" |
| BINS = [(0, 512), (512, 1024), (1024, 2048), (2048, 4096), (4096, 8192)] |
|
|
| _ORIG_SDPA = F.scaled_dot_product_attention |
| _SCHEME = None |
| _HITS = 0 |
|
|
|
|
| def _q_per_head(x, cell): |
| """Quantize [H, S, D] head-by-head (caps the MSE-search peak memory).""" |
| return torch.stack([roundtrip(x[h:h + 1], *cell, n_alphas=N_ALPHAS)[0] |
| for h in range(x.shape[0])]) |
|
|
|
|
| def _patched_sdpa(query, key, value, *a, **kw): |
| global _HITS |
| if _SCHEME is not None: |
| _HITS += 1 |
| key = _q_per_head(key[0], _SCHEME["k"]).unsqueeze(0) |
| value = _q_per_head(value[0], _SCHEME["v"]).unsqueeze(0) |
| return _ORIG_SDPA(query, key, value, *a, **kw) |
|
|
|
|
| def token_pool(tok): |
| files = sorted(glob.glob( |
| "/home/alex/poolside-hackathon-kv-quant/.venv/**/transformers/**/modeling_*.py", |
| recursive=True)) |
| texts, total = [], 0 |
| for f in files: |
| try: |
| t = open(f).read() |
| except OSError: |
| continue |
| texts.append(t) |
| total += len(t) |
| if total > 600_000: |
| break |
| ids = tok("\n\n".join(texts), return_tensors="pt").input_ids[0] |
| return ids |
|
|
|
|
| @torch.no_grad() |
| def logits_of(model, ids): |
| return model(input_ids=ids.unsqueeze(0)).logits[0] |
|
|
|
|
| def kl_top1(ref, lg, ref_arg, chunk=512): |
| """Per-position KL(bf16||scheme) and top-1 match. ref/lg bf16 [S,V].""" |
| S = ref.shape[0] |
| kls, t1 = [], [] |
| for i in range(0, S, chunk): |
| r, q = ref[i:i + chunk].float(), lg[i:i + chunk].float() |
| rlp, qlp = torch.log_softmax(r, -1), torch.log_softmax(q, -1) |
| kls.append((rlp.exp() * (rlp - qlp)).sum(-1)) |
| t1.append((q.argmax(-1) == ref_arg[i:i + chunk]).float()) |
| return torch.cat(kls), torch.cat(t1) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ctx", type=int, default=8000) |
| ap.add_argument("--windows", type=int, default=3) |
| args = ap.parse_args() |
| global _SCHEME, _HITS |
|
|
| F.scaled_dot_product_attention = _patched_sdpa |
| torch.nn.functional.scaled_dot_product_attention = _patched_sdpa |
|
|
| print(f"[load] {MODEL}", flush=True) |
| tok = AutoTokenizer.from_pretrained(MODEL) |
| model = AutoModelForCausalLM.from_pretrained( |
| MODEL, torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="sdpa") |
| model.eval() |
| device = next(model.parameters()).device |
|
|
| pool = token_pool(tok) |
| ctx = (args.ctx // 16) * 16 |
| nwin = min(args.windows, pool.shape[0] // ctx) |
| print(f"[seq] {nwin} windows x {ctx} tokens (pool={pool.shape[0]})", flush=True) |
|
|
| bins = [(lo, hi) for lo, hi in BINS if lo < ctx] |
| acc = {n: {b: {"kl": [], "t1": []} for b in bins} for n in SCHEMES} |
|
|
| for w in range(nwin): |
| ids = pool[w * ctx:(w + 1) * ctx].to(device) |
| _SCHEME = None |
| ref = logits_of(model, ids) |
| ref_arg = ref.argmax(-1) |
| for name, sch in SCHEMES.items(): |
| _HITS = 0 |
| _SCHEME = sch |
| lg = logits_of(model, ids) |
| _SCHEME = None |
| assert _HITS > 0, f"SDPA patch never fired for {name} — wrong attn path" |
| kl, t1 = kl_top1(ref, lg, ref_arg) |
| for lo, hi in bins: |
| acc[name][(lo, hi)]["kl"].append(kl[lo:min(hi, ctx)].mean().item()) |
| acc[name][(lo, hi)]["t1"].append(t1[lo:min(hi, ctx)].mean().item()) |
| del lg |
| print(f"[window {w}] done", flush=True) |
|
|
| avg = lambda xs: sum(xs) / max(len(xs), 1) |
| print("\n" + "=" * 74) |
| print(f"PER-POSITION KL(bf16||scheme), avg over {nwin} windows") |
| print(f" {'position':<12}" + "".join(f"{n[:13]:>14}" for n in SCHEMES)) |
| print(f" {'-'*12}" + "".join(f" {'-'*13}" for _ in SCHEMES)) |
| for lo, hi in bins: |
| row = f" {f'{lo}-{min(hi,ctx)}':<12}" |
| for n in SCHEMES: |
| row += f"{avg(acc[n][(lo,hi)]['kl']):>14.5f}" |
| print(row) |
|
|
| print(f"\nKIVI vs baseline: KL reduction by position (does the win grow?)") |
| print(f" {'position':<12}{'int4-kivi':>14}{'int3-kivi':>14}") |
| print(f" {'-'*12}{' '+'-'*13}{' '+'-'*13}") |
| for lo, hi in bins: |
| b = avg(acc[BASELINE][(lo, hi)]["kl"]) |
| i4 = avg(acc["int4-kivi"][(lo, hi)]["kl"]) |
| i3 = avg(acc["int3-kivi"][(lo, hi)]["kl"]) |
| row = f" {f'{lo}-{min(hi,ctx)}':<12}" |
| row += f"{100*(b-i4)/max(b,1e-12):>13.0f}%" |
| row += f"{100*(b-i3)/max(b,1e-12):>13.0f}%" |
| print(row) |
|
|
| print(f"\nTOP-1 agreement vs bf16 by position") |
| print(f" {'position':<12}" + "".join(f"{n[:13]:>14}" for n in SCHEMES)) |
| for lo, hi in bins: |
| row = f" {f'{lo}-{min(hi,ctx)}':<12}" |
| for n in SCHEMES: |
| row += f"{100*avg(acc[n][(lo,hi)]['t1']):>13.1f}%" |
| print(row) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
