Premchan369/Q-TensorFormer

1"""2Adaptive KV Cache Module for Q-TensorFormer.3 4Makes KV Cache memory and memory traffic first-class citizens in inference:5  - Multi-precision storage: FP16 (full), INT8 (quantized), INT4 (compressed)6  - Attention-aware rate-distortion compression7  - Budget-driven selective eviction (dropping lowest-utility tokens)8  - Detailed memory traffic instrumentation (bytes read, bytes written, peak MB)9"""10 11import torch12import torch.nn as nn13import torch.nn.functional as F14import math15from typing import Optional, Tuple, Dict, List, Union16from enum import Enum17from dataclasses import dataclass, field18 19 20class KVPrecision(str, Enum):21    FP16 = "fp16"22    INT8 = "int8"23    INT4 = "int4"24 25 26class KVResidencyTier(str, Enum):27    HOT_GPU = "hot_gpu"       # High-bandwidth GPU SRAM / HBM (immediate compute access)28    WARM_CPU = "warm_cpu"     # Host CPU-RAM (offloaded via PCIe with zero-copy pinning)29    COLD_EVICTED = "cold_evicted"  # Evicted or secondary storage30 31 32class QuantizedKVTensor:33    """34    Holds a quantized Key or Value tensor with per-channel scale and zero-point.35    """36 37    def __init__(self, tensor: torch.Tensor, precision: KVPrecision):38        self.precision = precision39        self.shape = tensor.shape40        self.device = tensor.device41 42        if precision == KVPrecision.FP16:43            self.data = tensor.to(torch.float16)44            self.scale = None45            self.zp = None46        elif precision == KVPrecision.INT8:47            # Symmetric 8-bit quantization along channel dimension48            max_val = tensor.abs().amax(dim=-1, keepdim=True).clamp(min=1e-5)49            self.scale = (max_val / 127.0).to(torch.float16)50            q = torch.round(tensor / self.scale).clamp(-128, 127).to(torch.int8)51            self.data = q52            self.zp = None53        elif precision == KVPrecision.INT4:54            # Asymmetric 4-bit packed quantization [0, 15]55            min_val = tensor.amin(dim=-1, keepdim=True)56            max_val = tensor.amax(dim=-1, keepdim=True).clamp(min=min_val + 1e-5)57            self.scale = ((max_val - min_val) / 15.0).to(torch.float16)58            self.zp = min_val.to(torch.float16)59            q = torch.round((tensor - self.zp) / self.scale).clamp(0, 15).to(torch.uint8)60            # Pack two 4-bit values into one 8-bit byte along last dimension if even61            last_dim = q.shape[-1]62            if last_dim % 2 == 0:63                q_packed = (q[..., 0::2] << 4) | (q[..., 1::2] & 0x0F)64                self.data = q_packed65                self.packed = True66            else:67                self.data = q68                self.packed = False69 70    def dequantize(self, target_dtype: torch.dtype = torch.float32) -> torch.Tensor:71        """Dequantize back to float tensor."""72        if self.precision == KVPrecision.FP16:73            return self.data.to(target_dtype)74 75        if self.precision == KVPrecision.INT8:76            return (self.data.to(target_dtype) * self.scale.to(target_dtype)).to(target_dtype)77 78        # INT479        if getattr(self, "packed", False):80            # Unpack high and low nibbles81            high = (self.data >> 4) & 0x0F82            low = self.data & 0x0F83            unpacked = torch.stack([high, low], dim=-1).reshape(self.shape)84            return (unpacked.to(target_dtype) * self.scale.to(target_dtype) + self.zp.to(target_dtype)).to(target_dtype)85        else:86            return (self.data.to(target_dtype) * self.scale.to(target_dtype) + self.zp.to(target_dtype)).to(target_dtype)87 88    @property89    def num_bytes(self) -> int:90        """Calculate physical memory in bytes."""91        total = self.data.numel() * self.data.element_size()92        if self.scale is not None:93            total += self.scale.numel() * self.scale.element_size()94        if self.zp is not None:95            total += self.zp.numel() * self.zp.element_size()96        return total97 98 99class AdaptiveKVCache:100    """101    Per-layer or per-model adaptive Key-Value cache.102 103    Supports:104      - Precision switching: FP16, INT8, INT4105      - Dynamic token retention and eviction based on attention utility106      - Memory traffic counters (bytes read/written)107    """108 109    def __init__(110        self,111        max_capacity: int = 4096,112        default_precision: KVPrecision = KVPrecision.FP16,113        eviction_policy: str = "attention_utility",  # 'attention_utility' or 'fifo'114        window_size: int = 128,  # protected recent tokens115    ):116        self.max_capacity = max_capacity117        self.precision = default_precision118        self.eviction_policy = eviction_policy119        self.window_size = window_size120 121        # Cached states: keys and values as list of tensors or QuantizedKVTensor122        self.k_cache: Optional[torch.Tensor] = None123        self.v_cache: Optional[torch.Tensor] = None124        self.quantized_k: Optional[QuantizedKVTensor] = None125        self.quantized_v: Optional[QuantizedKVTensor] = None126 127        # Attention utility score per cached token index: (seq_len,)128        self.utility_scores: Optional[torch.Tensor] = None129 130        # Traffic tracking131        self.total_bytes_written = 0132        self.total_bytes_read = 0133        self.evicted_tokens_count = 0134 135    def reset(self):136        """Clear cache state."""137        self.k_cache = None138        self.v_cache = None139        self.quantized_k = None140        self.quantized_v = None141        self.utility_scores = None142        self.total_bytes_written = 0143        self.total_bytes_read = 0144        self.evicted_tokens_count = 0145 146    @property147    def seq_len(self) -> int:148        if self.k_cache is not None:149            return self.k_cache.shape[-2]150        if self.quantized_k is not None:151            return self.quantized_k.shape[-2]152        return 0153 154    @property155    def current_bytes(self) -> int:156        """Return current memory footprint of cached keys and values in bytes."""157        if self.precision == KVPrecision.FP16 and self.k_cache is not None:158            return (self.k_cache.numel() + self.v_cache.numel()) * 2  # float16 = 2 bytes159        if self.quantized_k is not None and self.quantized_v is not None:160            return self.quantized_k.num_bytes + self.quantized_v.num_bytes161        return 0162 163    @property164    def current_mb(self) -> float:165        return self.current_bytes / (1024.0 * 1024.0)166 167    @property168    def memory_footprint_bytes(self) -> int:169        return self.current_bytes170 171    def set_precision(self, new_precision: Union[str, KVPrecision]):172        """Convert current cache in-place to new precision."""173        if isinstance(new_precision, str):174            new_precision = KVPrecision(new_precision.lower())175 176        if new_precision == self.precision:177            return178 179        if self.seq_len > 0:180            k, v = self.get_kv()181            self.precision = new_precision182            if self.precision == KVPrecision.FP16:183                self.k_cache = k.to(torch.float16)184                self.v_cache = v.to(torch.float16)185                self.quantized_k = None186                self.quantized_v = None187            else:188                self.quantized_k = QuantizedKVTensor(k, self.precision)189                self.quantized_v = QuantizedKVTensor(v, self.precision)190                self.k_cache = None191                self.v_cache = None192        else:193            self.precision = new_precision194 195    def update(196        self,197        key: torch.Tensor,198        value: torch.Tensor,199        attention_weights: Optional[torch.Tensor] = None,200    ) -> Tuple[torch.Tensor, torch.Tensor]:201        """202        Append new key/value tokens to the cache, evicting if necessary.203 204        Args:205            key: (batch, n_heads, seq_len_new, head_dim)206            value: (batch, n_heads, seq_len_new, head_dim)207            attention_weights: optional attention distribution to update utility scores208        Returns:209            full_keys: (batch, n_heads, total_seq_len, head_dim)210            full_values: (batch, n_heads, total_seq_len, head_dim)211        """212        # Count incoming memory traffic213        bytes_in = (key.numel() + value.numel()) * key.element_size()214        self.total_bytes_written += bytes_in215 216        # Retrieve existing float representation217        if self.seq_len == 0:218            curr_k = key219            curr_v = value220            new_tokens = key.shape[-2]221            self.utility_scores = torch.ones(new_tokens, device=key.device)222        else:223            old_k, old_v = self.get_kv(target_dtype=key.dtype)224            curr_k = torch.cat([old_k, key], dim=-2)225            curr_v = torch.cat([old_v, value], dim=-2)226            new_tokens = key.shape[-2]227            new_scores = torch.ones(new_tokens, device=key.device)228            self.utility_scores = torch.cat([self.utility_scores, new_scores], dim=0)229 230        # Update utility scores from attention weights if provided231        if attention_weights is not None:232            with torch.no_grad():233                # Incoming attention received by cached tokens: (B, H, Q, K) -> sum over Q, avg over B, H234                attn_importance = attention_weights.sum(dim=-2).mean(dim=(0, 1))  # (K,)235                n_tokens = min(len(self.utility_scores), len(attn_importance))236                self.utility_scores[:n_tokens] = 0.9 * self.utility_scores[:n_tokens] + 0.1 * attn_importance[:n_tokens]237 238        # Eviction if capacity exceeded239        curr_len = curr_k.shape[-2]240        if curr_len > self.max_capacity:241            curr_k, curr_v = self._evict(curr_k, curr_v, target_len=self.max_capacity)242 243        # Store in configured precision244        if self.precision == KVPrecision.FP16:245            self.k_cache = curr_k.to(torch.float16)246            self.v_cache = curr_v.to(torch.float16)247            self.quantized_k = None248            self.quantized_v = None249        else:250            self.quantized_k = QuantizedKVTensor(curr_k, self.precision)251            self.quantized_v = QuantizedKVTensor(curr_v, self.precision)252            self.k_cache = None253            self.v_cache = None254 255        # Return full dequantized tensors for attention computation256        out_k, out_v = self.get_kv(target_dtype=key.dtype)257 258        # Count read traffic259        bytes_out = (out_k.numel() + out_v.numel()) * out_k.element_size()260        self.total_bytes_read += bytes_out261 262        return out_k, out_v263 264    def _evict(self, k: torch.Tensor, v: torch.Tensor, target_len: int) -> Tuple[torch.Tensor, torch.Tensor]:265        """Evict lowest utility tokens while preserving sink tokens (first 4) and recent window."""266        total_len = k.shape[-2]267        num_to_evict = total_len - target_len268        if num_to_evict <= 0:269            return k, v270 271        # Always protect sink tokens (e.g. first 4) and recent window tokens272        sink_size = min(4, total_len)273        window_size = min(self.window_size, total_len - sink_size)274        candidate_end = total_len - window_size275 276        if candidate_end <= sink_size:277            # If sequence is mostly window, just do FIFO truncation from start278            self.evicted_tokens_count += num_to_evict279            self.utility_scores = self.utility_scores[num_to_evict:]280            return k[..., num_to_evict:, :], v[..., num_to_evict:, :]281 282        candidate_scores = self.utility_scores[sink_size:candidate_end]283        # Find indices with highest utility to KEEP284        num_candidates_to_keep = (candidate_end - sink_size) - num_to_evict285        if num_candidates_to_keep <= 0:286            keep_indices = torch.tensor([], dtype=torch.long, device=k.device)287        else:288            _, keep_rel = torch.topk(candidate_scores, k=num_candidates_to_keep, largest=True, sorted=True)289            keep_indices = sink_size + keep_rel.sort().values290 291        # Concatenate: sink + kept candidates + recent window292        sink_indices = torch.arange(0, sink_size, device=k.device)293        window_indices = torch.arange(candidate_end, total_len, device=k.device)294        final_indices = torch.cat([sink_indices, keep_indices, window_indices], dim=0)295 296        self.evicted_tokens_count += num_to_evict297        self.utility_scores = self.utility_scores[final_indices]298        return k[..., final_indices, :], v[..., final_indices, :]299 300    def get_kv(self, target_dtype: torch.dtype = torch.float32) -> Tuple[torch.Tensor, torch.Tensor]:301        """Retrieve dequantized full key and value tensors."""302        if self.seq_len == 0:303            raise ValueError("Cache is empty.")304 305        if self.precision == KVPrecision.FP16:306            return self.k_cache.to(target_dtype), self.v_cache.to(target_dtype)307        else:308            return self.quantized_k.dequantize(target_dtype), self.quantized_v.dequantize(target_dtype)309 310    def stats(self) -> Dict[str, Union[float, int, str]]:311        """Return diagnostic metrics for cache monitoring."""312        return {313            "seq_len": self.seq_len,314            "precision": self.precision.value,315            "footprint_mb": round(self.current_mb, 4),316            "bytes_written": self.total_bytes_written,317            "bytes_read": self.total_bytes_read,318            "evicted_tokens": self.evicted_tokens_count,319        }320 321 322@dataclass323class KVTransferStats:324    total_migrated_bytes: int = 0325    migration_latency_ms: float = 0.0326    cache_hits: int = 0327    cache_misses: int = 0328    hot_tokens: int = 0329    warm_tokens: int = 0330    evicted_tokens: int = 0331    fragmentation_ratio: float = 0.0332 333 334class HierarchicalAdaptiveKVCache(AdaptiveKVCache):335    """336    Hierarchical Adaptive Memory Hierarchy for Q-TensorFormer KV Cache.337    Aligned with MetaKV (prompt-level constrained budget selection) and SeKV (hierarchical GPU/CPU residency).338 339    Supports:340      - 3 residency tiers: HOT (GPU), WARM (Host CPU-RAM), COLD (evicted).341      - Explicit PCIe transfer latency modeling: T_mig = Bytes / BW_PCIe + T_launch.342      - Attention utility + recency + sink token protection.343      - Request-level adaptive budget policy selection.344      - Fragmentation and migration profiling.345    """346 347    def __init__(348        self,349        max_capacity: int = 2048,350        hot_capacity: int = 512,351        warm_capacity: int = 2048,352        window_size: int = 64,353        pcie_bandwidth_gb_s: float = 32.0,354        pcie_launch_overhead_ms: float = 0.015,355        default_precision: KVPrecision = KVPrecision.FP16,356        default_residency: KVResidencyTier = KVResidencyTier.HOT_GPU,357    ):358        super().__init__(359            max_capacity=max_capacity,360            window_size=window_size,361            default_precision=default_precision,362        )363        self.hot_capacity = hot_capacity364        self.warm_capacity = warm_capacity365        self.pcie_bandwidth_gb_s = pcie_bandwidth_gb_s366        self.pcie_launch_overhead_ms = pcie_launch_overhead_ms367        self.residency = default_residency368 369        self.transfer_stats = KVTransferStats()370        self.warm_k: Optional[QuantizedKVTensor] = None371        self.warm_v: Optional[QuantizedKVTensor] = None372 373    def update_hierarchical(374        self,375        key: torch.Tensor,376        value: torch.Tensor,377        attention_weights: Optional[torch.Tensor] = None,378        target_precision: Optional[KVPrecision] = None,379        target_residency: Optional[KVResidencyTier] = None,380        prompt_budget_ratio: Optional[float] = None,381    ) -> Tuple[torch.Tensor, torch.Tensor]:382        """383        Hierarchically updates cache with MetaKV-style prompt budget scaling384        and SeKV-style tiered GPU/CPU placement.385        """386        if target_precision is not None:387            self.set_precision(target_precision)388        if target_residency is not None:389            self.residency = target_residency390 391        # MetaKV adaptation: adjust active hot capacity based on prompt budget ratio392        effective_hot_cap = self.hot_capacity393        if prompt_budget_ratio is not None:394            effective_hot_cap = max(8, int(self.hot_capacity * prompt_budget_ratio))395 396        # Standard in-layer update397        out_k, out_v = self.update(key, value, attention_weights=attention_weights)398 399        curr_len = out_k.shape[-2]400        if curr_len > effective_hot_cap and self.residency == KVResidencyTier.WARM_CPU:401            # Demote overflow tokens from HOT to WARM CPU tier402            overflow_tokens = curr_len - effective_hot_cap403            migrated_bytes = overflow_tokens * out_k.shape[-1] * out_k.element_size() * 2404            # Explicit PCIe migration latency calculation405            transfer_ms = (migrated_bytes / (self.pcie_bandwidth_gb_s * 1e9)) * 1000.0 + self.pcie_launch_overhead_ms406            self.transfer_stats.total_migrated_bytes += migrated_bytes407            self.transfer_stats.migration_latency_ms += transfer_ms408            self.transfer_stats.warm_tokens += overflow_tokens409            self.transfer_stats.hot_tokens = effective_hot_cap410            self.transfer_stats.cache_misses += 1411        else:412            self.transfer_stats.hot_tokens = curr_len413            self.transfer_stats.cache_hits += 1414 415        self.transfer_stats.fragmentation_ratio = round(self.evicted_tokens_count / max(1, curr_len + self.evicted_tokens_count), 4)416        return out_k, out_v417 418    def stats(self) -> Dict[str, Union[float, int, str]]:419        base_stats = super().stats()420        total_accesses = max(1, self.transfer_stats.cache_hits + self.transfer_stats.cache_misses)421        base_stats.update({422            "residency_tier": self.residency.value,423            "hot_tokens": self.transfer_stats.hot_tokens,424            "warm_tokens": self.transfer_stats.warm_tokens,425            "migrated_bytes": self.transfer_stats.total_migrated_bytes,426            "migration_latency_ms": round(self.transfer_stats.migration_latency_ms, 3),427            "cache_hit_rate": round(self.transfer_stats.cache_hits / total_accesses, 3),428            "fragmentation_ratio": self.transfer_stats.fragmentation_ratio,429        })430        return base_stats431