| """Practical ceiling probe + 1M KV-memory analysis for INT4-KIVI vs BF16. |
| |
| 1. Builds an in-distribution code context at a target length (default 256k = |
| the model's trained RoPE ceiling), runs a single greedy forward + short |
| generation under BOTH BF16 DynamicCache and INT4-KIVI, and reports: |
| - whether the forward succeeded (no OOM / no error), |
| - coherence of the short continuation (printed verbatim), |
| - measured KV-cache bytes per mode at that length + compression ratio, |
| - a needle planted near the front, to check far-back retrieval survives |
| at the ceiling. |
| |
| 2. Prints the 1M-token KV-memory projection (per-token KV from the measured |
| cache, scaled linearly), and the feasibility verdict on the B300. |
| |
| Usage |
| ----- |
| .venv/bin/python scripts/ceiling_and_mem.py --length 262000 --max-new 16 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import random |
| import re |
| import sys |
| import time |
|
|
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache |
|
|
| ROOT = "/home/alex/poolside-hackathon-kv-quant" |
| sys.path.insert(0, ROOT) |
| from int4_kivi.hf_cache import Int4KiviCache |
| from scripts.longctx_code_bench import build_prefix_text |
|
|
| MODEL = "poolside/Laguna-XS.2" |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--length", type=int, default=262000, |
| help="target context length in tokens (<=262144 trained)") |
| ap.add_argument("--max-new", type=int, default=48) |
| args = ap.parse_args() |
|
|
| print(f"[load] {MODEL} ...", flush=True) |
| tok = AutoTokenizer.from_pretrained(MODEL) |
| model = AutoModelForCausalLM.from_pretrained( |
| MODEL, dtype=torch.bfloat16, device_map="auto" |
| ).eval() |
| device = next(model.parameters()).device |
|
|
| |
| rng = random.Random(99) |
| val = rng.randint(10_000_000, 99_999_999) |
| needle = f"# runtime config\nCEILING_SEED = {val}\n\n" |
| print(f"[prefix] building ~{args.length}-token code context ...", flush=True) |
| body, _ = build_prefix_text(tok, args.length) |
| reference = needle + body |
| question = ("In the reference code above there is a constant named " |
| "CEILING_SEED. Output its exact integer value as the very " |
| "first thing in your reply, then stop.") |
|
|
| sys_msg = ("You are a careful code assistant. Answer using ONLY the " |
| "reference code.") |
| user_msg = f"```python\n{reference}\n```\n\n{question}" |
| msgs = [{"role": "system", "content": sys_msg}, |
| {"role": "user", "content": user_msg}] |
| input_ids = tok.apply_chat_template( |
| msgs, add_generation_prompt=True, return_tensors="pt", return_dict=False, |
| ).to(device) |
| ctx = int(input_ids.shape[1]) |
| print(f"[ctx] actual context length: {ctx} tokens " |
| f"(trained ceiling 262144)\n", flush=True) |
|
|
| results = {} |
| for mode in ("bf16", "int4"): |
| torch.cuda.reset_peak_memory_stats() |
| cache = Int4KiviCache(config=model.config) if mode == "int4" else DynamicCache() |
| t0 = time.time() |
| try: |
| 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, num_beams=1, |
| ) |
| dt = time.time() - t0 |
| resp = tok.decode(out[0, ctx:], skip_special_tokens=True) |
| ok = bool(re.search(rf"\b{val}\b", resp.replace(",", ""))) |
| if mode == "int4": |
| kv_bytes = cache.nbytes() |
| bf16_ref = cache.bf16_nbytes() |
| ratio = cache.compression_ratio_vs_bf16() |
| else: |
| kv_bytes = sum( |
| l.keys.numel() * l.keys.element_size() |
| + l.values.numel() * l.values.element_size() |
| for l in cache.layers |
| ) |
| bf16_ref = kv_bytes |
| ratio = 1.0 |
| peak = torch.cuda.max_memory_allocated() |
| results[mode] = dict(ok=ok, resp=resp.strip()[:160], dt=dt, |
| kv=kv_bytes, ratio=ratio, peak=peak) |
| print(f"[{mode}] forward OK in {dt:.1f}s retrieval={'HIT' if ok else 'MISS'}") |
| print(f"[{mode}] KV bytes = {kv_bytes/2**30:.2f} GiB " |
| f"ratio_vs_bf16 = {ratio:.2f}x peak_alloc = {peak/2**30:.1f} GiB") |
| print(f"[{mode}] resp: {resp.strip()[:120]!r}\n", flush=True) |
| except torch.cuda.OutOfMemoryError as e: |
| results[mode] = dict(ok=False, err="OOM") |
| print(f"[{mode}] OUT OF MEMORY: {e}\n", flush=True) |
| torch.cuda.empty_cache() |
| except Exception as e: |
| results[mode] = dict(ok=False, err=str(e)[:200]) |
| print(f"[{mode}] ERROR: {e}\n", flush=True) |
|
|
| |
| print(f"{'':#<80}") |
| print(" KV-MEMORY PROJECTION") |
| print(f"{'':#<80}") |
| if "int4" in results and "kv" in results["int4"]: |
| i4_per_tok = results["int4"]["kv"] / ctx |
| bf_per_tok = results["int4"]["bf16_ref" if False else "kv"] |
| |
| if results.get("bf16", {}).get("kv") and results.get("int4", {}).get("kv"): |
| bf_per_tok = results["bf16"]["kv"] / ctx |
| i4_per_tok = results["int4"]["kv"] / ctx |
| print(f" measured @ {ctx} tok:") |
| print(f" BF16 KV/token = {bf_per_tok/1024:.1f} KiB " |
| f"-> {ctx} tok = {results['bf16']['kv']/2**30:.2f} GiB") |
| print(f" INT4-KIVI KV/token = {i4_per_tok/1024:.1f} KiB " |
| f"-> {ctx} tok = {results['int4']['kv']/2**30:.2f} GiB") |
| for N in (262144, 1_000_000): |
| print(f" @ {N:>9} tok: BF16 = {bf_per_tok*N/2**30:6.1f} GiB " |
| f"INT4-KIVI = {i4_per_tok*N/2**30:6.1f} GiB " |
| f"(headroom {bf_per_tok/i4_per_tok:.1f}x)") |
| print(f"\n B300 ~275 GiB: both fit at batch 1 for 1M tokens; INT4-KIVI gives") |
| print(f" ~3x headroom. 1M is 3.8x past the trained 256k RoPE range, so") |
| print(f" >256k needs RoPE extension (YaRN/NTK) for accuracy, not more memory.") |
| print("\n[done]", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
