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