diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index f08b983..d262039 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -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, diff --git a/tests/analytical_visualization/auto_explore.py b/tests/analytical_visualization/auto_explore.py index 31bbc3e..bf4b26f 100644 --- a/tests/analytical_visualization/auto_explore.py +++ b/tests/analytical_visualization/auto_explore.py @@ -419,6 +419,7 @@ def run_auto_explore( mode: str = "decode", include_attention: bool = True, include_ffn: bool = True, + b: int = 1, ) -> AutoExploreResult: """Enumerate all configs, score each, extract Pareto frontier. @@ -433,6 +434,8 @@ def run_auto_explore( all_scores: list[ConfigScore] = [] total_enumerated = 0 for topo in enumerate_configs(model, s_kv, mode): + # Inject the caller's batch B into every candidate. + topo.b = b total_enumerated += 1 cfg = FullConfig(model=model, topo=topo, machine=machine) all_scores.append(score_config( diff --git a/tests/analytical_visualization/auto_hardware.py b/tests/analytical_visualization/auto_hardware.py index 793edf3..b98f2af 100644 --- a/tests/analytical_visualization/auto_hardware.py +++ b/tests/analytical_visualization/auto_hardware.py @@ -237,6 +237,7 @@ def _best_parallelism_for_hw( model: ModelConfig, machine: MachineParams, s_kv: int, mode: str, include_attention: bool = True, include_ffn: bool = True, + b: int = 1, ) -> ConfigScore | None: """Return the latency-minimum feasible parallelism for this HW. @@ -247,6 +248,7 @@ def _best_parallelism_for_hw( best: ConfigScore | None = None best_key: tuple | None = None for topo in _iter_reduced_parallelism(model, s_kv, mode): + topo.b = b cfg = FullConfig(model=model, topo=topo, machine=machine) s = score_config(cfg, include_attention=include_attention, include_ffn=include_ffn) @@ -265,13 +267,14 @@ def _best_parallelism_two_stage( model: ModelConfig, machine: MachineParams, s_kv: int, mode: str, include_attention: bool = True, include_ffn: bool = True, + b: int = 1, ) -> ConfigScore | None: """Fast fallback: use autosuggest's memory-min (CP,TP,PP) then score it.""" sug = auto_suggest(model, machine, s_kv, mode) if not sug.fits: return None topo = TopologyConfig( - cp=sug.cp, tp=sug.tp, pp=sug.pp, dp=1, + cp=sug.cp, tp=sug.tp, pp=sug.pp, dp=1, b=b, s_kv=s_kv, mode=mode, kv_shard_mode="split", ffn_shard_scope="TP+CP", @@ -366,6 +369,7 @@ def joint_explore( depth: str = "balanced", include_attention: bool = True, include_ffn: bool = True, + b: int = 1, ) -> JointExploreResult: """Sweep HW candidates × parallelism, return joint Pareto + sensitivity. @@ -383,11 +387,13 @@ def joint_explore( par = _best_parallelism_two_stage( model, machine, s_kv, mode, include_attention=include_attention, include_ffn=include_ffn, + b=b, ) else: par = _best_parallelism_for_hw( model, machine, s_kv, mode, include_attention=include_attention, include_ffn=include_ffn, + b=b, ) if par is None: continue diff --git a/tests/analytical_visualization/autosuggest.py b/tests/analytical_visualization/autosuggest.py index 29e04ca..d6340ee 100644 --- a/tests/analytical_visualization/autosuggest.py +++ b/tests/analytical_visualization/autosuggest.py @@ -60,7 +60,8 @@ def _iter_candidates(model: ModelConfig) -> Iterable[tuple[int, int, int]]: def _score_candidate(cp: int, tp: int, pp: int, model: ModelConfig, machine: MachineParams, s_kv: int, mode: str, - slack_frac: float = 0.10) -> Suggestion: + slack_frac: float = 0.10, + b: int = 1) -> Suggestion: # For each triple, try both cp_placement options and keep the one # with fewer cubes (breaks the historical "CP always across cubes" # default when a smaller pack is possible). The pe placement is only @@ -73,7 +74,7 @@ def _score_candidate(cp: int, tp: int, pp: int, _placements_to_try.append("pe") for _place in _placements_to_try: _t = TopologyConfig( - cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode, + cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode, b=b, cp_placement=_place, ) if best_topo is None or _t.cubes_used < best_topo.cubes_used: @@ -108,7 +109,8 @@ def _score_candidate(cp: int, tp: int, pp: int, def auto_suggest(model: ModelConfig, machine: MachineParams, s_kv: int, mode: str = "decode", - slack_frac: float = 0.10) -> Suggestion: + slack_frac: float = 0.10, + b: int = 1) -> Suggestion: """Return the smallest deployment (fewer cubes, then fewer PEs) that fits. @@ -124,7 +126,8 @@ def auto_suggest(model: ModelConfig, machine: MachineParams, scored: list[Suggestion] = [] for cp, tp, pp in _iter_candidates(model): scored.append( - _score_candidate(cp, tp, pp, model, machine, s_kv, mode, slack_frac) + _score_candidate(cp, tp, pp, model, machine, s_kv, mode, + slack_frac, b=b) ) scored.sort(key=lambda s: (s.cubes_used, s.pes_used, s.pp, s.tp, s.cp))