""" TAF Browser — Pyodide-compatible TAF formulas + recipes. Pure-Python deterministic computations of TAF (Thermodynamic Attention Framework) formulas, plus 5 cross-section recipes for the most common viability questions. Author: Carles Marin License: Apache-2.0 """ from __future__ import annotations import math import json # ════════════════════════════════════════════════════════════════════════════ # §26 — γ-Thermodynamics (OUR contribution) # ════════════════════════════════════════════════════════════════════════════ def gamma_pade(theta: float, T_eval: int) -> float: """§26.1 — γ = (2θ - T√2)/(2θ + T√2)""" z_sqrt2 = T_eval * math.sqrt(2) return (2 * theta - z_sqrt2) / (2 * theta + z_sqrt2) def gamma_decompose(gamma_pade_val, has_GQA=False, has_SWA=False, n_params=0.0) -> dict: """§26.10 — 5-axis decomposition (n=23 OLS, paper sesión 28).""" delta_GQA = +0.11 if has_GQA else 0.0 delta_SWA = -0.21 if has_SWA else 0.0 delta_post_IH = -0.15 if n_params >= 4e8 else 0.0 return { "pade_centroid": gamma_pade_val, "delta_GQA": delta_GQA, "delta_SWA": delta_SWA, "delta_post_IH": delta_post_IH, "gamma_corrected": gamma_pade_val + delta_GQA + delta_SWA + delta_post_IH, } def d_horizon(theta: float, gamma: float): """§26.2 — d_h = θ(1-γ)√2/(1+γ). None if γ outside (0,1).""" if gamma <= 0 or gamma >= 1: return None return theta * (1 - gamma) * math.sqrt(2) / (1 + gamma) def l_niah_c(d_horizon_val): """§26.5 — L_NIAH^c = 2·d_horizon.""" return None if d_horizon_val is None else 2 * d_horizon_val def chi_susceptibility(gamma: float) -> float: """§26.16 — χ = 1/|γ-1|.""" return float('inf') if gamma == 1.0 else 1.0 / abs(gamma - 1.0) def p_hallucinate(L: int, theta: float, gamma: float): """§26.9 — Horizon-overshoot probability.""" dh = d_horizon(theta, gamma) if dh is None or L <= 0: return None chi = chi_susceptibility(gamma) if chi == float('inf'): return None geom = max(0.0, 1.0 - (dh / L) ** (1 - gamma)) return geom * (math.sqrt(chi) / (1 + math.sqrt(chi))) def theta_design(gamma_target: float, T_eval: int) -> float: """§26.3 — θ to land at γ_target at T_eval (Padé inverse).""" if gamma_target >= 1 or gamma_target <= -1: raise ValueError("gamma_target must be in (-1, 1)") return T_eval * math.sqrt(2) * (1 + gamma_target) / (2 * (1 - gamma_target)) def alpha_opt(gamma_target: float, T_eval: int, theta_nominal: float) -> float: """§26.4 — α = θ_design / θ_nominal.""" return theta_design(gamma_target, T_eval) / theta_nominal def df_window(gamma: float, N: int, f: float = 0.90): """§26.7 — KV compression window. None outside [0.65, 0.85] zone.""" if not (0.65 <= gamma <= 0.85): return None if gamma >= 1: return int(f * N) inner = (1 - f) + f * N ** (1 - gamma) return int(math.ceil(inner ** (1 / (1 - gamma)))) def kv_soft_decay_regime(theta: float, gamma: float, T_train: int) -> str: """§26.8 — Soft decay régimen-bound. d_h ≳ T_train/2 ⇒ applies.""" dh = d_horizon(theta, gamma) if dh is None: return "use-hard-cutoff" ratio = dh / max(1, T_train / 2) if ratio >= 1.2: return "applies" if ratio >= 0.8: return "borderline" return "use-hard-cutoff" # ════════════════════════════════════════════════════════════════════════════ # §17 — Pre-training viability formulas # ════════════════════════════════════════════════════════════════════════════ def chinchilla_optimal_tokens(N_params: float, ratio: float = 20.0) -> float: """§17.30 — Chinchilla 20:1 token budget. D = ratio · N.""" return ratio * N_params def chinchilla_optimal_N(D_tokens: float, ratio: float = 20.0) -> float: """§17.30 inverse — given D tokens, optimal N = D/20.""" return D_tokens / ratio def training_flops(N_params: float, D_tokens: float) -> float: """§17.10 — C ≈ 6·N·D total training FLOPs.""" return 6 * N_params * D_tokens def training_memory_16N(N_params: float) -> dict: """§17.20 — total memory ≈ 16·N bytes (model + grads + Adam moments).""" bytes_total = 16 * N_params return { "bytes": bytes_total, "GB": bytes_total / 1e9, } def emergent_threshold(N_params: float) -> str: """§17.60 — capability threshold heuristic (Wei 2022).""" if N_params >= 1e11: return "above 100B — strong reasoning capabilities expected" if N_params >= 1e10: return "above 10B — most emergent capabilities present" if N_params >= 1e9: return "above 1B — basic instruction-following, not strong reasoning" if N_params >= 1e8: return "above 100M — useful for narrow tasks, no emergence" return "below 100M — domain-specific tasks only" # ════════════════════════════════════════════════════════════════════════════ # §19 — Inference economics # ════════════════════════════════════════════════════════════════════════════ def kv_cache_memory(n_layers, n_kv_heads, d_head, seq_len, bytes_per_element=2.0) -> dict: """§19.1 — bytes = 2·L·n_kv·d_h·seq·B.""" bytes_total = 2 * n_layers * n_kv_heads * d_head * seq_len * bytes_per_element return {"bytes": bytes_total, "MB": bytes_total / 1e6, "GB": bytes_total / 1e9} def model_weights_memory(N_params, bytes_per_element=2.0) -> dict: """Inference memory for model weights only (BF16=2, INT8=1, INT4=0.5).""" return {"GB": N_params * bytes_per_element / 1e9} def inference_decode_throughput(N_params, hbm_GB_per_s, bytes_per_element=2.0) -> float: """§19.7 — memory-bound decode: tokens/sec = HBM_BW / model_size.""" model_GB = N_params * bytes_per_element / 1e9 return hbm_GB_per_s / model_GB # ════════════════════════════════════════════════════════════════════════════ # §20 — Hardware catalog (curated from vendor docs 2026) # ════════════════════════════════════════════════════════════════════════════ GPU_CATALOG = { # name: {bf16_TFLOPs, hbm_GB, hbm_GB_s, cloud_USD_per_h_spot, tdp_W} "H100 SXM": {"flops": 989, "vram_GB": 80, "bw_GB_s": 3350, "usd_h": 2.5, "tdp": 700}, "H100 PCIe": {"flops": 756, "vram_GB": 80, "bw_GB_s": 2000, "usd_h": 2.0, "tdp": 350}, "H200": {"flops": 989, "vram_GB": 141, "bw_GB_s": 4800, "usd_h": 3.5, "tdp": 700}, "B200": {"flops": 2250, "vram_GB": 192, "bw_GB_s": 8000, "usd_h": 5.0, "tdp": 1000}, "A100 80GB": {"flops": 312, "vram_GB": 80, "bw_GB_s": 2000, "usd_h": 1.2, "tdp": 400}, "A100 40GB": {"flops": 312, "vram_GB": 40, "bw_GB_s": 1555, "usd_h": 1.0, "tdp": 400}, "L40S": {"flops": 362, "vram_GB": 48, "bw_GB_s": 864, "usd_h": 0.7, "tdp": 350}, "MI300X": {"flops": 1307, "vram_GB": 192, "bw_GB_s": 5300, "usd_h": 2.1, "tdp": 750}, "RTX 4090": {"flops": 165, "vram_GB": 24, "bw_GB_s": 1008, "usd_h": 0.4, "tdp": 450}, "RTX 5090": {"flops": 419, "vram_GB": 32, "bw_GB_s": 1792, "usd_h": 0.7, "tdp": 575}, "RTX 5060Ti":{"flops": 36, "vram_GB": 16, "bw_GB_s": 448, "usd_h": 0.0, "tdp": 180}, # local } def cost_per_training_run(N_params: float, D_tokens: float, gpu: str = "H100 SXM", n_gpus: int = 8, mfu: float = 0.45) -> dict: """§20.11 — cost = (flops_total / (peak·MFU·n_gpus)) · USD/h.""" info = GPU_CATALOG.get(gpu) if info is None: return {"error": f"unknown gpu '{gpu}'", "available": list(GPU_CATALOG.keys())} total_flops = training_flops(N_params, D_tokens) # absolute FLOPs effective_flops_per_sec = info["flops"] * 1e12 * mfu * n_gpus seconds = total_flops / effective_flops_per_sec hours = seconds / 3600 usd = hours * info["usd_h"] * n_gpus return { "total_FLOPs": total_flops, "hours": hours, "days": hours / 24, "USD": usd, "gpu": gpu, "n_gpus": n_gpus, "mfu": mfu, } def cost_per_inference_token(model_GB: float, gpu: str, batch: int = 1) -> dict: """§19.9 / §20.12 — derived $/Mtok from memory-bound decode.""" info = GPU_CATALOG.get(gpu) if info is None: return {"error": f"unknown gpu '{gpu}'"} tok_per_sec = info["bw_GB_s"] / model_GB * batch sec_per_Mtok = 1e6 / tok_per_sec h_per_Mtok = sec_per_Mtok / 3600 usd_per_Mtok = h_per_Mtok * info["usd_h"] return { "tok_per_sec": tok_per_sec, "USD_per_Mtok": usd_per_Mtok, "gpu": gpu, "batch": batch, } # ════════════════════════════════════════════════════════════════════════════ # §24 — Cost / ROI # ════════════════════════════════════════════════════════════════════════════ API_PRICING = { # USD per million tokens (input/output blended typical) "GPT-4o": {"input": 2.5, "output": 10.0}, "GPT-4o-mini": {"input": 0.15, "output": 0.60}, "Claude-Opus-4": {"input": 15.0, "output": 75.0}, "Claude-Sonnet-4":{"input": 3.0, "output": 15.0}, "Claude-Haiku-4": {"input": 0.80, "output": 4.0}, "Gemini-1.5-Pro": {"input": 1.25, "output": 5.0}, "DeepSeek-V3": {"input": 0.27, "output": 1.10}, "Llama-3.3-70B (Together)": {"input": 0.88, "output": 0.88}, } def break_even_volume(training_cost: float, self_inference_per_Mtok: float, api_per_Mtok: float, blend_input_output: float = 0.5) -> dict: """§24.3 — monthly tokens at which custom training pays off.""" savings_per_Mtok = api_per_Mtok - self_inference_per_Mtok if savings_per_Mtok <= 0: return {"error": "self-host more expensive than API per token; never breaks even"} Mtok_breakeven = training_cost / savings_per_Mtok return { "savings_per_Mtok": savings_per_Mtok, "Mtok_breakeven": Mtok_breakeven, "tokens_breakeven": Mtok_breakeven * 1e6, } # ════════════════════════════════════════════════════════════════════════════ # RECIPES # ════════════════════════════════════════════════════════════════════════════ # ───────────────────────────────────────────────────────────────────── # X-2 — Long Context Viability # ───────────────────────────────────────────────────────────────────── def run_recipe_x2(theta, T_train, T_eval, n_attention_heads, n_kv_heads, d_head, n_layers, n_params, has_SWA=False, bytes_per_element=2.0, **_unused): """X-2: will model M serve length L doing NIAH retrieval?""" chain = [] g_pade = gamma_pade(theta, T_eval) chain.append(_step(1, "§26.1", "γ_Padé", "γ = (2θ - T√2)/(2θ + T√2)", {"theta": theta, "T_eval": T_eval}, g_pade, _phase_label(g_pade))) has_GQA = (n_kv_heads < n_attention_heads) decomp = gamma_decompose(g_pade, has_GQA=has_GQA, has_SWA=has_SWA, n_params=n_params) g_corr = decomp["gamma_corrected"] chain.append(_step(2, "§26.10", "γ-decomposition", "γ + δ_GQA + δ_SWA + δ_post_IH", {"has_GQA": has_GQA, "has_SWA": has_SWA, "n_params": n_params}, g_corr, breakdown=decomp)) dh = d_horizon(theta, g_corr) chain.append(_step(3, "§26.2", "d_horizon", "d_h = θ(1-γ)√2/(1+γ)", {"theta": theta, "gamma": g_corr}, dh, "n/a — γ outside (0,1)" if dh is None else f"horizon at d={dh:.0f}")) l_niah = l_niah_c(dh) chain.append(_step(4, "§26.5", "L_NIAH^c", "L_NIAH^c = 2·d_horizon", {"d_horizon": dh}, l_niah, "n/a" if l_niah is None else f"NIAH 50% at L={l_niah:.0f}")) p_hallu = p_hallucinate(T_eval, theta, g_corr) chain.append(_step(5, "§26.9", "P_hallucinate", "max(0,1-(d_h/L)^(1-γ))·√χ/(1+√χ)", {"L": T_eval, "theta": theta, "gamma": g_corr}, p_hallu, "n/a (Phase B)" if p_hallu is None else f"{p_hallu*100:.1f}% predicted")) kv = kv_cache_memory(n_layers, n_kv_heads, d_head, T_eval, bytes_per_element) chain.append(_step(6, "§19.1", "KV cache memory", "2·L·n_kv·d_h·seq·B", {"n_layers": n_layers, "n_kv_heads": n_kv_heads, "d_head": d_head, "seq_len": T_eval, "bytes_per_element": bytes_per_element}, kv, f"{kv['GB']:.2f} GB per request")) if g_corr <= 0 or g_corr >= 1: verdict, reason = "NO", "Phase B / geometric collapse (γ_corrected outside (0,1))" mit = (f"Apply NTK-aware extension. Required θ for γ=0.85: " f"{theta_design(0.85, T_eval):,.0f}. α_opt = {alpha_opt(0.85, T_eval, theta):.2f} " f"({'fine-tuning required' if alpha_opt(0.85, T_eval, theta) > 8 else 'zero-shot may work'}).") elif dh is not None and T_eval < dh: margin = (1 - T_eval / dh) * 100 verdict, reason = "YES", f"L={T_eval} inside d_horizon={dh:.0f} ({margin:.0f}% margin)." mit = "None required." elif dh is not None and T_eval < l_niah: verdict, reason = "DEGRADED", f"L between d_horizon ({dh:.0f}) and L_NIAH^c ({l_niah:.0f})." mit = "Consider context contraction OR NTK extension." else: verdict, reason = "NO", f"L={T_eval} exceeds NIAH ceiling {l_niah:.0f}." mit = f"Apply NTK extension; need θ ≈ {theta_design(0.85, T_eval):,.0f} for γ=0.85." return _wrap("X-2", "Long Context Viability", locals(), chain, verdict, reason, mit) # ───────────────────────────────────────────────────────────────────── # X-1 — Custom training vs API for a domain task # ───────────────────────────────────────────────────────────────────── def run_recipe_x1(N_params, D_tokens=None, gpu="H100 SXM", n_gpus=8, mfu=0.45, api_model="GPT-4o", monthly_tokens_M=10.0, **_unused): """X-1: custom training (Chinchilla optimal) vs API.""" chain = [] # Step 1: Chinchilla optimal D if D_tokens is None: D_tokens = chinchilla_optimal_tokens(N_params) chain.append(_step(1, "§17.30", "Chinchilla optimal D", "D = 20·N", {"N_params": N_params}, D_tokens, f"recommended D = {D_tokens:.2e} tokens")) # Step 2: training FLOPs flops = training_flops(N_params, D_tokens) chain.append(_step(2, "§17.10", "Training FLOPs", "C = 6·N·D", {"N": N_params, "D": D_tokens}, flops, f"{flops:.2e} FLOPs total")) # Step 3: training cost cost = cost_per_training_run(N_params, D_tokens, gpu=gpu, n_gpus=n_gpus, mfu=mfu) chain.append(_step(3, "§20.11", "Training cost", "hours·USD/h·n_gpus = total $", {"gpu": gpu, "n_gpus": n_gpus, "mfu": mfu}, cost, f"${cost['USD']:,.0f} over {cost['days']:.1f} days")) # Step 4: model_GB and decode throughput model_GB = N_params * 2 / 1e9 # BF16 inf = cost_per_inference_token(model_GB, gpu, batch=1) chain.append(_step(4, "§19.9 / §20.12", "Self-inference $/Mtok", "BW / model_GB → tok/s → $/Mtok", {"model_GB": model_GB, "gpu": gpu}, inf, f"${inf['USD_per_Mtok']:.2f} per million tokens (single user)")) # Step 5: API blended price api = API_PRICING.get(api_model, {"input": 2.0, "output": 8.0}) api_blend = (api["input"] + api["output"]) / 2 chain.append(_step(5, "§24.X", f"{api_model} blended price", "(input + output) / 2 USD/Mtok", {"api_model": api_model}, api_blend, f"${api_blend:.2f}/Mtok blended")) # Step 6: break-even be = break_even_volume(cost["USD"], inf["USD_per_Mtok"], api_blend) chain.append(_step(6, "§24.3", "Break-even tokens", "training$ / (api - self) = Mtok", {"training_cost": cost["USD"]}, be, _be_interp(be, monthly_tokens_M))) # Verdict if "error" in be: verdict, reason = "NO", be["error"] mit = f"Stick with {api_model} API." elif monthly_tokens_M >= be["Mtok_breakeven"]: verdict = "YES (custom)" months_to_payoff = be["Mtok_breakeven"] / monthly_tokens_M reason = (f"At {monthly_tokens_M} M tokens/month, break-even in " f"{months_to_payoff:.1f} months. Long-term custom is cheaper.") mit = f"Train at {gpu}×{n_gpus}; serve self-hosted." else: months = be["Mtok_breakeven"] / monthly_tokens_M verdict = "NO (API)" reason = (f"At {monthly_tokens_M} M tokens/month, break-even in " f"{months:.1f} months — too slow.") mit = f"Use {api_model} API (cheaper for your volume)." return _wrap("X-1", "Custom training vs API", locals(), chain, verdict, reason, mit) def _be_interp(be, monthly): if "error" in be: return be["error"] months = be["Mtok_breakeven"] / max(monthly, 0.001) return f"break-even at {be['Mtok_breakeven']:.0f} Mtok ({months:.1f} months at {monthly} M/mo)" # ───────────────────────────────────────────────────────────────────── # X-3 — Pre-flight check on $5K training budget # ───────────────────────────────────────────────────────────────────── def run_recipe_x3(USD_budget=5000.0, gpu="H100 SXM", mfu=0.45, n_gpus=1, **_unused): """X-3: given $ budget, what model can I train?""" chain = [] info = GPU_CATALOG[gpu] # Step 1: GPU-hours we can afford hours = USD_budget / (info["usd_h"] * n_gpus) chain.append(_step(1, "§20.11", "Affordable GPU-hours", "USD / ($/h·n_gpus)", {"USD": USD_budget, "gpu": gpu, "n_gpus": n_gpus}, hours, f"{hours:.0f} GPU-hours total ({hours/24:.1f} days at full use)")) # Step 2: max FLOPs max_flops = info["flops"] * 1e12 * mfu * n_gpus * hours * 3600 chain.append(_step(2, "§17.10", "Max training FLOPs", "peak·MFU·n_gpus·seconds", {"peak_TFLOPs": info["flops"], "MFU": mfu}, max_flops, f"{max_flops:.2e} effective FLOPs")) # Step 3: Chinchilla-optimal N (with D=20N) # 6·N·D = max_flops, D=20N → 120·N² = max_flops → N = sqrt(max_flops/120) N_chinchilla = math.sqrt(max_flops / 120) D_chinchilla = 20 * N_chinchilla chain.append(_step(3, "§17.30", "Chinchilla-optimal N", "N = √(C/120) at D=20N", {"max_FLOPs": max_flops}, N_chinchilla, f"N ≈ {N_chinchilla:.2e} params with D = {D_chinchilla:.2e} tokens")) # Step 4: emergence check emerg = emergent_threshold(N_chinchilla) chain.append(_step(4, "§17.60", "Emergence threshold", "Wei 2022 capability", {"N": N_chinchilla}, emerg, emerg)) # Step 5: memory budget check mem = training_memory_16N(N_chinchilla) fits = mem["GB"] <= info["vram_GB"] chain.append(_step(5, "§17.20", "16N training memory", "model + grads + AdamW", {"N": N_chinchilla}, mem, f"{mem['GB']:.1f} GB needed; " f"{'fits in ' if fits else 'EXCEEDS '}{info['vram_GB']} GB VRAM")) # Verdict if N_chinchilla < 1e8: verdict, reason = "TINY-MODEL", f"Budget supports only ~{N_chinchilla:.0e} params" mit = "Use LoRA fine-tuning of larger pretrained model instead." elif not fits: verdict, reason = "MEMORY-LIMITED", f"Chinchilla N ({N_chinchilla:.1e}) doesn't fit one {gpu}" mit = f"Use ZeRO-3 across multiple GPUs (need ≥{math.ceil(mem['GB']/info['vram_GB'])}× {gpu}) OR train smaller N undertrained." else: verdict = "GO" reason = (f"At ${USD_budget}, train {N_chinchilla:.1e}-param model on " f"{D_chinchilla:.1e} tokens in ~{hours/24:.1f} days. " f"Capability tier: {emerg.split('—')[0].strip()}.") mit = "None — proceed with Chinchilla-optimal recipe." return _wrap("X-3", "Budget pre-flight", locals(), chain, verdict, reason, mit) # ───────────────────────────────────────────────────────────────────── # X-5 — Hardware selection for serving # ───────────────────────────────────────────────────────────────────── def run_recipe_x5(N_params, T_eval=4096, n_layers=32, n_kv_heads=8, d_head=128, bytes_per_weight=2.0, target_tokens_per_day=10_000_000.0, concurrent_users=1, **_unused): """X-5: which GPU should I use to serve N-param model at L context?""" chain = [] # Step 1: weights memory w_mem = model_weights_memory(N_params, bytes_per_weight) chain.append(_step(1, "§19.X", "Model weights memory", "N · bytes_per_weight", {"N": N_params, "bytes": bytes_per_weight}, w_mem, f"{w_mem['GB']:.1f} GB for weights")) # Step 2: KV cache per request kv = kv_cache_memory(n_layers, n_kv_heads, d_head, T_eval, bytes_per_weight) chain.append(_step(2, "§19.1", "KV cache (per request)", "2·L·n_kv·d_h·seq·B", {"n_layers": n_layers, "n_kv": n_kv_heads, "d_head": d_head, "seq": T_eval}, kv, f"{kv['GB']:.2f} GB per concurrent request")) # Step 3: total memory needed total_GB = w_mem["GB"] + kv["GB"] * concurrent_users chain.append(_step(3, "§20.3", "Total GPU memory", "weights + KV·n_concurrent", {}, {"GB": total_GB}, f"{total_GB:.1f} GB for {concurrent_users} concurrent users")) # Step 4: scan GPU catalog candidates = [] for name, info in GPU_CATALOG.items(): if info["vram_GB"] < total_GB: continue # Decode throughput estimate (memory-bound) tok_per_s = info["bw_GB_s"] / w_mem["GB"] tok_per_day = tok_per_s * 86400 capacity_users = tok_per_day / target_tokens_per_day usd_per_day = info["usd_h"] * 24 usd_per_Mtok = (usd_per_day / (tok_per_day / 1e6)) if tok_per_day > 0 else float('inf') candidates.append({ "gpu": name, "vram_GB": info["vram_GB"], "bw_GB_s": info["bw_GB_s"], "tok_per_sec": tok_per_s, "tok_per_day": tok_per_day, "USD_per_day": usd_per_day, "USD_per_Mtok": usd_per_Mtok, "users_supported": capacity_users, }) candidates.sort(key=lambda c: c["USD_per_Mtok"]) chain.append(_step(4, "§20", f"Eligible GPUs (≥{total_GB:.0f}GB)", "filter + rank by $/Mtok", {"min_VRAM": total_GB}, candidates[:5], f"{len(candidates)} GPUs fit; cheapest: {candidates[0]['gpu'] if candidates else 'NONE'}")) # Verdict if not candidates: verdict, reason = "NO", f"No single GPU has ≥{total_GB:.0f} GB VRAM." mit = (f"Use tensor parallelism across multiple GPUs " f"(e.g. 2× H100 = 160GB), or quantize to INT8 (halves memory).") else: best = candidates[0] verdict = "YES" reason = (f"Best GPU: {best['gpu']} at ${best['USD_per_Mtok']:.2f}/Mtok. " f"Supports {best['users_supported']:.1f}× your daily target.") mit = f"Provision {best['gpu']}, expected {best['tok_per_sec']:.0f} tok/s decode." return _wrap("X-5", "Hardware selection for serving", locals(), chain, verdict, reason, mit) # ───────────────────────────────────────────────────────────────────── # X-19 — KV compression decision (ours vs literature) # ───────────────────────────────────────────────────────────────────── def run_recipe_x19(theta, T_train, T_eval, n_attention_heads, n_kv_heads, d_head, n_layers, n_params, has_SWA=False, **_unused): """X-19: should I use γ-soft KV decay, hard D_f, or literature methods?""" chain = [] # Step 1: γ_Padé g_pade = gamma_pade(theta, T_eval) chain.append(_step(1, "§26.1", "γ_Padé", "(2θ-T√2)/(2θ+T√2)", {"theta": theta, "T_eval": T_eval}, g_pade, _phase_label(g_pade))) # Step 2: γ-decomposition has_GQA = n_kv_heads < n_attention_heads decomp = gamma_decompose(g_pade, has_GQA, has_SWA, n_params) g_corr = decomp["gamma_corrected"] chain.append(_step(2, "§26.10", "γ-decomposition", "5-axis adjustment", {"has_GQA": has_GQA, "has_SWA": has_SWA, "n_params": n_params}, g_corr)) # Step 3: §26.7 D_f window applicability df = df_window(g_corr, T_eval, f=0.90) df_zone_ok = df is not None chain.append(_step(3, "§26.7", "D_f window (γ in [0.65, 0.85])", "[(1-f)+fN^(1-γ)]^(1/(1-γ))", {"gamma": g_corr, "N": T_eval, "f": 0.9}, df, f"D_f = {df}" if df_zone_ok else f"NOT applicable (γ={g_corr:.3f} outside [0.65, 0.85])")) # Step 4: §26.8 soft decay régimen regime = kv_soft_decay_regime(theta, g_corr, T_train) dh = d_horizon(theta, g_corr) dh_str = f"{dh:.0f}" if dh is not None else "n/a" chain.append(_step(4, "§26.8", "Soft decay régimen", "d_h ≳ T_train/2", {"theta": theta, "gamma": g_corr, "T_train": T_train}, regime, f"d_horizon={dh_str}; regime: {regime}")) # Step 5: KV cache memory baseline kv = kv_cache_memory(n_layers, n_kv_heads, d_head, T_eval) chain.append(_step(5, "§19.1", "Baseline KV memory", "2·L·n_kv·d_h·seq·B", {"L": n_layers, "n_kv": n_kv_heads, "d_h": d_head, "seq": T_eval}, kv, f"{kv['GB']:.2f} GB without compression")) # Verdict if regime == "applies" and df_zone_ok: verdict = "USE SOFT DECAY" reason = (f"d_horizon ≳ T_train/2 AND γ in compression zone. " f"Soft decay (1-d/d_h)^γ best (-21% PPL vs hard cutoff per F17).") mit = "Implement as 4D attention_mask additive bias with eager attention." elif df_zone_ok: verdict = "USE D_f HARD CUTOFF" reason = f"γ in [0.65, 0.85] zone but d_h < T_train/2. Hard truncation at D_f={df} works." mit = "Set cache_max_len = D_f." elif regime == "applies": verdict = "USE SOFT DECAY (caveat)" reason = "Régimen applies but γ outside D_f validity zone. Soft decay only." mit = "Soft decay; do not use D_f window." elif g_corr >= 1 or g_corr <= 0: verdict = "USE LITERATURE METHODS" reason = f"γ={g_corr:.3f} outside Phase A. Our formulas don't apply." mit = "Use SnapKV / PyramidKV / FastGen (literature heuristics)." else: verdict = "USE HARD T_train CUTOFF" reason = "Régimen not met AND γ outside zone. Cap context at T_train." mit = f"Set seq_len ≤ {T_train}, no extension." return _wrap("X-19", "KV compression decision", locals(), chain, verdict, reason, mit) # ════════════════════════════════════════════════════════════════════════════ # Helpers # ════════════════════════════════════════════════════════════════════════════ def _step(n, sec, name, formula, inputs, result, interpretation=None, breakdown=None): s = {"step": n, "section": sec, "name": name, "formula": formula, "inputs": inputs, "result": result} if interpretation: s["interpretation"] = interpretation if breakdown: s["breakdown"] = breakdown return s def _wrap(rid, rname, locals_dict, chain, verdict, reason, mitigation): # Clean inputs (drop chain/internal vars) inputs = {k: v for k, v in locals_dict.items() if not k.startswith("_") and k not in ("chain", "verdict", "reason", "mit", "info", "be", "kv", "g_pade", "g_corr", "decomp", "dh", "l_niah", "p_hallu", "cost", "model_GB", "inf", "api", "api_blend", "fits", "mem", "emerg", "max_flops", "hours", "N_chinchilla", "D_chinchilla", "candidates", "best", "tok_per_s", "tok_per_day", "capacity_users", "usd_per_day", "usd_per_Mtok", "total_GB", "w_mem", "df", "df_zone_ok", "regime", "has_GQA", "margin", "months", "months_to_payoff", "name")} return {"recipe_id": rid, "recipe_name": rname, "inputs": inputs, "chain": chain, "verdict": verdict, "reason": reason, "mitigation": mitigation} def _phase_label(g): if 0 < g < 1: return "Phase A (long-range OK)" if g >= 1: return "Phase B / Hagedorn" return "Phase B / catastrophic (negative γ — T too large for θ)" # ════════════════════════════════════════════════════════════════════════════ # Recipe registry # ════════════════════════════════════════════════════════════════════════════ RECIPES = { "X-1": { "name": "Custom Training vs API", "description": "Should I train a custom model or use a frontier API for my domain task?", "fn": run_recipe_x1, "params": ["N_params", "D_tokens", "gpu", "n_gpus", "mfu", "api_model", "monthly_tokens_M"], "category": "build-vs-buy", "uses_sections": ["§17", "§19", "§20", "§24"], }, "X-2": { "name": "Long Context Viability", "description": "Will model M serve length L doing Needle-in-a-Haystack retrieval?", "fn": run_recipe_x2, "params": ["theta", "T_train", "T_eval", "n_attention_heads", "n_kv_heads", "d_head", "n_layers", "n_params", "has_SWA"], "category": "long-context", "uses_sections": ["§26", "§19"], }, "X-3": { "name": "Budget Pre-flight", "description": "Given $ budget, what model is feasible to train?", "fn": run_recipe_x3, "params": ["USD_budget", "gpu", "mfu", "n_gpus"], "category": "training-budget", "uses_sections": ["§17", "§20"], }, "X-5": { "name": "Hardware Selection", "description": "Which GPU should I use to serve my model at target throughput?", "fn": run_recipe_x5, "params": ["N_params", "T_eval", "n_layers", "n_kv_heads", "d_head", "bytes_per_weight", "target_tokens_per_day", "concurrent_users"], "category": "serving", "uses_sections": ["§19", "§20"], }, "X-19": { "name": "KV Compression Decision", "description": "Should I use soft decay, D_f cutoff, or literature methods to compress KV?", "fn": run_recipe_x19, "params": ["theta", "T_train", "T_eval", "n_attention_heads", "n_kv_heads", "d_head", "n_layers", "n_params", "has_SWA"], "category": "kv-compression", "uses_sections": ["§26", "§19"], }, } def list_recipes() -> str: """Return JSON of all recipes for UI dropdown.""" return json.dumps([ {"id": rid, "name": r["name"], "description": r["description"], "category": r["category"], "params": r["params"], "uses_sections": r["uses_sections"]} for rid, r in RECIPES.items() ]) def run_recipe(recipe_id: str, **params) -> dict: """Dispatcher — execute recipe by id with given params.""" r = RECIPES.get(recipe_id) if r is None: return {"error": f"unknown recipe '{recipe_id}'", "available": list(RECIPES.keys())} return r["fn"](**params) # ════════════════════════════════════════════════════════════════════════════ # Known model presets # ════════════════════════════════════════════════════════════════════════════ PRESETS = { "EleutherAI/pythia-2.8b": { "theta": 10000, "T_train": 2048, "n_attention_heads": 32, "n_kv_heads": 32, "d_head": 80, "n_layers": 32, "n_params": 2.8e9, "has_SWA": False, }, "EleutherAI/pythia-1b": { "theta": 10000, "T_train": 2048, "n_attention_heads": 8, "n_kv_heads": 8, "d_head": 256, "n_layers": 16, "n_params": 1e9, "has_SWA": False, }, "EleutherAI/pythia-1.4b": { "theta": 10000, "T_train": 2048, "n_attention_heads": 16, "n_kv_heads": 16, "d_head": 128, "n_layers": 24, "n_params": 1.4e9, "has_SWA": False, }, "meta-llama/Meta-Llama-3-8B": { "theta": 500000, "T_train": 8192, "n_attention_heads": 32, "n_kv_heads": 8, "d_head": 128, "n_layers": 32, "n_params": 8e9, "has_SWA": False, }, "meta-llama/Llama-3.2-1B": { "theta": 500000, "T_train": 131072, "n_attention_heads": 32, "n_kv_heads": 8, "d_head": 64, "n_layers": 16, "n_params": 1.2e9, "has_SWA": False, }, "meta-llama/Llama-3.3-70B-Instruct": { "theta": 500000, "T_train": 131072, "n_attention_heads": 64, "n_kv_heads": 8, "d_head": 128, "n_layers": 80, "n_params": 70e9, "has_SWA": False, }, "mistralai/Mistral-7B-v0.1": { "theta": 10000, "T_train": 8192, "n_attention_heads": 32, "n_kv_heads": 8, "d_head": 128, "n_layers": 32, "n_params": 7e9, "has_SWA": True, }, "Qwen/Qwen2.5-7B": { "theta": 1000000, "T_train": 32768, "n_attention_heads": 28, "n_kv_heads": 4, "d_head": 128, "n_layers": 28, "n_params": 7.6e9, "has_SWA": False, }, "Qwen/Qwen2.5-1.5B": { "theta": 1000000, "T_train": 32768, "n_attention_heads": 12, "n_kv_heads": 2, "d_head": 128, "n_layers": 28, "n_params": 1.5e9, "has_SWA": False, }, "google/gemma-2-9b-it": { "theta": 10000, "T_train": 8192, "n_attention_heads": 16, "n_kv_heads": 8, "d_head": 256, "n_layers": 42, "n_params": 9e9, "has_SWA": True, }, "microsoft/phi-3-mini-4k-instruct": { "theta": 10000, "T_train": 4096, "n_attention_heads": 32, "n_kv_heads": 32, "d_head": 96, "n_layers": 32, "n_params": 3.8e9, "has_SWA": True, }, } def list_presets() -> str: return json.dumps([ {"id": k, "label": k.split("/")[-1], "theta": v["theta"], "T_train": v["T_train"]} for k, v in PRESETS.items() ]) def get_preset(model_id: str) -> dict: return PRESETS.get(model_id, {}) # Smoke test if __name__ == "__main__": print("─── X-2 Llama-3-8B @ 32K ───") r = run_recipe("X-2", theta=500_000, T_train=8192, T_eval=32_000, n_attention_heads=32, n_kv_heads=8, d_head=128, n_layers=32, n_params=8e9, has_SWA=False) print(f"Verdict: {r['verdict']} — {r['reason']}\n") print("─── X-1 Llama-3-8B vs GPT-4o (10M tok/mo) ───") r = run_recipe("X-1", N_params=8e9, monthly_tokens_M=10.0, api_model="GPT-4o") print(f"Verdict: {r['verdict']} — {r['reason']}\n") print("─── X-3 budget $5K ───") r = run_recipe("X-3", USD_budget=5000.0, gpu="H100 SXM", n_gpus=1) print(f"Verdict: {r['verdict']} — {r['reason']}\n") print("─── X-5 serve Llama-3-8B at 4K ───") r = run_recipe("X-5", N_params=8e9, T_eval=4096, n_layers=32, n_kv_heads=8, d_head=128, target_tokens_per_day=10e6, concurrent_users=1) print(f"Verdict: {r['verdict']} — {r['reason']}\n") print("─── X-19 KV compression for Llama-3-8B ───") r = run_recipe("X-19", theta=500_000, T_train=8192, T_eval=8192, n_attention_heads=32, n_kv_heads=8, d_head=128, n_layers=32, n_params=8e9) print(f"Verdict: {r['verdict']} — {r['reason']}\n")