diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 322f941..e5b6a97 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -6,6 +6,7 @@ Run: """ from __future__ import annotations +import dataclasses import importlib import sys @@ -53,6 +54,7 @@ from tests.analytical_visualization.stage_shapes import ( ffn_stage_shape_rows, ) from tests.analytical_visualization.chip_roofline import ( + ai_sensitivity_curve, arithmetic_intensity, balance_context, bound_regime, @@ -60,6 +62,8 @@ from tests.analytical_visualization.chip_roofline import ( good_batch, good_context, knee_batch, + memory_budget_curve_vs_batch, + memory_budget_curve_vs_skv, per_token_latency_curve, step_latency_curve, t_com, @@ -2179,37 +2183,64 @@ with tab_roofline: _l_star = balance_context(_default_machine, model) _n_active = total_active_params(model) - # ── Local knobs: S_kv and B — override sidebar for exploring the - # plots without changing the rest of the app's config. - st.markdown("**Explore: override S_kv and B for the plots below**") - _kn_l, _kn_r = st.columns(2) - with _kn_l: + # ── Local knobs: S_kv, B, FLOPs multiplier, BW multiplier ───── + # Override sidebar for exploring the plots without changing the + # rest of the app's config. + st.markdown("**Explore: local overrides for the plots below**") + _kn_a, _kn_b, _kn_c, _kn_d = st.columns(4) + with _kn_a: _skv_min = 128 - _skv_max = 1_000_000 # 1M tokens — the far end of interest + _skv_max = 1_000_000 _rf_skv = st.slider( "S_kv (context length)", min_value=_skv_min, max_value=_skv_max, value=min(int(s_kv), _skv_max), step=_skv_min, key="_rf_skv_override", - help=("Local to this tab. Drag to see how KV-read time grows " - "and where the knee disappears past L*."), + help="Local to this tab. Drives every plot's S_kv.", ) - with _kn_r: + with _kn_b: _b_min = 1 - _b_max = 256 # small scale — decode B rarely goes higher in practice + _b_max = 256 _rf_b = st.slider( "B (batch size)", min_value=_b_min, max_value=_b_max, value=min(max(1, int(b_batch)), _b_max), key="_rf_b_override", - help=("Local to this tab. Drag to move the vertical B " - "marker on plots and see which regime dominates."), + help="Local to this tab. Marker on plots + PE-memory KV.", ) + with _kn_c: + _rf_flops_mult = st.slider( + "FLOPs × (chip)", + min_value=0.25, max_value=8.0, + value=1.0, step=0.25, + key="_rf_flops_mult", + help=("Scale peak_tflops for the AI-sensitivity plot. " + "Doubling FLOPs doubles AI and B*."), + ) + with _kn_d: + _rf_bw_mult = st.slider( + "HBM BW × (chip)", + min_value=0.25, max_value=8.0, + value=1.0, step=0.25, + key="_rf_bw_mult", + help=("Scale bw_hbm for the AI-sensitivity plot. " + "Doubling BW halves AI and B*."), + ) + # Build a scaled machine for the AI-sensitivity view. + _scaled_machine = dataclasses.replace( + _default_machine, + peak_tflops_f16=_default_machine.peak_tflops_f16 * _rf_flops_mult, + bw_hbm_gbs=_default_machine.bw_hbm_gbs * _rf_bw_mult, + ) + _ai_scaled = arithmetic_intensity(_scaled_machine) + _b_star_scaled = critical_batch(_scaled_machine, model) st.caption( - f"Using **S_kv = {_rf_skv:,}** tokens, **B = {_rf_b}**. " - f"L* = **{_l_star:,.0f}** tokens → S_kv/L* = **{_rf_skv/_l_star:.2f}**. " - f"B* = **{_b_star:.0f}** → B/B* = **{_rf_b/_b_star:.2f}**." + f"**S_kv = {_rf_skv:,}**, **B = {_rf_b}**, " + f"**FLOPs × {_rf_flops_mult:.2f}**, **BW × {_rf_bw_mult:.2f}**. " + f"L* = **{_l_star:,.0f}** tokens (base chip) → S_kv/L* = " + f"**{_rf_skv/_l_star:.2f}**. Base B* = **{_b_star:.0f}**; " + f"scaled B* = **{_b_star_scaled:.0f}**." ) # ── Headline plots: step latency and cost per token (=÷B) ───── @@ -2224,7 +2255,11 @@ with tab_roofline: # Plot A — Latency per decode step (undivided by B) with _hcol1: - st.markdown("**Latency per decode step** (one forward pass)") + st.markdown( + "**Latency per decode step** (one forward pass)\n\n" + "`t_step = N·b/W + 2·N·B/C + B·S_kv·kv_bpt/W`\n\n" + "` = weight_fetch + compute·B + KV_read·B`" + ) _figA, _axA = plt.subplots(figsize=(6, 4.2)) _axA.plot(_xs_head, [p.weight_s * 1e3 for p in _step_pts], "^-", color="#ffbe0b", @@ -2295,6 +2330,143 @@ with tab_roofline: st.divider() + # ── PE memory budget (sharded, per-PE) ──────────────────────── + st.markdown( + "**PE memory budget — how batch and context fill the HBM**" + ) + _hbm_gb = _default_machine.pe_hbm_gb + _skv_sweep = [128, 512, 2048, 8192, 32768, 131072, + 524288, 1_000_000] + _bud_skv = memory_budget_curve_vs_skv(cfg, _skv_sweep, batch=_rf_b) + _b_sweep = [1, 2, 4, 8, 16, 32, 64, 128, 256] + _bud_b = memory_budget_curve_vs_batch(cfg, _b_sweep, s_kv=_rf_skv) + + _mem_L, _mem_R = st.columns(2) + + with _mem_L: + st.markdown(f"vs S_kv (at B = {_rf_b}) — sharded CP={topo.cp}, " + f"TP={topo.tp}, PP={topo.pp}") + _figM1, _axM1 = plt.subplots(figsize=(6, 4.2)) + _xM = [p.axis_val for p in _bud_skv] + _w = [p.weights_gb for p in _bud_skv] + _k = [p.kv_gb for p in _bud_skv] + _tr = [p.transient_gb for p in _bud_skv] + _axM1.stackplot(_xM, _w, _k, _tr, + labels=["Weights", "KV cache", "Transient"], + colors=["#ffbe0b", "#d90429", "#3a86ff"], + alpha=0.8) + _axM1.axhline(_hbm_gb, color="#212529", linestyle="--", + linewidth=1.5, label=f"HBM budget = {_hbm_gb} GB") + _axM1.axvline(_rf_skv, color="#7b1fa2", linewidth=1.5, + alpha=0.6, label=f"S_kv = {_rf_skv:,}") + _axM1.set_xscale("log") + _axM1.set_xlabel("S_kv (tokens)") + _axM1.set_ylabel("Per-PE memory (GB)") + _axM1.set_title("Weight is fixed; KV fills the budget as S_kv grows") + _axM1.grid(True, which="both", alpha=0.3) + _axM1.legend(fontsize=8, loc="upper left") + plt.tight_layout() + st.pyplot(_figM1, width='stretch') + plt.close(_figM1) + + with _mem_R: + st.markdown(f"vs B (at S_kv = {_rf_skv:,}) — same sharding") + _figM2, _axM2 = plt.subplots(figsize=(6, 4.2)) + _xM2 = [p.axis_val for p in _bud_b] + _w2 = [p.weights_gb for p in _bud_b] + _k2 = [p.kv_gb for p in _bud_b] + _tr2 = [p.transient_gb for p in _bud_b] + _axM2.stackplot(_xM2, _w2, _k2, _tr2, + labels=["Weights", "KV cache", "Transient"], + colors=["#ffbe0b", "#d90429", "#3a86ff"], + alpha=0.8) + _axM2.axhline(_hbm_gb, color="#212529", linestyle="--", + linewidth=1.5, label=f"HBM budget = {_hbm_gb} GB") + _axM2.axvline(_rf_b, color="#7b1fa2", linewidth=1.5, + alpha=0.6, label=f"B = {_rf_b}") + _axM2.set_xscale("log", base=2) + _axM2.set_xlabel("Batch size B") + _axM2.set_ylabel("Per-PE memory (GB)") + _axM2.set_title("Weight fixed; KV grows linearly with B") + _axM2.grid(True, which="both", alpha=0.3) + _axM2.legend(fontsize=8, loc="upper left") + plt.tight_layout() + st.pyplot(_figM2, width='stretch') + plt.close(_figM2) + + _bud_now = memory_budget_curve_vs_skv(cfg, [_rf_skv], batch=_rf_b)[0] + st.caption( + f"At **S_kv = {_rf_skv:,}, B = {_rf_b}**: weights = " + f"**{_bud_now.weights_gb:.2f} GB**, KV = **{_bud_now.kv_gb:.2f} GB**, " + f"transient = **{_bud_now.transient_gb:.2f} GB**, " + f"used = **{_bud_now.used_gb:.2f} GB / {_hbm_gb:.1f} GB** " + f"({'**OVER BUDGET**' if _bud_now.over_budget else f'free = {_bud_now.free_gb:.2f} GB'}). " + "Weights are fixed by sharding (CP·TP·PP·EP splits them across " + "PEs); KV per PE = B × S_kv × kv_bpt_per_PE." + ) + + st.divider() + + # ── AI sensitivity to FLOPs / BW ────────────────────────────── + st.markdown("**AI sensitivity — how scaling chip FLOPs or BW moves B\\***") + _ai_mults = [0.25, 0.5, 1.0, 2.0, 4.0, 8.0] + _ai_flops_pts = ai_sensitivity_curve(_default_machine, model, + _ai_mults, axis="flops") + _ai_bw_pts = ai_sensitivity_curve(_default_machine, model, + _ai_mults, axis="bw") + _figAI, (_axAI1, _axAI2) = plt.subplots(1, 2, figsize=(12, 4.2)) + + _axAI1.plot(_ai_mults, [p.ai for p in _ai_flops_pts], "o-", + color="#3a86ff", linewidth=2, label="Scale FLOPs") + _axAI1.plot(_ai_mults, [p.ai for p in _ai_bw_pts], "s-", + color="#d90429", linewidth=2, label="Scale HBM BW") + _axAI1.axhline(_ai, color="#2e7d32", linestyle=":", + label=f"Base AI = {_ai:.1f}") + _axAI1.axvline(1.0, color="#888", linestyle=":", alpha=0.5) + _axAI1.set_xscale("log", base=2) + _axAI1.set_yscale("log") + _axAI1.set_xlabel("Multiplier on chip parameter") + _axAI1.set_ylabel("AI (FLOPs/byte)") + _axAI1.set_title("AI = C / W. FLOPs ↑ raises AI, BW ↑ lowers AI") + _axAI1.grid(True, which="both", alpha=0.3) + _axAI1.legend(fontsize=8, loc="best") + + _axAI2.plot(_ai_mults, [p.b_star for p in _ai_flops_pts], "o-", + color="#3a86ff", linewidth=2, label="Scale FLOPs") + _axAI2.plot(_ai_mults, [p.b_star for p in _ai_bw_pts], "s-", + color="#d90429", linewidth=2, label="Scale HBM BW") + _axAI2.axhline(_b_star, color="#2e7d32", linestyle=":", + label=f"Base B* = {_b_star:.0f}") + _axAI2.axvline(_rf_flops_mult, color="#3a86ff", linewidth=1.2, + alpha=0.5, label=f"FLOPs × {_rf_flops_mult:.2f}") + _axAI2.axvline(_rf_bw_mult, color="#d90429", linewidth=1.2, + alpha=0.5, label=f"BW × {_rf_bw_mult:.2f}", + linestyle="--") + _axAI2.set_xscale("log", base=2) + _axAI2.set_yscale("log") + _axAI2.set_xlabel("Multiplier on chip parameter") + _axAI2.set_ylabel("B* (critical batch)") + _axAI2.set_title("B* = AI · b/2 — same shape as AI (for BF16)") + _axAI2.grid(True, which="both", alpha=0.3) + _axAI2.legend(fontsize=8, loc="best") + + plt.tight_layout() + st.pyplot(_figAI, width='stretch') + plt.close(_figAI) + + st.caption( + f"**Left**: AI vs chip-parameter multiplier. Scaling **FLOPs** " + f"raises AI linearly (more compute per byte). Scaling **HBM BW** " + f"lowers AI (more bytes to feed per FLOP). " + f"**Right**: B\\* moves in the same direction as AI. " + f"**Your knob position**: FLOPs × {_rf_flops_mult:.2f}, " + f"BW × {_rf_bw_mult:.2f} → **AI = {_ai_scaled:.2f}**, " + f"**B\\* = {_b_star_scaled:.0f}** (vs base AI {_ai:.1f}, " + f"B\\* {_b_star:.0f})." + ) + + st.divider() + _regime_now = bound_regime(_default_machine, model, batch=max(1, _rf_b), s_kv=_rf_skv) @@ -2347,84 +2519,6 @@ with tab_roofline: st.divider() - # ── Plot 2: cost curves at multiple context lengths (no-knee) ─ - st.markdown("**The 'no-knee' phenomenon — cost vs B at different S_kv**") - # Reuse the small batch range from the headline plots so all charts - # on this tab share the same 1..256 batch domain. - _fig2, _ax2 = plt.subplots(figsize=(9, 4.5)) - _l_ratios = [0.25, 0.5, 1.0, 2.0, 5.0] - _colors = ["#3a86ff", "#2e7d32", "#ffbe0b", "#ef6c00", "#d90429"] - for _r, _col in zip(_l_ratios, _colors): - _skv_at = max(1, int(_r * _l_star)) - _pts_r = per_token_latency_curve(_default_machine, model, _b_range, - s_kv=_skv_at) - _label = f"S_kv = {_r:.2g} × L* ({_skv_at:,} tok)" - _ax2.plot(_xs_head, [p.total_s * 1e3 for p in _pts_r], "-", - color=_col, label=_label, linewidth=1.6) - _ax2.axhline(_tok_pts[0].compute_s * 1e3, linestyle="--", - color="#888", label="Compute floor") - _ax2.axvline(_b_star, linestyle=":", color="#2e7d32", - label=f"B* = {_b_star:.0f}") - _ax2.axvline(_rf_b, linestyle="-", color="#7b1fa2", alpha=0.6, - label=f"B = {_rf_b}") - _ax2.set_xscale("log", base=2) - _ax2.set_yscale("log") - _ax2.set_xlabel("Batch size") - _ax2.set_ylabel("Per-token step time (ms)") - _ax2.set_title("Cost curves at several context lengths " - f"(L* = {_l_star:,.0f} tokens)") - _ax2.grid(True, which="both", alpha=0.3) - _ax2.legend(fontsize=8, loc="upper right") - plt.tight_layout() - st.pyplot(_fig2, width='stretch') - plt.close(_fig2) - - st.caption( - "**S_kv < L\\*:** curve bends at B_knee = B\\* / (1 − S_kv/L\\*), " - "then flattens onto the compute floor. **S_kv = L\\*:** curves " - "become parallel — you approach the floor but never touch it. " - "**S_kv > L\\*:** KV read is above the compute floor at all B; " - "no batch size recovers utilization. This is the algebraic core " - "of 'at 1M context, there's no B that gets you compute-bound'." - ) - - st.divider() - - # ── Plot 3: knee vs context ─────────────────────────────────── - st.markdown("**Effective knee B_knee(S_kv) = B\\* / (1 − S_kv/L\\*)**") - _skv_range = [int(_r * _l_star) - for _r in (0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, - 0.7, 0.8, 0.85, 0.9, 0.95, 0.98)] - _knees = [knee_batch(_default_machine, model, s) for s in _skv_range] - _fig3, _ax3 = plt.subplots(figsize=(9, 3.5)) - _ax3.plot(_skv_range, _knees, "o-", color="#7b1fa2", linewidth=2) - _ax3.axhline(_b_star, linestyle=":", color="#2e7d32", - label=f"B* = {_b_star:.0f}") - _ax3.axvline(_l_star, linestyle=":", color="#d90429", - label=f"L* = {_l_star:,.0f}") - _ax3.axvline(_rf_skv, linestyle="-", color="#7b1fa2", alpha=0.6, - label=f"S_kv = {_rf_skv:,}") - _ax3.axhline(_rf_b, linestyle="-", color="#7b1fa2", alpha=0.6, - label=f"B = {_rf_b}") - _ax3.set_yscale("log") - _ax3.set_xlabel("Context length S_kv (tokens)") - _ax3.set_ylabel("Effective knee batch size") - _ax3.set_title("How far right the cost-curve knee slides as context grows") - _ax3.grid(True, which="both", alpha=0.3) - _ax3.legend(fontsize=8, loc="upper left") - plt.tight_layout() - st.pyplot(_fig3, width='stretch') - plt.close(_fig3) - - st.caption( - "The knee starts at B\\* for short context and slides right as " - "S_kv/L\\* grows. It diverges at S_kv = L\\* and disappears " - "beyond that. Sparse attention (e.g. DeepSeek-style) replaces " - "S_kv with ~√S_kv in the KV term — bending this curve back down." - ) - - st.divider() - # ── Regime formulas (short vs long context) ─────────────────── st.markdown("**Cost-term formulas by regime**") _t_com_val = t_com(_default_machine, model) * 1e3 # ms diff --git a/tests/analytical_visualization/chip_roofline.py b/tests/analytical_visualization/chip_roofline.py index 3793772..b215891 100644 --- a/tests/analytical_visualization/chip_roofline.py +++ b/tests/analytical_visualization/chip_roofline.py @@ -27,9 +27,9 @@ exactly at B*. """ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, replace -from .model_config import MachineParams, ModelConfig +from .model_config import FullConfig, MachineParams, ModelConfig # Per-token bytes of BF16 MAC arithmetic: one multiply + one add = 2 FLOPs. @@ -304,3 +304,115 @@ def step_latency_curve(machine: MachineParams, model: ModelConfig, total_s=total, )) return points + + +# ── PE memory budget curves ─────────────────────────────────────── + + +@dataclass +class MemoryBudgetPoint: + axis_val: int # S_kv or B being swept + weights_gb: float + kv_gb: float + transient_gb: float + used_gb: float + free_gb: float # max(0, hbm_gb - used_gb) + over_budget: bool + + +def _budget_point(weights_bytes: int, kv_bytes: int, transient_bytes: int, + hbm_bytes: int, axis_val: int) -> MemoryBudgetPoint: + used = weights_bytes + kv_bytes + transient_bytes + return MemoryBudgetPoint( + axis_val=axis_val, + weights_gb=weights_bytes / 1e9, + kv_gb=kv_bytes / 1e9, + transient_gb=transient_bytes / 1e9, + used_gb=used / 1e9, + free_gb=max(0, hbm_bytes - used) / 1e9, + over_budget=(used > hbm_bytes), + ) + + +def memory_budget_curve_vs_skv(cfg: FullConfig, + s_kv_range: list[int], + batch: int) -> list[MemoryBudgetPoint]: + """Per-PE memory as S_kv sweeps. Uses cfg's current sharding + (CP, TP, PP). Sets topo.b = batch. Returns one point per S_kv.""" + from .memory_layout import ( + per_pe_kv_cache_bytes, per_pe_transient_bytes, per_pe_weight_bytes, + ) + hbm_bytes = int(cfg.machine.pe_budget_bytes) + weights_bytes = per_pe_weight_bytes(cfg) + transient_bytes = per_pe_transient_bytes(cfg) + points: list[MemoryBudgetPoint] = [] + for skv in s_kv_range: + swept_topo = replace(cfg.topo, s_kv=int(skv), b=max(1, int(batch))) + swept_cfg = FullConfig(model=cfg.model, topo=swept_topo, + machine=cfg.machine) + kv_bytes = per_pe_kv_cache_bytes(swept_cfg) + points.append(_budget_point(weights_bytes, kv_bytes, + transient_bytes, hbm_bytes, int(skv))) + return points + + +def memory_budget_curve_vs_batch(cfg: FullConfig, + b_range: list[int], + s_kv: int) -> list[MemoryBudgetPoint]: + """Per-PE memory as B sweeps. Sets topo.s_kv = s_kv.""" + from .memory_layout import ( + per_pe_kv_cache_bytes, per_pe_transient_bytes, per_pe_weight_bytes, + ) + hbm_bytes = int(cfg.machine.pe_budget_bytes) + weights_bytes = per_pe_weight_bytes(cfg) + transient_bytes = per_pe_transient_bytes(cfg) + points: list[MemoryBudgetPoint] = [] + for bs in b_range: + swept_topo = replace(cfg.topo, b=max(1, int(bs)), s_kv=int(s_kv)) + swept_cfg = FullConfig(model=cfg.model, topo=swept_topo, + machine=cfg.machine) + kv_bytes = per_pe_kv_cache_bytes(swept_cfg) + points.append(_budget_point(weights_bytes, kv_bytes, + transient_bytes, hbm_bytes, int(bs))) + return points + + +# ── AI / B* sensitivity to hardware knobs ────────────────────────── + + +@dataclass +class AISensitivityPoint: + multiplier: float # scale factor applied to the base machine + peak_tflops: float + bw_gbs: float + ai: float # C / W + b_star: float # C·b/(2·W) + + +def ai_sensitivity_curve(machine: MachineParams, model: ModelConfig, + multipliers: list[float], + axis: str = "flops") -> list[AISensitivityPoint]: + """Sweep FLOPs OR BW while holding the other fixed. Returns AI + B* + at each multiplier. + + axis='flops' → scales peak_tflops_f16 (AI grows linearly). + axis='bw' → scales bw_hbm_gbs (AI shrinks inversely). + """ + if axis not in ("flops", "bw"): + raise ValueError(f"axis must be 'flops' or 'bw', got {axis!r}") + base_flops = machine.peak_tflops_f16 + base_bw = machine.bw_hbm_gbs + points: list[AISensitivityPoint] = [] + for k in multipliers: + if axis == "flops": + m = replace(machine, peak_tflops_f16=base_flops * k) + else: + m = replace(machine, bw_hbm_gbs=base_bw * k) + points.append(AISensitivityPoint( + multiplier=k, + peak_tflops=m.peak_tflops_f16, + bw_gbs=m.bw_hbm_gbs, + ai=arithmetic_intensity(m), + b_star=critical_batch(m, model), + )) + return points diff --git a/tests/analytical_visualization/test_chip_roofline.py b/tests/analytical_visualization/test_chip_roofline.py index 1c5f8d3..443a0ef 100644 --- a/tests/analytical_visualization/test_chip_roofline.py +++ b/tests/analytical_visualization/test_chip_roofline.py @@ -6,6 +6,7 @@ import math import pytest from tests.analytical_visualization.chip_roofline import ( + ai_sensitivity_curve, arithmetic_intensity, balance_context, bound_regime, @@ -14,6 +15,8 @@ from tests.analytical_visualization.chip_roofline import ( good_context, knee_batch, kv_bytes_per_token, + memory_budget_curve_vs_batch, + memory_budget_curve_vs_skv, per_token_latency_curve, step_latency_curve, t_com, @@ -23,7 +26,7 @@ from tests.analytical_visualization.chip_roofline import ( utilization_at, ) from tests.analytical_visualization.model_config import ( - MachineParams, ModelConfig, + FullConfig, MachineParams, ModelConfig, TopologyConfig, ) from tests.analytical_visualization.model_presets import PRESETS @@ -335,6 +338,98 @@ def test_step_and_per_token_identical_at_b1(): assert sp.kv_s == pytest.approx(tp.kv_s) +# ── PE memory budget sweeps ──────────────────────────────────────── + + +def _budget_cfg(): + """Helper: real cfg with enough sharding to leave headroom.""" + return FullConfig( + model=PRESETS["Llama 3 8B"].model, + topo=TopologyConfig(cp=2, tp=8, pp=1, b=1, s_kv=8192), + machine=MachineParams(), + ) + + +def test_kv_grows_linear_with_skv(): + cfg = _budget_cfg() + pts = memory_budget_curve_vs_skv(cfg, s_kv_range=[1024, 2048, 8192], + batch=1) + assert pts[1].kv_gb == pytest.approx(2 * pts[0].kv_gb, rel=1e-3) + assert pts[2].kv_gb == pytest.approx(8 * pts[0].kv_gb, rel=1e-3) + + +def test_kv_grows_linear_with_batch(): + cfg = _budget_cfg() + pts = memory_budget_curve_vs_batch(cfg, b_range=[1, 2, 8], + s_kv=8192) + assert pts[1].kv_gb == pytest.approx(2 * pts[0].kv_gb, rel=1e-3) + assert pts[2].kv_gb == pytest.approx(8 * pts[0].kv_gb, rel=1e-3) + + +def test_weights_flat_across_skv_sweep(): + """Weight bytes are batch/S_kv-invariant — sharding fixed.""" + cfg = _budget_cfg() + pts = memory_budget_curve_vs_skv(cfg, s_kv_range=[128, 8192, 1_000_000], + batch=1) + ws = [p.weights_gb for p in pts] + assert all(w == pytest.approx(ws[0]) for w in ws), ws + + +def test_over_budget_flag_trips_past_hbm(): + """Push S_kv until KV blows the HBM budget; over_budget must flip.""" + cfg = _budget_cfg() + pts = memory_budget_curve_vs_skv( + cfg, s_kv_range=[1024, 10_000_000], batch=64, + ) + # 10M tokens × B=64 KV should overflow the 6 GB budget. + assert not pts[0].over_budget, pts[0] + assert pts[-1].over_budget, pts[-1] + + +def test_free_gb_never_negative(): + cfg = _budget_cfg() + pts = memory_budget_curve_vs_skv( + cfg, s_kv_range=[128, 10_000_000], batch=32, + ) + for p in pts: + assert p.free_gb >= 0, p + + +# ── AI sensitivity to FLOPs / BW ────────────────────────────────── + + +def test_ai_scales_linearly_with_flops(): + """Doubling FLOPs doubles AI.""" + m = MachineParams() + model = PRESETS["Llama 3 8B"].model + pts = ai_sensitivity_curve(m, model, [1.0, 2.0, 4.0], axis="flops") + assert pts[1].ai == pytest.approx(2 * pts[0].ai) + assert pts[2].ai == pytest.approx(4 * pts[0].ai) + + +def test_ai_scales_inversely_with_bw(): + """Doubling BW halves AI.""" + m = MachineParams() + model = PRESETS["Llama 3 8B"].model + pts = ai_sensitivity_curve(m, model, [1.0, 2.0], axis="bw") + assert pts[1].ai == pytest.approx(pts[0].ai / 2) + + +def test_ai_sensitivity_b_star_tracks_ai_for_bf16(): + """For BF16, B* == AI. Curve should mirror.""" + m = MachineParams() + model = PRESETS["Llama 3 8B"].model # BF16 + pts = ai_sensitivity_curve(m, model, [0.5, 1.0, 4.0], axis="flops") + for p in pts: + assert p.b_star == pytest.approx(p.ai) + + +def test_ai_sensitivity_bad_axis_raises(): + with pytest.raises(ValueError): + ai_sensitivity_curve(MachineParams(), PRESETS["Llama 3 8B"].model, + [1.0], axis="bogus") + + # ── App wiring: tab exists on the Streamlit app ────────────────────