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)