analytical-viz: roofline tab — PE memory budget, AI sensitivity, FLOPs/BW knobs

Three additions to the Chip Roofline tab:

1. **PE memory budget** section (two side-by-side stacked-area plots)
   - Left: per-PE memory vs S_kv (128..1M, log). Stacks: weights,
     KV, transient. HBM budget line + current-S_kv marker.
   - Right: per-PE memory vs B (1..256). Same stacks, current-B
     marker. KV grows linearly with either axis; weights are
     invariant of B and S_kv (fixed by CP·TP·PP sharding).
   Uses actual sharded per-PE quantities from memory_layout so the
   numbers match what the deployment would really allocate.

2. **AI sensitivity** section (two side-by-side line plots)
   - Left: AI vs chip-parameter multiplier. Two curves: scale FLOPs
     (AI grows linearly), scale HBM BW (AI shrinks inversely).
   - Right: B* vs multiplier. Same shape as AI (for BF16 where B*=AI).
   Vertical markers show current FLOPs/BW knob positions.

3. **Two new knobs** in the top row: FLOPs × multiplier and HBM BW ×
   multiplier (0.25..8, step 0.25). The AI-sensitivity plot markers
   move with them; the header caption shows scaled AI + B* for the
   selected point.

Also: formula annotation on the headline 'Latency per decode step'
plot header showing t_step = N·b/W + 2·N·B/C + B·S_kv·kv_bpt/W.

Pure additions in chip_roofline.py (all testable):
- MemoryBudgetPoint dataclass
- memory_budget_curve_vs_skv, memory_budget_curve_vs_batch
- AISensitivityPoint dataclass
- ai_sensitivity_curve(machine, model, mults, axis='flops'|'bw')

9 new tests cover linear scaling of KV with S_kv/B, weight
invariance across sweeps, over-budget flag flip, free_gb clamp,
FLOPs multiplier → AI linear, BW multiplier → AI inverse, B*
tracks AI for BF16, bad axis raises.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

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
+    # ── Local knobs: S_kv, B, FLOPs multiplier, BW multiplier ─────
-    # plots without changing the rest of the app's config.
+    # Override sidebar for exploring the plots without changing the
-    st.markdown("**Explore: override S_kv and B for the plots below**")
+    # rest of the app's config.
-    _kn_l, _kn_r = st.columns(2)
+    st.markdown("**Explore: local overrides for the plots below**")
-    with _kn_l:
+    _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 "
+            help="Local to this tab. Drives every plot's S_kv.",
-                  "and where the knee disappears past L*."),
         )
-    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 "
+            help="Local to this tab. Marker on plots + PE-memory KV.",
-                  "marker on plots and see which regime dominates."),
         )
+    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"**S_kv = {_rf_skv:,}**, **B = {_rf_b}**, "
-        f"L* = **{_l_star:,.0f}** tokens → S_kv/L* = **{_rf_skv/_l_star:.2f}**. "
+        f"**FLOPs × {_rf_flops_mult:.2f}**, **BW × {_rf_bw_mult:.2f}**. "
-        f"B* = **{_b_star:.0f}** → B/B* = **{_rf_b/_b_star:.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

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

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 ────────────────────