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)

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

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 ────────────────────