diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 7ae9a32..43ade29 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -237,8 +237,12 @@ with st.sidebar.expander("Parallelism", expanded=True): "`pe_hbm_gb` or reducing `s_kv`." ) else: + # Compound key: min latency, then fewer PEs, then lower HBM %. + # Prefer smaller-footprint configs when they tie on latency. _sb_best = min(_sb_res.pareto_scores, - key=lambda s: s.total_latency_ns) + key=lambda s: (s.total_latency_ns, + s.pes_used, + s.hbm_utilization)) st.session_state["cp"] = _snap(_sb_best.cp, _CP_OPTS) st.session_state["tp"] = _snap(_sb_best.tp, _TP_OPTS) st.session_state["pp"] = _snap(_sb_best.pp, _PP_OPTS) @@ -580,8 +584,11 @@ with tab_layout: "`s_kv`." ) else: + # Compound tiebreaker: min latency, then fewer PEs, then lower HBM %. _pl_best = min(_pl_res.pareto_scores, - key=lambda s: s.total_latency_ns) + key=lambda s: (s.total_latency_ns, + s.pes_used, + s.hbm_utilization)) # Load the winning config into the sidebar's session state. st.session_state["cp"] = _pl_best.cp st.session_state["tp"] = _pl_best.tp @@ -1587,7 +1594,10 @@ def _render_auto_explore_tab(): _m2.metric("Feasible", f"{_res.total_feasible:,}") _m3.metric("Pareto configs", len(_res.pareto_scores)) if _res.pareto_scores: - _best = min(_res.pareto_scores, key=lambda s: s.total_latency_ns) + _best = min(_res.pareto_scores, + key=lambda s: (s.total_latency_ns, + s.pes_used, + s.hbm_utilization)) _m4.metric("Best latency", f"{_best.latency_ms:.2f} ms") if not _res.pareto_scores: @@ -1884,7 +1894,9 @@ def _render_auto_hardware_tab(): _hm3.metric("Pareto configs", len(_hw_res.pareto_scores)) if _hw_res.pareto_scores: _best = min(_hw_res.pareto_scores, - key=lambda s: s.total_latency_ns) + key=lambda s: (s.total_latency_ns, + s.parallelism.pes_used, + s.parallelism.hbm_utilization)) _hm4.metric("Best latency", f"{_best.latency_ms:.2f} ms") if not _hw_res.pareto_scores: diff --git a/tests/analytical_visualization/auto_hardware.py b/tests/analytical_visualization/auto_hardware.py index 219e36c..793edf3 100644 --- a/tests/analytical_visualization/auto_hardware.py +++ b/tests/analytical_visualization/auto_hardware.py @@ -245,14 +245,19 @@ def _best_parallelism_for_hw( matches the caller's attention/FFN/full choice. """ best: ConfigScore | None = None + best_key: tuple | None = None for topo in _iter_reduced_parallelism(model, s_kv, mode): cfg = FullConfig(model=model, topo=topo, machine=machine) s = score_config(cfg, include_attention=include_attention, include_ffn=include_ffn) if not (s.fits_memory and s.placement_valid): continue - if best is None or s.total_latency_ns < best.total_latency_ns: + # Compound tiebreaker: prefer fewer PEs and lower HBM % when latency + # ties across variant knobs (cp_ring, placement, kv_shard_mode, etc.). + _k = (s.total_latency_ns, s.pes_used, s.hbm_utilization) + if best is None or _k < best_key: best = s + best_key = _k return best