""" Unified Transformer Block: Information-Value Driven Hybrid Block. Integrates: - MultiHeadAttention with optional QKSAM and Adaptive KV Cache - TokenInformationState vector construction (8D) - InformationValueAllocator for dynamic rank, depth, and attention path - Nested TTFeedForward with instantaneous rank switching - Layer-skipping depth management """ import torch import torch.nn as nn from typing import Optional, Dict, Tuple from .attention import MultiHeadAttention from .tensor_layers import TTFeedForward from .information_state import TokenInformationState from .resource_allocator import InformationValueAllocator, AllocationBudget from .kv_cache import AdaptiveKVCache from .quantum_backend import QuantumBackend, BackendType class HybridBlock(nn.Module): """ A single Q-TensorFormer block with closed-loop resource allocation. """ def __init__( self, d_model: int = 128, n_heads: int = 4, n_kv_heads: Optional[int] = None, ff_multiplier: int = 4, tt_rank: int = 8, tt_min_rank: int = 2, use_quantum: bool = True, n_qubits: int = 4, backend_type: str = "classical_surrogate", dropout: float = 0.1, max_seq_len: int = 128, hysteresis_tau: float = 0.15, default_preset: str = "balanced", enable_early_exit: bool = False, early_exit_threshold: float = 0.20, vocab_size: Optional[int] = None, ): super().__init__() self.d_model = d_model self.use_quantum = use_quantum self.max_seq_len = max_seq_len self.tt_max_rank = tt_rank self.n_heads = n_heads self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads self.enable_early_exit = enable_early_exit self.early_exit_threshold = early_exit_threshold if enable_early_exit and vocab_size is not None: self.early_exit_head = nn.Linear(d_model, vocab_size, bias=False) else: self.early_exit_head = None # Attention sublayer (supports GQA / MQA) self.attention = MultiHeadAttention( d_model=d_model, n_heads=n_heads, n_kv_heads=self.n_kv_heads, dropout=dropout, max_seq_len=max_seq_len, use_quantum_kernel=False, n_qubits=n_qubits, backend_type=backend_type, ) # Layer norms self.ln1 = nn.LayerNorm(d_model) self.ln2 = nn.LayerNorm(d_model) # Token Information State Engine self.info_engine = TokenInformationState(d_model=d_model, n_heads=n_heads) # Information-Value Resource Allocator self.allocator = InformationValueAllocator( info_dim=8, hidden_dim=32, hysteresis_tau=hysteresis_tau, default_preset=default_preset, ) # Nested Tensor-Train FFN self.tt_ffn = TTFeedForward( hidden_dim=d_model, ff_multiplier=ff_multiplier, rank=tt_rank, ) # Quantum Feature Transformer (for hard tokens routed to quantum feature map) self.quantum_backend = QuantumBackend( backend_type=backend_type, n_qubits=n_qubits, n_layers=2, d_model=d_model, ) self.dropout = nn.Dropout(dropout) def forward( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, kv_cache: Optional[AdaptiveKVCache] = None, budget: Optional[AllocationBudget] = None, preset: Optional[str] = None, force_classical: bool = False, ) -> Tuple[torch.Tensor, Dict]: """ Closed-loop block forward pass. Flow: 1. Construct initial information state z_t 2. Resource Allocator chooses rank, attention mode, depth, KV precision 3. Attention with chosen mode & KV cache 4. FFN with dynamically sliced TT cores 5. Selective quantum enrichment for hard tokens 6. Residuals and layer outputs """ B, T, D = x.shape stats = {} # 1. Probe information state z_t = self.info_engine(hidden_states=x) # 2. Allocate resources based on marginal utility decisions, diagnostics = self.allocator( z_t, budget=budget, preset=preset, force_classical=force_classical or (not self.use_quantum), ) chosen_rank = decisions["rank"] depth_mode = decisions["depth_mode"] is_quantum_token = decisions["is_quantum_token"] # (B, T) kv_precision = decisions["kv_precision"] stats.update(diagnostics) # Handle depth skipping if depth_mode == "skip": # Zero computation: skip layer via identity residual return x, stats # Dynamic KV Cache precision adjustment if kv_cache is not None: kv_cache.set_precision(kv_precision) # 3. Attention Sublayer normed_x = self.ln1(x) # Determine whether any token in sequence requests quantum kernel has_quantum_attn = (decisions["attn_mode_idx"] == 2).any().item() attn_out, attn_weights = self.attention( normed_x, mask=mask, kv_cache=kv_cache, use_quantum=has_quantum_attn and not force_classical, return_attn_weights=True, ) x = x + self.dropout(attn_out) # Early exit check based on epistemic uncertainty if self.enable_early_exit and self.early_exit_head is not None: mean_uncertainty = z_t[..., 2].mean().item() if mean_uncertainty < self.early_exit_threshold: stats["early_exit_triggered"] = True stats["early_exit_logits"] = self.early_exit_head(self.ln2(x)) return x, stats # If partial depth mode: skip FFN if depth_mode == "partial": return x, stats # 4. FFN Sublayer with Nested TT rank slicing self.tt_ffn.set_rank(chosen_rank) normed_ffn_in = self.ln2(x) token_ranks = decisions.get("token_ranks", None) ffn_out = self.tt_ffn(normed_ffn_in, token_ranks=token_ranks) # 5. Quantum feature enhancement on selected hard tokens if is_quantum_token.any() and not force_classical: q_out, q_meta = self.quantum_backend(normed_ffn_in) mask_expanded = is_quantum_token.unsqueeze(-1).float() # (B, T, 1) ffn_out = ffn_out + mask_expanded * q_out x = x + self.dropout(ffn_out) # Update memory traffic stats ffn_traffic = self.tt_ffn.get_memory_traffic() stats["memory_traffic_bytes"] = ffn_traffic["total_bytes"] return x, stats def set_rank(self, rank: int): """Manual rank override.""" self.tt_ffn.set_rank(rank) def reset_scheduler(self): """Reset allocator stability state.""" self.allocator.reset_stability_counters() @property def total_params(self) -> int: return sum(p.numel() for p in self.parameters()) def flops_estimate(self, batch_size: int = 1, seq_len: int = 32) -> Dict[str, int]: attn_flops = self.attention.flops(batch_size, seq_len)["total"] ffn_flops = self.tt_ffn.flops(batch_size, seq_len) return { "attention": attn_flops, "tt_ffn": ffn_flops, "total": attn_flops + ffn_flops, }