From bccc90db82a1c966e69fd4c4a2bc640ecc35c65f Mon Sep 17 00:00:00 2001
From: Mukesh Garg <mukeshgarg81@gmail.com>
Date: Tue, 28 Jul 2026 14:21:27 -0700
Subject: [PATCH] analytical-viz: Physical Layout scope filter for stage tables
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Clicking one of the three Physical Layout tab buttons now persists the
chosen scope in session_state["_pl_active_scope"] and filters the
per-stage latency tables accordingly:

  - attn scope → show only "Attention" stage table
  - ffn scope  → show only "FFN" stage table
  - full scope → show both (default; also matches a fresh sidebar-driven
                  config with no button click yet)

Added a "Layout scope: {label}" header so the user can tell at a glance
which scope's config is loaded.

The pipeline diagram + topology map + weight/tensor sharding + PE
layout continue to show the full model layout — those diagrams give
context that stays useful even when the user is focusing on attn or
FFN. Deeper filtering (e.g., attention-only pipeline stripe) can come
later if needed.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 tests/analytical_visualization/app.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py
index bf4e97b..c7b4c50 100644
--- a/tests/analytical_visualization/app.py
+++ b/tests/analytical_visualization/app.py
@@ -521,6 +521,9 @@ with tab_layout:
             }
             st.session_state["ffn_scope_label"] = \
                 _pl_ffn_map[_pl_best.ffn_shard_scope]
+            # Persist the chosen scope so subsequent layout renders filter
+            # stage tables (attn/ffn/full).
+            st.session_state["_pl_active_scope"] = _pl_scope
             st.success(
                 f"Loaded latency-optimal {_pl_lbl} config: CP={_pl_best.cp} "
                 f"TP={_pl_best.tp} PP={_pl_best.pp} DP={_pl_best.dp} → "
@@ -529,6 +532,14 @@ with tab_layout:
             )
             st.rerun()
 
+    # Determine the effective scope for filtering the tables below. Default
+    # to "full" so a fresh page load / sidebar-driven config shows everything.
+    _pl_active_scope = st.session_state.get("_pl_active_scope", "full")
+    _pl_show_attn = _pl_active_scope in ("attn", "full")
+    _pl_show_ffn = _pl_active_scope in ("ffn", "full")
+    _pl_active_label = _pl_scope_meta[_pl_active_scope][2]
+    st.markdown(f"### Layout scope: **{_pl_active_label}**")
+
     st.divider()
 
     # Row A: pipeline diagram (color-coded) + per-stage table (formulas).
@@ -688,8 +699,10 @@ with tab_layout:
                 f"blue=compute, orange=memory, red=comm, grey=trivial."
             )
 
-        _render_table(_stages_attn, "Attention")
-        _render_table(_stages_ffn, "FFN")
+        if _pl_show_attn:
+            _render_table(_stages_attn, "Attention")
+        if _pl_show_ffn:
+            _render_table(_stages_ffn, "FFN")
 
         with st.expander("Symbol glossary (what does T_q, d_h, div... mean?)",
                           expanded=False):