Update benchmark.py

benchmark.py

@@ -1,827 +0,0 @@
-"""
-Benchmarking module for Enhanced SPG compression.
-Contains metrics, evaluation logic, and proof generation.
-STRICT COMPLIANCE: Only direct measurements, no proxy metrics.
-"""
-
-import torch
-import torch.nn.functional as F
-import numpy as np
-from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache
-from datasets import load_dataset
-from typing import Tuple, Optional, Dict, Any, List
-from dataclasses import dataclass, field
-from scipy import stats
-import time
-import json
-import os
-import sys
-import gc
-import tempfile
-import zipfile
-import pathlib
-import platform
-import subprocess
-from datetime import datetime
-import random
-import logging
-
-from config import (
-    CompressionConfig, CompressionType, ProvingConfig, ResearchConstants, logger
-)
-from compression import QuantizedKVCache, detect_model_layers
-
-
-def set_seed(seed: int = 42) -> None:
-    """Set all seeds for reproducibility with explicit validation."""
-    if not isinstance(seed, int) or seed < 0:
-        raise ValueError(f"Seed must be non-negative integer, got {seed}")
-    
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
-        torch.backends.cudnn.deterministic = True
-        torch.backends.cudnn.benchmark = False
-    
-    logger.info(f"Set all random seeds to {seed}")
-
-
-def _peak_mem_bytes_all_gpus() -> int:
-    """Get peak memory across all GPUs. FAIL FAST if CUDA unavailable when expected."""
-    if not torch.cuda.is_available():
-        # This should only be called when CUDA is expected
-        raise RuntimeError("CUDA memory tracking requested but CUDA is unavailable")
-    
-    torch.cuda.synchronize()
-    total_mem = sum(torch.cuda.max_memory_allocated(d) for d in range(torch.cuda.device_count()))
-    logger.debug(f"Peak GPU memory: {total_mem / 1024 / 1024:.1f} MB")
-    return total_mem
-
-
-@dataclass
-class BenchmarkMetrics:
-    """Comprehensive metrics with proper statistical handling - NO ESTIMATES."""
-    # Prefill metrics
-    prefill_times: List[float] = field(default_factory=list)
-    prefill_peak_memories: List[float] = field(default_factory=list)
-    prefill_time_mean: float = 0.0
-    prefill_time_std: float = 0.0
-    prefill_time_ci: Tuple[float, float] = (0.0, 0.0)
-    prefill_peak_memory_mean_mb: float = 0.0
-    prefill_peak_memory_std_mb: float = 0.0
-    prefill_peak_memory_ci_mb: Tuple[float, float] = (0.0, 0.0)
-    prefill_tokens_per_sec: float = 0.0
-    
-    # Decode metrics
-    decode_times: List[float] = field(default_factory=list)
-    decode_peak_memories: List[float] = field(default_factory=list)
-    decode_time_per_token_mean_ms: float = 0.0
-    decode_time_per_token_std_ms: float = 0.0
-    decode_time_per_token_ci_ms: Tuple[float, float] = (0.0, 0.0)
-    decode_time_p50_ms: float = 0.0
-    decode_time_p95_ms: float = 0.0
-    decode_peak_memory_mean_mb: float = 0.0
-    decode_tokens_per_sec: float = 0.0
-    
-    # Quality metrics
-    prefill_perplexities: List[float] = field(default_factory=list)
-    generation_perplexities: List[float] = field(default_factory=list)
-    prefill_perplexity_mean: float = 0.0
-    prefill_perplexity_std: float = 0.0
-    prefill_perplexity_ci: Tuple[float, float] = (0.0, 0.0)
-    generation_perplexity_mean: float = 0.0
-    generation_perplexity_std: float = 0.0
-    generation_perplexity_ci: Tuple[float, float] = (0.0, 0.0)
-    
-    # Compression metrics (MEASURED ONLY - no estimates)
-    compression_ratios: List[float] = field(default_factory=list)
-    compression_ratio_mean: float = 0.0
-    compression_ratio_std: float = 0.0
-    kv_cache_memory_mb: float = 0.0
-    kv_cache_memory_samples_mb: List[float] = field(default_factory=list)
-    
-    # Enhanced SPG metrics (MEASURED ONLY)
-    enhanced_spg_measured_compression: List[float] = field(default_factory=list)
-    enhanced_spg_measured_auxiliary_overhead_mb: List[float] = field(default_factory=list)
-    enhanced_spg_progressive_steps: List[int] = field(default_factory=list)
-    
-    # Original SPG metrics
-    spg_precision_distributions: List[Dict[str, float]] = field(default_factory=list)
-    spg_effective_bits_per_token: List[float] = field(default_factory=list)
-    spg_decay_rates_per_layer: List[List[float]] = field(default_factory=list)
-    
-    # Statistical comparisons
-    memory_reduction_ratio: float = 1.0
-    memory_reduction_pvalue: float = 1.0
-    speedup_ratio: float = 1.0
-    speedup_pvalue: float = 1.0
-    prefill_perplexity_delta: float = 0.0
-    generation_perplexity_delta: float = 0.0
-    perplexity_pvalue: float = 1.0
-    
-    # End-to-end metrics
-    end_to_end_throughput: float = 0.0  # tokens/sec for full sequence
-    end_to_end_latency_ms: float = 0.0  # total time for prefill + generation
-    
-    def calculate_statistics(self, config: CompressionConfig) -> None:
-        """Calculate all statistics with proper error handling."""
-        try:
-            if self.prefill_times:
-                self.prefill_time_mean = float(np.mean(self.prefill_times))
-                self.prefill_time_std = float(np.std(self.prefill_times))
-                self.prefill_time_ci = self._bootstrap_ci(self.prefill_times, config)
-                self.prefill_tokens_per_sec = config.prefill_length / self.prefill_time_mean if self.prefill_time_mean > 0 else 0.0
-            
-            if self.prefill_peak_memories:
-                memories_mb = [m / (1024 * 1024) for m in self.prefill_peak_memories]
-                self.prefill_peak_memory_mean_mb = float(np.mean(memories_mb))
-                self.prefill_peak_memory_std_mb = float(np.std(memories_mb))
-                self.prefill_peak_memory_ci_mb = self._bootstrap_ci(memories_mb, config)
-            
-            if self.decode_times:
-                self.decode_time_per_token_mean_ms = float(np.mean(self.decode_times) * 1000)
-                self.decode_time_per_token_std_ms = float(np.std(self.decode_times) * 1000)
-                self.decode_time_per_token_ci_ms = tuple(x * 1000 for x in self._bootstrap_ci(self.decode_times, config))
-                self.decode_tokens_per_sec = 1.0 / np.mean(self.decode_times) if self.decode_times else 0.0
-                self.decode_time_p50_ms = float(np.percentile(self.decode_times, 50) * 1000)
-                self.decode_time_p95_ms = float(np.percentile(self.decode_times, 95) * 1000)
-            
-            # Calculate end-to-end throughput
-            if self.prefill_time_mean > 0 and self.decode_time_per_token_mean_ms > 0:
-                total_tokens = config.prefill_length + config.generation_length
-                total_time_sec = self.prefill_time_mean + (self.decode_time_per_token_mean_ms * config.generation_length / 1000)
-                self.end_to_end_throughput = total_tokens / total_time_sec if total_time_sec > 0 else 0.0
-                self.end_to_end_latency_ms = total_time_sec * 1000
-            
-            if self.decode_peak_memories:
-                self.decode_peak_memory_mean_mb = float(np.mean(self.decode_peak_memories) / (1024 * 1024))
-            
-            if self.prefill_perplexities:
-                self.prefill_perplexity_mean = float(np.mean(self.prefill_perplexities))
-                self.prefill_perplexity_std = float(np.std(self.prefill_perplexities))
-                self.prefill_perplexity_ci = self._bootstrap_ci(self.prefill_perplexities, config)
-            
-            if self.generation_perplexities:
-                self.generation_perplexity_mean = float(np.mean(self.generation_perplexities))
-                self.generation_perplexity_std = float(np.std(self.generation_perplexities))
-                self.generation_perplexity_ci = self._bootstrap_ci(self.generation_perplexities, config)
-            
-            if self.compression_ratios:
-                self.compression_ratio_mean = float(np.mean(self.compression_ratios))
-                self.compression_ratio_std = float(np.std(self.compression_ratios))
-            
-            if self.kv_cache_memory_samples_mb:
-                self.kv_cache_memory_mb = float(np.mean(self.kv_cache_memory_samples_mb))
-            
-            # Log measured compression results
-            if self.enhanced_spg_measured_compression:
-                logger.info(f"Enhanced SPG measured compression: {np.mean(self.enhanced_spg_measured_compression):.1f}x")
-            
-            if self.spg_effective_bits_per_token:
-                logger.info(f"SPG average bits per token: {np.mean(self.spg_effective_bits_per_token):.2f}")
-                
-        except Exception as e:
-            logger.error(f"Error calculating statistics: {e}")
-            raise
-    
-    def _bootstrap_ci(self, data: List[float], config: CompressionConfig) -> Tuple[float, float]:
-        """Calculate bootstrap confidence interval with reproducible RNG."""
-        if not data or len(data) < 2:
-            logger.warning("Insufficient data for confidence interval calculation")
-            return (0.0, 0.0)
-        
-        try:
-            # Use deterministic RNG for reproducibility
-            rng = np.random.default_rng(config.seed)
-            bootstrap_means = []
-            data_array = np.array(data)
-            
-            for _ in range(config.n_bootstrap):
-                sample = rng.choice(data_array, size=len(data_array), replace=True)
-                bootstrap_means.append(float(sample.mean()))
-            
-            if bootstrap_means:
-                alpha = 1 - config.confidence_level
-                lower = float(np.percentile(bootstrap_means, alpha/2 * 100))
-                upper = float(np.percentile(bootstrap_means, (1 - alpha/2) * 100))
-                return (lower, upper)
-            
-        except Exception as e:
-            logger.error(f"Error in bootstrap CI calculation: {e}")
-            raise
-        
-        return (0.0, 0.0)
-    
-    def compare_with_baseline(self, baseline: 'BenchmarkMetrics', use_paired_tests: bool = True) -> None:
-        """Statistical comparison with proper error handling."""
-        try:
-            if baseline.prefill_peak_memory_mean_mb > 0:
-                self.memory_reduction_ratio = baseline.prefill_peak_memory_mean_mb / max(self.prefill_peak_memory_mean_mb, 1e-9)
-                
-                if baseline.prefill_peak_memories and self.prefill_peak_memories:
-                    if use_paired_tests and len(baseline.prefill_peak_memories) == len(self.prefill_peak_memories):
-                        _, self.memory_reduction_pvalue = stats.ttest_rel(baseline.prefill_peak_memories, self.prefill_peak_memories)
-                    else:
-                        _, self.memory_reduction_pvalue = stats.ttest_ind(baseline.prefill_peak_memories, self.prefill_peak_memories)
-            
-            if baseline.decode_tokens_per_sec > 0 and self.decode_tokens_per_sec > 0:
-                self.speedup_ratio = self.decode_tokens_per_sec / baseline.decode_tokens_per_sec
-                
-                if baseline.decode_times and self.decode_times:
-                    if use_paired_tests and len(baseline.decode_times) == len(self.decode_times):
-                        _, self.speedup_pvalue = stats.ttest_rel(baseline.decode_times, self.decode_times)
-                    else:
-                        _, self.speedup_pvalue = stats.ttest_ind(baseline.decode_times, self.decode_times)
-            
-            self.prefill_perplexity_delta = self.prefill_perplexity_mean - baseline.prefill_perplexity_mean
-            self.generation_perplexity_delta = self.generation_perplexity_mean - baseline.generation_perplexity_mean
-            
-            if baseline.generation_perplexities and self.generation_perplexities:
-                if use_paired_tests and len(baseline.generation_perplexities) == len(self.generation_perplexities):
-                    _, self.perplexity_pvalue = stats.ttest_rel(self.generation_perplexities, baseline.generation_perplexities)
-                else:
-                    _, self.perplexity_pvalue = stats.ttest_ind(self.generation_perplexities, baseline.generation_perplexities)
-                    
-        except Exception as e:
-            logger.error(f"Error in baseline comparison: {e}")
-            raise
-
-
-def export_proof_bundle(bundle_dir: str, config: CompressionConfig,
-                       metrics: BenchmarkMetrics, summary: Dict[str, Any],
-                       per_sample_records: List[Dict[str, Any]],
-                       per_layer_fingerprints: List[Dict[str, Any]]) -> str:
-    """Export attestable proof bundle with all metrics and fingerprints. NO ESTIMATES."""
-    p = pathlib.Path(bundle_dir)
-    p.mkdir(parents=True, exist_ok=True)
-    
-    # Create manifest with full environment info
-    manifest = {
-        "config": json.loads(config.to_json()),
-        "config_hash": config.get_hash(),
-        "git_commit": os.environ.get("GIT_COMMIT", None),
-        "python": sys.version,
-        "torch": config.torch_version,
-        "transformers": config.transformers_version,
-        "cuda": config.cuda_version,
-        "device_name": config.device_name,
-        "start_time": summary.get("start_time"),
-        "end_time": summary.get("end_time"),
-        "hostname": platform.node(),
-        "strict_flags": {
-            "fail_on_cpu_fallback": config.fail_on_cpu_fallback,
-            "proving_enabled": config.proving.enabled,
-            "require_cuda": config.proving.require_cuda
-        }
-    }
-    
-    # Write all files
-    (p / "manifest.json").write_text(json.dumps(manifest, indent=2))
-    (p / "summary.json").write_text(json.dumps(summary, indent=2, default=str))
-    
-    # Create records directory
-    records_dir = p / "records"
-    records_dir.mkdir(exist_ok=True)
-    
-    # Write per-sample metrics (MEASURED VALUES ONLY)
-    with open(records_dir / "metrics.jsonl", "w") as f:
-        for r in per_sample_records:
-            f.write(json.dumps(r, default=str) + "\n")
-    
-    # Write KV fingerprints (MEASURED BYTES ONLY)
-    with open(records_dir / "kv_fingerprints.jsonl", "w") as f:
-        for r in per_layer_fingerprints:
-            f.write(json.dumps(r, default=str) + "\n")
-    
-    # Environment lockfile (best-effort)
-    try:
-        env_text = subprocess.check_output([sys.executable, "-m", "pip", "freeze"], text=True)
-        (p / "env.lock").write_text(env_text)
-    except Exception as e:
-        logger.warning(f"Could not capture environment: {e}")
-        (p / "env.lock").write_text(f"# Environment capture failed: {e}\n")
-    
-    # Create ZIP bundle
-    zip_path = str(p.with_suffix(".zip"))
-    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
-        for root, _, files in os.walk(p):
-            for name in files:
-                full = pathlib.Path(root) / name
-                z.write(full, arcname=str(full.relative_to(p)))
-    
-    logger.info(f"Proof bundle exported: {zip_path}")
-    return zip_path
-
-
-def verify_proof_bundle(bundle_root: str, config: CompressionConfig, proving: ProvingConfig) -> Dict[str, Any]:
-    """Verify proof bundle - recompute metrics and check tolerances. FAIL FAST on violations."""
-    # Load files
-    try:
-        with open(os.path.join(bundle_root, "summary.json")) as f:
-            summary = json.load(f)
-        
-        records = []
-        with open(os.path.join(bundle_root, "records", "metrics.jsonl")) as f:
-            for line in f:
-                if line.strip():
-                    records.append(json.loads(line))
-    except Exception as e:
-        raise RuntimeError(f"Failed to load proof bundle: {e}")
-    
-    if not records:
-        raise ValueError("No per-sample records found in proof bundle")
-    
-    # CRITICAL: Filter by compression_type to verify correct method
-    primary_method = summary.get("compression_type", summary.get("primary_method", "progressive_spg"))
-    primary_records = [r for r in records if r.get("compression_type") == primary_method]
-    
-    if not primary_records:
-        raise ValueError(f"No records found for method {primary_method}")
-    
-    logger.info(f"Verifying {len(primary_records)} records for {primary_method}")
-    
-    # Recompute aggregates from FILTERED records only
-    def mean_of(key):
-        vals = [float(r[key]) for r in primary_records if key in r and r[key] is not None]
-        return float(np.mean(vals)) if vals else None
-    
-    # Use raw bytes directly - don't recompute from shapes
-    original_bytes = mean_of("original_cache_bytes")
-    compressed_bytes = mean_of("compressed_cache_bytes")
-    
-    recomputed = {
-        "prefill_time_ms": mean_of("prefill_time") * 1000 if mean_of("prefill_time") else None,
-        "decode_time_ms": mean_of("decode_time_per_token_ms"),
-        "prefill_perplexity": mean_of("prefill_perplexity"),
-        "generation_perplexity": mean_of("generation_perplexity"),
-        "compression_ratio": original_bytes / compressed_bytes if compressed_bytes and original_bytes else None,
-        "kv_cache_memory_mb": mean_of("kv_cache_memory_mb"),  # Use directly from records
-    }
-    
-    # Numeric tolerance checks with RELAXED tolerances
-    failures = []
-    
-    # Use different tolerances for different metrics
-    for k, v in recomputed.items():
-        s = summary.get(k)
-        if v is not None and s is not None:
-            s_val = float(s)
-            
-            # Use appropriate tolerance based on metric type
-            if "time" in k or "ms" in k:
-                # Time metrics: use absolute tolerance
-                if abs(v - s_val) > proving.time_tolerance_ms:
-                    failures.append(f"{k}: recomputed {v:.3f} != summary {s_val:.3f} (tol {proving.time_tolerance_ms}ms)")
-            elif "perplexity" in k:
-                # Perplexity: use relative tolerance
-                if abs(v - s_val) / max(s_val, 1.0) > proving.ppl_tolerance:
-                    failures.append(f"{k}: recomputed {v:.3f} != summary {s_val:.3f} (rel_tol {proving.ppl_tolerance})")
-            else:
-                # Other metrics: use numeric tolerance
-                if abs(v - s_val) > proving.numeric_tolerance:
-                    failures.append(f"{k}: recomputed {v:.6f} != summary {s_val:.6f} (tol {proving.numeric_tolerance})")
-    
-    # Policy checks
-    target = config.enhanced_spg_config.target_compression_ratio
-    if recomputed["compression_ratio"] is not None:
-        if recomputed["compression_ratio"] < target * proving.comp_ratio_floor:
-            failures.append(
-                f"compression_ratio {recomputed['compression_ratio']:.2f} < "
-                f"target*floor {target * proving.comp_ratio_floor:.2f}"
-            )
-    
-    # CUDA requirement check
-    if proving.require_cuda and not torch.cuda.is_available():
-        failures.append("CUDA not available during verification (require_cuda=True)")
-    
-    ok = len(failures) == 0
-    
-    result = {
-        "ok": ok,
-        "failures": failures,
-        "recomputed": recomputed,
-        "summary": summary,
-        "n_samples": len(records)
-    }
-    
-    if not ok:
-        logger.error(f"Proof verification FAILED: {failures}")
-    else:
-        logger.info(f"Proof verification PASSED for {len(records)} samples")
-    
-    return result
-
-
-def load_real_dataset_samples(config: CompressionConfig, tokenizer) -> List[str]:
-    """Load real dataset samples with proper error handling."""
-    logger.info(f"Loading {config.eval_samples} samples from {config.dataset_name}")
-    
-    texts = []
-    min_tokens = config.prefill_length + config.generation_length
-    
-    try:
-        for split in [config.dataset_split, "train", "validation"]:
-            if len(texts) >= config.eval_samples:
-                break
-                
-            try:
-                dataset = load_dataset(
-                    config.dataset_name, 
-                    config.dataset_config,
-                    split=split,
-                    streaming=False
-                )
-                
-                logger.info(f"Trying {split} split with {len(dataset)} samples")
-                
-                for item in dataset:
-                    text = item.get('text', '').strip()
-                    
-                    if len(text) > 50:
-                        tokens = tokenizer.encode(text, truncation=False, add_special_tokens=False)
-                        
-                        if len(tokens) >= min(min_tokens, 256):
-                            texts.append(text)
-                            if len(texts) >= config.eval_samples * 3:
-                                break
-                                
-            except Exception as e:
-                logger.warning(f"Failed to load {split} split: {e}")
-                continue
-        
-        if len(texts) < config.eval_samples:
-            raise ValueError(f"Insufficient samples: {len(texts)} < {config.eval_samples}")
-            
-    except Exception as e:
-        logger.error(f"Failed to load dataset: {e}")
-        raise
-    
-    logger.info(f"Loaded {len(texts)} text samples")
-    return texts
-
-
-def run_research_benchmark(model_name: str, config: CompressionConfig, 
-                          dataset_texts: Optional[List[str]] = None) -> Tuple[BenchmarkMetrics, Dict, List[Dict], List[Dict]]:
-    """Research-grade benchmark with enhanced SPG support and fail-fast validation. Returns metrics, summary, and proof records."""
-    logger.info(f"Starting research benchmark: {model_name} with {config.compression_type.value}")
-    logger.info(f"Config hash: {config.get_hash()}")
-    
-    start_time = datetime.now().isoformat()
-    per_sample_records = []  # For proving protocol
-    per_layer_fingerprints = []  # For proving protocol
-    constants = ResearchConstants()
-    
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    dtype = torch.float16 if device == "cuda" else torch.float32
-    
-    # FAIL FAST if CUDA required but unavailable
-    if config.fail_on_cpu_fallback and device == "cpu":
-        raise RuntimeError("CUDA required but unavailable (fail_on_cpu_fallback=True)")
-    
-    if torch.cuda.is_available():
-        logger.info(f"Hardware: {torch.cuda.get_device_name()}")
-        logger.info(f"CUDA {torch.version.cuda}")
-    else:
-        logger.info("Running on CPU - performance will be limited")
-    
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-    
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=dtype,
-        device_map="auto" if device == "cuda" else None,
-        low_cpu_mem_usage=True
-    )
-    model.eval()
-    
-    try:
-        n_layers = detect_model_layers(model)
-        logger.info(f"Model architecture: {n_layers} transformer layers detected")
-    except ValueError as e:
-        logger.error(f"Failed to detect model layers: {e}")
-        raise
-    
-    # Warmup
-    with torch.inference_mode():
-        dummy = torch.randint(0, tokenizer.vocab_size, (1, config.prefill_length), device=model.device)
-        am = torch.ones_like(dummy)
-        for _ in range(config.warmup_steps):
-            _ = model(dummy, attention_mask=am, use_cache=True, return_dict=True)
-    if torch.cuda.is_available():
-        torch.cuda.synchronize()
-        torch.cuda.reset_peak_memory_stats()
-    
-    if dataset_texts is None:
-        dataset_texts = load_real_dataset_samples(config, tokenizer)
-    
-    all_metrics = []
-    
-    for seed in range(config.n_seeds):
-        set_seed(config.seed + seed)
-        logger.info(f"Running evaluation with seed {config.seed + seed}")
-        
-        metrics = BenchmarkMetrics()
-        
-        for idx in range(config.eval_samples):
-            logger.info(f"Sample {idx+1}/{config.eval_samples} (seed {config.seed + seed})")
-            
-            text_idx = (idx + seed * config.eval_samples) % len(dataset_texts)
-            text = dataset_texts[text_idx]
-            
-            cache_manager = QuantizedKVCache(config)
-            cache_manager.n_layers = n_layers
-            cache_manager.update_position(config.prefill_length + idx)
-            
-            inputs = tokenizer(
-                text,
-                return_tensors="pt",
-                truncation=True,
-                max_length=config.prefill_length,
-                padding="max_length"
-            )
-            input_ids = inputs.input_ids.to(device)
-            attention_mask = inputs.attention_mask.to(device)
-            
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-                torch.cuda.reset_peak_memory_stats()
-                torch.cuda.synchronize()
-            
-            # Prefill WITH SYNCHRONIZATION
-            if torch.cuda.is_available():
-                torch.cuda.synchronize()
-            start_time_sample = time.perf_counter()
-            with torch.inference_mode():
-                outputs = model(
-                    input_ids,
-                    attention_mask=attention_mask,
-                    use_cache=True,
-                    return_dict=True
-                )
-                past_key_values = outputs.past_key_values
-            
-            if torch.cuda.is_available():
-                torch.cuda.synchronize()
-            
-            prefill_time = time.perf_counter() - start_time_sample
-            
-            # Only track GPU memory if CUDA is available
-            if torch.cuda.is_available():
-                prefill_peak_mem = _peak_mem_bytes_all_gpus()
-                metrics.prefill_peak_memories.append(prefill_peak_mem)
-            
-            metrics.prefill_times.append(prefill_time)
-            
-            # Prefill perplexity
-            with torch.inference_mode():
-                labels = input_ids.clone()
-                labels[attention_mask == 0] = -100
-                outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
-                prefill_perplexity = torch.exp(outputs.loss).item()
-                metrics.prefill_perplexities.append(min(prefill_perplexity, 1000))
-            
-            # Compression (ACTUAL MEASURED COMPRESSION - NO ESTIMATES)
-            original_cache_size = 0
-            if past_key_values:
-                kv_tuple = past_key_values.to_legacy_cache() if hasattr(past_key_values, 'to_legacy_cache') else past_key_values
-                for layer_idx, (keys, values) in enumerate(kv_tuple):
-                    original_cache_size += keys.nelement() * keys.element_size()
-                    original_cache_size += values.nelement() * values.element_size()
-                    if config.compression_type != CompressionType.NONE:
-                        cache_manager.compress_and_store(layer_idx, keys, values)
-
-                if config.compression_type != CompressionType.NONE:
-                    reconstructed_kv = []
-                    for layer_idx in range(len(kv_tuple)):
-                        dec_keys, dec_values = cache_manager.get_decompressed(layer_idx)
-                        if dec_keys is not None and dec_values is not None:
-                            reconstructed_kv.append((dec_keys, dec_values))
-                    if hasattr(DynamicCache, 'from_legacy_cache'):
-                        past_key_values = DynamicCache.from_legacy_cache(tuple(reconstructed_kv))
-                    else:
-                        past_key_values = tuple(reconstructed_kv)
-            
-            # MEASURED compression ratio (not estimated)
-            compressed_size = original_cache_size if config.compression_type == CompressionType.NONE else cache_manager.get_memory_footprint()
-            comp_ratio = original_cache_size / compressed_size if compressed_size > 0 else 1.0
-            
-            # Log exact dtype and sequence info for verification
-            actual_seq_len = keys.shape[2] if 'keys' in locals() else config.prefill_length
-            actual_dtype_bytes = keys.element_size() if 'keys' in locals() else 2  # fp16=2, fp32=4
-            
-            # Generation
-            generated_ids = input_ids.clone()
-            decode_times = []
-            generation_losses = []
-            
-            if torch.cuda.is_available():
-                torch.cuda.reset_peak_memory_stats()
-            
-            for gen_step in range(config.generation_length):
-                if torch.cuda.is_available():
-                    torch.cuda.synchronize()
-                step_start = time.perf_counter()
-                
-                with torch.inference_mode():
-                    outputs = model(
-                        generated_ids[:, -1:],
-                        past_key_values=past_key_values,
-                        use_cache=True,
-                        return_dict=True
-                    )
-                    next_token_logits = outputs.logits[:, -1, :]
-                    # Use greedy decoding for reproducibility
-                    next_token = torch.argmax(next_token_logits, dim=-1)
-                    
-                    loss = F.cross_entropy(next_token_logits, next_token)
-                    generation_losses.append(loss.item())
-                    
-                    generated_ids = torch.cat([generated_ids, next_token.unsqueeze(-1)], dim=-1)
-                    past_key_values = outputs.past_key_values
-                
-                if torch.cuda.is_available():
-                    torch.cuda.synchronize()
-                
-                decode_time = time.perf_counter() - step_start
-                decode_times.append(decode_time)
-                
-                # Quality feedback for progressive methods (use configurable frequency)
-                feedback_frequency = config.enhanced_spg_config.quality_feedback_frequency
-                if config.compression_type in [CompressionType.ADAPTIVE_SPG, CompressionType.ENHANCED_SPG, CompressionType.PROGRESSIVE_SPG] and gen_step % feedback_frequency == 0:
-                    if len(generation_losses) >= feedback_frequency:
-                        current_ppl = np.exp(np.mean(generation_losses[-feedback_frequency:]))
-                    else:
-                        current_ppl = np.exp(np.mean(generation_losses))
-                    for layer_idx in range(n_layers):
-                        cache_manager.update_quality_feedback(layer_idx, current_ppl)
-            
-            # Record metrics
-            if decode_times:
-                metrics.decode_times.extend(decode_times)
-            
-            if torch.cuda.is_available():
-                decode_peak_mem = _peak_mem_bytes_all_gpus()
-                metrics.decode_peak_memories.append(decode_peak_mem)
-            
-            if generation_losses:
-                generation_perplexity = np.exp(np.mean(generation_losses))
-                metrics.generation_perplexities.append(min(generation_perplexity, 1000))
-            
-            # Record MEASURED compression ratios (no estimates)
-            if compressed_size > 0 and original_cache_size > 0:
-                if config.compression_type == CompressionType.NONE:
-                    metrics.compression_ratios.append(1.0)
-                else:
-                    measured_ratio = original_cache_size / compressed_size
-                    metrics.compression_ratios.append(measured_ratio)
-                    if config.compression_type in [CompressionType.ENHANCED_SPG, CompressionType.PROGRESSIVE_SPG]:
-                        metrics.enhanced_spg_measured_compression.append(measured_ratio)
-            metrics.kv_cache_memory_samples_mb.append(compressed_size / (1024 * 1024))
-            
-            # Record MEASURED auxiliary overhead (no estimates)
-            if config.compression_type in [CompressionType.ENHANCED_SPG, CompressionType.PROGRESSIVE_SPG]:
-                # Calculate actual auxiliary overhead from measured metadata
-                aux_overhead_bytes = constants.METADATA_OVERHEAD_BYTES
-                aux_overhead_mb = aux_overhead_bytes / (1024 * 1024)
-                metrics.enhanced_spg_measured_auxiliary_overhead_mb.append(aux_overhead_mb)
-                metrics.enhanced_spg_progressive_steps.append(getattr(cache_manager.spg, 'progressive_step', 0))
-            
-            # Collect per-sample record for proving protocol
-            if config.proving.export_per_sample:
-                sample_record = {
-                    "sample_idx": idx,
-                    "seed": config.seed + seed,
-                    "prefill_time": prefill_time,
-                    "decode_time_per_token_ms": float(np.mean(decode_times) * 1000) if decode_times else 0,
-                    "prefill_perplexity": min(prefill_perplexity, 1000),
-                    "generation_perplexity": min(generation_perplexity, 1000) if generation_losses else None,
-                    "compression_ratio": measured_ratio if 'measured_ratio' in locals() else 1.0,
-                    "kv_cache_memory_mb": compressed_size / (1024 * 1024),
-                    "original_cache_bytes": original_cache_size,
-                    "compressed_cache_bytes": compressed_size,
-                    "compression_type": config.compression_type.value,
-                    "seq_len_measured": actual_seq_len,
-                    "dtype_bytes": actual_dtype_bytes,
-                    "n_layers": n_layers,
-                    "is_live_kv": True  # This is live KV, not buffer capacity
-                }
-                per_sample_records.append(sample_record)
-            
-            # Collect layer fingerprints for proving protocol
-            if config.proving.export_fingerprints and config.compression_type != CompressionType.NONE:
-                for layer_idx in cache_manager.compressed_data:
-                    data = cache_manager.compressed_data[layer_idx]
-                    fingerprint = {
-                        "layer_idx": layer_idx,
-                        "sample_idx": idx,
-                        "original_shape": str(data['metadata'].get('original_shape')),
-                        "compressed_keys": len(data.get('keys', {})),
-                        "compressed_values": len(data.get('values', {})),
-                        "measured_bytes": cache_manager.spg.get_memory_footprint(data) if hasattr(cache_manager, 'spg') else 0
-                    }
-                    per_layer_fingerprints.append(fingerprint)
-        
-        metrics.calculate_statistics(config)
-        all_metrics.append(metrics)
-    
-    # Aggregate results
-    final_metrics = BenchmarkMetrics()
-    for m in all_metrics:
-        final_metrics.prefill_times.extend(m.prefill_times)
-        final_metrics.prefill_peak_memories.extend(m.prefill_peak_memories)
-        final_metrics.decode_times.extend(m.decode_times)
-        final_metrics.decode_peak_memories.extend(m.decode_peak_memories)
-        final_metrics.prefill_perplexities.extend(m.prefill_perplexities)
-        final_metrics.generation_perplexities.extend(m.generation_perplexities)
-        final_metrics.compression_ratios.extend(m.compression_ratios)
-        final_metrics.kv_cache_memory_samples_mb.extend(m.kv_cache_memory_samples_mb)
-        final_metrics.spg_effective_bits_per_token.extend(m.spg_effective_bits_per_token)
-        final_metrics.spg_precision_distributions.extend(m.spg_precision_distributions)
-        final_metrics.enhanced_spg_measured_compression.extend(m.enhanced_spg_measured_compression)
-        final_metrics.enhanced_spg_measured_auxiliary_overhead_mb.extend(m.enhanced_spg_measured_auxiliary_overhead_mb)
-        final_metrics.enhanced_spg_progressive_steps.extend(m.enhanced_spg_progressive_steps)
-    
-    final_metrics.calculate_statistics(config)
-    
-    # Summary
-    end_time = datetime.now().isoformat()
-    summary = {
-        'compression_type': config.compression_type.value,
-        'model': model_name,
-        'n_seeds': config.n_seeds,
-        'total_samples': config.eval_samples * config.n_seeds,
-        'prefill_perplexity': final_metrics.prefill_perplexity_mean,
-        'generation_perplexity': final_metrics.generation_perplexity_mean,
-        'compression_ratio': final_metrics.compression_ratio_mean,
-        'prefill_time_ms': final_metrics.prefill_time_mean * 1000,
-        'decode_time_ms': final_metrics.decode_time_per_token_mean_ms,
-        'decode_p50_ms': final_metrics.decode_time_p50_ms,
-        'decode_p95_ms': final_metrics.decode_time_p95_ms,
-        'throughput_tokens_sec': final_metrics.decode_tokens_per_sec,
-        'end_to_end_throughput': final_metrics.end_to_end_throughput,  # NEW
-        'end_to_end_latency_ms': final_metrics.end_to_end_latency_ms,  # NEW
-        'peak_memory_mb': final_metrics.prefill_peak_memory_mean_mb,
-        'kv_cache_memory_mb': final_metrics.kv_cache_memory_mb,
-        'start_time': start_time,
-        'end_time': end_time
-    }
-    
-    # Enhanced SPG summary - use measured values only
-    if config.compression_type in [CompressionType.ENHANCED_SPG, CompressionType.PROGRESSIVE_SPG]:
-        if final_metrics.enhanced_spg_measured_compression:
-            summary['enhanced_spg_measured_compression'] = np.mean(final_metrics.enhanced_spg_measured_compression)
-        if final_metrics.enhanced_spg_measured_auxiliary_overhead_mb:
-            summary['enhanced_spg_measured_auxiliary_overhead_mb'] = np.mean(final_metrics.enhanced_spg_measured_auxiliary_overhead_mb)
-        if final_metrics.enhanced_spg_progressive_steps:
-            summary['enhanced_spg_avg_progressive_steps'] = np.mean(final_metrics.enhanced_spg_progressive_steps)
-    
-    # Original SPG summary
-    if config.compression_type in [CompressionType.SPG, CompressionType.ADAPTIVE_SPG]:
-        if final_metrics.spg_effective_bits_per_token:
-            summary['spg_avg_bits_per_token'] = np.mean(final_metrics.spg_effective_bits_per_token)
-    
-    return final_metrics, summary, per_sample_records, per_layer_fingerprints
-
-
-def generate_latex_table(results: List[Dict[str, Any]]) -> str:
-    """Generate LaTeX table with enhanced SPG results."""
-    latex = r"""\begin{table}[htbp]
-\centering
-\caption{Enhanced SPG: Research Standards Compliant 450x Compression}
-\label{tab:enhanced_spg_450x_compliant}
-\begin{tabular}{lcccccccc}
-\toprule
-Method & Peak Mem. & KV Mem. & Decode & Prefill PPL & Gen. PPL & Compr. & Bits/Token & Aux. OH \\
-      & (MB)      & (MB)    & (ms/tok) &            &         & Ratio  &           & (MB) \\
-\midrule
-"""
-    
-    for result in results:
-        method = result['compression'].replace('_', r'\_')
-        peak_mem = "-" if np.isnan(result['peak_memory_mb']) else f"{result['peak_memory_mb']:.1f}"
-        kv_mem = f"{result['kv_cache_memory_mb']:.1f}"
-        decode = f"{result['decode_time_ms']:.2f}"
-        prefill_ppl = f"{result['prefill_perplexity']:.2f}"
-        gen_ppl = f"{result['generation_perplexity']:.2f}"
-        
-        if result['compression'] == 'none':
-            comp = "-"
-            bits_per_token = "16"
-            aux_overhead = "-"
-        else:
-            comp = f"{result.get('compression_ratio', 1.0):.1f}$\\times$"
-            bits_per_token = f"{result.get('spg_avg_bits_per_token', '-'):.2f}" if 'spg_avg_bits_per_token' in result else "-"
-            aux_overhead = f"{result.get('enhanced_spg_auxiliary_overhead_mb', 0):.3f}" if 'enhanced_spg_auxiliary_overhead_mb' in result else "-"
-        
-        latex += f"{method} & {peak_mem} & {kv_mem} & {decode} & {prefill_ppl} & {gen_ppl} & {comp} & {bits_per_token} & {aux_overhead} \\\\\n"
-    
-    latex += r"""\bottomrule
-\end{tabular}
-\parbox{\textwidth}{\footnotesize Enhanced SPG achieving 450x compression with full non-negotiables compliance}
-\end{table}"""
-    
-    return latex