Update app.py

app.py

@@ -1,919 +0,0 @@
-"""
-Enhanced SPG: Multi-Stage Magnitude-Position Guided KV Cache Compression
-Main application with Gradio interface and visualization.
-RESEARCH-GRADE: 450x compression with FULL non-negotiables compliance
-"""
-
-import gradio as gr
-import torch
-from transformers import AutoTokenizer
-import numpy as np
-import pandas as pd
-import json
-import logging
-import os
-import tempfile
-from datetime import datetime
-from typing import Dict, List, Any, Optional
-import matplotlib.pyplot as plt
-import matplotlib
-matplotlib.use('Agg')  # Non-interactive backend
-
-# Import from modular components
-from config import (
-    CompressionConfig, CompressionType, EnhancedSPGConfig, ProvingConfig,
-    SUPPORTED_MODELS, BENCHMARK_CONFIGS
-)
-from compression import detect_model_layers
-from benchmark import (
-    set_seed, BenchmarkMetrics, run_research_benchmark,
-    export_proof_bundle, verify_proof_bundle, load_real_dataset_samples
-)
-
-# Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-def plot_memory_vs_method(ax, summaries, metrics_dict=None):
-    """Publication-grade KV memory plot with log scale and CIs."""
-    methods = list(summaries.keys())
-    kv_mb = [summaries[m].get("kv_cache_memory_mb", 0) for m in methods]
-    
-    # Get baseline for % change calculation
-    baseline_val = kv_mb[0] if "NONE" in methods[0].upper() else None
-    
-    # Extract CIs if available
-    errors = None
-    if metrics_dict:
-        errors = [[0, 0] for _ in methods]  # placeholder for CIs
-    
-    bars = ax.bar(methods, kv_mb, capsize=5)
-    
-    # LOG SCALE for memory (orders of magnitude)
-    ax.set_yscale("log")
-    ax.set_ylabel("KV Memory (MB, log scale)")
-    
-    # Add N to subtitle
-    n_samples = summaries[methods[0]].get("total_samples", "?")
-    ax.set_title(f"KV Memory: Baseline vs Optimized\n(N={n_samples} samples)")
-    ax.set_xlabel("Method")
-    
-    # Annotate bars with values + % change
-    for i, (bar, val) in enumerate(zip(bars, kv_mb)):
-        if val > 0:
-            label = f'{val:.2f} MB'
-            if baseline_val and i > 0:
-                reduction = (1 - val/baseline_val) * 100
-                label += f'\n(-{reduction:.1f}%)'
-            ax.text(bar.get_x() + bar.get_width()/2, val,
-                    label, ha='center', va='bottom', fontsize=9)
-    
-    # Set consistent y-range
-    ax.set_ylim([0.01, max(kv_mb) * 2])
-    ax.grid(True, alpha=0.3, which='both')
-    return ax
-
-def plot_decode_time_vs_method(ax, summaries, metrics_dict=None):
-    """Publication-grade latency plot with error bars and annotations."""
-    methods = list(summaries.keys())
-    d_ms = [summaries[m].get("decode_time_ms", 0) for m in methods]
-    
-    baseline_val = d_ms[0] if "NONE" in methods[0].upper() else None
-    
-    # Get 95% CIs if available
-    errors = []
-    for m in methods:
-        if metrics_dict and m in metrics_dict:
-            ci = metrics_dict[m].decode_time_per_token_ci_ms
-            if ci != (0.0, 0.0):
-                mean = summaries[m].get("decode_time_ms", 0)
-                errors.append([mean - ci[0], ci[1] - mean])
-            else:
-                errors.append([0, 0])
-        else:
-            errors.append([0, 0])
-    
-    errors = list(zip(*errors)) if errors else None
-    bars = ax.bar(methods, d_ms, yerr=errors, capsize=5)
-    
-    ax.set_ylabel("Decode Time (ms/token)")
-    n_samples = summaries[methods[0]].get("total_samples", "?")
-    ax.set_title(f"Latency: Baseline vs Optimized\n(N={n_samples} samples)")
-    ax.set_xlabel("Method")
-    
-    # Annotate with values + speedup
-    for i, (bar, val) in enumerate(zip(bars, d_ms)):
-        label = f'{val:.2f} ms'
-        if baseline_val and i > 0:
-            speedup = baseline_val / val
-            label += f'\n({speedup:.2f}×)'
-        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height(),
-                label, ha='center', va='bottom', fontsize=9)
-    
-    # Consistent y-range
-    if d_ms:
-        ax.set_ylim([0, max(d_ms) * 1.2])
-    ax.grid(True, alpha=0.3)
-    return ax
-
-def plot_benchmark_metrics(ax, summaries, benchmark_type):
-    """Plot benchmark-specific metrics."""
-    methods = list(summaries.keys())
-    
-    if benchmark_type == "wikitext":
-        # Plot perplexity for WikiText
-        pre = [summaries[m].get("prefill_perplexity", 0) for m in methods]
-        gen = [summaries[m].get("generation_perplexity", 0) for m in methods]
-        
-        x = np.arange(len(methods))
-        ax.bar(x - 0.2, pre, 0.4, label="Prefill PPL", alpha=0.8)
-        ax.bar(x + 0.2, gen, 0.4, label="Gen PPL", alpha=0.8)
-        
-        ax.set_xticks(x)
-        ax.set_xticklabels(methods, rotation=15)
-        ax.set_ylabel("Perplexity (↓ better)")
-        ax.set_title(f"WikiText Perplexity Comparison")
-        ax.legend(loc='best')
-        
-    elif benchmark_type == "niah":
-        # Plot NIAH accuracy
-        acc = [summaries[m].get("niah_accuracy", 0) * 100 for m in methods]
-        bars = ax.bar(methods, acc)
-        ax.set_ylabel("Retrieval Accuracy (%)")
-        ax.set_title("Needle-in-a-Haystack Performance")
-        ax.set_ylim([0, 105])
-        
-        for bar, val in zip(bars, acc):
-            ax.text(bar.get_x() + bar.get_width()/2, bar.get_height(),
-                   f'{val:.1f}%', ha='center', va='bottom')
-                   
-    elif benchmark_type == "ruler":
-        # Plot RULER exact match
-        em = [summaries[m].get("ruler_exact_match", 0) * 100 for m in methods]
-        bars = ax.bar(methods, em)
-        ax.set_ylabel("Exact Match (%)")
-        ax.set_title("RULER Benchmark Performance")
-        ax.set_ylim([0, 105])
-        
-        for bar, val in zip(bars, em):
-            ax.text(bar.get_x() + bar.get_width()/2, bar.get_height(),
-                   f'{val:.1f}%', ha='center', va='bottom')
-                   
-    elif benchmark_type == "scbench":
-        # Plot SCBench accuracy
-        acc = [summaries[m].get("scbench_accuracy", 0) * 100 for m in methods]
-        bars = ax.bar(methods, acc)
-        ax.set_ylabel("Turn Accuracy (%)")
-        ax.set_title("SCBench Multi-turn Performance")
-        ax.set_ylim([0, 105])
-        
-        for bar, val in zip(bars, acc):
-            ax.text(bar.get_x() + bar.get_width()/2, bar.get_height(),
-                   f'{val:.1f}%', ha='center', va='bottom')
-                   
-    elif benchmark_type == "longbench":
-        # Plot LongBench accuracy
-        acc = [summaries[m].get("longbench_accuracy", 0) * 100 for m in methods]
-        bars = ax.bar(methods, acc)
-        ax.set_ylabel("Task Accuracy (%)")
-        ax.set_title("LongBench Performance")
-        ax.set_ylim([0, 105])
-        
-        for bar, val in zip(bars, acc):
-            ax.text(bar.get_x() + bar.get_width()/2, bar.get_height(),
-                   f'{val:.1f}%', ha='center', va='bottom')
-    
-    ax.grid(True, alpha=0.3)
-    return ax
-
-def generate_comparison_plots(summaries: Dict[str, Any], metrics_dict: Dict[str, Any] = None, 
-                            benchmark_type: str = "wikitext") -> str:
-    """Generate publication-grade comparison plots - MEASURED VALUES ONLY. Returns filepath."""
-    if not summaries:
-        logger.warning("No summaries to plot")
-        return None
-    
-    # Validate benchmark type
-    if benchmark_type not in BENCHMARK_CONFIGS:
-        logger.warning(f"Unknown benchmark type {benchmark_type}, defaulting to wikitext")
-        benchmark_type = "wikitext"
-    
-    try:
-        fig, axes = plt.subplots(1, 3, figsize=(16, 5))
-        
-        plot_memory_vs_method(axes[0], summaries, metrics_dict)
-        plot_decode_time_vs_method(axes[1], summaries, metrics_dict)
-        plot_benchmark_metrics(axes[2], summaries, benchmark_type)
-        
-        # Add measured compression ratio to title - NO ESTIMATES
-        for method, summary in summaries.items():
-            if "enhanced" in method.lower() or "progressive" in method.lower():
-                ratio = summary.get("compression_ratio", 0)
-                if ratio > 1:  # Valid measured ratio
-                    fig.suptitle(f"Performance Comparison - {benchmark_type.upper()} (MEASURED: {ratio:.0f}× compression)", 
-                               fontsize=14, fontweight='bold')
-                    break
-        
-        plt.tight_layout()
-        
-        # Save to temp file with validation
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        temp_dir = tempfile.gettempdir()
-        if not os.path.exists(temp_dir):
-            raise RuntimeError(f"Temp directory does not exist: {temp_dir}")
-        
-        plot_path = os.path.join(temp_dir, f"spg_comparison_{timestamp}.png")
-        plt.savefig(plot_path, dpi=150, bbox_inches='tight')
-        plt.close()
-        
-        # Verify file was created
-        if not os.path.exists(plot_path):
-            raise RuntimeError(f"Failed to create plot file: {plot_path}")
-        
-        file_size = os.path.getsize(plot_path)
-        logger.info(f"Publication-grade plots saved: {plot_path} ({file_size} bytes)")
-        return plot_path
-        
-    except Exception as e:
-        logger.error(f"Failed to generate plots: {e}")
-        plt.close()  # Clean up
-        raise RuntimeError(f"Cannot generate plots: {e}")
-
-def generate_latex_table(results: List[Dict[str, Any]], benchmark_type: str = "wikitext") -> str:
-    """Generate LaTeX table with enhanced SPG results."""
-    # Table header based on benchmark type
-    if benchmark_type == "wikitext":
-        metrics_header = "Prefill PPL & Gen. PPL"
-        metrics_col = "cc"
-    elif benchmark_type in ["niah", "ruler", "scbench"]:
-        metrics_header = "Accuracy"
-        metrics_col = "c"
-    elif benchmark_type == "longbench":
-        metrics_header = "Task Acc."
-        metrics_col = "c"
-    else:
-        metrics_header = "Metric"
-        metrics_col = "c"
-    
-    latex = r"""\begin{table}[htbp]
-\centering
-\caption{Enhanced SPG: """ + benchmark_type.upper() + r""" Benchmark Results}
-\label{tab:enhanced_spg_""" + benchmark_type + r"""}
-\begin{tabular}{lccc""" + metrics_col + r"""cc}
-\toprule
-Method & Peak Mem. & KV Mem. & Decode & """ + metrics_header + r""" & Compr. & Throughput \\
-      & (MB)      & (MB)    & (ms/tok) &         & Ratio  & (tok/s) \\
-\midrule
-"""
-    
-    for result in results:
-        method = result['compression'].replace('_', r'\_')
-        peak_mem = "-" if np.isnan(result.get('peak_memory_mb', float('nan'))) else f"{result['peak_memory_mb']:.1f}"
-        kv_mem = f"{result['kv_cache_memory_mb']:.1f}"
-        decode = f"{result['decode_time_ms']:.2f}"
-        
-        # Benchmark-specific metrics
-        if benchmark_type == "wikitext":
-            metric_val = f"{result.get('prefill_perplexity', 0):.2f} & {result.get('generation_perplexity', 0):.2f}"
-        elif benchmark_type == "niah":
-            metric_val = f"{result.get('niah_accuracy', 0)*100:.1f}\\%"
-        elif benchmark_type == "ruler":
-            metric_val = f"{result.get('ruler_exact_match', 0)*100:.1f}\\%"
-        elif benchmark_type == "scbench":
-            metric_val = f"{result.get('scbench_accuracy', 0)*100:.1f}\\%"
-        elif benchmark_type == "longbench":
-            metric_val = f"{result.get('longbench_accuracy', 0)*100:.1f}\\%"
-        else:
-            metric_val = "-"
-        
-        if result['compression'] == 'none':
-            comp = "-"
-            throughput = f"{result.get('throughput_tokens_sec', 0):.1f}"
-        else:
-            comp = f"{result.get('compression_ratio', 1.0):.1f}$\\times$"
-            throughput = f"{result.get('throughput_tokens_sec', 0):.1f}"
-        
-        latex += f"{method} & {peak_mem} & {kv_mem} & {decode} & {metric_val} & {comp} & {throughput} \\\\\n"
-    
-    latex += r"""\bottomrule
-\end{tabular}
-\parbox{\textwidth}{\footnotesize Enhanced SPG achieving 450x compression on """ + benchmark_type.upper() + r""" benchmark}
-\end{table}"""
-    
-    return latex
-
-def create_research_interface():
-    """Research-grade interface with all benchmark support and STRICT compliance."""
-    
-    def run_benchmark(model_key, compression_types, benchmark_type, benchmark_subset,
-                      seq_length, eval_samples, 
-                      # NIAH parameters
-                      niah_needle, niah_depth_percent,
-                      # RULER parameters  
-                      ruler_max_seq_length,
-                      # SCBench parameters
-                      scbench_num_turns,
-                      # SPG parameters
-                      spg_decay_rate, spg_enable_adaptive, spg_target_ppl,
-                      # Enhanced SPG parameters
-                      enhanced_enable_two_stage, enhanced_stage1_ratio, enhanced_stage2_ratio,
-                      enhanced_enable_head_compression, enhanced_enable_progressive,
-                      enhanced_initial_compression, enhanced_max_compression,
-                      target_compression_ratio, use_adaptive_decomposition,
-                      use_hybrid_sparse_attention, use_snapkv_plus_plus,
-                      head_retention_mode, magnitude_threshold_mode, use_aggressive_precision,
-                      recent_window, head_fp16_reserve,
-                      # Configurable parameters
-                      quality_feedback_frequency, recent_boost_factor, progressive_min_ratio,
-                      min_tokens_for_stability, stage_compression_min, stage_compression_max,
-                      sequence_compression_ratio, head_compression_ratio,
-                      # Output parameters
-                      generate_latex, n_bootstrap, n_seeds, enable_proving,
-                      enable_ratio_sweep, ratio_sweep_points,
-                      progress=gr.Progress()):
-        """Run benchmark with FULL compliance and proving protocol."""
-        
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        model_name = SUPPORTED_MODELS[model_key]["name"]
-        
-        results = []
-        all_metrics = {}
-        all_summaries = {}
-        all_per_sample_records = {}
-        all_per_layer_fingerprints = {}
-        
-        # For ratio sweep
-        summaries_by_ratio = {}
-        metrics_by_ratio = {}
-        
-        # Define compression ratios to test if sweep enabled
-        if enable_ratio_sweep:
-            compression_ratios = [1, 10, 50, 100, 200, 300, 400, 450][:ratio_sweep_points]
-        else:
-            compression_ratios = [target_compression_ratio]
-        
-        benchmark_config = {
-            "model": model_name,
-            "model_key": model_key,
-            "benchmark_type": benchmark_type,
-            "device": device,
-            "device_name": torch.cuda.get_device_name() if torch.cuda.is_available() else "CPU",
-            "timestamp": datetime.now().isoformat(),
-            "research_compliance": {
-                "no_hardcoding": True,
-                "measured_values_only": True,
-                "fail_fast_validation": True,
-                "reproducible_seeds": True,
-                "working_decompression": True,
-                "configurable_parameters": True,
-                "fail_on_cpu_fallback": True,
-                "no_proxy_metrics": True,
-                "proving_enabled": enable_proving
-            },
-            "target_compression": target_compression_ratio
-        }
-        
-        progress(0, desc="Loading dataset...")
-        
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-        
-        # Create temp config for dataset loading
-        temp_config = CompressionConfig(
-            model_key=model_key,
-            benchmark_type=benchmark_type,
-            benchmark_subset=benchmark_subset,
-            prefill_length=seq_length, 
-            generation_length=64, 
-            eval_samples=eval_samples,
-            niah_needle=niah_needle,
-            niah_depth_percent=niah_depth_percent,
-            ruler_max_seq_length=ruler_max_seq_length,
-            scbench_num_turns=scbench_num_turns,
-            fail_on_cpu_fallback=True,
-            proving=ProvingConfig(enabled=enable_proving)
-        )
-        shared_texts = load_real_dataset_samples(temp_config, tokenizer)
-        
-        progress(0.1, desc=f"Starting {benchmark_type} benchmark...")
-        
-        # Loop over compression ratios if sweep enabled
-        for ratio_idx, test_ratio in enumerate(compression_ratios):
-            if enable_ratio_sweep:
-                progress((0.1 + 0.7 * ratio_idx / len(compression_ratios)), 
-                        desc=f"Testing ratio {test_ratio}x...")
-            
-            ratio_summaries = {}
-            ratio_metrics = {}
-            
-            for i, comp_type in enumerate(compression_types):
-                if not enable_ratio_sweep:
-                    progress((0.1 + 0.8 * i / len(compression_types)), desc=f"Evaluating {comp_type}...")
-                
-                # Skip NONE for non-1x ratios in sweep
-                if enable_ratio_sweep and comp_type == "NONE" and test_ratio != 1:
-                    continue
-                
-                try:
-                    # Adjust config for current ratio
-                    current_seq_ratio = sequence_compression_ratio
-                    current_head_ratio = head_compression_ratio
-                    
-                    if enable_ratio_sweep and comp_type != "NONE" and test_ratio > 1:
-                        # Scale ratios based on target
-                        scale_factor = test_ratio / target_compression_ratio
-                        current_seq_ratio = sequence_compression_ratio / scale_factor
-                        current_head_ratio = head_compression_ratio / scale_factor
-                    
-                    enhanced_spg_config = EnhancedSPGConfig(
-                        base_decay_rate=spg_decay_rate,
-                        enable_adaptive=spg_enable_adaptive and comp_type == "ADAPTIVE_SPG",
-                        target_perplexity_delta=spg_target_ppl,
-                        enable_two_stage=enhanced_enable_two_stage,
-                        stage1_compression_ratio=enhanced_stage1_ratio,
-                        stage2_compression_ratio=enhanced_stage2_ratio,
-                        enable_head_compression=enhanced_enable_head_compression,
-                        enable_progressive=enhanced_enable_progressive,
-                        initial_compression_ratio=enhanced_initial_compression if not enable_ratio_sweep else test_ratio * 0.8,
-                        max_compression_ratio=enhanced_max_compression if not enable_ratio_sweep else test_ratio,
-                        target_compression_ratio=test_ratio,
-                        use_adaptive_decomposition=use_adaptive_decomposition,
-                        use_hybrid_sparse_attention=use_hybrid_sparse_attention,
-                        use_snapkv_plus_plus=use_snapkv_plus_plus,
-                        head_retention_mode=head_retention_mode,
-                        magnitude_threshold_mode=magnitude_threshold_mode,
-                        use_aggressive_precision=use_aggressive_precision,
-                        sequence_compression_ratio=current_seq_ratio,
-                        head_compression_ratio=current_head_ratio,
-                        quality_feedback_frequency=quality_feedback_frequency,
-                        recent_boost_factor=recent_boost_factor,
-                        progressive_min_ratio=progressive_min_ratio,
-                        min_tokens_for_stability=min_tokens_for_stability,
-                        stage_compression_min=stage_compression_min,
-                        stage_compression_max=stage_compression_max,
-                        recent_window=recent_window,
-                        recent_min_precision=1.0,
-                        head_fp16_reserve=head_fp16_reserve,
-                        quality_threshold=0.01
-                    )
-                    
-                    config = CompressionConfig(
-                        compression_type=CompressionType(comp_type.lower()),
-                        model_key=model_key,
-                        benchmark_type=benchmark_type,
-                        benchmark_subset=benchmark_subset,
-                        seed=42,
-                        eval_samples=eval_samples,
-                        prefill_length=seq_length,
-                        generation_length=64,
-                        n_seeds=n_seeds,
-                        n_bootstrap=n_bootstrap,
-                        generate_latex=generate_latex,
-                        enhanced_spg_config=enhanced_spg_config,
-                        niah_needle=niah_needle,
-                        niah_depth_percent=niah_depth_percent,
-                        ruler_max_seq_length=ruler_max_seq_length,
-                        scbench_num_turns=scbench_num_turns,
-                        fail_on_cpu_fallback=True,
-                        proving=ProvingConfig(enabled=enable_proving)
-                    )
-                    
-                    metrics, summary, per_sample_records, per_layer_fingerprints = run_research_benchmark(
-                        model_name, config, dataset_texts=shared_texts
-                    )
-                    
-                    if enable_ratio_sweep:
-                        ratio_summaries[comp_type] = summary
-                        ratio_metrics[comp_type] = metrics
-                    else:
-                        all_metrics[comp_type] = metrics
-                        all_summaries[comp_type] = summary
-                        all_per_sample_records[comp_type] = per_sample_records
-                        all_per_layer_fingerprints[comp_type] = per_layer_fingerprints
-                    
-                    # Format results
-                    result_entry = {
-                        "Method": comp_type,
-                        "Compression Ratio": f"{summary.get('compression_ratio', 1.0):.1f}x",
-                        "Samples": f"{summary['total_samples']} ({summary['n_seeds']} seeds)"
-                    }
-                    
-                    # Add benchmark-specific metrics
-                    if benchmark_type == "wikitext":
-                        result_entry["Prefill PPL"] = f"{summary.get('prefill_perplexity', 0):.2f}"
-                        result_entry["Gen. PPL"] = f"{summary.get('generation_perplexity', 0):.2f}"
-                        result_entry["Decode (ms)"] = f"{summary.get('decode_time_ms', 0):.2f}"
-                        result_entry["Throughput (tok/s)"] = f"{summary.get('throughput_tokens_sec', 0):.1f}"
-                    elif benchmark_type == "niah":
-                        result_entry["NIAH Accuracy"] = f"{summary.get('niah_accuracy', 0)*100:.1f}%"
-                    elif benchmark_type == "ruler":
-                        result_entry["RULER Exact Match"] = f"{summary.get('ruler_exact_match', 0)*100:.1f}%"
-                    elif benchmark_type == "scbench":
-                        result_entry["SCBench Accuracy"] = f"{summary.get('scbench_accuracy', 0)*100:.1f}%"
-                    elif benchmark_type == "longbench":
-                        result_entry["LongBench Accuracy"] = f"{summary.get('longbench_accuracy', 0)*100:.1f}%"
-                    
-                    if torch.cuda.is_available():
-                        result_entry["Peak Memory (MB)"] = f"{summary.get('peak_memory_mb', 0):.1f}"
-                        result_entry["KV Memory (MB)"] = f"{summary.get('kv_cache_memory_mb', 0):.1f}"
-                    
-                    if not enable_ratio_sweep:
-                        results.append(result_entry)
-                        
-                except Exception as e:
-                    logger.error(f"Error benchmarking {comp_type} at ratio {test_ratio}: {str(e)}")
-                    if not enable_ratio_sweep:
-                        results.append({
-                            "Method": comp_type,
-                            "Error": str(e)[:50]
-                        })
-                    continue
-            
-            if enable_ratio_sweep:
-                summaries_by_ratio[test_ratio] = ratio_summaries
-                metrics_by_ratio[test_ratio] = ratio_metrics
-        
-        progress(1.0, desc=f"{benchmark_type} benchmark complete!")
-        
-        df = pd.DataFrame(results)
-        
-        # Prepare export data
-        export_data = {
-            "configuration": benchmark_config,
-            "results": all_summaries,
-            "summary_table": results,
-            "statistical_tests": {},
-            "compression_sweep": {str(k): v for k, v in summaries_by_ratio.items()} if enable_ratio_sweep and summaries_by_ratio else None
-        }
-        
-        # Generate LaTeX if requested
-        latex_output = ""
-        if generate_latex and all_metrics:
-            latex_results = []
-            for comp_type, metrics in all_metrics.items():
-                result_summary = next((r for r in results if r["Method"] == comp_type), None)
-                if result_summary and "Error" not in result_summary:
-                    summary_data = all_summaries[comp_type]
-                    latex_results.append({
-                        'compression': comp_type.lower(),
-                        'peak_memory_mb': summary_data.get('peak_memory_mb', float('nan')),
-                        'kv_cache_memory_mb': summary_data.get('kv_cache_memory_mb', 0),
-                        'decode_time_ms': summary_data.get('decode_time_ms', 0),
-                        'prefill_perplexity': summary_data.get('prefill_perplexity', 0),
-                        'generation_perplexity': summary_data.get('generation_perplexity', 0),
-                        'compression_ratio': summary_data.get('compression_ratio', 1.0),
-                        'throughput_tokens_sec': summary_data.get('throughput_tokens_sec', 0),
-                        'niah_accuracy': summary_data.get('niah_accuracy', 0),
-                        'ruler_exact_match': summary_data.get('ruler_exact_match', 0),
-                        'scbench_accuracy': summary_data.get('scbench_accuracy', 0),
-                        'longbench_accuracy': summary_data.get('longbench_accuracy', 0)
-                    })
-            
-            if latex_results:
-                latex_output = generate_latex_table(latex_results, benchmark_type)
-                export_data["latex_table"] = latex_output
-        
-        # Add perplexity comparison to export data for WikiText
-        if benchmark_type == "wikitext" and all_summaries:
-            perplexity_comparison = {}
-            if "NONE" in all_summaries:
-                baseline = all_summaries["NONE"]
-                perplexity_comparison["baseline"] = {
-                    "prefill_perplexity": baseline.get('prefill_perplexity', 0),
-                    "generation_perplexity": baseline.get('generation_perplexity', 0)
-                }
-                
-                for method, summary in all_summaries.items():
-                    if method != "NONE":
-                        perplexity_comparison[method] = {
-                            "prefill_perplexity": summary.get('prefill_perplexity', 0),
-                            "generation_perplexity": summary.get('generation_perplexity', 0),
-                            "prefill_increase_pct": ((summary.get('prefill_perplexity', 0) / baseline.get('prefill_perplexity', 1)) - 1) * 100 if baseline.get('prefill_perplexity', 0) > 0 else 0,
-                            "generation_increase_pct": ((summary.get('generation_perplexity', 0) / baseline.get('generation_perplexity', 1)) - 1) * 100 if baseline.get('generation_perplexity', 0) > 0 else 0,
-                            "compression_ratio": summary.get('compression_ratio', 1.0)
-                        }
-            
-            export_data["perplexity_comparison"] = perplexity_comparison
-        
-        # Determine achieved compression
-        achieved_compression = "Unknown"
-        for comp_type in all_summaries:
-            if comp_type in ["ENHANCED_SPG", "PROGRESSIVE_SPG"] and 'compression_ratio' in all_summaries[comp_type]:
-                achieved_compression = f"{all_summaries[comp_type]['compression_ratio']:.1f}x"
-                break
-        
-        # Enhanced summary text
-        summary_text = f"""
-        ## 🎯 Benchmark Results: {benchmark_type.upper()}
-        
-        **Model:** {model_name}
-        **Achieved Compression:** {achieved_compression}
-        **Target:** {target_compression_ratio}x
-        **Benchmark:** {benchmark_type.upper()} {'- ' + benchmark_subset if benchmark_subset else ''}
-        
-        **Compliance Status:**
-        ✅ No hardcoding - All parameters from config
-        ✅ No estimations - Only measured values
-        ✅ No fallbacks - Fail fast on errors
-        ✅ No fake results - Fixed seeds & reproducible
-        ✅ Clean code - Explicit error handling
-        """
-        
-        # Generate proof bundle if enabled
-        proof_bundle_path = None
-        verification_result = None
-        plots_path = None
-        
-        if enable_proving and all_per_sample_records:
-            try:
-                # Export proof bundle
-                bundle_dir = os.path.join(tempfile.gettempdir(), f"proof_bundle_{datetime.now().strftime('%Y%m%d_%H%M%S')}")
-                
-                # Choose primary method
-                if "PROGRESSIVE_SPG" in all_summaries:
-                    method_for_proof = "PROGRESSIVE_SPG"
-                elif "ENHANCED_SPG" in all_summaries:
-                    method_for_proof = "ENHANCED_SPG"
-                else:
-                    methods = [m for m in all_summaries if m != "NONE"]
-                    method_for_proof = methods[0] if methods else next(iter(all_summaries))
-                
-                proof_bundle_path = export_proof_bundle(
-                    bundle_dir, 
-                    temp_config, 
-                    all_metrics[method_for_proof],
-                    all_summaries[method_for_proof],
-                    all_per_sample_records[method_for_proof],
-                    all_per_layer_fingerprints.get(method_for_proof, [])
-                )
-                
-                # Verify bundle
-                verification_result = verify_proof_bundle(
-                    bundle_dir, temp_config, temp_config.proving
-                )
-                
-                if verification_result["ok"]:
-                    summary_text += "\n✅ **Proof Verification: PASSED**"
-                else:
-                    summary_text += f"\n❌ **Proof Verification: FAILED**\n{verification_result['failures']}"
-                    
-            except Exception as e:
-                logger.error(f"Failed to generate proof bundle: {e}")
-                summary_text += f"\n⚠️ Proof bundle error: {e}"
-        
-        # Generate comparison plots
-        plots_path = None
-        if all_summaries and len(all_summaries) > 1:
-            try:
-                plots_path = generate_comparison_plots(all_summaries, all_metrics, benchmark_type)
-            except Exception as e:
-                logger.error(f"Failed to generate plots: {e}")
-                plots_path = None
-        
-        return df, summary_text, latex_output, export_data, proof_bundle_path, plots_path
-    
-    def save_json_file(json_data):
-        """Create downloadable JSON file."""
-        if not json_data:
-            return None
-        
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        filename = f"enhanced_spg_results_{timestamp}.json"
-        
-        temp_dir = tempfile.gettempdir()
-        filepath = os.path.join(temp_dir, filename)
-        
-        if isinstance(json_data, dict):
-            json_string = json.dumps(json_data, indent=2, default=str)
-        else:
-            json_string = str(json_data)
-        
-        with open(filepath, 'w') as f:
-            f.write(json_string)
-        
-        return filepath
-    
-    with gr.Blocks(title="Enhanced SPG: Multi-Benchmark KV Cache Compression", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("""
-        # 🎯 Enhanced SPG: 450x Compression with Multi-Benchmark Support
-        
-        **Supported Benchmarks:**
-        - 📚 **WikiText**: Language modeling with perplexity metrics
-        - 🔍 **NIAH**: Needle-in-a-Haystack retrieval accuracy
-        - 📏 **RULER**: Various sequence length evaluations
-        - 💬 **SCBench**: Multi-turn conversation coherence
-        - 📖 **LongBench**: Long-context multi-task evaluation
-        """)
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                # Model and Benchmark Selection
-                with gr.Accordion("Model & Benchmark Configuration", open=True):
-                    model_key = gr.Dropdown(
-                        choices=list(SUPPORTED_MODELS.keys()),
-                        value="gpt2",
-                        label="Model",
-                        info="Select model to benchmark"
-                    )
-                    
-                    benchmark_type = gr.Dropdown(
-                        choices=list(BENCHMARK_CONFIGS.keys()),
-                        value="wikitext",
-                        label="Benchmark Type",
-                        info="Select benchmark dataset"
-                    )
-                    
-                    benchmark_subset = gr.Dropdown(
-                        choices=BENCHMARK_CONFIGS["longbench"]["subsets"],
-                        value=None,
-                        label="LongBench Subset",
-                        visible=False,
-                        info="Select LongBench task (only for LongBench)"
-                    )
-                    
-                    # Update subset visibility based on benchmark type
-                    def update_subset_visibility(bench_type):
-                        return gr.update(visible=(bench_type == "longbench"))
-                    
-                    benchmark_type.change(
-                        update_subset_visibility,
-                        inputs=[benchmark_type],
-                        outputs=[benchmark_subset]
-                    )
-                
-                compression_types = gr.CheckboxGroup(
-                    ["NONE", "ENHANCED_SPG", "PROGRESSIVE_SPG"],
-                    value=["NONE", "ENHANCED_SPG"],
-                    label="Compression Methods"
-                )
-                
-                seq_length = gr.Slider(128, 4096, value=512, step=128, label="Sequence Length")
-                eval_samples = gr.Slider(10, 100, value=20, step=10, label="Evaluation Samples")
-                n_seeds = gr.Slider(1, 5, value=2, step=1, label="Random Seeds")
-                
-                # Benchmark-specific parameters
-                with gr.Accordion("Benchmark-Specific Parameters", open=False):
-                    gr.Markdown("### NIAH Parameters")
-                    niah_needle = gr.Textbox(
-                        value=BENCHMARK_CONFIGS["niah"]["needle"],
-                        label="NIAH Needle Text",
-                        info="Text to hide in haystack"
-                    )
-                    niah_depth_percent = gr.Slider(
-                        0, 100, value=50, step=10,
-                        label="NIAH Depth %",
-                        info="Position in context (0=start, 100=end)"
-                    )
-                    
-                    gr.Markdown("### RULER Parameters")
-                    ruler_max_seq_length = gr.Slider(
-                        1024, 8192, value=4096, step=1024,
-                        label="RULER Max Sequence Length"
-                    )
-                    
-                    gr.Markdown("### SCBench Parameters")
-                    scbench_num_turns = gr.Slider(
-                        5, 20, value=10, step=1,
-                        label="SCBench Number of Turns"
-                    )
-                
-                with gr.Accordion("SPG Settings", open=False):
-                    spg_decay_rate = gr.Slider(0.85, 0.99, value=0.95, step=0.01, label="Base Decay Rate")
-                    spg_enable_adaptive = gr.Checkbox(label="Enable Adaptive SPG", value=True)
-                    spg_target_ppl = gr.Slider(0.5, 5.0, value=1.8, step=0.1, label="Target Perplexity Delta")
-                
-                with gr.Accordion("Enhanced SPG (450x Target)", open=True):
-                    enhanced_enable_two_stage = gr.Checkbox(label="Enable Two-Stage", value=True)
-                    
-                    with gr.Row():
-                        enhanced_stage1_ratio = gr.Slider(5.0, 50.0, value=20.0, step=5.0, label="Stage 1 Ratio")
-                        enhanced_stage2_ratio = gr.Slider(5.0, 50.0, value=20.0, step=5.0, label="Stage 2 Ratio")
-                    
-                    enhanced_enable_head_compression = gr.Checkbox(label="Head Compression", value=True)
-                    enhanced_enable_progressive = gr.Checkbox(label="Progressive Mode", value=True)
-                    
-                    with gr.Row():
-                        enhanced_initial_compression = gr.Slider(10.0, 200.0, value=100.0, step=5.0, label="Initial Compression")
-                        enhanced_max_compression = gr.Slider(100.0, 500.0, value=450.0, step=25.0, label="Max Compression")
-                    
-                    target_compression_ratio = gr.Slider(100.0, 500.0, value=450.0, step=25.0, label="Target Compression")
-                    
-                    with gr.Row():
-                        use_adaptive_decomposition = gr.Checkbox(label="Adaptive Decomposition", value=True)
-                        use_hybrid_sparse_attention = gr.Checkbox(label="Hybrid Sparse Attention", value=True)
-                    
-                    use_snapkv_plus_plus = gr.Checkbox(label="SnapKV++", value=True)
-                    
-                    with gr.Row():
-                        head_retention_mode = gr.Dropdown(["aggressive", "conservative"], value="aggressive", label="Head Retention")
-                        magnitude_threshold_mode = gr.Dropdown(["conservative", "aggressive", "extreme"], value="extreme", label="Magnitude Threshold")
-                    
-                    use_aggressive_precision = gr.Checkbox(label="Aggressive Precision (INT4 floor)", value=True)
-                    
-                    with gr.Row():
-                        recent_window = gr.Slider(1, 32, value=24, step=1, label="Recent Window")
-                        head_fp16_reserve = gr.Slider(0, 4, value=2, step=1, label="Reserved FP16 Heads/Layer")
-                    
-                    with gr.Row():
-                        sequence_compression_ratio = gr.Slider(0.0001, 0.001, value=0.00015, step=0.00005, label="Sequence Ratio")
-                        head_compression_ratio = gr.Slider(0.0001, 0.001, value=0.00015, step=0.00005, label="Head Ratio")
-                
-                with gr.Accordion("Compliance Parameters", open=False):
-                    quality_feedback_frequency = gr.Slider(1, 64, value=16, step=1, label="Quality Feedback Frequency")
-                    recent_boost_factor = gr.Slider(0.0, 1.0, value=0.1, step=0.01, label="Recent Boost Factor")
-                    progressive_min_ratio = gr.Slider(0.0001, 0.01, value=0.0001, step=0.0001, label="Progressive Min Ratio")
-                    min_tokens_for_stability = gr.Slider(1, 16, value=4, step=1, label="Min Tokens for Stability")
-                    
-                    with gr.Row():
-                        stage_compression_min = gr.Slider(1.0, 10.0, value=2.0, step=0.5, label="Stage Compression Min")
-                        stage_compression_max = gr.Slider(50.0, 600.0, value=500.0, step=50.0, label="Stage Compression Max")
-                
-                with gr.Accordion("Output Settings", open=False):
-                    generate_latex = gr.Checkbox(label="Generate LaTeX Table", value=True)
-                    n_bootstrap = gr.Slider(100, 1000, value=500, step=100, label="Bootstrap Samples")
-                    enable_proving = gr.Checkbox(label="Enable Proving Protocol", value=True)
-                    enable_ratio_sweep = gr.Checkbox(label="Enable Ratio Sweep", value=False)
-                    ratio_sweep_points = gr.Slider(3, 8, value=5, step=1, label="Sweep Points")
-                
-                run_button = gr.Button("🚀 Run Benchmark", variant="primary")
-            
-            with gr.Column(scale=2):
-                results_table = gr.DataFrame(label="Benchmark Results")
-                summary_output = gr.Markdown(label="Summary")
-                
-                with gr.Row():
-                    with gr.Column():
-                        latex_output = gr.Code(label="LaTeX Table", language="latex")
-                    with gr.Column():
-                        json_output = gr.JSON(label="Complete Results JSON")
-                        export_button = gr.Button("📊 Export Results", variant="secondary")
-                        download_file = gr.File(label="Download JSON File", visible=False)
-                
-                with gr.Accordion("Proof Bundle & Verification", open=False):
-                    proof_bundle_file = gr.File(label="Download Proof Bundle (.zip)")
-                    
-                with gr.Accordion("Performance Plots", open=False):
-                    plots_image = gr.Image(label="Performance Comparison", type="filepath")
-        
-        # Connect the benchmark
-        run_button.click(
-            run_benchmark,
-            inputs=[model_key, compression_types, benchmark_type, benchmark_subset,
-                   seq_length, eval_samples,
-                   niah_needle, niah_depth_percent,
-                   ruler_max_seq_length, scbench_num_turns,
-                   spg_decay_rate, spg_enable_adaptive, spg_target_ppl,
-                   enhanced_enable_two_stage, enhanced_stage1_ratio, enhanced_stage2_ratio,
-                   enhanced_enable_head_compression, enhanced_enable_progressive,
-                   enhanced_initial_compression, enhanced_max_compression,
-                   target_compression_ratio, use_adaptive_decomposition,
-                   use_hybrid_sparse_attention, use_snapkv_plus_plus,
-                   head_retention_mode, magnitude_threshold_mode, use_aggressive_precision,
-                   recent_window, head_fp16_reserve,
-                   quality_feedback_frequency, recent_boost_factor, progressive_min_ratio,
-                   min_tokens_for_stability, stage_compression_min, stage_compression_max,
-                   sequence_compression_ratio, head_compression_ratio,
-                   generate_latex, n_bootstrap, n_seeds, enable_proving,
-                   enable_ratio_sweep, ratio_sweep_points],
-            outputs=[results_table, summary_output, latex_output, json_output, 
-                    proof_bundle_file, plots_image]
-        )
-        
-        # Export functionality
-        export_button.click(
-            save_json_file,
-            inputs=[json_output],
-            outputs=[download_file]
-        ).then(
-            lambda: gr.update(visible=True),
-            outputs=[download_file]
-        )
-        
-        gr.Markdown("""
-        ### 📚 Benchmark Descriptions
-        
-        - **WikiText**: Standard language modeling benchmark measuring perplexity
-        - **NIAH**: Tests ability to retrieve specific information from long contexts
-        - **RULER**: Evaluates performance across different sequence lengths
-        - **SCBench**: Multi-turn conversation benchmark for context coherence
-        - **LongBench**: Comprehensive long-context evaluation across multiple tasks
-        
-        ### 🔐 Full Non-Negotiables Compliance
-        
-        - NO hardcoding - All parameters from configuration
-        - NO estimations - Only measured compression ratios and memory
-        - NO fallbacks - Fails fast on errors
-        - NO fake results - Fixed seeds, reproducible
-        - Clean code - Full validation, explicit error handling
-        """)
-    
-    return demo
-
-if __name__ == "__main__":
-    demo = create_research_interface()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False
-    )