feat: adding benchmark for longer context

benchmark_long_context.py

@@ -0,0 +1,124 @@
+"""
+Long context benchmarks at 16K and 32K.
+This is where KV cache compression matters most.
+"""
+import torch
+import json
+import os
+import sys
+import time
+import math
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from datasets import load_dataset
+
+sys.path.append(os.path.expanduser("~/kv-hack"))
+from kernel.quant_cache import MixedPrecisionKVCache
+
+MODEL_NAME = sys.argv[1] if len(sys.argv) > 1 else "mistral-7b"
+MODEL_PATHS = {
+    "mistral-7b": "~/kv-hack/mistral-model",
+    "llama-3-8b": "~/kv-hack/llama-model",
+}
+model_path  = os.path.expanduser(MODEL_PATHS[MODEL_NAME])
+results_dir = os.path.expanduser(f"~/kv-hack/results/{MODEL_NAME}")
+
+with open(f"{results_dir}/bit_allocation.json") as f:
+    raw = json.load(f)
+bit_alloc = {
+    int(l): [raw[l][str(h)] for h in range(len(raw[l]))]
+    for l in raw
+}
+num_layers = len(bit_alloc)
+avg_bits   = sum(b for l in bit_alloc.values() for b in l) / \
+             sum(len(l) for l in bit_alloc.values())
+
+print(f"Model: {MODEL_NAME}")
+print(f"Avg bits: {avg_bits:.2f}")
+
+print("Loading model...")
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model     = AutoModelForCausalLM.from_pretrained(
+    model_path, dtype=torch.float16, device_map="cuda"
+)
+model.eval()
+print(f"Loaded: {torch.cuda.memory_allocated()/1e9:.2f} GB")
+
+def measure_context(context_len: int):
+    print(f"\n  Context {context_len} tokens...")
+    input_ids = torch.randint(1, 1000, (1, context_len)).cuda()
+
+    # peak memory
+    torch.cuda.reset_peak_memory_stats()
+    with torch.no_grad():
+        out = model(input_ids, use_cache=True)
+        kv  = out.past_key_values
+    torch.cuda.synchronize()
+    peak_mem = torch.cuda.max_memory_allocated() / 1e9
+
+    # KV compression
+    fp16_bytes       = 0
+    uniform8_bytes   = 0
+    compressed_bytes = 0
+
+    for layer_idx in range(num_layers):
+        k = kv.layers[layer_idx].keys
+        v = kv.layers[layer_idx].values
+        fp16_bytes     += k.numel() * 2 + v.numel() * 2
+        uniform8_bytes += k.numel() + v.numel()
+        cache = MixedPrecisionKVCache(bit_alloc[layer_idx])
+        cache.store(k, v)
+        compressed_bytes += cache.memory_bytes()
+
+    # prefill speed
+    times = []
+    for _ in range(3):
+        torch.cuda.synchronize()
+        t0 = time.time()
+        with torch.no_grad():
+            _ = model(input_ids, use_cache=True)
+        torch.cuda.synchronize()
+        times.append(time.time() - t0)
+    prefill_ms = round(sum(times) / len(times) * 1000, 1)
+
+    return {
+        "context_len":          context_len,
+        "peak_memory_gb":       round(peak_mem, 2),
+        "fp16_mb":              round(fp16_bytes / 1e6, 2),
+        "uniform8_mb":          round(uniform8_bytes / 1e6, 2),
+        "mixed_precision_mb":   round(compressed_bytes / 1e6, 2),
+        "compression_vs_fp16":  round(fp16_bytes / compressed_bytes, 2),
+        "compression_vs_8bit":  round(uniform8_bytes / compressed_bytes, 2),
+        "prefill_ms":           prefill_ms,
+    }
+
+print("\n" + "="*60)
+print("LONG CONTEXT BENCHMARK")
+print("="*60)
+
+results = []
+for ctx in [512, 1024, 2048, 4096, 8192, 16384, 32768]:
+    try:
+        r = measure_context(ctx)
+        results.append(r)
+        print(f"  ctx={ctx:6d} | "
+              f"mem={r['peak_memory_gb']:.2f}GB | "
+              f"FP16={r['fp16_mb']:.0f}MB | "
+              f"Ours={r['mixed_precision_mb']:.0f}MB | "
+              f"{r['compression_vs_fp16']}x | "
+              f"prefill={r['prefill_ms']}ms")
+    except torch.cuda.OutOfMemoryError:
+        print(f"  ctx={ctx:6d} | OOM — FP16 ran out of memory ✓")
+        # still measure our compressed version
+        results.append({
+            "context_len": ctx,
+            "peak_memory_gb": "OOM",
+            "fp16_mb": ctx * num_layers * 8 * 128 * 4 / 1e6,
+            "note": "FP16 OOM, compressed might fit"
+        })
+        break
+
+# save
+out_path = f"{results_dir}/long_context_results.json"
+with open(out_path, "w") as f:
+    json.dump({"model": MODEL_NAME, "results": results}, f, indent=2)
+print(f"\n✅ Saved to {out_path}")