anujsarker/kv-cache-visualizer

1import gradio as gr2import matplotlib.pyplot as plt3import numpy as np4 5def calculate_kv_cache(model_choice, batch_size, max_seq, precision):6    # Model Architecture Logic7    configs = {8        "Llama-3-8B": {"layers": 32, "heads": 32, "d_head": 128},9        "Llama-3-70B": {"layers": 80, "heads": 64, "d_head": 128},10        "DeepSeek-V3 (MLA)": {"layers": 60, "heads": 128, "d_head": 128}11    }12    13    cfg = configs[model_choice]14    prec_val = 2.0 if "FP16" in precision else 1.0 if "FP8" in precision else 0.515    16    # Calculate Data for Plot17    x_range = np.linspace(512, max_seq, 20)18    # Formula: 2 * layers * heads * d_head * seq * batch * precision19    y_vram = [(2 * cfg['layers'] * cfg['heads'] * cfg['d_head'] * s * batch_size * prec_val) / (1024**3) for s in x_range]20    21    # Optimization estimates22    final_vram = y_vram[-1]23    gqa_vram = final_vram / 8  # Rough GQA 8:1 ratio24    mla_vram = final_vram * 0.15 # MLA compression estimate25    26    # Create Plot27    plt.figure(figsize=(10, 5))28    plt.plot(x_range, y_vram, label="Standard MHA", color="#FF4B4B", linewidth=2)29    plt.fill_between(x_range, y_vram, alpha=0.2, color="#FF4B4B")30    plt.title(f"KV Cache VRAM Growth: {model_choice}")31    plt.xlabel("Sequence Length")32    plt.ylabel("VRAM (GB)")33    plt.grid(True, linestyle='--', alpha=0.6)34    plt.legend()35    36    summary = (37        f"### 📊 Analysis for {model_choice}\n"38        f"- **Peak KV Cache Usage:** {final_vram:.2f} GB\n"39        f"- **With GQA Optimization:** ~{gqa_vram:.2f} GB\n"40        f"- **With MLA (DeepSeek) Optimization:** ~{mla_vram:.2f} GB\n\n"41        "**Research Note:** In the Decode phase, this memory must be loaded for every single token. "42        "High VRAM usage here directly reduces your maximum throughput (tokens/sec)."43    )44    45    return plt, summary46 47# UI Setup48with gr.Blocks(theme=gr.themes.Soft()) as demo:49    gr.Markdown("# 🧠 KV Cache Research Dashboard")50    gr.Markdown("Interactive visualizer for the paper: *'Demystifying KV Cache'*")51    52    with gr.Row():53        with gr.Column():54            model_drop = gr.Dropdown(choices=["Llama-3-8B", "Llama-3-70B", "DeepSeek-V3 (MLA)"], value="Llama-3-8B", label="Select Model Architecture")55            batch_sld = gr.Slider(1, 128, value=32, step=1, label="Batch Size")56            seq_sld = gr.Slider(512, 131072, value=8192, step=512, label="Max Context Window")57            prec_rad = gr.Radio(["FP16", "FP8", "INT4"], value="FP16", label="Cache Precision")58            btn = gr.Button("Analyze Bottleneck", variant="primary")59            60        with gr.Column():61            plot_out = gr.Plot(label="Memory Scaling")62            text_out = gr.Markdown()63 64    btn.click(calculate_kv_cache, inputs=[model_drop, batch_sld, seq_sld, prec_rad], outputs=[plot_out, text_out])65 66demo.launch()