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:
2026-07-28 23:20:32 -07:00
parent 50aab8d591
commit 47c4f40f4d
2 changed files with 109 additions and 67 deletions
+13 -12
View File
@@ -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..."