Premchan369/Q-TensorFormer

1"""2Information-Value Resource Allocator Module for Q-TensorFormer.3 4The Intellectual Core:5  Instead of asking "How difficult is this token?", the allocator asks:6  "Given token information state z_t, hardware device state H_device, and current7   budgets, what is the cheapest additional computation that produces the greatest8   expected marginal improvement?"9 10Action Space:11  - TT Rank: 1, 2, 4, 812  - Attention Pathway: classical_fast (SDPA), classical_standard, quantum_qksam13  - Computation Depth: skip, partial, full14  - KV Cache Precision: FP16, INT8, INT4, evict15 16Includes:17  - MarginalValueModel: empirical neural predictor of Delta Q and Delta Costs18  - InformationValueAllocator: dimensionally consistent marginal utility routing19  - Hysteresis & Anti-Chattering stabilization20  - Routing Churn & Stability tracking21  - PIDDualSubgradientController: closed-loop empirical SLA tracking with diagnostics22"""23 24import torch25import torch.nn as nn26import torch.nn.functional as F27import math28from typing import Dict, Optional, Tuple, List, NamedTuple, Any, Union29from dataclasses import dataclass, field30 31 32class AllocatorAction(NamedTuple):33    rank: int              # 1, 2, 4, 834    attention_mode: str    # "classical_fast", "classical_standard", "quantum_qksam"35    depth_mode: str        # "skip", "partial", "full"36    kv_precision: str      # "fp16", "int8", "int4", "evict"37    kv_residency: str = "hot_gpu"  # "hot_gpu", "warm_cpu", "cold_evicted"38 39 40@dataclass41class AllocationBudget:42    max_latency_ms: Optional[float] = None43    max_memory_mb: Optional[float] = None44    max_peak_memory_mb: Optional[float] = None45    max_energy_uj: Optional[float] = None46    max_energy_per_token_j: Optional[float] = None47    max_kv_mb: Optional[float] = None48    max_ttft_ms: Optional[float] = None49    max_tpot_ms: Optional[float] = None50    max_bandwidth_gb_s: Optional[float] = None51    max_cost_usd_1m: Optional[float] = None52    min_quality_target: float = 0.053    min_quality_fidelity: float = 0.9054    risk_tolerance: float = 0.555    phase: str = "decode"                 # "prefill" or "decode"56    workload_type: str = "general"        # "reasoning", "math", "code", "dialogue", "long_context"57    lambda_latency: float = 1.058    lambda_memory: float = 0.559    lambda_energy: float = 0.260    lambda_bandwidth: float = 0.361    lambda_cost: float = 0.162 63 64class MarginalValueModel(nn.Module):65    """66    Learned Marginal-Value Predictor for Q-TensorFormer.67 68    Predicts expected quality gain (Delta Q) and hardware resource increments69    (Delta latency, Delta memory, Delta energy, Delta bandwidth) for candidate70    actions conditioned on token information state z_t, hardware state, and budget state.71 72    Mathematical Formulation:73      Delta Q_hat, Delta C_hat = f_theta(z_t, a, H_device, B_state)74    """75 76    CANDIDATE_RANKS = [1, 2, 4, 8]77    ATTENTION_MODES = ["classical_fast", "classical_standard", "quantum_qksam"]78    DEPTH_MODES = ["skip", "partial", "full"]79    KV_MODES = ["fp16", "int8", "int4"]80 81    def __init__(self, info_dim: int = 8, hidden_dim: int = 64):82        super().__init__()83        self.info_dim = info_dim84        self.hidden_dim = hidden_dim85 86        # Action encoding: rank (4), attention (3), depth (3), kv (3) = 13 dims87        self.action_dim = 4 + 3 + 3 + 388        self.hw_dim = 3       # latency_pressure, memory_pressure, bandwidth_pressure89        self.budget_dim = 4   # lambda_l, lambda_m, lambda_e, lambda_b90 91        in_dim = info_dim + self.action_dim + self.hw_dim + self.budget_dim92 93        self.backbone = nn.Sequential(94            nn.Linear(in_dim, hidden_dim),95            nn.LayerNorm(hidden_dim),96            nn.SiLU(),97            nn.Linear(hidden_dim, hidden_dim),98            nn.SiLU(),99        )100 101        # 7 output heads: Quality, Latency, Memory, Energy, Bandwidth, Financial Cost, Epistemic Uncertainty102        self.head_quality = nn.Linear(hidden_dim, 1)      # Delta Q in [0, 1]103        self.head_latency = nn.Linear(hidden_dim, 1)      # Delta Latency in ms >= 0104        self.head_memory = nn.Linear(hidden_dim, 1)       # Delta Memory in MB >= 0105        self.head_energy = nn.Linear(hidden_dim, 1)       # Delta Energy in uJ >= 0106        self.head_bandwidth = nn.Linear(hidden_dim, 1)    # Delta Bandwidth in Bytes >= 0107        self.head_cost = nn.Linear(hidden_dim, 1)         # Delta Cost in $/1M tokens >= 0108        self.head_uncertainty = nn.Linear(hidden_dim, 1)  # Epistemic Uncertainty sigma in [0, 1]109 110    def encode_action(111        self,112        rank_idx: int,113        attn_idx: int,114        depth_idx: int,115        kv_idx: int,116        device: torch.device,117    ) -> torch.Tensor:118        """One-hot encodes the action components into a 13-dim vector."""119        vec = torch.zeros(self.action_dim, device=device)120        vec[rank_idx] = 1.0121        vec[4 + attn_idx] = 1.0122        vec[7 + depth_idx] = 1.0123        vec[10 + kv_idx] = 1.0124        return vec125 126    def forward(127        self,128        z_t: torch.Tensor,129        action_enc: torch.Tensor,130        hw_state: torch.Tensor,131        budget_state: torch.Tensor,132    ) -> Dict[str, torch.Tensor]:133        """134        Forward pass predicting marginal outcomes.135 136        Args:137            z_t: (B, T, 8) or (N, 8)138            action_enc: (B, T, 13) or (N, 13)139            hw_state: (B, T, 3) or (N, 3)140            budget_state: (B, T, 4) or (N, 4)141 142        Returns:143            Dict containing predicted delta_q, delta_latency_ms, delta_memory_mb,144            delta_energy_uj, delta_bandwidth_bytes, delta_cost_usd, uncertainty145        """146        x = torch.cat([z_t, action_enc, hw_state, budget_state], dim=-1)147        h = self.backbone(x)148 149        dq = torch.sigmoid(self.head_quality(h)).squeeze(-1)150        dlat = F.softplus(self.head_latency(h)).squeeze(-1) * 5.0151        dmem = F.softplus(self.head_memory(h)).squeeze(-1) * 2.0152        dnrg = F.softplus(self.head_energy(h)).squeeze(-1) * 5000.0153        dbw = F.softplus(self.head_bandwidth(h)).squeeze(-1) * 30000.0154        dcost = F.softplus(self.head_cost(h)).squeeze(-1) * 2.50155        dunc = torch.sigmoid(self.head_uncertainty(h)).squeeze(-1)156 157        return {158            "delta_q": dq,159            "delta_latency_ms": dlat,160            "delta_memory_mb": dmem,161            "delta_energy_uj": dnrg,162            "delta_bandwidth_bytes": dbw,163            "delta_cost_usd": dcost,164            "uncertainty": dunc,165        }166 167    def compute_marginal_utility(168        self,169        z_t: torch.Tensor,170        rank_idx: int,171        attn_idx: int,172        depth_idx: int,173        kv_idx: int,174        hw_state: torch.Tensor,175        budget_state: torch.Tensor,176        eps: float = 1e-4,177    ) -> torch.Tensor:178        """179        Compute dimensionally consistent marginal utility:180          Value(a | z_t) = Delta Q / (Delta C_eff + eps)181        """182        device = z_t.device183        act_enc = self.encode_action(rank_idx, attn_idx, depth_idx, kv_idx, device)184        # Expand act_enc, hw_state, budget_state to match z_t shape185        shape = z_t.shape[:-1]186        act_enc_expanded = act_enc.reshape(*([1] * len(shape)), self.action_dim).expand(*shape, self.action_dim)187        hw_expanded = hw_state.reshape(*([1] * len(shape)), self.hw_dim).expand(*shape, self.hw_dim)188        b_expanded = budget_state.reshape(*([1] * len(shape)), self.budget_dim).expand(*shape, self.budget_dim)189 190        preds = self.forward(z_t, act_enc_expanded, hw_expanded, b_expanded)191 192        # Dimensionless normalized cost combination193        # Normalize: latency (ms / 10ms), memory (MB / 2MB), energy (uJ / 20000uJ), bandwidth (Bytes / 65600B)194        norm_lat = preds["delta_latency_ms"] / 10.0195        norm_mem = preds["delta_memory_mb"] / 2.0196        norm_nrg = preds["delta_energy_uj"] / 20000.0197        norm_bw = preds["delta_bandwidth_bytes"] / 65600.0198 199        lambda_l = b_expanded[..., 0]200        lambda_m = b_expanded[..., 1]201        lambda_e = b_expanded[..., 2]202        lambda_b = b_expanded[..., 3]203 204        eff_cost = norm_lat * (1.0 + lambda_l) + norm_mem * lambda_m + norm_nrg * lambda_e + norm_bw * lambda_b205        utility = preds["delta_q"] / (eff_cost + eps)206 207        return utility208 209 210class InformationValueAllocator(nn.Module):211    """212    Closed-loop resource allocation controller with marginal utility modeling.213    """214 215    CANDIDATE_RANKS = [1, 2, 4, 8]216    ATTENTION_MODES = ["classical_fast", "classical_standard", "quantum_qksam"]217    DEPTH_MODES = ["skip", "partial", "full"]218    KV_MODES = ["fp16", "int8", "int4"]219 220    def __init__(221        self,222        info_dim: int = 8,223        hidden_dim: int = 32,224        hysteresis_tau: float = 0.15,225        default_preset: str = "balanced",226        marginal_value_model: Optional[MarginalValueModel] = None,227    ):228        super().__init__()229        self.info_dim = info_dim230        self.hysteresis_tau = hysteresis_tau231        self.default_preset = default_preset232        self.marginal_value_model = marginal_value_model233 234        # Learned quality gain estimator: predicts Delta Q(a | z_t) for each candidate rank235        self.quality_rank_net = nn.Sequential(236            nn.Linear(info_dim, hidden_dim),237            nn.SiLU(),238            nn.Linear(hidden_dim, len(self.CANDIDATE_RANKS)),239        )240 241        # Learned gate for attention mode: [fast, standard, quantum]242        self.quality_attn_net = nn.Sequential(243            nn.Linear(info_dim, hidden_dim),244            nn.SiLU(),245            nn.Linear(hidden_dim, len(self.ATTENTION_MODES)),246        )247 248        # Learned gate for depth: [skip, partial, full]249        self.quality_depth_net = nn.Sequential(250            nn.Linear(info_dim, hidden_dim),251            nn.SiLU(),252            nn.Linear(hidden_dim, len(self.DEPTH_MODES)),253        )254 255        # Learned gate for KV precision: [fp16, int8, int4]256        self.quality_kv_net = nn.Sequential(257            nn.Linear(info_dim, hidden_dim),258            nn.SiLU(),259            nn.Linear(hidden_dim, len(self.KV_MODES)),260        )261 262        # Inductive bias calibration:263        # High entropy (z[1]) and uncertainty (z[2]) scale quality gains for ranks 4 and 8264        with torch.no_grad():265            self.quality_rank_net[0].weight.data.normal_(0, 0.05)266            self.quality_rank_net[0].weight.data[0, 1] += 2.0  # H_t267            self.quality_rank_net[0].weight.data[1, 2] += 2.5  # U_t268            self.quality_rank_net[2].weight.data.normal_(0, 0.05)269            self.quality_rank_net[2].weight.data[2, 0] += 1.8  # rank 4270            self.quality_rank_net[2].weight.data[3, 0] += 2.8  # rank 8271            self.quality_rank_net[2].weight.data[3, 1] += 2.2  # rank 8 on uncertainty272            self.quality_rank_net[2].bias.data.copy_(torch.tensor([-0.2, 0.2, 0.6, 1.1]))273 274            # Quantum attention boosted when uncertainty is high275            self.quality_attn_net[0].weight.data[0, 2] += 3.0276            self.quality_attn_net[2].weight.data[2, 0] += 2.5277 278            # Base depth preference: [skip, partial, full]279            self.quality_depth_net[2].bias.data.copy_(torch.tensor([-2.0, 0.0, 2.0]))280 281        # Anti-chattering / Hysteresis state tracking282        self.register_buffer("prev_rank_idx", torch.tensor(2, dtype=torch.long))  # default rank 4 (idx 2)283        self.register_buffer("prev_attn_idx", torch.tensor(0, dtype=torch.long))  # default classical_fast284        self.register_buffer("total_decisions", torch.tensor(0, dtype=torch.long))285        self.register_buffer("churn_count", torch.tensor(0, dtype=torch.long))286 287        # Base nominal costs for actions (empirically normalized relative units)288        self.cost_ranks = [0.15, 0.30, 0.60, 1.00]      # ranks 1, 2, 4, 8289        self.cost_attn = [0.20, 0.50, 1.80]             # fast, standard, quantum290        self.cost_depth = [0.05, 0.40, 1.00]            # skip, partial, full291        self.cost_kv = [1.00, 0.50, 0.25]               # fp16, int8, int4292 293    def reset_stability_counters(self):294        """Reset churn and total decision counters."""295        self.total_decisions.zero_()296        self.churn_count.zero_()297        self.prev_rank_idx.fill_(2)298        self.prev_attn_idx.fill_(0)299 300    @property301    def routing_churn_rate(self) -> float:302        """Percentage of token steps where routing changed between successive steps."""303        tot = max(1, self.total_decisions.item())304        return self.churn_count.item() / tot305 306    def forward(307        self,308        z_t: torch.Tensor,309        budget: Optional[AllocationBudget] = None,310        preset: Optional[str] = None,311        force_classical: bool = False,312    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, float]]:313        """314        Evaluate marginal utility and select optimal action per token.315        """316        B, T, _ = z_t.shape317        device = z_t.device318        mode = (preset or self.default_preset).lower()319 320        # Extract weights from budget321        b = budget or AllocationBudget()322        lambda_l = b.lambda_latency323        lambda_m = b.lambda_memory324        lambda_e = b.lambda_energy325        lambda_b = b.lambda_bandwidth326 327        # Adjust lambda weights based on deployment preset328        if mode == "latency":329            lambda_l *= 2.5330        elif mode == "memory":331            lambda_m *= 3.0332        elif mode == "energy":333            lambda_e *= 3.0334        elif mode == "edge":335            lambda_l *= 2.0336            lambda_m *= 2.5337            lambda_e *= 2.5338            force_classical = True339        elif mode == "classical_only":340            force_classical = True341        elif mode == "full":342            lambda_l *= 0.2343            lambda_m *= 0.2344            lambda_e *= 0.2345 346        # 1. Rank Selection: Value(r | z_t) = Delta Q_r / (Cost_r * (1 + lambda_l * L + lambda_m * M) + eps)347        rank_logits = self.quality_rank_net(z_t)  # (B, T, 4)348        est_dq_rank = torch.sigmoid(rank_logits)349 350        # Resource pressure: z_t[..., 5]=L, z_t[..., 6]=M, z_t[..., 7]=B351        L_pressure = z_t[..., 5].unsqueeze(-1)352        M_pressure = z_t[..., 6].unsqueeze(-1)353        B_pressure = z_t[..., 7].unsqueeze(-1)354        E_pressure = (L_pressure + M_pressure) / 2.0355 356        rank_costs = torch.tensor(self.cost_ranks, device=device).reshape(1, 1, 4)357        cost_multiplier = 0.40 * (358            1.0 + lambda_l * L_pressure + lambda_m * M_pressure + lambda_b * B_pressure + lambda_e * E_pressure359        )360        effective_rank_cost = cost_multiplier * rank_costs361        utility_rank = est_dq_rank - effective_rank_cost  # (B, T, 4) Lagrangian dual objective362 363        # Per-token best rank indices364        token_best_rank_idx = torch.argmax(utility_rank, dim=-1)  # (B, T)365 366        # Hysteresis stabilization on sequence level367        mean_utility_rank = utility_rank.mean(dim=(0, 1))  # (4,)368        best_rank_idx = int(torch.argmax(mean_utility_rank).item())369        prev_idx = self.prev_rank_idx.item()370 371        delta_u = mean_utility_rank[best_rank_idx] - mean_utility_rank[prev_idx]372        if delta_u < self.hysteresis_tau:373            chosen_rank_idx = prev_idx374        else:375            chosen_rank_idx = best_rank_idx376            if self.training or not torch.is_grad_enabled():377                if chosen_rank_idx != prev_idx:378                    self.churn_count.add_(1)379            self.prev_rank_idx.fill_(chosen_rank_idx)380 381        self.total_decisions.add_(1)382        chosen_rank = self.CANDIDATE_RANKS[chosen_rank_idx]383 384        # 2. Attention Pathway Selection (Token-level granularity)385        attn_logits = self.quality_attn_net(z_t)  # (B, T, 3)386        est_dq_attn = torch.sigmoid(attn_logits)387        attn_costs = torch.tensor(self.cost_attn, device=device).reshape(1, 1, 3)388        cost_multiplier_attn = 0.35 * (1.0 + lambda_l * L_pressure + lambda_e * E_pressure)389        utility_attn = est_dq_attn - cost_multiplier_attn * attn_costs  # (B, T, 3)390 391        if force_classical:392            utility_attn[..., 2] = -float("inf")393 394        chosen_attn_idx = torch.argmax(utility_attn, dim=-1)  # (B, T)395 396        # 3. Depth Execution (Layer-level or sequence-level)397        depth_logits = self.quality_depth_net(z_t)  # (B, T, 3)398        depth_scores = torch.softmax(depth_logits, dim=-1)  # (B, T, 3)399        avg_depth = depth_scores.mean(dim=(0, 1))400        chosen_depth_idx = int(torch.argmax(avg_depth).item())401        chosen_depth = self.DEPTH_MODES[chosen_depth_idx]402 403        # 4. KV Cache Policy Selection404        kv_logits = self.quality_kv_net(z_t)  # (B, T, 3)405        kv_costs = torch.tensor(self.cost_kv, device=device).reshape(1, 1, 3)406        cost_multiplier_kv = 0.40 * (1.0 + lambda_m * M_pressure * 2.0)407        utility_kv = torch.sigmoid(kv_logits) - cost_multiplier_kv * kv_costs408        mean_kv_u = utility_kv.mean(dim=(0, 1))409        chosen_kv_idx = int(torch.argmax(mean_kv_u).item())410        chosen_kv = self.KV_MODES[chosen_kv_idx]411 412        # Calculate diagnostics413        q_routed_tokens = (chosen_attn_idx == 2).sum().item()414        total_tokens = B * T415        q_usage_pct = (q_routed_tokens / max(1, total_tokens)) * 100.0416 417        # Per-token ranks mapped418        token_ranks = torch.tensor(self.CANDIDATE_RANKS, device=device)[token_best_rank_idx]  # (B, T)419 420        diagnostics = {421            "chosen_rank": chosen_rank,422            "mean_rank": float(token_ranks.float().mean().item()),423            "quantum_usage_pct": round(q_usage_pct, 2),424            "chosen_depth": chosen_depth,425            "chosen_kv_precision": chosen_kv,426            "routing_churn_rate": round(self.routing_churn_rate, 4),427            "effective_rank_cost": round(effective_rank_cost.mean().item(), 3),428        }429 430        decisions = {431            "rank": chosen_rank,432            "token_ranks": token_ranks,                # (B, T) per-token rank433            "attn_mode_idx": chosen_attn_idx,          # (B, T)434            "depth_mode": chosen_depth,435            "kv_precision": chosen_kv,436            "is_quantum_token": (chosen_attn_idx == 2),  # (B, T) bool437        }438 439        return decisions, diagnostics440 441 442class PIDDualSubgradientController:443    """444    Online Closed-Loop Dual Multiplier Controller for Q-TensorFormer.445 446    Tunes Lagrange multipliers lambda_k for latency, memory, energy, and bandwidth447    to empirically track user-specified SLA targets.448 449    Empirically Verifiable Control Metrics:450      - Settling time (t_settle): tokens until error enters +/- 5% tolerance band451      - Maximum overshoot (M_p): peak violation percentage above target452      - Steady-state error (e_ss): mean absolute error in the terminal window453      - Violation rate: percentage of steps where measured > budget454    """455 456    def __init__(457        self,458        target_latency_ms: Optional[float] = None,459        target_memory_mb: Optional[float] = None,460        target_energy_uj: Optional[float] = None,461        target_bandwidth_bytes: Optional[float] = None,462        kp: float = 0.05,463        ki: float = 0.01,464        kd: float = 0.005,465        lambda_min: float = 0.05,466        lambda_max: float = 10.0,467    ):468        self.target_latency_ms = target_latency_ms469        self.target_memory_mb = target_memory_mb470        self.target_energy_uj = target_energy_uj471        self.target_bandwidth_bytes = target_bandwidth_bytes472        self.kp = kp473        self.ki = ki474        self.kd = kd475        self.lambda_min = lambda_min476        self.lambda_max = lambda_max477 478        # Current multiplier states479        self.lambda_latency = 1.0480        self.lambda_memory = 0.5481        self.lambda_energy = 0.2482        self.lambda_bandwidth = 0.3483 484        # Integrals and previous errors485        self.integral_errors = {"latency": 0.0, "memory": 0.0, "energy": 0.0, "bandwidth": 0.0}486        self.prev_errors = {"latency": 0.0, "memory": 0.0, "energy": 0.0, "bandwidth": 0.0}487        self.history: List[Dict[str, float]] = []488 489    def update(490        self,491        measured_latency_ms: Optional[float] = None,492        measured_memory_mb: Optional[float] = None,493        measured_energy_uj: Optional[float] = None,494        measured_bandwidth_bytes: Optional[float] = None,495    ) -> AllocationBudget:496        """497        Update dual multipliers given observed empirical hardware metrics.498        Returns an updated AllocationBudget.499        """500        err_lat = 0.0501        err_mem = 0.0502        err_nrg = 0.0503        err_bw = 0.0504 505        if self.target_latency_ms is not None and measured_latency_ms is not None:506            err_lat = measured_latency_ms - self.target_latency_ms507            self.integral_errors["latency"] = max(-5.0, min(5.0, self.integral_errors["latency"] + err_lat))508            deriv = err_lat - self.prev_errors["latency"]509            self.prev_errors["latency"] = err_lat510            delta = self.kp * err_lat + self.ki * self.integral_errors["latency"] + self.kd * deriv511            self.lambda_latency = max(self.lambda_min, min(self.lambda_max, self.lambda_latency + delta))512 513        if self.target_memory_mb is not None and measured_memory_mb is not None:514            err_mem = measured_memory_mb - self.target_memory_mb515            self.integral_errors["memory"] = max(-5.0, min(5.0, self.integral_errors["memory"] + err_mem))516            deriv = err_mem - self.prev_errors["memory"]517            self.prev_errors["memory"] = err_mem518            delta = self.kp * err_mem + self.ki * self.integral_errors["memory"] + self.kd * deriv519            self.lambda_memory = max(self.lambda_min, min(self.lambda_max, self.lambda_memory + delta))520 521        if self.target_energy_uj is not None and measured_energy_uj is not None:522            err_nrg = measured_energy_uj - self.target_energy_uj523            self.integral_errors["energy"] = max(-5.0, min(5.0, self.integral_errors["energy"] + err_nrg))524            deriv = err_nrg - self.prev_errors["energy"]525            self.prev_errors["energy"] = err_nrg526            delta = self.kp * err_nrg + self.ki * self.integral_errors["energy"] + self.kd * deriv527            self.lambda_energy = max(self.lambda_min, min(self.lambda_max, self.lambda_energy + delta))528 529        if self.target_bandwidth_bytes is not None and measured_bandwidth_bytes is not None:530            err_bw = (measured_bandwidth_bytes - self.target_bandwidth_bytes) / 1000.0531            self.integral_errors["bandwidth"] = max(-5.0, min(5.0, self.integral_errors["bandwidth"] + err_bw))532            deriv = err_bw - self.prev_errors["bandwidth"]533            self.prev_errors["bandwidth"] = err_bw534            delta = self.kp * err_bw + self.ki * self.integral_errors["bandwidth"] + self.kd * deriv535            self.lambda_bandwidth = max(self.lambda_min, min(self.lambda_max, self.lambda_bandwidth + delta))536 537        record = {538            "step": len(self.history) + 1,539            "measured_latency_ms": measured_latency_ms or 0.0,540            "measured_memory_mb": measured_memory_mb or 0.0,541            "measured_energy_uj": measured_energy_uj or 0.0,542            "measured_bandwidth_bytes": measured_bandwidth_bytes or 0.0,543            "error_latency": err_lat,544            "error_memory": err_mem,545            "error_energy": err_nrg,546            "error_bandwidth": err_bw,547            "lambda_latency": round(self.lambda_latency, 4),548            "lambda_memory": round(self.lambda_memory, 4),549            "lambda_energy": round(self.lambda_energy, 4),550            "lambda_bandwidth": round(self.lambda_bandwidth, 4),551        }552        self.history.append(record)553 554        return AllocationBudget(555            max_latency_ms=self.target_latency_ms,556            max_memory_mb=self.target_memory_mb,557            max_energy_uj=self.target_energy_uj,558            lambda_latency=self.lambda_latency,559            lambda_memory=self.lambda_memory,560            lambda_energy=self.lambda_energy,561            lambda_bandwidth=self.lambda_bandwidth,562        )563 564    def get_diagnostics(self) -> Dict[str, float]:565        """Compute control metrics over history."""566        if not self.history:567            return {}568 569        lat_errors = [h["error_latency"] for h in self.history if self.target_latency_ms is not None]570        if not lat_errors:571            return {"total_steps": len(self.history)}572 573        violations = sum(1 for e in lat_errors if e > 0)574        violation_rate = (violations / len(lat_errors)) * 100.0575 576        target = self.target_latency_ms or 1.0577        overshoot_pct = max(0.0, max(lat_errors) / target * 100.0)578 579        # Steady-state error over final 20%580        w = max(1, int(len(lat_errors) * 0.2))581        ss_error = sum(abs(e) for e in lat_errors[-w:]) / w582 583        # Settling time: step index where error remains within +/- 5% of target584        band = 0.05 * target585        settling_step = len(lat_errors)586        for i in range(len(lat_errors)):587            if all(abs(e) <= band for e in lat_errors[i:]):588                settling_step = i + 1589                break590 591        return {592            "total_steps": len(self.history),593            "violation_rate_pct": round(violation_rate, 2),594            "max_overshoot_pct": round(overshoot_pct, 2),595            "steady_state_error": round(ss_error, 4),596            "settling_step": settling_step,597        }598 599    def get_budget(self) -> AllocationBudget:600        return AllocationBudget(601            max_latency_ms=self.target_latency_ms,602            max_memory_mb=self.target_memory_mb,603            max_energy_uj=self.target_energy_uj,604            lambda_latency=self.lambda_latency,605            lambda_memory=self.lambda_memory,606            lambda_energy=self.lambda_energy,607            lambda_bandwidth=self.lambda_bandwidth,608        )609 610 611class ConstrainedDecisionEngine(nn.Module):612    """613    Unified Closed-Loop Constrained Decision Engine for Q-TensorFormer.614 615    Integrates:616      1. Information state evaluation (z_t)617      2. Multi-objective marginal value & resource cost prediction (Delta Q, Delta L, Delta M, Delta B, Delta E, Delta $)618      3. Risk-aware utility penalization with conservative fallback when confidence is low619      4. Binding constraint detection (identifies which SLA ceiling is throttling inference)620      5. Phase-aware execution modes (Prefill vs Decode)621      6. Workload-specific adaptation (Reasoning, Math, Code, Dialogue, Long-Context)622      7. Closed-loop PID feedback adaptation of Lagrangian shadow prices623    """624 625    CANDIDATE_RANKS = [1, 2, 4, 8]626    ATTENTION_MODES = ["classical_fast", "classical_standard", "quantum_qksam"]627    DEPTH_MODES = ["skip", "partial", "full"]628    KV_MODES = ["fp16", "int8", "int4"]629    KV_RESIDENCIES = ["hot_gpu", "warm_cpu", "cold_evicted"]630 631    def __init__(632        self,633        info_dim: int = 8,634        marginal_value_model: Optional[MarginalValueModel] = None,635        pid_controller: Optional[PIDDualSubgradientController] = None,636        risk_gamma: float = 0.35,637        uncertainty_threshold: float = 0.65,638        default_preset: str = "balanced",639    ):640        super().__init__()641        self.info_dim = info_dim642        self.marginal_value_model = marginal_value_model or MarginalValueModel(info_dim=info_dim)643        self.pid_controller = pid_controller or PIDDualSubgradientController()644        self.risk_gamma = risk_gamma645        self.uncertainty_threshold = uncertainty_threshold646        self.default_preset = default_preset647 648        # Historical state tracking649        self.last_action: Optional[AllocatorAction] = None650        self.step_count = 0651        self.binding_constraints_history: List[str] = []652 653    def evaluate_candidates(654        self,655        z_t: torch.Tensor,656        hw_state: Optional[torch.Tensor] = None,657        budget: Optional[AllocationBudget] = None,658    ) -> Tuple[AllocatorAction, Dict[str, Any]]:659        """660        Evaluates candidate actions under multi-budget constraints and risk penalties.661        Returns selected AllocatorAction and detailed diagnostics.662        """663        device = z_t.device664        budget = budget or self.pid_controller.get_budget()665 666        if hw_state is None:667            hw_state = torch.tensor([0.2, 0.3, 0.25], device=device)668 669        budget_vec = torch.tensor([670            budget.lambda_latency,671            budget.lambda_memory,672            budget.lambda_energy,673            budget.lambda_bandwidth,674        ], device=device)675 676        # Build candidate action space677        phase = getattr(budget, "phase", "decode").lower()678        candidates = []679        for r_idx, r in enumerate(self.CANDIDATE_RANKS):680            for a_idx, attn in enumerate(self.ATTENTION_MODES):681                for d_idx, depth in enumerate(self.DEPTH_MODES):682                    for k_idx, kv in enumerate(self.KV_MODES):683                        # In prefill phase, do not skip entire layer to preserve context representation684                        if phase == "prefill" and depth == "skip":685                            continue686                        candidates.append((r_idx, a_idx, d_idx, k_idx, AllocatorAction(r, attn, depth, kv, "hot_gpu")))687 688        best_action = None689        best_score = -float("inf")690        diagnostics: Dict[str, Any] = {}691 692        # Evaluate candidate utility693        for r_idx, a_idx, d_idx, k_idx, action in candidates:694            act_enc = self.marginal_value_model.encode_action(r_idx, a_idx, d_idx, k_idx, device)695            shape = z_t.shape[:-1]696            act_expanded = act_enc.reshape(*([1] * len(shape)), -1).expand(*shape, -1)697            hw_expanded = hw_state.reshape(*([1] * len(shape)), -1).expand(*shape, -1)698            b_expanded = budget_vec.reshape(*([1] * len(shape)), -1).expand(*shape, -1)699 700            preds = self.marginal_value_model(z_t, act_expanded, hw_expanded, b_expanded)701 702            dq = preds["delta_q"].mean().item()703            dlat = preds["delta_latency_ms"].mean().item()704            dmem = preds["delta_memory_mb"].mean().item()705            dnrg = preds["delta_energy_uj"].mean().item()706            dbw = preds["delta_bandwidth_bytes"].mean().item()707            dcost = preds["delta_cost_usd"].mean().item()708            dunc = preds["uncertainty"].mean().item()709 710            # Workload-specific inductive scaling711            workload = getattr(budget, "workload_type", "general").lower()712            if workload in ["reasoning", "math"]:713                if action.rank >= 4 and action.depth_mode == "full":714                    dq *= 1.25715            elif workload == "code":716                if action.rank >= 4:717                    dq *= 1.15718            elif workload == "dialogue":719                if action.rank <= 2 and action.kv_precision == "int4":720                    dlat *= 0.85721 722            # Dimensionless effective constraint penalty723            c_eff = (724                (dlat / 10.0) * budget.lambda_latency +725                (dmem / 2.0) * budget.lambda_memory +726                (dnrg / 20000.0) * budget.lambda_energy +727                (dbw / 65600.0) * budget.lambda_bandwidth +728                (dcost / 2.50) * getattr(budget, "lambda_cost", 0.1)729            )730 731            # Risk-penalized Lagrangian dual objective732            risk_penalty = self.risk_gamma * dunc733            score = dq - c_eff - risk_penalty734 735            # Check hard SLA limits if specified736            violates_budget = False737            if budget.max_tpot_ms is not None and dlat > budget.max_tpot_ms:738                violates_budget = True739            if budget.max_memory_mb is not None and dmem > budget.max_memory_mb:740                violates_budget = True741            if budget.max_energy_uj is not None and dnrg > budget.max_energy_uj:742                violates_budget = True743 744            if not violates_budget and score > best_score:745                best_score = score746                best_action = action747                diagnostics = {748                    "score": round(score, 4),749                    "expected_delta_q": round(dq, 4),750                    "expected_latency_ms": round(dlat, 3),751                    "expected_memory_mb": round(dmem, 3),752                    "expected_energy_uj": round(dnrg, 2),753                    "expected_bandwidth_bytes": int(dbw),754                    "expected_cost_usd": round(dcost, 4),755                    "uncertainty": round(dunc, 4),756                    "is_fallback": False,757                }758 759        # Safe fallback if high uncertainty or no feasible candidate found760        if best_action is None or diagnostics.get("uncertainty", 0.0) > self.uncertainty_threshold:761            best_action = AllocatorAction(rank=4, attention_mode="classical_standard", depth_mode="full", kv_precision="int8", kv_residency="hot_gpu")762            diagnostics["is_fallback"] = True763            diagnostics["fallback_reason"] = "uncertainty_exceeded" if best_action else "no_feasible_budget_candidate"764 765        # Detect binding constraint766        binding = "none"767        if budget.max_tpot_ms and diagnostics.get("expected_latency_ms", 0) >= 0.85 * budget.max_tpot_ms:768            binding = "latency"769        elif budget.max_memory_mb and diagnostics.get("expected_memory_mb", 0) >= 0.85 * budget.max_memory_mb:770            binding = "memory"771        elif budget.max_energy_uj and diagnostics.get("expected_energy_uj", 0) >= 0.85 * budget.max_energy_uj:772            binding = "energy"773        diagnostics["binding_constraint"] = binding774        self.binding_constraints_history.append(binding)775 776        self.last_action = best_action777        self.step_count += 1778        return best_action, diagnostics779 780    def step_feedback(781        self,782        measured_latency_ms: Optional[float] = None,783        measured_memory_mb: Optional[float] = None,784        measured_energy_uj: Optional[float] = None,785        measured_bandwidth_bytes: Optional[float] = None,786    ) -> AllocationBudget:787        """Closed-loop feedback update for PID dual subgradient multipliers."""788        return self.pid_controller.update(789            measured_latency_ms=measured_latency_ms,790            measured_memory_mb=measured_memory_mb,791            measured_energy_uj=measured_energy_uj,792            measured_bandwidth_bytes=measured_bandwidth_bytes,793        )794 795    def forward(796        self,797        z_t: torch.Tensor,798        budget: Optional[AllocationBudget] = None,799        hw_state: Optional[torch.Tensor] = None,800    ) -> Tuple[AllocatorAction, Dict[str, Any]]:801        """Forward pass delegating to evaluate_candidates for PyTorch module compliance."""802        return self.evaluate_candidates(z_t, hw_state=hw_state, budget=budget)803