""" Q-TensorFormer: Complete Model Architecture. Unified Information-to-Resource Allocation Architecture. Integrates: - Nested TT-FFN (ranks 1, 2, 4, 8 with instantaneous slicing) - Token Information State (8D feature tracking) - Information-Value Resource Allocator (marginal utility optimization) - Hardware-Aware Cost Profiling - Adaptive KV Cache with multi-precision & eviction - Deployment Presets: QTF_FULL, QTF_BALANCED, QTF_LATENCY, QTF_MEMORY, QTF_ENERGY, QTF_EDGE, QTF_CLASSICAL_ONLY """ import torch import torch.nn as nn import torch.nn.functional as F import math from typing import Optional, Dict, List, Tuple, Union from .blocks import HybridBlock from .config import ModelConfig, validate_model_config from .kv_cache import AdaptiveKVCache from .resource_allocator import AllocationBudget from .hardware_cost_model import HardwareCostModel class PositionalEncoding(nn.Module): def __init__(self, d_model: int, max_len: int = 512, dropout: float = 0.1): super().__init__() self.dropout = nn.Dropout(dropout) self.d_model = d_model pe = torch.zeros(max_len, d_model) pos = torch.arange(0, max_len).float().unsqueeze(1) div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(pos * div) pe[:, 1::2] = torch.cos(pos * div) self.register_buffer("pe", pe.unsqueeze(0)) def forward(self, x: torch.Tensor, start_pos: int = 0) -> torch.Tensor: seq_len = x.size(1) end_pos = start_pos + seq_len if end_pos <= self.pe.size(1): return self.dropout(x + self.pe[:, start_pos:end_pos, :]) pos = torch.arange(start_pos, end_pos, device=x.device, dtype=torch.float32).unsqueeze(1) div = torch.exp( torch.arange(0, self.d_model, 2, device=x.device, dtype=torch.float32) * (-math.log(10000.0) / self.d_model) ) pe_dyn = torch.zeros(1, seq_len, self.d_model, device=x.device, dtype=x.dtype) pe_dyn[0, :, 0::2] = torch.sin(pos * div) pe_dyn[0, :, 1::2] = torch.cos(pos * div) return self.dropout(x + pe_dyn) class QTensorFormer(nn.Module): """ Q-TensorFormer: Closed-Loop Information-to-Resource Adaptive Transformer. """ DEPLOYMENT_PRESETS = [ "QTF_FULL", "QTF_BALANCED", "QTF_LATENCY", "QTF_MEMORY", "QTF_ENERGY", "QTF_EDGE", "QTF_CLASSICAL_ONLY", ] def __init__(self, config: ModelConfig, preset: str = "QTF_BALANCED"): super().__init__() validate_model_config(config, self.__class__) self.config = config self.current_preset = preset.upper() self.embedding = nn.Embedding(config.vocab_size, config.d_model) self.pos_encoding = PositionalEncoding(config.d_model, config.max_seq_len, config.dropout) # Transformer blocks backend_type = getattr(config, "backend_type", "classical_surrogate") self.blocks = nn.ModuleList([ HybridBlock( d_model=config.d_model, n_heads=config.n_heads, n_kv_heads=getattr(config, "n_kv_heads", None), ff_multiplier=config.ff_multiplier, tt_rank=config.tt_rank, tt_min_rank=config.tt_min_rank, use_quantum=config.use_quantum, n_qubits=config.n_qubits, backend_type=backend_type, dropout=config.dropout, max_seq_len=config.max_seq_len, hysteresis_tau=getattr(config, "hysteresis_tau", 0.15), default_preset=self.current_preset.replace("QTF_", "").lower(), enable_early_exit=getattr(config, "enable_early_exit", False), early_exit_threshold=getattr(config, "early_exit_threshold", 0.20), vocab_size=config.vocab_size, ) for _ in range(config.n_layers) ]) self.ln_f = nn.LayerNorm(config.d_model) self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) self.lm_head.weight = self.embedding.weight # Weight tying # Hardware cost model self.hardware_model = HardwareCostModel() self._init_weights() def _init_weights(self): for name, p in self.named_parameters(): if "weight" in name and p.dim() >= 2: nn.init.xavier_uniform_(p) elif "bias" in name: nn.init.zeros_(p) def set_preset(self, preset_name: str): """Switch deployment preset.""" preset = preset_name.upper() if preset not in self.DEPLOYMENT_PRESETS: raise ValueError(f"Unknown preset {preset}. Choose from {self.DEPLOYMENT_PRESETS}") self.current_preset = preset def forward( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, kv_caches: Optional[List[AdaptiveKVCache]] = None, budget: Optional[AllocationBudget] = None, return_stats: bool = False, preset: Optional[str] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[Dict]]]: """ Forward pass with dynamic information-to-resource allocation. """ if preset is not None: self.set_preset(preset) B, T = input_ids.shape start_pos = kv_caches[0].seq_len if (kv_caches is not None and kv_caches[0].seq_len > 0) else 0 x = self.embedding(input_ids) x = self.pos_encoding(x, start_pos=start_pos) preset_str = self.current_preset.replace("QTF_", "").lower() force_classical = (self.current_preset in ["QTF_CLASSICAL_ONLY", "QTF_EDGE"] or not self.config.use_quantum) all_stats = [] for i, block in enumerate(self.blocks): layer_kv = kv_caches[i] if kv_caches is not None else None x, stats = block( x, mask=attention_mask, kv_cache=layer_kv, budget=budget, preset=preset_str, force_classical=force_classical, ) all_stats.append(stats) if stats.get("early_exit_triggered", False) and "early_exit_logits" in stats: logits = stats["early_exit_logits"] if return_stats: return logits, all_stats return logits x = self.ln_f(x) logits = self.lm_head(x) if return_stats: return logits, all_stats return logits @torch.no_grad() def generate( self, input_ids: torch.Tensor, max_new_tokens: int = 20, temperature: float = 1.0, top_k: int = 50, use_cache: bool = True, do_sample: bool = True, eos_token_id: Optional[int] = None, pad_token_id: Optional[int] = 0, attention_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Autoregressive generation with KV Cache, context-window boundary handling, deterministic greedy mode (`do_sample=False` or `temperature=0.0`), and EOS/PAD stopping. """ if input_ids is None or input_ids.dim() != 2 or input_ids.size(0) == 0 or input_ids.size(1) == 0: raise ValueError("input_ids must be a non-empty 2D tensor of shape (batch_size, seq_len).") if (input_ids < 0).any() or (input_ids >= self.config.vocab_size).any(): raise ValueError( f"input_ids contains token IDs outside valid vocabulary range [0, {self.config.vocab_size - 1}]." ) if max_new_tokens < 0: raise ValueError(f"max_new_tokens must be >= 0, got {max_new_tokens}.") if max_new_tokens == 0: return input_ids.clone() self.eval() B = input_ids.shape[0] max_ctx = max(1, int(self.config.max_seq_len)) kv_caches = None if use_cache: kv_caches = [ AdaptiveKVCache(max_capacity=max_ctx) for _ in range(self.config.n_layers) ] curr_input = input_ids curr_mask = attention_mask.clone() if attention_mask is not None else None finished = torch.zeros(B, dtype=torch.bool, device=input_ids.device) for _ in range(max_new_tokens): if kv_caches is not None and kv_caches[0].seq_len > 0: step_input = curr_input[:, -1:] # When KV cache is active and may evict to max_ctx, align mask to cached seq_len + 1 if curr_mask is not None: expected_k_len = min(max_ctx, kv_caches[0].seq_len + 1) step_mask = curr_mask[:, -expected_k_len:] else: step_mask = None else: # Prefill or uncached step: enforce context-window boundary max_seq_len step_input = curr_input[:, -max_ctx:] step_mask = curr_mask[:, -max_ctx:] if curr_mask is not None else None logits = self(step_input, attention_mask=step_mask, kv_caches=kv_caches) next_logits = logits[:, -1, :] if not do_sample or temperature <= 0.0: next_token = torch.argmax(next_logits, dim=-1, keepdim=True) else: scaled_logits = next_logits / max(1e-5, float(temperature)) if top_k > 0: v, _ = torch.topk(scaled_logits, min(top_k, scaled_logits.size(-1))) scaled_logits = scaled_logits.masked_fill(scaled_logits < v[:, [-1]], float("-inf")) probs = F.softmax(scaled_logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) if eos_token_id is not None: if pad_token_id is not None: next_token = torch.where( finished.unsqueeze(-1), torch.full_like(next_token, pad_token_id), next_token, ) finished = finished | (next_token.squeeze(-1) == eos_token_id) curr_input = torch.cat([curr_input, next_token], dim=-1) if curr_mask is not None: new_mask_col = (~finished).long().unsqueeze(-1) if eos_token_id is not None else torch.ones( (B, 1), dtype=curr_mask.dtype, device=curr_mask.device ) curr_mask = torch.cat([curr_mask, new_mask_col.to(curr_mask.dtype)], dim=-1) if eos_token_id is not None and finished.all(): break return curr_input def reset_schedulers(self): for block in self.blocks: block.reset_scheduler() @property def total_params(self) -> int: return sum(p.numel() for p in self.parameters()) @property def active_params(self) -> int: tt_active = sum(b.tt_ffn.active_params for b in self.blocks) base = self.total_params - sum(b.tt_ffn.total_params for b in self.blocks) return base + tt_active @property def compression_ratio(self) -> float: dense_per_block = 2 * self.config.d_model * self.config.d_model * self.config.ff_multiplier base = self.total_params - sum(b.tt_ffn.total_params for b in self.blocks) active_tt = sum(b.tt_ffn.active_params for b in self.blocks) return (base + dense_per_block * self.config.n_layers) / max(base + active_tt, 1) def flops_estimate(self, batch_size: int = 1, seq_len: int = 32) -> Dict: total = 0 breakdown = {} for i, block in enumerate(self.blocks): b = block.flops_estimate(batch_size, seq_len) total += b["total"] breakdown[f"block_{i}"] = b return {"total": total, "breakdown": breakdown} class DenseBaseline(nn.Module): """ Standard dense transformer baseline with identical hyperparameters. """ def __init__(self, config: ModelConfig): super().__init__() validate_model_config(config, self.__class__) self.config = config self.embedding = nn.Embedding(config.vocab_size, config.d_model) self.pos_encoding = PositionalEncoding(config.d_model, config.max_seq_len, config.dropout) self.blocks = nn.ModuleList([ nn.ModuleDict({ "ln1": nn.LayerNorm(config.d_model), "attn": nn.MultiheadAttention(config.d_model, config.n_heads, dropout=config.dropout, batch_first=True), "ln2": nn.LayerNorm(config.d_model), "ffn": nn.Sequential( nn.Linear(config.d_model, config.d_model * config.ff_multiplier), nn.GELU(), nn.Linear(config.d_model * config.ff_multiplier, config.d_model), ), "dropout": nn.Dropout(config.dropout), }) for _ in range(config.n_layers) ]) self.ln_f = nn.LayerNorm(config.d_model) self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) self.lm_head.weight = self.embedding.weight self._init_weights() def _init_weights(self): for name, p in self.named_parameters(): if "weight" in name and p.dim() >= 2: nn.init.xavier_uniform_(p) elif "bias" in name: nn.init.zeros_(p) def forward(self, input_ids, attention_mask=None, return_stats=False): B, T = input_ids.shape x = self.embedding(input_ids) x = self.pos_encoding(x) causal_hide = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1) if attention_mask is None: attn_mask = torch.zeros((T, T), device=x.device, dtype=x.dtype).masked_fill(causal_hide, float("-inf")) else: m = attention_mask.to(device=x.device) pad_hide = (~m) if m.dtype == torch.bool else (m == 0) combined_hide = causal_hide.unsqueeze(0) | pad_hide[:, None, :T] row_all_masked = combined_hide.all(dim=-1, keepdim=True) diag = torch.eye(T, device=x.device, dtype=torch.bool).unsqueeze(0) safe_hide = torch.where(row_all_masked, ~diag, combined_hide) attn_mask = torch.zeros((B, T, T), device=x.device, dtype=x.dtype).masked_fill(safe_hide, float("-inf")) attn_mask = attn_mask.repeat_interleave(self.config.n_heads, dim=0) for block in self.blocks: normed = block["ln1"](x) attn_out, _ = block["attn"](normed, normed, normed, attn_mask=attn_mask, need_weights=False) x = x + block["dropout"](attn_out) ffn_out = block["ffn"](block["ln2"](x)) x = x + block["dropout"](ffn_out) x = self.ln_f(x) logits = self.lm_head(x) if return_stats: return logits, [] return logits @property def total_params(self) -> int: return sum(p.numel() for p in self.parameters()) def create_model(config: ModelConfig, model_type: str = "qtensor", preset: str = "QTF_BALANCED"): """Factory helper to build models.""" if model_type in ["qtensor", "hybrid"]: return QTensorFormer(config, preset=preset) elif model_type in ["dense", "baseline"]: return DenseBaseline(config) elif model_type == "tensor_only": cfg = ModelConfig(**config.__dict__) cfg.use_quantum = False return QTensorFormer(cfg, preset="QTF_CLASSICAL_ONLY") else: raise ValueError(f"Unknown model_type: {model_type}")