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)