diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index d733c7b..1c859c1 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -419,9 +419,9 @@ if _warnings: # ── Tabs ───────────────────────────────────────────────────────── -tab_layout, tab_memory, tab_stages, tab_compare = st.tabs([ +tab_layout, tab_memory, tab_stages, tab_compare, tab_auto = st.tabs([ "Physical layout", "Memory breakdown", "Per-stage latency", - "Save & compare", + "Save & compare", "Auto Explore", ]) @@ -1282,3 +1282,181 @@ with tab_compare: "'-' means the stage is absent for that config (e.g. C1 in " "decode, since the O/m/l all-reduce is folded into S8)." ) + + +# ── TAB 5: Auto Explore ───────────────────────────────────────── +with tab_auto: + from tests.analytical_visualization.auto_explore import run_auto_explore + + st.markdown( + "Sweep every valid parallelism config (CP, TP, PP, DP, kv_shard, " + "ffn_scope, tp_placement, cp_placement, cp_ring) for the currently-" + "selected model + workload and extract the 3D Pareto frontier on " + "(latency ↓, PEs ↓, efficiency ↑). Throughput = 1/latency for a " + "single request so it collapses with latency and is shown for info only." + ) + + ax_l, ax_r = st.columns([3, 1]) + with ax_l: + st.markdown( + f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}** | " + f"per-PE HBM = **{machine.pe_hbm_gb:.1f} GB**" + ) + with ax_r: + _run_sweep = st.button("Run sweep", type="primary", width='stretch') + + if _run_sweep or st.session_state.get("_auto_explore_result") is not None: + if _run_sweep: + with st.spinner("Sweeping ~28k configs..."): + _res = run_auto_explore(model, machine, s_kv=s_kv, mode=mode) + st.session_state["_auto_explore_result"] = _res + st.session_state["_auto_explore_ctx"] = ( + model.name, s_kv, mode, machine.pe_hbm_gb, + ) + _res = st.session_state["_auto_explore_result"] + + # Warn if the cached result is stale relative to the current config. + _cached_ctx = st.session_state.get("_auto_explore_ctx") + _cur_ctx = (model.name, s_kv, mode, machine.pe_hbm_gb) + if _cached_ctx != _cur_ctx: + st.warning( + "Sweep result is from a different model/workload. " + "Click **Run sweep** to refresh." + ) + + _m1, _m2, _m3, _m4 = st.columns(4) + _m1.metric("Enumerated", f"{_res.total_enumerated:,}") + _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) + _m4.metric("Best latency", f"{_best.latency_ms:.2f} ms") + + if not _res.pareto_scores: + st.error( + "No feasible configs — every parallelism choice exceeds " + "per-PE HBM. Try raising `pe_hbm_gb` in the sidebar or " + "reducing `s_kv`." + ) + else: + # ── Pareto scatter: latency vs PEs, coloured by efficiency ─ + import matplotlib.pyplot as _plt + _fig, _axs = _plt.subplots(1, 2, figsize=(10, 3.5)) + + # All feasible = grey points; Pareto = coloured by efficiency. + _feas = [s for s in _res.all_scores + if s.fits_memory and s.placement_valid] + _axs[0].scatter( + [s.pes_used for s in _feas], + [s.latency_ms for s in _feas], + c="lightgrey", s=8, alpha=0.4, label="feasible", + ) + _pareto_sorted = sorted(_res.pareto_scores, + key=lambda s: s.pes_used) + _sc = _axs[0].scatter( + [s.pes_used for s in _pareto_sorted], + [s.latency_ms for s in _pareto_sorted], + c=[s.efficiency_score for s in _pareto_sorted], + cmap="viridis", s=60, edgecolors="black", linewidths=0.8, + label="Pareto", + ) + # Connect Pareto in PE order to show the trade-off curve. + _axs[0].plot( + [s.pes_used for s in _pareto_sorted], + [s.latency_ms for s in _pareto_sorted], + "k--", linewidth=0.8, alpha=0.5, + ) + _axs[0].set_xlabel("PEs used"); _axs[0].set_ylabel("Latency (ms)") + _axs[0].set_xscale("log", base=2) + _axs[0].set_yscale("log") + _axs[0].set_title("Pareto: latency vs PEs") + _axs[0].grid(alpha=0.3) + _axs[0].legend(loc="upper right", fontsize=8) + _fig.colorbar(_sc, ax=_axs[0], label="efficiency") + + # HBM utilization vs latency for Pareto — where does the config land? + _pareto_by_lat = sorted(_res.pareto_scores, + key=lambda s: s.total_latency_ns) + _axs[1].scatter( + [s.hbm_utilization * 100 for s in _pareto_by_lat], + [s.latency_ms for s in _pareto_by_lat], + c=[s.pes_used for s in _pareto_by_lat], + cmap="plasma", s=60, edgecolors="black", linewidths=0.8, + ) + _axs[1].set_xlabel("HBM utilization (%)") + _axs[1].set_ylabel("Latency (ms)") + _axs[1].set_yscale("log") + _axs[1].set_title("HBM usage vs latency (Pareto)") + _axs[1].grid(alpha=0.3) + _sc2 = _axs[1].collections[0] + _fig.colorbar(_sc2, ax=_axs[1], label="PEs") + + _fig.tight_layout() + st.pyplot(_fig, width='stretch') + _plt.close(_fig) + + # ── Pareto table + Load button ───────────────────────────── + st.markdown("**Pareto configurations** (sorted by latency)") + _rows = [] + for i, s in enumerate(_pareto_by_lat): + _rows.append({ + "#": i, + "CP": s.cp, "TP": s.tp, "PP": s.pp, "DP": s.dp, + "kv": s.kv_shard_mode, + "ffn": s.ffn_shard_scope, + "tp_place": s.tp_placement, + "cp_place": s.cp_placement, + "cp_ring": s.cp_ring_variant, + "lat (ms)": round(s.latency_ms, 3), + "eff": round(s.efficiency_score, 4), + "PEs": s.pes_used, + "SIPs": s.sips_used, + "HBM %": round(s.hbm_utilization * 100, 1), + }) + _df = pd.DataFrame(_rows) + st.dataframe(_df, width='stretch', hide_index=True) + + # ── Load-into-sidebar ───────────────────────────────────── + st.markdown("**Load a Pareto config into the main sidebar sliders**") + _lc1, _lc2 = st.columns([1, 3]) + with _lc1: + _pick = st.number_input( + "Row #", min_value=0, max_value=len(_pareto_by_lat) - 1, + value=0, step=1, key="_auto_pick_row", + ) + with _lc2: + if st.button("Load into sidebar", type="secondary"): + _s = _pareto_by_lat[int(_pick)] + # Set the sidebar's session_state keys directly. The + # sidebar reads these on the next rerun. + st.session_state["cp"] = _s.cp + st.session_state["tp"] = _s.tp + st.session_state["pp"] = _s.pp + st.session_state["dp"] = _s.dp + st.session_state["tp_placement"] = _s.tp_placement + st.session_state["cp_placement"] = _s.cp_placement + st.session_state["cp_ring_variant"] = _s.cp_ring_variant + st.session_state["kv_mode"] = _s.kv_shard_mode + # ffn_scope label is dynamic: "TP+CP (div=…)" etc. + # Reconstruct with the target cp/tp/dp values. + _tp, _cp, _dp = _s.tp, _s.cp, _s.dp + _ffn_label_map = { + "TP": f"TP only (div={max(1, _tp)})", + "TP+CP": f"TP*CP (div={max(1, _tp * _cp)})", + "TP+CP+DP": f"TP*CP*DP (div={max(1, _tp * _cp * _dp)})", + } + st.session_state["ffn_scope_label"] = \ + _ffn_label_map[_s.ffn_shard_scope] + st.success( + f"Loaded row {_pick}: CP={_s.cp} TP={_s.tp} PP={_s.pp} " + f"DP={_s.dp} → latency {_s.latency_ms:.2f} ms, " + f"{_s.pes_used} PEs. Flip to another tab to see the " + f"full breakdown." + ) + st.rerun() + else: + st.info( + "Click **Run sweep** to enumerate all valid parallelism " + "configurations and rank them on the Pareto frontier. " + "Takes ~5–10 s for the full search." + )