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