analytical-viz: add sidebar Batch B selectbox + wire through all suggesters

New sidebar widget in the Workload expander:
  Batch B (concurrent requests) — selectbox 1/2/4/…/256, default 1.

Wired end-to-end:
  - app.py sidebar → forwards b_batch when constructing the main topo
  - auto_suggest(..., b=1) — passes b to _score_candidate, which builds
    TopologyConfig with b=b
  - run_auto_explore(..., b=1) — sets topo.b = b for every enumerated
    candidate before scoring
  - joint_explore(..., b=1) — forwards b to _best_parallelism_for_hw
    and _best_parallelism_two_stage; both set topo.b = b before scoring

All button handlers (sidebar Apply-scope, Physical Layout scope,
Auto Suggest tab sweep, Auto Hardware tab sweep) now pass b_batch.

Combined with the earlier partial-scaffold commits (memory_layout
scales KV by B; stage_rmsnorm / stage_wq / stage_wkv / stage_kv_append /
_per_hop_qkT_pv scale flops + activation memory by B), changing B in
the sidebar now affects reported latency and per-PE memory footprint
in the visible parts of the pipeline. The remaining FFN + comm-AR
stages still ignore B (they'll be next); their contribution is small
for the memory-bound decode case that matters most, but latency for
FFN-heavy configs at high B will be slightly under-reported until
those are scaled too.

Verified: 24 pytest tests pass at default B=1 (backward compat).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
2026-07-28 22:51:56 -07:00
parent 674e194dc9
commit fd45bd2408
4 changed files with 35 additions and 12 deletions
+18 -7
View File
@@ -161,11 +161,21 @@ with st.sidebar.expander("Workload", expanded=True):
with w2:
mode = st.selectbox("Mode", ["decode", "prefill"],
index=0, key="mode")
b_batch = st.selectbox(
"Batch B (concurrent requests)",
[1, 2, 4, 8, 16, 32, 64, 128, 256],
index=0, key="b_batch",
help=("Number of concurrent requests processed in one decode/prefill "
"step. KV cache scales linearly with B (each request holds its "
"own slice). Compute + AR-comm bytes also scale with B, but "
"weights are shared. Batching is the primary way to raise "
"HBM-bandwidth efficiency for decode."),
)
# Compute auto-suggest for defaults / apply button
_default_machine = MachineParams()
_default_suggestion = auto_suggest(model, _default_machine, s_kv, mode)
_default_suggestion = auto_suggest(model, _default_machine, s_kv, mode, b=b_batch)
def _snap(v: int, opts: tuple) -> int:
return min(opts, key=lambda o: abs(o - v))
@@ -235,7 +245,7 @@ with st.sidebar.expander("Parallelism", expanded=True):
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,
include_attention=_sb_ia, include_ffn=_sb_ff, b=b_batch,
)
if not _sb_res.pareto_scores:
st.error(
@@ -440,7 +450,7 @@ machine = MachineParams(
compute_util=compute_util,
)
topo = TopologyConfig(cp=cp, tp=tp, pp=pp, dp=dp, ep=ep,
topo = TopologyConfig(cp=cp, tp=tp, pp=pp, dp=dp, ep=ep, b=b_batch,
s_kv=s_kv, mode=mode,
kv_shard_mode=kv_shard_mode,
ffn_shard_scope=ffn_shard_scope,
@@ -464,7 +474,7 @@ if preset.note:
# Recompute auto-suggest with the *current* machine (in case user changed HBM etc.)
suggestion = auto_suggest(model, machine, s_kv, mode)
suggestion = auto_suggest(model, machine, s_kv, mode, b=b_batch)
# Compact top summary: auto-suggest + current in one row of tight bullets.
_sug_status = "OK" if suggestion.fits else "NO FIT"
@@ -584,7 +594,7 @@ with tab_layout:
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,
include_attention=_pl_ia, include_ffn=_pl_ff, b=b_batch,
)
if not _pl_res.pareto_scores:
st.error(
@@ -1558,7 +1568,7 @@ def _render_auto_explore_tab():
with st.spinner(f"Sweeping ~28k configs ({_lbl})..."):
_r = run_auto_explore(
model, machine, s_kv=s_kv, mode=mode,
include_attention=_ia, include_ffn=_ff,
include_attention=_ia, include_ffn=_ff, b=b_batch,
)
st.session_state[f"_auto_explore_result_{_scope}"] = _r
st.session_state[f"_auto_explore_ctx_{_scope}"] = (
@@ -1891,7 +1901,8 @@ def _render_auto_hardware_tab():
_ia, _ff, _lbl = _scope_meta_hw[_scope]
with st.spinner(f"Sweeping HW × parallelism ({_depth}, {_lbl})..."):
_r = joint_explore(model, s_kv=s_kv, mode=mode, depth=_depth,
include_attention=_ia, include_ffn=_ff)
include_attention=_ia, include_ffn=_ff,
b=b_batch)
st.session_state[f"_hw_result_{_scope}"] = _r
st.session_state[f"_hw_ctx_{_scope}"] = (
model.name, s_kv, mode, _depth,