analytical-viz: capacity planning — symbol legend + plain-English formulas

Section 1 rewritten so a layman can read it without a glossary:

1. **Symbol legend table** listing every letter that appears in the
   three-axis formulas (N, b, N·b, HBM_per_PE, S_kv, kv_bpt,
   users_per_replica, TPOT SLO, step_latency, B_at_SLO, n_users,
   N_replicas, ⌈ ⌉). Each row: Symbol / Plain English / Value now
   for the current sidebar model + chip. For example:
     - N = 'Total number of model parameters (attention weights +
       FFN weights, across ALL layers)' → 6.98 billion params for
       Llama 3 8B
     - kv_bpt = 'KV cache bytes per token per user, summed across all
       layers (= 2·H_kv·d_head·b·layers)' → 128 KB/token

2. **Three-axis table** widened from (Formula / Meaning / Grows with)
   to (In symbols / In plain English / What it says / Value now).
   Axis A now shows the actual substituted arithmetic:
   ⌈ 13.96 GB / 6.0 GB ⌉ = ⌈ 2.33 ⌉ = 3 PEs.
   B and C flag 'see calculator' because they depend on user inputs.

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

tests/analytical_visualization/auto_explore.py

@@ -419,6 +419,7 @@ def run_auto_explore(
     mode: str = "decode",
     include_attention: bool = True,
     include_ffn: bool = True,
+    b: int = 1,
 ) -> AutoExploreResult:
     """Enumerate all configs, score each, extract Pareto frontier.
 
@@ -433,6 +434,8 @@ def run_auto_explore(
     all_scores: list[ConfigScore] = []
     total_enumerated = 0
     for topo in enumerate_configs(model, s_kv, mode):
+        # Inject the caller's batch B into every candidate.
+        topo.b = b
         total_enumerated += 1
         cfg = FullConfig(model=model, topo=topo, machine=machine)
         all_scores.append(score_config(

tests/analytical_visualization/auto_hardware.py

@@ -237,6 +237,7 @@ def _best_parallelism_for_hw(
     model: ModelConfig, machine: MachineParams,
     s_kv: int, mode: str,
     include_attention: bool = True, include_ffn: bool = True,
+    b: int = 1,
 ) -> ConfigScore | None:
     """Return the latency-minimum feasible parallelism for this HW.
 
