diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 1c859c1..a8236b3 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, tab_auto = st.tabs([ +tab_layout, tab_memory, tab_stages, tab_compare, tab_auto, tab_hw = st.tabs([ "Physical layout", "Memory breakdown", "Per-stage latency", - "Save & compare", "Auto Explore", + "Save & compare", "Auto Explore", "Auto Hardware", ]) @@ -1460,3 +1460,226 @@ with tab_auto: "configurations and rank them on the Pareto frontier. " "Takes ~5–10 s for the full search." ) + + +# ── TAB 6: Auto Hardware ───────────────────────────────────────── +with tab_hw: + from tests.analytical_visualization.auto_hardware import ( + _HW_KNOB_DEFAULTS, joint_explore, + ) + + st.markdown( + "Sweep hardware knobs (PE HBM, HBM BW, TFLOPs, PE↔PE / D2D / C2C " + "interconnect BW) **together** with parallelism. For each hardware " + "candidate, we pick the latency-minimum parallelism that fits, then " + "Pareto-rank the joint (latency ↓, hardware cost ↓) space. Also " + "computes per-knob sensitivity: which HW investment gives the biggest " + "speedup?" + ) + + hx_l, hx_m, hx_r = st.columns([2, 2, 1]) + with hx_l: + st.markdown( + f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}**" + ) + with hx_m: + _depth = st.radio( + "Sweep depth", + options=["two_stage", "balanced", "coarse"], + index=1, key="_hw_depth", horizontal=True, + help=("two_stage: 1 HW candidate (defaults) + autosuggest " + "parallelism, ~1s. Great for sensitivity alone.\n" + "balanced (default): 64 HW × ~2k parallelism, ~10-20s.\n" + "coarse: 729 HW × ~2k parallelism, ~2-5 min."), + ) + with hx_r: + _run_hw = st.button("Run joint sweep", type="primary", + width='stretch', key="_hw_run_btn") + + if _run_hw or st.session_state.get("_hw_result") is not None: + if _run_hw: + with st.spinner(f"Sweeping HW × parallelism ({_depth})..."): + _hw_res = joint_explore(model, s_kv=s_kv, + mode=mode, depth=_depth) + st.session_state["_hw_result"] = _hw_res + st.session_state["_hw_ctx"] = ( + model.name, s_kv, mode, _depth, + ) + _hw_res = st.session_state["_hw_result"] + + _cached_ctx = st.session_state.get("_hw_ctx") + _cur_ctx = (model.name, s_kv, mode, _depth) + if _cached_ctx != _cur_ctx: + st.warning( + "Sweep result is from a different model/workload/depth. " + "Click **Run joint sweep** to refresh." + ) + + _hm1, _hm2, _hm3, _hm4 = st.columns(4) + _hm1.metric("HW candidates", _hw_res.total_hw) + _hm2.metric("Feasible joint", _hw_res.total_joint) + _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) + _hm4.metric("Best latency", f"{_best.latency_ms:.2f} ms") + + if not _hw_res.pareto_scores: + st.error( + "No feasible joint configurations — every HW+parallelism " + "combo exceeds per-PE HBM at this workload." + ) + else: + # ── Panel 1: latency vs cost Pareto ──────────────────────── + import matplotlib.pyplot as _plt + _fig1, _ax1 = _plt.subplots(figsize=(8, 4)) + _feas = _hw_res.all_scores + _ax1.scatter( + [s.cost_score for s in _feas], + [s.latency_ms for s in _feas], + c="lightgrey", s=15, alpha=0.5, label="feasible", + ) + _pareto_by_cost = sorted(_hw_res.pareto_scores, + key=lambda s: s.cost_score) + _sc = _ax1.scatter( + [s.cost_score for s in _pareto_by_cost], + [s.latency_ms for s in _pareto_by_cost], + c=[s.parallelism.pes_used for s in _pareto_by_cost], + cmap="viridis", s=90, edgecolors="black", linewidths=1.0, + label="Pareto", + ) + _ax1.plot( + [s.cost_score for s in _pareto_by_cost], + [s.latency_ms for s in _pareto_by_cost], + "k--", linewidth=1.0, alpha=0.6, + ) + _ax1.set_xlabel("Hardware cost proxy (unitless — 6.0 = all defaults)") + _ax1.set_ylabel("Latency (ms)") + _ax1.set_yscale("log") + _ax1.set_title("Optimal HW+parallelism: latency vs cost") + _ax1.grid(alpha=0.3) + _ax1.legend(loc="upper right", fontsize=9) + _fig1.colorbar(_sc, ax=_ax1, label="PEs") + _fig1.tight_layout() + st.pyplot(_fig1, width='stretch') + _plt.close(_fig1) + + # ── Panel 2: per-knob sensitivity ────────────────────────── + st.markdown( + "**Per-knob sensitivity** (doubling each HW knob from the " + "latency-optimal baseline). Higher bar = more valuable " + "investment for this workload." + ) + _fig2, _ax2 = _plt.subplots(figsize=(8, 3.5)) + _knobs = [r.knob for r in _hw_res.sensitivity] + _speedups = [r.rel_speedup * 100 for r in _hw_res.sensitivity] + _colors = ["#2ecc71" if v > 0 else "#95a5a6" for v in _speedups] + _ax2.barh(_knobs, _speedups, color=_colors, + edgecolor="black", linewidth=0.5) + _ax2.invert_yaxis() # biggest speedup on top + _ax2.set_xlabel("Speedup when knob is doubled (%)") + _ax2.set_title("Sensitivity — where to invest next") + _ax2.grid(alpha=0.3, axis="x") + for i, r in enumerate(_hw_res.sensitivity): + _ax2.text( + r.rel_speedup * 100 + 0.5, i, + f"{r.baseline_value:g} → {r.doubled_value:g}", + va="center", fontsize=8, + ) + _fig2.tight_layout() + st.pyplot(_fig2, width='stretch') + _plt.close(_fig2) + + # ── Panel 3: Pareto table (HW spec + parallelism) ────────── + st.markdown("**Pareto joint configurations** (sorted by latency)") + _pareto_by_lat = sorted(_hw_res.pareto_scores, + key=lambda s: s.total_latency_ns) + _rows = [] + for i, s in enumerate(_pareto_by_lat): + p = s.parallelism + h = s.hardware + _rows.append({ + "#": i, + "lat (ms)": round(s.latency_ms, 3), + "cost": round(s.cost_score, 2), + "PEs": p.pes_used, + "CP": p.cp, "TP": p.tp, "PP": p.pp, "DP": p.dp, + "kv": p.kv_shard_mode, + "PE HBM": f"{h.pe_hbm_gb:g} GB", + "HBM BW": f"{h.bw_hbm_gbs:g} GB/s", + "TFLOPs": f"{h.peak_tflops_f16:g}", + "PE↔PE": f"{h.bw_intra_gbs:g}", + "D2D": f"{h.bw_inter_gbs:g}", + "C2C": f"{h.bw_intersip_gbs:g}", + }) + _df_hw = pd.DataFrame(_rows) + st.dataframe(_df_hw, width='stretch', hide_index=True) + + # ── Load hardware into sidebar ───────────────────────────── + st.markdown( + "**Load a Pareto HW config into the main sidebar** " + "(overwrites the Hardware→Per-PE and Hardware→Interconnect " + "sliders; also loads the paired parallelism)." + ) + _lch1, _lch2 = st.columns([1, 3]) + with _lch1: + _pick_hw = st.number_input( + "Row #", min_value=0, + max_value=len(_pareto_by_lat) - 1, + value=0, step=1, key="_hw_pick_row", + ) + with _lch2: + if st.button("Load into sidebar", type="secondary", + key="_hw_load_btn"): + _s = _pareto_by_lat[int(_pick_hw)] + _h = _s.hardware + _p = _s.parallelism + # HW knobs (only load if the value is in the selectbox + # options; otherwise leave as-is). + _apply_if_in = { + "pe_hbm": ([3.0, 6.0, 12.0, 24.0, 48.0, 96.0], + _h.pe_hbm_gb), + "bw_hbm": ([128.0, 256.0, 512.0, 1024.0, 2048.0], + _h.bw_hbm_gbs), + "tflops": ([2.0, 4.0, 8.0, 16.0, 32.0, 64.0], + _h.peak_tflops_f16), + "bw_intra": ([128.0, 256.0, 512.0, 1024.0, 2048.0, + 4096.0], _h.bw_intra_gbs), + "bw_inter": ([32.0, 64.0, 128.0, 256.0, 512.0, + 900.0], _h.bw_inter_gbs), + "bw_intersip": ([12.5, 25.0, 50.0, 100.0, 200.0, + 400.0], _h.bw_intersip_gbs), + } + for _key, (_opts, _val) in _apply_if_in.items(): + if _val in _opts: + st.session_state[_key] = _val + # Parallelism knobs + st.session_state["cp"] = _p.cp + st.session_state["tp"] = _p.tp + st.session_state["pp"] = _p.pp + st.session_state["dp"] = _p.dp + st.session_state["tp_placement"] = _p.tp_placement + st.session_state["cp_placement"] = _p.cp_placement + st.session_state["cp_ring_variant"] = _p.cp_ring_variant + st.session_state["kv_mode"] = _p.kv_shard_mode + # Ffn_scope_label reconstruction (same as auto_explore tab) + _tp2, _cp2, _dp2 = _p.tp, _p.cp, _p.dp + _ffn_label_map = { + "TP": f"TP only (div={max(1, _tp2)})", + "TP+CP": f"TP*CP (div={max(1, _tp2 * _cp2)})", + "TP+CP+DP": f"TP*CP*DP (div={max(1, _tp2 * _cp2 * _dp2)})", + } + st.session_state["ffn_scope_label"] = \ + _ffn_label_map[_p.ffn_shard_scope] + st.success( + f"Loaded row {_pick_hw}: latency {_s.latency_ms:.2f} " + f"ms, cost {_s.cost_score:.2f}, " + f"{_p.pes_used} PEs. Flip to another tab to see the " + f"full breakdown." + ) + st.rerun() + else: + st.info( + "Click **Run joint sweep** to explore hardware and parallelism " + "co-design space for this model + workload." + )