Text Generation
Transformers
Safetensors
PyTorch
ONNX
English
qtensorformer
q-tensorformer
high-throughput-inference
tensor-train
tensor-networks
tensorrt-llm
vllm
continuous-batching
paged-attention
speculative-decoding
zero-svd
adaptive-computation
dynamic-neural-networks
kv-cache
kv-cache-quantization
attention-sinks
lyapunov-stability
control-theory
contextual-bandits
edge-ai
edge-inference
on-device-ai
model-compression
hopper-h100
blackwell-b200
jetson-orin
jetson-thor
grouped-query-attention
rotary-position-embeddings
custom_code
academic-paper
ieee
neurips
vit-ap-university
deepseek
deepseek-r1
reasoning
chain-of-thought
gsm8k
openai-compatible
fastapi
rest-api
sse-streaming
onnx-runtime
embedded-ai
hardware-telemetry
power-efficiency
green-ai
raspberry-pi
smollm2
qwen2.5
llama-3.2
apple-silicon
Premchandyadav369
feat: complete 6 frontier research extensions, empirical benchmarks, and 101-test suite
b5cb5a0 | """ | |
| src/bandit_allocator.py | |
| Contextual Bandit & Reinforcement Learning Dynamic Resource Allocator for Q-TensorFormer. | |
| Replaces reactive PID loops with an online-learning LinUCB / Contextual Policy that maps | |
| the 8D token information state z_t directly to the optimal Pareto decision tuple | |
| (TT-rank, KV-precision, early-exit). | |
| """ | |
| from typing import Dict, Any, List, Tuple, Optional | |
| import math | |
| import torch | |
| import numpy as np | |
| class ContextualBanditAllocator: | |
| """ | |
| LinUCB (Linear Upper Confidence Bound) Contextual Bandit Allocator. | |
| Observes 8D context vector z_t in [0, 1]^8 and selects from discrete action space: | |
| - Action 0: r=1, INT4 KV, Early Exit Allowed (Ultra-Edge) | |
| - Action 1: r=2, INT4 KV, No Exit (Edge) | |
| - Action 2: r=4, INT8 KV, No Exit (Moderate) | |
| - Action 3: r=8, INT8 KV, No Exit (Balanced) | |
| - Action 4: r=16, FP16 KV, No Exit (High Capacity) | |
| """ | |
| def __init__(self, context_dim: int = 8, n_actions: int = 5, alpha: float = 0.5): | |
| self.context_dim = context_dim | |
| self.n_actions = n_actions | |
| self.alpha = alpha # Exploration parameter | |
| # LinUCB state parameters: A_a in R^{d x d}, b_a in R^{d} | |
| self.A = [np.identity(context_dim) for _ in range(n_actions)] | |
| self.b = [np.zeros((context_dim, 1)) for _ in range(n_actions)] | |
| # Action mapping | |
| self.action_specs = [ | |
| {"rank": 1, "kv_precision": "INT4", "allow_exit": True, "cost_multiplier": 0.05}, | |
| {"rank": 2, "kv_precision": "INT4", "allow_exit": False, "cost_multiplier": 0.15}, | |
| {"rank": 4, "kv_precision": "INT8", "allow_exit": False, "cost_multiplier": 0.35}, | |
| {"rank": 8, "kv_precision": "INT8", "allow_exit": False, "cost_multiplier": 0.65}, | |
| {"rank": 16, "kv_precision": "FP16", "allow_exit": False, "cost_multiplier": 1.00}, | |
| ] | |
| # Statistics | |
| self.step_count = 0 | |
| self.cumulative_reward = 0.0 | |
| self.action_counts = [0] * n_actions | |
| def select_action(self, z_t: torch.Tensor) -> Tuple[int, Dict[str, Any]]: | |
| """ | |
| Selects an action given 8D context vector z_t via LinUCB. | |
| Args: | |
| z_t: Tensor of shape (context_dim,) or (1, context_dim) | |
| """ | |
| z = z_t.detach().cpu().numpy().flatten().reshape(-1, 1) | |
| p = np.zeros(self.n_actions) | |
| for a in range(self.n_actions): | |
| A_inv = np.linalg.inv(self.A[a]) | |
| theta_a = A_inv @ self.b[a] | |
| # UCB score = expected reward + confidence bound | |
| expected_r = float(np.asarray(theta_a.T @ z).item()) | |
| conf_bound = float(self.alpha * np.sqrt(max(0.0, float(np.asarray(z.T @ A_inv @ z).item())))) | |
| p[a] = expected_r + conf_bound | |
| chosen_action = int(np.argmax(p)) | |
| self.action_counts[chosen_action] += 1 | |
| self.step_count += 1 | |
| spec = self.action_specs[chosen_action] | |
| telemetry = { | |
| "step": self.step_count, | |
| "chosen_action": chosen_action, | |
| "chosen_rank": spec["rank"], | |
| "chosen_kv_precision": spec["kv_precision"], | |
| "allow_exit": spec["allow_exit"], | |
| "predicted_score": float(p[chosen_action]), | |
| "action_counts": list(self.action_counts), | |
| } | |
| return chosen_action, telemetry | |
| def update(self, z_t: torch.Tensor, action: int, reward: float): | |
| """ | |
| Updates LinUCB parameters after observing the actual reward. | |
| R_t = Delta_Q - lambda_L * Latency - lambda_M * Memory | |
| """ | |
| z = z_t.detach().cpu().numpy().flatten().reshape(-1, 1) | |
| self.A[action] += z @ z.T | |
| self.b[action] += reward * z | |
| self.cumulative_reward += reward | |
| def compute_reward( | |
| token_surprise: float, | |
| action_rank: int, | |
| latency_budget_ratio: float = 1.0, | |
| memory_budget_ratio: float = 1.0, | |
| ) -> float: | |
| """ | |
| Multi-objective reward function balancing prediction quality vs hardware pressure. | |
| """ | |
| # Quality gain saturates if surprise is high and rank is high | |
| quality_gain = min(1.0, (action_rank / 16.0) * (token_surprise / 2.0)) | |
| # Resource cost penalty | |
| resource_penalty = (action_rank / 16.0) * (0.5 * latency_budget_ratio + 0.5 * memory_budget_ratio) | |
| return float(quality_gain - resource_penalty) | |