@@ -245,14 +246,20 @@ def _best_parallelism_for_hw(
     matches the caller's attention/FFN/full choice.
     """
     best: ConfigScore | None = None
+    best_key: tuple | None = None
     for topo in _iter_reduced_parallelism(model, s_kv, mode):
+        topo.b = b
         cfg = FullConfig(model=model, topo=topo, machine=machine)
         s = score_config(cfg, include_attention=include_attention,
                           include_ffn=include_ffn)
         if not (s.fits_memory and s.placement_valid):
             continue
-        if best is None or s.total_latency_ns < best.total_latency_ns:
+        # Compound tiebreaker: prefer fewer PEs and lower HBM % when latency
+        # ties across variant knobs (cp_ring, placement, kv_shard_mode, etc.).
+        _k = (s.total_latency_ns, s.pes_used, s.hbm_utilization)
+        if best is None or _k < best_key:
             best = s
+            best_key = _k
     return best
 
 
@@ -260,13 +267,14 @@ def _best_parallelism_two_stage(
     model: ModelConfig, machine: MachineParams,
     s_kv: int, mode: str,
     include_attention: bool = True, include_ffn: bool = True,
+    b: int = 1,
 ) -> ConfigScore | None:
     """Fast fallback: use autosuggest's memory-min (CP,TP,PP) then score it."""
     sug = auto_suggest(model, machine, s_kv, mode)
     if not sug.fits:
         return None
     topo = TopologyConfig(
-        cp=sug.cp, tp=sug.tp, pp=sug.pp, dp=1,
+        cp=sug.cp, tp=sug.tp, pp=sug.pp, dp=1, b=b,
         s_kv=s_kv, mode=mode,
         kv_shard_mode="split",
         ffn_shard_scope="TP+CP",
@@ -361,6 +369,7 @@ def joint_explore(
     depth: str = "balanced",
     include_attention: bool = True,
     include_ffn: bool = True,
+    b: int = 1,
 ) -> JointExploreResult:
     """Sweep HW candidates × parallelism, return joint Pareto + sensitivity.
 
@@ -378,11 +387,13 @@ def joint_explore(
             par = _best_parallelism_two_stage(
                 model, machine, s_kv, mode,
                 include_attention=include_attention, include_ffn=include_ffn,
+                b=b,
             )
         else:
             par = _best_parallelism_for_hw(
                 model, machine, s_kv, mode,
                 include_attention=include_attention, include_ffn=include_ffn,
+                b=b,
             )
         if par is None:
             continue

tests/analytical_visualization/autosuggest.py

@@ -32,6 +32,8 @@ class Suggestion:
     fits: bool
     pes_used: int
     sips_used: int
+    cubes_used: int = 0            # picked to minimize this (see _score_candidate)
+    cp_placement: str = "cube"     # "pe" packs CP into intra-cube PEs when it fits
     reason: str = ""
 
 
@@ -58,10 +60,33 @@ def _iter_candidates(model: ModelConfig) -> Iterable[tuple[int, int, int]]:
 def _score_candidate(cp: int, tp: int, pp: int,
                      model: ModelConfig, machine: MachineParams,
                      s_kv: int, mode: str,
-                     slack_frac: float = 0.10) -> Suggestion:
-    topo = TopologyConfig(cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode)
+                     slack_frac: float = 0.10,
+                     b: int = 1,
+                     include_attention: bool = True,
+                     include_ffn: bool = True) -> Suggestion:
+    # For each triple, try both cp_placement options and keep the one
+    # with fewer cubes (breaks the historical "CP always across cubes"
+    # default when a smaller pack is possible). The pe placement is only
+    # valid when CP·TP fits within a single cube's PE count.
+    best_topo: TopologyConfig | None = None
+    best_placement = "cube"
+    _pe_per_cube_hw = TopologyConfig().pes_per_cube_hw  # instance-invariant HW const
+    _placements_to_try = ["cube"]
+    if cp * tp <= _pe_per_cube_hw:
+        _placements_to_try.append("pe")
+    for _place in _placements_to_try:
+        _t = TopologyConfig(
+            cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode, b=b,
+            cp_placement=_place,
+        )
+        if best_topo is None or _t.cubes_used < best_topo.cubes_used:
+            best_topo = _t
+            best_placement = _place
+    topo = best_topo
     cfg = FullConfig(model=model, topo=topo, machine=machine)
-    mem = compute_memory(cfg)
+    mem = compute_memory(
+        cfg, include_attention=include_attention, include_ffn=include_ffn,
+    )
 
     fits = (not mem.over_budget
             and mem.slack_bytes >= slack_frac * mem.budget_bytes)
@@ -80,27 +105,42 @@ def _score_candidate(cp: int, tp: int, pp: int,
         fits=fits,
         pes_used=topo.total_pes,
         sips_used=topo.sips_used,
+        cubes_used=topo.cubes_used,
+        cp_placement=best_placement,
         reason=reason,
     )
 
 
 def auto_suggest(model: ModelConfig, machine: MachineParams,
                  s_kv: int, mode: str = "decode",
-                 slack_frac: float = 0.10) -> Suggestion:
-    """Return the smallest-PE-count (CP, TP, PP) that fits.
+                 slack_frac: float = 0.10,
+                 b: int = 1,
+                 include_attention: bool = True,
+                 include_ffn: bool = True) -> Suggestion:
+    """Return the smallest deployment (fewer cubes, then fewer PEs)
+    that fits.
+
+    Sort key: (cubes_used ↑, pes_used ↑, pp ↑, tp ↑, cp ↑).
+    Fewer cubes wins first because a cube is the physical die-level
+    hardware unit; PE count is the tiebreaker. Preferring fewer PP then
+    TP then CP keeps the earlier historical bias (avoid pipeline
+    bubbles / head-dim splits) among equal-cube-and-PE ties.
 
     If no candidate fits, returns the best-effort one (highest slack,
     even if negative) with fits=False.
     """
-    best_fit: Suggestion | None = None
-    best_effort: Suggestion | None = None
-
+    scored: list[Suggestion] = []
     for cp, tp, pp in _iter_candidates(model):
-        s = _score_candidate(cp, tp, pp, model, machine, s_kv, mode, slack_frac)
-        if s.fits and best_fit is None:
-            best_fit = s
-            break  # candidates are pre-sorted by PE count
-        if best_effort is None or s.slack_gb > best_effort.slack_gb:
-            best_effort = s
+        scored.append(
+            _score_candidate(cp, tp, pp, model, machine, s_kv, mode,
+                              slack_frac, b=b,
+                              include_attention=include_attention,
+                              include_ffn=include_ffn)
+        )
+    scored.sort(key=lambda s: (s.cubes_used, s.pes_used, s.pp, s.tp, s.cp))
 
-    return best_fit or best_effort  # type: ignore
+    for s in scored:
+        if s.fits:
+            return s
+    # No fit — return the best-effort (largest slack, closest to fitting).
+    return max(scored, key=lambda s: s.slack_gb)

tests/analytical_visualization/chip_roofline.py

@@ -0,0 +1,518 @@
+"""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, replace
+
+from .model_config import FullConfig, 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)
+
+
+# ── Regime-dependent cost terms ────────────────────────────────────
+
+
+def t_mem_short(machine: MachineParams, model: ModelConfig,
+                 batch: int) -> float:
+    """Per-token weight-fetch time (short-context regime term).
+
+    N_active · b / (W · B). Shrinks as B grows — this is what
+    batching amortizes.
+    """
+    return (total_active_params(model) * model.bytes_per_elem
+            / (machine.bw_hbm * max(1, batch)))
+
+
+def t_mem_long(machine: MachineParams, model: ModelConfig,
+                s_kv: int) -> float:
+    """Per-token KV-read time (long-context regime term).
+
+    S_kv · kv_bpt / W. Independent of B — each sequence reads its
+    own KV cache; batching doesn't help.
+    """
+    return s_kv * kv_bytes_per_token(model) / machine.bw_hbm
+
+
+def t_com(machine: MachineParams, model: ModelConfig) -> float:
+    """Per-token compute time. Same in both regimes: 2·N/C, peak."""
+    return _FLOPS_PER_PARAM_PER_TOKEN * total_active_params(model) / machine.peak_flops
+
+
+# ── "Good" batch / context recommendations ─────────────────────────
+
+
+@dataclass
+class BatchRecommendation:
+    target: float          # Pope's rule: 2 × B*
+    b_star: float          # B* itself
+    effective: float       # what we recommend using
+    reason: str            # short explanation
+
+
+@dataclass
+class ContextRecommendation:
+    l_star: float          # balance context length
+    max_efficient: float   # same as l_star (compute-friendly ceiling)
+    utilization_at: float  # utilization at current s_kv
+    reason: str
+
+
+def good_batch(machine: MachineParams, model: ModelConfig,
+                sparsity: float = 1.0) -> BatchRecommendation:
+    """Recommended batch size: 2 × B* (Pope's rule of thumb).
+
+    Below B*: memory-bound, doubling B halves cost/token.
+    At 2×B*: 50% excess over compute floor — the sweet spot.
+    Beyond 3×B*: diminishing returns; latency keeps growing linearly.
+    """
+    b_star = critical_batch(machine, model, sparsity)
+    target = 2 * b_star
+    return BatchRecommendation(
+        target=target, b_star=b_star, effective=target,
+        reason=(f"2·B* = 2 · {b_star:.0f} = {target:.0f}. "
+                "Below B*: memory-bound (doubling B halves cost/token). "
+                "Beyond 3·B*: diminishing returns."),
+    )
+
+
+def good_context(machine: MachineParams, model: ModelConfig,
+                  s_kv: int) -> ContextRecommendation:
+    """Recommended max context: L* — the compute-friendly ceiling.
+
+    Below L*: KV read is cheap relative to compute → good utilization.
+    Above L*: KV bandwidth wall → utilization = 1/(1 + S_kv/L*).
+    """
+    l_star = balance_context(machine, model)
+    util = utilization_at(s_kv, l_star)
+    return ContextRecommendation(
+        l_star=l_star, max_efficient=l_star, utilization_at=util,
+        reason=(f"L* = {l_star:,.0f} tokens. Below L*: compute-bound "
+                f"(good util). At {s_kv:,} tokens: peak utilization ≈ "
+                f"{util*100:.1f}% (1 / (1 + S_kv/L*))."),
+    )
+
+
+def utilization_at(s_kv: int, l_star: float) -> float:
+    """Peak compute utilization at context length s_kv, given L*.
+
+    util = compute / (compute + KV_read) = 1 / (1 + S_kv/L*).
+    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
+
+
+# ── PE memory budget curves ───────────────────────────────────────
+
+
+@dataclass
+class MemoryBudgetPoint:
+    axis_val: int          # S_kv or B being swept
+    weights_gb: float
+    kv_gb: float
+    transient_gb: float
+    used_gb: float
+    free_gb: float         # max(0, hbm_gb - used_gb)
+    over_budget: bool
+
+
+def _budget_point(weights_bytes: int, kv_bytes: int, transient_bytes: int,
+                   hbm_bytes: int, axis_val: int) -> MemoryBudgetPoint:
+    used = weights_bytes + kv_bytes + transient_bytes
+    return MemoryBudgetPoint(
+        axis_val=axis_val,
+        weights_gb=weights_bytes / 1e9,
+        kv_gb=kv_bytes / 1e9,
+        transient_gb=transient_bytes / 1e9,
+        used_gb=used / 1e9,
+        free_gb=max(0, hbm_bytes - used) / 1e9,
+        over_budget=(used > hbm_bytes),
+    )
+
+
+def memory_budget_curve_vs_skv(cfg: FullConfig,
+                                s_kv_range: list[int],
+                                batch: int) -> list[MemoryBudgetPoint]:
+    """Per-PE memory as S_kv sweeps. Uses cfg's current sharding
+    (CP, TP, PP). Sets topo.b = batch. Returns one point per S_kv."""
+    from .memory_layout import (
+        per_pe_kv_cache_bytes, per_pe_transient_bytes, per_pe_weight_bytes,
+    )
+    hbm_bytes = int(cfg.machine.pe_budget_bytes)
+    weights_bytes = per_pe_weight_bytes(cfg)
+    transient_bytes = per_pe_transient_bytes(cfg)
+    points: list[MemoryBudgetPoint] = []
+    for skv in s_kv_range:
+        swept_topo = replace(cfg.topo, s_kv=int(skv), b=max(1, int(batch)))
+        swept_cfg = FullConfig(model=cfg.model, topo=swept_topo,
+                                machine=cfg.machine)
+        kv_bytes = per_pe_kv_cache_bytes(swept_cfg)
+        points.append(_budget_point(weights_bytes, kv_bytes,
+                                     transient_bytes, hbm_bytes, int(skv)))
+    return points
+
+
+def memory_budget_curve_vs_batch(cfg: FullConfig,
+                                   b_range: list[int],
+                                   s_kv: int) -> list[MemoryBudgetPoint]:
+    """Per-PE memory as B sweeps. Sets topo.s_kv = s_kv."""
+    from .memory_layout import (
+        per_pe_kv_cache_bytes, per_pe_transient_bytes, per_pe_weight_bytes,
+    )
+    hbm_bytes = int(cfg.machine.pe_budget_bytes)
+    weights_bytes = per_pe_weight_bytes(cfg)
+    transient_bytes = per_pe_transient_bytes(cfg)
+    points: list[MemoryBudgetPoint] = []
+    for bs in b_range:
+        swept_topo = replace(cfg.topo, b=max(1, int(bs)), s_kv=int(s_kv))
+        swept_cfg = FullConfig(model=cfg.model, topo=swept_topo,
+                                machine=cfg.machine)
+        kv_bytes = per_pe_kv_cache_bytes(swept_cfg)
+        points.append(_budget_point(weights_bytes, kv_bytes,
+                                     transient_bytes, hbm_bytes, int(bs)))
+    return points
+
+
+# ── AI / B* sensitivity to hardware knobs ──────────────────────────
+
+
+@dataclass
+class AISensitivityPoint:
+    multiplier: float          # scale factor applied to the base machine
+    peak_tflops: float
+    bw_gbs: float
+    ai: float                  # C / W
+    b_star: float              # C·b/(2·W)
+
+
+# ── GPU count sizing (three-axis) ─────────────────────────────────
+
+
+@dataclass
+class GpuSizingResult:
+    # Inputs echoed back for display / debugging
+    n_users: int
+    avg_ctx_tokens: int
+    tpot_slo_s: float
+    # Per-replica values
+    pes_axis_a_capacity: int    # PEs to hold weights alone (bare floor)
+    pes_axis_b_kv: int          # PEs to hold weights + KV of this replica's users
+    pes_per_replica: int        # max(A, B)
+    # Throughput axis
+    b_at_slo: int               # largest per-replica B satisfying SLO
+    users_per_replica: int      # min(n_users, b_at_slo), at least 1
+    n_replicas: int             # ceil(n_users / users_per_replica)
+    # Grand total
+    total_pes: int
+    binding_axis: str           # 'capacity' | 'kv' | 'throughput'
+
+
+def max_batch_within_slo(machine: MachineParams, model: ModelConfig,
+                          s_kv: int, tpot_slo_s: float) -> int:
+    """Largest per-replica B such that decode step latency ≤ SLO.
+
+    step_latency(B) = N·b/W + 2·N·B/C + B·S_kv·kv_bpt/W
+                    = weight_fetch + B · (compute + kv_read)
+
+    Returns 0 if even B=1 exceeds SLO (weight fetch alone too large,
+    or per-sequence term already blows the budget).
+    """
+    n = total_active_params(model)
+    b_elem = model.bytes_per_elem
+    weight_time = n * b_elem / machine.bw_hbm
+    per_seq_time = (_FLOPS_PER_PARAM_PER_TOKEN * n / machine.peak_flops
+                    + s_kv * kv_bytes_per_token(model) / machine.bw_hbm)
+    if weight_time + per_seq_time > tpot_slo_s:
+        return 0
+    remaining = tpot_slo_s - weight_time
+    b_max = int(remaining / per_seq_time)
+    return max(1, b_max)
+
+
+def size_deployment(machine: MachineParams, model: ModelConfig,
+                     n_users: int, avg_ctx: int,
+                     tpot_slo_s: float) -> GpuSizingResult:
+    """Three-axis GPU count for a target workload.
+
+    Axis A — capacity floor: PEs to hold one replica's weights.
+    Axis B — KV headroom:    PEs to hold weights + all KV of the
+                             users assigned to this replica.
+    Axis C — throughput SLO: replicas needed so per-replica B ≤ b_at_slo.
+
+    Result: pes_per_replica = max(A, B); total = pes_per_replica × replicas.
+    Binding axis is whichever grew the count the most.
+    """
+    n = total_active_params(model)
+    b_elem = model.bytes_per_elem
+    weight_bytes = n * b_elem
+    hbm_pe = int(machine.pe_hbm_gb * 1e9)
+    kv_bpt = kv_bytes_per_token(model)
+
+    # Axis C: throughput
+    b_at_slo = max_batch_within_slo(machine, model, avg_ctx, tpot_slo_s)
+    users_per_replica = max(1, min(n_users, b_at_slo)) if b_at_slo > 0 else 1
+    n_replicas = (n_users + users_per_replica - 1) // users_per_replica
+
+    # Axis A: bare weights
+    pes_a = (weight_bytes + hbm_pe - 1) // hbm_pe
+
+    # Axis B: weights + KV load for this replica's users
+    kv_per_replica = users_per_replica * avg_ctx * kv_bpt
+    pes_b = (weight_bytes + kv_per_replica + hbm_pe - 1) // hbm_pe
+
+    pes_per_replica = max(pes_a, pes_b)
+    total = pes_per_replica * n_replicas
+
+    if n_replicas > 1:
+        binding = "throughput"
+    elif pes_b > pes_a:
+        binding = "kv"
+    else:
+        binding = "capacity"
+
+    return GpuSizingResult(
+        n_users=n_users,
+        avg_ctx_tokens=avg_ctx,
+        tpot_slo_s=tpot_slo_s,
+        pes_axis_a_capacity=int(pes_a),
+        pes_axis_b_kv=int(pes_b),
+        pes_per_replica=int(pes_per_replica),
+        b_at_slo=b_at_slo,
+        users_per_replica=int(users_per_replica),
+        n_replicas=int(n_replicas),
+        total_pes=int(total),
+        binding_axis=binding,
+    )
+
+
+def ai_sensitivity_curve(machine: MachineParams, model: ModelConfig,
+                          multipliers: list[float],
+                          axis: str = "flops") -> list[AISensitivityPoint]:
+    """Sweep FLOPs OR BW while holding the other fixed. Returns AI + B*
+    at each multiplier.
+
+    axis='flops' → scales peak_tflops_f16 (AI grows linearly).
+    axis='bw'    → scales bw_hbm_gbs (AI shrinks inversely).
+    """
+    if axis not in ("flops", "bw"):
+        raise ValueError(f"axis must be 'flops' or 'bw', got {axis!r}")
+    base_flops = machine.peak_tflops_f16
+    base_bw = machine.bw_hbm_gbs
+    points: list[AISensitivityPoint] = []
+    for k in multipliers:
+        if axis == "flops":
+            m = replace(machine, peak_tflops_f16=base_flops * k)
+        else:
+            m = replace(machine, bw_hbm_gbs=base_bw * k)
+        points.append(AISensitivityPoint(
+            multiplier=k,
+            peak_tflops=m.peak_tflops_f16,
+            bw_gbs=m.bw_hbm_gbs,
+            ai=arithmetic_intensity(m),
+            b_star=critical_batch(m, model),
+        ))
+    return points

tests/analytical_visualization/memory_layout.py

@@ -30,12 +30,19 @@ class MemoryBreakdown:
         return self.used_bytes > self.budget_bytes
 
 
-def per_pe_weight_bytes(cfg: FullConfig) -> int:
+def per_pe_weight_bytes(cfg: FullConfig,
+                        include_attention: bool = True,
+                        include_ffn: bool = True) -> int:
     """Attention + FFN weights per PE, bf16. Divided by TP, PP, and EP.
 
     EP divides the FFN (experts) across ranks; attention weights are
     unaffected. When kv_shard_mode='replicate' and TP > H_kv, each KV
     head is replicated (per-PE W_K/W_V size doesn't drop below one head).
+
+    Scope flags let callers count only attention or only FFN weights —
+    useful for "smallest deployment for just this block" sizing. Both
+    True (default) matches the traditional full-model per-PE weight
+    footprint.
     """
     m = cfg.model
     tp = cfg.topo.tp
@@ -50,18 +57,23 @@ def per_pe_weight_bytes(cfg: FullConfig) -> int:
         # Head-dim split: fractional head allowed (< 1 head bytes when TP > H_kv).
         hkv_per_pe_bytes = m.h_kv / tp
 
+    per_layer_attn = 0
+    if include_attention:
         per_layer_attn = (
             m.hidden * hq_per_pe * m.d_head +           # W_Q
             m.hidden * hkv_per_pe_bytes * m.d_head +    # W_K
             m.hidden * hkv_per_pe_bytes * m.d_head +    # W_V
             hq_per_pe * m.d_head * m.hidden             # W_O
         )
+
+    per_layer_ffn = 0
+    if include_ffn:
         # FFN divisor: TP (default) or TP*CP or TP*CP*DP if user opts in.
         # EP further divides (MoE experts).
         ffn_div = cfg.ffn_shard_divisor * ep
         per_layer_ffn = 3 * m.hidden * (m.ffn_dim // ffn_div)
-    per_layer = (per_layer_attn + per_layer_ffn) * m.bytes_per_elem
 
+    per_layer = (per_layer_attn + per_layer_ffn) * m.bytes_per_elem
     layers_per_stage = (m.layers + pp - 1) // pp
     return int(per_layer * layers_per_stage)
 
@@ -74,17 +86,19 @@ def per_pe_kv_cache_bytes(cfg: FullConfig) -> int:
       TP > H_kv, each KV head is duplicated across TP/H_kv ranks so
       per-PE storage doesn't shrink below 1 head.
     - PP: only layers_per_stage layers stored per PE.
+    - Batch (B): each concurrent request keeps its own KV cache slice.
     """
     m = cfg.model
     pp = cfg.topo.pp
     tp = cfg.topo.tp
+    B = max(1, cfg.topo.b)
     if cfg.topo.kv_shard_mode == "replicate":
         hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
     else:
         hkv_per_pe_bytes = m.h_kv / tp
     layers_per_stage = (m.layers + pp - 1) // pp
     per_layer = 2 * cfg.topo.s_local * hkv_per_pe_bytes * m.d_head * m.bytes_per_elem
-    return int(per_layer * layers_per_stage)
+    return int(per_layer * layers_per_stage * B)
 
 
 def per_pe_transient_bytes(cfg: FullConfig) -> int:
@@ -101,10 +115,19 @@ def per_pe_transient_bytes(cfg: FullConfig) -> int:
         return int(2 * tile_score + T_q * m.hidden * m.bytes_per_elem)
 
 
-def compute_memory(cfg: FullConfig) -> MemoryBreakdown:
+def compute_memory(cfg: FullConfig,
+                   include_attention: bool = True,
+                   include_ffn: bool = True) -> MemoryBreakdown:
+    """Per-PE memory breakdown. Scope flags let callers count only the
+    attention or only the FFN block (KV cache goes with the attention
+    block; transient activation buffer is small either way and left in).
+    """
+    kv_bytes = per_pe_kv_cache_bytes(cfg) if include_attention else 0
     return MemoryBreakdown(
-        weights_bytes=per_pe_weight_bytes(cfg),
-        kv_cache_bytes=per_pe_kv_cache_bytes(cfg),
+        weights_bytes=per_pe_weight_bytes(
+            cfg, include_attention=include_attention, include_ffn=include_ffn,
+        ),
+        kv_cache_bytes=kv_bytes,
         transient_bytes=per_pe_transient_bytes(cfg),
         budget_bytes=cfg.machine.pe_budget_bytes,
     )

tests/analytical_visualization/model_config.py

@@ -43,6 +43,7 @@ class TopologyConfig:
     pp: int = 1                 # pipeline stages (layer shard)
     dp: int = 1                 # data-parallel replicas
     ep: int = 1                 # expert-parallel degree (MoE only)
+    b: int = 1                  # batch size (concurrent requests per PP stage)
     s_kv: int = 4096
     mode: str = "decode"        # "decode" or "prefill"
     kv_shard_mode: str = "split"  # "split" or "replicate" (used when TP > H_kv)

tests/analytical_visualization/pe_weight_layout.py

@@ -114,6 +114,14 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
     cube_h = 6.0
     cube_gap = 0.6
 
+    # When TP fills the cube (TP>=8), cp_placement is "cube" and CP
+    # ranks span cubes — the intra-cube CP-color branch below is
+    # skipped. Surface which CP rank this figure is showing so users
+    # aren't guessing which of `CP` groups is drawn.
+    _cube_place_tag = ""
+    if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1:
+        _cube_place_tag = f"  [CP rank 0 of {cfg.topo.cp}]"
+
     for cube_idx in range(n_cubes):
         x0 = cube_idx * (cube_w + cube_gap)
         y0 = 0
@@ -125,7 +133,8 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
         )
         ax.add_patch(rect)
         ax.text(x0 + cube_w / 2, y0 + cube_h + 0.1,
-                f"Cube {cube_idx}  (PEs {cube_idx*8}-{cube_idx*8+7})",
+                f"Cube {cube_idx}  (PEs {cube_idx*8}-{cube_idx*8+7})"
+                + _cube_place_tag,
                 ha="center", va="bottom", fontsize=11, fontweight="bold")
 
         pe_pad_x = 0.15
@@ -136,6 +145,26 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
         pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
         pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
 
+        # When cp_placement=pe, CP ranks are packed intra-cube along with
+        # TP ranks. Color-code by CP rank so it's obvious which PEs belong
+        # to which sequence-shard group. Palette wraps beyond 8 CP ranks.
+        _cp_place = cfg.topo.cp_placement
+        _cp_val = cfg.topo.cp
+        _cp_palette = [
+            "#f0f7ff",  # baseline blue (also the fallback for cp_placement=cube)
+            "#e8f5e9",  # light green
+            "#fff3e0",  # light orange
+            "#f3e5f5",  # light purple
+            "#ffebee",  # light red
+            "#e0f7fa",  # light cyan
+            "#fce4ec",  # light pink
+            "#f9fbe7",  # light lime
+        ]
+        _cp_edge_palette = [
+            "#3a86ff", "#2e7d32", "#ef6c00", "#7b1fa2",
+            "#c62828", "#0097a7", "#c2185b", "#9e9d24",
+        ]
+
         for pr in range(PE_ROWS):
             for pc in range(PE_COLS):
                 pe_local = pr * PE_COLS + pc
@@ -145,17 +174,43 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
                 px = x0 + pe_pad_x + pc * (pe_w + pe_gap)
                 py = y0 + pe_pad_y + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
 
+                # Which (cp_rank, tp_rank) does this PE hold?
+                # When cp_placement=pe: cp_rank = pe_id // tp, tp_rank = pe_id % tp.
+                # When cp_placement=cube: every PE in this cube is one CP rank
+                # (drawn on its own tile), tp_rank = pe_id % tp.
+                if _cp_place == "pe" and _cp_val > 1:
+                    _cp_rank = pe_id_in_group // tp
+                    _tp_rank = pe_id_in_group % tp
+                    _fc = _cp_palette[_cp_rank % len(_cp_palette)]
+                    _ec = _cp_edge_palette[_cp_rank % len(_cp_edge_palette)]
+                    _cp_label = f"  |  CP={_cp_rank}"
+                else:
+                    _cp_rank = None
+                    _tp_rank = pe_id_in_group % tp
+                    _fc = "#f0f7ff"
+                    _ec = "#3a86ff"
+                    # cp_placement=="cube" with CP>1: the whole cube is
+                    # one CP rank (rank 0 by convention — see cube title
+                    # tag above). Add the locator so users know what
+                    # rank this per-PE data belongs to.
+                    if _cp_place == "cube" and _cp_val > 1:
+                        _cp_label = "  |  CP=0"
+                    else:
+                        _cp_label = ""
+
                 pe_rect = patches.Rectangle(
                     (px, py), pe_w, pe_h,
-                    facecolor="#f0f7ff", edgecolor="#3a86ff", linewidth=1.0,
+                    facecolor=_fc, edgecolor=_ec, linewidth=1.2,
                 )
                 ax.add_patch(pe_rect)
 
-                q_heads = _q_heads_for_pe(pe_id_in_group, tp, cfg.model.h_q)
+                # Head assignment uses TP rank (not raw pe_id) so that under
+                # cp_placement=pe, all CP ranks share the same head split.
+                q_heads = _q_heads_for_pe(_tp_rank, tp, cfg.model.h_q)
                 kv_heads, kv_note = _kv_heads_for_pe(
-                    pe_id_in_group, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
+                    _tp_rank, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
                 )
-                bytes_ = _per_pe_bytes(cfg, pe_id_in_group)
+                bytes_ = _per_pe_bytes(cfg, _tp_rank)
                 weights_gb = sum(v for k, v in bytes_.items()
                                  if k not in ("KV cache", "Transient")) / 1e9
                 kv_gb = bytes_["KV cache"] / 1e9
@@ -172,7 +227,7 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
                 ffn_mb = (bytes_["W_gate"] + bytes_["W_up"]
                           + bytes_["W_down"]) / 1e6
                 # Header (PE id and heads) bigger, then per-tensor breakdown
-                header = (f"PE {pe_id_in_group}\n"
+                header = (f"PE {pe_id_in_group}  TP={_tp_rank}{_cp_label}\n"
                           f"Q heads: {q_str}\n"
                           f"KV head: {kv_str}")
                 ax.text(px + 0.08, py + pe_h - 0.05, header,
@@ -211,9 +266,12 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
     ax.set_aspect("equal")
     ax.axis("off")
 
+    _title_cp = f", CP={cfg.topo.cp}"
+    if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1:
+        _title_cp += " (showing 1 of CP groups; others identical)"
     title = (
         f"Per-PE layout of one CP group  "
-        f"(TP={tp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})"
+        f"(TP={tp}{_title_cp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})"
         f" | KV mode: {cfg.topo.kv_shard_mode}"
     )
     ax.set_title(title, fontsize=11, fontweight="bold")

tests/analytical_visualization/stage_latencies.py

@@ -50,19 +50,22 @@ def stage_rmsnorm(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    bytes_ = T_q * d * b * 2  # load x + load weight
-    flops = 4 * T_q * d
+    B = max(1, cfg.topo.b)
+    # Activation memory (x) scales with B; weight (once) is fixed.
+    bytes_ = B * T_q * d * b + T_q * d * b
+    flops = 4 * B * T_q * d
     mem_s = bytes_ / cfg.machine.bw_hbm
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name="S1 RMSNorm",
-        formula=f"bytes = {T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM",
+        formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
         bound=bnd, visible_s=vis,
         flops=flops, mem_bytes=bytes_,
-        flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}",
-        mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B",
+        flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}",
+        mem_formula=f"B*T_q*d*b + T_q*d*b (weight) "
+                     f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B",
     )
 
 
@@ -76,21 +79,26 @@ def stage_wq(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     hq_per_pe = cfg.h_q_per_pe
     dh = cfg.model.d_head
-    flops = 2 * T_q * d * (hq_per_pe * dh)
+    # FLOPs scale with batch; weight bytes fixed (shared across batch).
+    flops = 2 * B * T_q * d * (hq_per_pe * dh)
     weight_B = d * (hq_per_pe * dh) * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name="S2 W_Q GEMM",
-        formula=f"FLOPs = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; "
-                f"weight = {weight_B/1e6:.1f} MB",
+        formula=f"FLOPs = 2*B*T_q*d*(H_q/TP*d_h) "
+                f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
         bound=bnd, visible_s=vis,
         flops=int(flops), mem_bytes=int(weight_B),
-        flops_formula=f"2*T_q*d*(H_q/TP*d_h) = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}",
-        mem_formula=f"d*(H_q/TP*d_h)*b = {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB",
+        flops_formula=f"2*B*T_q*d*(H_q/TP*d_h) "
+                       f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}",
+        mem_formula=f"d*(H_q/TP*d_h)*b (weight, B-invariant) "
+                     f"= {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB",
     )
 
 
@@ -98,9 +106,10 @@ def stage_wkv(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
     dh = cfg.model.d_head
-    flops_one = 2 * T_q * d * (hkv_per_pe * dh)
+    flops_one = 2 * B * T_q * d * (hkv_per_pe * dh)
     weight_B_one = d * (hkv_per_pe * dh) * b
     flops = 2 * flops_one
     weight_B = 2 * weight_B_one
@@ -108,30 +117,36 @@ def stage_wkv(cfg: FullConfig) -> StageCost:
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name="S3 W_K + W_V GEMM",
-        formula=f"FLOPs per proj = 2*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2",
+        formula=f"FLOPs per proj = 2*B*T_q*d*(H_kv/TP*d_h) "
+                f"= 2*{B}*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
         bound=bnd, visible_s=vis,
         flops=int(flops), mem_bytes=int(weight_B),
-        flops_formula=f"2*(2*T_q*d*(H_kv/TP*d_h)) = 2*(2*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}",
-        mem_formula=f"2*(d*(H_kv/TP*d_h)*b) = 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB",
+        flops_formula=f"2*(2*B*T_q*d*(H_kv/TP*d_h)) "
+                       f"= 2*(2*{B}*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}",
+        mem_formula=f"2*(d*(H_kv/TP*d_h)*b) (weights, B-invariant) "
+                     f"= 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB",
     )
 
 
 def stage_kv_append(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
     dh = cfg.model.d_head
-    bytes_ = 2 * T_q * hkv_per_pe * dh * b
+    bytes_ = 2 * B * T_q * hkv_per_pe * dh * b
     mem_s = bytes_ / cfg.machine.bw_hbm
     return StageCost(
         name="S4 KV cache append",
-        formula=f"bytes = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
+        formula=f"bytes = 2*B*T_q*(H_kv/TP)*d_h*b "
+                f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
         compute_s=0, memory_s=mem_s, comm_s=0,
         bound="memory", visible_s=mem_s,
         flops=0, mem_bytes=bytes_,
         flops_formula="0 (no matmul, just cache write)",
-        mem_formula=f"2*T_q*(H_kv/TP)*d_h*b = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
+        mem_formula=f"2*B*T_q*(H_kv/TP)*d_h*b "
+                     f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
     )
 
 
@@ -141,7 +156,8 @@ def _per_hop_qkT_pv(cfg: FullConfig) -> tuple[float, str]:
     S_local = cfg.topo.s_local
     dh = cfg.model.d_head
     hq_per_pe = cfg.h_q_per_pe
-    flops = 2 * T_q * S_local * dh * hq_per_pe
+    B = max(1, cfg.topo.b)
+    flops = 2 * B * T_q * S_local * dh * hq_per_pe
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
     formula = f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops:.2g} FLOPs/hop"
     return cmp_s, formula
@@ -165,19 +181,23 @@ def stage_qkT(cfg: FullConfig) -> StageCost:
     S_local = cfg.topo.s_local
     dh = cfg.model.d_head
     hq_per_pe = cfg.h_q_per_pe
-    flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe
+    B = max(1, cfg.topo.b)
+    flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe
     total_flops = flops_per_hop * passes
     _hop_word = "pass" if passes == 1 else "hops"
     return StageCost(
         name=f"S5 Q.K^T (x{passes} {_hop_word})",
-        formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop",
+        formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = "
+                f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = "
+                f"{flops_per_hop:.2g} FLOPs/hop",
         compute_s=total, memory_s=0, comm_s=0,
         bound="compute", visible_s=total,
         hop_multiplier=passes,
         flops=int(total_flops), mem_bytes=0,
         flops_formula=(
-            f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = "
-            f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}"
+            f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = "
+            f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = "
+            f"{total_flops:.3g}"
         ),
         mem_formula="0 (scores accumulated on-chip)",
     )
@@ -188,7 +208,8 @@ def stage_softmax(cfg: FullConfig) -> StageCost:
     S_local = cfg.topo.s_local
     hq_per_pe = cfg.h_q_per_pe
     b = cfg.model.bytes_per_elem
-    elems = hq_per_pe * T_q * S_local
+    B = max(1, cfg.topo.b)
+    elems = B * hq_per_pe * T_q * S_local
     bytes_ = elems * b * 2
     mem_s_per_hop = bytes_ / cfg.machine.bw_hbm
     passes = _cp_compute_passes(cfg)
@@ -197,15 +218,16 @@ def stage_softmax(cfg: FullConfig) -> StageCost:
     _hop_word = "pass" if passes == 1 else "hops"
     return StageCost(
         name=f"S6 softmax (x{passes} {_hop_word})",
-        formula=f"elts/hop = {hq_per_pe}*{T_q}*{S_local} = {elems:.2g}",
+        formula=f"elts/hop = B*(H_q/TP)*T_q*S_local "
+                f"= {B}*{hq_per_pe}*{T_q}*{S_local} = {elems:.2g}",
         compute_s=0, memory_s=total, comm_s=0,
         bound="memory", visible_s=total,
         hop_multiplier=passes,
         flops=0, mem_bytes=int(total_bytes),
         flops_formula="~O(elts) (negligible)",
         mem_formula=(
-            f"{passes}*2*b*(H_q/TP)*T_q*S_local = "
-            f"{passes}*2*{b}*{hq_per_pe}*{T_q}*{S_local} = "
+            f"{passes}*2*b*B*(H_q/TP)*T_q*S_local = "
+            f"{passes}*2*{b}*{B}*{hq_per_pe}*{T_q}*{S_local} = "
             f"{total_bytes/1e6:.2f} MB"
         ),
     )
@@ -219,19 +241,23 @@ def stage_pv(cfg: FullConfig) -> StageCost:
     S_local = cfg.topo.s_local
     dh = cfg.model.d_head
     hq_per_pe = cfg.h_q_per_pe
-    flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe
+    B = max(1, cfg.topo.b)
+    flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe
     total_flops = flops_per_hop * passes
     _hop_word = "pass" if passes == 1 else "hops"
     return StageCost(
         name=f"S7 P.V (x{passes} {_hop_word})",
-        formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop",
+        formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = "
+                f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = "
+                f"{flops_per_hop:.2g} FLOPs/hop",
         compute_s=total, memory_s=0, comm_s=0,
         bound="compute", visible_s=total,
         hop_multiplier=passes,
         flops=int(total_flops), mem_bytes=0,
         flops_formula=(
-            f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = "
-            f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}"
+            f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = "
+            f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = "
+            f"{total_flops:.3g}"
         ),
         mem_formula="0 (accumulated on-chip)",
     )
@@ -247,7 +273,8 @@ def stage_merge(cfg: FullConfig) -> StageCost:
     hq_per_pe = cfg.h_q_per_pe
     dh = cfg.model.d_head
     b = cfg.model.bytes_per_elem
-    flops = 6 * T_q * hq_per_pe * dh * max(0, cfg.topo.cp - 1)
+    B = max(1, cfg.topo.b)
+    flops = 6 * B * T_q * hq_per_pe * dh * max(0, cfg.topo.cp - 1)
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
 
     comm_s = 0.0
@@ -257,8 +284,8 @@ def stage_merge(cfg: FullConfig) -> StageCost:
 
     if cfg.topo.mode == "decode" and cfg.topo.cp > 1:
         cp = cfg.topo.cp
-        # (O + m + l) bytes per rank
-        M = (T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+        # (O + m + l) bytes per rank — scales with B (per-request partials).
+        M = B * ((T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
         if cfg.topo.cp_placement == "pe":
             bw, alpha, tier = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
         elif cfg.topo.sips_used > 1:
@@ -293,13 +320,15 @@ def stage_merge(cfg: FullConfig) -> StageCost:
 
     return StageCost(
         name=f"S8 online-softmax merge (x{max(0, cfg.topo.cp-1)}){name_suffix}",
-        formula=f"~6*{T_q}*{hq_per_pe}*{dh}*(C-1) = {flops:.2g} FLOPs",
+        formula=f"~6*B*T_q*(H_q/TP)*d_h*(C-1) "
+                f"= 6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} "
+                f"= {flops:.2g} FLOPs",
         compute_s=cmp_s, memory_s=0, comm_s=comm_s,
         bound=bound, visible_s=visible_s,
         flops=int(flops), mem_bytes=0, comm_bytes=comm_bytes,
         flops_formula=(
-            f"6*T_q*(H_q/TP)*d_h*(CP-1) = "
-            f"6*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}"
+            f"6*B*T_q*(H_q/TP)*d_h*(CP-1) = "
+            f"6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}"
         ),
         mem_formula="0 (in-register)",
         comm_formula=comm_formula or "0",
@@ -310,15 +339,18 @@ def stage_normalize(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     hq_per_pe = cfg.h_q_per_pe
     dh = cfg.model.d_head
-    flops = T_q * hq_per_pe * dh
+    B = max(1, cfg.topo.b)
+    flops = B * T_q * hq_per_pe * dh
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
     return StageCost(
         name="S9 normalize O/l",
-        formula=f"{T_q}*{hq_per_pe}*{dh} = {flops} divisions",
+        formula=f"B*T_q*(H_q/TP)*d_h = {B}*{T_q}*{hq_per_pe}*{dh} "
+                f"= {flops} divisions",
         compute_s=cmp_s, memory_s=0, comm_s=0,
         bound="trivial", visible_s=cmp_s,
         flops=int(flops), mem_bytes=0,
-        flops_formula=f"T_q*(H_q/TP)*d_h = {T_q}*{hq_per_pe}*{dh} = {flops}",
+        flops_formula=f"B*T_q*(H_q/TP)*d_h "
+                       f"= {B}*{T_q}*{hq_per_pe}*{dh} = {flops}",
         mem_formula="0",
     )
 
@@ -327,21 +359,25 @@ def stage_wo(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     hq_per_pe = cfg.h_q_per_pe
     dh = cfg.model.d_head
-    flops = 2 * T_q * (hq_per_pe * dh) * d
+    flops = 2 * B * T_q * (hq_per_pe * dh) * d
     weight_B = (hq_per_pe * dh) * d * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name="S10 W_O GEMM",
-        formula=f"FLOPs = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; "
-                f"weight = {weight_B/1e6:.1f} MB",
+        formula=f"FLOPs = 2*B*T_q*(H_q/TP*d_h)*d "
+                f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
         bound=bnd, visible_s=vis,
         flops=int(flops), mem_bytes=int(weight_B),
-        flops_formula=f"2*T_q*(H_q/TP*d_h)*d = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}",
-        mem_formula=f"(H_q/TP*d_h)*d*b = {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB",
+        flops_formula=f"2*B*T_q*(H_q/TP*d_h)*d "
+                       f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}",
+        mem_formula=f"(H_q/TP*d_h)*d*b (weight, B-invariant) "
+                     f"= {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB",
     )
 
 
@@ -369,12 +405,13 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     dh = cfg.model.d_head
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
 
     # ── Decode: single O/m/l all-reduce after local compute ──
     if cfg.topo.mode == "decode":
         cp = cfg.topo.cp
-        # per-rank O + m + l bytes
-        M = (T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+        # per-rank O + m + l bytes — scales with B (per-request partials).
+        M = B * ((T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
         if cfg.topo.cp_placement == "pe":
             bw = cfg.machine.bw_intra
             alpha = cfg.machine.alpha_intra
@@ -392,7 +429,7 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost:
         return StageCost(
             name=f"C1 O/m/l all-reduce ({tier}, x{cp} ranks)",
             formula=(f"decode: gather partial (O,m,l) once at end; "
-                     f"M = T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; "
+                     f"M = B*T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; "
                      f"AR: 2*(CP-1)/CP*M / BW + 2*(CP-1)*alpha"),
             compute_s=0, memory_s=0, comm_s=ar_time,
             bound="comm", visible_s=ar_time,
@@ -417,19 +454,19 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost:
 
     # ── Prefill: per-hop ring (K/V or Q+O/m/l) concurrent with S5-S7 ──
     if cfg.topo.cp_ring_variant == "qoml":
-        M_KV = (2 * T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+        M_KV = B * ((2 * T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
         variant_desc = "Q+O/m/l ring"
         formula_bytes = (
-            f"2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b "
-            f"= 2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b} "
+            f"B*(2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b) "
+            f"= {B}*(2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b}) "
             f"= {M_KV} B/hop"
         )
     else:
-        M_KV = 2 * S_local * hkv_per_pe * dh * b
+        M_KV = 2 * B * S_local * hkv_per_pe * dh * b
         variant_desc = "K/V ring"
         formula_bytes = (
-            f"2*S_local*(H_kv/TP)*d_h*b "
-            f"= 2*{S_local}*{hkv_per_pe}*{dh}*{b} "
+            f"2*B*S_local*(H_kv/TP)*d_h*b "
+            f"= 2*{B}*{S_local}*{hkv_per_pe}*{dh}*{b} "
             f"= {M_KV/1e6:.3f} MB/hop"
         )
 
@@ -501,7 +538,8 @@ def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    bytes_ = T_q * d * b
+    B = max(1, cfg.topo.b)
+    bytes_ = B * T_q * d * b
     tp = cfg.topo.tp
     tier = cfg.topo.tp_link_tier()  # "intra" | "inter" | "intersip"
     if tier == "intra":
@@ -514,7 +552,8 @@ def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
     total_comm_bytes = int(2 * (tp - 1) / tp * bytes_)
     return StageCost(
         name=f"C2 TP AllReduce W_O (TP={tp} ranks, {scope})",
-        formula=f"2*(TP-1)/TP * {T_q}*{d}*{b} B / BW + 2(TP-1)*alpha [{scope}]",
+        formula=f"2*(TP-1)/TP * B*T_q*d*b B / BW + 2(TP-1)*alpha "
+                f"[B={B}, {scope}]",
         compute_s=0, memory_s=0, comm_s=comm_time,
         bound="comm", visible_s=comm_time,
         flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
@@ -528,7 +567,8 @@ def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
             f"every rank ends up with the full hidden vector for the next\n"
             f"stage (RMSNorm -> FFN). Fires ONCE per layer over {scope} links.\n"
             f"---\n"
-            f"2*(TP-1)/TP * T_q*d*b = 2*({tp}-1)/{tp} * {T_q}*{d}*{b} "
+            f"2*(TP-1)/TP * B*T_q*d*b "
+            f"= 2*({tp}-1)/{tp} * {B}*{T_q}*{d}*{b} "
             f"= {total_comm_bytes/1e6:.2f} MB over {scope} at "
             f"{bw/1e9:.0f} GB/s"
         ),
@@ -552,8 +592,9 @@ def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost:
     S_local = cfg.topo.s_local
     hq_per_pe = cfg.h_q_per_pe
     b = cfg.model.bytes_per_elem
-    # score bytes per rank per hop
-    bytes_per_hop = hq_per_pe * T_q * S_local * b
+    B = max(1, cfg.topo.b)
+    # score bytes per rank per hop — scales with B (per-request scores).
+    bytes_per_hop = B * hq_per_pe * T_q * S_local * b
     # split-group AllReduce (intra-cube assumed if group fits)
     bw = cfg.machine.bw_intra
     alpha = cfg.machine.alpha_intra
@@ -562,7 +603,9 @@ def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost:
     total_comm_bytes = int(2 * (split - 1) / split * bytes_per_hop * cfg.topo.cp)
     return StageCost(
         name=f"C3 Score AllReduce ({split}-way, xCP hops)",
-        formula=f"2*({split}-1)/{split} * {hq_per_pe}*{T_q}*{S_local}*{b} / BW + latency",
+        formula=f"2*({split}-1)/{split} * B*(H_q/TP)*T_q*S_local*b / BW "
+                f"= 2*({split}-1)/{split} * {B}*{hq_per_pe}*{T_q}*{S_local}*{b} "
+                f"/ BW + latency",
         compute_s=0, memory_s=0, comm_s=total,
         bound="comm", visible_s=total,
         flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
@@ -626,17 +669,20 @@ def stage_ffn_rmsnorm(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    bytes_ = T_q * d * b * 2
-    flops = 4 * T_q * d
+    B = max(1, cfg.topo.b)
+    # Activation memory scales with B; weight (once) is fixed.
+    bytes_ = B * T_q * d * b + T_q * d * b
+    flops = 4 * B * T_q * d
     mem_s = bytes_ / cfg.machine.bw_hbm
     return StageCost(
         name="F1 RMSNorm (pre-FFN)",
-        formula=f"{T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM",
+        formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM",
         compute_s=0, memory_s=mem_s, comm_s=0,
         bound="memory", visible_s=mem_s,
         flops=flops, mem_bytes=bytes_,
-        flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}",
-        mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B",
+        flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}",
+        mem_formula=f"B*T_q*d*b + T_q*d*b (weight) "
+                     f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B",
     )
 
 
@@ -644,18 +690,24 @@ def _ffn_gemm(cfg: FullConfig, name: str, ffn_per_pe: int) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    flops = 2 * T_q * d * ffn_per_pe
+    B = max(1, cfg.topo.b)
+    # FLOPs scale with batch; weight (shared) is fixed.
+    flops = 2 * B * T_q * d * ffn_per_pe
     weight_B = d * ffn_per_pe * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name=name,
-        formula=f"FLOPs = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB",
+        formula=f"FLOPs = 2*B*T_q*d*(ffn/div) "
+                f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0,
         bound=bnd, visible_s=vis,
         flops=int(flops), mem_bytes=int(weight_B),
-        flops_formula=f"2*T_q*d*(ffn/div) = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}",
-        mem_formula=f"d*(ffn/div)*b = {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB",
+        flops_formula=f"2*B*T_q*d*(ffn/div) "
+                       f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}",
+        mem_formula=f"d*(ffn/div)*b (weight, B-invariant) "
+                     f"= {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB",
     )
 
 
@@ -674,17 +726,20 @@ def stage_ffn_swiglu(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
     b = cfg.model.bytes_per_elem
-    bytes_ = 3 * T_q * ffn_per_pe * b
-    flops = 3 * T_q * ffn_per_pe  # gate * silu(up) approximated as 3 flops/elt
+    B = max(1, cfg.topo.b)
+    bytes_ = 3 * B * T_q * ffn_per_pe * b
+    flops = 3 * B * T_q * ffn_per_pe  # gate * silu(up) approximated as 3 flops/elt
     mem_s = bytes_ / cfg.machine.bw_hbm
     return StageCost(
         name="F4 SwiGLU act",
-        formula=f"3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM",
+        formula=f"3*B*T_q*(ffn/div)*b "
+                f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM",
         compute_s=0, memory_s=mem_s, comm_s=0,
         bound="memory", visible_s=mem_s,
         flops=flops, mem_bytes=bytes_,
-        flops_formula=f"~3*T_q*(ffn/div) = 3*{T_q}*{ffn_per_pe} = {flops}",
-        mem_formula=f"3*T_q*(ffn/div)*b = 3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B",
+        flops_formula=f"~3*B*T_q*(ffn/div) = 3*{B}*{T_q}*{ffn_per_pe} = {flops}",
+        mem_formula=f"3*B*T_q*(ffn/div)*b "
+                     f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B",
     )
 
 
@@ -693,19 +748,24 @@ def stage_ffn_down(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
-    flops = 2 * T_q * ffn_per_pe * d
+    flops = 2 * B * T_q * ffn_per_pe * d
     weight_B = ffn_per_pe * d * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name="F5 W_down GEMM",
-        formula=f"FLOPs = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB",
+        formula=f"FLOPs = 2*B*T_q*(ffn/div)*d "
+                f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0,
         bound=bnd, visible_s=vis,
         flops=int(flops), mem_bytes=int(weight_B),
-        flops_formula=f"2*T_q*(ffn/div)*d = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}",
-        mem_formula=f"(ffn/div)*d*b = {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB",
+        flops_formula=f"2*B*T_q*(ffn/div)*d "
+                       f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}",
+        mem_formula=f"(ffn/div)*d*b (weight, B-invariant) "
+                     f"= {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB",
     )
 
 
@@ -723,7 +783,9 @@ def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    bytes_ = T_q * d * b
+    B = max(1, cfg.topo.b)
+    # AR reduces the batched FFN output — bytes scale with B.
+    bytes_ = B * T_q * d * b
     # Choose BW/alpha tier based on scope (rough): TP=intra-cube (or inter-cube),
     # +CP=inter-SIP possible, +DP=inter-SIP always.
     if "DP" in scope or "CP" in scope:
@@ -736,7 +798,8 @@ def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
     total_comm_bytes = int(2 * (divisor - 1) / divisor * bytes_)
     return StageCost(
         name=f"CF1 FFN AllReduce (scope={scope}, x{divisor})",
-        formula=f"2*({divisor}-1)/{divisor} * {T_q}*{d}*{b} B / BW + latency",
+        formula=f"2*({divisor}-1)/{divisor} * B*T_q*d*b B / BW + latency "
+                f"[B={B}]",
         compute_s=0, memory_s=0, comm_s=comm_time,
         bound="comm", visible_s=comm_time,
         flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
@@ -749,8 +812,8 @@ def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
             f"rank has the full FFN output for the next layer. Larger scope\n"
             f"= less FFN weight memory per PE but bigger AR bill.\n"
             f"---\n"
-            f"2*(div-1)/div * T_q*d*b = 2*({divisor}-1)/{divisor} * "
-            f"{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at "
+            f"2*(div-1)/div * B*T_q*d*b = 2*({divisor}-1)/{divisor} * "
+            f"{B}*{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at "
             f"{bw/1e9:.0f} GB/s (scope={scope})"
         ),
     )

tests/analytical_visualization/stage_shapes.py

@@ -0,0 +1,195 @@
+"""Per-stage tensor shape rows for the analytical visualization.
+
+Complements stage_latencies.py: while that module reports FLOPs/bytes/
+time per stage, this one reports the INPUT / WEIGHT / OUTPUT tensor
+shapes each stage operates on, per PE. Used for the 'per-stage shape'
+tables in the Streamlit app.
+
+Shape strings substitute the current cfg's numeric values so the table
+reads like an inspection of the actual deployment (no symbolic-only
+form).
+"""
+from __future__ import annotations
+
+from .model_config import FullConfig
+
+
+def _s(shape: tuple) -> str:
+    """Render a tuple of dim values as '(a, b, c)'."""
+    return "(" + ", ".join(str(x) for x in shape) + ")"
+
+
+def attn_stage_shape_rows(cfg: FullConfig) -> list[dict]:
+    """Per-PE attention shapes, one row per stage.
+
+    Rows follow the same S1..S10, C1..C3 order the latency table uses,
+    with the same conditional insertions (C1 only in prefill+CP>1, C3
+    only when kv_shard_mode='split' and TP > H_kv, C2 only when TP>1,
+    S8 merge only when CP>1).
+    """
+    m = cfg.model
+    t = cfg.topo
+    B = max(1, t.b)
+    T_q = t.T_q
+    S_local = t.s_local
+    d = m.hidden
+    dh = m.d_head
+    hq_pe = cfg.h_q_per_pe
+    hkv_pe = max(1, m.h_kv // t.tp)
+
+    rows: list[dict] = []
+
+    rows.append({
+        "Stage": "S1", "Op": "RMSNorm",
+        "Input (per PE)": _s((B, T_q, d)),
+        "Weight (per PE)": _s((d,)),
+        "Output (per PE)": _s((B, T_q, d)),
+    })
+    rows.append({
+        "Stage": "S2", "Op": "W_Q GEMM",
+        "Input (per PE)": _s((B, T_q, d)),
+        "Weight (per PE)": _s((d, hq_pe * dh)),
+        "Output (per PE)": _s((B, T_q, hq_pe * dh)),
+    })
+    rows.append({
+        "Stage": "S3", "Op": "W_K + W_V GEMM",
+        "Input (per PE)": _s((B, T_q, d)),
+        "Weight (per PE)": f"{_s((d, hkv_pe * dh))} x2 (K, V)",
+        "Output (per PE)": f"{_s((B, T_q, hkv_pe * dh))} x2 (K, V)",
+    })
+    rows.append({
+        "Stage": "S4", "Op": "KV cache append",
+        "Input (per PE)": f"{_s((B, T_q, hkv_pe * dh))} x2 (K, V)",
+        "Weight (per PE)": "-",
+        "Output (per PE)": (
+            f"KV cache: {_s((B, S_local, hkv_pe * dh))} x2 "
+            f"(extends by T_q={T_q} tokens)"
+        ),
+    })
+    rows.append({
+        "Stage": "S5", "Op": "Q · K^T",
+        "Input (per PE)": (
+            f"Q={_s((B, hq_pe, T_q, dh))}, "
+            f"K^T={_s((B, hkv_pe, dh, S_local))}"
+        ),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, S_local)),
+    })
+    rows.append({
+        "Stage": "S6", "Op": "softmax",
+        "Input (per PE)": _s((B, hq_pe, T_q, S_local)),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, S_local)),
+    })
+    rows.append({
+        "Stage": "S7", "Op": "P · V",
+        "Input (per PE)": (
+            f"P={_s((B, hq_pe, T_q, S_local))}, "
+            f"V={_s((B, hkv_pe, S_local, dh))}"
+        ),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, dh)),
+    })
+
+    if t.mode == "prefill" and t.cp > 1:
+        variant = t.cp_ring_variant
+        if variant == "kv":
+            payload = (
+                f"K,V shards rotate: {_s((B, S_local, hkv_pe, dh))} x2 per hop"
+            )
+        else:
+            payload = (
+                f"Q + running (O,m,l) rotate: {_s((B, hq_pe, T_q, dh))} per hop"
+            )
+        rows.append({
+            "Stage": "C1", "Op": f"CP ring ({variant})",
+            "Input (per PE)": payload,
+            "Weight (per PE)": "-",
+            "Output (per PE)": "(circulating; final owner reduces)",
+        })
+
+    if t.kv_shard_mode == "split" and t.tp > m.h_kv:
+        rows.append({
+            "Stage": "C3", "Op": "Score AllReduce (head-split)",
+            "Input (per PE)": _s((B, hq_pe, T_q, S_local)),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, hq_pe, T_q, S_local)),
+        })
+
+    if t.cp > 1:
+        rows.append({
+            "Stage": "S8", "Op": "online-softmax merge",
+            "Input (per PE)": (
+                f"O={_s((B, hq_pe, T_q, dh))}, "
+                f"m,l={_s((B, hq_pe, T_q))} each"
+            ),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, hq_pe, T_q, dh)),
+        })
+
+    rows.append({
+        "Stage": "S9", "Op": "normalize O / l",
+        "Input (per PE)": _s((B, hq_pe, T_q, dh)),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, dh)),
+    })
+    rows.append({
+        "Stage": "S10", "Op": "W_O GEMM",
+        "Input (per PE)": _s((B, T_q, hq_pe * dh)),
+        "Weight (per PE)": _s((hq_pe * dh, d)),
+        "Output (per PE)": _s((B, T_q, d)),
+    })
+
+    if t.tp > 1:
+        rows.append({
+            "Stage": "C2", "Op": f"TP AllReduce (W_O, TP={t.tp})",
+            "Input (per PE)": _s((B, T_q, d)),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, T_q, d)),
+        })
+
+    return rows
+
+
+def ffn_stage_shape_rows(cfg: FullConfig) -> list[dict]:
+    """Per-PE FFN shapes, one row per stage."""
+    m = cfg.model
+    t = cfg.topo
+    B = max(1, t.b)
+    T_q = t.T_q
+    d = m.hidden
+    divisor = cfg.ffn_shard_divisor * max(1, t.ep)
+    ffn_pe = max(1, m.ffn_dim // divisor)
+    scope = t.ffn_shard_scope
+
+    rows: list[dict] = [
+        {"Stage": "F1", "Op": "RMSNorm (pre-FFN)",
+         "Input (per PE)": _s((B, T_q, d)),
+         "Weight (per PE)": _s((d,)),
+         "Output (per PE)": _s((B, T_q, d))},
+        {"Stage": "F2", "Op": "W_gate GEMM",
+         "Input (per PE)": _s((B, T_q, d)),
+         "Weight (per PE)": _s((d, ffn_pe)),
+         "Output (per PE)": _s((B, T_q, ffn_pe))},
+        {"Stage": "F3", "Op": "W_up GEMM",
+         "Input (per PE)": _s((B, T_q, d)),
+         "Weight (per PE)": _s((d, ffn_pe)),
+         "Output (per PE)": _s((B, T_q, ffn_pe))},
+        {"Stage": "F4", "Op": "SwiGLU (gate * silu(up))",
+         "Input (per PE)": f"gate, up = {_s((B, T_q, ffn_pe))} each",
+         "Weight (per PE)": "-",
+         "Output (per PE)": _s((B, T_q, ffn_pe))},
+        {"Stage": "F5", "Op": "W_down GEMM",
+         "Input (per PE)": _s((B, T_q, ffn_pe)),
+         "Weight (per PE)": _s((ffn_pe, d)),
+         "Output (per PE)": _s((B, T_q, d))},
+    ]
+    if cfg.ffn_shard_divisor > 1:
+        rows.append({
+            "Stage": "CF1",
+            "Op": f"FFN AllReduce (scope={scope}, x{cfg.ffn_shard_divisor})",
+            "Input (per PE)": _s((B, T_q, d)),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, T_q, d)),
+        })
+    return rows

tests/analytical_visualization/tensor_sharding.py

@@ -279,12 +279,38 @@ def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0,
                   row_label=f"rows split by {ffn_txt}",
                   note="row-parallel")
 
-    # ── KV cache: (S_kv, H_kv*d_head), 2D shard ─
+    # ── KV cache: (B, S_kv, H_kv*d_head), 2D shard on S_kv x heads ─
+    # Batch (B) is a stacking dimension — each concurrent request holds
+    # its own KV slice per PE, so total KV per PE = B * (per-slice size).
     kv_row_splits = cp
     kv_col_splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
+    _B = max(1, cfg.topo.b)
+    _batch_suffix = f", B={_B}" if _B > 1 else ""
+    _batch_note = (f"; batch B={_B} => {_B}x KV bytes per PE"
+                    if _B > 1 else "")
+    # Draw B-1 offset "shadow" rectangles behind the KV cache to give
+    # a visual hint of the batch stacking dimension. Cap at 5 shadows
+    # so a large B doesn't overwhelm the diagram.
+    if _B > 1:
+        _kv_shadow_x = px(3)
+        _kv_shadow_y = py(0)
+        _n_shadows = min(_B - 1, 5)
+        _offset_step = 0.15
+        for _si in range(_n_shadows, 0, -1):
+            _dx = _si * _offset_step
+            _dy = _si * _offset_step
+            _shadow = patches.Rectangle(
+                (_kv_shadow_x + _dx, _kv_shadow_y + _dy),
+                cell_w, cell_h,
+                facecolor="#f5f5f5",
+                edgecolor="#adb5bd",
+                linewidth=0.8, linestyle="--", alpha=0.55,
+            )
+            ax.add_patch(_shadow)
     _draw_tensor(ax, px(3), py(0), cell_w, cell_h,
                   "KV cache (K and V)",
-                  f"({m.h_kv * m.d_head}, S_kv={cfg.topo.s_kv:,})",
+                  f"({m.h_kv * m.d_head}, S_kv={cfg.topo.s_kv:,}"
+                  f"{_batch_suffix})",
                   row_splits=kv_row_splits,
                   col_splits=kv_col_splits if show_physical else 1,
                   my_row_shard=kv_row_shard,
@@ -292,7 +318,8 @@ def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0,
                   row_label=f"rows: S split by {cp_txt}",
                   col_label=f"cols: heads split by {tp_txt}",
                   col_line_color=_TP_LINE_COLOR,
-                  note="2D shard: seq axis (CP) x head axis (TP)",
+                  note=f"2D shard: seq axis (CP) x head axis (TP)"
+                        f"{_batch_note}",
                   cell_annot=_kv_ann)
     # In non-physical mode, KV was drawn with col_splits=1; overlay TP col
     # splits as dotted lines to hint at the 2D nature. In physical mode
@@ -311,8 +338,11 @@ def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0,
     ax.set_aspect("auto")
     ax.axis("off")
 
+    _B_title = max(1, cfg.topo.b)
+    _b_title_suffix = f", B={_B_title}" if _B_title > 1 else ""
     ax.set_title(
-        f"Tensor sharding view (PE {my_pe})  |  TP={tp}, CP={cp}, "
+        f"Tensor sharding view (PE {my_pe})  |  TP={tp}, CP={cp}"
+        f"{_b_title_suffix}, "
         f"FFN scope={cfg.topo.ffn_shard_scope}  |  "
         f"blue = this PE's shard",
         fontsize=11, fontweight="bold",

tests/analytical_visualization/test_auto_explore.py

@@ -19,8 +19,9 @@ from tests.analytical_visualization.auto_explore import (
     run_auto_explore,
     score_config,
 )
+from tests.analytical_visualization.autosuggest import auto_suggest
 from tests.analytical_visualization.model_config import (
-    FullConfig, MachineParams,
+    FullConfig, MachineParams, TopologyConfig,
 )
 from tests.analytical_visualization.model_presets import PRESETS
 
@@ -208,6 +209,39 @@ def test_ffn_only_latency_less_than_full_and_different_from_attn():
     )
 
 
+def test_auto_suggest_cubes_match_default_topology():
+    """A TopologyConfig built from auto_suggest's returned (cp, tp, pp,
+    cp_placement) with all other fields left at their dataclass defaults
+    must produce the same cubes_used the Suggestion reports. Guards the
+    sidebar apply: any session_state key auto_suggest doesn't explore
+    (tp_placement, kv_shard_mode, cp_ring_variant, ep, ffn_shard_scope)
+    must be reset to the default before rerun, else the applied topology
+    won't match what the memory-min search actually scored.
+    """
+    machine = MachineParams()
+    matrix = [
+        ("Llama 3 70B", 8192),
+        ("Llama 3 70B", 32768),
+        ("Llama 3 8B",  131072),
+        ("Qwen 3 32B",  32768),
+    ]
+    for name, skv in matrix:
+        m = PRESETS[name].model
+        for ia, ff in [(True, False), (False, True), (True, True)]:
+            sug = auto_suggest(m, machine, s_kv=skv, mode="decode",
+                               include_attention=ia, include_ffn=ff)
+            topo = TopologyConfig(
+                cp=sug.cp, tp=sug.tp, pp=sug.pp,
+                s_kv=skv, mode="decode", b=1,
+                cp_placement=sug.cp_placement,
+            )
+            assert topo.cubes_used == sug.cubes_used, (
+                f"{name} skv={skv} scope=(ia={ia},ff={ff}): "
+                f"topo.cubes_used={topo.cubes_used} != "
+                f"sug.cubes_used={sug.cubes_used}"
+            )
+
+
 def test_parallelism_sensitivity_includes_baseline_value():
     """Each row's values contain the baseline_value."""
     from tests.analytical_visualization.auto_explore import (

tests/analytical_visualization/test_chip_roofline.py

@@ -0,0 +1,519 @@
+"""Tests for chip_roofline: AI, B*, L*, per-token latency curves."""
+from __future__ import annotations
+
+import math
+
+import pytest
+
+from tests.analytical_visualization.chip_roofline import (
+    ai_sensitivity_curve,
+    arithmetic_intensity,
+    balance_context,
+    bound_regime,
+    critical_batch,
+    good_batch,
+    good_context,
+    knee_batch,
+    kv_bytes_per_token,
+    max_batch_within_slo,
+    memory_budget_curve_vs_batch,
+    memory_budget_curve_vs_skv,
+    per_token_latency_curve,
+    size_deployment,
+    step_latency_curve,
+    t_com,
+    t_mem_long,
+    t_mem_short,
+    total_active_params,
+    utilization_at,
+)
+from tests.analytical_visualization.model_config import (
+    FullConfig, MachineParams, ModelConfig, TopologyConfig,
+)
+from tests.analytical_visualization.model_presets import PRESETS
+
+
+# ── Arithmetic intensity ────────────────────────────────────────────
+
+
+def test_arithmetic_intensity_matches_ratio():
+    """AI = peak_flops / bw_hbm — pure ratio, no util factor."""
+    m = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
+    assert arithmetic_intensity(m) == pytest.approx(8e12 / 256e9)
+
+
+def test_ai_scales_with_flops_and_bandwidth():
+    m1 = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
+    m2 = MachineParams(peak_tflops_f16=16.0, bw_hbm_gbs=256.0)
+    m3 = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=128.0)
+    assert arithmetic_intensity(m2) == pytest.approx(2 * arithmetic_intensity(m1))
+    assert arithmetic_intensity(m3) == pytest.approx(2 * arithmetic_intensity(m1))
+
+
+# ── Critical batch B* ───────────────────────────────────────────────
+
+
+def test_critical_batch_h100_reference():
+    """H100-class: 989 TFLOPs BF16, 3.35 TB/s HBM3 → B* ≈ 295."""
+    m = MachineParams(peak_tflops_f16=989.0, bw_hbm_gbs=3350.0)
+    model = PRESETS["Llama 3 8B"].model  # bf16 (b=2)
+    b_star = critical_batch(m, model)
+    assert b_star == pytest.approx(295, rel=0.05), b_star
+
+
+def test_critical_batch_scales_with_sparsity():
+    """MoE 8× sparsity gives 8× B*."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    b_dense = critical_batch(m, model, sparsity=1)
+    b_moe8 = critical_batch(m, model, sparsity=8)
+    assert b_moe8 == pytest.approx(8 * b_dense)
+
+
+def test_critical_batch_bf16_formula():
+    """For BF16 (b=2), B* = AI in flops-per-byte units → numerically."""
+    m = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
+    model = ModelConfig(bytes_per_elem=2)
+    ai = arithmetic_intensity(m)
+    assert critical_batch(m, model) == pytest.approx(ai), (
+        critical_batch(m, model), ai,
+    )
+
+
+# ── Balance context L* ──────────────────────────────────────────────
+
+
+def test_balance_context_positive_finite():
+    """L* for a real model on real machine is a positive finite number."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    l_star = balance_context(m, model)
+    assert 0 < l_star < 1e9, l_star
+
+
+def test_balance_context_scales_with_bandwidth():
+    """L* = 2N*W/(C*kv_bpt): doubling HBM BW doubles L*.
+    Slower memory hits the KV wall at a shorter context. Machine-only
+    change so N_active and kv_bpt stay fixed."""
+    fast = MachineParams(bw_hbm_gbs=512.0)
+    slow = MachineParams(bw_hbm_gbs=256.0)
+    model = PRESETS["Llama 3 8B"].model
+    assert (balance_context(fast, model)
+            == pytest.approx(2 * balance_context(slow, model)))
+
+
+# ── Knee ────────────────────────────────────────────────────────────
+
+
+def test_knee_equals_b_star_at_short_context():
+    """S_kv → 0: B_knee → B*."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    b_star = critical_batch(m, model)
+    assert knee_batch(m, model, s_kv=1) == pytest.approx(b_star, rel=1e-3)
+
+
+def test_knee_diverges_at_balance_context():
+    """S_kv = L*: knee is infinite; S_kv > L*: no knee (None)."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    l_star = balance_context(m, model)
+    assert knee_batch(m, model, s_kv=int(2 * l_star)) is None
+    # At 0.9 * L*, knee ~= 10 * B*; assert at least 5x to survive rounding.
+    below = knee_batch(m, model, s_kv=int(0.9 * l_star))
+    assert below is not None and below > 5 * critical_batch(m, model)
+
+
+def test_knee_slides_right_with_context():
+    """S_kv = L*/2 → B_knee = 2 * B*."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    b_star = critical_batch(m, model)
+    l_star = balance_context(m, model)
+    got = knee_batch(m, model, s_kv=int(l_star / 2))
+    assert got == pytest.approx(2 * b_star, rel=0.01), (got, 2 * b_star)
+
+
+# ── Per-token latency curve ─────────────────────────────────────────
+
+
+def test_latency_curve_monotonically_decreasing():
+    """total_s must strictly decrease as batch increases (weight/B shrinks)."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    pts = per_token_latency_curve(m, model, [1, 2, 4, 8, 16, 32, 64],
+                                    s_kv=1024)
+    totals = [p.total_s for p in pts]
+    for a, b in zip(totals, totals[1:]):
+        assert b < a, (a, b)
+
+
+def test_latency_curve_asymptotes_to_compute_plus_kv():
+    """At very large B, total → compute + KV (weight/B → 0)."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    pts = per_token_latency_curve(m, model, [10_000_000], s_kv=1024)
+    p = pts[0]
+    assert p.total_s == pytest.approx(p.compute_s + p.kv_s, rel=1e-3)
+
+
+def test_latency_at_b_star_is_roughly_2x_floor():
+    """At B*, weight = compute (dense, S_kv small), so total ≈ 2 * compute."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    b_star = int(round(critical_batch(m, model)))
+    pts = per_token_latency_curve(m, model, [b_star], s_kv=1)
+    p = pts[0]
+    # weight_s ≈ compute_s at B*
+    assert p.weight_s == pytest.approx(p.compute_s, rel=0.05), (
+        p.weight_s, p.compute_s,
+    )
+
+
+# ── Regime classification ───────────────────────────────────────────
+
+
+def test_regime_memory_bound_at_b1():
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    assert bound_regime(m, model, batch=1, s_kv=1024) == "memory-bound"
+
+
+def test_regime_kv_bound_past_balance_context():
+    """S_kv > L* with reasonable batch → KV dominates."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    l_star = balance_context(m, model)
+    b_star = int(round(critical_batch(m, model)))
+    assert bound_regime(m, model, batch=b_star * 4,
+                         s_kv=int(3 * l_star)) == "kv-bound"
+
+
+# ── Component sanity ────────────────────────────────────────────────
+
+
+def test_total_active_params_llama3_8b_ballpark():
+    """Llama 3 8B — attn+FFN alone (no embeddings/LM head) is ~6.9B,
+    HF-reported total 8.03B. Guard the ballpark."""
+    n = total_active_params(PRESETS["Llama 3 8B"].model)
+    assert 6.5e9 < n < 8e9, n
+
+
+def test_kv_bytes_per_token_llama3_8b():
+    """Llama 3 8B: 2*8*128*2 bytes/layer/token × 32 layers = 131072 bytes."""
+    kv_bpt = kv_bytes_per_token(PRESETS["Llama 3 8B"].model)
+    assert kv_bpt == 2 * 8 * 128 * 2 * 32
+
+
+# ── Regime terms + good-batch / good-context recommendations ──────
+
+
+def test_t_mem_short_shrinks_with_batch():
+    """t_mem_short = N·b/(W·B) — doubling B halves it."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    assert t_mem_short(m, model, 2) == pytest.approx(
+        t_mem_short(m, model, 1) / 2
+    )
+
+
+def test_t_mem_long_equals_t_com_at_l_star():
+    """L* is defined by t_mem_long(L*) == t_com. Cross-check."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    l_star = balance_context(m, model)
+    assert t_mem_long(m, model, int(round(l_star))) == pytest.approx(
+        t_com(m, model), rel=0.01,
+    )
+
+
+def test_t_com_is_batch_independent():
+    """t_com = 2·N/C. Constant — doesn't depend on B or S_kv."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    tc = t_com(m, model)
+    assert tc == pytest.approx(2 * total_active_params(model) / m.peak_flops)
+
+
+def test_good_batch_is_two_b_star():
+    """Recommendation is 2 × B* (Pope's rule)."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    rec = good_batch(m, model)
+    assert rec.effective == pytest.approx(2 * critical_batch(m, model))
+    assert rec.b_star == pytest.approx(critical_batch(m, model))
+
+
+def test_good_batch_moe_scales_with_sparsity():
+    """MoE sparsity=8 → recommended B is 8× dense recommendation."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    r_dense = good_batch(m, model, sparsity=1.0)
+    r_moe = good_batch(m, model, sparsity=8.0)
+    assert r_moe.effective == pytest.approx(8 * r_dense.effective)
+
+
+def test_good_context_returns_l_star_as_ceiling():
+    """Recommendation is L*; utilization drops past it."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    rec = good_context(m, model, s_kv=100)
+    assert rec.l_star == pytest.approx(balance_context(m, model))
+    assert rec.max_efficient == pytest.approx(rec.l_star)
+
+
+def test_utilization_at_l_star_is_half():
+    """At S_kv = L*, compute and KV read take equal time → 50% util."""
+    l_star = 3000.0
+    assert utilization_at(int(l_star), l_star) == pytest.approx(0.5)
+
+
+def test_utilization_at_2_l_star_is_one_third():
+    """At S_kv = 2·L*, util = 1/(1+2) = 33.3%."""
+    l_star = 3000.0
+    assert utilization_at(int(2 * l_star), l_star) == pytest.approx(1 / 3, rel=1e-3)
+
+
+def test_utilization_at_zero_context_is_100_percent():
+    """Zero context, no KV to read → all time is compute."""
+    assert utilization_at(0, 3000.0) == pytest.approx(1.0)
+
+
+def test_utilization_drops_monotonically_with_context():
+    """util(S_kv) is strictly decreasing."""
+    l_star = 3000.0
+    vals = [utilization_at(s, l_star) for s in [100, 1000, 3000, 6000, 30000]]
+    for a, b in zip(vals, vals[1:]):
+        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)
+
+
+# ── PE memory budget sweeps ────────────────────────────────────────
+
+
+def _budget_cfg():
+    """Helper: real cfg with enough sharding to leave headroom."""
+    return FullConfig(
+        model=PRESETS["Llama 3 8B"].model,
+        topo=TopologyConfig(cp=2, tp=8, pp=1, b=1, s_kv=8192),
+        machine=MachineParams(),
+    )
+
+
+def test_kv_grows_linear_with_skv():
+    cfg = _budget_cfg()
+    pts = memory_budget_curve_vs_skv(cfg, s_kv_range=[1024, 2048, 8192],
+                                       batch=1)
+    assert pts[1].kv_gb == pytest.approx(2 * pts[0].kv_gb, rel=1e-3)
+    assert pts[2].kv_gb == pytest.approx(8 * pts[0].kv_gb, rel=1e-3)
+
+
+def test_kv_grows_linear_with_batch():
+    cfg = _budget_cfg()
+    pts = memory_budget_curve_vs_batch(cfg, b_range=[1, 2, 8],
+                                          s_kv=8192)
+    assert pts[1].kv_gb == pytest.approx(2 * pts[0].kv_gb, rel=1e-3)
+    assert pts[2].kv_gb == pytest.approx(8 * pts[0].kv_gb, rel=1e-3)
+
+
+def test_weights_flat_across_skv_sweep():
+    """Weight bytes are batch/S_kv-invariant — sharding fixed."""
+    cfg = _budget_cfg()
+    pts = memory_budget_curve_vs_skv(cfg, s_kv_range=[128, 8192, 1_000_000],
+                                       batch=1)
+    ws = [p.weights_gb for p in pts]
+    assert all(w == pytest.approx(ws[0]) for w in ws), ws
+
+
+def test_over_budget_flag_trips_past_hbm():
+    """Push S_kv until KV blows the HBM budget; over_budget must flip."""
+    cfg = _budget_cfg()
+    pts = memory_budget_curve_vs_skv(
+        cfg, s_kv_range=[1024, 10_000_000], batch=64,
+    )
+    # 10M tokens × B=64 KV should overflow the 6 GB budget.
+    assert not pts[0].over_budget, pts[0]
+    assert pts[-1].over_budget, pts[-1]
+
+
+def test_free_gb_never_negative():
+    cfg = _budget_cfg()
+    pts = memory_budget_curve_vs_skv(
+        cfg, s_kv_range=[128, 10_000_000], batch=32,
+    )
+    for p in pts:
+        assert p.free_gb >= 0, p
+
+
+# ── AI sensitivity to FLOPs / BW ──────────────────────────────────
+
+
+def test_ai_scales_linearly_with_flops():
+    """Doubling FLOPs doubles AI."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    pts = ai_sensitivity_curve(m, model, [1.0, 2.0, 4.0], axis="flops")
+    assert pts[1].ai == pytest.approx(2 * pts[0].ai)
+    assert pts[2].ai == pytest.approx(4 * pts[0].ai)
+
+
+def test_ai_scales_inversely_with_bw():
+    """Doubling BW halves AI."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    pts = ai_sensitivity_curve(m, model, [1.0, 2.0], axis="bw")
+    assert pts[1].ai == pytest.approx(pts[0].ai / 2)
+
+
+def test_ai_sensitivity_b_star_tracks_ai_for_bf16():
+    """For BF16, B* == AI. Curve should mirror."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model  # BF16
+    pts = ai_sensitivity_curve(m, model, [0.5, 1.0, 4.0], axis="flops")
+    for p in pts:
+        assert p.b_star == pytest.approx(p.ai)
+
+
+def test_ai_sensitivity_bad_axis_raises():
+    with pytest.raises(ValueError):
+        ai_sensitivity_curve(MachineParams(), PRESETS["Llama 3 8B"].model,
+                              [1.0], axis="bogus")
+
+
+# ── GPU sizing (three-axis) ────────────────────────────────────────
+
+
+def test_capacity_floor_scales_with_model_size():
+    """A model with 2× params needs ~2× PEs for weights alone (axis A)."""
+    m = MachineParams()
+    small = PRESETS["Llama 3 8B"].model
+    big = PRESETS["Llama 3 70B"].model
+    r_small = size_deployment(m, small, n_users=1, avg_ctx=1024,
+                                tpot_slo_s=1.0)
+    r_big = size_deployment(m, big, n_users=1, avg_ctx=1024,
+                              tpot_slo_s=1.0)
+    ratio = r_big.pes_axis_a_capacity / r_small.pes_axis_a_capacity
+    # 70B / 8B ≈ 8.75× params — allow generous tolerance
+    assert 5 < ratio < 12, (ratio, r_small, r_big)
+
+
+def test_kv_headroom_scales_with_users_and_context():
+    """Doubling n_users OR avg_ctx grows KV load roughly linearly."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    base = size_deployment(m, model, n_users=64, avg_ctx=4096,
+                             tpot_slo_s=1.0)
+    twice_users = size_deployment(m, model, n_users=128, avg_ctx=4096,
+                                    tpot_slo_s=1.0)
+    twice_ctx = size_deployment(m, model, n_users=64, avg_ctx=8192,
+                                  tpot_slo_s=1.0)
+    # KV bytes doubled either way → axis-B PE count must strictly grow.
+    assert twice_users.pes_axis_b_kv > base.pes_axis_b_kv
+    assert twice_ctx.pes_axis_b_kv > base.pes_axis_b_kv
+
+
+def test_binding_axis_flips_with_workload():
+    """Huge N_users + short ctx → throughput binds.
+    Huge ctx + few users → kv (or capacity) binds."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    big_users = size_deployment(m, model, n_users=10_000, avg_ctx=1024,
+                                  tpot_slo_s=0.05)
+    long_ctx = size_deployment(m, model, n_users=1, avg_ctx=500_000,
+                                 tpot_slo_s=1.0)
+    assert big_users.binding_axis == "throughput", big_users
+    assert long_ctx.binding_axis in ("kv", "capacity"), long_ctx
+
+
+def test_max_batch_shrinks_with_shorter_slo():
+    """Tighter SLO → smaller B fits."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    loose = max_batch_within_slo(m, model, s_kv=1024, tpot_slo_s=1.0)
+    tight = max_batch_within_slo(m, model, s_kv=1024, tpot_slo_s=0.05)
+    assert tight < loose, (tight, loose)
+
+
+def test_max_batch_returns_zero_when_weight_time_exceeds_slo():
+    """If weight-fetch alone > SLO, no batch fits."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 70B"].model
+    # 70B in BF16 = 140 GB. Reading at 256 GB/s takes ~547 ms.
+    # SLO of 100 ms cannot be met.
+    assert max_batch_within_slo(m, model, s_kv=100,
+                                  tpot_slo_s=0.1) == 0
+
+
+def test_total_pes_equals_per_replica_times_replicas():
+    """Sanity: total_pes = pes_per_replica × n_replicas."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    r = size_deployment(m, model, n_users=500, avg_ctx=4096,
+                          tpot_slo_s=0.1)
+    assert r.total_pes == r.pes_per_replica * r.n_replicas
+
+
+# ── App wiring: tab exists on the Streamlit app ────────────────────
+
+
+def test_roofline_tab_registered_in_app():
+    """The new 'Chip roofline & B*' tab is listed in st.tabs and gets a
+    corresponding `with tab_roofline:` block."""
+    from pathlib import Path
+    src = (Path(__file__).parent / "app.py").resolve().read_text(
+        encoding="utf-8"
+    )
+    assert src.count('"Chip roofline & B*"') == 1
+    assert "with tab_roofline:" in src

tests/analytical_visualization/test_pe_weight_layout.py

@@ -0,0 +1,101 @@
+"""Tests for draw_pe_layout — CP rank locator visibility across TP tiers.
+
+The PE-level view has always color-labeled CP ranks when they're packed
+intra-cube (`cp_placement="pe"`). When TP >= 8, TP fills the cube and CP
+gets pushed to cube-level, so the intra-cube color branch is skipped —
+which used to silently drop the CP rank locator entirely. These tests
+guard the fix: a `[CP rank 0 of N]` cube tag + `CP=0` per-PE label are
+always present when CP > 1, regardless of placement.
+"""
+from __future__ import annotations
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from tests.analytical_visualization.model_config import (
+    FullConfig, MachineParams, TopologyConfig,
+)
+from tests.analytical_visualization.model_presets import PRESETS
+from tests.analytical_visualization.pe_weight_layout import draw_pe_layout
+
+
+def _cfg(cp: int, tp: int, cp_placement: str | None = None):
+    model = PRESETS["Llama 3 8B"].model
+    kwargs = dict(cp=cp, tp=tp, pp=1, s_kv=8192, mode="decode")
+    if cp_placement is not None:
+        kwargs["cp_placement"] = cp_placement
+    return FullConfig(
+        model=model, topo=TopologyConfig(**kwargs), machine=MachineParams(),
+    )
+
+
+def _all_text(fig) -> str:
+    """Concatenate all text artists (labels + title) into one blob."""
+    ax = fig.gca()
+    parts = [t.get_text() for t in ax.texts]
+    parts.append(ax.get_title())
+    return "\n".join(parts)
+
+
+def test_high_tp_cp_gt_1_shows_rank_locator():
+    """TP=8, CP=4 -> cp_placement forced to 'cube'; CP rank locator
+    must still appear in cube tag + per-PE header + figure title."""
+    cfg = _cfg(cp=4, tp=8)
+    assert cfg.topo.cp_placement == "cube"  # sanity
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        assert "CP rank 0 of 4" in text, text
+        assert "CP=0" in text, text
+        assert "CP=4" in text, text  # figure title dim summary
+    finally:
+        plt.close(fig)
+
+
+def test_intra_cube_cp_placement_still_shows_per_rank_labels():
+    """cp_placement='pe' (small TP, CP>1) — the old branch already
+    colors + labels each CP rank. Regression check: the new tag does
+    NOT replace those, and each CP=0..CP-1 label is still emitted."""
+    cfg = _cfg(cp=4, tp=2)  # 4*2=8 fits intra-cube
+    assert cfg.topo.cp_placement == "cube"  # default, but auto/manual can differ
+    # Force pe placement to exercise the intra-cube branch.
+    cfg.topo.cp_placement = "pe"
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        for r in range(4):
+            assert f"CP={r}" in text, (r, text)
+        # The 'CP rank 0 of N' tag is only for the cube-placement fallback.
+        assert "CP rank 0 of" not in text, text
+    finally:
+        plt.close(fig)
+
+
+def test_cp_1_no_locator_added():
+    """CP=1: nothing to locate; no CP tag or label should appear."""
+    cfg = _cfg(cp=1, tp=8)
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        assert "CP rank" not in text, text
+        # The header 'CP=0' should not be added when CP=1.
+        # (Figure title mentions 'CP=1' — that's an OK dim summary,
+        # not a locator per-PE.)
+        assert "  |  CP=0" not in text, text
+    finally:
+        plt.close(fig)
+
+
+def test_tp_16_still_shows_cp_rank_across_spilled_cubes():
+    """TP=16 (spills to 2 cubes per CP group), CP=2 -> both cubes
+    for CP rank 0 must carry the '[CP rank 0 of 2]' tag."""
+    cfg = _cfg(cp=2, tp=16)
+    assert cfg.topo.cp_placement == "cube"
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        assert text.count("CP rank 0 of 2") >= 2, text
+        assert "CP=0" in text, text
+    finally:
+        plt.close(fig)

tests/analytical_visualization/test_stage_shapes.py

@@ -0,0 +1,167 @@
+"""Tests for stage_shapes: per-PE shape rows for attention + FFN stages."""
+from __future__ import annotations
+
+from tests.analytical_visualization.model_config import (
+    FullConfig, MachineParams, TopologyConfig,
+)
+from tests.analytical_visualization.model_presets import PRESETS
+from tests.analytical_visualization.stage_shapes import (
+    attn_stage_shape_rows,
+    ffn_stage_shape_rows,
+)
+
+
+def _cfg(cp=2, tp=8, pp=1, b=1, mode="decode", s_kv=8192):
+    model = PRESETS["Llama 3 8B"].model
+    topo = TopologyConfig(cp=cp, tp=tp, pp=pp, b=b, s_kv=s_kv, mode=mode)
+    return FullConfig(model=model, topo=topo, machine=MachineParams())
+
+
+# ── Attention shapes ────────────────────────────────────────────────
+
+
+def test_attn_rows_have_required_columns():
+    """Every attention row must define Stage, Op, Input, Weight, Output."""
+    rows = attn_stage_shape_rows(_cfg())
+    required = {"Stage", "Op", "Input (per PE)",
+                "Weight (per PE)", "Output (per PE)"}
+    assert rows, "empty attention shape rows"
+    for r in rows:
+        missing = required - r.keys()
+        assert not missing, f"row {r.get('Stage')} missing {missing}"
+        for k in required:
+            assert r[k], f"row {r['Stage']} col {k} is empty"
+
+
+def test_attn_row_prefixes_present_in_order():
+    """S1..S10 always appear; C1 only when prefill+CP>1; C2 when TP>1;
+    C3 when kv_shard_mode=split and TP>H_kv; S8 only when CP>1."""
+    cfg = _cfg(cp=2, tp=8, mode="decode")
+    prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
+    # decode with CP=2, TP=8 (h_kv=8 for Llama 3 8B) → no C1, no C3; C2 yes; S8 yes.
+    assert "C1" not in prefixes, prefixes
+    assert "C3" not in prefixes, prefixes
+    assert "S8" in prefixes, prefixes
+    assert "C2" in prefixes, prefixes
+    # Ensure S1..S10 sequence appears in relative order.
+    s_order = [p for p in prefixes if p.startswith("S")]
+    assert s_order == sorted(s_order, key=lambda s: int(s[1:])), s_order
+
+
+def test_attn_cp1_omits_s8_merge():
+    """CP=1 → no online-softmax merge stage."""
+    cfg = _cfg(cp=1, tp=8, mode="decode")
+    prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
+    assert "S8" not in prefixes
+
+
+def test_attn_prefill_cp_gt_1_inserts_c1():
+    """Prefill mode with CP>1 shows the CP-ring stage."""
+    cfg = _cfg(cp=4, tp=2, mode="prefill", s_kv=8192)
+    prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
+    assert "C1" in prefixes
+
+
+def test_attn_shape_reflects_batch():
+    """B appears as the leading dim in each activation shape."""
+    for b in (1, 4, 32):
+        rows = attn_stage_shape_rows(_cfg(b=b))
+        s1 = next(r for r in rows if r["Stage"] == "S1")
+        assert s1["Input (per PE)"].startswith(f"({b},"), s1
+        assert s1["Output (per PE)"].startswith(f"({b},"), s1
+
+
+def test_attn_s5_shape_uses_tq_and_slocal():
+    """S5 (Q·K^T) shape must reference both T_q and S_local values."""
+    cfg = _cfg(cp=2, tp=8, mode="decode", s_kv=8192)
+    T_q = cfg.topo.T_q
+    S_local = cfg.topo.s_local
+    rows = attn_stage_shape_rows(cfg)
+    s5 = next(r for r in rows if r["Stage"] == "S5")
+    # Output shape (B, H_q/TP, T_q, S_local)
+    assert str(T_q) in s5["Output (per PE)"], s5
+    assert str(S_local) in s5["Output (per PE)"], s5
+
+
+# ── FFN shapes ──────────────────────────────────────────────────────
+
+
+def test_ffn_rows_have_required_columns():
+    """Every FFN row must define Stage, Op, Input, Weight, Output."""
+    rows = ffn_stage_shape_rows(_cfg())
+    required = {"Stage", "Op", "Input (per PE)",
+                "Weight (per PE)", "Output (per PE)"}
+    assert rows, "empty ffn shape rows"
+    for r in rows:
+        missing = required - r.keys()
+        assert not missing, f"row {r.get('Stage')} missing {missing}"
+
+
+def test_ffn_row_prefixes_f1_to_f5_always_present():
+    """F1..F5 always emitted; CF1 only when FFN divisor > 1."""
+    rows = ffn_stage_shape_rows(_cfg(cp=2, tp=8))
+    prefixes = [r["Stage"] for r in rows]
+    for p in ("F1", "F2", "F3", "F4", "F5"):
+        assert p in prefixes, prefixes
+
+
+def test_ffn_gate_up_output_matches_ffn_per_pe():
+    """F2/F3 output must have ffn_per_pe as the last dim."""
+    cfg = _cfg(cp=2, tp=8)
+    ffn_pe = max(1, cfg.model.ffn_dim
+                 // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+    rows = ffn_stage_shape_rows(cfg)
+    f2 = next(r for r in rows if r["Stage"] == "F2")
+    assert str(ffn_pe) in f2["Output (per PE)"], (ffn_pe, f2)
+
+
+# ── App wiring: consolidated shapes section lives on layout tab ───
+
+
+def test_shapes_section_is_on_physical_layout_tab():
+    """The 'All tensor shapes (per PE)' expander must live inside
+    tab_layout (Physical layout), not tab_stages. Guards the move so
+    the section doesn't accidentally end up back on the latency tab.
+    """
+    from pathlib import Path
+    app_path = (
+        Path(__file__).parent / "app.py"
+    ).resolve()
+    src = app_path.read_text(encoding="utf-8")
+    # Section appears exactly once.
+    assert src.count('"All tensor shapes (per PE)"') == 1, src.count(
+        '"All tensor shapes (per PE)"'
+    )
+    # And it sits inside tab_layout, before tab_memory begins.
+    section_pos = src.index('"All tensor shapes (per PE)"')
+    tab_memory_pos = src.index("with tab_memory:")
+    tab_stages_pos = src.index("with tab_stages:")
+    assert section_pos < tab_memory_pos, (
+        "shapes section must be inside tab_layout (before tab_memory)"
+    )
+    assert section_pos < tab_stages_pos, (
+        "shapes section must be inside tab_layout (before tab_stages)"
+    )
+    # The old per-stage shape rendering under tab_stages must be gone —
+    # no attn_stage_shape_rows / ffn_stage_shape_rows call may sit
+    # between 'with tab_stages:' and the next 'with tab_' block.
+    stages_block = src[tab_stages_pos:src.index("# ── TAB", tab_stages_pos + 1)]
+    assert "attn_stage_shape_rows(cfg)" not in stages_block, (
+        "per-stage shape table must not remain on Per-stage latency tab"
+    )
+    assert "ffn_stage_shape_rows(cfg)" not in stages_block, (
+        "per-stage FFN shape table must not remain on Per-stage latency tab"
+    )
+
+
+def test_ttft_tpot_section_wired_on_layout_tab():
+    """The 'Full-model latency (TTFT & TPOT)' subheader must live on
+    tab_layout, before tab_memory. Guard the section presence + placement."""
+    from pathlib import Path
+    src = (
+        Path(__file__).parent / "app.py"
+    ).resolve().read_text(encoding="utf-8")
+    assert src.count('"Full-model latency (TTFT & TPOT)"') == 1
+    pos = src.index('"Full-model latency (TTFT & TPOT)"')
+    tab_memory_pos = src.index("with tab_memory:")
+    assert pos < tab_memory_pos, "TTFT section must live inside tab_layout"

tests/analytical_visualization/topology_map.py

@@ -258,12 +258,27 @@ def _draw_one_sip(ax, cfg: FullConfig, sip_idx: int,
         pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
         pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
         pes_used = cfg.topo.pes_per_cube_used if is_used else 0
+        # When cp_placement=pe, multiple CP ranks live inside this cube's PEs
+        # (cp_rank = pe_id // tp). Color each PE by its cp_rank so all four
+        # groups are visible; otherwise fall back to the whole-cube pe_fill.
+        _cp_packed = is_used and cfg.topo.cp_placement == "pe" and cfg.topo.cp > 1
         for pr in range(PE_ROWS):
             for pc in range(PE_COLS):
                 pe_id = pr * PE_COLS + pc
                 px = x + pe_pad + pc * (pe_w + pe_gap)
                 py = y + pe_pad + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
-                fill = pe_fill if pe_id < pes_used else _INACTIVE_PE
+                if pe_id >= pes_used:
+                    fill = _INACTIVE_PE
+                elif _cp_packed:
+                    _pe_cp_rank = pe_id // cfg.topo.tp
+                    _pe_color, _pe_pe_fill = _cp_color(
+                        info[0] if is_used else 0,
+                        _pe_cp_rank,
+                        cfg.topo.cp,
+                    )
+                    fill = _pe_pe_fill
+                else:
+                    fill = pe_fill
                 pe_rect = patches.Rectangle(
                     (px, py), pe_w, pe_h,
                     facecolor=fill, edgecolor="#666", linewidth=0.3,