analytical-viz: interactive Streamlit tool for SIP transformer analysis

New tests/analytical_visualization/ module - an interactive dashboard
for exploring memory / latency tradeoffs of transformer inference on
the SIP architecture.

Highlights:
- 30+ model presets (Qwen, Llama 2/3/3.1, Mistral, Gemma 2, Phi 3,
  DeepSeek MLA, Mixtral/Qwen 3 MoE, Grok-1, ...)
- Placement toggles: TP and CP each on PE-level vs cube-level
- CP ring variant: K/V ring vs Q+O/m/l ring (prefill); in decode the
  O/m/l all-reduce is folded into S8 (no separate C1 row)
- SIP interconnect: ring / mesh2d / torus2d with matching link drawing
- Per-stage latency table with compute + memory + comm formulas,
  auto-scaled ns/us/ms, colored by dominant bound
- Ring attention loop indicator on the pipeline diagram (purple arc
  over S5-S8 with 'xN hops' badge)
- Tensor sharding view with optional physical PE/cube annotations
- Replication-waste + optimization-hints panel
- Save & compare configurations (config1, config2, ...): summary table
  plus side-by-side per-stage attention and FFN latency, best-in-row
  highlighting
- Symbol glossary with current values for every symbol used in formulas

Not tied to production sim_engine or runtime API; purely analytical
tooling for design-space exploration.

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

tests/analytical_visualization/autosuggest.py

