diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py
index 0e54dea..3e63c6e 100644
--- a/tests/analytical_visualization/app.py
+++ b/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():
diff --git a/tests/analytical_visualization/stage_shapes.py b/tests/analytical_visualization/stage_shapes.py
new file mode 100644
index 0000000..8d0d9b8
--- /dev/null
+++ b/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
diff --git a/tests/analytical_visualization/test_stage_shapes.py b/tests/analytical_visualization/test_stage_shapes.py
new file mode 100644
index 0000000..aeae835
--- /dev/null
+++ b/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)