diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 5a8d1aa..0285b02 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -172,12 +172,18 @@ if st.session_state.get("_last_preset") != preset_name: st.session_state["pp"] = _snap(_default_suggestion.pp, _PP_OPTS) st.session_state["dp"] = 1 st.session_state["ep"] = 1 + # Apply the picked cp_placement so the sidebar reflects the packed + # config (autosuggest may have chosen "pe" packing to minimize cubes). + st.session_state["cp_placement"] = _default_suggestion.cp_placement st.session_state["_last_preset"] = preset_name with st.sidebar.expander("Parallelism", expanded=True): st.caption( f"Auto-suggest (memory-min): CP={_default_suggestion.cp}, " - f"TP={_default_suggestion.tp}, PP={_default_suggestion.pp}" + f"TP={_default_suggestion.tp}, PP={_default_suggestion.pp} " + f"→ **{_default_suggestion.cubes_used} cube(s)**, " + f"{_default_suggestion.pes_used} PEs " + f"(cp_placement={_default_suggestion.cp_placement})" ) # Three latency-optimal apply buttons — pick a scope, load the # latency-min Pareto config into the sidebar sliders. diff --git a/tests/analytical_visualization/autosuggest.py b/tests/analytical_visualization/autosuggest.py index 7d8f7c4..29e04ca 100644 --- a/tests/analytical_visualization/autosuggest.py +++ b/tests/analytical_visualization/autosuggest.py @@ -32,6 +32,8 @@ class Suggestion: fits: bool pes_used: int sips_used: int + cubes_used: int = 0 # picked to minimize this (see _score_candidate) + cp_placement: str = "cube" # "pe" packs CP into intra-cube PEs when it fits reason: str = "" @@ -59,7 +61,25 @@ def _score_candidate(cp: int, tp: int, pp: int, model: ModelConfig, machine: MachineParams, s_kv: int, mode: str, slack_frac: float = 0.10) -> Suggestion: - topo = TopologyConfig(cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode) + # 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 + # valid when CP·TP fits within a single cube's PE count. + best_topo: TopologyConfig | None = None + best_placement = "cube" + _pe_per_cube_hw = TopologyConfig().pes_per_cube_hw # instance-invariant HW const + _placements_to_try = ["cube"] + if cp * tp <= _pe_per_cube_hw: + _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_placement=_place, + ) + if best_topo is None or _t.cubes_used < best_topo.cubes_used: + best_topo = _t + best_placement = _place + topo = best_topo cfg = FullConfig(model=model, topo=topo, machine=machine) mem = compute_memory(cfg) @@ -80,6 +100,8 @@ def _score_candidate(cp: int, tp: int, pp: int, fits=fits, pes_used=topo.total_pes, sips_used=topo.sips_used, + cubes_used=topo.cubes_used, + cp_placement=best_placement, reason=reason, ) @@ -87,20 +109,27 @@ 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: - """Return the smallest-PE-count (CP, TP, PP) that fits. + """Return the smallest deployment (fewer cubes, then fewer PEs) + that fits. + + Sort key: (cubes_used ↑, pes_used ↑, pp ↑, tp ↑, cp ↑). + Fewer cubes wins first because a cube is the physical die-level + hardware unit; PE count is the tiebreaker. Preferring fewer PP then + TP then CP keeps the earlier historical bias (avoid pipeline + bubbles / head-dim splits) among equal-cube-and-PE ties. If no candidate fits, returns the best-effort one (highest slack, even if negative) with fits=False. """ - best_fit: Suggestion | None = None - best_effort: Suggestion | None = None - + scored: list[Suggestion] = [] for cp, tp, pp in _iter_candidates(model): - s = _score_candidate(cp, tp, pp, model, machine, s_kv, mode, slack_frac) - if s.fits and best_fit is None: - best_fit = s - break # candidates are pre-sorted by PE count - if best_effort is None or s.slack_gb > best_effort.slack_gb: - best_effort = s + scored.append( + _score_candidate(cp, tp, pp, model, machine, s_kv, mode, slack_frac) + ) + scored.sort(key=lambda s: (s.cubes_used, s.pes_used, s.pp, s.tp, s.cp)) - return best_fit or best_effort # type: ignore + for s in scored: + if s.fits: + return s + # No fit — return the best-effort (largest slack, closest to fitting). + return max(scored, key=lambda s: s.slack_gb)