@@ -0,0 +1,106 @@
+"""Auto-suggest (CP, TP, PP) that fits the model in the per-PE budget.
+
+Strategy: iterate over candidate (CP, TP, PP) triples in order of
+increasing total PE count and return the first one that satisfies
+memory + kernel-support constraints with >10% slack. If none fits, return
+the best-effort configuration and flag over-budget.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Iterable
+
+from .model_config import FullConfig, ModelConfig, TopologyConfig, MachineParams
+from .memory_layout import compute_memory
+
+
+# Candidate parallelism dimensions (powers of 2 mostly, up to sensible limits).
+_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
+_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
+_PP_OPTIONS = (1, 2, 4, 8, 16, 32)
+
+
+@dataclass
+class Suggestion:
+    cp: int
+    tp: int
+    pp: int
+    weights_gb: float
+    kv_gb: float
+    transient_gb: float
+    slack_gb: float
+    fits: bool
+    pes_used: int
+    sips_used: int
+    reason: str = ""
+
+
+def _iter_candidates(model: ModelConfig) -> Iterable[tuple[int, int, int]]:
+    """Yield (CP, TP, PP) in order of increasing PE count."""
+    triples: list[tuple[int, int, int]] = []
+    for tp in _TP_OPTIONS:
+        for cp in _CP_OPTIONS:
+            for pp in _PP_OPTIONS:
+                # PP must not exceed layer count.
+                if pp > model.layers:
+                    continue
+                # Skip TP > H_q * some factor (unrealistic).
+                if tp > model.h_q * 4:
+                    continue
+                triples.append((cp, tp, pp))
+    # Sort by pe_count then by (pp, tp, cp) — prefer smaller PP first
+    # (avoids pipeline bubbles), then smaller TP (avoids head-dim split).
+    triples.sort(key=lambda t: (t[0] * t[1] * t[2], t[2], t[1], t[0]))
+    for t in triples:
+        yield t
+
+
+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)
+    cfg = FullConfig(model=model, topo=topo, machine=machine)
+    mem = compute_memory(cfg)
+
+    fits = (not mem.over_budget
+            and mem.slack_bytes >= slack_frac * mem.budget_bytes)
+    reason = ""
+    if mem.over_budget:
+        reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
+                  f"exceeds budget ({mem.budget_bytes/1e9:.2f} GB)")
+    elif not fits:
+        reason = f"slack ({mem.slack_bytes/1e9:.2f} GB) below 10% of budget"
+    return Suggestion(
+        cp=cp, tp=tp, pp=pp,
+        weights_gb=mem.weights_bytes / 1e9,
+        kv_gb=mem.kv_cache_bytes / 1e9,
+        transient_gb=mem.transient_bytes / 1e9,
+        slack_gb=mem.slack_bytes / 1e9,
+        fits=fits,
+        pes_used=topo.total_pes,
+        sips_used=topo.sips_used,
+        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.
+
+    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
+
+    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
+
+    return best_fit or best_effort  # type: ignore

tests/analytical_visualization/memory_layout.py

@@ -0,0 +1,225 @@
+"""Per-PE memory footprint: weights + KV cache + activations + slack.
+
+All formulas are per-PE, per-PP-stage. If PP > 1, each PE holds only
+its stage's layers (layers/PP).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .model_config import FullConfig
+
+
+@dataclass
+class MemoryBreakdown:
+    weights_bytes: int
+    kv_cache_bytes: int
+    transient_bytes: int
+    budget_bytes: int
+
+    @property
+    def used_bytes(self) -> int:
+        return self.weights_bytes + self.kv_cache_bytes + self.transient_bytes
+
+    @property
+    def slack_bytes(self) -> int:
+        return max(0, self.budget_bytes - self.used_bytes)
+
+    @property
+    def over_budget(self) -> bool:
+        return self.used_bytes > self.budget_bytes
+
+
+def per_pe_weight_bytes(cfg: FullConfig) -> 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).
+    """
+    m = cfg.model
+    tp = cfg.topo.tp
+    pp = cfg.topo.pp
+    ep = max(1, cfg.topo.ep)
+
+    hq_per_pe = cfg.h_q_per_pe
+    if cfg.topo.kv_shard_mode == "replicate":
+        # Each KV head held fully; replicated across ranks if TP > H_kv.
+        hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
+    else:  # "split"
+        # Head-dim split: fractional head allowed (< 1 head bytes when TP > H_kv).
+        hkv_per_pe_bytes = m.h_kv / tp
+
+    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
+    )
+    # 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
+
+    layers_per_stage = (m.layers + pp - 1) // pp
+    return int(per_layer * layers_per_stage)
+
+
+def per_pe_kv_cache_bytes(cfg: FullConfig) -> int:
+    """K + V per PE, across all layers this stage holds.
+
+    - CP shards the sequence dim -> S_local = S_kv/CP tokens per PE.
+    - TP splits KV heads across ranks. If kv_shard_mode='replicate' and
+      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.
+    """
+    m = cfg.model
+    pp = cfg.topo.pp
+    tp = cfg.topo.tp
+    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)
+
+
+def per_pe_transient_bytes(cfg: FullConfig) -> int:
+    """Rough peak transient (activations, GEMM outputs) per PE."""
+    m = cfg.model
+    T_q = cfg.topo.T_q
+    hq_per_pe = cfg.h_q_per_pe
+
+    if cfg.topo.mode == "decode":
+        return 4 * (m.hidden + m.head_dim_total_q // cfg.topo.tp) * m.bytes_per_elem
+    else:
+        TILE = 1024
+        tile_score = hq_per_pe * T_q * TILE * m.bytes_per_elem
+        return int(2 * tile_score + T_q * m.hidden * m.bytes_per_elem)
+
+
+def compute_memory(cfg: FullConfig) -> MemoryBreakdown:
+    return MemoryBreakdown(
+        weights_bytes=per_pe_weight_bytes(cfg),
+        kv_cache_bytes=per_pe_kv_cache_bytes(cfg),
+        transient_bytes=per_pe_transient_bytes(cfg),
+        budget_bytes=cfg.machine.pe_budget_bytes,
+    )
+
+
+def _one_layer_row(name: str,
+                   global_shape: tuple[int, int],
+                   per_pe_shape: tuple[int, int],
+                   bytes_per_elem: int) -> dict:
+    """Per-tensor row for ONE layer (both global and per-PE shard)."""
+    p_global = global_shape[0] * global_shape[1]
+    p_per_pe = per_pe_shape[0] * per_pe_shape[1]
+    return {
+        "Tensor": name,
+        "Global shape": f"({global_shape[0]}, {global_shape[1]})",
+        "Params/layer": f"{p_global/1e6:.2f} M",
+        "Bytes/layer (global)": f"{p_global * bytes_per_elem / 1e6:.2f} MB",
+        "Per-PE shape": f"({per_pe_shape[0]}, {per_pe_shape[1]})",
+        "Bytes/layer (per PE)": f"{p_per_pe * bytes_per_elem / 1e6:.2f} MB",
+        "_p_global": p_global,
+        "_p_per_pe": p_per_pe,
+    }
+
+
+def attention_weight_rows(cfg: FullConfig) -> list[dict]:
+    """Per-tensor rows for attention weights (one layer)."""
+    m = cfg.model
+    hq_per_pe = cfg.h_q_per_pe
+    hkv_per_pe = max(1, m.h_kv // cfg.topo.tp)
+    return [
+        _one_layer_row("W_Q", (m.hidden, m.h_q * m.d_head),
+                       (m.hidden, hq_per_pe * m.d_head),
+                       m.bytes_per_elem),
+        _one_layer_row("W_K", (m.hidden, m.h_kv * m.d_head),
+                       (m.hidden, hkv_per_pe * m.d_head),
+                       m.bytes_per_elem),
+        _one_layer_row("W_V", (m.hidden, m.h_kv * m.d_head),
+                       (m.hidden, hkv_per_pe * m.d_head),
+                       m.bytes_per_elem),
+        _one_layer_row("W_O", (m.h_q * m.d_head, m.hidden),
+                       (hq_per_pe * m.d_head, m.hidden),
+                       m.bytes_per_elem),
+    ]
+
+
+def ffn_weight_rows(cfg: FullConfig) -> list[dict]:
+    """Per-tensor rows for FFN weights (one layer, activated for MoE)."""
+    m = cfg.model
+    ffn_per_pe = m.ffn_dim // cfg.topo.tp
+    return [
+        _one_layer_row("W_gate", (m.hidden, m.ffn_dim),
+                       (m.hidden, ffn_per_pe), m.bytes_per_elem),
+        _one_layer_row("W_up", (m.hidden, m.ffn_dim),
+                       (m.hidden, ffn_per_pe), m.bytes_per_elem),
+        _one_layer_row("W_down", (m.ffn_dim, m.hidden),
+                       (ffn_per_pe, m.hidden), m.bytes_per_elem),
+    ]
+
+
+def kv_cache_rows(cfg: FullConfig) -> list[dict]:
+    """Per-tensor row for K and V cache (one layer)."""
+    m = cfg.model
+    hkv_per_pe = max(1, m.h_kv // cfg.topo.tp)
+    b = m.bytes_per_elem
+    global_k = cfg.topo.s_kv * m.h_kv * m.d_head
+    per_pe_k = cfg.topo.s_local * hkv_per_pe * m.d_head
+    return [
+        {
+            "Tensor": "K cache",
+            "Global shape": f"({cfg.topo.s_kv:,}, {m.h_kv * m.d_head})",
+            "Params/layer": f"{global_k/1e6:.2f} M",
+            "Bytes/layer (global)": f"{global_k * b / 1e6:.2f} MB",
+            "Per-PE shape": f"({cfg.topo.s_local:,}, {hkv_per_pe * m.d_head})",
+            "Bytes/layer (per PE)": f"{per_pe_k * b / 1e6:.2f} MB",
+            "_p_global": global_k,
+            "_p_per_pe": per_pe_k,
+        },
+        {
+            "Tensor": "V cache",
+            "Global shape": f"({cfg.topo.s_kv:,}, {m.h_kv * m.d_head})",
+            "Params/layer": f"{global_k/1e6:.2f} M",
+            "Bytes/layer (global)": f"{global_k * b / 1e6:.2f} MB",
+            "Per-PE shape": f"({cfg.topo.s_local:,}, {hkv_per_pe * m.d_head})",
+            "Bytes/layer (per PE)": f"{per_pe_k * b / 1e6:.2f} MB",
+            "_p_global": global_k,
+            "_p_per_pe": per_pe_k,
+        },
+    ]
+
+
+def sum_bytes_all_layers(rows: list[dict], cfg: FullConfig,
+                          per_pe: bool = False) -> int:
+    """Sum bytes across all rows × N layers (or layers/PP for per_pe)."""
+    b = cfg.model.bytes_per_elem
+    layers = cfg.model.layers
+    if per_pe:
+        layers = (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp
+    key = "_p_per_pe" if per_pe else "_p_global"
+    return sum(r[key] for r in rows) * b * layers
+
+
+def total_weight_bytes_full_model(cfg: FullConfig) -> int:
+    """Total bytes of ALL weights (attention + FFN) for the full unsharded model."""
+    m = cfg.model
+    attn_per_layer = (
+        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_per_layer = 3 * m.hidden * m.ffn_dim
+    return (attn_per_layer + ffn_per_layer) * m.bytes_per_elem * m.layers
+
+
+def total_kv_bytes_full_model(cfg: FullConfig) -> int:
+    """Total bytes of KV cache for full unsharded model at cfg.topo.s_kv."""
+    m = cfg.model
+    per_layer = 2 * cfg.topo.s_kv * m.h_kv * m.d_head * m.bytes_per_elem
+    return per_layer * m.layers

tests/analytical_visualization/model_config.py

@@ -0,0 +1,310 @@
+"""Model + machine parameters for the analytical visualization tool.
+
+All numbers are per-rank / per-PE unless noted. bytes are counted in
+bytes (b), FLOPs in raw ops, bandwidth in bytes/second, latency in
+seconds. Convert to μs / GB / GFLOP at display time only.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ModelConfig:
+    """Transformer model dimensions."""
+    name: str = "Qwen 3 8B"
+    hidden: int = 4096
+    ffn_dim: int = 12288
+    h_q: int = 32               # Q heads
+    h_kv: int = 8               # KV heads
+    d_head: int = 128
+    layers: int = 36
+    bytes_per_elem: int = 2     # bf16
+
+    @property
+    def group_size(self) -> int:
+        """GQA group size = H_q / H_kv."""
+        return self.h_q // self.h_kv
+
+    @property
+    def head_dim_total_q(self) -> int:
+        return self.h_q * self.d_head
+
+    @property
+    def head_dim_total_kv(self) -> int:
+        return self.h_kv * self.d_head
+
+
+@dataclass
+class TopologyConfig:
+    """Physical mapping — DP × PP × CP × TP (× EP for MoE)."""
+    cp: int = 4                 # inter-cube ring size (sequence shard)
+    tp: int = 8                 # PEs per cube (head/hidden shard)
+    pp: int = 1                 # pipeline stages (layer shard)
+    dp: int = 1                 # data-parallel replicas
+    ep: int = 1                 # expert-parallel degree (MoE only)
+    s_kv: int = 4096
+    mode: str = "decode"        # "decode" or "prefill"
+    kv_shard_mode: str = "split"  # "split" or "replicate" (used when TP > H_kv)
+    # FFN sharding scope: default TP+CP (across cube group) trades a small
+    # extra AllReduce for lower per-PE weight bytes.
+    ffn_shard_scope: str = "TP+CP"   # "TP" | "TP+CP" | "TP+CP+DP"
+    # Inter-SIP interconnect: how physical SIPs are wired together.
+    # "ring" (1D wrap), "mesh2d" (grid, no wrap), "torus2d" (grid + wrap).
+    sip_topology: str = "ring"
+    # Placement of parallelism dims on the hierarchy.
+    # "pe"   -> dim lives on PEs within a single cube (intra-cube).
+    # "cube" -> dim lives across cubes (inter-cube).
+    tp_placement: str = "pe"     # default: TP fills PEs of a cube
+    cp_placement: str = "cube"   # default: CP fills cubes of a SIP
+    # CP ring variant: what gets passed around the CP ring each hop.
+    # "kv"    -> K and V shards rotate (Q,O,m,l stay local). Bytes/hop scales
+    #            with S_local * H_kv * d_h. Good for prefill (large T_q).
+    # "qoml"  -> Q and running (O, m, l) rotate; K,V stay local. Bytes/hop
+    #            scales with T_q * H_q * d_h. Cheap for decode (T_q=1).
+    cp_ring_variant: str = "kv"
+    pes_per_cube_hw: int = 8    # SIP hardware constant
+    cubes_per_sip_hw: int = 16
+
+    @property
+    def pes_per_stage(self) -> int:
+        """PEs per pipeline stage (all stages have the same size)."""
+        return self.cp * self.tp
+
+    @property
+    def pes_per_replica(self) -> int:
+        """PEs per one full model replica (CP*TP*PP)."""
+        return self.cp * self.tp * self.pp
+
+    @property
+    def total_pes(self) -> int:
+        return self.pes_per_replica * self.dp
+
+    # ── Placement-aware layout ────────────────────────────────
+    @property
+    def intra_cube_dims(self) -> int:
+        """Product of dims placed on PE level (must fit in one cube, else spills)."""
+        n = 1
+        if self.tp_placement == "pe":
+            n *= self.tp
+        if self.cp_placement == "pe":
+            n *= self.cp
+        return n
+
+    @property
+    def inter_cube_dims(self) -> int:
+        """Product of dims placed at cube level = distinct cube groups per stage."""
+        n = 1
+        if self.tp_placement == "cube":
+            n *= self.tp
+        if self.cp_placement == "cube":
+            n *= self.cp
+        return n
+
+    @property
+    def pes_per_cube_used(self) -> int:
+        """How many PEs are live in each cube."""
+        return min(self.intra_cube_dims, self.pes_per_cube_hw)
+
+    @property
+    def cubes_per_stage(self) -> int:
+        """Cubes needed per PP stage under current placement.
+        If intra-cube demand exceeds PEs/cube, the intra dim spills to cubes."""
+        spill = max(1, (self.intra_cube_dims + self.pes_per_cube_hw - 1)
+                    // self.pes_per_cube_hw)
+        return self.inter_cube_dims * spill
+
+    @property
+    def cubes_used(self) -> int:
+        """Physical cubes used across all PP stages and DP replicas."""
+        return self.cubes_per_stage * self.pp * self.dp
+
+    @property
+    def sips_used(self) -> int:
+        """Number of SIPs used (each SIP = pes_per_cube_hw × cubes_per_sip_hw)."""
+        cubes_per_sip = self.cubes_per_sip_hw
+        return max(1, (self.cubes_used + cubes_per_sip - 1) // cubes_per_sip)
+
+    @property
+    def placement_valid(self) -> bool:
+        """False when intra-cube demand exceeds one cube (dim spills to cubes)."""
+        return self.intra_cube_dims <= self.pes_per_cube_hw
+
+    @property
+    def tp_spans_cubes(self) -> int:
+        """How many cubes one TP group spans (>=1). Depends on placement:
+        - tp_placement=pe: 1 if TP fits in a cube, else spill.
+        - tp_placement=cube: TP itself is the number of cubes per TP group.
+        """
+        if self.tp_placement == "cube":
+            return max(1, self.tp)
+        return max(1, (self.tp + self.pes_per_cube_hw - 1) // self.pes_per_cube_hw)
+
+    @property
+    def cp_intra_sip_hops(self) -> int:
+        """CP ring hops that stay within one SIP."""
+        if self.cp <= 1:
+            return 0
+        sips_in_ring = self.sips_used  # one PP stage's CP ring may span SIPs
+        return max(0, (self.cp - 1) - (sips_in_ring - 1))
+
+    @property
+    def cp_inter_sip_hops(self) -> int:
+        """CP ring hops that cross SIP boundaries."""
+        if self.cp <= 1:
+            return 0
+        # Approximation: one inter-SIP hop per SIP boundary in the ring.
+        # For CP=32 across 2 SIPs → 1 inter-SIP + 30 intra-SIP hops.
+        return max(0, self.sips_used - 1)
+
+    @property
+    def T_q(self) -> int:
+        return 1 if self.mode == "decode" else self.s_kv // self.cp
+
+    @property
+    def s_local(self) -> int:
+        return self.s_kv // self.cp
+
+    @property
+    def layers_per_stage(self) -> str:
+        """Layers per PP stage. Requires model.layers to render — leave as fn."""
+        return "N_layers / PP"
+
+    def tp_link_tier(self) -> str:
+        """Return which BW tier a TP AllReduce uses: 'intra' | 'inter' | 'intersip'.
+        - tp_placement=pe and TP fits in one cube -> 'intra'
+        - tp_placement=cube and TP fits in one SIP -> 'inter'
+        - otherwise -> 'intersip'
+        """
+        if self.tp_placement == "pe":
+            if self.tp <= self.pes_per_cube_hw:
+                return "intra"
+            # spilled to multiple cubes: cross-cube
+            if self.tp_spans_cubes <= self.cubes_per_sip_hw:
+                return "inter"
+            return "intersip"
+        else:  # tp_placement == "cube"
+            if self.tp <= self.cubes_per_sip_hw:
+                return "inter"
+            return "intersip"
+
+    def cp_link_tier(self) -> str:
+        """Return BW tier for the CP ring: 'intra' | 'inter' | 'intersip'."""
+        if self.cp_placement == "pe":
+            return "intra"
+        # cp_placement == "cube"
+        if self.cp <= self.cubes_per_sip_hw:
+            return "inter"
+        return "intersip"
+
+
+@dataclass
+class MachineParams:
+    """SIP hardware parameters — all user-adjustable via sliders.
+
+    Bandwidth tiers, from fastest to slowest:
+      HBM (per-PE local memory)
+        > intra-cube (PE↔PE on same die)
+        > inter-cube (UCIe D2D between dies on same SIP)
+        > inter-SIP (C2C / RDMA-class, across SIPs / servers)
+    """
+    # Per-PE compute peak (TFLOPs f16).
+    peak_tflops_f16: float = 8.0
+    # Per-PE HBM budget (GB) — the memory ceiling per PE.
+    pe_hbm_gb: float = 6.0
+    # HBM per-PE bandwidth (GB/s).
+    bw_hbm_gbs: float = 256.0
+    # Intra-cube PE↔PE link BW (GB/s).
+    bw_intra_gbs: float = 512.0
+    # Inter-cube (UCIe D2D) BW (GB/s).
+    bw_inter_gbs: float = 128.0
+    # Inter-SIP (C2C / RDMA) BW (GB/s).
+    bw_intersip_gbs: float = 50.0
+    # Per-hop latencies (ns).
+    alpha_intra_ns: float = 20.0
+    alpha_inter_ns: float = 100.0
+    alpha_intersip_ns: float = 1000.0
+    # Achievable utilization for compute-bound stages (0-1).
+    compute_util: float = 0.8
+
+    @property
+    def peak_flops(self) -> float:
+        return self.peak_tflops_f16 * 1e12
+
+    @property
+    def bw_hbm(self) -> float:
+        return self.bw_hbm_gbs * 1e9
+
+    @property
+    def bw_intra(self) -> float:
+        return self.bw_intra_gbs * 1e9
+
+    @property
+    def bw_inter(self) -> float:
+        return self.bw_inter_gbs * 1e9
+
+    @property
+    def bw_intersip(self) -> float:
+        return self.bw_intersip_gbs * 1e9
+
+    @property
+    def alpha_intra(self) -> float:
+        return self.alpha_intra_ns * 1e-9
+
+    @property
+    def alpha_inter(self) -> float:
+        return self.alpha_inter_ns * 1e-9
+
+    @property
+    def alpha_intersip(self) -> float:
+        return self.alpha_intersip_ns * 1e-9
+
+    @property
+    def pe_budget_bytes(self) -> int:
+        return int(self.pe_hbm_gb * 1e9)
+
+
+@dataclass
+class FullConfig:
+    """Bundled config for one analysis run."""
+    model: ModelConfig = field(default_factory=ModelConfig)
+    topo: TopologyConfig = field(default_factory=TopologyConfig)
+    machine: MachineParams = field(default_factory=MachineParams)
+
+    @property
+    def h_q_per_pe(self) -> int:
+        return max(1, self.model.h_q // self.topo.tp)
+
+    @property
+    def h_kv_per_pe(self) -> float:
+        """Fractional if TP > H_kv (head-dim split needed)."""
+        return self.model.h_kv / self.topo.tp
+
+    @property
+    def d_per_pe(self) -> int:
+        return self.model.hidden // self.topo.tp
+
+    @property
+    def ffn_per_pe(self) -> int:
+        return self.model.ffn_dim // self.topo.tp
+
+    @property
+    def kv_replication_needed(self) -> bool:
+        """True if TP > H_kv (each KV head shared across TP/H_kv ranks)."""
+        return self.topo.tp > self.model.h_kv
+
+    @property
+    def head_dim_split_factor(self) -> int:
+        """Number of TP ranks sharing one head (via d_head split)."""
+        return max(1, self.topo.tp // self.model.h_kv)
+
+    @property
+    def ffn_shard_divisor(self) -> int:
+        """How much the FFN dim is sharded across ranks."""
+        scope = self.topo.ffn_shard_scope
+        d = self.topo.tp
+        if "CP" in scope:
+            d *= self.topo.cp
+        if "DP" in scope:
+            d *= self.topo.dp
+        return max(1, d)

tests/analytical_visualization/model_presets.py

@@ -0,0 +1,126 @@
+"""Model preset library for the analytical visualization tool.
+
+Comprehensive coverage of popular MHA, GQA, MQA, and MLA models across
+sizes. MoE models are approximated as dense with (ffn_dim = activated
+experts × per_expert_ffn) for compute-side estimates; the note field
+flags this. MLA models are approximated as GQA with matching H_kv until
+we add proper MLA support.
+
+Source of numeric params: model cards / config.json from HuggingFace.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .model_config import ModelConfig
+
+
+@dataclass
+class Preset:
+    label: str
+    model: ModelConfig
+    family: str = ""      # e.g. "Llama 3", "Qwen 3"
+    attn_type: str = ""   # "MHA" | "GQA" | "MQA" | "MLA"
+    note: str = ""
+
+
+def _mk(name: str, hidden: int, ffn: int, hq: int, hkv: int,
+        dh: int, layers: int, *, family: str = "", attn: str = "GQA",
+        note: str = "") -> Preset:
+    return Preset(
+        label=name,
+        model=ModelConfig(
+            name=name, hidden=hidden, ffn_dim=ffn,
+            h_q=hq, h_kv=hkv, d_head=dh, layers=layers,
+        ),
+        family=family, attn_type=attn, note=note,
+    )
+
+
+PRESETS: dict[str, Preset] = {
+    # ── MHA (H_q == H_kv) ─────────────────────────────────────────
+    "GPT-2 (117M)":       _mk("GPT-2 117M",       768,  3072,  12, 12, 64,  12, family="GPT-2", attn="MHA"),
+    "GPT-2 XL (1.5B)":    _mk("GPT-2 XL 1.5B",    1600, 6400,  25, 25, 64,  48, family="GPT-2", attn="MHA"),
+    "Llama 2 7B":         _mk("Llama 2 7B",       4096, 11008, 32, 32, 128, 32, family="Llama 2", attn="MHA"),
+    "Llama 2 13B":        _mk("Llama 2 13B",      5120, 13824, 40, 40, 128, 40, family="Llama 2", attn="MHA"),
+    "Yi 6B":              _mk("Yi 6B",            4096, 11008, 32, 4,  128, 32, family="Yi",     attn="GQA"),
+    "Yi 9B":              _mk("Yi 9B",            4096, 11008, 32, 4,  128, 48, family="Yi",     attn="GQA"),
+    "OPT 30B":            _mk("OPT 30B",          7168, 28672, 56, 56, 128, 48, family="OPT",    attn="MHA"),
+
+    # ── GQA (H_q > H_kv > 1) ──────────────────────────────────────
+    "Qwen 2 0.5B":        _mk("Qwen 2 0.5B",      896,  4864,  14, 2,  64,  24, family="Qwen 2"),
+    "Qwen 2 1.5B":        _mk("Qwen 2 1.5B",      1536, 8960,  12, 2,  128, 28, family="Qwen 2"),
+    "Qwen 2 7B":          _mk("Qwen 2 7B",        3584, 18944, 28, 4,  128, 28, family="Qwen 2"),
+    "Qwen 2 72B":         _mk("Qwen 2 72B",       8192, 29568, 64, 8,  128, 80, family="Qwen 2"),
+
+    "Qwen 3 0.6B":        _mk("Qwen 3 0.6B",      1024, 3072,  16, 8,  128, 28, family="Qwen 3"),
+    "Qwen 3 1.7B":        _mk("Qwen 3 1.7B",      2048, 6144,  16, 8,  128, 28, family="Qwen 3"),
+    "Qwen 3 4B":          _mk("Qwen 3 4B",        2560, 9728,  32, 8,  128, 36, family="Qwen 3"),
+    "Qwen 3 8B":          _mk("Qwen 3 8B",        4096, 12288, 32, 8,  128, 36, family="Qwen 3"),
+    "Qwen 3 14B":         _mk("Qwen 3 14B",       5120, 17408, 40, 8,  128, 40, family="Qwen 3"),
+    "Qwen 3 32B":         _mk("Qwen 3 32B",       5120, 25600, 64, 8,  128, 64, family="Qwen 3"),
+
+    "Llama 3 8B":         _mk("Llama 3 8B",       4096, 14336, 32, 8,  128, 32, family="Llama 3"),
+    "Llama 3 70B":        _mk("Llama 3 70B",      8192, 28672, 64, 8,  128, 80, family="Llama 3"),
+    "Llama 3.1 8B":       _mk("Llama 3.1 8B",     4096, 14336, 32, 8,  128, 32, family="Llama 3.1"),
+    "Llama 3.1 70B":      _mk("Llama 3.1 70B",    8192, 28672, 64, 8,  128, 80, family="Llama 3.1"),
+    "Llama 3.1 405B":     _mk("Llama 3.1 405B",   16384,53248, 128,8,  128, 126,family="Llama 3.1"),
+    "Llama 3.2 1B":       _mk("Llama 3.2 1B",     2048, 8192,  32, 8,  64,  16, family="Llama 3.2"),
+    "Llama 3.2 3B":       _mk("Llama 3.2 3B",     3072, 8192,  24, 8,  128, 28, family="Llama 3.2"),
+
+    "Mistral 7B":         _mk("Mistral 7B v0.3",  4096, 14336, 32, 8,  128, 32, family="Mistral"),
+    "Mistral Nemo 12B":   _mk("Mistral Nemo 12B", 5120, 14336, 32, 8,  128, 40, family="Mistral"),
+    "Mistral Small 22B":  _mk("Mistral Small 22B",6144, 16384, 48, 8,  128, 56, family="Mistral"),
+    "Mistral Large 123B": _mk("Mistral Large 123B",12288,28672,96, 8,  128, 88, family="Mistral"),
+
+    "Gemma 2 2B":         _mk("Gemma 2 2B",       2304, 9216,  8,  4,  256, 26, family="Gemma 2"),
+    "Gemma 2 9B":         _mk("Gemma 2 9B",       3584, 14336, 16, 8,  256, 42, family="Gemma 2"),
+    "Gemma 2 27B":        _mk("Gemma 2 27B",      4608, 36864, 32, 16, 128, 46, family="Gemma 2"),
+
+    "Phi 3 mini (3.8B)":  _mk("Phi 3 mini 3.8B",  3072, 8192,  32, 32, 96,  32, family="Phi 3", attn="MHA"),
+    "Phi 3 small (7B)":   _mk("Phi 3 small 7B",   4096, 14336, 32, 8,  128, 32, family="Phi 3"),
+    "Phi 3 medium (14B)": _mk("Phi 3 medium 14B", 5120, 17920, 40, 10, 128, 40, family="Phi 3"),
+
+    "Yi 34B":             _mk("Yi 34B",           7168, 20480, 56, 8,  128, 60, family="Yi"),
+    "Command R+ 104B":    _mk("Command R+ 104B",  12288,33792, 96, 8,  128, 64, family="Cohere"),
+
+    # ── MoE (activated-only approximation) ────────────────────────
+    "Mixtral 8x7B (MoE)": _mk("Mixtral 8x7B",     4096, 14336*2, 32, 8, 128, 32,
+                              family="Mistral MoE",
+                              note="MoE: 8 experts x 2 activated. FFN ~ 2 x per_expert."),
+    "Mixtral 8x22B (MoE)":_mk("Mixtral 8x22B",    6144, 16384*2, 48, 8, 128, 56,
+                              family="Mistral MoE",
+                              note="MoE: 8 experts x 2 activated."),
+    "Qwen 3 30B (MoE)":   _mk("Qwen 3 30B (MoE)", 2048, 768*8,   32, 4, 128, 48,
+                              family="Qwen 3 MoE",
+                              note="128 experts x 8 activated per token."),
+    "Qwen 3 235B (MoE)":  _mk("Qwen 3 235B (MoE)",4096, 1536*8,  64, 4, 128, 94,
+                              family="Qwen 3 MoE",
+                              note="128 experts x 8 activated. Dense-approx overstates weight."),
+    "DeepSeek V2 (dense-approx)":
+        _mk("DeepSeek V2",  5120, 12288, 128, 128, 128, 60,
+            family="DeepSeek", attn="MLA",
+            note="MLA compresses KV to ~576 B/token; d_c=512 latent. Dense-approx."),
+    "DeepSeek V3 (dense-approx)":
+        _mk("DeepSeek V3",  7168, 16384, 128, 128, 128, 61,
+            family="DeepSeek", attn="MLA",
+            note="MoE 256 x 8 activated + MLA. Both approximations."),
+    "Grok-1 314B (MoE)":  _mk("Grok-1 314B",     6144, 32768*2, 48, 8,  128, 64,
+                              family="xAI",
+                              note="MoE 8 experts x 2 activated."),
+
+    # ── MQA (H_kv == 1) ───────────────────────────────────────────
+    "Falcon 7B (MQA)":    _mk("Falcon 7B",       4544, 18176, 71, 1,  64,  32, family="Falcon", attn="MQA"),
+    "Falcon 40B (MQA)":   _mk("Falcon 40B",      8192, 32768, 128,1,  64,  60, family="Falcon", attn="MQA"),
+
+    # ── Custom slot ──────────────────────────────────────────────
+    "Custom": Preset(label="Custom", model=ModelConfig(name="Custom")),
+}
+
+
+def preset_names_by_family() -> dict[str, list[str]]:
+    """Return preset names grouped by family (for a nested dropdown)."""
+    grouped: dict[str, list[str]] = {}
+    for name, p in PRESETS.items():
+        grouped.setdefault(p.family or "Other", []).append(name)
+    return grouped

tests/analytical_visualization/optimization_report.py

@@ -0,0 +1,235 @@
+"""Replication (waste) accounting + optimization opportunities.
+
+Two things this module produces:
+
+1. `replication_report(cfg)` — where memory is being *duplicated* across
+   the deployment, per category (attention weights, FFN weights, KV cache
+   under replicate mode, DP replicas). Answers "if I lifted this
+   duplication, how much would I save?"
+
+2. `optimization_hints(cfg)` — a list of actionable hints. Each hint has
+   a `category` (space | comm | compute), a `severity` (info | warn |
+   good), and a short human message. UI code renders these as bullets.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .model_config import FullConfig
+from .memory_layout import per_pe_weight_bytes, per_pe_kv_cache_bytes
+
+
+# ── Replication accounting ────────────────────────────────────────
+
+@dataclass
+class ReplicationEntry:
+    tensor: str
+    per_pe_bytes: int
+    replicated_across: str    # e.g. "CP (x4 copies)"
+    copies: int               # total # copies globally
+    wasted_bytes: int         # (copies - 1) * per_pe_bytes  (per PE view)
+
+
+def _attn_weight_bytes_per_pe(cfg: FullConfig) -> int:
+    """W_Q + W_K + W_V + W_O per PE, one layer."""
+    m = cfg.model
+    tp = cfg.topo.tp
+    hq_per_pe = cfg.h_q_per_pe
+    if cfg.topo.kv_shard_mode == "replicate":
+        hkv_per_pe = max(1.0, m.h_kv / tp)
+    else:
+        hkv_per_pe = m.h_kv / tp
+    per_layer = (
+        m.hidden * hq_per_pe * m.d_head +
+        m.hidden * hkv_per_pe * m.d_head +
+        m.hidden * hkv_per_pe * m.d_head +
+        hq_per_pe * m.d_head * m.hidden
+    ) * m.bytes_per_elem
+    layers_per_stage = (m.layers + cfg.topo.pp - 1) // cfg.topo.pp
+    return int(per_layer * layers_per_stage)
+
+
+def _ffn_weight_bytes_per_pe(cfg: FullConfig) -> int:
+    """W_gate + W_up + W_down per PE, all local layers."""
+    m = cfg.model
+    ep = max(1, cfg.topo.ep)
+    ffn_div = cfg.ffn_shard_divisor * ep
+    per_layer = 3 * m.hidden * (m.ffn_dim // max(1, ffn_div)) * m.bytes_per_elem
+    layers_per_stage = (m.layers + cfg.topo.pp - 1) // cfg.topo.pp
+    return int(per_layer * layers_per_stage)
+
+
+def _kv_bytes_per_pe(cfg: FullConfig) -> int:
+    return per_pe_kv_cache_bytes(cfg)
+
+
+def replication_report(cfg: FullConfig) -> list[ReplicationEntry]:
+    """Enumerate replicated (duplicated) tensors and their waste."""
+    entries: list[ReplicationEntry] = []
+    topo = cfg.topo
+
+    # Attention weights are NOT sharded by CP - so each CP rank holds a copy.
+    attn_pe = _attn_weight_bytes_per_pe(cfg)
+    cp_copies = topo.cp
+    if cp_copies > 1:
+        entries.append(ReplicationEntry(
+            tensor="Attn weights (W_Q/W_K/W_V/W_O)",
+            per_pe_bytes=attn_pe,
+            replicated_across=f"CP (x{cp_copies} identical group copies)",
+            copies=cp_copies,
+            wasted_bytes=attn_pe * (cp_copies - 1),
+        ))
+
+    # FFN weights: replicated across scope levels NOT included.
+    # Scope=TP  -> replicated across CP AND DP.
+    # Scope=TP+CP -> replicated across DP only.
+    # Scope=TP+CP+DP -> no replication (perfect share).
+    ffn_pe = _ffn_weight_bytes_per_pe(cfg)
+    scope = topo.ffn_shard_scope
+    ffn_cp_copies = 1 if "CP" in scope else topo.cp
+    ffn_dp_copies = 1 if "DP" in scope else topo.dp
+    ffn_total_copies = ffn_cp_copies * ffn_dp_copies
+    if ffn_total_copies > 1:
+        reasons = []
+        if ffn_cp_copies > 1:
+            reasons.append(f"CP (x{ffn_cp_copies})")
+        if ffn_dp_copies > 1:
+            reasons.append(f"DP (x{ffn_dp_copies})")
+        entries.append(ReplicationEntry(
+            tensor=f"FFN weights (scope={scope})",
+            per_pe_bytes=ffn_pe,
+            replicated_across=" & ".join(reasons),
+            copies=ffn_total_copies,
+            wasted_bytes=ffn_pe * (ffn_total_copies - 1),
+        ))
+
+    # KV cache: if kv_shard_mode='replicate' and TP > H_kv, each KV head
+    # is duplicated tp // h_kv times across TP ranks.
+    if (cfg.kv_replication_needed
+            and topo.kv_shard_mode == "replicate"):
+        kv_pe = _kv_bytes_per_pe(cfg)
+        rep = cfg.head_dim_split_factor  # tp // h_kv
+        entries.append(ReplicationEntry(
+            tensor="KV cache (replicate mode)",
+            per_pe_bytes=kv_pe,
+            replicated_across=f"TP (x{rep} copies of each KV head)",
+            copies=rep,
+            wasted_bytes=kv_pe * (rep - 1),
+        ))
+
+    # DP replicas duplicate the entire per-PE footprint (attn + ffn + KV).
+    if topo.dp > 1:
+        total_pe = attn_pe + ffn_pe + _kv_bytes_per_pe(cfg)
+        entries.append(ReplicationEntry(
+            tensor="Full per-PE footprint (attn + FFN + KV)",
+            per_pe_bytes=total_pe,
+            replicated_across=f"DP (x{topo.dp} model replicas)",
+            copies=topo.dp,
+            wasted_bytes=total_pe * (topo.dp - 1),
+        ))
+
+    return entries
+
+
+# ── Optimization hints ────────────────────────────────────────────
+
+@dataclass
+class Hint:
+    category: str   # "space" | "comm" | "compute" | "layout"
+    severity: str   # "info" | "warn" | "good"
+    message: str
+
+
+def optimization_hints(cfg: FullConfig) -> list[Hint]:
+    hints: list[Hint] = []
+    m = cfg.model
+    topo = cfg.topo
+
+    # SPACE hints -------------------------------------------------
+    scope = topo.ffn_shard_scope
+    ffn_pe = _ffn_weight_bytes_per_pe(cfg)
+    if scope == "TP" and topo.cp > 1:
+        current = ffn_pe
+        new = current // topo.cp
+        hints.append(Hint(
+            "space", "warn",
+            f"FFN scope=TP replicates FFN weights across all {topo.cp} CP "
+            f"groups. Switching to **TP+CP** would drop per-PE FFN weight "
+            f"from {current/1e9:.2f} GB to {new/1e9:.2f} GB "
+            f"(saves {(current-new)/1e9:.2f} GB/PE) at the cost of one "
+            f"extra AllReduce per layer."
+        ))
+    if scope in ("TP", "TP+CP") and topo.dp > 1:
+        base = ffn_pe if scope == "TP+CP" else ffn_pe // topo.cp
+        new = base // topo.dp
+        hints.append(Hint(
+            "space", "info",
+            f"With DP={topo.dp}, scope=**TP+CP+DP** would shard FFN across "
+            f"replicas too, dropping per-PE FFN from {base/1e9:.2f} GB to "
+            f"{new/1e9:.2f} GB (comes with an inter-replica AllReduce)."
+        ))
+    if (cfg.kv_replication_needed
+            and topo.kv_shard_mode == "replicate"):
+        rep = cfg.head_dim_split_factor
+        kv_pe = _kv_bytes_per_pe(cfg)
+        hints.append(Hint(
+            "space", "warn",
+            f"KV mode=replicate holds each KV head {rep}x. Switching to "
+            f"**split** saves ~{kv_pe*(rep-1)/rep/1e9:.2f} GB KV per PE "
+            f"but adds a Score AllReduce over {rep} ranks per hop."
+        ))
+    if topo.cp == 1 and topo.tp <= m.h_kv:
+        # Suggest CP if s_kv is large and KV dominates.
+        kv_pe = _kv_bytes_per_pe(cfg)
+        attn_pe = _attn_weight_bytes_per_pe(cfg)
+        if kv_pe > attn_pe:
+            hints.append(Hint(
+                "space", "info",
+                f"KV cache ({kv_pe/1e9:.2f} GB) dominates weights "
+                f"({attn_pe/1e9:.2f} GB) at S_kv={topo.s_kv:,}. Adding CP "
+                f"shards the sequence axis (S_local = S_kv/CP)."
+            ))
+
+    # COMM hints --------------------------------------------------
+    if topo.cp_inter_sip_hops > 0:
+        hints.append(Hint(
+            "comm", "warn",
+            f"CP ring crosses {topo.cp_inter_sip_hops} SIP boundary(ies) "
+            f"at {cfg.machine.bw_intersip_gbs:.0f} GB/s vs "
+            f"{cfg.machine.bw_inter_gbs:.0f} GB/s intra-SIP - "
+            f"consider a smaller CP or a topology that keeps CP inside one SIP."
+        ))
+    if topo.tp_spans_cubes > 1:
+        hints.append(Hint(
+            "comm", "warn",
+            f"TP={topo.tp} spans {topo.tp_spans_cubes} cubes, so the W_O "
+            f"AllReduce runs cross-cube "
+            f"({cfg.machine.bw_intra_gbs:.0f} -> "
+            f"{cfg.machine.bw_inter_gbs:.0f} GB/s). A TP that fits in one "
+            f"cube (TP<={topo.pes_per_cube_hw}) keeps it intra-cube."
+        ))
+    if topo.sips_used > 1 and topo.sip_topology == "ring":
+        hints.append(Hint(
+            "layout", "info",
+            f"With {topo.sips_used} SIPs on a ring, worst-case inter-SIP "
+            f"distance is {topo.sips_used - 1} hops. Switching to **torus2d** "
+            f"cuts worst-case hops to ~sqrt({topo.sips_used})."
+        ))
+    if topo.tp <= m.h_kv:
+        hints.append(Hint(
+            "comm", "good",
+            f"TP ({topo.tp}) <= H_kv ({m.h_kv}): each KV head lives on "
+            f"exactly one TP rank - no head-split AllReduce needed."
+        ))
+
+    # COMPUTE hints -----------------------------------------------
+    if topo.mode == "decode" and topo.T_q == 1:
+        # Decode is memory-bound almost everywhere; note it.
+        hints.append(Hint(
+            "compute", "info",
+            "Decode (T_q=1): most GEMMs are memory-bound. FLOPs matter far "
+            "less than HBM BW here - budget attention around "
+            "weight-read cost, not compute peak."
+        ))
+
+    return hints

tests/analytical_visualization/pe_weight_layout.py

@@ -0,0 +1,221 @@
+"""Draw a PE-level view of one CP group (one TP group of cubes).
+
+Shows how weights + KV cache are distributed across the PEs of ONE CP
+rank. Attention weights are sharded by head (H_q, H_kv split across TP);
+FFN is sharded by dim (ffn_dim / TP / EP); KV cache is sharded by CP
+(sequence) × TP (head).
+
+If TP > 8 the group spans multiple cubes — layout draws all of them.
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+
+from .model_config import FullConfig
+
+
+PE_COLS = 4
+PE_ROWS = 2
+PES_PER_CUBE = PE_COLS * PE_ROWS
+
+
+def layers_per_stage_val(cfg: FullConfig) -> int:
+    return (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp
+
+
+def _q_heads_for_pe(pe_id_in_group: int, tp: int, h_q: int) -> list[int]:
+    """Return the list of Q-head indices assigned to this PE."""
+    heads_per_pe = h_q / tp
+    start = int(round(pe_id_in_group * heads_per_pe))
+    end = int(round((pe_id_in_group + 1) * heads_per_pe))
+    return list(range(start, end))
+
+
+def _kv_heads_for_pe(pe_id_in_group: int, tp: int, h_kv: int,
+                      kv_shard_mode: str) -> tuple[list[int], str]:
+    """KV heads assigned to this PE and a note about replication/split."""
+    if h_kv >= tp:
+        heads_per_pe = h_kv // tp
+        start = pe_id_in_group * heads_per_pe
+        end = start + heads_per_pe
+        return list(range(start, end)), ""
+    # TP > H_kv
+    if kv_shard_mode == "replicate":
+        rep_factor = tp // h_kv
+        head_idx = pe_id_in_group // rep_factor
+        return [head_idx], f"replicated x{rep_factor}"
+    # split mode
+    split_factor = tp // h_kv
+    head_idx = pe_id_in_group // split_factor
+    part = pe_id_in_group % split_factor
+    return [head_idx], f"head-split ({part+1}/{split_factor})"
+
+
+def _per_pe_bytes(cfg: FullConfig, pe_id_in_group: int) -> dict:
+    """Return per-tensor bytes for one PE within a TP group."""
+    m = cfg.model
+    tp = cfg.topo.tp
+    pp = cfg.topo.pp
+    ep = max(1, cfg.topo.ep)
+    b = m.bytes_per_elem
+    layers_per_stage = (m.layers + pp - 1) // pp
+
+    hq_per_pe = m.h_q / tp
+    if cfg.topo.kv_shard_mode == "replicate":
+        hkv_per_pe = max(1.0, m.h_kv / tp)
+    else:
+        hkv_per_pe = m.h_kv / tp
+    ffn_div = cfg.ffn_shard_divisor * ep
+    ffn_per_pe = m.ffn_dim // ffn_div
+
+    per_layer = {
+        "W_Q": int(m.hidden * hq_per_pe * m.d_head * b),
+        "W_K": int(m.hidden * hkv_per_pe * m.d_head * b),
+        "W_V": int(m.hidden * hkv_per_pe * m.d_head * b),
+        "W_O": int(hq_per_pe * m.d_head * m.hidden * b),
+        "W_gate": int(m.hidden * ffn_per_pe * b),
+        "W_up": int(m.hidden * ffn_per_pe * b),
+        "W_down": int(ffn_per_pe * m.hidden * b),
+    }
+    kv_per_layer = int(2 * cfg.topo.s_local * hkv_per_pe * m.d_head * b)
+
+    all_layers = {k: v * layers_per_stage for k, v in per_layer.items()}
+    all_layers["KV cache"] = kv_per_layer * layers_per_stage
+    # transient estimate
+    if cfg.topo.mode == "decode":
+        all_layers["Transient"] = int(4 * (m.hidden + hq_per_pe * m.d_head) * b)
+    else:
+        TILE = 1024
+        all_layers["Transient"] = int(2 * hq_per_pe * cfg.topo.T_q * TILE * b
+                                       + cfg.topo.T_q * m.hidden * b)
+    return all_layers
+
+
+def draw_pe_layout(cfg: FullConfig, ax=None):
+    """Draw one CP group's PEs with per-tensor weight + KV breakdown.
+
+    Each PE shows: PE id, Q heads, KV heads, W_Q / W_K / W_V / W_O in
+    MB, FFN (gate+up+down) in MB, KV cache in GB, total in GB.
+
+    TP=8 -> 1 cube. TP=16 -> 2 cubes. TP=32 -> 4 cubes.
+    Under placement, n_cubes reflects cubes-per-group (may include CP spill).
+    """
+    tp = cfg.topo.tp
+    # cubes per one cube-level group (= intra-cube dims / PEs-per-cube, rounded up).
+    n_cubes = max(1, (cfg.topo.intra_cube_dims + PES_PER_CUBE - 1) // PES_PER_CUBE)
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(7.5 * n_cubes, 8.5))
+    else:
+        fig = ax.figure
+
+    cube_w = 7.0
+    cube_h = 6.0
+    cube_gap = 0.6
+
+    for cube_idx in range(n_cubes):
+        x0 = cube_idx * (cube_w + cube_gap)
+        y0 = 0
+
+        rect = patches.FancyBboxPatch(
+            (x0, y0), cube_w, cube_h,
+            boxstyle="round,pad=0.05",
+            facecolor="#f8f9fa", edgecolor="#212529", linewidth=1.5,
+        )
+        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})",
+                ha="center", va="bottom", fontsize=11, fontweight="bold")
+
+        pe_pad_x = 0.15
+        pe_pad_y = 0.2
+        pe_gap = 0.10
+        inner_w = cube_w - 2 * pe_pad_x
+        inner_h = cube_h - 2 * pe_pad_y
+        pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
+        pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
+
+        for pr in range(PE_ROWS):
+            for pc in range(PE_COLS):
+                pe_local = pr * PE_COLS + pc
+                pe_id_in_group = cube_idx * PES_PER_CUBE + pe_local
+                if pe_id_in_group >= cfg.topo.intra_cube_dims:
+                    continue
+                px = x0 + pe_pad_x + pc * (pe_w + pe_gap)
+                py = y0 + pe_pad_y + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
+
+                pe_rect = patches.Rectangle(
+                    (px, py), pe_w, pe_h,
+                    facecolor="#f0f7ff", edgecolor="#3a86ff", linewidth=1.0,
+                )
+                ax.add_patch(pe_rect)
+
+                q_heads = _q_heads_for_pe(pe_id_in_group, 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,
+                )
+                bytes_ = _per_pe_bytes(cfg, pe_id_in_group)
+                weights_gb = sum(v for k, v in bytes_.items()
+                                 if k not in ("KV cache", "Transient")) / 1e9
+                kv_gb = bytes_["KV cache"] / 1e9
+                trans_mb = bytes_["Transient"] / 1e6
+                total_gb = sum(bytes_.values()) / 1e9
+
+                q_str = (f"{q_heads[0]}-{q_heads[-1]}"
+                         if len(q_heads) > 1 else str(q_heads[0])
+                         if q_heads else "-")
+                kv_str = str(kv_heads[0]) if kv_heads else "-"
+                if kv_note:
+                    kv_str += f"({kv_note})"
+
+                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"
+                          f"Q heads: {q_str}\n"
+                          f"KV head: {kv_str}")
+                ax.text(px + 0.08, py + pe_h - 0.05, header,
+                        ha="left", va="top",
+                        fontsize=7.0, fontweight="bold",
+                        family="monospace", color="#0d47a1")
+
+                # Per-tensor bytes with explicit division divisors.
+                # Divisor is always the PRODUCT of the parallel dims listed
+                # (TP*CP means TP times CP, i.e. multiplication).
+                tp_d = cfg.topo.tp
+                ffn_div = cfg.ffn_shard_divisor
+                ffn_scope = cfg.topo.ffn_shard_scope.replace("+", "*")
+                kv_div = cfg.topo.cp * cfg.topo.tp
+                L = layers_per_stage_val(cfg)
+                detail = (
+                    f"weights (all {L} layers):\n"
+                    f" W_Q /TP={tp_d} : {bytes_['W_Q']/1e6:7.1f}MB\n"
+                    f" W_K /TP={tp_d} : {bytes_['W_K']/1e6:7.1f}MB\n"
+                    f" W_V /TP={tp_d} : {bytes_['W_V']/1e6:7.1f}MB\n"
+                    f" W_O /TP={tp_d} : {bytes_['W_O']/1e6:7.1f}MB\n"
+                    f" FFN /({ffn_scope})={ffn_div}: {ffn_mb:7.1f}MB\n"
+                    f"runtime:\n"
+                    f" KV /(CP*TP)={kv_div}: {kv_gb*1000:7.1f}MB\n"
+                    f" trans     : {trans_mb:7.1f}MB\n"
+                    f"W total: {weights_gb*1000:7.1f}MB\n"
+                    f"TOTAL  : {total_gb*1000:7.1f}MB"
+                )
+                ax.text(px + pe_w / 2, py + 0.06, detail,
+                        ha="center", va="bottom",
+                        fontsize=5.8, family="monospace", color="#1a1a1a")
+
+    total_w = n_cubes * cube_w + (n_cubes - 1) * cube_gap
+    ax.set_xlim(-0.2, total_w + 0.2)
+    ax.set_ylim(-0.3, cube_h + 0.6)
+    ax.set_aspect("equal")
+    ax.axis("off")
+
+    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" | KV mode: {cfg.topo.kv_shard_mode}"
+    )
+    ax.set_title(title, fontsize=11, fontweight="bold")
+
+    return fig

tests/analytical_visualization/pipeline_diagram.py

@@ -0,0 +1,300 @@
+"""Pipeline diagram: attention + FFN steps, colored by bound type.
+
+Each stage is drawn as a rounded box in a horizontal flow:
+  Attention: S1 -> S2 -> S3 -> ... -> S10 -> C2 (AllReduce)
+  FFN:       F1 -> F2 -> F3 -> F4 -> F5 -> CF1 (AllReduce)
+
+CP ring communication (C1) and score AllReduce (C3) are drawn as
+"overhead" boxes attached above the QKt/PV region (they run
+concurrently with per-hop compute in the ring).
+
+Box fill color = the stage's bound classification:
+  compute-bound  blue
+  memory-bound   orange
+  comm-bound     red
+  trivial        grey
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+
+from .model_config import FullConfig
+from .stage_latencies import (
+    all_stages, all_ffn_stages, StageCost,
+)
+
+
+BOUND_COLOR = {
+    "compute": "#4682b4",   # blue
+    "memory":  "#e37400",   # orange
+    "comm":    "#c62828",   # red
+    "trivial": "#a9a9a9",   # grey
+}
+BOUND_LABEL = {
+    "compute": "compute-bound",
+    "memory":  "memory-bound",
+    "comm":    "communication",
+    "trivial": "trivial / negligible",
+}
+
+
+def _fmt_us(sec: float) -> str:
+    us = sec * 1e6
+    if us < 1:
+        return f"{us*1e3:.1f}ns"
+    if us < 100:
+        return f"{us:.2f}us"
+    return f"{us:.0f}us"
+
+
+def _draw_block(ax, x, y, w, h, label, cost: StageCost, edge="#333"):
+    color = BOUND_COLOR.get(cost.bound, "#a9a9a9")
+    box = patches.FancyBboxPatch(
+        (x, y), w, h,
+        boxstyle="round,pad=0.05",
+        facecolor=color, edgecolor=edge, linewidth=1.2, alpha=0.85,
+    )
+    ax.add_patch(box)
+    # Two lines: stage name (bold) + visible latency (small)
+    ax.text(x + w / 2, y + h * 0.62, label,
+            ha="center", va="center",
+            fontsize=8.0, fontweight="bold", color="white")
+    ax.text(x + w / 2, y + h * 0.25, _fmt_us(cost.visible_s),
+            ha="center", va="center",
+            fontsize=7.0, color="white")
+
+
+_SHORT_NAMES = {
+    "S1":  "RMSNorm",
+    "S2":  "W_Q",
+    "S3":  "W_K+W_V",
+    "S4":  "KV append",
+    "S5":  "Q.K^T",
+    "S6":  "softmax",
+    "S7":  "P.V",
+    "S8":  "merge",
+    "S9":  "norm O/l",
+    "S10": "W_O",
+    "C1":  "CP ring",
+    "C2":  "TP AR",
+    "C3":  "Score AR",
+    "F1":  "RMSNorm",
+    "F2":  "W_gate",
+    "F3":  "W_up",
+    "F4":  "SwiGLU",
+    "F5":  "W_down",
+    "CF1": "FFN AR",
+}
+
+
+def _short_name(name: str) -> str:
+    """Map 'S1 RMSNorm' -> 'RMSNorm'; fallback to first token."""
+    tok = name.split()[0]
+    return _SHORT_NAMES.get(tok, tok)
+
+
+def _draw_arrow(ax, x0, x1, y):
+    ax.annotate("", xy=(x1, y), xytext=(x0, y),
+                arrowprops=dict(arrowstyle="->", color="#333", lw=1.1))
+
+
+def _draw_ring_loop(ax, x_left, x_right, y_top, label: str, n_iters: int,
+                     arc_lift: float = 0.9):
+    """Draw a curved back-arrow above [x_left..x_right] with a repeat label.
+
+    Visualizes 'repeat this block n_iters times' (like a loop in a diagram).
+    arc_lift: how high above y_top the arc sits.
+    """
+    color = "#5a189a"
+    lw = 1.6
+    arc_y = y_top + arc_lift
+    # Vertical tick up from the right edge
+    ax.plot([x_right, x_right], [y_top + 0.02, arc_y],
+            color=color, linewidth=lw)
+    # Horizontal top of the loop
+    ax.plot([x_right, x_left], [arc_y, arc_y],
+            color=color, linewidth=lw)
+    # Down-arrow into the left edge (a back-edge indicating the repeat)
+    ax.annotate("", xy=(x_left, y_top + 0.02), xytext=(x_left, arc_y),
+                arrowprops=dict(arrowstyle="->", color=color, lw=lw))
+    # Loop label centered above the arc
+    ax.text((x_left + x_right) / 2, arc_y + 0.08,
+            label,
+            ha="center", va="bottom",
+            fontsize=9, fontweight="bold", color=color)
+    # Iteration count as a badge on the right
+    ax.text(x_right + 0.08, arc_y - 0.03,
+            f"x {n_iters}",
+            ha="left", va="top",
+            fontsize=10, fontweight="bold", color=color,
+            bbox=dict(boxstyle="round,pad=0.15", facecolor="white",
+                      edgecolor=color, linewidth=1.2, alpha=0.95))
+
+
+def _draw_row(ax, y_base, row_title, stages: list[StageCost],
+              x_start: float, box_w: float, box_h: float, gap: float):
+    """Draw one row of stage boxes with arrows between them."""
+    ax.text(x_start - 0.4, y_base + box_h / 2, row_title,
+            ha="right", va="center",
+            fontsize=11, fontweight="bold", color="#212529")
+    x = x_start
+    for i, st in enumerate(stages):
+        _draw_block(ax, x, y_base, box_w, box_h,
+                    _short_name(st.name), st)
+        if i < len(stages) - 1:
+            _draw_arrow(ax, x + box_w + 0.02, x + box_w + gap - 0.02,
+                        y_base + box_h / 2)
+        x += box_w + gap
+    return x - gap
+
+
+def _draw_overlap_boxes(ax, hidden_stages, x_left, x_right, y_top,
+                          gap_between: float = 0.15):
+    """Draw dashed 'overlapped comm' boxes spanning x_left..x_right above y_top.
+
+    Boxes are widened to cover the horizontal range of the stages they
+    overlap with (e.g. C1 CP ring sits above S5-S7).
+    """
+    if not hidden_stages:
+        return
+    hy = y_top + 0.10
+    n = len(hidden_stages)
+    total_span = x_right - x_left
+    box_span = (total_span - (n - 1) * gap_between) / max(1, n)
+    box_h_local = 0.42
+    hx = x_left
+    for st in hidden_stages:
+        box = patches.FancyBboxPatch(
+            (hx, hy), box_span, box_h_local,
+            boxstyle="round,pad=0.03",
+            facecolor=BOUND_COLOR["comm"], edgecolor="#333",
+            linewidth=1.0, alpha=0.35, linestyle="--",
+        )
+        ax.add_patch(box)
+        ax.text(hx + box_span / 2, hy + box_h_local * 0.65,
+                f"{_short_name(st.name)} (concurrent)",
+                ha="center", va="center",
+                fontsize=7, color="#7a0e0e", fontweight="bold")
+        ax.text(hx + box_span / 2, hy + box_h_local * 0.2,
+                _fmt_us(st.visible_s),
+                ha="center", va="center",
+                fontsize=6.5, color="#7a0e0e")
+        hx += box_span + gap_between
+
+
+def draw_pipeline(cfg: FullConfig, ax=None):
+    """Draw attention + FFN pipeline. Boxes colored by bound type."""
+    attn = all_stages(cfg)
+    ffn = all_ffn_stages(cfg)
+
+    # Split attention: main stages, then C2 (TP AllReduce). C1/C3 = overlapped.
+    attn_main = [s for s in attn if not s.name.startswith("C")]
+    attn_c2 = next((s for s in attn if s.name.startswith("C2")), None)
+    attn_hidden = [s for s in attn
+                   if s.name.startswith("C1") or s.name.startswith("C3")]
+    attn_hidden = [s for s in attn_hidden if s.visible_s > 0]
+    attn_row = attn_main + ([attn_c2] if attn_c2 and attn_c2.visible_s > 0 else [])
+
+    # FFN row: all F* + CF1 if non-trivial
+    ffn_main = [s for s in ffn if not s.name.startswith("CF")]
+    ffn_cf = next((s for s in ffn if s.name.startswith("CF1")), None)
+    ffn_row = ffn_main + ([ffn_cf] if ffn_cf and ffn_cf.visible_s > 0 else [])
+
+    n_max = max(len(attn_row), len(ffn_row))
+    box_w = 1.35
+    box_h = 0.78
+    gap = 0.18
+    x_start = 1.1
+
+    total_w = x_start + n_max * box_w + (n_max - 1) * gap + 0.5
+    total_h = 4.2
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(max(11, total_w * 1.1),
+                                         max(4.5, total_h * 1.0)))
+    else:
+        fig = ax.figure
+
+    _draw_row(ax, y_base=2.6, row_title="Attention",
+              stages=attn_row, x_start=x_start,
+              box_w=box_w, box_h=box_h, gap=gap)
+
+    _draw_row(ax, y_base=0.6, row_title="FFN",
+              stages=ffn_row, x_start=x_start,
+              box_w=box_w, box_h=box_h, gap=gap)
+
+    # Ring loop indicator + concurrent-comm boxes anchored over S5..S8.
+    _idx_by_prefix = {s.name.split()[0]: i for i, s in enumerate(attn_row)}
+    _s5_idx = _idx_by_prefix.get("S5")
+    _s7_idx = _idx_by_prefix.get("S7")
+    _s8_idx = _idx_by_prefix.get("S8")
+
+    if cfg.topo.cp > 1 and cfg.topo.mode == "prefill" and _s5_idx is not None:
+        _right_idx = _s8_idx if _s8_idx is not None else _s7_idx
+        _lx = x_start + _s5_idx * (box_w + gap)
+        _rx = x_start + _right_idx * (box_w + gap) + box_w
+        _y_top = 2.6 + box_h
+        # Concurrent-comm boxes (C1 CP ring, C3 score AR) directly above S5-S7
+        _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
+        # Ring loop arc above those
+        _draw_ring_loop(
+            ax, _lx, _rx, _y_top,
+            label="ring attention loop (K/V ring per hop; S8 merges partial O,m,l)",
+            n_iters=cfg.topo.cp,
+            arc_lift=1.05,
+        )
+    elif cfg.topo.cp > 1 and cfg.topo.mode == "decode" and _s5_idx is not None:
+        _right_idx = _s7_idx if _s7_idx is not None else _s5_idx
+        _lx = x_start + _s5_idx * (box_w + gap)
+        _rx = x_start + _right_idx * (box_w + gap) + box_w
+        _y_top = 2.6 + box_h
+        # Any straggling concurrent comm (e.g. C3 if head-split active) above S5-S7
+        _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
+        _draw_ring_loop(
+            ax, _lx, _rx, _y_top,
+            label="decode: local pass only; S8 fires O/m/l all-reduce",
+            n_iters=1,
+            arc_lift=1.05 if attn_hidden else 0.85,
+        )
+    elif attn_hidden:
+        # Fallback: CP=1 but still have C3 or similar; anchor above S5-S7 if present
+        if _s5_idx is not None and _s7_idx is not None:
+            _lx = x_start + _s5_idx * (box_w + gap)
+            _rx = x_start + _s7_idx * (box_w + gap) + box_w
+            _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, 2.6 + box_h)
+
+    # Totals
+    attn_total = sum(s.visible_s for s in attn_row) \
+        + sum(s.visible_s for s in attn_hidden)
+    ffn_total = sum(s.visible_s for s in ffn_row)
+    ax.text(total_w - 0.3, 2.6 + box_h / 2 + 1.15,
+            f"Attn total: {_fmt_us(attn_total)}",
+            ha="right", va="bottom", fontsize=9,
+            fontweight="bold", color="#212529")
+    ax.text(total_w - 0.3, 0.6 + box_h + 0.15,
+            f"FFN total:  {_fmt_us(ffn_total)}",
+            ha="right", va="bottom", fontsize=9,
+            fontweight="bold", color="#212529")
+
+    # Legend
+    handles = []
+    for k, lbl in BOUND_LABEL.items():
+        handles.append(patches.Patch(
+            facecolor=BOUND_COLOR[k], edgecolor="#333",
+            alpha=0.85, label=lbl,
+        ))
+    ax.legend(handles=handles, loc="upper center",
+              bbox_to_anchor=(0.5, -0.02),
+              ncol=4, fontsize=8, frameon=False)
+
+    ax.set_xlim(0, total_w + 0.6)
+    ax.set_ylim(0, 4.9)
+    ax.set_aspect("auto")
+    ax.axis("off")
+    ax.set_title(
+        f"Per-layer pipeline (one PE, {cfg.topo.mode}, T_q={cfg.topo.T_q}, "
+        f"S_local={cfg.topo.s_local}) - color = dominant bound",
+        fontsize=10, fontweight="bold",
+    )
+    return fig

tests/analytical_visualization/stage_latencies.py

@@ -0,0 +1,767 @@
+"""Per-stage cost formulas for attention forward pass.
+
+Each stage returns:
+    name, formula (str with the substituted numeric form),
+    compute_time, memory_time, comm_time, bound, visible_time.
+
+Bound is the max of the three; "hidden" comm can appear as 0 when
+overlapped with compute.
+
+All times are in seconds; convert to μs at display.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+from .model_config import FullConfig
+
+
+Bound = Literal["compute", "memory", "comm", "trivial"]
+
+
+@dataclass
+class StageCost:
+    name: str
+    formula: str
+    compute_s: float
+    memory_s: float
+    comm_s: float
+    bound: Bound
+    visible_s: float
+    hop_multiplier: int = 1  # some stages scale with N_cp hops
+    # Detailed breakdown for the per-stage table.
+    flops: int = 0                  # total FLOPs (numeric)
+    mem_bytes: int = 0              # HBM bytes moved (numeric)
+    comm_bytes: int = 0             # comm bytes over the wire (numeric)
+    flops_formula: str = ""         # human-readable FLOPs formula
+    mem_formula: str = ""           # human-readable memory formula
+    comm_formula: str = ""          # human-readable comm formula
+
+
+def _visible(compute: float, memory: float, comm: float) -> tuple[float, Bound]:
+    """Pick the dominant time as the visible stage cost."""
+    parts = {"compute": compute, "memory": memory, "comm": comm}
+    dominant = max(parts, key=parts.get)  # type: ignore
+    return parts[dominant], dominant  # type: ignore
+
+
+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
+    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",
+        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",
+    )
+
+
+def _gemm_time(flops: int, weight_bytes: int, cfg: FullConfig) -> tuple[float, float]:
+    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+    mem_s = weight_bytes / cfg.machine.bw_hbm
+    return cmp_s, mem_s
+
+
+def stage_wq(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    hq_per_pe = cfg.h_q_per_pe
+    dh = cfg.model.d_head
+    flops = 2 * 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",
+        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",
+    )
+
+
+def stage_wkv(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    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)
+    weight_B_one = d * (hkv_per_pe * dh) * b
+    flops = 2 * flops_one
+    weight_B = 2 * weight_B_one
+    cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
+    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",
+        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",
+    )
+
+
+def stage_kv_append(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    b = cfg.model.bytes_per_elem
+    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
+    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",
+        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",
+    )
+
+
+def _per_hop_qkT_pv(cfg: FullConfig) -> tuple[float, str]:
+    """Q·Kᵀ (and P·V has same FLOPs) per hop, per PE."""
+    T_q = cfg.topo.T_q
+    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
+    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
+
+
+def _cp_compute_passes(cfg: FullConfig) -> int:
+    """How many local S5/S6/S7 passes happen per token step under CP.
+    - decode: 1 (each PE computes once against its local K,V, then all-reduce O/m/l)
+    - prefill: CP  (K/V or Q rotates through the ring; one local pass per hop)
+    """
+    if cfg.topo.cp <= 1:
+        return 1
+    return 1 if cfg.topo.mode == "decode" else cfg.topo.cp
+
+
+def stage_qkT(cfg: FullConfig) -> StageCost:
+    cmp_hop, _ = _per_hop_qkT_pv(cfg)
+    passes = _cp_compute_passes(cfg)
+    total = cmp_hop * passes
+    T_q = cfg.topo.T_q
+    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
+    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",
+        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}"
+        ),
+        mem_formula="0 (scores accumulated on-chip)",
+    )
+
+
+def stage_softmax(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    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
+    bytes_ = elems * b * 2
+    mem_s_per_hop = bytes_ / cfg.machine.bw_hbm
+    passes = _cp_compute_passes(cfg)
+    total = mem_s_per_hop * passes
+    total_bytes = bytes_ * passes
+    _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}",
+        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"{total_bytes/1e6:.2f} MB"
+        ),
+    )
+
+
+def stage_pv(cfg: FullConfig) -> StageCost:
+    cmp_hop, _ = _per_hop_qkT_pv(cfg)
+    passes = _cp_compute_passes(cfg)
+    total = cmp_hop * passes
+    T_q = cfg.topo.T_q
+    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
+    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",
+        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}"
+        ),
+        mem_formula="0 (accumulated on-chip)",
+    )
+
+
+def stage_merge(cfg: FullConfig) -> StageCost:
+    """S8: online-softmax merge of partial (o, m, l).
+    - prefill (K/V ring): merge happens in-place across CP hops; no comm here.
+    - decode: merge is the moment we all-reduce partial (o, m, l) across CP
+      ranks; comm is folded into this stage (no separate C1).
+    """
+    T_q = cfg.topo.T_q
+    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)
+    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+
+    comm_s = 0.0
+    comm_bytes = 0
+    comm_formula = ""
+    name_suffix = ""
+
+    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)
+        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:
+            bw, alpha, tier = cfg.machine.bw_intersip, cfg.machine.alpha_intersip, "inter-SIP"
+        else:
+            bw, alpha, tier = cfg.machine.bw_inter, cfg.machine.alpha_inter, "inter-cube"
+        ar_bytes = 2 * (cp - 1) / cp * M
+        comm_s = ar_bytes / bw + 2 * (cp - 1) * alpha
+        comm_bytes = int(ar_bytes)
+        comm_formula = (
+            f"PURPOSE (folded into S8 in decode): each CP rank computed\n"
+            f"attention against its LOCAL slice of KV and holds a partial\n"
+            f"(O, m, l). This all-reduce merges those partials across all\n"
+            f"{cp} CP ranks (online-softmax combine). No separate C1 row\n"
+            f"in decode because this is the only CP comm and it happens\n"
+            f"exactly here, at the end of the local attention body.\n"
+            f"---\n"
+            f"AR of (O + m + l): 2*(CP-1)/CP * M "
+            f"= 2*({cp}-1)/{cp} * {M} B "
+            f"= {ar_bytes:.0f} B over {tier} ({bw/1e9:.0f} GB/s) "
+            f"+ 2*({cp}-1)*alpha"
+        )
+        name_suffix = f" + O/m/l AR (CP={cp} ranks, {tier})"
+
+    visible_s = max(cmp_s, comm_s)
+    if comm_s > cmp_s:
+        bound = "comm"
+    elif cmp_s > 0:
+        bound = "compute"
+    else:
+        bound = "trivial"
+
+    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",
+        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}"
+        ),
+        mem_formula="0 (in-register)",
+        comm_formula=comm_formula or "0",
+    )
+
+
+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
+    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",
+        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}",
+        mem_formula="0",
+    )
+
+
+def stage_wo(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    hq_per_pe = cfg.h_q_per_pe
+    dh = cfg.model.d_head
+    flops = 2 * 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",
+        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",
+    )
+
+
+def comm_cp_ring(cfg: FullConfig) -> StageCost:
+    """CP comm — depends on mode:
+
+    - decode: single O/m/l all-reduce over CP ranks AFTER local S5-S8 finish.
+      No per-hop ring; each PE computes attention against its local K,V once
+      and then contributes (o, m, l) to the reduce.
+    - prefill: per-hop ring during S5/S7 compute (either K/V or Q+O/m/l per
+      cp_ring_variant).
+
+    BW tier depends on cp_placement:
+    - cp_placement=pe:   intra-cube BW for all hops
+    - cp_placement=cube: inter-cube BW; long rings cross SIP boundaries
+    """
+    if cfg.topo.cp <= 1:
+        return StageCost(
+            name="C1 CP comm", formula="C=1 -> no comm",
+            compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
+        )
+    S_local = cfg.topo.s_local
+    hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
+    hq_per_pe = cfg.h_q_per_pe
+    T_q = cfg.topo.T_q
+    dh = cfg.model.d_head
+    b = cfg.model.bytes_per_elem
+
+    # ── 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)
+        if cfg.topo.cp_placement == "pe":
+            bw = cfg.machine.bw_intra
+            alpha = cfg.machine.alpha_intra
+            tier = "intra-cube"
+        elif cfg.topo.sips_used > 1:
+            bw = cfg.machine.bw_intersip
+            alpha = cfg.machine.alpha_intersip
+            tier = "inter-SIP"
+        else:
+            bw = cfg.machine.bw_inter
+            alpha = cfg.machine.alpha_inter
+            tier = "inter-cube"
+        ar_bytes = 2 * (cp - 1) / cp * M
+        ar_time = ar_bytes / bw + 2 * (cp - 1) * alpha
+        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"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,
+            flops=0, mem_bytes=0, comm_bytes=int(ar_bytes),
+            flops_formula="0",
+            mem_formula="0",
+            comm_formula=(
+                f"PURPOSE: In decode, each CP rank computed attention against\n"
+                f"its OWN slice of the KV cache. Each rank now holds a "
+                f"partial\n"
+                f"(O, m, l). This single all-reduce merges those partials "
+                f"across\n"
+                f"all {cp} CP ranks (using online-softmax math) to get the "
+                f"final O.\n"
+                f"---\n"
+                f"AR of (O + m + l): 2*(CP-1)/CP * M "
+                f"= 2*({cp}-1)/{cp} * {M} "
+                f"= {ar_bytes:.0f} B over {tier} at {bw/1e9:.0f} GB/s "
+                f"+ 2*({cp}-1)*alpha"
+            ),
+        )
+
+    # ── 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)
+        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"= {M_KV} B/hop"
+        )
+    else:
+        M_KV = 2 * 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"= {M_KV/1e6:.3f} MB/hop"
+        )
+
+    if cfg.topo.cp_placement == "pe":
+        # All hops intra-cube (fastest).
+        intra_hops = cfg.topo.cp - 1
+        inter_hops = 0
+        intra_time = (intra_hops * M_KV / cfg.machine.bw_intra
+                      + intra_hops * cfg.machine.alpha_intra)
+        inter_time = 0.0
+    else:  # cp_placement == "cube"
+        intra_hops = cfg.topo.cp_intra_sip_hops
+        inter_hops = cfg.topo.cp_inter_sip_hops
+        intra_time = (intra_hops * M_KV / cfg.machine.bw_inter
+                      + intra_hops * cfg.machine.alpha_inter)
+        inter_time = (inter_hops * M_KV / cfg.machine.bw_intersip
+                      + inter_hops * cfg.machine.alpha_intersip)
+    comm_time = intra_time + inter_time
+
+    # Overlap check: per-hop compute (S5+S6+S7) hides intra-SIP hops well;
+    # inter-SIP hops usually dominate.
+    per_hop_cmp, _ = _per_hop_qkT_pv(cfg)
+    per_hop_mem_softmax = (cfg.h_q_per_pe * cfg.topo.T_q
+                          * S_local * cfg.model.bytes_per_elem * 2
+                          / cfg.machine.bw_hbm)
+    per_hop_compute = 2 * per_hop_cmp + per_hop_mem_softmax
+    per_intra_ring = M_KV / cfg.machine.bw_inter + cfg.machine.alpha_inter
+    per_inter_ring = M_KV / cfg.machine.bw_intersip + cfg.machine.alpha_intersip
+    visible_intra = intra_hops * max(0.0, per_intra_ring - per_hop_compute)
+    visible_inter = inter_hops * max(0.0, per_inter_ring - per_hop_compute)
+    visible_total = visible_intra + visible_inter
+
+    tier_desc = f"{intra_hops} intra-SIP + {inter_hops} inter-SIP hops"
+    total_ring_bytes = M_KV * (intra_hops + inter_hops)
+    _var_purpose = (
+        "K, V shards rotate between CP ranks each hop"
+        if cfg.topo.cp_ring_variant == "kv"
+        else "Q + running (O, m, l) rotate between CP ranks each hop"
+    )
+    return StageCost(
+        name=f"C1 CP {variant_desc} ({tier_desc})",
+        formula=f"{formula_bytes}; "
+                f"intra: {intra_hops}*M/BW + inter: {inter_hops}*M/BW",
+        compute_s=0, memory_s=0, comm_s=comm_time,
+        bound="comm", visible_s=visible_total,
+        flops=0, mem_bytes=0, comm_bytes=int(total_ring_bytes),
+        flops_formula="0",
+        mem_formula="0",
+        comm_formula=(
+            f"PURPOSE: CP shards the sequence axis. Each CP rank holds only\n"
+            f"1/{cfg.topo.cp} of the KV cache, so to compute full attention\n"
+            f"we must move data between CP ranks. Variant: {_var_purpose}.\n"
+            f"Runs concurrently with S5/S6/S7 - each hop, compute of the\n"
+            f"just-arrived shard overlaps with comm of the next shard.\n"
+            f"---\n"
+            f"M*(CP-1) with M = {formula_bytes}; "
+            f"total = {total_ring_bytes/1e6:.3f} MB over the ring"
+        ),
+    )
+
+
+def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
+    """TP AllReduce on W_O output."""
+    if cfg.topo.tp <= 1:
+        return StageCost(
+            name="C2 TP AllReduce W_O", formula="TP=1 -> no AR",
+            compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
+        )
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    bytes_ = T_q * d * b
+    tp = cfg.topo.tp
+    tier = cfg.topo.tp_link_tier()  # "intra" | "inter" | "intersip"
+    if tier == "intra":
+        bw, alpha, scope = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
+    elif tier == "inter":
+        bw, alpha, scope = cfg.machine.bw_inter, cfg.machine.alpha_inter, "cross-cube"
+    else:
+        bw, alpha, scope = cfg.machine.bw_intersip, cfg.machine.alpha_intersip, "cross-SIP"
+    comm_time = 2 * (tp - 1) / tp * bytes_ / bw + 2 * (tp - 1) * alpha
+    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}]",
+        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,
+        flops_formula="0",
+        mem_formula="0",
+        comm_formula=(
+            f"PURPOSE: W_O is row-parallel across TP ranks (each rank holds\n"
+            f"1/{tp} of W_O's input rows). After the local W_O GEMM, each of\n"
+            f"the {tp} TP ranks holds a PARTIAL hidden vector (sum over its\n"
+            f"own Q-head slice only). This all-reduce sums those partials so\n"
+            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"= {total_comm_bytes/1e6:.2f} MB over {scope} at "
+            f"{bw/1e9:.0f} GB/s"
+        ),
+    )
+
+
+def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost:
+    """Extra AllReduce on attention scores when TP > H_kv with head-dim split.
+
+    Ranks sharing a KV head have partial scores; must AllReduce across the
+    split_factor group to get final scores per hop.
+    """
+    if not cfg.kv_replication_needed or cfg.topo.kv_shard_mode != "split":
+        return StageCost(
+            name="C3 Score AllReduce (head-split)",
+            formula="TP <= H_kv or replicate mode: not needed",
+            compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
+        )
+    split = cfg.head_dim_split_factor  # ranks sharing one head
+    T_q = cfg.topo.T_q
+    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
+    # split-group AllReduce (intra-cube assumed if group fits)
+    bw = cfg.machine.bw_intra
+    alpha = cfg.machine.alpha_intra
+    per_hop = 2 * (split - 1) / split * bytes_per_hop / bw + 2 * (split - 1) * alpha
+    total = per_hop * cfg.topo.cp
+    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",
+        compute_s=0, memory_s=0, comm_s=total,
+        bound="comm", visible_s=total,
+        flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
+        flops_formula="0",
+        mem_formula="0",
+        comm_formula=(
+            f"PURPOSE: TP={cfg.topo.tp} > H_kv={cfg.model.h_kv}, so each KV\n"
+            f"head is split across {split} TP ranks along the head-dim (d_h\n"
+            f"/{split} per rank). Each rank's Q.K^T is therefore only a\n"
+            f"PARTIAL dot product; the {split} ranks sharing one KV head\n"
+            f"must AllReduce their partial scores to get the true score.\n"
+            f"This fires per hop of the ring, so {cfg.topo.cp}x per layer.\n"
+            f"To eliminate C3: set KV mode = 'replicate' (costs {split}x KV\n"
+            f"memory but no per-hop score AR), or lower TP so TP <= H_kv.\n"
+            f"---\n"
+            f"CP * 2*(split-1)/split * (H_q/TP)*T_q*S_local*b "
+            f"= {cfg.topo.cp} * 2*({split}-1)/{split} * {hq_per_pe}*{T_q}*{S_local}*{b} "
+            f"= {total_comm_bytes/1e6:.2f} MB total"
+        ),
+    )
+
+
+def all_stages(cfg: FullConfig) -> list[StageCost]:
+    """Ordered per-layer attention stages.
+
+    In prefill with CP > 1, C1 (the CP ring) is inserted between S7 and S8
+    because it runs concurrently with S5-S7 per hop. Placing it there in the
+    stage list makes the per-stage table read left-to-right in the order
+    things actually happen. C3 (score AR) sits with C1 for the same reason.
+    C2 (TP AllReduce on W_O) sits right after S10.
+
+    In decode, C1's comm is folded into S8, so no separate C1 row.
+    """
+    stages = [
+        stage_rmsnorm(cfg),
+        stage_wq(cfg),
+        stage_wkv(cfg),
+        stage_kv_append(cfg),
+        stage_qkT(cfg),
+        stage_softmax(cfg),
+        stage_pv(cfg),
+    ]
+    # Concurrent comm (with S5-S7 in prefill) goes here, right after S7.
+    if cfg.topo.mode == "prefill" and cfg.topo.cp > 1:
+        stages.append(comm_cp_ring(cfg))
+    _c3 = comm_kv_split_allreduce(cfg)
+    if _c3.visible_s > 0 or _c3.comm_s > 0:
+        stages.append(_c3)
+    stages.extend([
+        stage_merge(cfg),
+        stage_normalize(cfg),
+        stage_wo(cfg),
+    ])
+    # C2 fires right after S10 (W_O produces per-TP-rank output; AR combines).
+    stages.append(comm_tp_allreduce(cfg))
+    return stages
+
+
+# ── FFN block stages (lightweight; per PE, per layer) ────────────
+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
+    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",
+        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",
+    )
+
+
+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
+    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",
+        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",
+    )
+
+
+def stage_ffn_gate(cfg: FullConfig) -> StageCost:
+    ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+    return _ffn_gemm(cfg, "F2 W_gate GEMM", ffn_per_pe)
+
+
+def stage_ffn_up(cfg: FullConfig) -> StageCost:
+    ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+    return _ffn_gemm(cfg, "F3 W_up GEMM", ffn_per_pe)
+
+
+def stage_ffn_swiglu(cfg: FullConfig) -> StageCost:
+    """SwiGLU element-wise activation on the intermediate FFN tensor."""
+    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
+    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",
+        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",
+    )
+
+
+def stage_ffn_down(cfg: FullConfig) -> StageCost:
+    """W_down: (ffn_per_pe, hidden). Row-parallel over FFN scope."""
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    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
+    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",
+        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",
+    )
+
+
+def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
+    """AllReduce on FFN output across the FFN sharding scope."""
+    scope = cfg.topo.ffn_shard_scope
+    divisor = cfg.ffn_shard_divisor  # TP or TP*CP or TP*CP*DP
+    if divisor <= 1:
+        return StageCost(
+            name="CF1 FFN AllReduce",
+            formula="FFN scope = 1: not needed",
+            compute_s=0, memory_s=0, comm_s=0,
+            bound="trivial", visible_s=0,
+        )
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    bytes_ = 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:
+        bw = cfg.machine.bw_intersip if cfg.topo.sips_used > 1 else cfg.machine.bw_inter
+        alpha = cfg.machine.alpha_intersip if cfg.topo.sips_used > 1 else cfg.machine.alpha_inter
+    else:  # TP only
+        bw = cfg.machine.bw_intra if cfg.topo.tp_spans_cubes == 1 else cfg.machine.bw_inter
+        alpha = cfg.machine.alpha_intra if cfg.topo.tp_spans_cubes == 1 else cfg.machine.alpha_inter
+    comm_time = 2 * (divisor - 1) / divisor * bytes_ / bw + 2 * (divisor - 1) * alpha
+    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",
+        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,
+        flops_formula="0",
+        mem_formula="0",
+        comm_formula=(
+            f"PURPOSE: W_down is row-parallel across the {divisor} ranks in\n"
+            f"the FFN scope (scope={scope}). Each rank produced a PARTIAL\n"
+            f"hidden vector after W_down; this all-reduce sums them so every\n"
+            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"{bw/1e9:.0f} GB/s (scope={scope})"
+        ),
+    )
+
+
+def all_ffn_stages(cfg: FullConfig) -> list[StageCost]:
+    return [
+        stage_ffn_rmsnorm(cfg),
+        stage_ffn_gate(cfg),
+        stage_ffn_up(cfg),
+        stage_ffn_swiglu(cfg),
+        stage_ffn_down(cfg),
+        comm_ffn_allreduce(cfg),
+    ]

tests/analytical_visualization/tensor_sharding.py

@@ -0,0 +1,330 @@
+"""Pictorial view of how each weight/KV tensor is sharded across ranks.
+
+Every tensor is drawn as a rectangle labelled with its (rows, cols)
+shape. Grid lines show shard boundaries; one shard (this PE's share)
+is highlighted so it's obvious which dim the split is along.
+
+Sharding pattern:
+  - W_Q  (hidden, H_q*d_head) : split on **output cols** across TP
+  - W_K  (hidden, H_kv*d_head): split on **output cols** across TP
+  - W_V  (hidden, H_kv*d_head): split on **output cols** across TP
+  - W_O  (H_q*d_head, hidden) : split on **input rows** across TP
+  - W_gate/up (hidden, ffn_dim): split on **output cols** across FFN scope
+  - W_down (ffn_dim, hidden)   : split on **input rows** across FFN scope
+  - K/V cache (S_kv, H_kv*d_head): split on **rows (S axis)** by CP AND
+    **cols (head axis)** by TP -> 2D shard grid
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+
+from .model_config import FullConfig
+
+
+_UNSHARDED_FACE = "#e9ecef"
+_UNSHARDED_EDGE = "#adb5bd"
+_MY_SHARD_FACE = "#3a86ff"
+_MY_SHARD_EDGE = "#0057b3"
+_TP_LINE_COLOR = "#d90429"     # TP split lines
+_CP_LINE_COLOR = "#0f9d58"     # CP split lines
+_FFN_LINE_COLOR = "#e37400"    # FFN-scope split lines
+
+
+def _draw_tensor(ax, x, y, w, h, name, shape_str,
+                  row_splits: int = 1, col_splits: int = 1,
+                  my_row_shard: int = 0, my_col_shard: int = 0,
+                  row_line_color: str = _TP_LINE_COLOR,
+                  col_line_color: str = _TP_LINE_COLOR,
+                  row_label: str = "", col_label: str = "",
+                  note: str = "",
+                  cell_annot=None):
+    """Draw a tensor with sharding grid.
+
+    cell_annot: optional callable (row, col) -> str returning the physical
+    label to render inside shard cell (r, c). If it returns "", cell is not
+    annotated.
+    """
+    """Draw a tensor as a rectangle with grid lines showing shards."""
+    # Base rectangle (unsharded background)
+    rect = patches.Rectangle(
+        (x, y), w, h,
+        facecolor=_UNSHARDED_FACE, edgecolor=_UNSHARDED_EDGE, linewidth=1.0,
+    )
+    ax.add_patch(rect)
+
+    # Highlight this PE's shard
+    shard_w = w / max(1, col_splits)
+    shard_h = h / max(1, row_splits)
+    sx = x + my_col_shard * shard_w
+    sy = y + (row_splits - 1 - my_row_shard) * shard_h
+    highlight = patches.Rectangle(
+        (sx, sy), shard_w, shard_h,
+        facecolor=_MY_SHARD_FACE, edgecolor=_MY_SHARD_EDGE, linewidth=1.5,
+        alpha=0.8,
+    )
+    ax.add_patch(highlight)
+
+    # Horizontal grid lines (row splits)
+    for r in range(1, row_splits):
+        yy = y + r * shard_h
+        ax.plot([x, x + w], [yy, yy],
+                color=row_line_color, linewidth=1.2, linestyle="--")
+    # Vertical grid lines (col splits)
+    for c in range(1, col_splits):
+        xx = x + c * shard_w
+        ax.plot([xx, xx], [y, y + h],
+                color=col_line_color, linewidth=1.2, linestyle="--")
+
+    # Per-cell physical annotations (e.g. "cube3 PE5") if provided.
+    if cell_annot is not None:
+        for r in range(row_splits):
+            for c in range(col_splits):
+                text = cell_annot(r, c)
+                if not text:
+                    continue
+                cx = x + c * shard_w + shard_w / 2
+                cy = y + (row_splits - 1 - r) * shard_h + shard_h / 2
+                ax.text(cx, cy, text,
+                        ha="center", va="center",
+                        fontsize=6.0, color="#0d1b2a",
+                        alpha=0.85)
+
+    # Stacked labels ABOVE rectangle: note (italic, top), shape, name (bold, closest)
+    if note:
+        ax.text(x + w / 2, y + h + 0.55, note,
+                ha="center", va="bottom", fontsize=6.5,
+                color="#555", style="italic")
+    ax.text(x + w / 2, y + h + 0.30, shape_str,
+            ha="center", va="bottom", fontsize=7, color="#666")
+    ax.text(x + w / 2, y + h + 0.08, name,
+            ha="center", va="bottom",
+            fontsize=10, fontweight="bold")
+
+    # Split labels BELOW rectangle (stacked if both present)
+    if row_label and col_label:
+        ax.text(x + w / 2, y - 0.10, row_label,
+                ha="center", va="top", fontsize=6.5,
+                color=row_line_color, fontweight="bold")
+        ax.text(x + w / 2, y - 0.35, col_label,
+                ha="center", va="top", fontsize=6.5,
+                color=col_line_color, fontweight="bold")
+    elif row_label:
+        ax.text(x + w / 2, y - 0.10, row_label,
+                ha="center", va="top", fontsize=6.5,
+                color=row_line_color, fontweight="bold")
+    elif col_label:
+        ax.text(x + w / 2, y - 0.10, col_label,
+                ha="center", va="top", fontsize=6.5,
+                color=col_line_color, fontweight="bold")
+
+
+def _physical_label(cfg: FullConfig, dim: str, rank: int) -> str:
+    """Where does rank `rank` of dim (`tp`|`cp`|`ffn`) live physically?
+    Returns a short "cubeX PEy" style label based on the placement config.
+    """
+    topo = cfg.topo
+    if dim == "tp":
+        if topo.tp_placement == "pe":
+            # TP fills PEs within one cube (spilling if tp > 8)
+            pe_in_cube = rank % topo.pes_per_cube_hw
+            cube = rank // topo.pes_per_cube_hw
+            return f"c{cube}p{pe_in_cube}"
+        else:  # tp on cube
+            return f"cube{rank}"
+    if dim == "cp":
+        if topo.cp_placement == "pe":
+            pe_in_cube = rank % topo.pes_per_cube_hw
+            cube = rank // topo.pes_per_cube_hw
+            return f"c{cube}p{pe_in_cube}"
+        else:  # cp on cube
+            return f"cube{rank}"
+    return ""
+
+
+def _kv_cell_annot(cfg: FullConfig):
+    """Return an annotation function for KV cache cells: (cp_r, tp_r) -> label."""
+    def _ann(r, c):
+        cp_lbl = _physical_label(cfg, "cp", r)
+        tp_lbl = _physical_label(cfg, "tp", c)
+        # Two lines if both are non-trivial; otherwise a single line
+        if cp_lbl and tp_lbl:
+            return f"{cp_lbl}\n{tp_lbl}"
+        return cp_lbl or tp_lbl
+    return _ann
+
+
+def _tp_cell_annot(cfg: FullConfig):
+    """Col-parallel tensors (W_Q/W_K/W_V/W_gate/W_up): each col = one TP rank."""
+    def _ann(r, c):
+        return _physical_label(cfg, "tp", c)
+    return _ann
+
+
+def _tp_row_cell_annot(cfg: FullConfig):
+    """Row-parallel tensors split by TP (W_O): each row = one TP rank."""
+    def _ann(r, c):
+        return _physical_label(cfg, "tp", r)
+    return _ann
+
+
+def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0,
+                          show_physical: bool = False):
+    """Draw all tensors with their sharding overlays for PE `my_pe`.
+
+    show_physical: annotate each shard cell with the owning cube/PE.
+    """
+    m = cfg.model
+    tp = cfg.topo.tp
+    cp = cfg.topo.cp
+
+    # PE's assigned shard index for each dim
+    q_shard = my_pe % tp  # W_Q col shard
+    kv_shard_col = min(my_pe % tp, m.h_kv - 1) if m.h_kv >= tp else 0
+    ffn_shard = my_pe % cfg.ffn_shard_divisor
+    kv_row_shard = 0  # per-CP; showing CP=0 view for clarity
+
+    # Enlarged mode gives more room for physical annotations.
+    _fig_w, _fig_h = (22, 13) if show_physical else (15, 9)
+    _cell_w = 3.4 if show_physical else 2.6
+    _cell_h = 3.0 if show_physical else 2.2
+    _gap_x = 0.8 if show_physical else 0.55
+    _gap_y = 2.0 if show_physical else 1.7
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(_fig_w, _fig_h))
+    else:
+        fig = ax.figure
+
+    # Grid layout: 2 rows x 4 cols of tensor rects
+    # Attention row: W_Q, W_K, W_V, W_O
+    # FFN + KV row: W_gate, W_up, W_down, KV cache
+    cell_w = _cell_w
+    cell_h = _cell_h
+    gap_x = _gap_x
+    gap_y = _gap_y
+
+    _kv_ann = _kv_cell_annot(cfg) if show_physical else None
+    _tp_col_ann = _tp_cell_annot(cfg) if show_physical else None
+    _tp_row_ann = _tp_row_cell_annot(cfg) if show_physical else None
+
+    positions = [
+        (0, 1),  # W_Q at row 0
+        (1, 1),  # W_K
+        (2, 1),  # W_V
+        (3, 1),  # W_O
+        (0, 0),  # W_gate at row 1
+        (1, 0),  # W_up
+        (2, 0),  # W_down
+        (3, 0),  # KV cache
+    ]
+
+    # ── attention tensors ────────────────────────
+    tp_txt = f"TP={tp}"
+    ffn_txt = f"{cfg.topo.ffn_shard_scope}={cfg.ffn_shard_divisor}"
+    cp_txt = f"CP={cp}"
+
+    def px(col): return col * (cell_w + gap_x)
+    def py(row): return row * (cell_h + gap_y)
+
+    # W_Q: (hidden, H_q*d_head), split on COL by TP
+    _draw_tensor(ax, px(0), py(1), cell_w, cell_h,
+                  "W_Q", f"({m.hidden}, {m.h_q * m.d_head})",
+                  row_splits=1, col_splits=tp,
+                  my_col_shard=q_shard,
+                  col_line_color=_TP_LINE_COLOR,
+                  col_label=f"cols split by {tp_txt}",
+                  note="col-parallel (output dim)",
+                  cell_annot=_tp_col_ann)
+
+    # W_K, W_V: (hidden, H_kv*d_head)
+    for (name, col_idx) in [("W_K", 1), ("W_V", 2)]:
+        splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
+        _draw_tensor(ax, px(col_idx), py(1), cell_w, cell_h,
+                      name, f"({m.hidden}, {m.h_kv * m.d_head})",
+                      row_splits=1, col_splits=splits,
+                      my_col_shard=min(q_shard, splits - 1),
+                      col_line_color=_TP_LINE_COLOR,
+                      col_label=f"cols split by {tp_txt}",
+                      note="col-parallel (heads)",
+                      cell_annot=_tp_col_ann)
+
+    # W_O: (H_q*d_head, hidden), split on ROW by TP
+    _draw_tensor(ax, px(3), py(1), cell_w, cell_h,
+                  "W_O", f"({m.h_q * m.d_head}, {m.hidden})",
+                  row_splits=tp, col_splits=1,
+                  my_row_shard=q_shard,
+                  row_line_color=_TP_LINE_COLOR,
+                  row_label=f"rows split by {tp_txt}",
+                  note="row-parallel (input dim)",
+                  cell_annot=_tp_row_ann)
+
+    # ── FFN tensors ──────────────────────────────
+    # W_gate, W_up: (hidden, ffn_dim), split on COL by FFN scope
+    for (name, col_idx) in [("W_gate", 0), ("W_up", 1)]:
+        _draw_tensor(ax, px(col_idx), py(0), cell_w, cell_h,
+                      name, f"({m.hidden}, {m.ffn_dim})",
+                      row_splits=1, col_splits=cfg.ffn_shard_divisor,
+                      my_col_shard=ffn_shard,
+                      col_line_color=_FFN_LINE_COLOR,
+                      col_label=f"cols split by {ffn_txt}",
+                      note="col-parallel")
+
+    # W_down: (ffn_dim, hidden), split on ROW by FFN scope
+    _draw_tensor(ax, px(2), py(0), cell_w, cell_h,
+                  "W_down", f"({m.ffn_dim}, {m.hidden})",
+                  row_splits=cfg.ffn_shard_divisor, col_splits=1,
+                  my_row_shard=ffn_shard,
+                  row_line_color=_FFN_LINE_COLOR,
+                  row_label=f"rows split by {ffn_txt}",
+                  note="row-parallel")
+
+    # ── KV cache: (S_kv, H_kv*d_head), 2D shard ─
+    kv_row_splits = cp
+    kv_col_splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
+    _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:,})",
+                  row_splits=kv_row_splits,
+                  col_splits=kv_col_splits if show_physical else 1,
+                  my_row_shard=kv_row_shard,
+                  row_line_color=_CP_LINE_COLOR,
+                  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)",
+                  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
+    # the _draw_tensor call already drew proper vertical dashes.
+    if not show_physical:
+        kv_x = px(3)
+        kv_y = py(0)
+        shard_col_w = cell_w / max(1, kv_col_splits)
+        for c in range(1, kv_col_splits):
+            xx = kv_x + c * shard_col_w
+            ax.plot([xx, xx], [kv_y, kv_y + cell_h],
+                    color=_TP_LINE_COLOR, linewidth=1.2, linestyle=":")
+
+    ax.set_xlim(-0.3, 4 * (cell_w + gap_x) - gap_x + 0.3)
+    ax.set_ylim(-0.9, 2 * (cell_h + gap_y) - gap_y + 0.9)
+    ax.set_aspect("auto")
+    ax.axis("off")
+
+    ax.set_title(
+        f"Tensor sharding view (PE {my_pe})  |  TP={tp}, CP={cp}, "
+        f"FFN scope={cfg.topo.ffn_shard_scope}  |  "
+        f"blue = this PE's shard",
+        fontsize=11, fontweight="bold",
+    )
+    # Legend for line colors
+    ax.plot([], [], color=_TP_LINE_COLOR, linewidth=1.2, linestyle="--",
+            label=f"TP split ({tp} ranks)")
+    ax.plot([], [], color=_CP_LINE_COLOR, linewidth=1.2, linestyle="--",
+            label=f"CP split ({cp} ranks, sequence dim)")
+    ax.plot([], [], color=_FFN_LINE_COLOR, linewidth=1.2, linestyle="--",
+            label=f"FFN scope split ({cfg.ffn_shard_divisor} ranks)")
+    ax.legend(loc="lower center", bbox_to_anchor=(0.5, -0.05),
+              ncol=3, fontsize=8, frameon=False)
+
+    return fig

