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>
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user