""" Kalpana RIF — Real Empirical Benchmark Harness ================================================ Measures ACTUAL GPU memory, latency, and recall at multiple context lengths. Compares: Standard DynamicCache vs KalpanaDynamicCache vs SinkCache (StreamingLLM). All numbers are measured, not estimated. CRITICAL NOTE on what is measured: - persistent_cache_mb: The stored cache state size (O(1) for Kalpana) - peak_vram_mb: PEAK GPU allocation including intermediate tensors during forward pass — this includes reconstruction intermediates for Kalpana - prefill_time_s: Wall clock to process all input tokens - ttft_ms: Time to generate the FIRST output token after prefill - avg_token_ms: Average time per generated token - reconstruction_cosine_sim: Cosine similarity of Kalpana's reconstructed K/V vs ground-truth standard cache K/V (measures information loss) """ import torch import torch.nn.functional as F import time import json import gc import os import traceback MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" # --------------------------------------------------------------------------- # Haystack builder: long filler text with a planted "needle" fact # --------------------------------------------------------------------------- FILLER = ( "System telemetry block {i}: harmonic sensor reading at {f:.4f} MHz " "with phase offset {p} degrees in monitoring sector {s}. " "All parameters within nominal operating range. " ) NEEDLE_TEMPLATE = ( "CRITICAL CLASSIFIED FINDING: The secret authorization passkey " "for Project Nightingale is {code}. This information is top-secret. " ) NEEDLE_QUERY = ( "What is the secret authorization passkey for Project Nightingale? " "Reply with ONLY the passkey code, nothing else." ) def build_haystack(tokenizer, target_tokens, needle_code, needle_depth_pct=0.5): """Build input_ids with a needle fact embedded at specified depth percentage.""" # Generate filler chunks chunks = [] for i in range(30000): chunks.append(FILLER.format(i=i, f=i * 0.31416, p=(i * 37) % 360, s=i % 16)) # Estimate tokens per filler chunk sample_enc = tokenizer.encode(chunks[0], add_special_tokens=False) toks_per_chunk = max(1, len(sample_enc)) # Calculate chunks needed (leave room for needle + query + template) overhead_tokens = 120 # chat template + query + needle content_tokens = max(10, target_tokens - overhead_tokens) n_chunks = max(1, content_tokens // toks_per_chunk) # Insert needle at target depth needle_idx = max(0, int(n_chunks * needle_depth_pct)) needle_text = NEEDLE_TEMPLATE.format(code=needle_code) chunks_to_use = chunks[:n_chunks] chunks_to_use.insert(needle_idx, needle_text) context = " ".join(chunks_to_use) full_prompt = context + "\n\nQuestion: " + NEEDLE_QUERY messages = [{"role": "user", "content": full_prompt}] formatted = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) input_ids = tokenizer( formatted, return_tensors="pt", truncation=True, max_length=target_tokens ).input_ids return input_ids # --------------------------------------------------------------------------- # Core measurement function # --------------------------------------------------------------------------- def measure_one(model, tokenizer, input_ids, cache, cache_name, device, num_gen=10): """ Measure one benchmark point: prefill + generation. Returns dict with all measured metrics. """ N = input_ids.shape[1] # Clean slate gc.collect() torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats(device) baseline_vram = torch.cuda.memory_allocated(device) # === PREFILL === t_prefill_start = time.perf_counter() try: with torch.inference_mode(): out = model( input_ids.to(device), past_key_values=cache, use_cache=True ) torch.cuda.synchronize() except Exception as e: gc.collect() torch.cuda.empty_cache() return { "cache_type": cache_name, "context_length": N, "error": f"Prefill failed: {type(e).__name__}: {e}", } t_prefill_end = time.perf_counter() peak_vram_prefill = torch.cuda.max_memory_allocated(device) alloc_after_prefill = torch.cuda.memory_allocated(device) # Persistent cache size if hasattr(cache, "get_total_memory_mb"): persist_mb = cache.get_total_memory_mb() elif hasattr(cache, "key_cache"): b = 0 for t in getattr(cache, "key_cache", []): if isinstance(t, torch.Tensor): b += t.nelement() * t.element_size() for t in getattr(cache, "value_cache", []): if isinstance(t, torch.Tensor): b += t.nelement() * t.element_size() persist_mb = b / (1024 * 1024) elif hasattr(cache, "layers"): b = 0 for layer in cache.layers: if hasattr(layer, "keys") and isinstance(layer.keys, torch.Tensor): b += layer.keys.nelement() * layer.keys.element_size() if hasattr(layer, "values") and isinstance(layer.values, torch.Tensor): b += layer.values.nelement() * layer.values.element_size() persist_mb = b / (1024 * 1024) else: persist_mb = -1 # === GENERATION (token by token) === torch.cuda.reset_peak_memory_stats(device) nxt = out.logits[:, -1:, :].argmax(dim=-1) generated_ids = [] gen_times = [] for _ in range(num_gen): t0g = time.perf_counter() try: with torch.inference_mode(): out = model(nxt, past_key_values=cache, use_cache=True) torch.cuda.synchronize() except Exception: break gen_times.append(time.perf_counter() - t0g) nxt = out.logits[:, -1:, :].argmax(dim=-1) generated_ids.append(nxt.item()) peak_vram_gen = torch.cuda.max_memory_allocated(device) gen_text = tokenizer.decode(generated_ids, skip_special_tokens=True) del out, nxt return { "cache_type": cache_name, "context_length": N, "persistent_cache_mb": round(persist_mb, 3), "peak_vram_prefill_mb": round(peak_vram_prefill / (1024 ** 2), 2), "peak_vram_generation_mb": round(peak_vram_gen / (1024 ** 2), 2), "vram_delta_after_prefill_mb": round( (alloc_after_prefill - baseline_vram) / (1024 ** 2), 2 ), "prefill_time_s": round(t_prefill_end - t_prefill_start, 4), "prefill_tok_per_s": round(N / max(1e-6, t_prefill_end - t_prefill_start), 1), "ttft_ms": round(gen_times[0] * 1000, 2) if gen_times else None, "avg_token_ms": round( sum(gen_times) / max(1, len(gen_times)) * 1000, 2 ) if gen_times else None, "tokens_generated": len(generated_ids), "generated_text": gen_text[:300], } # --------------------------------------------------------------------------- # Reconstruction fidelity: compare Kalpana K/V vs ground-truth # --------------------------------------------------------------------------- def measure_reconstruction_fidelity(model, tokenizer, input_ids, device, num_layers): """ Compare K/V tensors from standard DynamicCache vs KalpanaDynamicCache. Returns per-layer cosine similarity. """ from transformers import DynamicCache from kalpana_embed_to_kv import KalpanaDynamicCache N = input_ids.shape[1] # Run standard gc.collect() torch.cuda.empty_cache() std_cache = DynamicCache() with torch.inference_mode(): model(input_ids.to(device), past_key_values=std_cache, use_cache=True) torch.cuda.synchronize() # Capture standard K/V if hasattr(std_cache, "key_cache"): std_keys = [k.detach().clone() for k in getattr(std_cache, "key_cache", []) if isinstance(k, torch.Tensor)] std_vals = [v.detach().clone() for v in getattr(std_cache, "value_cache", []) if isinstance(v, torch.Tensor)] elif hasattr(std_cache, "layers"): std_keys = [layer.keys.detach().clone() for layer in std_cache.layers if hasattr(layer, "keys") and isinstance(layer.keys, torch.Tensor)] std_vals = [layer.values.detach().clone() for layer in std_cache.layers if hasattr(layer, "values") and isinstance(layer.values, torch.Tensor)] else: std_keys, std_vals = [], [] del std_cache gc.collect() torch.cuda.empty_cache() # Run Kalpana kal_cache = KalpanaDynamicCache( num_layers=num_layers, bands=2048, sliding_window=128 ) with torch.inference_mode(): model(input_ids.to(device), past_key_values=kal_cache, use_cache=True) torch.cuda.synchronize() kal_keys = [k.detach().clone() for k in kal_cache.key_cache] kal_vals = [v.detach().clone() for v in kal_cache.value_cache] del kal_cache gc.collect() torch.cuda.empty_cache() # Compare layer_sims = [] for layer_idx in range(min(len(std_keys), len(kal_keys))): sk = std_keys[layer_idx].float().flatten() kk = kal_keys[layer_idx].float().flatten() sv = std_vals[layer_idx].float().flatten() kv = kal_vals[layer_idx].float().flatten() # Shapes might differ if Kalpana hybrid has window + prefix min_len_k = min(sk.shape[0], kk.shape[0]) min_len_v = min(sv.shape[0], kv.shape[0]) key_sim = F.cosine_similarity(sk[:min_len_k].unsqueeze(0), kk[:min_len_k].unsqueeze(0)).item() val_sim = F.cosine_similarity(sv[:min_len_v].unsqueeze(0), kv[:min_len_v].unsqueeze(0)).item() layer_sims.append({ "layer": layer_idx, "key_cosine_sim": round(key_sim, 6), "val_cosine_sim": round(val_sim, 6), "std_key_shape": list(std_keys[layer_idx].shape), "kal_key_shape": list(kal_keys[layer_idx].shape), }) del std_keys, std_vals, kal_keys, kal_vals gc.collect() torch.cuda.empty_cache() avg_key_sim = sum(l["key_cosine_sim"] for l in layer_sims) / max(1, len(layer_sims)) avg_val_sim = sum(l["val_cosine_sim"] for l in layer_sims) / max(1, len(layer_sims)) return { "context_length": N, "avg_key_cosine_sim": round(avg_key_sim, 6), "avg_val_cosine_sim": round(avg_val_sim, 6), "per_layer": layer_sims, } # --------------------------------------------------------------------------- # Main benchmark runner # --------------------------------------------------------------------------- def run_benchmark( context_lengths=None, num_gen_tokens=10, run_fidelity=True, fidelity_lengths=None, ): """ Run the full benchmark suite. Args: context_lengths: list of int, token counts to test (default: [128..4096]) num_gen_tokens: how many tokens to generate per test run_fidelity: whether to run reconstruction fidelity comparison fidelity_lengths: context lengths for fidelity test (default: [128, 256, 512]) Returns: dict with metadata and results """ if context_lengths is None: context_lengths = [128, 256, 512, 1024, 2048, 4096] if fidelity_lengths is None: fidelity_lengths = [128, 256, 512] device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.float16 if device == "cuda" else torch.float32 gpu_name = torch.cuda.get_device_name(0) if device == "cuda" else "CPU" total_vram = ( torch.cuda.get_device_properties(0).total_memory / (1024 ** 3) if device == "cuda" else 0 ) from transformers import AutoModelForCausalLM, AutoTokenizer print(f"[Benchmark] Loading {MODEL_NAME}...") tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) if tokenizer.pad_token_id is None: tokenizer.pad_token_id = tokenizer.eos_token_id model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, torch_dtype=dtype, low_cpu_mem_usage=True ).to(device) model.eval() model_vram = ( torch.cuda.memory_allocated(device) / (1024 ** 2) if device == "cuda" else 0 ) num_layers = getattr(model.config, "num_hidden_layers", 24) num_kv_heads = getattr(model.config, "num_key_value_heads", 2) head_dim = getattr(model.config, "head_dim", 64) elem_bytes = 2 if dtype == torch.float16 else 4 # Theoretical KV bytes per token for standard cache kv_bytes_per_token = num_layers * num_kv_heads * head_dim * 2 * elem_bytes # Theoretical Kalpana persistent state size # layers * (K+V) * heads * bands * dim * (real+imag) * fp32 kalpana_state_bytes = num_layers * 2 * num_kv_heads * 2048 * head_dim * 2 * 4 kalpana_state_mb = kalpana_state_bytes / (1024 ** 2) meta = { "gpu": gpu_name, "total_vram_gb": round(total_vram, 1), "model": MODEL_NAME, "model_vram_mb": round(model_vram, 1), "num_layers": num_layers, "num_kv_heads": num_kv_heads, "head_dim": head_dim, "dtype": str(dtype), "kv_bytes_per_token_standard": kv_bytes_per_token, "kalpana_theoretical_state_mb": round(kalpana_state_mb, 2), "timestamp": time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime()), } print(f"[Benchmark] GPU: {gpu_name}, VRAM: {total_vram:.1f} GB") print(f"[Benchmark] Model VRAM: {model_vram:.1f} MB") print(f"[Benchmark] KV bytes/token (standard): {kv_bytes_per_token}") print(f"[Benchmark] Kalpana theoretical state: {kalpana_state_mb:.2f} MB") needle_code = "NIGHTINGALE-7749" results = [] # ── Main scaling benchmark ── for ctx_len in context_lengths: print(f"\n{'=' * 60}") print(f"CONTEXT LENGTH: {ctx_len} tokens") print(f"{'=' * 60}") input_ids = build_haystack(tokenizer, ctx_len, needle_code, needle_depth_pct=0.5) actual = input_ids.shape[1] print(f" Actual input tokens: {actual}") # --- Standard DynamicCache --- print(" [1/3] Standard DynamicCache...") from transformers import DynamicCache cache = DynamicCache() r = measure_one(model, tokenizer, input_ids, cache, "Standard_DynamicCache", device, num_gen_tokens) r["needle_code"] = needle_code r["needle_found"] = needle_code.lower() in r.get("generated_text", "").lower() r["theoretical_cache_mb"] = round(actual * kv_bytes_per_token / (1024 ** 2), 3) results.append(r) del cache gc.collect() torch.cuda.empty_cache() print(f" prefill={r.get('prefill_time_s')}s peak={r.get('peak_vram_prefill_mb')}MB cache={r.get('persistent_cache_mb')}MB needle={r.get('needle_found')}") # --- KalpanaDynamicCache --- print(" [2/3] KalpanaDynamicCache (bands=2048, window=128)...") try: from kalpana_embed_to_kv import KalpanaDynamicCache cache = KalpanaDynamicCache( num_layers=num_layers, bands=2048, sliding_window=128 ) r = measure_one(model, tokenizer, input_ids, cache, "Kalpana_RIF", device, num_gen_tokens) r["needle_code"] = needle_code r["needle_found"] = needle_code.lower() in r.get("generated_text", "").lower() r["kalpana_theoretical_state_mb"] = round(kalpana_state_mb, 3) results.append(r) del cache gc.collect() torch.cuda.empty_cache() print(f" prefill={r.get('prefill_time_s')}s peak={r.get('peak_vram_prefill_mb')}MB persist={r.get('persistent_cache_mb')}MB needle={r.get('needle_found')}") except Exception as e: err_r = { "cache_type": "Kalpana_RIF", "context_length": actual, "error": f"{type(e).__name__}: {e}", } results.append(err_r) print(f" ERROR: {e}") gc.collect() torch.cuda.empty_cache() # --- SinkCache (StreamingLLM) --- print(" [3/3] SinkCache (StreamingLLM, window=128, sinks=4)...") try: from transformers import SinkCache cache = SinkCache(window_length=128, num_sink_tokens=4) r = measure_one(model, tokenizer, input_ids, cache, "SinkCache_StreamingLLM", device, num_gen_tokens) r["needle_code"] = needle_code r["needle_found"] = needle_code.lower() in r.get("generated_text", "").lower() results.append(r) del cache gc.collect() torch.cuda.empty_cache() print(f" prefill={r.get('prefill_time_s')}s peak={r.get('peak_vram_prefill_mb')}MB cache={r.get('persistent_cache_mb')}MB needle={r.get('needle_found')}") except ImportError: results.append({ "cache_type": "SinkCache_StreamingLLM", "context_length": actual, "error": "SinkCache not available in this transformers version", }) print(" SKIPPED (SinkCache not available)") except Exception as e: results.append({ "cache_type": "SinkCache_StreamingLLM", "context_length": actual, "error": f"{type(e).__name__}: {e}", }) print(f" ERROR: {e}") gc.collect() torch.cuda.empty_cache() # ── Reconstruction fidelity test ── fidelity_results = [] if run_fidelity: print(f"\n{'=' * 60}") print("RECONSTRUCTION FIDELITY TEST") print(f"{'=' * 60}") for fl in fidelity_lengths: if fl > max(context_lengths): continue print(f" Fidelity test at {fl} tokens...") try: input_ids = build_haystack(tokenizer, fl, needle_code) fr = measure_reconstruction_fidelity( model, tokenizer, input_ids, device, num_layers ) fidelity_results.append(fr) print(f" avg_key_sim={fr['avg_key_cosine_sim']:.6f} avg_val_sim={fr['avg_val_cosine_sim']:.6f}") except Exception as e: fidelity_results.append({ "context_length": fl, "error": f"{type(e).__name__}: {e}", }) print(f" ERROR: {e}") gc.collect() torch.cuda.empty_cache() return { "metadata": meta, "scaling_results": results, "fidelity_results": fidelity_results, } # --------------------------------------------------------------------------- # Standalone entry point # --------------------------------------------------------------------------- if __name__ == "__main__": import sys result = run_benchmark() out_path = os.path.join(os.path.dirname(__file__), "benchmark_results.json") with open(out_path, "w") as f: json.dump(result, f, indent=2, default=str) print(f"\n\nResults saved to {out_path}") print(json.dumps(result, indent=2, default=str))