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

analytical-viz: roofline tab — trim to 4 plots, cap B at 256 / S_kv at 1M

Browse Source 
- Slider caps: B_max = 256, S_kv_max = 1M (matches practical decode
  operating ranges).
- Batch axis on every plot capped at 256 to match the slider.
- Removed the two redundant plots (old Plot 1 and old Plot 4). Both
  were re-drawing the same weight/compute/KV decomposition as the new
  headline 'Cost per token' plot at the top — the only distinguishing
  bits (extra 'Memory total' curve, B* marker) have been folded into
  the headline plot already, so keeping them was pure clutter.
- Every remaining plot now reacts to the S_kv / B knobs:
    * Plot A (step latency): S_kv drives curves, B marker drawn.
    * Plot B (cost/token): same.
    * Plot 2 (no-knee overlay): S_kv drives the L* ratios shown,
      B marker drawn.
    * Plot 3 (knee vs context): S_kv AND B markers drawn on both axes.

Net: 4 plots (headline pair + no-knee + knee-vs-context), all live-
interactive against both sliders.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-07-29 13:52:40 -07:00
parent 221097ef08
commit 9b4aee83da
1 changed files with 17 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/analytical_visualization/app.py
+17 -91
View File
@@ -2185,23 +2185,23 @@ with tab_roofline:
_kn_l, _kn_r = st.columns(2)
with _kn_l:
_skv_min = 128
_skv_max = max(int(20 * _l_star), 4 * s_kv, 1_000_000)
_skv_max = 1_000_000 # 1M tokens — the far end of interest
_rf_skv = st.slider(
"S_kv (context length)",
min_value=_skv_min, max_value=_skv_max,
value=int(s_kv),
step=max(1, _skv_min),
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*."),
)
with _kn_r:
_b_min = 1
_b_max = max(int(8 * _b_star), 1024, 4 * max(1, b_batch))
_b_max = 256 # small scale — decode B rarely goes higher in practice
_rf_b = st.slider(
"B (batch size)",
min_value=_b_min, max_value=_b_max,
value=max(1, int(b_batch)),
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."),
@@ -2213,7 +2213,7 @@ with tab_roofline:
)
# ── Headline plots: step latency and cost per token (=÷B) ─────
_b_range = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
_b_range = [1, 2, 4, 8, 16, 32, 64, 128, 256]
_step_pts = step_latency_curve(_default_machine, model, _b_range,
s_kv=_rf_skv)
_tok_pts = per_token_latency_curve(_default_machine, model, _b_range,
@@ -2347,45 +2347,10 @@ with tab_roofline:
st.divider()
# ── Plot 1: cost vs B at current S_kv ─────────────────────────
st.markdown(f"**Per-token latency vs batch size** (S_kv = {_rf_skv:,})")
_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=_rf_skv)
_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 = {_rf_skv:,} 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 ───────────
# ── 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"]
@@ -2394,12 +2359,14 @@ with tab_roofline:
_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], "-",
_ax2.plot(_xs_head, [p.total_s * 1e3 for p in _pts_r], "-",
color=_col, label=_label, linewidth=1.6)
_ax2.axhline(_pts[0].compute_s * 1e3, linestyle="--",
_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")
@@ -2435,6 +2402,10 @@ with tab_roofline:
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")
@@ -2454,51 +2425,6 @@ with tab_roofline:
st.divider()
# ── Plot 4: decode-step latency decomposition ─────────────────
st.markdown(f"**Latency per decode step — decomposition** "
f"(S_kv = {_rf_skv:,})")
_fig4, _ax4 = plt.subplots(figsize=(9, 4.5))
_weight_vals = [p.weight_s * 1e3 for p in _pts]
_compute_vals = [p.compute_s * 1e3 for p in _pts]
_kv_vals = [p.kv_s * 1e3 for p in _pts]
_mem_total = [w + k for w, k in zip(_weight_vals, _kv_vals)]
_total_vals = [p.total_s * 1e3 for p in _pts]
_ax4.plot(_xs, _compute_vals, "--", color="#3a86ff",
label="Compute", linewidth=1.6)
_ax4.plot(_xs, _weight_vals, "^-", color="#ffbe0b",
label="Weight fetch (shrinks 1/B)", linewidth=1.6)
_ax4.plot(_xs, _kv_vals, "s-", color="#d90429",
label="KV fetch (flat)", linewidth=1.6)
_ax4.plot(_xs, _mem_total, "-", color="#7b1fa2",
label="Memory total (weights + KV)", linewidth=2)
_ax4.plot(_xs, _total_vals, "-", color="#212529",
label="Total (compute + memory)", linewidth=2.5)
_ax4.axvline(_b_star, linestyle=":", color="#2e7d32",
label=f"B* = {_b_star:.0f}")
_ax4.set_xscale("log", base=2)
_ax4.set_yscale("log")
_ax4.set_xlabel("Batch size (sequences in flight)")
_ax4.set_ylabel("Per-token step time (ms)")
_ax4.set_title("Where the time goes: compute vs weight fetch vs KV fetch")
_ax4.grid(True, which="both", alpha=0.3)
_ax4.legend(fontsize=8, loc="upper right")
plt.tight_layout()
st.pyplot(_fig4, width='stretch')
plt.close(_fig4)
st.caption(
"**Compute** is flat in B. **Weight fetch** shrinks like 1/B "
"(the amortization win). **KV fetch** is flat in B — each "
"sequence reads its own KV, so batching does nothing for it. "
"**Memory total** is weight+KV summed. Which of compute or "
"memory-total is bigger at your operating B tells you the "
"regime; if memory-total's floor (= KV alone at large B) "
"already sits above compute, you're past L*."
)
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