From 09721040caa906eb580921f096807c06d65c1eec Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Tue, 28 Jul 2026 14:11:30 -0700 Subject: [PATCH] analytical-viz: 3 sweep buttons at top of Physical Layout tab MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Shortcut: click any of {Attention, FFN/MoE, Attn+FFN/MoE} → runs run_auto_explore for that scope, picks the latency-minimum Pareto config, loads it into the sidebar's session_state (cp, tp, pp, dp, tp_placement, cp_placement, cp_ring_variant, kv_mode, ffn_scope_label), then st.rerun(). Because the Physical Layout tab reads model + machine + cfg from the sidebar, the sidebar update automatically redraws the pipeline diagram, per-stage table, and everything below. No preview / parallel-display state — WYSIWYG. Ffn_scope_label reconstruction handles the dynamic "(div=…)" suffix the sidebar shows (same pattern as auto_explore + auto_hardware tabs' "Load into sidebar" widgets). Verified: app.py parses; 24 pytest tests still pass. Co-Authored-By: Claude Opus 4.7 --- tests/analytical_visualization/app.py | 79 +++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 7146600..9ab82df 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -432,6 +432,85 @@ tab_auto, tab_layout, tab_memory, tab_stages, tab_compare, tab_hw = st.tabs([ # ── TAB 1: Physical layout ────────────────────────────────────── with tab_layout: + # Quick shortcut: three buttons that pick the latency-optimal Pareto + # config for a chosen scope and load it into the sidebar. Layout + # re-renders immediately on the next Streamlit run with the new config. + from tests.analytical_visualization.auto_explore import ( + run_auto_explore as _run_auto_explore_layout, + ) + + st.markdown( + "**Quick auto-suggest layout** — pick a scope; the latency-optimal " + "Pareto config for that scope is loaded into the sidebar and this " + "layout redraws around it." + ) + _pl1, _pl2, _pl3, _pl4 = st.columns([1, 1, 1, 1]) + with _pl1: + _pl_run_attn = st.button("Run sweep — Attention", + type="primary", width='stretch', + key="_pl_run_attn") + with _pl2: + _pl_run_ffn = st.button("Run sweep — FFN/MoE", + type="primary", width='stretch', + key="_pl_run_ffn") + with _pl3: + _pl_run_full = st.button("Run sweep — Attn + FFN/MoE", + type="primary", width='stretch', + key="_pl_run_full") + + _pl_scope_meta = { + "attn": (True, False, "Attention only"), + "ffn": (False, True, "FFN / MoE only"), + "full": (True, True, "Attn + FFN/MoE"), + } + _pl_button_click = { + "attn": _pl_run_attn, "ffn": _pl_run_ffn, "full": _pl_run_full, + } + for _pl_scope, _pl_clicked in _pl_button_click.items(): + if not _pl_clicked: + continue + _pl_ia, _pl_ff, _pl_lbl = _pl_scope_meta[_pl_scope] + with st.spinner(f"Finding latency-optimal parallelism ({_pl_lbl})..."): + _pl_res = _run_auto_explore_layout( + model, machine, s_kv=s_kv, mode=mode, + include_attention=_pl_ia, include_ffn=_pl_ff, + ) + if not _pl_res.pareto_scores: + st.error( + f"No feasible parallelism for {_pl_lbl} — every config " + "exceeds per-PE HBM. Try raising `pe_hbm_gb` or reducing " + "`s_kv`." + ) + else: + _pl_best = min(_pl_res.pareto_scores, + key=lambda s: s.total_latency_ns) + # Load the winning config into the sidebar's session state. + st.session_state["cp"] = _pl_best.cp + st.session_state["tp"] = _pl_best.tp + st.session_state["pp"] = _pl_best.pp + st.session_state["dp"] = _pl_best.dp + st.session_state["tp_placement"] = _pl_best.tp_placement + st.session_state["cp_placement"] = _pl_best.cp_placement + st.session_state["cp_ring_variant"] = _pl_best.cp_ring_variant + st.session_state["kv_mode"] = _pl_best.kv_shard_mode + _pl_tp, _pl_cp, _pl_dp = _pl_best.tp, _pl_best.cp, _pl_best.dp + _pl_ffn_map = { + "TP": f"TP only (div={max(1, _pl_tp)})", + "TP+CP": f"TP*CP (div={max(1, _pl_tp * _pl_cp)})", + "TP+CP+DP": f"TP*CP*DP (div={max(1, _pl_tp * _pl_cp * _pl_dp)})", + } + st.session_state["ffn_scope_label"] = \ + _pl_ffn_map[_pl_best.ffn_shard_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} → " + f"{_pl_best.latency_ms:.2f} ms, {_pl_best.pes_used} PEs. " + "Redrawing layout..." + ) + st.rerun() + + st.divider() + # Row A: pipeline diagram (color-coded) + per-stage table (formulas). st.markdown("**Per-layer pipeline (color = dominant bound)**") fig_pipe = draw_pipeline(cfg)