""" Information-Value Resource Allocator Module for Q-TensorFormer. The Intellectual Core: Instead of asking "How difficult is this token?", the allocator asks: "Given token information state z_t, hardware device state H_device, and current budgets, what is the cheapest additional computation that produces the greatest expected marginal improvement?" Action Space: - TT Rank: 1, 2, 4, 8 - Attention Pathway: classical_fast (SDPA), classical_standard, quantum_qksam - Computation Depth: skip, partial, full - KV Cache Precision: FP16, INT8, INT4, evict Includes: - MarginalValueModel: empirical neural predictor of Delta Q and Delta Costs - InformationValueAllocator: dimensionally consistent marginal utility routing - Hysteresis & Anti-Chattering stabilization - Routing Churn & Stability tracking - PIDDualSubgradientController: closed-loop empirical SLA tracking with diagnostics """ import torch import torch.nn as nn import torch.nn.functional as F import math from typing import Dict, Optional, Tuple, List, NamedTuple, Any, Union from dataclasses import dataclass, field class AllocatorAction(NamedTuple): rank: int # 1, 2, 4, 8 attention_mode: str # "classical_fast", "classical_standard", "quantum_qksam" depth_mode: str # "skip", "partial", "full" kv_precision: str # "fp16", "int8", "int4", "evict" kv_residency: str = "hot_gpu" # "hot_gpu", "warm_cpu", "cold_evicted" @dataclass class AllocationBudget: max_latency_ms: Optional[float] = None max_memory_mb: Optional[float] = None max_peak_memory_mb: Optional[float] = None max_energy_uj: Optional[float] = None max_energy_per_token_j: Optional[float] = None max_kv_mb: Optional[float] = None max_ttft_ms: Optional[float] = None max_tpot_ms: Optional[float] = None max_bandwidth_gb_s: Optional[float] = None max_cost_usd_1m: Optional[float] = None min_quality_target: float = 0.0 min_quality_fidelity: float = 0.90 risk_tolerance: float = 0.5 phase: str = "decode" # "prefill" or "decode" workload_type: str = "general" # "reasoning", "math", "code", "dialogue", "long_context" lambda_latency: float = 1.0 lambda_memory: float = 0.5 lambda_energy: float = 0.2 lambda_bandwidth: float = 0.3 lambda_cost: float = 0.1 class MarginalValueModel(nn.Module): """ Learned Marginal-Value Predictor for Q-TensorFormer. Predicts expected quality gain (Delta Q) and hardware resource increments (Delta latency, Delta memory, Delta energy, Delta bandwidth) for candidate actions conditioned on token information state z_t, hardware state, and budget state. Mathematical Formulation: Delta Q_hat, Delta C_hat = f_theta(z_t, a, H_device, B_state) """ CANDIDATE_RANKS = [1, 2, 4, 8] ATTENTION_MODES = ["classical_fast", "classical_standard", "quantum_qksam"] DEPTH_MODES = ["skip", "partial", "full"] KV_MODES = ["fp16", "int8", "int4"] def __init__(self, info_dim: int = 8, hidden_dim: int = 64): super().__init__() self.info_dim = info_dim self.hidden_dim = hidden_dim # Action encoding: rank (4), attention (3), depth (3), kv (3) = 13 dims self.action_dim = 4 + 3 + 3 + 3 self.hw_dim = 3 # latency_pressure, memory_pressure, bandwidth_pressure self.budget_dim = 4 # lambda_l, lambda_m, lambda_e, lambda_b in_dim = info_dim + self.action_dim + self.hw_dim + self.budget_dim self.backbone = nn.Sequential( nn.Linear(in_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim), nn.SiLU(), ) # 7 output heads: Quality, Latency, Memory, Energy, Bandwidth, Financial Cost, Epistemic Uncertainty self.head_quality = nn.Linear(hidden_dim, 1) # Delta Q in [0, 1] self.head_latency = nn.Linear(hidden_dim, 1) # Delta Latency in ms >= 0 self.head_memory = nn.Linear(hidden_dim, 1) # Delta Memory in MB >= 0 self.head_energy = nn.Linear(hidden_dim, 1) # Delta Energy in uJ >= 0 self.head_bandwidth = nn.Linear(hidden_dim, 1) # Delta Bandwidth in Bytes >= 0 self.head_cost = nn.Linear(hidden_dim, 1) # Delta Cost in $/1M tokens >= 0 self.head_uncertainty = nn.Linear(hidden_dim, 1) # Epistemic Uncertainty sigma in [0, 1] def encode_action( self, rank_idx: int, attn_idx: int, depth_idx: int, kv_idx: int, device: torch.device, ) -> torch.Tensor: """One-hot encodes the action components into a 13-dim vector.""" vec = torch.zeros(self.action_dim, device=device) vec[rank_idx] = 1.0 vec[4 + attn_idx] = 1.0 vec[7 + depth_idx] = 1.0 vec[10 + kv_idx] = 1.0 return vec def forward( self, z_t: torch.Tensor, action_enc: torch.Tensor, hw_state: torch.Tensor, budget_state: torch.Tensor, ) -> Dict[str, torch.Tensor]: """ Forward pass predicting marginal outcomes. Args: z_t: (B, T, 8) or (N, 8) action_enc: (B, T, 13) or (N, 13) hw_state: (B, T, 3) or (N, 3) budget_state: (B, T, 4) or (N, 4) Returns: Dict containing predicted delta_q, delta_latency_ms, delta_memory_mb, delta_energy_uj, delta_bandwidth_bytes, delta_cost_usd, uncertainty """ x = torch.cat([z_t, action_enc, hw_state, budget_state], dim=-1) h = self.backbone(x) dq = torch.sigmoid(self.head_quality(h)).squeeze(-1) dlat = F.softplus(self.head_latency(h)).squeeze(-1) * 5.0 dmem = F.softplus(self.head_memory(h)).squeeze(-1) * 2.0 dnrg = F.softplus(self.head_energy(h)).squeeze(-1) * 5000.0 dbw = F.softplus(self.head_bandwidth(h)).squeeze(-1) * 30000.0 dcost = F.softplus(self.head_cost(h)).squeeze(-1) * 2.50 dunc = torch.sigmoid(self.head_uncertainty(h)).squeeze(-1) return { "delta_q": dq, "delta_latency_ms": dlat, "delta_memory_mb": dmem, "delta_energy_uj": dnrg, "delta_bandwidth_bytes": dbw, "delta_cost_usd": dcost, "uncertainty": dunc, } def compute_marginal_utility( self, z_t: torch.Tensor, rank_idx: int, attn_idx: int, depth_idx: int, kv_idx: int, hw_state: torch.Tensor, budget_state: torch.Tensor, eps: float = 1e-4, ) -> torch.Tensor: """ Compute dimensionally consistent marginal utility: Value(a | z_t) = Delta Q / (Delta C_eff + eps) """ device = z_t.device act_enc = self.encode_action(rank_idx, attn_idx, depth_idx, kv_idx, device) # Expand act_enc, hw_state, budget_state to match z_t shape shape = z_t.shape[:-1] act_enc_expanded = act_enc.reshape(*([1] * len(shape)), self.action_dim).expand(*shape, self.action_dim) hw_expanded = hw_state.reshape(*([1] * len(shape)), self.hw_dim).expand(*shape, self.hw_dim) b_expanded = budget_state.reshape(*([1] * len(shape)), self.budget_dim).expand(*shape, self.budget_dim) preds = self.forward(z_t, act_enc_expanded, hw_expanded, b_expanded) # Dimensionless normalized cost combination # Normalize: latency (ms / 10ms), memory (MB / 2MB), energy (uJ / 20000uJ), bandwidth (Bytes / 65600B) norm_lat = preds["delta_latency_ms"] / 10.0 norm_mem = preds["delta_memory_mb"] / 2.0 norm_nrg = preds["delta_energy_uj"] / 20000.0 norm_bw = preds["delta_bandwidth_bytes"] / 65600.0 lambda_l = b_expanded[..., 0] lambda_m = b_expanded[..., 1] lambda_e = b_expanded[..., 2] lambda_b = b_expanded[..., 3] eff_cost = norm_lat * (1.0 + lambda_l) + norm_mem * lambda_m + norm_nrg * lambda_e + norm_bw * lambda_b utility = preds["delta_q"] / (eff_cost + eps) return utility class InformationValueAllocator(nn.Module): """ Closed-loop resource allocation controller with marginal utility modeling. """ CANDIDATE_RANKS = [1, 2, 4, 8] ATTENTION_MODES = ["classical_fast", "classical_standard", "quantum_qksam"] DEPTH_MODES = ["skip", "partial", "full"] KV_MODES = ["fp16", "int8", "int4"] def __init__( self, info_dim: int = 8, hidden_dim: int = 32, hysteresis_tau: float = 0.15, default_preset: str = "balanced", marginal_value_model: Optional[MarginalValueModel] = None, ): super().__init__() self.info_dim = info_dim self.hysteresis_tau = hysteresis_tau self.default_preset = default_preset self.marginal_value_model = marginal_value_model # Learned quality gain estimator: predicts Delta Q(a | z_t) for each candidate rank self.quality_rank_net = nn.Sequential( nn.Linear(info_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, len(self.CANDIDATE_RANKS)), ) # Learned gate for attention mode: [fast, standard, quantum] self.quality_attn_net = nn.Sequential( nn.Linear(info_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, len(self.ATTENTION_MODES)), ) # Learned gate for depth: [skip, partial, full] self.quality_depth_net = nn.Sequential( nn.Linear(info_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, len(self.DEPTH_MODES)), ) # Learned gate for KV precision: [fp16, int8, int4] self.quality_kv_net = nn.Sequential( nn.Linear(info_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, len(self.KV_MODES)), ) # Inductive bias calibration: # High entropy (z[1]) and uncertainty (z[2]) scale quality gains for ranks 4 and 8 with torch.no_grad(): self.quality_rank_net[0].weight.data.normal_(0, 0.05) self.quality_rank_net[0].weight.data[0, 1] += 2.0 # H_t self.quality_rank_net[0].weight.data[1, 2] += 2.5 # U_t self.quality_rank_net[2].weight.data.normal_(0, 0.05) self.quality_rank_net[2].weight.data[2, 0] += 1.8 # rank 4 self.quality_rank_net[2].weight.data[3, 0] += 2.8 # rank 8 self.quality_rank_net[2].weight.data[3, 1] += 2.2 # rank 8 on uncertainty self.quality_rank_net[2].bias.data.copy_(torch.tensor([-0.2, 0.2, 0.6, 1.1])) # Quantum attention boosted when uncertainty is high self.quality_attn_net[0].weight.data[0, 2] += 3.0 self.quality_attn_net[2].weight.data[2, 0] += 2.5 # Base depth preference: [skip, partial, full] self.quality_depth_net[2].bias.data.copy_(torch.tensor([-2.0, 0.0, 2.0])) # Anti-chattering / Hysteresis state tracking self.register_buffer("prev_rank_idx", torch.tensor(2, dtype=torch.long)) # default rank 4 (idx 2) self.register_buffer("prev_attn_idx", torch.tensor(0, dtype=torch.long)) # default classical_fast self.register_buffer("total_decisions", torch.tensor(0, dtype=torch.long)) self.register_buffer("churn_count", torch.tensor(0, dtype=torch.long)) # Base nominal costs for actions (empirically normalized relative units) self.cost_ranks = [0.15, 0.30, 0.60, 1.00] # ranks 1, 2, 4, 8 self.cost_attn = [0.20, 0.50, 1.80] # fast, standard, quantum self.cost_depth = [0.05, 0.40, 1.00] # skip, partial, full self.cost_kv = [1.00, 0.50, 0.25] # fp16, int8, int4 def reset_stability_counters(self): """Reset churn and total decision counters.""" self.total_decisions.zero_() self.churn_count.zero_() self.prev_rank_idx.fill_(2) self.prev_attn_idx.fill_(0) @property def routing_churn_rate(self) -> float: """Percentage of token steps where routing changed between successive steps.""" tot = max(1, self.total_decisions.item()) return self.churn_count.item() / tot def forward( self, z_t: torch.Tensor, budget: Optional[AllocationBudget] = None, preset: Optional[str] = None, force_classical: bool = False, ) -> Tuple[Dict[str, torch.Tensor], Dict[str, float]]: """ Evaluate marginal utility and select optimal action per token. """ B, T, _ = z_t.shape device = z_t.device mode = (preset or self.default_preset).lower() # Extract weights from budget b = budget or AllocationBudget() lambda_l = b.lambda_latency lambda_m = b.lambda_memory lambda_e = b.lambda_energy lambda_b = b.lambda_bandwidth # Adjust lambda weights based on deployment preset if mode == "latency": lambda_l *= 2.5 elif mode == "memory": lambda_m *= 3.0 elif mode == "energy": lambda_e *= 3.0 elif mode == "edge": lambda_l *= 2.0 lambda_m *= 2.5 lambda_e *= 2.5 force_classical = True elif mode == "classical_only": force_classical = True elif mode == "full": lambda_l *= 0.2 lambda_m *= 0.2 lambda_e *= 0.2 # 1. Rank Selection: Value(r | z_t) = Delta Q_r / (Cost_r * (1 + lambda_l * L + lambda_m * M) + eps) rank_logits = self.quality_rank_net(z_t) # (B, T, 4) est_dq_rank = torch.sigmoid(rank_logits) # Resource pressure: z_t[..., 5]=L, z_t[..., 6]=M, z_t[..., 7]=B L_pressure = z_t[..., 5].unsqueeze(-1) M_pressure = z_t[..., 6].unsqueeze(-1) B_pressure = z_t[..., 7].unsqueeze(-1) E_pressure = (L_pressure + M_pressure) / 2.0 rank_costs = torch.tensor(self.cost_ranks, device=device).reshape(1, 1, 4) cost_multiplier = 0.40 * ( 1.0 + lambda_l * L_pressure + lambda_m * M_pressure + lambda_b * B_pressure + lambda_e * E_pressure ) effective_rank_cost = cost_multiplier * rank_costs utility_rank = est_dq_rank - effective_rank_cost # (B, T, 4) Lagrangian dual objective # Per-token best rank indices token_best_rank_idx = torch.argmax(utility_rank, dim=-1) # (B, T) # Hysteresis stabilization on sequence level mean_utility_rank = utility_rank.mean(dim=(0, 1)) # (4,) best_rank_idx = int(torch.argmax(mean_utility_rank).item()) prev_idx = self.prev_rank_idx.item() delta_u = mean_utility_rank[best_rank_idx] - mean_utility_rank[prev_idx] if delta_u < self.hysteresis_tau: chosen_rank_idx = prev_idx else: chosen_rank_idx = best_rank_idx if self.training or not torch.is_grad_enabled(): if chosen_rank_idx != prev_idx: self.churn_count.add_(1) self.prev_rank_idx.fill_(chosen_rank_idx) self.total_decisions.add_(1) chosen_rank = self.CANDIDATE_RANKS[chosen_rank_idx] # 2. Attention Pathway Selection (Token-level granularity) attn_logits = self.quality_attn_net(z_t) # (B, T, 3) est_dq_attn = torch.sigmoid(attn_logits) attn_costs = torch.tensor(self.cost_attn, device=device).reshape(1, 1, 3) cost_multiplier_attn = 0.35 * (1.0 + lambda_l * L_pressure + lambda_e * E_pressure) utility_attn = est_dq_attn - cost_multiplier_attn * attn_costs # (B, T, 3) if force_classical: utility_attn[..., 2] = -float("inf") chosen_attn_idx = torch.argmax(utility_attn, dim=-1) # (B, T) # 3. Depth Execution (Layer-level or sequence-level) depth_logits = self.quality_depth_net(z_t) # (B, T, 3) depth_scores = torch.softmax(depth_logits, dim=-1) # (B, T, 3) avg_depth = depth_scores.mean(dim=(0, 1)) chosen_depth_idx = int(torch.argmax(avg_depth).item()) chosen_depth = self.DEPTH_MODES[chosen_depth_idx] # 4. KV Cache Policy Selection kv_logits = self.quality_kv_net(z_t) # (B, T, 3) kv_costs = torch.tensor(self.cost_kv, device=device).reshape(1, 1, 3) cost_multiplier_kv = 0.40 * (1.0 + lambda_m * M_pressure * 2.0) utility_kv = torch.sigmoid(kv_logits) - cost_multiplier_kv * kv_costs mean_kv_u = utility_kv.mean(dim=(0, 1)) chosen_kv_idx = int(torch.argmax(mean_kv_u).item()) chosen_kv = self.KV_MODES[chosen_kv_idx] # Calculate diagnostics q_routed_tokens = (chosen_attn_idx == 2).sum().item() total_tokens = B * T q_usage_pct = (q_routed_tokens / max(1, total_tokens)) * 100.0 # Per-token ranks mapped token_ranks = torch.tensor(self.CANDIDATE_RANKS, device=device)[token_best_rank_idx] # (B, T) diagnostics = { "chosen_rank": chosen_rank, "mean_rank": float(token_ranks.float().mean().item()), "quantum_usage_pct": round(q_usage_pct, 2), "chosen_depth": chosen_depth, "chosen_kv_precision": chosen_kv, "routing_churn_rate": round(self.routing_churn_rate, 4), "effective_rank_cost": round(effective_rank_cost.mean().item(), 3), } decisions = { "rank": chosen_rank, "token_ranks": token_ranks, # (B, T) per-token rank "attn_mode_idx": chosen_attn_idx, # (B, T) "depth_mode": chosen_depth, "kv_precision": chosen_kv, "is_quantum_token": (chosen_attn_idx == 2), # (B, T) bool } return decisions, diagnostics class PIDDualSubgradientController: """ Online Closed-Loop Dual Multiplier Controller for Q-TensorFormer. Tunes Lagrange multipliers lambda_k for latency, memory, energy, and bandwidth to empirically track user-specified SLA targets. Empirically Verifiable Control Metrics: - Settling time (t_settle): tokens until error enters +/- 5% tolerance band - Maximum overshoot (M_p): peak violation percentage above target - Steady-state error (e_ss): mean absolute error in the terminal window - Violation rate: percentage of steps where measured > budget """ def __init__( self, target_latency_ms: Optional[float] = None, target_memory_mb: Optional[float] = None, target_energy_uj: Optional[float] = None, target_bandwidth_bytes: Optional[float] = None, kp: float = 0.05, ki: float = 0.01, kd: float = 0.005, lambda_min: float = 0.05, lambda_max: float = 10.0, ): self.target_latency_ms = target_latency_ms self.target_memory_mb = target_memory_mb self.target_energy_uj = target_energy_uj self.target_bandwidth_bytes = target_bandwidth_bytes self.kp = kp self.ki = ki self.kd = kd self.lambda_min = lambda_min self.lambda_max = lambda_max # Current multiplier states self.lambda_latency = 1.0 self.lambda_memory = 0.5 self.lambda_energy = 0.2 self.lambda_bandwidth = 0.3 # Integrals and previous errors self.integral_errors = {"latency": 0.0, "memory": 0.0, "energy": 0.0, "bandwidth": 0.0} self.prev_errors = {"latency": 0.0, "memory": 0.0, "energy": 0.0, "bandwidth": 0.0} self.history: List[Dict[str, float]] = [] def update( self, measured_latency_ms: Optional[float] = None, measured_memory_mb: Optional[float] = None, measured_energy_uj: Optional[float] = None, measured_bandwidth_bytes: Optional[float] = None, ) -> AllocationBudget: """ Update dual multipliers given observed empirical hardware metrics. Returns an updated AllocationBudget. """ err_lat = 0.0 err_mem = 0.0 err_nrg = 0.0 err_bw = 0.0 if self.target_latency_ms is not None and measured_latency_ms is not None: err_lat = measured_latency_ms - self.target_latency_ms self.integral_errors["latency"] = max(-5.0, min(5.0, self.integral_errors["latency"] + err_lat)) deriv = err_lat - self.prev_errors["latency"] self.prev_errors["latency"] = err_lat delta = self.kp * err_lat + self.ki * self.integral_errors["latency"] + self.kd * deriv self.lambda_latency = max(self.lambda_min, min(self.lambda_max, self.lambda_latency + delta)) if self.target_memory_mb is not None and measured_memory_mb is not None: err_mem = measured_memory_mb - self.target_memory_mb self.integral_errors["memory"] = max(-5.0, min(5.0, self.integral_errors["memory"] + err_mem)) deriv = err_mem - self.prev_errors["memory"] self.prev_errors["memory"] = err_mem delta = self.kp * err_mem + self.ki * self.integral_errors["memory"] + self.kd * deriv self.lambda_memory = max(self.lambda_min, min(self.lambda_max, self.lambda_memory + delta)) if self.target_energy_uj is not None and measured_energy_uj is not None: err_nrg = measured_energy_uj - self.target_energy_uj self.integral_errors["energy"] = max(-5.0, min(5.0, self.integral_errors["energy"] + err_nrg)) deriv = err_nrg - self.prev_errors["energy"] self.prev_errors["energy"] = err_nrg delta = self.kp * err_nrg + self.ki * self.integral_errors["energy"] + self.kd * deriv self.lambda_energy = max(self.lambda_min, min(self.lambda_max, self.lambda_energy + delta)) if self.target_bandwidth_bytes is not None and measured_bandwidth_bytes is not None: err_bw = (measured_bandwidth_bytes - self.target_bandwidth_bytes) / 1000.0 self.integral_errors["bandwidth"] = max(-5.0, min(5.0, self.integral_errors["bandwidth"] + err_bw)) deriv = err_bw - self.prev_errors["bandwidth"] self.prev_errors["bandwidth"] = err_bw delta = self.kp * err_bw + self.ki * self.integral_errors["bandwidth"] + self.kd * deriv self.lambda_bandwidth = max(self.lambda_min, min(self.lambda_max, self.lambda_bandwidth + delta)) record = { "step": len(self.history) + 1, "measured_latency_ms": measured_latency_ms or 0.0, "measured_memory_mb": measured_memory_mb or 0.0, "measured_energy_uj": measured_energy_uj or 0.0, "measured_bandwidth_bytes": measured_bandwidth_bytes or 0.0, "error_latency": err_lat, "error_memory": err_mem, "error_energy": err_nrg, "error_bandwidth": err_bw, "lambda_latency": round(self.lambda_latency, 4), "lambda_memory": round(self.lambda_memory, 4), "lambda_energy": round(self.lambda_energy, 4), "lambda_bandwidth": round(self.lambda_bandwidth, 4), } self.history.append(record) return AllocationBudget( max_latency_ms=self.target_latency_ms, max_memory_mb=self.target_memory_mb, max_energy_uj=self.target_energy_uj, lambda_latency=self.lambda_latency, lambda_memory=self.lambda_memory, lambda_energy=self.lambda_energy, lambda_bandwidth=self.lambda_bandwidth, ) def get_diagnostics(self) -> Dict[str, float]: """Compute control metrics over history.""" if not self.history: return {} lat_errors = [h["error_latency"] for h in self.history if self.target_latency_ms is not None] if not lat_errors: return {"total_steps": len(self.history)} violations = sum(1 for e in lat_errors if e > 0) violation_rate = (violations / len(lat_errors)) * 100.0 target = self.target_latency_ms or 1.0 overshoot_pct = max(0.0, max(lat_errors) / target * 100.0) # Steady-state error over final 20% w = max(1, int(len(lat_errors) * 0.2)) ss_error = sum(abs(e) for e in lat_errors[-w:]) / w # Settling time: step index where error remains within +/- 5% of target band = 0.05 * target settling_step = len(lat_errors) for i in range(len(lat_errors)): if all(abs(e) <= band for e in lat_errors[i:]): settling_step = i + 1 break return { "total_steps": len(self.history), "violation_rate_pct": round(violation_rate, 2), "max_overshoot_pct": round(overshoot_pct, 2), "steady_state_error": round(ss_error, 4), "settling_step": settling_step, } def get_budget(self) -> AllocationBudget: return AllocationBudget( max_latency_ms=self.target_latency_ms, max_memory_mb=self.target_memory_mb, max_energy_uj=self.target_energy_uj, lambda_latency=self.lambda_latency, lambda_memory=self.lambda_memory, lambda_energy=self.lambda_energy, lambda_bandwidth=self.lambda_bandwidth, ) class ConstrainedDecisionEngine(nn.Module): """ Unified Closed-Loop Constrained Decision Engine for Q-TensorFormer. Integrates: 1. Information state evaluation (z_t) 2. Multi-objective marginal value & resource cost prediction (Delta Q, Delta L, Delta M, Delta B, Delta E, Delta $) 3. Risk-aware utility penalization with conservative fallback when confidence is low 4. Binding constraint detection (identifies which SLA ceiling is throttling inference) 5. Phase-aware execution modes (Prefill vs Decode) 6. Workload-specific adaptation (Reasoning, Math, Code, Dialogue, Long-Context) 7. Closed-loop PID feedback adaptation of Lagrangian shadow prices """ CANDIDATE_RANKS = [1, 2, 4, 8] ATTENTION_MODES = ["classical_fast", "classical_standard", "quantum_qksam"] DEPTH_MODES = ["skip", "partial", "full"] KV_MODES = ["fp16", "int8", "int4"] KV_RESIDENCIES = ["hot_gpu", "warm_cpu", "cold_evicted"] def __init__( self, info_dim: int = 8, marginal_value_model: Optional[MarginalValueModel] = None, pid_controller: Optional[PIDDualSubgradientController] = None, risk_gamma: float = 0.35, uncertainty_threshold: float = 0.65, default_preset: str = "balanced", ): super().__init__() self.info_dim = info_dim self.marginal_value_model = marginal_value_model or MarginalValueModel(info_dim=info_dim) self.pid_controller = pid_controller or PIDDualSubgradientController() self.risk_gamma = risk_gamma self.uncertainty_threshold = uncertainty_threshold self.default_preset = default_preset # Historical state tracking self.last_action: Optional[AllocatorAction] = None self.step_count = 0 self.binding_constraints_history: List[str] = [] def evaluate_candidates( self, z_t: torch.Tensor, hw_state: Optional[torch.Tensor] = None, budget: Optional[AllocationBudget] = None, ) -> Tuple[AllocatorAction, Dict[str, Any]]: """ Evaluates candidate actions under multi-budget constraints and risk penalties. Returns selected AllocatorAction and detailed diagnostics. """ device = z_t.device budget = budget or self.pid_controller.get_budget() if hw_state is None: hw_state = torch.tensor([0.2, 0.3, 0.25], device=device) budget_vec = torch.tensor([ budget.lambda_latency, budget.lambda_memory, budget.lambda_energy, budget.lambda_bandwidth, ], device=device) # Build candidate action space phase = getattr(budget, "phase", "decode").lower() candidates = [] for r_idx, r in enumerate(self.CANDIDATE_RANKS): for a_idx, attn in enumerate(self.ATTENTION_MODES): for d_idx, depth in enumerate(self.DEPTH_MODES): for k_idx, kv in enumerate(self.KV_MODES): # In prefill phase, do not skip entire layer to preserve context representation if phase == "prefill" and depth == "skip": continue candidates.append((r_idx, a_idx, d_idx, k_idx, AllocatorAction(r, attn, depth, kv, "hot_gpu"))) best_action = None best_score = -float("inf") diagnostics: Dict[str, Any] = {} # Evaluate candidate utility for r_idx, a_idx, d_idx, k_idx, action in candidates: act_enc = self.marginal_value_model.encode_action(r_idx, a_idx, d_idx, k_idx, device) shape = z_t.shape[:-1] act_expanded = act_enc.reshape(*([1] * len(shape)), -1).expand(*shape, -1) hw_expanded = hw_state.reshape(*([1] * len(shape)), -1).expand(*shape, -1) b_expanded = budget_vec.reshape(*([1] * len(shape)), -1).expand(*shape, -1) preds = self.marginal_value_model(z_t, act_expanded, hw_expanded, b_expanded) dq = preds["delta_q"].mean().item() dlat = preds["delta_latency_ms"].mean().item() dmem = preds["delta_memory_mb"].mean().item() dnrg = preds["delta_energy_uj"].mean().item() dbw = preds["delta_bandwidth_bytes"].mean().item() dcost = preds["delta_cost_usd"].mean().item() dunc = preds["uncertainty"].mean().item() # Workload-specific inductive scaling workload = getattr(budget, "workload_type", "general").lower() if workload in ["reasoning", "math"]: if action.rank >= 4 and action.depth_mode == "full": dq *= 1.25 elif workload == "code": if action.rank >= 4: dq *= 1.15 elif workload == "dialogue": if action.rank <= 2 and action.kv_precision == "int4": dlat *= 0.85 # Dimensionless effective constraint penalty c_eff = ( (dlat / 10.0) * budget.lambda_latency + (dmem / 2.0) * budget.lambda_memory + (dnrg / 20000.0) * budget.lambda_energy + (dbw / 65600.0) * budget.lambda_bandwidth + (dcost / 2.50) * getattr(budget, "lambda_cost", 0.1) ) # Risk-penalized Lagrangian dual objective risk_penalty = self.risk_gamma * dunc score = dq - c_eff - risk_penalty # Check hard SLA limits if specified violates_budget = False if budget.max_tpot_ms is not None and dlat > budget.max_tpot_ms: violates_budget = True if budget.max_memory_mb is not None and dmem > budget.max_memory_mb: violates_budget = True if budget.max_energy_uj is not None and dnrg > budget.max_energy_uj: violates_budget = True if not violates_budget and score > best_score: best_score = score best_action = action diagnostics = { "score": round(score, 4), "expected_delta_q": round(dq, 4), "expected_latency_ms": round(dlat, 3), "expected_memory_mb": round(dmem, 3), "expected_energy_uj": round(dnrg, 2), "expected_bandwidth_bytes": int(dbw), "expected_cost_usd": round(dcost, 4), "uncertainty": round(dunc, 4), "is_fallback": False, } # Safe fallback if high uncertainty or no feasible candidate found if best_action is None or diagnostics.get("uncertainty", 0.0) > self.uncertainty_threshold: best_action = AllocatorAction(rank=4, attention_mode="classical_standard", depth_mode="full", kv_precision="int8", kv_residency="hot_gpu") diagnostics["is_fallback"] = True diagnostics["fallback_reason"] = "uncertainty_exceeded" if best_action else "no_feasible_budget_candidate" # Detect binding constraint binding = "none" if budget.max_tpot_ms and diagnostics.get("expected_latency_ms", 0) >= 0.85 * budget.max_tpot_ms: binding = "latency" elif budget.max_memory_mb and diagnostics.get("expected_memory_mb", 0) >= 0.85 * budget.max_memory_mb: binding = "memory" elif budget.max_energy_uj and diagnostics.get("expected_energy_uj", 0) >= 0.85 * budget.max_energy_uj: binding = "energy" diagnostics["binding_constraint"] = binding self.binding_constraints_history.append(binding) self.last_action = best_action self.step_count += 1 return best_action, diagnostics def step_feedback( self, measured_latency_ms: Optional[float] = None, measured_memory_mb: Optional[float] = None, measured_energy_uj: Optional[float] = None, measured_bandwidth_bytes: Optional[float] = None, ) -> AllocationBudget: """Closed-loop feedback update for PID dual subgradient multipliers.""" return self.pid_controller.update( measured_latency_ms=measured_latency_ms, measured_memory_mb=measured_memory_mb, measured_energy_uj=measured_energy_uj, measured_bandwidth_bytes=measured_bandwidth_bytes, ) def forward( self, z_t: torch.Tensor, budget: Optional[AllocationBudget] = None, hw_state: Optional[torch.Tensor] = None, ) -> Tuple[AllocatorAction, Dict[str, Any]]: """Forward pass delegating to evaluate_candidates for PyTorch module compliance.""" return self.evaluate_candidates(z_t, hw_state=hw_state, budget=budget)