analytical-viz: 3 sweep buttons at top of Physical Layout tab
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 <noreply@anthropic.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user