analytical-viz: 'Chip roofline & B*' tab — AI, B*, L*, cost curves
New tab surfaces the arithmetic-intensity story from LLM-serving
practice for the current sidebar chip + model. All derived from
MachineParams + ModelConfig; no new configuration.
chip_roofline module (pure functions):
- arithmetic_intensity = C / W (FLOPs per byte HBM BW)
- critical_batch (B*) = C * b / (2 * W) * sparsity
- balance_context (L*) = 2 * N_active * W / (C * kv_bpt)
- knee_batch (B_knee) = B* / (1 - S_kv/L*) (None past L*)
- per_token_latency_curve = weight/B + compute + KV, per B
- bound_regime = which term dominates now
Peak roofline convention (no compute_util factor) so weight_s ==
compute_s exactly at B*. Comm + TP/CP sharding intentionally
excluded — stage_latencies is the full latency model; this is the
back-of-envelope chip-vs-model view.
Tab shows:
- 4 KPI cards: AI, B*, L*, current regime at (B, S_kv)
- MoE hint when preset flags MoE ('300 x sparsity' rule)
- Plot 1: cost vs B at current S_kv (weight, compute floor, KV
floor, total) with B* marker line
- Plot 2: cost curves at S_kv/L* = 0.25, 0.5, 1, 2, 5 — shows the
no-knee regime past L*
- Plot 3: B_knee vs S_kv — knee slides right, diverges at L*
- Formulas + interpretation table
test_chip_roofline covers B* against H100 reference (295), sparsity
scaling, L* scaling with HBM BW, knee divergence at L*, per-token
curve monotonicity + asymptote to compute+KV floor, regime
classification. Smoke test guards the tab is wired in app.py.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -30,6 +30,7 @@ for _mod_name in (
|
||||
"tests.analytical_visualization.memory_layout",
|
||||
"tests.analytical_visualization.stage_latencies",
|
||||
"tests.analytical_visualization.stage_shapes",
|
||||
"tests.analytical_visualization.chip_roofline",
|
||||
"tests.analytical_visualization.auto_explore",
|
||||
"tests.analytical_visualization.auto_hardware",
|
||||
"tests.analytical_visualization.pe_weight_layout",
|
||||
@@ -51,6 +52,15 @@ from tests.analytical_visualization.stage_shapes import (
|
||||
attn_stage_shape_rows,
|
||||
ffn_stage_shape_rows,
|
||||
)
|
||||
from tests.analytical_visualization.chip_roofline import (
|
||||
arithmetic_intensity,
|
||||
balance_context,
|
||||
bound_regime,
|
||||
critical_batch,
|
||||
knee_batch,
|
||||
per_token_latency_curve,
|
||||
total_active_params,
|
||||
)
|
||||
from tests.analytical_visualization.memory_layout import (
|
||||
compute_memory,
|
||||
attention_weight_rows,
|
||||
@@ -531,11 +541,13 @@ if _warnings:
|
||||
# Attn + FFN/MoE) so the user can toggle scope without switching tabs.
|
||||
# Auto Suggest is renamed "Parallelism" since it only varies parallelism
|
||||
# knobs (hardware is held fixed at the sidebar values).
|
||||
tab_layout, tab_auto, tab_memory, tab_stages, tab_compare, tab_hw = st.tabs([
|
||||
(tab_layout, tab_auto, tab_memory, tab_stages, tab_compare, tab_hw,
|
||||
tab_roofline) = st.tabs([
|
||||
"Physical layout",
|
||||
"Auto Suggest Parallelism",
|
||||
"Memory breakdown", "Per-stage latency",
|
||||
"Save & compare", "Auto Hardware",
|
||||
"Chip roofline & B*",
|
||||
])
|
||||
|
||||
|
||||
@@ -2142,3 +2154,199 @@ def _render_auto_hardware_tab():
|
||||
|
||||
with tab_hw:
|
||||
_render_auto_hardware_tab()
|
||||
|
||||
|
||||
# ── TAB 7: Chip roofline & B* ────────────────────────────────────
|
||||
with tab_roofline:
|
||||
st.subheader("Chip roofline & critical batch size (B*)")
|
||||
st.caption(
|
||||
"Back-of-envelope answer to 'is my chip well-matched to this "
|
||||
"model?' — see how big the batch must be before decode becomes "
|
||||
"compute-bound, and how big the context can grow before the KV "
|
||||
"bandwidth wall kills your utilization. All numbers per PE, "
|
||||
"peak roofline (no compute-util factor), no comm."
|
||||
)
|
||||
|
||||
_ai = arithmetic_intensity(_default_machine)
|
||||
_b_star = critical_batch(_default_machine, model)
|
||||
_l_star = balance_context(_default_machine, model)
|
||||
_n_active = total_active_params(model)
|
||||
_regime_now = bound_regime(_default_machine, model,
|
||||
batch=max(1, b_batch), s_kv=s_kv)
|
||||
|
||||
# ── KPI cards ─────────────────────────────────────────────────
|
||||
_r1, _r2, _r3, _r4 = st.columns(4)
|
||||
_r1.metric("AI (FLOPs/byte)", f"{_ai:.1f}",
|
||||
help="C/W. Peak FLOPs per byte of HBM bandwidth. "
|
||||
"Chip-only — no model-dependent factor.")
|
||||
_r2.metric("B* (dense)", f"{_b_star:.0f}",
|
||||
help="C*b/(2*W). Batch size where weight-fetch time "
|
||||
"equals compute time. Below this you're memory-bound.")
|
||||
_r3.metric("L* (tokens)", f"{_l_star:,.0f}",
|
||||
help="Context length where KV-read time equals compute "
|
||||
"time. Above this, no batch size gets you back to "
|
||||
"compute-bound.")
|
||||
_r4.metric(f"Regime @ (B={b_batch}, S_kv={s_kv:,})",
|
||||
_regime_now.replace("-bound", ""),
|
||||
help="Which term dominates cost at your current sidebar "
|
||||
"settings.")
|
||||
|
||||
# Optional MoE B*: only surface if the preset flags MoE.
|
||||
_is_moe = "MoE" in (preset.family + preset.note)
|
||||
if _is_moe:
|
||||
st.info(
|
||||
"This preset is dense-approximated as an MoE. Reiner Pope's "
|
||||
"'300 × sparsity' rule: for real sparsity k = N_total/N_active, "
|
||||
f"B*_moe ≈ {_b_star:.0f} × k. E.g. DeepSeek 32-of-256 experts "
|
||||
f"(k=8) → B* ≈ {_b_star*8:.0f}."
|
||||
)
|
||||
|
||||
st.caption(
|
||||
f"**Full-model attention+FFN params:** {_n_active/1e9:.2f} B. "
|
||||
f"**Chip:** {_default_machine.peak_tflops_f16:.0f} TFLOPs BF16, "
|
||||
f"{_default_machine.bw_hbm_gbs:.0f} GB/s HBM. "
|
||||
f"Change these in the sidebar Hardware panel to see the roofline shift."
|
||||
)
|
||||
|
||||
st.divider()
|
||||
|
||||
# ── Plot 1: cost vs B at current S_kv ─────────────────────────
|
||||
st.markdown("**Per-token latency vs batch size** (current S_kv = "
|
||||
f"{s_kv:,})")
|
||||
_b_range = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
|
||||
_pts = per_token_latency_curve(_default_machine, model, _b_range,
|
||||
s_kv=s_kv)
|
||||
_fig1, _ax1 = plt.subplots(figsize=(9, 4.5))
|
||||
_xs = [p.batch for p in _pts]
|
||||
_ax1.plot(_xs, [p.weight_s * 1e3 for p in _pts],
|
||||
"o-", label="Weight fetch (1/B)", color="#ffbe0b")
|
||||
_ax1.axhline(_pts[0].compute_s * 1e3, linestyle="--",
|
||||
color="#3a86ff", label="Compute floor (peak)")
|
||||
_ax1.axhline(_pts[0].kv_s * 1e3, linestyle=":",
|
||||
color="#d90429", label="KV read floor")
|
||||
_ax1.plot(_xs, [p.total_s * 1e3 for p in _pts],
|
||||
"-", color="#212529", linewidth=2, label="Total")
|
||||
_ax1.axvline(_b_star, linestyle=":", color="#2e7d32",
|
||||
label=f"B* = {_b_star:.0f}")
|
||||
_ax1.set_xscale("log", base=2)
|
||||
_ax1.set_yscale("log")
|
||||
_ax1.set_xlabel("Batch size (sequences in flight)")
|
||||
_ax1.set_ylabel("Per-token step time (ms)")
|
||||
_ax1.set_title(f"Cost curve at S_kv = {s_kv:,} tokens")
|
||||
_ax1.grid(True, which="both", alpha=0.3)
|
||||
_ax1.legend(fontsize=8, loc="upper right")
|
||||
plt.tight_layout()
|
||||
st.pyplot(_fig1, width='stretch')
|
||||
plt.close(_fig1)
|
||||
|
||||
st.caption(
|
||||
"**Below B\\*:** cost is dominated by weight fetch; doubling B "
|
||||
"halves cost per token. **At B\\*:** weight fetch = compute. "
|
||||
"**Above 2-3× B\\*:** you've captured most amortization; further "
|
||||
"batching mostly buys latency you don't want."
|
||||
)
|
||||
|
||||
st.divider()
|
||||
|
||||
# ── Plot 2: cost curves at multiple context lengths ───────────
|
||||
st.markdown("**The 'no-knee' phenomenon — cost vs B at different S_kv**")
|
||||
_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, [p.total_s * 1e3 for p in _pts_r], "-",
|
||||
color=_col, label=_label, linewidth=1.6)
|
||||
_ax2.axhline(_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.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.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()
|
||||
|
||||
# ── Interpretation table ───────────────────────────────────────
|
||||
st.markdown("**Formulas & interpretation**")
|
||||
_formula_rows = [
|
||||
{"Symbol": "AI",
|
||||
"Formula": "C / W",
|
||||
"Meaning": "FLOPs per byte of HBM bandwidth. Chip-only.",
|
||||
"Your chip": f"{_ai:.1f}"},
|
||||
{"Symbol": "B*",
|
||||
"Formula": "C · b / (2 · W)",
|
||||
"Meaning": ("Batch size where weight-fetch time = compute time. "
|
||||
"Below: memory-bound. Above: amortized."),
|
||||
"Your chip": f"{_b_star:.0f}"},
|
||||
{"Symbol": "B*_moe",
|
||||
"Formula": "B* · (N_total / N_active)",
|
||||
"Meaning": "MoE fetches all experts but computes on active only.",
|
||||
"Your chip": ("N/A (dense preset)" if not _is_moe
|
||||
else f"{_b_star * 8:.0f} (k=8 example)")},
|
||||
{"Symbol": "L*",
|
||||
"Formula": "2 · N_active · W / (C · kv_bpt)",
|
||||
"Meaning": ("Context where KV-read time = compute time. "
|
||||
"Above: KV-bound regardless of B."),
|
||||
"Your chip": f"{_l_star:,.0f} tok"},
|
||||
{"Symbol": "B_knee",
|
||||
"Formula": "B* / (1 − S_kv/L*)",
|
||||
"Meaning": ("Batch size where cost curve bends onto its floor. "
|
||||
"Diverges at S_kv = L*."),
|
||||
"Your chip": (f"{knee_batch(_default_machine, model, s_kv):.0f}"
|
||||
if knee_batch(_default_machine, model, s_kv) is not None
|
||||
else "no knee (S_kv >= L*)")},
|
||||
]
|
||||
st.dataframe(pd.DataFrame(_formula_rows),
|
||||
width='stretch', hide_index=True)
|
||||
|
||||
Reference in New Issue
Block a user