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>
mukesh
a1c3c28f62