analytical-viz: roofline tab — headline step-latency + cost-per-token plots, B knob

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 <noreply@anthropic.com>

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)