Spaces:
Sleeping
Sleeping
Premchandyadav369
feat: deploy interactive Q-TensorFormer demo with trained weights and Gradio 4 UI
67ded87 | """ | |
| 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() | |
| 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, | |
| } | |