analytical-viz: per-stage tensor shape tables (attention + FFN)

New stage_shapes module reports per-PE input / weight / output tensor
shapes for every attention (S1..S10, C1..C3) and FFN (F1..F5, CF1)
stage. Streamlit's Stages tab picks these up below the existing
latency table + bar chart as two separate dataframes ('Attention
stages' and 'FFN stages'), so anyone who wants to see the exact
sharded tensor dims per stage can, without cluttering the latency
view.

Rows mirror the conditional structure of all_stages / all_ffn_stages:
S8 merge only when CP>1; C1 CP ring only in prefill with CP>1; C3
Score AR only when kv_shard_mode='split' and TP>H_kv; C2 TP AR only
when TP>1; CF1 FFN AR only when the FFN divisor > 1. Batch B flows
into every activation dim; weights stay B-invariant.

test_stage_shapes covers column presence, conditional row insertion
(decode omits C1; CP=1 omits S8; prefill+CP>1 inserts C1), batch
scaling of the leading dim, and that S5 output references both T_q
and S_local for the current cfg.

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

tests/analytical_visualization/app.py

@@ -29,6 +29,7 @@ for _mod_name in (
     "tests.analytical_visualization.autosuggest",
     "tests.analytical_visualization.memory_layout",
     "tests.analytical_visualization.stage_latencies",
+    "tests.analytical_visualization.stage_shapes",
     "tests.analytical_visualization.auto_explore",
     "tests.analytical_visualization.auto_hardware",
     "tests.analytical_visualization.pe_weight_layout",
@@ -46,6 +47,10 @@ from tests.analytical_visualization.model_config import (
 )
 from tests.analytical_visualization.model_presets import PRESETS
 from tests.analytical_visualization.stage_latencies import all_stages, all_ffn_stages
+from tests.analytical_visualization.stage_shapes import (
+    attn_stage_shape_rows,
+    ffn_stage_shape_rows,
+)
 from tests.analytical_visualization.memory_layout import (
     compute_memory,
     attention_weight_rows,
@@ -1222,6 +1227,22 @@ with tab_stages:
     st.pyplot(fig_bar, width='stretch')
     plt.close(fig_bar)
 
+    # ── Per-stage tensor shapes (per PE) ───────────────────────────
+    st.divider()
+    st.subheader("Per-stage tensor shapes (per PE)")
+    st.caption(
+        "Input / weight / output tensor dimensions each stage operates on, "
+        "*per PE* (already sharded by CP, TP, PP, EP as configured). Rows "
+        "match the latency table above; comm-only stages show payload in "
+        "the Input column."
+    )
+    st.markdown("**Attention stages**")
+    st.dataframe(pd.DataFrame(attn_stage_shape_rows(cfg)),
+                 width='stretch', hide_index=True)
+    st.markdown("**FFN stages**")
+    st.dataframe(pd.DataFrame(ffn_stage_shape_rows(cfg)),
+                 width='stretch', hide_index=True)
+
 
 # ── TAB 4: Save & compare ───────────────────────────────────────
 def _snapshot_current_config():

tests/analytical_visualization/stage_shapes.py

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

tests/analytical_visualization/test_stage_shapes.py

@@ -0,0 +1,115 @@
+"""Tests for stage_shapes: per-PE shape rows for attention + FFN stages."""
+from __future__ import annotations
+
+from tests.analytical_visualization.model_config import (
+    FullConfig, MachineParams, TopologyConfig,
+)
+from tests.analytical_visualization.model_presets import PRESETS
+from tests.analytical_visualization.stage_shapes import (
+    attn_stage_shape_rows,
+    ffn_stage_shape_rows,
+)
+
+
+def _cfg(cp=2, tp=8, pp=1, b=1, mode="decode", s_kv=8192):
+    model = PRESETS["Llama 3 8B"].model
+    topo = TopologyConfig(cp=cp, tp=tp, pp=pp, b=b, s_kv=s_kv, mode=mode)
+    return FullConfig(model=model, topo=topo, machine=MachineParams())
+
+
+# ── Attention shapes ────────────────────────────────────────────────
+
+
+def test_attn_rows_have_required_columns():
+    """Every attention row must define Stage, Op, Input, Weight, Output."""
+    rows = attn_stage_shape_rows(_cfg())
+    required = {"Stage", "Op", "Input (per PE)",
+                "Weight (per PE)", "Output (per PE)"}
+    assert rows, "empty attention shape rows"
+    for r in rows:
+        missing = required - r.keys()
+        assert not missing, f"row {r.get('Stage')} missing {missing}"
+        for k in required:
+            assert r[k], f"row {r['Stage']} col {k} is empty"
+
+
+def test_attn_row_prefixes_present_in_order():
+    """S1..S10 always appear; C1 only when prefill+CP>1; C2 when TP>1;
+    C3 when kv_shard_mode=split and TP>H_kv; S8 only when CP>1."""
+    cfg = _cfg(cp=2, tp=8, mode="decode")
+    prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
+    # decode with CP=2, TP=8 (h_kv=8 for Llama 3 8B) → no C1, no C3; C2 yes; S8 yes.
+    assert "C1" not in prefixes, prefixes
+    assert "C3" not in prefixes, prefixes
+    assert "S8" in prefixes, prefixes
+    assert "C2" in prefixes, prefixes
+    # Ensure S1..S10 sequence appears in relative order.
+    s_order = [p for p in prefixes if p.startswith("S")]
+    assert s_order == sorted(s_order, key=lambda s: int(s[1:])), s_order
+
+
+def test_attn_cp1_omits_s8_merge():
+    """CP=1 → no online-softmax merge stage."""
+    cfg = _cfg(cp=1, tp=8, mode="decode")
+    prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
+    assert "S8" not in prefixes
+
+
+def test_attn_prefill_cp_gt_1_inserts_c1():
+    """Prefill mode with CP>1 shows the CP-ring stage."""
+    cfg = _cfg(cp=4, tp=2, mode="prefill", s_kv=8192)
+    prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
+    assert "C1" in prefixes
+
+
+def test_attn_shape_reflects_batch():
+    """B appears as the leading dim in each activation shape."""
+    for b in (1, 4, 32):
+        rows = attn_stage_shape_rows(_cfg(b=b))
+        s1 = next(r for r in rows if r["Stage"] == "S1")
+        assert s1["Input (per PE)"].startswith(f"({b},"), s1
+        assert s1["Output (per PE)"].startswith(f"({b},"), s1
+
+
+def test_attn_s5_shape_uses_tq_and_slocal():
+    """S5 (Q·K^T) shape must reference both T_q and S_local values."""
+    cfg = _cfg(cp=2, tp=8, mode="decode", s_kv=8192)
+    T_q = cfg.topo.T_q
+    S_local = cfg.topo.s_local
+    rows = attn_stage_shape_rows(cfg)
+    s5 = next(r for r in rows if r["Stage"] == "S5")
+    # Output shape (B, H_q/TP, T_q, S_local)
+    assert str(T_q) in s5["Output (per PE)"], s5
+    assert str(S_local) in s5["Output (per PE)"], s5
+
+
+# ── FFN shapes ──────────────────────────────────────────────────────
+
+
+def test_ffn_rows_have_required_columns():
+    """Every FFN row must define Stage, Op, Input, Weight, Output."""
+    rows = ffn_stage_shape_rows(_cfg())
+    required = {"Stage", "Op", "Input (per PE)",
+                "Weight (per PE)", "Output (per PE)"}
+    assert rows, "empty ffn shape rows"
+    for r in rows:
+        missing = required - r.keys()
+        assert not missing, f"row {r.get('Stage')} missing {missing}"
+
+
+def test_ffn_row_prefixes_f1_to_f5_always_present():
+    """F1..F5 always emitted; CF1 only when FFN divisor > 1."""
+    rows = ffn_stage_shape_rows(_cfg(cp=2, tp=8))
+    prefixes = [r["Stage"] for r in rows]
+    for p in ("F1", "F2", "F3", "F4", "F5"):
+        assert p in prefixes, prefixes
+
+
+def test_ffn_gate_up_output_matches_ffn_per_pe():
+    """F2/F3 output must have ffn_per_pe as the last dim."""
+    cfg = _cfg(cp=2, tp=8)
+    ffn_pe = max(1, cfg.model.ffn_dim
+                 // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+    rows = ffn_stage_shape_rows(cfg)
+    f2 = next(r for r in rows if r["Stage"] == "F2")
+    assert str(ffn_pe) in f2["Output (per PE)"], (ffn_pe, f2)