tests/analytical_visualization/topology_map.py

@@ -0,0 +1,611 @@
+"""Draw the SIP topology diagram, color-coded by parallelism group.
+
+Color scheme:
+  - CP ring:   orange arrows (unchanged)
+  - Inter-SIP: red arrows across SIP boundaries
+  - PP stage:  cube border color (each PP stage gets a distinct hue)
+  - TP group:  PE fill color (each TP group within a cube gets one hue,
+               matches its cube's PP stage but slightly desaturated)
+  - EP group:  a small badge on the cube (for MoE)
+  - DP replica: cubes of each DP replica get a hatched pattern
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from matplotlib.lines import Line2D
+
+from .model_config import FullConfig
+
+
+MESH_W = 4
+MESH_H = 4
+PE_COLS = 4
+PE_ROWS = 2
+
+
+def _group_label(cfg, group_idx: int) -> str:
+    """Label for a cube-level group under the current placement.
+    - cp_placement=cube -> CP rank (or CP/TP if both on cube)
+    - cp_placement=pe, tp_placement=cube -> TP rank
+    - both on pe -> single group ("PE")
+    """
+    parts = []
+    if cfg.topo.cp_placement == "cube":
+        cp_r = group_idx % max(1, cfg.topo.cp)
+        parts.append(f"CP{cp_r}")
+    if cfg.topo.tp_placement == "cube":
+        div = cfg.topo.cp if cfg.topo.cp_placement == "cube" else 1
+        tp_r = (group_idx // max(1, div)) % max(1, cfg.topo.tp)
+        parts.append(f"TP{tp_r}")
+    return "/".join(parts) if parts else "PE"
+
+# Fully-distinct palette for (PP stage x CP rank) — 24 unique colors.
+# Each cube's (pp, cp) pair maps to one entry; wraps if you exceed 24 groups.
+_GROUP_PALETTE = [
+    "#e6194b", "#3cb44b", "#ffe119", "#4363d8",
+    "#f58231", "#911eb4", "#42d4f4", "#f032e6",
+    "#bfef45", "#fabed4", "#469990", "#dcbeff",
+    "#9a6324", "#fffac8", "#800000", "#aaffc3",
+    "#808000", "#ffd8b1", "#000075", "#a9a9a9",
+    "#004d40", "#c62828", "#4527a0", "#00695c",
+]
+
+# Kept for backward-compat / PP-stage arrows
+_PP_HUES = _GROUP_PALETTE[:8]
+
+
+def _shade(hue_hex: str, t: float) -> str:
+    """Mix hue_hex toward white by t in [0, 1] (0 = original, 1 = white)."""
+    h = hue_hex.lstrip("#")
+    r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)
+    rn = int(r + (255 - r) * t)
+    gn = int(g + (255 - g) * t)
+    bn = int(b + (255 - b) * t)
+    return f"#{rn:02x}{gn:02x}{bn:02x}"
+
+
+def _cp_color(pp_stage: int, cp_rank: int, cp_size: int) -> tuple[str, str]:
+    """Return (border_color, pe_fill_color) for a cube at (pp_stage, cp_rank).
+
+    Each unique (pp, cp) pair gets a **fully distinct color** from the
+    24-entry palette. Border is the pure color; PE fill is a lightened
+    version so the cube border stays visible around the PE grid.
+    """
+    idx = pp_stage * max(1, cp_size) + cp_rank
+    color = _GROUP_PALETTE[idx % len(_GROUP_PALETTE)]
+    pe_fill = _shade(color, 0.30)   # slightly lightened for PE fill
+    return color, pe_fill
+
+# Inactive (unused) colors.
+_INACTIVE_CUBE_FACE = "#f7f7f7"
+_INACTIVE_CUBE_EDGE = "#c8c8c8"
+_INACTIVE_PE = "#eaeaea"
+
+
+def _snake_path(n: int, mesh_w: int, mesh_h: int) -> list[int]:
+    path: list[int] = []
+    for r in range(mesh_h):
+        cols = range(mesh_w) if r % 2 == 0 else range(mesh_w - 1, -1, -1)
+        for c in cols:
+            path.append(r * mesh_w + c)
+            if len(path) == n:
+                return path
+    return path[:n]
+
+
+def _rect_shape(n: int, max_w: int, max_h: int) -> tuple[int, int]:
+    """Return (rows, cols) for n cubes packed as square as possible in max_wxmax_h."""
+    if n <= 0:
+        return (0, 0)
+    # Prefer near-square shapes; cap at mesh dims.
+    best = None
+    for cols in range(1, min(n, max_w) + 1):
+        rows = (n + cols - 1) // cols
+        if rows > max_h:
+            continue
+        # Score: penalise non-squareness + wasted cells
+        waste = rows * cols - n
+        aspect = abs(rows - cols)
+        score = aspect * 10 + waste
+        if best is None or score < best[0]:
+            best = (score, rows, cols)
+    if best is None:
+        # Fallback: single row
+        return (1, min(n, max_w))
+    return best[1], best[2]
+
+
+def _pack_groups_2d(items_per_group: int, n_groups: int,
+                     mesh_w: int, mesh_h: int,
+                     layout_mode: str = "compact") -> list[tuple[int, int]]:
+    """Return list of (row, col) positions for n_groups × items_per_group
+    cubes in a mesh_w x mesh_h mesh.
+
+    layout_mode:
+      - "compact": each group is a near-square rectangle; groups are also
+        tiled 2D as near-square. Ideal for 2x2, 2x4 etc.
+      - "linear": each group is a single row of cubes; groups tile down
+        as rows (original sequential layout).
+    """
+    if n_groups == 0 or items_per_group == 0:
+        return []
+
+    if layout_mode == "linear":
+        # Each group is a single horizontal row of cubes.
+        tp_rows, tp_cols = 1, min(items_per_group, mesh_w)
+    else:  # compact
+        tp_rows, tp_cols = _rect_shape(items_per_group, mesh_w, mesh_h)
+    if tp_rows == 0:
+        return []
+
+    # Group grid: how many groups per row/col of GROUPS
+    max_group_cols = max(1, mesh_w // tp_cols)
+    max_group_rows = max(1, mesh_h // tp_rows)
+    if layout_mode == "linear":
+        # Row-major: fill each mesh row with groups, then wrap down.
+        g_cols = max_group_cols
+        g_rows = max_group_rows
+    else:
+        g_rows, g_cols = _rect_shape(n_groups, max_group_cols, max_group_rows)
+
+    positions: list[tuple[int, int]] = []
+    for g in range(n_groups):
+        gr = g // g_cols
+        gc = g % g_cols
+        base_row = gr * tp_rows
+        base_col = gc * tp_cols
+        for i in range(items_per_group):
+            r = i // tp_cols
+            c = i % tp_cols
+            row = base_row + r
+            col = base_col + c
+            if row >= mesh_h or col >= mesh_w:
+                fallback = (g * items_per_group + i)
+                row = fallback // mesh_w
+                col = fallback % mesh_w
+            positions.append((row, col))
+    return positions
+
+
+def _draw_one_sip(ax, cfg: FullConfig, sip_idx: int,
+                   sip_x0: float, sip_y0: float,
+                   sip_width: float, sip_height: float,
+                   cube_pp_cp: dict[int, tuple[int, int]],
+                   ring_paths: dict[int, list[int]],
+                   tp_group_cubes: dict[int, list[int]]):
+    """Draw one SIP with PP hue + CP shade + TP group boundaries."""
+    cube_size = min(sip_width / MESH_W, sip_height / MESH_H) * 0.85
+    cube_gap = cube_size * 0.15
+    total_w = MESH_W * (cube_size + cube_gap) - cube_gap
+    total_h = MESH_H * (cube_size + cube_gap) - cube_gap
+    ox = sip_x0 + (sip_width - total_w) / 2
+    oy = sip_y0 + (sip_height - total_h) / 2
+
+    sip_rect = patches.FancyBboxPatch(
+        (sip_x0, sip_y0), sip_width, sip_height,
+        boxstyle="round,pad=0.02", facecolor="#f8f9fa",
+        edgecolor="#212529", linewidth=1.5,
+    )
+    ax.add_patch(sip_rect)
+    ax.text(sip_x0 + sip_width / 2, sip_y0 + sip_height + 0.05,
+            f"SIP {sip_idx}", ha="center", va="bottom",
+            fontsize=10, fontweight="bold")
+
+    def _cube_center(cube_local: int) -> tuple[float, float]:
+        r = cube_local // MESH_W
+        c = cube_local % MESH_W
+        x = ox + c * (cube_size + cube_gap)
+        y = oy + (MESH_H - 1 - r) * (cube_size + cube_gap)
+        return x, y
+
+    for cube_local in range(MESH_W * MESH_H):
+        x, y = _cube_center(cube_local)
+        info = cube_pp_cp.get(cube_local, None)
+        is_used = info is not None
+
+        if is_used:
+            pp_stage, cp_rank = info
+            cube_edge, pe_fill = _cp_color(
+                pp_stage, cp_rank,
+                max(1, cfg.topo.inter_cube_dims),
+            )
+            cube_face = _shade(cube_edge, 0.85)
+            cube_lw = 2.0
+        else:
+            cube_edge = _INACTIVE_CUBE_EDGE
+            cube_face = _INACTIVE_CUBE_FACE
+            cube_lw = 0.6
+            pe_fill = _INACTIVE_PE
+
+        rect = patches.FancyBboxPatch(
+            (x, y), cube_size, cube_size,
+            boxstyle="round,pad=0.01",
+            facecolor=cube_face, edgecolor=cube_edge, linewidth=cube_lw,
+        )
+        ax.add_patch(rect)
+        if is_used:
+            pp_stage, cp_rank = info
+            # Small cube ID above
+            ax.text(x + cube_size / 2, y + cube_size + 0.01,
+                    f"cube {cube_local}", ha="center", va="bottom",
+                    fontsize=5.5, color=cube_edge, fontweight="bold")
+            # Group label (CP/TP depending on placement)
+            _lbl = _group_label(cfg, cp_rank)
+            ax.text(x + 0.03, y + cube_size - 0.03,
+                    _lbl, ha="left", va="top",
+                    fontsize=11, color=cube_edge, fontweight="bold",
+                    bbox=dict(boxstyle="round,pad=0.15",
+                              facecolor="white", edgecolor=cube_edge,
+                              linewidth=1.0, alpha=0.9))
+            if cfg.topo.pp > 1:
+                ax.text(x + cube_size - 0.03, y + cube_size - 0.03,
+                        f"PP{pp_stage}", ha="right", va="top",
+                        fontsize=9, color=cube_edge, fontweight="bold",
+                        bbox=dict(boxstyle="round,pad=0.1",
+                                  facecolor="white", edgecolor=cube_edge,
+                                  linewidth=0.8, alpha=0.9))
+        else:
+            ax.text(x + cube_size / 2, y + cube_size + 0.008,
+                    f"{cube_local}", ha="center", va="bottom",
+                    fontsize=6, color=cube_edge)
+
+        # PEs
+        pe_pad = cube_size * 0.08
+        pe_gap = cube_size * 0.02
+        inner_w = cube_size - 2 * pe_pad
+        inner_h = cube_size - 2 * pe_pad
+        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
+        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
+                pe_rect = patches.Rectangle(
+                    (px, py), pe_w, pe_h,
+                    facecolor=fill, edgecolor="#666", linewidth=0.3,
+                )
+                ax.add_patch(pe_rect)
+
+        if is_used and cfg.topo.ep > 1:
+            ax.text(x + cube_size - 0.02, y + cube_size - 0.02,
+                    f"EP{cfg.topo.ep}", ha="right", va="top",
+                    fontsize=6, color="#5f0f40", fontweight="bold")
+
+    # (CP ring arrows removed per user request)
+
+    # Draw TP group bounding boxes (RED dashed) when TP > 8
+    for tp_g, cubes in tp_group_cubes.items():
+        if len(cubes) <= 1:
+            continue  # TP fits in one cube, redundant with cube border
+        # Bounding box around all cubes in this TP group
+        xs, ys = [], []
+        for cube_local in cubes:
+            x, y = _cube_center(cube_local)
+            xs.extend([x, x + cube_size])
+            ys.extend([y, y + cube_size])
+        pad = 0.03
+        bbox_x = min(xs) - pad
+        bbox_y = min(ys) - pad
+        bbox_w = (max(xs) - min(xs)) + 2 * pad
+        bbox_h = (max(ys) - min(ys)) + 2 * pad
+        tp_rect = patches.Rectangle(
+            (bbox_x, bbox_y), bbox_w, bbox_h,
+            fill=False, edgecolor="#d90429", linewidth=1.8,
+            linestyle="--",
+        )
+        ax.add_patch(tp_rect)
+        ax.text(bbox_x + bbox_w / 2, bbox_y - 0.05,
+                f"TP group {tp_g}", ha="center", va="top",
+                color="#d90429", fontsize=7, fontweight="bold")
+
+
+def _sip_grid_layout(n_sips: int, topology: str) -> tuple[int, int]:
+    """(rows, cols) for arranging SIPs on the page.
+    - ring: horizontal chain (1 x N)
+    - mesh2d / torus2d: near-square grid
+    """
+    if n_sips <= 0:
+        return 0, 0
+    if topology == "ring" or n_sips <= 2:
+        return 1, n_sips
+    import math
+    cols = int(math.ceil(math.sqrt(n_sips)))
+    rows = (n_sips + cols - 1) // cols
+    return rows, cols
+
+
+def _draw_sip_links(ax, sip_xy: list[tuple[float, float]],
+                     topology: str, sip_w: float, sip_h: float,
+                     grid_rows: int, grid_cols: int):
+    """Draw inter-SIP interconnect lines.
+
+    Solid lines: adjacent (grid-neighbor) links.
+    Dashed lines: wrap-around links (ring, torus).
+    """
+    n = len(sip_xy)
+    if n <= 1:
+        return
+    color = "#c62828"
+    lw = 2.0
+
+    def center(i):
+        x, y = sip_xy[i]
+        return x + sip_w / 2, y + sip_h / 2
+
+    if topology == "ring":
+        # 1D chain (adjacent solid)
+        for i in range(n - 1):
+            x1, y1 = center(i)
+            x2, y2 = center(i + 1)
+            ax.plot([x1 + sip_w / 2, x2 - sip_w / 2], [y1, y2],
+                    color=color, linewidth=lw)
+        # wrap-around from last back to first (dashed arc below)
+        if n > 2:
+            x_first, y_first = center(0)
+            x_last, y_last = center(n - 1)
+            arc_y = min(y_first, y_last) - sip_h / 2 - 0.35
+            ax.plot([x_first, x_first], [y_first - sip_h / 2, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x_first, x_last], [arc_y, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x_last, x_last], [arc_y, y_last - sip_h / 2],
+                    color=color, linewidth=lw, linestyle="--")
+        elif n == 2:
+            # 2 SIPs: single link
+            x1, y1 = center(0)
+            x2, y2 = center(1)
+            # wrap link: dashed loop below
+            arc_y = y1 - sip_h / 2 - 0.35
+            ax.plot([x1, x1], [y1 - sip_h / 2, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x1, x2], [arc_y, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x2, x2], [arc_y, y2 - sip_h / 2],
+                    color=color, linewidth=lw, linestyle="--")
+        return
+
+    # mesh2d / torus2d: grid neighbours
+    def idx(r, c):
+        return r * grid_cols + c
+
+    for r in range(grid_rows):
+        for c in range(grid_cols):
+            i = idx(r, c)
+            if i >= n:
+                continue
+            # right
+            if c + 1 < grid_cols and idx(r, c + 1) < n:
+                x1, y1 = center(i)
+                x2, y2 = center(idx(r, c + 1))
+                ax.plot([x1 + sip_w / 2, x2 - sip_w / 2], [y1, y2],
+                        color=color, linewidth=lw)
+            # down (visually — grid row+1 is drawn below)
+            if r + 1 < grid_rows and idx(r + 1, c) < n:
+                x1, y1 = center(i)
+                x2, y2 = center(idx(r + 1, c))
+                ax.plot([x1, x2], [y1 - sip_h / 2, y2 + sip_h / 2],
+                        color=color, linewidth=lw)
+
+    if topology == "torus2d":
+        # row wrap: last col -> first col in each row (dashed arc BELOW that row)
+        for r in range(grid_rows):
+            left_i = idx(r, 0)
+            right_i = idx(r, grid_cols - 1)
+            if left_i >= n or right_i >= n or left_i == right_i:
+                continue
+            x1, y1 = center(left_i)
+            x2, y2 = center(right_i)
+            arc_y = min(y1, y2) - sip_h / 2 - 0.35
+            ax.plot([x1, x1], [y1 - sip_h / 2, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x1, x2], [arc_y, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x2, x2], [arc_y, y2 - sip_h / 2],
+                    color=color, linewidth=lw, linestyle="--")
+        # col wrap: last row -> first row in each col (dashed arc to the RIGHT)
+        for c in range(grid_cols):
+            top_i = idx(0, c)
+            bot_i = idx(grid_rows - 1, c)
+            if top_i >= n or bot_i >= n or top_i == bot_i:
+                continue
+            x1, y1 = center(top_i)   # top row (higher y)
+            x2, y2 = center(bot_i)   # bottom row (lower y)
+            arc_x = max(x1, x2) + sip_w / 2 + 0.35
+            ax.plot([x1 + sip_w / 2, arc_x], [y1, y1],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([arc_x, arc_x], [y1, y2],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([arc_x, x2 + sip_w / 2], [y2, y2],
+                    color=color, linewidth=lw, linestyle="--")
+
+
+def draw_topology(cfg: FullConfig, ax=None, layout_mode: str = "linear"):
+    """Draw one or more SIPs. Each (PP,CP) group gets a unique color.
+
+    SIP grid arrangement + inter-SIP links depend on cfg.topo.sip_topology:
+      "ring"    -> 1D horizontal chain, wrap arc below
+      "mesh2d"  -> near-square grid, no wrap
+      "torus2d" -> grid + row/col wrap arcs
+    """
+    sips_used = max(1, cfg.topo.sips_used)
+    total_pes = cfg.topo.total_pes
+    pes_per_stage = cfg.topo.pes_per_stage    # CP*TP
+    n_replicas = cfg.topo.dp
+    n_pp_stages = cfg.topo.pp
+    pes_per_cube = cfg.topo.pes_per_cube_hw
+    # Placement-aware: how many cubes per "group" (was: cubes-per-TP).
+    # A group is one cube-level unit (one CP rank if cp on cube, or one TP
+    # rank if only tp on cube, etc.). Cubes per group == intra-cube spill.
+    cubes_per_group = max(
+        1, (cfg.topo.intra_cube_dims + pes_per_cube - 1) // pes_per_cube
+    )
+    n_groups_per_stage = max(1, cfg.topo.inter_cube_dims)
+    cubes_per_stage = cfg.topo.cubes_per_stage
+    cubes_used = cfg.topo.cubes_used
+    cubes_per_sip = MESH_W * MESH_H
+
+    sip_topo = cfg.topo.sip_topology
+    grid_rows, grid_cols = _sip_grid_layout(sips_used, sip_topo)
+
+    sip_w = 4.0
+    sip_h = 4.0
+    sip_gap = 0.6
+
+    fig_w = grid_cols * sip_w + (grid_cols - 1) * sip_gap + 1.5
+    fig_h = grid_rows * sip_h + (grid_rows - 1) * sip_gap + 2.0
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(max(6, fig_w), max(6, fig_h)))
+    else:
+        fig = ax.figure
+
+    # 2D packing: each group's cubes form a rectangle inside the SIP mesh.
+    # If more groups than fit in one SIP, spill to next SIP.
+    tp_rows, tp_cols = _rect_shape(cubes_per_group, MESH_W, MESH_H)
+    if tp_rows == 0:
+        tp_rows, tp_cols = 1, 1
+    groups_per_sip = max(1, (MESH_W // tp_cols) * (MESH_H // tp_rows))
+
+    cube_pp_cp_global: dict[int, tuple[int, int]] = {}
+    ring_globals_by_pp: dict[int, list[int]] = {}
+    tp_group_cubes_global: dict[int, list[int]] = {}
+
+    current_sip = 0
+    for rep in range(n_replicas):
+        for pp_s in range(n_pp_stages):
+            # Chunk cube-level groups into per-SIP batches.
+            for chunk_start in range(0, n_groups_per_stage, groups_per_sip):
+                chunk = min(groups_per_sip, n_groups_per_stage - chunk_start)
+                positions = _pack_groups_2d(cubes_per_group, chunk,
+                                             MESH_W, MESH_H,
+                                             layout_mode=layout_mode)
+                for local_gi in range(chunk):
+                    group_idx = chunk_start + local_gi
+                    tp_gid = ((rep * n_pp_stages) + pp_s) * n_groups_per_stage + group_idx
+                    cubes_here: list[int] = []
+                    for i in range(cubes_per_group):
+                        idx_in_group = local_gi * cubes_per_group + i
+                        r, c = positions[idx_in_group]
+                        local_cube = r * MESH_W + c
+                        gc = current_sip * cubes_per_sip + local_cube
+                        cube_pp_cp_global[gc] = (pp_s, group_idx)
+                        cubes_here.append(gc)
+                    tp_group_cubes_global[tp_gid] = cubes_here
+                    ring_globals_by_pp.setdefault(pp_s, []).append(cubes_here[0])
+                current_sip += 1  # next SIP for the next chunk (or (rep, pp))
+
+    # Partition (pp, cp) map by SIP using cube_positions_global
+    cubes_by_sip: list[dict[int, tuple[int, int]]] = [dict() for _ in range(sips_used)]
+    for gc, ppcp in cube_pp_cp_global.items():
+        sip_idx = gc // cubes_per_sip
+        local_cube = gc % cubes_per_sip
+        if sip_idx < sips_used:
+            cubes_by_sip[sip_idx][local_cube] = ppcp
+
+    # Per-SIP snake-path CP rings (organized by pp_stage)
+    ring_by_sip: list[dict[int, list[int]]] = [dict() for _ in range(sips_used)]
+    for pp_s, gc_list in ring_globals_by_pp.items():
+        by_sip_local: dict[int, list[int]] = {}
+        for gc in gc_list:
+            si = gc // cubes_per_sip
+            lc = gc % cubes_per_sip
+            by_sip_local.setdefault(si, []).append(lc)
+        for si, locals_list in by_sip_local.items():
+            if si < sips_used:
+                # Snake through the packed positions in order
+                ring_by_sip[si][pp_s] = locals_list
+
+    # TP group cubes by SIP
+    tp_by_sip: list[dict[int, list[int]]] = [dict() for _ in range(sips_used)]
+    for tp_g, gc_list in tp_group_cubes_global.items():
+        by_sip_local: dict[int, list[int]] = {}
+        for gc in gc_list:
+            si = gc // cubes_per_sip
+            lc = gc % cubes_per_sip
+            by_sip_local.setdefault(si, []).append(lc)
+        for si, locals_list in by_sip_local.items():
+            if si < sips_used:
+                tp_by_sip[si][tp_g] = locals_list
+
+    # Compute (x, y) top-left for each SIP based on grid layout.
+    # Row 0 is drawn at the TOP visually, so y decreases with row index.
+    sip_xy: list[tuple[float, float]] = []
+    for sip_idx in range(sips_used):
+        r = sip_idx // grid_cols
+        c = sip_idx % grid_cols
+        x = c * (sip_w + sip_gap)
+        y = (grid_rows - 1 - r) * (sip_h + sip_gap)
+        sip_xy.append((x, y))
+
+    for sip_idx in range(sips_used):
+        x0, y0 = sip_xy[sip_idx]
+        _draw_one_sip(
+            ax, cfg, sip_idx,
+            sip_x0=x0, sip_y0=y0,
+            sip_width=sip_w, sip_height=sip_h,
+            cube_pp_cp=cubes_by_sip[sip_idx],
+            ring_paths=ring_by_sip[sip_idx],
+            tp_group_cubes=tp_by_sip[sip_idx],
+        )
+
+    # Inter-SIP interconnect links per user-selected topology.
+    _draw_sip_links(ax, sip_xy, sip_topo, sip_w, sip_h, grid_rows, grid_cols)
+
+    total_w = grid_cols * sip_w + (grid_cols - 1) * sip_gap
+    total_h = grid_rows * sip_h + (grid_rows - 1) * sip_gap
+    # Extra room: bottom for wrap arcs, right for torus col-wrap arcs, top for SIP labels
+    ax.set_xlim(-0.3, total_w + 0.9)
+    ax.set_ylim(-0.9, total_h + 0.7)
+    ax.set_aspect("equal")
+    ax.axis("off")
+
+    # Legend — one entry per unique (PP, cube-group) + TP boundary + SIP links
+    legend_handles = []
+    for pp_s in range(n_pp_stages):
+        for gi in range(n_groups_per_stage):
+            color, _ = _cp_color(pp_s, gi, n_groups_per_stage)
+            legend_handles.append(Line2D(
+                [], [], color=color, marker="s", linestyle="",
+                markersize=10,
+                label=(f"PP{pp_s} " if n_pp_stages > 1 else "")
+                       + _group_label(cfg, gi),
+            ))
+    if cfg.topo.tp > cfg.topo.pes_per_cube_hw:
+        legend_handles.append(Line2D(
+            [], [], color="#d90429", marker="s", linestyle="--",
+            markersize=10, markerfacecolor="none", markeredgewidth=1.5,
+            label="TP group (dashed red)",
+        ))
+    if cfg.topo.ep > 1:
+        legend_handles.append(Line2D(
+            [], [], color="#5f0f40", marker="s", linestyle="",
+            markersize=8, label=f"EP={cfg.topo.ep}",
+        ))
+    if sips_used > 1:
+        _topo_label = {
+            "ring": "Inter-SIP: ring (chain + wrap)",
+            "mesh2d": "Inter-SIP: 2D mesh (no wrap)",
+            "torus2d": "Inter-SIP: 2D torus (grid + wrap)",
+        }.get(sip_topo, f"Inter-SIP: {sip_topo}")
+        legend_handles.append(Line2D(
+            [], [], color="#c62828", linewidth=2.0,
+            label=_topo_label,
+        ))
+    # Legend below the figure; wrap ncols so it's readable
+    n_items = len(legend_handles)
+    ncol = min(6, max(3, n_items))
+    ax.legend(handles=legend_handles, loc="upper center",
+              bbox_to_anchor=(0.5, -0.02), ncol=ncol,
+              fontsize=8, frameon=False)
+
+    title = (
+        f"Layout: {total_pes} PEs = {n_replicas} DP replica(s) x "
+        f"{n_pp_stages} PP stage(s) x {cubes_per_stage} cubes (CP={cfg.topo.cp}) "
+        f"x TP={cfg.topo.tp} | cubes: {cubes_used}, SIPs: {sips_used}"
+    )
+    ax.set_title(title, fontsize=10)
+
+    return fig