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():