analytical-viz: batch B in every formula string + rename 'latency-optimal' → 'memory-optimal'
Two related fixes. 1) Formula strings show B alongside every other symbol. Previously the numeric FLOPs/mem_bytes/comm_bytes were scaled by B, but the printed formula strings still read like B=1 amounts. That made the per-stage table confusing at high B: numbers moved but formulas didn't. Updated every attention + FFN stage's `formula` / `flops_formula` / `mem_formula` / `comm_formula` to include B explicitly. Weight-bytes lines now say "(shared across batch)" / "(weight, B-invariant)" so it's obvious what does and doesn't scale. Stages touched: S1..S10 (attention), C1/C2/C3 (attn comm), F1..F5 (FFN), CF1 (FFN AR). 2) Renamed the button/spinner/help copy from "latency-optimal" to "memory-optimal" to match the smallest-fit pick semantics I put in place earlier. The buttons already picked smallest-fit; the labels just still said "latency-optimal" from the previous iteration. Left the Auto Hardware sensitivity chart's "latency-optimal baseline" label alone — that panel is a HW co-design view where "how fast can each HW go?" is the intended question. Verified: 24 pytest tests still pass. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -218,8 +218,8 @@ with st.sidebar.expander("Parallelism", expanded=True):
|
||||
st.session_state["cp_placement"] = _default_suggestion.cp_placement
|
||||
st.session_state["_pl_active_scope"] = "full"
|
||||
st.rerun()
|
||||
# Three latency-optimal apply buttons — pick a scope, load the
|
||||
# latency-min Pareto config into the sidebar sliders.
|
||||
# Three memory-optimal apply buttons — pick a scope, load the
|
||||
# smallest-fit Pareto config into the sidebar sliders.
|
||||
from tests.analytical_visualization.auto_explore import (
|
||||
run_auto_explore as _run_auto_explore_sb,
|
||||
)
|
||||
@@ -235,14 +235,14 @@ with st.sidebar.expander("Parallelism", expanded=True):
|
||||
_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 "
|
||||
help=f"Apply memory-optimal (smallest-fit) 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})..."):
|
||||
with st.spinner(f"Finding memory-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, b=b_batch,
|
||||
@@ -553,17 +553,18 @@ tab_layout, tab_auto, 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.
|
||||
# Quick shortcut: three buttons that pick the memory-optimal
|
||||
# (smallest-fit) 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."
|
||||
"**Quick auto-suggest layout** — pick a scope; the memory-optimal "
|
||||
"(smallest-fit) 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:
|
||||
@@ -591,7 +592,7 @@ with tab_layout:
|
||||
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})..."):
|
||||
with st.spinner(f"Finding memory-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, b=b_batch,
|
||||
@@ -629,7 +630,7 @@ with tab_layout:
|
||||
# stage tables (attn/ffn/full).
|
||||
st.session_state["_pl_active_scope"] = _pl_scope
|
||||
st.success(
|
||||
f"Loaded latency-optimal {_pl_lbl} config: CP={_pl_best.cp} "
|
||||
f"Loaded memory-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..."
|
||||
|
||||
Reference in New Issue
Block a user