analytical-viz: 'Chip roofline & B*' tab — AI, B*, L*, cost curves

New tab surfaces the arithmetic-intensity story from LLM-serving
practice for the current sidebar chip + model. All derived from
MachineParams + ModelConfig; no new configuration.

chip_roofline module (pure functions):
- arithmetic_intensity      = C / W   (FLOPs per byte HBM BW)
- critical_batch (B*)       = C * b / (2 * W) * sparsity
- balance_context (L*)      = 2 * N_active * W / (C * kv_bpt)
- knee_batch (B_knee)       = B* / (1 - S_kv/L*)  (None past L*)
- per_token_latency_curve   = weight/B + compute + KV, per B
- bound_regime              = which term dominates now

Peak roofline convention (no compute_util factor) so weight_s ==
compute_s exactly at B*. Comm + TP/CP sharding intentionally
excluded — stage_latencies is the full latency model; this is the
back-of-envelope chip-vs-model view.

Tab shows:
- 4 KPI cards: AI, B*, L*, current regime at (B, S_kv)
- MoE hint when preset flags MoE ('300 x sparsity' rule)
- Plot 1: cost vs B at current S_kv (weight, compute floor, KV
  floor, total) with B* marker line
- Plot 2: cost curves at S_kv/L* = 0.25, 0.5, 1, 2, 5 — shows the
  no-knee regime past L*
- Plot 3: B_knee vs S_kv — knee slides right, diverges at L*
- Formulas + interpretation table

test_chip_roofline covers B* against H100 reference (295), sparsity
scaling, L* scaling with HBM BW, knee divergence at L*, per-token
curve monotonicity + asymptote to compute+KV floor, regime
classification. Smoke test guards the tab is wired in app.py.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

tests/analytical_visualization/chip_roofline.py

