From 221097ef08c89050945115a83a2624c5a625b49a Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Wed, 29 Jul 2026 13:40:37 -0700 Subject: [PATCH] =?UTF-8?q?analytical-viz:=20roofline=20tab=20=E2=80=94=20?= =?UTF-8?q?headline=20step-latency=20+=20cost-per-token=20plots,=20B=20kno?= =?UTF-8?q?b?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two new plots at the top of the Chip Roofline tab (right after the knob row), side by side: 1. **Latency per decode step** — one forward pass's total time as B grows. weight_fetch (flat in B), compute (linear), KV (linear), total. The SLO / TTFT view — bigger B = longer step. 2. **Cost per token** = step latency ÷ B. weight (shrinks 1/B), compute (flat), KV (flat), total. The efficiency / cost view — bigger B (up to ~2·B*) = cheaper per token. Same underlying decomposition, two divisors. Puts the classic throughput ↔ latency tradeoff in one glance. Also added a **B (batch size) slider** next to the existing S_kv slider. Both plots draw a vertical marker at the current B; the regime KPI ("memory-bound / compute-bound / kv-bound") now uses the slider B, so you can sweep B live and watch the label flip when you cross B*. step_latency_curve + StepLatencyPoint added to chip_roofline; 5 new tests cover: weight_s batch-invariant, compute_s/kv_s linear in B, step_total == per_token_total × B (definitional), and step == per_token at B=1. Co-Authored-By: Claude Opus 4.7 --- tests/analytical_visualization/app.py | 137 +++++++++++++++--- .../analytical_visualization/chip_roofline.py | 48 ++++++ .../test_chip_roofline.py | 54 +++++++ 3 files changed, 219 insertions(+), 20 deletions(-) diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 792cfa3..b70dcf9 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -61,6 +61,7 @@ from tests.analytical_visualization.chip_roofline import ( good_context, knee_batch, per_token_latency_curve, + step_latency_curve, t_com, t_mem_long, t_mem_short, @@ -2178,28 +2179,124 @@ with tab_roofline: _l_star = balance_context(_default_machine, model) _n_active = total_active_params(model) - # ── Local S_kv knob — override sidebar for exploring the plots - # without changing the rest of the app's config. - st.markdown("**Explore: override S_kv for the plots below**") - _skv_min = 128 - _skv_max = max(int(20 * _l_star), 4 * s_kv, 1_000_000) - _rf_skv = st.slider( - "S_kv (context length) for roofline plots", - min_value=_skv_min, max_value=_skv_max, - value=int(s_kv), - step=max(1, _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*."), - ) + # ── 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: + _skv_min = 128 + _skv_max = max(int(20 * _l_star), 4 * s_kv, 1_000_000) + _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), + 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)) + _rf_b = st.slider( + "B (batch size)", + min_value=_b_min, max_value=_b_max, + value=max(1, int(b_batch)), + key="_rf_b_override", + help=("Local to this tab. Drag to move the vertical B " + "marker on plots and see which regime dominates."), + ) st.caption( - f"Using S_kv = **{_rf_skv:,}** tokens for all plots on this tab. " - f"L* on your chip is **{_l_star:,.0f}** tokens → " - f"S_kv / L* = **{_rf_skv/_l_star:.2f}**." + 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}**." ) + # ── Headline plots: step latency and cost per token (=÷B) ───── + _b_range = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096] + _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, + s_kv=_rf_skv) + _xs_head = [p.batch for p in _step_pts] + + _hcol1, _hcol2 = st.columns(2) + + # Plot A — Latency per decode step (undivided by B) + with _hcol1: + st.markdown("**Latency per decode step** (one forward pass)") + _figA, _axA = plt.subplots(figsize=(6, 4.2)) + _axA.plot(_xs_head, [p.weight_s * 1e3 for p in _step_pts], + "^-", color="#ffbe0b", + label="Weight fetch (flat in B)") + _axA.plot(_xs_head, [p.compute_s * 1e3 for p in _step_pts], + "o-", color="#3a86ff", + label="Compute (linear ↑)") + _axA.plot(_xs_head, [p.kv_s * 1e3 for p in _step_pts], + "s-", color="#d90429", + label="KV fetch (linear ↑)") + _axA.plot(_xs_head, [p.total_s * 1e3 for p in _step_pts], + "-", color="#212529", linewidth=2.5, label="Total") + _axA.axvline(_rf_b, linestyle="-", color="#7b1fa2", alpha=0.6, + label=f"B = {_rf_b}") + _axA.set_xscale("log", base=2) + _axA.set_yscale("log") + _axA.set_xlabel("Batch size B") + _axA.set_ylabel("Step time (ms)") + _axA.set_title("Step latency = weight + compute·B + KV·B") + _axA.grid(True, which="both", alpha=0.3) + _axA.legend(fontsize=8, loc="upper left") + plt.tight_layout() + st.pyplot(_figA, width='stretch') + plt.close(_figA) + + # Plot B — Cost per token = latency ÷ B + with _hcol2: + st.markdown("**Cost per token** (= step latency ÷ B)") + _figB, _axB = plt.subplots(figsize=(6, 4.2)) + _axB.plot(_xs_head, [p.weight_s * 1e3 for p in _tok_pts], + "^-", color="#ffbe0b", + label="Weight fetch (shrinks 1/B)") + _axB.plot(_xs_head, [p.compute_s * 1e3 for p in _tok_pts], + "o--", color="#3a86ff", + label="Compute (flat)") + _axB.plot(_xs_head, [p.kv_s * 1e3 for p in _tok_pts], + "s--", color="#d90429", + label="KV fetch (flat)") + _axB.plot(_xs_head, [p.total_s * 1e3 for p in _tok_pts], + "-", color="#212529", linewidth=2.5, label="Total") + _axB.axvline(_b_star, linestyle=":", color="#2e7d32", + label=f"B* = {_b_star:.0f}") + _axB.axvline(_rf_b, linestyle="-", color="#7b1fa2", alpha=0.6, + label=f"B = {_rf_b}") + _axB.set_xscale("log", base=2) + _axB.set_yscale("log") + _axB.set_xlabel("Batch size B") + _axB.set_ylabel("Per-token time (ms)") + _axB.set_title("Cost/token = weight/B + compute + KV") + _axB.grid(True, which="both", alpha=0.3) + _axB.legend(fontsize=8, loc="upper right") + plt.tight_layout() + st.pyplot(_figB, width='stretch') + plt.close(_figB) + + st.caption( + "**Left (step latency)**: how long one forward pass takes. " + "Grows with B because compute and KV per-sequence grow linearly; " + "weight fetch is loaded once per step so it's flat. This is the " + "**SLO / TTFT view** — bigger B → longer step.\n\n" + "**Right (cost per token)**: step latency divided by B tokens " + "produced. Weight fetch amortizes (shrinks 1/B); compute and KV " + "per-sequence stay flat per token. This is the **efficiency / " + "$-per-token view** — bigger B (up to ~2·B*) → cheaper per token.\n\n" + "Same underlying decomposition, two divisors — the classic " + "throughput ↔ latency tradeoff." + ) + + st.divider() + _regime_now = bound_regime(_default_machine, model, - batch=max(1, b_batch), s_kv=_rf_skv) + batch=max(1, _rf_b), s_kv=_rf_skv) # ── KPI cards ───────────────────────────────────────────────── _r1, _r2, _r3, _r4 = st.columns(4) @@ -2213,10 +2310,10 @@ with tab_roofline: 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={_rf_skv:,})", + _r4.metric(f"Regime @ (B={_rf_b}, S_kv={_rf_skv:,})", _regime_now.replace("-bound", ""), help="Which term dominates cost at the current (B, S_kv) " - "you've selected.") + "you've selected via the sliders above.") # ── Recommended B and L (heuristics) ────────────────────────── _rec_b = good_batch(_default_machine, model, sparsity=1.0) diff --git a/tests/analytical_visualization/chip_roofline.py b/tests/analytical_visualization/chip_roofline.py index 984b1ba..3793772 100644 --- a/tests/analytical_visualization/chip_roofline.py +++ b/tests/analytical_visualization/chip_roofline.py @@ -256,3 +256,51 @@ def utilization_at(s_kv: int, l_star: float) -> float: At S_kv=0: 100%. At S_kv=L*: 50%. At 2·L*: 33.3%. At 5·L*: 16.7%. """ return 1.0 / (1.0 + s_kv / l_star) + + +# ── Per-step latency (undivided by B) ───────────────────────────── + + +@dataclass +class StepLatencyPoint: + batch: int + weight_s: float # N·b / W — flat in B (loaded once per step) + compute_s: float # 2·N·B / C — linear in B + kv_s: float # B · S_kv · kv_bpt / W — linear in B + total_s: float + + +def step_latency_curve(machine: MachineParams, model: ModelConfig, + batch_range: list[int], + s_kv: int) -> list[StepLatencyPoint]: + """Total time of one decode step (one forward pass), across a + range of batch sizes. NOT divided by B — this is the SLO view. + + step_weight = N·b / W (batch-invariant) + step_compute = 2·N·B / C (linear in B) + step_kv = B · S_kv · kv_bpt / W (linear in B) + + Per-token cost = step_total / B — the two views trade off: + bigger B lowers cost/token but raises step latency. + """ + n_active = total_active_params(model) + b = model.bytes_per_elem + weight_bytes = n_active * b + kv_bpt = kv_bytes_per_token(model) + + step_weight = weight_bytes / machine.bw_hbm + + points: list[StepLatencyPoint] = [] + for bs in batch_range: + B = max(1, bs) + step_compute = _FLOPS_PER_PARAM_PER_TOKEN * n_active * B / machine.peak_flops + step_kv = B * s_kv * kv_bpt / machine.bw_hbm + total = step_weight + step_compute + step_kv + points.append(StepLatencyPoint( + batch=bs, + weight_s=step_weight, + compute_s=step_compute, + kv_s=step_kv, + total_s=total, + )) + return points diff --git a/tests/analytical_visualization/test_chip_roofline.py b/tests/analytical_visualization/test_chip_roofline.py index 61236ff..1c5f8d3 100644 --- a/tests/analytical_visualization/test_chip_roofline.py +++ b/tests/analytical_visualization/test_chip_roofline.py @@ -15,6 +15,7 @@ from tests.analytical_visualization.chip_roofline import ( knee_batch, kv_bytes_per_token, per_token_latency_curve, + step_latency_curve, t_com, t_mem_long, t_mem_short, @@ -281,6 +282,59 @@ def test_utilization_drops_monotonically_with_context(): assert b < a +# ── Step latency (undivided by B) ────────────────────────────────── + + +def test_step_weight_is_batch_invariant(): + """step_weight = N·b/W — same for every B (loaded once per step).""" + m = MachineParams() + model = PRESETS["Llama 3 8B"].model + pts = step_latency_curve(m, model, [1, 8, 64, 512], s_kv=1024) + ws = [p.weight_s for p in pts] + assert all(w == pytest.approx(ws[0]) for w in ws), ws + + +def test_step_compute_linear_in_batch(): + """step_compute = 2·N·B/C — doubling B doubles compute.""" + m = MachineParams() + model = PRESETS["Llama 3 8B"].model + pts = step_latency_curve(m, model, [1, 2, 8], s_kv=1024) + assert pts[1].compute_s == pytest.approx(2 * pts[0].compute_s) + assert pts[2].compute_s == pytest.approx(8 * pts[0].compute_s) + + +def test_step_kv_linear_in_batch(): + """step_kv = B · S_kv · kv_bpt / W — linear in B.""" + m = MachineParams() + model = PRESETS["Llama 3 8B"].model + pts = step_latency_curve(m, model, [1, 4], s_kv=1024) + assert pts[1].kv_s == pytest.approx(4 * pts[0].kv_s) + + +def test_step_total_equals_per_token_times_batch(): + """The step and per-token views are just ÷ B / × B of each other.""" + m = MachineParams() + model = PRESETS["Llama 3 8B"].model + B_range = [1, 4, 16, 64] + s_kv = 1024 + step_pts = step_latency_curve(m, model, B_range, s_kv=s_kv) + tok_pts = per_token_latency_curve(m, model, B_range, s_kv=s_kv) + for sp, tp in zip(step_pts, tok_pts): + assert sp.total_s == pytest.approx(tp.total_s * sp.batch) + + +def test_step_and_per_token_identical_at_b1(): + """At B=1, step latency == per-token cost (dividing by 1).""" + m = MachineParams() + model = PRESETS["Llama 3 8B"].model + sp = step_latency_curve(m, model, [1], s_kv=1024)[0] + tp = per_token_latency_curve(m, model, [1], s_kv=1024)[0] + assert sp.total_s == pytest.approx(tp.total_s) + assert sp.weight_s == pytest.approx(tp.weight_s) + assert sp.compute_s == pytest.approx(tp.compute_s) + assert sp.kv_s == pytest.approx(tp.kv_s) + + # ── App wiring: tab exists on the Streamlit app ────────────────────