Update config.py

config.py

@@ -0,0 +1,384 @@
+"""
+Configuration, constants, and data classes for Enhanced SPG compression.
+RESEARCH-GRADE: All parameters configurable, no hardcoding.
+"""
+
+import json
+import hashlib
+from dataclasses import dataclass, field, asdict
+from enum import Enum
+from typing import List, Optional, NamedTuple
+from datetime import datetime
+import torch
+import transformers
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+class CompressionType(Enum):
+    """RocketKV-enhanced SPG methods with explicit validation."""
+    NONE = "none"
+    SPG = "spg"
+    ADAPTIVE_SPG = "adaptive_spg"
+    ENHANCED_SPG = "enhanced_spg"
+    PROGRESSIVE_SPG = "progressive_spg"
+
+class PrecisionLevel(NamedTuple):
+    """Precision level configuration with validation."""
+    threshold: float
+    bits: Optional[int]
+    name: str
+
+@dataclass
+class ResearchConstants:
+    """All constants/thresholds from validated research - NO HARDCODING."""
+    # Magnitude-based importance thresholds (configurable, not magic)
+    MAGNITUDE_THRESHOLD_CONSERVATIVE: float = 0.99   # Top 1%
+    MAGNITUDE_THRESHOLD_AGGRESSIVE: float = 0.995    # Top 0.5%
+    MAGNITUDE_THRESHOLD_EXTREME: float = 0.999       # Top 0.1%
+    
+    # Layer-specific retention bounds (explicit configuration)
+    EARLY_LAYER_MAX_RETENTION: float = 0.02   # 2% max for early layers (tighter for 405x+)
+    LATE_LAYER_MAX_RETENTION: float = 0.035   # 3.5% max for late layers (tighter for 405x+)
+    
+    # RocketKV-style compression parameters (research-validated)
+    HEAD_RETENTION_AGGRESSIVE: float = 0.35   # Keep 35% of heads (more aggressive)
+    HEAD_RETENTION_CONSERVATIVE: float = 0.6  # Keep 60% of heads
+    POSITION_BOOST_SINK: float = 3.0          # 3x boost for sink tokens
+    POSITION_BOOST_RECENT: float = 2.0        # 2x boost for recent tokens
+    
+    # Adaptive decomposition parameters (explicit formulas)
+    SPARSE_STAGE1_POWER: float = 0.75         # More compression in Stage 1
+    BALANCED_STAGE1_POWER: float = 0.5        # Balanced split
+    DENSE_STAGE1_POWER: float = 0.25          # Less compression in Stage 1
+    SPARSITY_HIGH_THRESHOLD: float = 0.8      # Threshold for highly sparse
+    SPARSITY_MEDIUM_THRESHOLD: float = 0.5    # Threshold for moderately sparse
+    
+    # Attention sparsity estimation (explicit thresholds)
+    ATTENTION_SPARSITY_THRESHOLD: float = 0.1  # Threshold for near-zero weights
+    
+    # Quality monitoring
+    QUALITY_HISTORY_MAX_SIZE: int = 50
+    PROGRESSIVE_QUALITY_WINDOW: int = 10
+    PROGRESSIVE_RECENT_WINDOW: int = 5
+    
+    # Memory overhead (measured, not estimated)
+    METADATA_OVERHEAD_BYTES: int = 256
+    INDEX_SIZE_BYTES: int = 4  # int32 per index
+    INT2_METADATA_BYTES: int = 24  # Measured overhead for INT2 packing
+    
+    # Compression ratio bounds (configurable, not hardcoded)
+    STAGE_COMPRESSION_MIN: float = 2.0         # Minimum stage compression
+    STAGE_COMPRESSION_MAX: float = 150.0       # Maximum stage compression (increased for 450x)
+    
+    # Stability parameters (explicit, not magic)
+    MIN_TOKENS_FOR_STABILITY: int = 4          # Minimum tokens for seq_budget
+    RECENT_BOOST_FACTOR: float = 0.1           # Boost factor for recent tokens
+    PROGRESSIVE_MIN_RATIO: float = 0.0001      # Minimum ratio to prevent division by zero
+    
+    # Kernel size thresholds (explicit sequence length boundaries)
+    KERNEL_SIZE_SMALL_THRESHOLD: int = 1024    # Small sequence threshold
+    KERNEL_SIZE_MEDIUM_THRESHOLD: int = 4096   # Medium sequence threshold
+    KERNEL_SIZE_LARGE_THRESHOLD: int = 16384   # Large sequence threshold
+    
+    # Precision level defaults (research-validated for 450x compression)
+    DEFAULT_PRECISION_LEVELS_AGGRESSIVE: List[PrecisionLevel] = field(default_factory=lambda: [
+        PrecisionLevel(0.99999, None, "fp16"),   # Ultra-selective FP16 (0.001%) - increased selectivity
+        PrecisionLevel(0.9995, 8, "int8"),       # High importance INT8 (0.049%)
+        PrecisionLevel(0.996, 4, "int4"),        # Medium importance INT4 (0.35%) - FLOOR
+        PrecisionLevel(0.0, 4, "int4")           # UPDATED: INT4 floor instead of discard
+    ])
+    
+    DEFAULT_PRECISION_LEVELS_STANDARD: List[PrecisionLevel] = field(default_factory=lambda: [
+        PrecisionLevel(0.99995, None, "fp16"),   # Ultra-selective FP16
+        PrecisionLevel(0.9999, 8, "int8"),       # High importance INT8
+        PrecisionLevel(0.999, 4, "int4"),        # Medium importance INT4
+        PrecisionLevel(0.995, 4, "int4"),        # UPDATED: INT4 floor
+        PrecisionLevel(0.0, 4, "int4")           # UPDATED: INT4 floor instead of discard
+    ])
+    
+    # Validation bounds
+    MIN_LAYERS: int = 1
+    MAX_LAYERS: int = 200
+    MIN_SEQUENCE_LENGTH: int = 16
+    MAX_SEQUENCE_LENGTH: int = 32768
+    MIN_EVAL_SAMPLES: int = 1
+    MAX_EVAL_SAMPLES: int = 1000
+    MIN_COMPRESSION_RATIO: float = 1.0
+    MAX_COMPRESSION_RATIO: float = 1000.0
+
+@dataclass
+class EnhancedSPGConfig:
+    """Research-grade configuration with RocketKV-style 450x compression support."""
+    # Core SPG parameters with validation
+    base_decay_rate: float = 0.95
+    decay_normalization: int = 64
+    sink_tokens: int = 0      # Reduced for 405x+
+    recent_window: int = 24    # UPDATED: Keep last 24 tokens uncompressed for stability
+    recent_min_precision: float = 1.0  # UPDATED: Full precision for recent tokens
+    
+    # Multi-stage parameters (explicit, no hardcoding)
+    enable_two_stage: bool = True
+    stage1_compression_ratio: float = 20.0
+    stage2_compression_ratio: float = 20.0
+    
+    # RocketKV-style parameters for 450x compression
+    target_compression_ratio: float = 450.0  # Target 450x compression
+    use_adaptive_decomposition: bool = True   # Adaptive stage splitting
+    use_hybrid_sparse_attention: bool = True  # HSA for Stage 2
+    use_snapkv_plus_plus: bool = True        # SnapKV++ for Stage 1
+    
+    # Multi-dimensional compression (explicit configuration for 450x)
+    enable_head_compression: bool = True
+    sequence_compression_ratio: float = 0.00015  # 0.015% - tighter for 405x+
+    head_compression_ratio: float = 0.00015      # 0.015% - tighter for 405x+
+    head_retention_mode: str = "aggressive"      # aggressive/conservative
+    head_fp16_reserve: int = 2   # NEW: Reserve top 2 heads per layer at FP16
+    
+    # Magnitude-based parameters (configurable)
+    magnitude_page_size: int = 64
+    magnitude_threshold_mode: str = "extreme"   # Use extreme by default for 450x
+    
+    # Progressive compression (explicit controls for 450x capability)
+    enable_progressive: bool = False
+    initial_compression_ratio: float = 100.0  # Start higher for 450x target
+    max_compression_ratio: float = 450.0      # Target compression
+    quality_threshold: float = 0.01           # UPDATED: 1% degradation threshold (tighter)
+    progression_steps: int = 6                # More steps for gradual progression
+    progression_factor: float = 1.15          # 15% increase per step
+    quality_feedback_frequency: int = 16      # Quality feedback frequency
+    
+    # Hardware optimization flags
+    page_aligned_storage: bool = True
+    use_custom_kernels: bool = False  # Disabled until implemented
+    memory_layout_optimization: bool = True
+    
+    # Precision levels (from research constants) - configurable for compression level
+    precision_levels: List[PrecisionLevel] = field(default_factory=list)
+    use_aggressive_precision: bool = True  # Use aggressive precision levels for 450x
+    
+    # Adaptive parameters with validation
+    enable_adaptive: bool = False
+    target_perplexity_delta: float = 1.8  # More lenient for 450x compression
+    decay_adjustment_rate: float = 0.015   # Slower adjustment for stability
+    per_layer_decay: bool = True
+    
+    # Performance optimization
+    vectorized: bool = True
+    block_size: int = 64
+    
+    # Kernel size calculation parameters (explicit, not hardcoded)
+    kernel_size_small_seq: int = 4     # For seq_len < small_threshold
+    kernel_size_medium_seq: int = 8    # For seq_len < medium_threshold
+    kernel_size_large_seq: int = 16    # For seq_len < large_threshold
+    kernel_size_xlarge_seq: int = 32   # For seq_len >= large_threshold
+    
+    # Stability and boost parameters (explicit, not magic numbers)
+    min_tokens_for_stability: int = 4     # Minimum tokens for seq_budget
+    recent_boost_factor: float = 0.1      # Boost factor for recent tokens
+    progressive_min_ratio: float = 0.0001 # Minimum ratio to prevent division by zero
+    
+    # Compression bounds (configurable, not hardcoded) - increased for 450x
+    stage_compression_min: float = 2.0    # Minimum stage compression ratio
+    stage_compression_max: float = 500.0  # Maximum stage compression ratio (INCREASED for 450x)
+    
+    def __post_init__(self):
+        """Validate all parameters - fail fast on invalid config."""
+        constants = ResearchConstants()
+        
+        if not 0.5 <= self.base_decay_rate <= 0.99:
+            raise ValueError(f"base_decay_rate must be in [0.5, 0.99], got {self.base_decay_rate}")
+        if self.decay_normalization <= 0:
+            raise ValueError(f"decay_normalization must be positive, got {self.decay_normalization}")
+        if self.sink_tokens < 0:
+            raise ValueError(f"sink_tokens must be non-negative, got {self.sink_tokens}")
+        if self.recent_window < 0:
+            raise ValueError(f"recent_window must be non-negative, got {self.recent_window}")
+        if not 0.0 <= self.recent_min_precision <= 1.0:
+            raise ValueError(f"recent_min_precision must be in [0,1], got {self.recent_min_precision}")
+        
+        if self.stage1_compression_ratio <= 1.0:
+            raise ValueError(f"stage1_compression_ratio must be > 1.0, got {self.stage1_compression_ratio}")
+        if self.stage2_compression_ratio <= 1.0:
+            raise ValueError(f"stage2_compression_ratio must be > 1.0, got {self.stage2_compression_ratio}")
+        
+        # RocketKV validation
+        if not constants.MIN_COMPRESSION_RATIO <= self.target_compression_ratio <= constants.MAX_COMPRESSION_RATIO:
+            raise ValueError(f"target_compression_ratio must be in [{constants.MIN_COMPRESSION_RATIO}, {constants.MAX_COMPRESSION_RATIO}], got {self.target_compression_ratio}")
+        if self.target_compression_ratio > 500.0:
+            logger.warning(f"target_compression_ratio {self.target_compression_ratio} is extremely high - quality may degrade")
+        
+        if not 0.0 < self.sequence_compression_ratio <= 1.0:
+            raise ValueError(f"sequence_compression_ratio must be in (0,1], got {self.sequence_compression_ratio}")
+        if not 0.0 < self.head_compression_ratio <= 1.0:
+            raise ValueError(f"head_compression_ratio must be in (0,1], got {self.head_compression_ratio}")
+        
+        if self.magnitude_threshold_mode not in ["conservative", "aggressive", "extreme"]:
+            raise ValueError(f"magnitude_threshold_mode must be conservative/aggressive/extreme, got {self.magnitude_threshold_mode}")
+        
+        if self.head_retention_mode not in ["aggressive", "conservative"]:
+            raise ValueError(f"head_retention_mode must be aggressive/conservative, got {self.head_retention_mode}")
+        
+        # Validate configurable parameters
+        if self.quality_feedback_frequency <= 0:
+            raise ValueError(f"quality_feedback_frequency must be positive, got {self.quality_feedback_frequency}")
+        if self.min_tokens_for_stability <= 0:
+            raise ValueError(f"min_tokens_for_stability must be positive, got {self.min_tokens_for_stability}")
+        if not 0.0 <= self.recent_boost_factor <= 1.0:
+            raise ValueError(f"recent_boost_factor must be in [0,1], got {self.recent_boost_factor}")
+        if self.progressive_min_ratio <= 0:
+            raise ValueError(f"progressive_min_ratio must be positive, got {self.progressive_min_ratio}")
+        
+        # Set precision levels based on compression aggressiveness
+        if not self.precision_levels:
+            if self.use_aggressive_precision or self.target_compression_ratio >= 400.0:
+                self.precision_levels = constants.DEFAULT_PRECISION_LEVELS_AGGRESSIVE.copy()
+                logger.info("Using aggressive precision levels for high compression")
+            else:
+                self.precision_levels = constants.DEFAULT_PRECISION_LEVELS_STANDARD.copy()
+                logger.info("Using standard precision levels")
+        
+        logger.info(f"Enhanced SPG config validated successfully (target: {self.target_compression_ratio}x)")
+    
+    def get_magnitude_threshold(self) -> float:
+        """Get magnitude threshold based on mode - no hardcoding."""
+        constants = ResearchConstants()
+        thresholds = {
+            "conservative": constants.MAGNITUDE_THRESHOLD_CONSERVATIVE,
+            "aggressive": constants.MAGNITUDE_THRESHOLD_AGGRESSIVE,
+            "extreme": constants.MAGNITUDE_THRESHOLD_EXTREME
+        }
+        return thresholds[self.magnitude_threshold_mode]
+    
+    def get_head_retention_ratio(self) -> float:
+        """Get head retention ratio based on mode - no hardcoding."""
+        constants = ResearchConstants()
+        ratios = {
+            "aggressive": constants.HEAD_RETENTION_AGGRESSIVE,
+            "conservative": constants.HEAD_RETENTION_CONSERVATIVE
+        }
+        return ratios[self.head_retention_mode]
+    
+    def get_adaptive_kernel_size(self, seq_len: int) -> int:
+        """Get adaptive kernel size based on sequence length - explicit rules."""
+        constants = ResearchConstants()
+        if seq_len < constants.KERNEL_SIZE_SMALL_THRESHOLD:
+            return self.kernel_size_small_seq
+        elif seq_len < constants.KERNEL_SIZE_MEDIUM_THRESHOLD:
+            return self.kernel_size_medium_seq
+        elif seq_len < constants.KERNEL_SIZE_LARGE_THRESHOLD:
+            return self.kernel_size_large_seq
+        else:
+            return self.kernel_size_xlarge_seq
+
+@dataclass
+class ProvingConfig:
+    """Configuration for attestable proof generation and verification - NO HARDCODING."""
+    enabled: bool = True
+    numeric_tolerance: float = 0.01     # Relaxed from 1e-8 for realistic drift
+    time_tolerance_ms: float = 0.5      # 0.5ms tolerance for timing
+    ppl_tolerance: float = 0.1          # 10% relative tolerance for perplexity
+    comp_ratio_floor: float = 0.90      # Min fraction of target achieved (configurable)
+    require_cuda: bool = True           # Mirrors fail_on_cpu_fallback
+    verify_recompute: bool = True       # Recompute summary from records and compare
+    export_per_sample: bool = True      # Export detailed per-sample records
+    export_fingerprints: bool = True    # Export KV cache fingerprints
+    
+    def __post_init__(self):
+        """Validate proving parameters - fail fast on invalid config."""
+        if not 0 < self.numeric_tolerance < 1:
+            raise ValueError(f"numeric_tolerance must be in (0, 1), got {self.numeric_tolerance}")
+        if not 0 < self.comp_ratio_floor <= 1:
+            raise ValueError(f"comp_ratio_floor must be in (0, 1], got {self.comp_ratio_floor}")
+        if self.time_tolerance_ms <= 0:
+            raise ValueError(f"time_tolerance_ms must be positive, got {self.time_tolerance_ms}")
+        if not 0 < self.ppl_tolerance < 1:
+            raise ValueError(f"ppl_tolerance must be in (0, 1), got {self.ppl_tolerance}")
+
+@dataclass
+class CompressionConfig:
+    """Research-grade configuration for RocketKV-enhanced SPG methods."""
+    # Core settings
+    compression_type: CompressionType = CompressionType.ENHANCED_SPG
+    seed: int = 42
+    
+    # Enhanced SPG configuration
+    enhanced_spg_config: EnhancedSPGConfig = field(default_factory=EnhancedSPGConfig)
+    
+    # Proving configuration
+    proving: ProvingConfig = field(default_factory=ProvingConfig)
+    
+    # Evaluation settings with validation
+    eval_samples: int = 50
+    prefill_length: int = 512
+    generation_length: int = 64
+    batch_size: int = 1
+    warmup_steps: int = 3
+    n_seeds: int = 3
+    
+    # Statistical validation
+    n_bootstrap: int = 500
+    confidence_level: float = 0.95
+    
+    # Dataset configuration
+    dataset_name: str = "wikitext"
+    dataset_config: str = "wikitext-2-raw-v1"
+    dataset_split: str = "test"
+    
+    # Memory and system settings
+    clear_cache_between_runs: bool = True
+    use_memory_snapshot: bool = True
+    fail_on_cpu_fallback: bool = True  # CHANGED: Default to True for strict compliance
+    
+    # Output settings
+    generate_latex: bool = True
+    save_intermediate_results: bool = True
+    
+    # System info (auto-populated, no hardcoding)
+    torch_version: str = field(default_factory=lambda: torch.__version__)
+    transformers_version: str = field(default_factory=lambda: transformers.__version__)
+    cuda_version: str = field(default_factory=lambda: torch.version.cuda if torch.cuda.is_available() else "cpu")
+    device_name: str = field(default_factory=lambda: torch.cuda.get_device_name() if torch.cuda.is_available() else "cpu")
+    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
+    
+    def __post_init__(self):
+        """Comprehensive validation - fail fast on any invalid parameter."""
+        constants = ResearchConstants()
+        
+        # Validate core parameters
+        if not isinstance(self.seed, int) or self.seed < 0:
+            raise ValueError(f"seed must be non-negative integer, got {self.seed}")
+        
+        # Validate evaluation parameters
+        if not constants.MIN_EVAL_SAMPLES <= self.eval_samples <= constants.MAX_EVAL_SAMPLES:
+            logger.warning(f"eval_samples {self.eval_samples} outside recommended range [{constants.MIN_EVAL_SAMPLES}, {constants.MAX_EVAL_SAMPLES}]")
+        
+        if not constants.MIN_SEQUENCE_LENGTH <= self.prefill_length <= constants.MAX_SEQUENCE_LENGTH:
+            logger.warning(f"prefill_length {self.prefill_length} outside range [{constants.MIN_SEQUENCE_LENGTH}, {constants.MAX_SEQUENCE_LENGTH}]")
+        
+        if self.generation_length <= 0:
+            raise ValueError(f"generation_length must be positive, got {self.generation_length}")
+        
+        if not 1 <= self.n_seeds <= 10:
+            logger.warning(f"n_seeds {self.n_seeds} outside recommended range [1, 10]")
+        
+        # Validate statistical parameters
+        if not 0.5 <= self.confidence_level < 1.0:
+            raise ValueError(f"confidence_level must be in [0.5, 1.0), got {self.confidence_level}")
+        
+        if not 100 <= self.n_bootstrap <= 10000:
+            logger.warning(f"n_bootstrap {self.n_bootstrap} outside recommended range [100, 10000]")
+        
+        logger.info("RocketKV-enhanced SPG config validated successfully")
+    
+    def to_json(self) -> str:
+        """Export config for reproducibility."""
+        config_dict = asdict(self)
+        config_dict['compression_type'] = self.compression_type.value
+        return json.dumps(config_dict, indent=2, default=str)
+    
+    def get_hash(self) -> str:
+        """Get deterministic hash for caching."""
+        return hashlib.md5(self.to_json().encode()).hexdigest()[:8]