File size: 6,213 Bytes
6e668dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | """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] # [S, V], bf16
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()
|