"""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 # Plant a needle near the very front of an in-distribution code context. 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) # ---- 1M projection from measured per-token KV ---- 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"] # placeholder # Use measured numbers explicitly. 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()