From 9b4aee83da2bc1a601d86315b1c30f0c63eb49b9 Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Wed, 29 Jul 2026 13:52:40 -0700 Subject: [PATCH] =?UTF-8?q?analytical-viz:=20roofline=20tab=20=E2=80=94=20?= =?UTF-8?q?trim=20to=204=20plots,=20cap=20B=20at=20256=20/=20S=5Fkv=20at?= =?UTF-8?q?=201M?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- tests/analytical_visualization/app.py | 108 ++++---------------------- 1 file changed, 17 insertions(+), 91 deletions(-) diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index b70dcf9..322f941 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -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