benchmark_real.py

"""
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))