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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user