@@ -0,0 +1,166 @@
+"""Chip-level roofline math: AI, B*, L*, per-token latency curves.
+
+Surfaces the arithmetic-intensity story from LLM-serving practice:
+
+- **AI**  = C / W  (peak FLOPs per byte of HBM bandwidth).
+- **B\***  = C * b / (2 * W) * sparsity  — the critical batch size at
+  which weight-fetch time equals compute time for one decode step.
+  Sparsity = N_total / N_active (MoE factor; 1 for dense).
+- **L\***  = 2 * N_active / (AI * kv_bytes_per_token) — the balance
+  context length at which KV-read time equals compute time.
+- **B_knee(S_kv)** = B* / (1 - S_kv/L*) — the batch size where the
+  cost curve bends (weight-fetch drops below the compute+KV floor).
+  Diverges at S_kv = L* and no knee exists past it.
+
+Per-token decode-step latency, per PE (dense-approx, no comm):
+
+    t(B, S_kv)  =  N_active * b / (W * B)     <-- weight fetch, 1/B
+                 + 2 * N_active / C            <-- compute (peak), flat
+                 + S_kv * kv_bpt / W           <-- KV read, flat
+
+All numbers per PE / per one forward pass. **Peak roofline — no
+utilization factor.** Comm cost and TP/CP sharding are intentionally
+NOT in the roofline — this is the back-of-envelope chip-vs-model view
+the transcript talks about, not the full latency model that
+stage_latencies.py builds. With this convention weight_s == compute_s
+exactly at B*.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .model_config import MachineParams, ModelConfig
+
+
+# Per-token bytes of BF16 MAC arithmetic: one multiply + one add = 2 FLOPs.
+_FLOPS_PER_PARAM_PER_TOKEN = 2
+
+
+# ── Chip / model derived quantities ────────────────────────────────
+
+
+def arithmetic_intensity(machine: MachineParams) -> float:
+    """FLOPs per byte of HBM bandwidth. Peak roofline; utilization
+    is applied only in the compute-time formula, not here."""
+    return machine.peak_flops / machine.bw_hbm
+
+
+def total_active_params(model: ModelConfig) -> int:
+    """Full-model parameter count (attention + FFN, all layers).
+
+    Attention: 4 projections × hidden × H_q * d_head effective per layer.
+    (W_Q hidden×H_q*d_h, W_O H_q*d_h×hidden, W_K/W_V hidden×H_kv*d_h.)
+    FFN: 3 × hidden × ffn_dim per layer (gate, up, down).
+    """
+    m = model
+    attn = (
+        m.hidden * m.h_q * m.d_head            # W_Q
+        + m.hidden * m.h_kv * m.d_head * 2     # W_K + W_V
+        + m.h_q * m.d_head * m.hidden          # W_O
+    )
+    ffn = 3 * m.hidden * m.ffn_dim
+    return (attn + ffn) * m.layers
+
+
+def kv_bytes_per_token(model: ModelConfig) -> int:
+    """Bytes of KV cache one new token adds across ALL layers, per one
+    sequence, un-sharded (K + V, H_kv heads * d_h * bytes)."""
+    m = model
+    return 2 * m.h_kv * m.d_head * m.bytes_per_elem * m.layers
+
+
+def critical_batch(machine: MachineParams, model: ModelConfig,
+                    sparsity: float = 1.0) -> float:
+    """B* = C * b / (2 * W) * sparsity.
+
+    Sparsity = N_total / N_active (>= 1). Dense = 1. MoE 8-of-256 = 8.
+    """
+    b = model.bytes_per_elem
+    ai = arithmetic_intensity(machine)
+    return ai * b / _FLOPS_PER_PARAM_PER_TOKEN * sparsity
+
+
+def balance_context(machine: MachineParams, model: ModelConfig) -> float:
+    """L* = 2 * N_active / (AI * kv_bpt).
+
+    Context length (in tokens) at which per-step KV read matches the
+    per-step compute cost. Beyond L*, the KV term is dominant and no
+    batch size gets you compute-bound.
+    """
+    n_active = total_active_params(model)
+    ai = arithmetic_intensity(machine)
+    kv_bpt = kv_bytes_per_token(model)
+    return _FLOPS_PER_PARAM_PER_TOKEN * n_active / (ai * kv_bpt)
+
+
+def knee_batch(machine: MachineParams, model: ModelConfig,
+                s_kv: int) -> float | None:
+    """B_knee(S_kv) = B* / (1 - S_kv/L*).
+
+    Returns None when S_kv >= L* (no knee exists — the total-cost
+    curve never touches the compute floor).
+    """
+    b_star = critical_batch(machine, model)
+    l_star = balance_context(machine, model)
+    r = s_kv / l_star
+    if r >= 1.0:
+        return None
+    return b_star / (1.0 - r)
+
+
+# ── Per-token latency curves ───────────────────────────────────────
+
+
+@dataclass
+class RooflinePoint:
+    batch: int
+    weight_s: float   # weight fetch time, 1/B
+    compute_s: float  # compute time, flat
+    kv_s: float       # KV read time, flat
+    total_s: float
+
+
+def per_token_latency_curve(machine: MachineParams, model: ModelConfig,
+                             batch_range: list[int],
+                             s_kv: int) -> list[RooflinePoint]:
+    """Per-token decode-step latency curve across a range of batch sizes.
+
+    Returns one point per batch. All times per PE, dense-approx,
+    utilization from machine.compute_util. Comm and TP/CP sharding are
+    excluded — this is the roofline model.
+    """
+    n_active = total_active_params(model)
+    b = model.bytes_per_elem
+    weight_bytes = n_active * b
+    compute_flops = _FLOPS_PER_PARAM_PER_TOKEN * n_active
+    kv_read_bytes = s_kv * kv_bytes_per_token(model)
+
+    compute_s = compute_flops / machine.peak_flops
+    kv_s = kv_read_bytes / machine.bw_hbm
+
+    points: list[RooflinePoint] = []
+    for bs in batch_range:
+        weight_s = weight_bytes / machine.bw_hbm / max(1, bs)
+        points.append(RooflinePoint(
+            batch=bs,
+            weight_s=weight_s,
+            compute_s=compute_s,
+            kv_s=kv_s,
+            total_s=weight_s + compute_s + kv_s,
+        ))
+    return points
+
+
+def bound_regime(machine: MachineParams, model: ModelConfig,
+                  batch: int, s_kv: int) -> str:
+    """Which term dominates at the current (batch, S_kv) point.
+
+    Returns 'memory-bound' if weight_fetch is the largest term,
+    'kv-bound' if KV read is largest, 'compute-bound' if compute.
+    """
+    pts = per_token_latency_curve(machine, model, [batch], s_kv)
+    p = pts[0]
+    parts = {"memory-bound": p.weight_s,
+             "kv-bound": p.kv_s,
+             "compute-bound": p.compute_s}
+    return max(parts, key=parts.get)