File size: 4,560 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 | """INT4-vs-NVFP4 KV-cache quantization sweep on Laguna-XS.2.
For each layer's real K and V activations, measures reconstruction RMSE across
the grid {nvfp4, int4} × {headdim-block, per-channel-block} × {absmax, mse}.
The baseline cell (nvfp4 / headdim / absmax) is what vLLM's NVFP4 KV kernel ships
today. Every cell has identical memory (4-bit data + one scale per 16 elements),
so this isolates reconstruction quality. K and V are reported separately to
expose any KIVI-style asymmetry (per-channel K, per-token V).
Usage:
python -m scripts.quant_sweep [--max-new 200] [--n-alphas 32]
"""
from __future__ import annotations
import argparse
import sys
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
sys.path.insert(0, "/home/alex/poolside-hackathon-kv-quant")
from kv_quant import BLOCK, SWEEP_CELLS, SWEEP_BASELINE, sweep_tensor
MODEL = "poolside/Laguna-XS.2"
PROMPT = (
"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."
)
def _cell_label(cell) -> str:
return f"{cell[0]:<6} {cell[1]:<8} {cell[2]:<6}"
def _print_table(title: str, abs_rmse: dict, ratio: dict) -> None:
base = SWEEP_BASELINE
print(f"\n{title}")
print(f" {'format':<6} {'layout':<8} {'calib':<6} {'avg RMSE':>10} {'vs vLLM':>10}")
print(f" {'-'*6} {'-'*8} {'-'*6} {'-'*10} {'-'*10}")
for c in [base] + [c for c in SWEEP_CELLS if c != base]:
tag = "baseline" if c == base else f"{(1.0 - ratio[c]) * 100.0:+.1f}%"
print(f" {_cell_label(c)} {abs_rmse[c]:>10.5f} {tag:>10}")
best = min(SWEEP_CELLS, key=lambda c: ratio[c])
print(f" best: {best[0]}/{best[1]}/{best[2]} "
f"({(1.0 - ratio[best]) * 100.0:+.1f}% vs baseline)")
def run_sweep(per_layer_kv, n_alphas, device):
"""per_layer_kv: list of (K, V) [n_kv, seq, head_dim]. Returns means + ratios."""
base = SWEEP_BASELINE
k_abs = {c: [] for c in SWEEP_CELLS}
k_rat = {c: [] for c in SWEEP_CELLS}
v_abs = {c: [] for c in SWEEP_CELLS}
v_rat = {c: [] for c in SWEEP_CELLS}
n_used = 0
for K, V in per_layer_kv:
if K.shape[1] < BLOCK:
continue
n_used += 1
ks = sweep_tensor(K.to(device), n_alphas)
vs = sweep_tensor(V.to(device), n_alphas)
for c in SWEEP_CELLS:
k_abs[c].append(ks[c]); k_rat[c].append(ks[c] / max(ks[base], 1e-12))
v_abs[c].append(vs[c]); v_rat[c].append(vs[c] / max(vs[base], 1e-12))
mean = lambda d: {c: sum(x) / max(len(x), 1) for c, x in d.items()}
return n_used, mean(k_abs), mean(k_rat), mean(v_abs), mean(v_rat)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--max-new", type=int, default=200)
ap.add_argument("--n-alphas", type=int, default=32)
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
cfg = model.config
print(f"[load] layers={cfg.num_hidden_layers} kv_heads={cfg.num_key_value_heads} "
f"head_dim={getattr(cfg, 'head_dim', None)} device={device}", flush=True)
msgs = [{"role": "user", "content": PROMPT}]
input_ids = tok.apply_chat_template(
msgs, add_generation_prompt=True, return_tensors="pt", return_dict=False,
).to(device)
cache = DynamicCache()
with torch.no_grad():
model.generate(input_ids, max_new_tokens=args.max_new,
past_key_values=cache, use_cache=True, do_sample=False)
per_layer = [(layer.keys[0].float().cpu(), layer.values[0].float().cpu())
for layer in cache.layers]
seq = per_layer[0][0].shape[1] if per_layer else 0
print(f"[seq] {seq} tokens cached across {len(per_layer)} layers", flush=True)
n_used, k_abs, k_rat, v_abs, v_rat = run_sweep(per_layer, args.n_alphas, device)
print(f"\n{'='*56}")
print(f"INT4 vs NVFP4 KV sweep — {n_used} layers, seq cropped to /{BLOCK}")
print("All cells: identical memory (4-bit + 1 scale per 16 elems, ~3.56x vs BF16).")
_print_table("KEY cache (avg RMSE over layers):", k_abs, k_rat)
_print_table("VALUE cache (avg RMSE over layers):", v_abs, v_rat)
if __name__ == "__main__":
main()
|