ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

analytical-viz: roofline knobs — absolute TFLOPS / GB/s, drive whole tab

Browse Source 
Two changes:

1. **Knob units and ranges**
   - FLOPs slider: absolute TFLOPS, 1..32 (step 0.5). Was 0.25..8× multiplier.
   - HBM BW slider: absolute GB/s, 128..1024 (step 16). Was 0.25..8× multiplier.
   Defaults still initialize to the sidebar Hardware panel's base chip.
   Caption also shows the implied × multipliers of the base chip.

2. **Knobs now drive every roofline number on the tab.**
   Previously FLOPs/BW only affected the AI-sensitivity plot. Now the
   scaled machine feeds:
     - AI / B* / L* KPI cards
     - Regime KPI
     - Good-B / Good-L cards + utilization
     - Headline plots (step latency + cost/token)
     - Regime formulas table (t_com, t_mem_short, t_mem_long values)
     - Formula table (C, W, knee substitution)
   Left unchanged (intentionally):
     - PE memory budget stacks — depend on HBM capacity, not BW/FLOPs.
     - AI-sensitivity plot itself — sweeps around the SIDEBAR base
       chip so the curves are comparable across knob positions.

Header caption now shows both effective chip (from knobs) and base
chip (from sidebar) side by side.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-07-29 14:47:30 -07:00
parent bda76cbd66
commit 97e765b58f
1 changed files with 56 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/analytical_visualization/app.py
+56 -42
View File
@@ -2178,14 +2178,13 @@ with tab_roofline:
"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)
# ── Local knobs: S_kv, B, FLOPs multiplier, BW multiplier ─────
# ── Local knobs: S_kv, B, FLOPs, HBM BW ───────────────────────
# Override sidebar for exploring the plots without changing the
# rest of the app's config.
# rest of the app's config. FLOPs/BW knobs drive ALL roofline
# math and every plot/formula on this tab (memory-budget stacks
# excluded — those are HBM-capacity, not BW).
st.markdown("**Explore: local overrides for the plots below**")
_kn_a, _kn_b, _kn_c, _kn_d = st.columns(4)
with _kn_a:
@@ -2210,44 +2209,56 @@ with tab_roofline:
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*."),
_rf_flops_tflops = st.slider(
"FLOPs (TFLOPS)",
min_value=1.0, max_value=32.0,
value=float(_default_machine.peak_tflops_f16),
step=0.5,
key="_rf_flops_tflops",
help=("Chip peak TFLOPS for the AI-sensitivity plot. "
"Higher FLOPs → higher AI and higher 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*."),
_rf_bw_gbs = st.slider(
"HBM BW (GB/s)",
min_value=128.0, max_value=1024.0,
value=float(_default_machine.bw_hbm_gbs),
step=16.0,
key="_rf_bw_gbs",
help=("Chip HBM bandwidth for the AI-sensitivity plot. "
"Higher BW → lower AI (more bytes to feed per FLOP)."),
)
# Build a scaled machine for the AI-sensitivity view.
# Build the effective chip from the knobs; every roofline number
# on this tab is derived from it.
_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,
peak_tflops_f16=_rf_flops_tflops,
bw_hbm_gbs=_rf_bw_gbs,
)
_ai_scaled = arithmetic_intensity(_scaled_machine)
_b_star_scaled = critical_batch(_scaled_machine, model)
_ai = arithmetic_intensity(_scaled_machine)
_b_star = critical_batch(_scaled_machine, model)
_l_star = balance_context(_scaled_machine, model)
_flops_mult_now = _rf_flops_tflops / _default_machine.peak_tflops_f16
_bw_mult_now = _rf_bw_gbs / _default_machine.bw_hbm_gbs
_ai_base = arithmetic_intensity(_default_machine)
_b_star_base = critical_batch(_default_machine, model)
st.caption(
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}**."
f"**FLOPs = {_rf_flops_tflops:.1f} TFLOPS** "
f"(× {_flops_mult_now:.2f} of base), "
f"**BW = {_rf_bw_gbs:.0f} GB/s** "
f"(× {_bw_mult_now:.2f} of base). "
f"Effective **AI = {_ai:.1f}** (base {_ai_base:.1f}), "
f"**B\\* = {_b_star:.0f}** (base {_b_star_base:.0f}), "
f"**L\\* = {_l_star:,.0f}** tokens → S_kv/L\\* = "
f"**{_rf_skv/_l_star:.2f}**."
)
# ── Headline plots: step latency and cost per token (=÷B) ─────
_b_range = [1, 2, 4, 8, 16, 32, 64, 128, 256]
_step_pts = step_latency_curve(_default_machine, model, _b_range,
_step_pts = step_latency_curve(_scaled_machine, model, _b_range,
s_kv=_rf_skv)
_tok_pts = per_token_latency_curve(_default_machine, model, _b_range,
_tok_pts = per_token_latency_curve(_scaled_machine, model, _b_range,
s_kv=_rf_skv)
_xs_head = [p.batch for p in _step_pts]
@@ -2463,7 +2474,7 @@ with tab_roofline:
st.divider()
_regime_now = bound_regime(_default_machine, model,
_regime_now = bound_regime(_scaled_machine, model,
batch=max(1, _rf_b), s_kv=_rf_skv)
# ── KPI cards ─────────────────────────────────────────────────
@@ -2484,8 +2495,8 @@ with tab_roofline:
"you've selected via the sliders above.")
# ── Recommended B and L (heuristics) ──────────────────────────
_rec_b = good_batch(_default_machine, model, sparsity=1.0)
_rec_l = good_context(_default_machine, model, s_kv=_rf_skv)
_rec_b = good_batch(_scaled_machine, model, sparsity=1.0)
_rec_l = good_context(_scaled_machine, model, s_kv=_rf_skv)
_g1, _g2, _g3 = st.columns(3)
_g1.metric("Good B (dense)", f"{_rec_b.effective:.0f}",
help=_rec_b.reason)
@@ -2508,19 +2519,22 @@ with tab_roofline:
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."
f"**Effective chip (from knobs):** "
f"{_scaled_machine.peak_tflops_f16:.1f} TFLOPs BF16, "
f"{_scaled_machine.bw_hbm_gbs:.0f} GB/s HBM. "
f"Base chip from sidebar Hardware: "
f"{_default_machine.peak_tflops_f16:.0f} TFLOPs, "
f"{_default_machine.bw_hbm_gbs:.0f} GB/s."
)
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
_t_mem_s_at_bstar = t_mem_short(_default_machine, model,
_t_com_val = t_com(_scaled_machine, model) * 1e3 # ms
_t_mem_s_at_bstar = t_mem_short(_scaled_machine, model,
int(round(_b_star))) * 1e3
_t_mem_l_at_skv = t_mem_long(_default_machine, model, _rf_skv) * 1e3
_t_mem_l_at_skv = t_mem_long(_scaled_machine, model, _rf_skv) * 1e3
_regime_rows = [
{"Term": "t_com (compute)",
"Short context (S_kv < L*)": "2 · N / C",
@@ -2571,12 +2585,12 @@ with tab_roofline:
# ── Interpretation table ───────────────────────────────────────
st.markdown("**Formulas & interpretation**")
_C = _default_machine.peak_flops # FLOPs / s
_W = _default_machine.bw_hbm # bytes / s
_C = _scaled_machine.peak_flops # FLOPs / s (from knobs)
_W = _scaled_machine.bw_hbm # bytes / s (from knobs)
_b_elem = model.bytes_per_elem # bytes / weight
_kv_bpt_full = 2 * model.h_kv * model.d_head * model.bytes_per_elem * model.layers
_knee_now = knee_batch(_default_machine, model, _rf_skv)
_knee_now = knee_batch(_scaled_machine, model, _rf_skv)
if _knee_now is None:
_sub_bknee = (
f"{_b_star:.0f} / (1 − {_rf_skv:,}/{_l_star:,.0f}) = "