| """INT4 KV-cache quantization benchmark on Laguna-XS.2. |
| |
| Runs standard greedy generation, then snapshots the final KV cache, quantizes |
| it with MSE-optimal blockwise scaling, and reports memory savings vs BF16. |
| Also measures per-layer reconstruction quality (absmax vs MSE-optimal scale). |
| |
| Usage: |
| python -m scripts.quant_inference [--max-new 256] [--prompt "..."] |
| """ |
| 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 QuantizedKVCache, measure_page_error |
|
|
| 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 main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--max-new", type=int, default=256) |
| ap.add_argument("--prompt", type=str, default=PROMPT) |
| 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() |
| cfg = model.config |
| print(f"[load] {cfg.num_hidden_layers} layers, kv_heads={cfg.num_key_value_heads}, " |
| f"head_dim={cfg.head_dim}, device={next(model.parameters()).device}") |
|
|
| msgs = [{"role": "user", "content": args.prompt}] |
| input_ids = tok.apply_chat_template( |
| msgs, add_generation_prompt=True, return_tensors="pt", return_dict=False, |
| ).to(next(model.parameters()).device) |
| print(f"[prompt] {input_ids.shape[1]} tokens", flush=True) |
|
|
| |
| cache = DynamicCache() |
| t0 = time.time() |
| with torch.no_grad(): |
| out = model.generate( |
| input_ids, |
| max_new_tokens=args.max_new, |
| past_key_values=cache, |
| use_cache=True, |
| do_sample=False, |
| ) |
| gen_t = time.time() - t0 |
| n_gen = out.shape[1] - input_ids.shape[1] |
|
|
| print("\n========== OUTPUT ==========") |
| print(tok.decode(out[0, input_ids.shape[1]:], skip_special_tokens=True)[:1200]) |
|
|
| |
| n_layers = len(cache.layers) |
| bf16_bytes = sum( |
| (layer.keys.numel() + layer.values.numel()) * 2 |
| for layer in cache.layers |
| ) |
| seq_len = cache.layers[0].keys.shape[-2] if n_layers else 0 |
|
|
| qcache = QuantizedKVCache() |
| for i, layer in enumerate(cache.layers): |
| qcache.update(i, layer.keys, layer.values) |
| int4_bytes = qcache.mem_bytes() |
|
|
| |
| |
| sample_layer_idx = 0 |
| sample_k = cache.layers[sample_layer_idx].keys[0] |
| err = measure_page_error(sample_k.float().cpu()) |
|
|
| bf16_mb = bf16_bytes / 1e6 |
| int4_mb = int4_bytes / 1e6 |
| ratio = bf16_bytes / max(int4_bytes, 1) |
|
|
| print("\n========== KV CACHE ==========") |
| print(f" layers: {n_layers}") |
| print(f" seq_len: {seq_len}") |
| print(f" BF16: {bf16_mb:.1f} MB") |
| print(f" INT4 (optimal): {int4_mb:.1f} MB") |
| print(f" ratio: {ratio:.2f}x") |
| print(f" absmax MSE: {err['absmax_mse']:.6f}") |
| print(f" optimal MSE: {err['optimal_mse']:.6f}") |
| print(f" MSE reduction: {err['reduction_pct']:.1f}% (MSE-optimal vs absmax)") |
| print(f"\n generated {n_gen} tokens in {gen_t:.1f}s ({n_gen/gen_t:.1f} tok/s)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
