Premchan369/Q-TensorFormer

1"""2Unified Transformer Block: Information-Value Driven Hybrid Block.3 4Integrates:5  - MultiHeadAttention with optional QKSAM and Adaptive KV Cache6  - TokenInformationState vector construction (8D)7  - InformationValueAllocator for dynamic rank, depth, and attention path8  - Nested TTFeedForward with instantaneous rank switching9  - Layer-skipping depth management10"""11 12import torch13import torch.nn as nn14from typing import Optional, Dict, Tuple15from .attention import MultiHeadAttention16from .tensor_layers import TTFeedForward17from .information_state import TokenInformationState18from .resource_allocator import InformationValueAllocator, AllocationBudget19from .kv_cache import AdaptiveKVCache20from .quantum_backend import QuantumBackend, BackendType21 22 23class HybridBlock(nn.Module):24    """25    A single Q-TensorFormer block with closed-loop resource allocation.26    """27 28    def __init__(29        self,30        d_model: int = 128,31        n_heads: int = 4,32        n_kv_heads: Optional[int] = None,33        ff_multiplier: int = 4,34        tt_rank: int = 8,35        tt_min_rank: int = 2,36        use_quantum: bool = True,37        n_qubits: int = 4,38        backend_type: str = "classical_surrogate",39        dropout: float = 0.1,40        max_seq_len: int = 128,41        hysteresis_tau: float = 0.15,42        default_preset: str = "balanced",43        enable_early_exit: bool = False,44        early_exit_threshold: float = 0.20,45        vocab_size: Optional[int] = None,46    ):47        super().__init__()48        self.d_model = d_model49        self.use_quantum = use_quantum50        self.max_seq_len = max_seq_len51        self.tt_max_rank = tt_rank52        self.n_heads = n_heads53        self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads54        self.enable_early_exit = enable_early_exit55        self.early_exit_threshold = early_exit_threshold56 57        if enable_early_exit and vocab_size is not None:58            self.early_exit_head = nn.Linear(d_model, vocab_size, bias=False)59        else:60            self.early_exit_head = None61 62        # Attention sublayer (supports GQA / MQA)63        self.attention = MultiHeadAttention(64            d_model=d_model,65            n_heads=n_heads,66            n_kv_heads=self.n_kv_heads,67            dropout=dropout,68            max_seq_len=max_seq_len,69            use_quantum_kernel=False,70            n_qubits=n_qubits,71            backend_type=backend_type,72        )73 74        # Layer norms75        self.ln1 = nn.LayerNorm(d_model)76        self.ln2 = nn.LayerNorm(d_model)77 78        # Token Information State Engine79        self.info_engine = TokenInformationState(d_model=d_model, n_heads=n_heads)80 81        # Information-Value Resource Allocator82        self.allocator = InformationValueAllocator(83            info_dim=8,84            hidden_dim=32,85            hysteresis_tau=hysteresis_tau,86            default_preset=default_preset,87        )88 89        # Nested Tensor-Train FFN90        self.tt_ffn = TTFeedForward(91            hidden_dim=d_model,92            ff_multiplier=ff_multiplier,93            rank=tt_rank,94        )95 96        # Quantum Feature Transformer (for hard tokens routed to quantum feature map)97        self.quantum_backend = QuantumBackend(98            backend_type=backend_type,99            n_qubits=n_qubits,100            n_layers=2,101            d_model=d_model,102        )103 104        self.dropout = nn.Dropout(dropout)105 106    def forward(107        self,108        x: torch.Tensor,109        mask: Optional[torch.Tensor] = None,110        kv_cache: Optional[AdaptiveKVCache] = None,111        budget: Optional[AllocationBudget] = None,112        preset: Optional[str] = None,113        force_classical: bool = False,114    ) -> Tuple[torch.Tensor, Dict]:115        """116        Closed-loop block forward pass.117 118        Flow:119          1. Construct initial information state z_t120          2. Resource Allocator chooses rank, attention mode, depth, KV precision121          3. Attention with chosen mode & KV cache122          4. FFN with dynamically sliced TT cores123          5. Selective quantum enrichment for hard tokens124          6. Residuals and layer outputs125        """126        B, T, D = x.shape127        stats = {}128 129        # 1. Probe information state130        z_t = self.info_engine(hidden_states=x)131 132        # 2. Allocate resources based on marginal utility133        decisions, diagnostics = self.allocator(134            z_t,135            budget=budget,136            preset=preset,137            force_classical=force_classical or (not self.use_quantum),138        )139 140        chosen_rank = decisions["rank"]141        depth_mode = decisions["depth_mode"]142        is_quantum_token = decisions["is_quantum_token"]  # (B, T)143        kv_precision = decisions["kv_precision"]144 145        stats.update(diagnostics)146 147        # Handle depth skipping148        if depth_mode == "skip":149            # Zero computation: skip layer via identity residual150            return x, stats151 152        # Dynamic KV Cache precision adjustment153        if kv_cache is not None:154            kv_cache.set_precision(kv_precision)155 156        # 3. Attention Sublayer157        normed_x = self.ln1(x)158        # Determine whether any token in sequence requests quantum kernel159        has_quantum_attn = (decisions["attn_mode_idx"] == 2).any().item()160 161        attn_out, attn_weights = self.attention(162            normed_x,163            mask=mask,164            kv_cache=kv_cache,165            use_quantum=has_quantum_attn and not force_classical,166            return_attn_weights=True,167        )168 169        x = x + self.dropout(attn_out)170 171        # Early exit check based on epistemic uncertainty172        if self.enable_early_exit and self.early_exit_head is not None:173            mean_uncertainty = z_t[..., 2].mean().item()174            if mean_uncertainty < self.early_exit_threshold:175                stats["early_exit_triggered"] = True176                stats["early_exit_logits"] = self.early_exit_head(self.ln2(x))177                return x, stats178 179        # If partial depth mode: skip FFN180        if depth_mode == "partial":181            return x, stats182 183        # 4. FFN Sublayer with Nested TT rank slicing184        self.tt_ffn.set_rank(chosen_rank)185        normed_ffn_in = self.ln2(x)186        token_ranks = decisions.get("token_ranks", None)187        ffn_out = self.tt_ffn(normed_ffn_in, token_ranks=token_ranks)188 189        # 5. Quantum feature enhancement on selected hard tokens190        if is_quantum_token.any() and not force_classical:191            q_out, q_meta = self.quantum_backend(normed_ffn_in)192            mask_expanded = is_quantum_token.unsqueeze(-1).float()  # (B, T, 1)193            ffn_out = ffn_out + mask_expanded * q_out194 195        x = x + self.dropout(ffn_out)196 197        # Update memory traffic stats198        ffn_traffic = self.tt_ffn.get_memory_traffic()199        stats["memory_traffic_bytes"] = ffn_traffic["total_bytes"]200 201        return x, stats202 203    def set_rank(self, rank: int):204        """Manual rank override."""205        self.tt_ffn.set_rank(rank)206 207    def reset_scheduler(self):208        """Reset allocator stability state."""209        self.allocator.reset_stability_counters()210 211    @property212    def total_params(self) -> int:213        return sum(p.numel() for p in self.parameters())214 215    def flops_estimate(self, batch_size: int = 1, seq_len: int = 32) -> Dict[str, int]:216        attn_flops = self.attention.flops(batch_size, seq_len)["total"]217        ffn_flops = self.tt_ffn.flops(batch_size, seq_len)218        return {219            "attention": attn_flops,220            "tt_ffn": ffn_flops,221            "total": attn_flops + ffn_flops,222        }223