From 9afeb6bc1650afdc67befadac29444e3778d127c Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Tue, 28 Jul 2026 14:24:06 -0700 Subject: [PATCH] analytical-viz: 3 sidebar apply-buttons (Attn / FFN / Attn+FFN) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the single memory-only "Apply auto-suggest" button in the sidebar Parallelism expander with three latency-optimal buttons: "Attention", "FFN/MoE", "Attn+FFN". Each clicked button runs run_auto_explore for its scope, picks the latency-minimum Pareto config, snaps the parallelism knobs to the sidebar's selectbox option sets (CP/TP/PP/DP), and loads the other knobs directly (tp_placement, cp_placement, cp_ring_variant, kv_mode, ffn_scope_label). Also sets _pl_active_scope so the Physical Layout tab's stage-table filter follows the scope automatically. The caption above the buttons still shows the memory-only autosuggest values as a reference — separately labeled "(memory-min)" to avoid confusion with the latency-optimal buttons below. Co-Authored-By: Claude Opus 4.7 --- tests/analytical_visualization/app.py | 66 ++++++++++++++++++++++++--- 1 file changed, 60 insertions(+), 6 deletions(-) diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index c7b4c50..4aca6bc 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -176,14 +176,68 @@ if st.session_state.get("_last_preset") != preset_name: with st.sidebar.expander("Parallelism", expanded=True): st.caption( - f"Auto-suggest: CP={_default_suggestion.cp}, " + f"Auto-suggest (memory-min): CP={_default_suggestion.cp}, " f"TP={_default_suggestion.tp}, PP={_default_suggestion.pp}" ) - if st.button("Apply auto-suggest", width='stretch'): - st.session_state["cp"] = _snap(_default_suggestion.cp, _CP_OPTS) - st.session_state["tp"] = _snap(_default_suggestion.tp, _TP_OPTS) - st.session_state["pp"] = _snap(_default_suggestion.pp, _PP_OPTS) - st.rerun() + # Three latency-optimal apply buttons — pick a scope, load the + # latency-min Pareto config into the sidebar sliders. + from tests.analytical_visualization.auto_explore import ( + run_auto_explore as _run_auto_explore_sb, + ) + _sb_scope_meta = { + "attn": (True, False, "Attention"), + "ffn": (False, True, "FFN/MoE"), + "full": (True, True, "Attn+FFN"), + } + _sb_cols = st.columns(3) + _sb_clicks = {} + for _sb_scope, _sb_col in zip(("attn", "ffn", "full"), _sb_cols): + with _sb_col: + _sb_clicks[_sb_scope] = st.button( + _sb_scope_meta[_sb_scope][2], width='stretch', + key=f"_sb_apply_{_sb_scope}", + help=f"Apply latency-optimal parallelism for " + f"{_sb_scope_meta[_sb_scope][2]} scope.", + ) + for _sb_scope, _sb_click in _sb_clicks.items(): + if not _sb_click: + continue + _sb_ia, _sb_ff, _sb_lbl = _sb_scope_meta[_sb_scope] + with st.spinner(f"Finding latency-optimal parallelism ({_sb_lbl})..."): + _sb_res = _run_auto_explore_sb( + model, _default_machine, s_kv=s_kv, mode=mode, + include_attention=_sb_ia, include_ffn=_sb_ff, + ) + if not _sb_res.pareto_scores: + st.error( + f"No feasible parallelism for {_sb_lbl}. Try raising " + "`pe_hbm_gb` or reducing `s_kv`." + ) + else: + _sb_best = min(_sb_res.pareto_scores, + key=lambda s: s.total_latency_ns) + st.session_state["cp"] = _snap(_sb_best.cp, _CP_OPTS) + st.session_state["tp"] = _snap(_sb_best.tp, _TP_OPTS) + st.session_state["pp"] = _snap(_sb_best.pp, _PP_OPTS) + st.session_state["dp"] = _snap(_sb_best.dp, _DP_OPTS) + st.session_state["tp_placement"] = _sb_best.tp_placement + st.session_state["cp_placement"] = _sb_best.cp_placement + st.session_state["cp_ring_variant"] = _sb_best.cp_ring_variant + st.session_state["kv_mode"] = _sb_best.kv_shard_mode + _sb_tp2 = _snap(_sb_best.tp, _TP_OPTS) + _sb_cp2 = _snap(_sb_best.cp, _CP_OPTS) + _sb_dp2 = _snap(_sb_best.dp, _DP_OPTS) + _sb_ffn_map = { + "TP": f"TP only (div={max(1, _sb_tp2)})", + "TP+CP": f"TP*CP (div={max(1, _sb_tp2 * _sb_cp2)})", + "TP+CP+DP": f"TP*CP*DP (div={max(1, _sb_tp2 * _sb_cp2 * _sb_dp2)})", + } + st.session_state["ffn_scope_label"] = \ + _sb_ffn_map[_sb_best.ffn_shard_scope] + # Track for the Physical Layout scope filter, so the tables + # below also filter to this scope. + st.session_state["_pl_active_scope"] = _sb_scope + st.rerun() # Sharding dims — arrange in 2x3 grid to save vertical space p_row1 = st.columns(3)