diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 57e4a71..7146600 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -1324,51 +1324,56 @@ def _render_auto_explore_tab(): f"per-PE HBM = **{machine.pe_hbm_gb:.1f} GB**" ) - _b1, _b2, _b3 = st.columns([1, 1, 2]) + _b1, _b2, _b3, _b4 = st.columns([1, 1, 1, 1]) with _b1: _run_attn = st.button("Run sweep — Attention", type="primary", width='stretch', key="_auto_run_attn") with _b2: + _run_ffn = st.button("Run sweep — FFN/MoE", + type="primary", width='stretch', + key="_auto_run_ffn") + with _b3: _run_full = st.button("Run sweep — Attn + FFN/MoE", type="primary", width='stretch', key="_auto_run_full") - if _run_attn: - with st.spinner("Sweeping ~28k configs (Attention only)..."): + _scope_meta = { + "attn": (True, False, "Attention only"), + "ffn": (False, True, "FFN / MoE only"), + "full": (True, True, "Attn + FFN/MoE"), + } + _button_click = {"attn": _run_attn, "ffn": _run_ffn, "full": _run_full} + + for _scope, _clicked in _button_click.items(): + if not _clicked: + continue + _ia, _ff, _lbl = _scope_meta[_scope] + with st.spinner(f"Sweeping ~28k configs ({_lbl})..."): _r = run_auto_explore( - model, machine, s_kv=s_kv, mode=mode, include_ffn=False, + model, machine, s_kv=s_kv, mode=mode, + include_attention=_ia, include_ffn=_ff, ) - st.session_state["_auto_explore_result_attn"] = _r - st.session_state["_auto_explore_ctx_attn"] = ( + st.session_state[f"_auto_explore_result_{_scope}"] = _r + st.session_state[f"_auto_explore_ctx_{_scope}"] = ( model.name, s_kv, mode, machine.pe_hbm_gb, ) - st.session_state["_auto_explore_active"] = "attn" - if _run_full: - with st.spinner("Sweeping ~28k configs (Attn + FFN/MoE)..."): - _r = run_auto_explore( - model, machine, s_kv=s_kv, mode=mode, include_ffn=True, - ) - st.session_state["_auto_explore_result_full"] = _r - st.session_state["_auto_explore_ctx_full"] = ( - model.name, s_kv, mode, machine.pe_hbm_gb, - ) - st.session_state["_auto_explore_active"] = "full" + st.session_state["_auto_explore_active"] = _scope _active = st.session_state.get("_auto_explore_active") if _active is None: st.info( "Click one of the run buttons above to enumerate valid " "parallelism configurations and rank them on the Pareto " - "frontier. Takes ~5–10 s. Both button results are cached so " - "you can flip between the two scopes without re-running." + "frontier. Takes ~5–10 s per scope. All button results are " + "cached so you can flip between scopes without re-running." ) return - # Downstream code expects _suffix, _label, include_ffn, _res, _ctx_key. - include_ffn = (_active == "full") - _suffix = "_full" if include_ffn else "_attn" - _label = "Attn + FFN/MoE" if include_ffn else "Attention only" + # Downstream code expects _suffix, _label, include_attention, include_ffn, + # _res, _ctx_key. + include_attention, include_ffn, _label = _scope_meta[_active] + _suffix = f"_{_active}" _result_key = f"_auto_explore_result{_suffix}" _ctx_key = f"_auto_explore_ctx{_suffix}" _res = st.session_state.get(_result_key) @@ -1387,206 +1392,191 @@ def _render_auto_explore_tab(): f"Click **Run sweep — {_label.split(' ')[0]}** 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") + _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) - - # ── Parallelism sensitivity subplots ────────────────────── - st.markdown( - "**Parallelism sensitivity** — for the *baseline* config " - "picked below, vary each knob one at a time (holding others " - "fixed). Answers 'if I only tweak CP / TP / PP / DP / EP, " - "what happens?' Solid line = fits memory; red X = infeasible." - ) - _sens_c1, _sens_c2 = st.columns([1, 4]) - with _sens_c1: - _sens_row = st.number_input( - "Baseline row #", - min_value=0, max_value=len(_pareto_by_lat) - 1, - value=0, step=1, key=f"_auto_sens_row{_suffix}", - help="Which Pareto row is the sensitivity computed around?", - ) - _baseline_for_sens = _pareto_by_lat[int(_sens_row)] - _sens_rows = compute_parallelism_sensitivity( - _baseline_for_sens, model, machine, s_kv=s_kv, mode=mode, - include_ffn=include_ffn, - ) - - import matplotlib.pyplot as _plt - _fig_sens, _axes_sens = _plt.subplots(1, 5, figsize=(14, 3.2)) - for _ax, _row in zip(_axes_sens, _sens_rows): - _fit_vals, _fit_lats = [], [] - _bad_vals = [] - for _v, _lat_ns, _fit in zip(_row.values, _row.latencies_ns, - _row.fits_flags): - if _lat_ns != _lat_ns: # NaN — skip out-of-bounds - continue - if _fit: - _fit_vals.append(_v) - _fit_lats.append(_lat_ns / 1e6) # ms - else: - _bad_vals.append((_v, _lat_ns / 1e6)) - if _fit_vals: - _ax.plot(_fit_vals, _fit_lats, "o-", color="#2c3e50", - linewidth=1.4, markersize=4) - if _bad_vals: - _ax.scatter( - [v for v, _ in _bad_vals], - [l for _, l in _bad_vals], - marker="x", color="#c0392b", s=40, alpha=0.6, - label="no fit", - ) - # Baseline marker - _ax.axvline(_row.baseline_value, color="#3498db", - linestyle=":", linewidth=1.0) - _ax.set_xscale("log") - _ax.set_yscale("log") - _ax.set_title( - f"{_row.knob.upper()} (baseline={_row.baseline_value})", - fontsize=10, - ) - _ax.set_xlabel("value") - _ax.set_ylabel("lat (ms)") - _ax.grid(alpha=0.3) - if _bad_vals: - _ax.legend(loc="best", fontsize=7) - _fig_sens.tight_layout() - st.pyplot(_fig_sens, width='stretch') - _plt.close(_fig_sens) - - # ── 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=f"_auto_pick_row{_suffix}", - ) - with _lc2: - if st.button("Load into sidebar", type="secondary", - key=f"_auto_load_btn{_suffix}"): - _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." + 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`." ) + return + + # ── Pareto scatter: latency vs PEs, coloured by efficiency ─ + import matplotlib.pyplot as _plt + _fig, _axs = _plt.subplots(1, 2, figsize=(10, 3.5)) + + _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", + ) + _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") + + _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 ──────────────────────────────────────────── + 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) + + # ── Parallelism sensitivity subplots ────────────────────── + st.markdown( + "**Parallelism sensitivity** — for the *baseline* config picked " + "below, vary each knob one at a time (holding others fixed). " + "Solid line = fits memory; red X = infeasible." + ) + _sens_c1, _sens_c2 = st.columns([1, 4]) + with _sens_c1: + _sens_row = st.number_input( + "Baseline row #", + min_value=0, max_value=len(_pareto_by_lat) - 1, + value=0, step=1, key=f"_auto_sens_row{_suffix}", + help="Which Pareto row is the sensitivity computed around?", + ) + _baseline_for_sens = _pareto_by_lat[int(_sens_row)] + _sens_rows = compute_parallelism_sensitivity( + _baseline_for_sens, model, machine, s_kv=s_kv, mode=mode, + include_attention=include_attention, include_ffn=include_ffn, + ) + + _fig_sens, _axes_sens = _plt.subplots(1, 5, figsize=(14, 3.2)) + for _ax, _row in zip(_axes_sens, _sens_rows): + _fit_vals, _fit_lats = [], [] + _bad_vals = [] + for _v, _lat_ns, _fit in zip(_row.values, _row.latencies_ns, + _row.fits_flags): + if _lat_ns != _lat_ns: + continue + if _fit: + _fit_vals.append(_v) + _fit_lats.append(_lat_ns / 1e6) + else: + _bad_vals.append((_v, _lat_ns / 1e6)) + if _fit_vals: + _ax.plot(_fit_vals, _fit_lats, "o-", color="#2c3e50", + linewidth=1.4, markersize=4) + if _bad_vals: + _ax.scatter( + [v for v, _ in _bad_vals], + [l for _, l in _bad_vals], + marker="x", color="#c0392b", s=40, alpha=0.6, + label="no fit", + ) + _ax.axvline(_row.baseline_value, color="#3498db", + linestyle=":", linewidth=1.0) + _ax.set_xscale("log") + _ax.set_yscale("log") + _ax.set_title( + f"{_row.knob.upper()} (baseline={_row.baseline_value})", + fontsize=10, + ) + _ax.set_xlabel("value") + _ax.set_ylabel("lat (ms)") + _ax.grid(alpha=0.3) + if _bad_vals: + _ax.legend(loc="best", fontsize=7) + _fig_sens.tight_layout() + st.pyplot(_fig_sens, width='stretch') + _plt.close(_fig_sens) + + # ── 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=f"_auto_pick_row{_suffix}", + ) + with _lc2: + if st.button("Load into sidebar", type="secondary", + key=f"_auto_load_btn{_suffix}"): + _s = _pareto_by_lat[int(_pick)] + 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 + _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() with tab_auto: @@ -1634,47 +1624,53 @@ def _render_auto_hardware_tab(): "coarse: 729 HW × ~2k parallelism, ~2-5 min."), ) - _bh1, _bh2, _bh3 = st.columns([1, 1, 2]) + _bh1, _bh2, _bh3, _bh4 = st.columns([1, 1, 1, 1]) with _bh1: _run_hw_attn = st.button("Run joint sweep — Attention", type="primary", width='stretch', key="_hw_run_attn") with _bh2: + _run_hw_ffn = st.button("Run joint sweep — FFN/MoE", + type="primary", width='stretch', + key="_hw_run_ffn") + with _bh3: _run_hw_full = st.button("Run joint sweep — Attn + FFN/MoE", type="primary", width='stretch', key="_hw_run_full") - if _run_hw_attn: - with st.spinner(f"Sweeping HW × parallelism ({_depth}, Attention only)..."): - _r = joint_explore(model, s_kv=s_kv, mode=mode, - depth=_depth, include_ffn=False) - st.session_state["_hw_result_attn"] = _r - st.session_state["_hw_ctx_attn"] = ( + _scope_meta_hw = { + "attn": (True, False, "Attention only"), + "ffn": (False, True, "FFN / MoE only"), + "full": (True, True, "Attn + FFN/MoE"), + } + _button_click_hw = { + "attn": _run_hw_attn, "ffn": _run_hw_ffn, "full": _run_hw_full, + } + + for _scope, _clicked in _button_click_hw.items(): + if not _clicked: + continue + _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) + st.session_state[f"_hw_result_{_scope}"] = _r + st.session_state[f"_hw_ctx_{_scope}"] = ( model.name, s_kv, mode, _depth, ) - st.session_state["_hw_active"] = "attn" - if _run_hw_full: - with st.spinner(f"Sweeping HW × parallelism ({_depth}, Attn + FFN/MoE)..."): - _r = joint_explore(model, s_kv=s_kv, mode=mode, - depth=_depth, include_ffn=True) - st.session_state["_hw_result_full"] = _r - st.session_state["_hw_ctx_full"] = ( - model.name, s_kv, mode, _depth, - ) - st.session_state["_hw_active"] = "full" + st.session_state["_hw_active"] = _scope _active = st.session_state.get("_hw_active") if _active is None: st.info( "Click one of the run buttons to explore hardware × parallelism " - "co-design for this model + workload. Both button results are " - "cached so you can flip between the two scopes without re-running." + "co-design for this model + workload. All button results are " + "cached so you can flip between scopes without re-running." ) return - include_ffn = (_active == "full") - _suffix = "_full" if include_ffn else "_attn" - _label = "Attn + FFN/MoE" if include_ffn else "Attention only" + include_attention, include_ffn, _label = _scope_meta_hw[_active] + _suffix = f"_{_active}" _hw_result_key = f"_hw_result{_suffix}" _hw_ctx_key = f"_hw_ctx{_suffix}" _hw_res = st.session_state.get(_hw_result_key) @@ -1692,21 +1688,22 @@ def _render_auto_hardware_tab(): f"Click **Run joint sweep — {_label.split(' ')[0]}** 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") + _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: + if not _hw_res.pareto_scores: + st.error( + "No feasible joint configurations — every HW+parallelism " + "combo exceeds per-PE HBM at this workload." + ) + return + if True: # ── Panel 1: latency vs cost Pareto ──────────────────────── import matplotlib.pyplot as _plt _fig1, _ax1 = _plt.subplots(figsize=(8, 4)) @@ -1855,11 +1852,6 @@ def _render_auto_hardware_tab(): f"full breakdown." ) st.rerun() - else: - st.info( - "Click **Run joint sweep** to explore hardware and parallelism " - "co-design space for this model + workload." - ) with tab_hw: diff --git a/tests/analytical_visualization/auto_explore.py b/tests/analytical_visualization/auto_explore.py index 4b60f4f..31bbc3e 100644 --- a/tests/analytical_visualization/auto_explore.py +++ b/tests/analytical_visualization/auto_explore.py @@ -164,7 +164,11 @@ def enumerate_configs( # ── Scoring ────────────────────────────────────────────────────────── -def _sum_visible_latency(cfg: FullConfig, include_ffn: bool = True) -> float: +def _sum_visible_latency( + cfg: FullConfig, + include_attention: bool = True, + include_ffn: bool = True, +) -> float: """Total single-request latency (seconds) across all model layers. A single request traverses every layer sequentially, whether the layers @@ -177,28 +181,31 @@ def _sum_visible_latency(cfg: FullConfig, include_ffn: bool = True) -> float: throughput under batching. This function is the single-request cost, so we multiply by full model.layers regardless of PP. - ``include_ffn=False`` restricts to attention stages only — useful for - isolating attention-kernel tuning from FFN cost. + Scope selection (both default True — full per-token cost): + - ``include_attention=True, include_ffn=True`` → full transformer + - ``include_attention=True, include_ffn=False`` → attention only + - ``include_attention=False, include_ffn=True`` → FFN / MoE only + - ``include_attention=False, include_ffn=False`` → zero (rejected caller-side) """ - attn = sum(s.visible_s for s in all_stages(cfg)) + attn = sum(s.visible_s for s in all_stages(cfg)) if include_attention else 0.0 ffn = sum(s.visible_s for s in all_ffn_stages(cfg)) if include_ffn else 0.0 per_layer = attn + ffn return per_layer * cfg.model.layers def _efficiency(cfg: FullConfig, latency_s: float, + include_attention: bool = True, include_ffn: bool = True) -> float: """Geo-mean of compute-util and BW-util. Range ~ (0, 1]. - compute_util = achieved_flops / (peak_flops × pes × latency) - bw_util = achieved_bytes / (peak_bw × pes × latency) - ``include_ffn=False`` mirrors the same restriction as - :func:`_sum_visible_latency` — attention stages only. + Scope flags mirror :func:`_sum_visible_latency`. """ if latency_s <= 0: return 0.0 - attn = all_stages(cfg) + attn = all_stages(cfg) if include_attention else [] ffn = all_ffn_stages(cfg) if include_ffn else [] layers = math.ceil(cfg.model.layers / cfg.topo.pp) total_flops = layers * sum(s.flops for s in attn + ffn) @@ -216,24 +223,32 @@ def _efficiency(cfg: FullConfig, latency_s: float, return math.sqrt(compute_util * bw_util) -def score_config(cfg: FullConfig, include_ffn: bool = True) -> ConfigScore: +def score_config(cfg: FullConfig, + include_attention: bool = True, + include_ffn: bool = True) -> ConfigScore: """Compute all 4 objectives + info fields for one config. Feasibility (memory + placement) is stored but does NOT gate scoring — infeasible configs get returned with fits_memory=False so callers can filter or display them. - ``include_ffn=False`` restricts latency + efficiency to attention stages - only. Memory feasibility is unchanged — still checks weights+KV+transient - fit in per-PE HBM since the model still exists physically. + Scope flags (default: full transformer) restrict *latency + efficiency* + to attention only, FFN only, or both. Memory feasibility is unchanged — + still checks weights+KV+transient fit in per-PE HBM since the model + still exists physically regardless of what the caller is scoring. """ mem = compute_memory(cfg) placement_ok = cfg.topo.placement_valid - latency_s = _sum_visible_latency(cfg, include_ffn=include_ffn) + latency_s = _sum_visible_latency( + cfg, include_attention=include_attention, include_ffn=include_ffn, + ) throughput = 1.0 / latency_s if latency_s > 0 else 0.0 - efficiency = (_efficiency(cfg, latency_s, include_ffn=include_ffn) - if latency_s > 0 else 0.0) + efficiency = ( + _efficiency(cfg, latency_s, + include_attention=include_attention, include_ffn=include_ffn) + if latency_s > 0 else 0.0 + ) fits = not mem.over_budget reason = "" @@ -345,6 +360,7 @@ def compute_parallelism_sensitivity( machine: MachineParams, s_kv: int, mode: str, + include_attention: bool = True, include_ffn: bool = True, ) -> list[ParallelismSensitivityRow]: """For each parallelism knob (CP, TP, PP, DP, EP), sweep its values @@ -377,7 +393,9 @@ def compute_parallelism_sensitivity( # Build a variant TopologyConfig with just this one knob changed. trial = replace(baseline_topo, **{knob: v}) cfg = FullConfig(model=model, topo=trial, machine=machine) - score = score_config(cfg, include_ffn=include_ffn) + score = score_config( + cfg, include_attention=include_attention, include_ffn=include_ffn, + ) latencies.append(score.total_latency_ns) fits.append(score.fits_memory and score.placement_valid) rows.append(ParallelismSensitivityRow( @@ -399,6 +417,7 @@ def run_auto_explore( machine: MachineParams, s_kv: int, mode: str = "decode", + include_attention: bool = True, include_ffn: bool = True, ) -> AutoExploreResult: """Enumerate all configs, score each, extract Pareto frontier. @@ -407,15 +426,18 @@ def run_auto_explore( ``pareto_scores`` subset (for the scatter/highlight view). Both are sorted by total_latency_ns ascending. - ``include_ffn=False`` restricts to attention stages only — useful for - isolating attention-kernel tuning independent of FFN cost. + Scope: at least one of ``include_attention`` / ``include_ffn`` must be + True. Setting both False would give zero latency for every config — + the caller should not do that. """ all_scores: list[ConfigScore] = [] total_enumerated = 0 for topo in enumerate_configs(model, s_kv, mode): total_enumerated += 1 cfg = FullConfig(model=model, topo=topo, machine=machine) - all_scores.append(score_config(cfg, include_ffn=include_ffn)) + all_scores.append(score_config( + cfg, include_attention=include_attention, include_ffn=include_ffn, + )) all_scores.sort(key=lambda s: s.total_latency_ns) pareto = pareto_frontier(all_scores) diff --git a/tests/analytical_visualization/auto_hardware.py b/tests/analytical_visualization/auto_hardware.py index d57aa87..219e36c 100644 --- a/tests/analytical_visualization/auto_hardware.py +++ b/tests/analytical_visualization/auto_hardware.py @@ -235,18 +235,20 @@ def _iter_reduced_parallelism( def _best_parallelism_for_hw( model: ModelConfig, machine: MachineParams, - s_kv: int, mode: str, include_ffn: bool = True, + s_kv: int, mode: str, + include_attention: bool = True, include_ffn: bool = True, ) -> ConfigScore | None: """Return the latency-minimum feasible parallelism for this HW. Feasibility: fits memory + placement_valid. Returns None if nothing fits. - ``include_ffn`` is forwarded to score_config so the "latency" ranked on - matches the caller's attention-only vs full-model choice. + Scope flags are forwarded to score_config so the "latency" ranked on + matches the caller's attention/FFN/full choice. """ best: ConfigScore | 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_ffn=include_ffn) + 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: @@ -256,7 +258,8 @@ def _best_parallelism_for_hw( def _best_parallelism_two_stage( model: ModelConfig, machine: MachineParams, - s_kv: int, mode: str, include_ffn: bool = True, + s_kv: int, mode: str, + include_attention: bool = True, include_ffn: bool = True, ) -> ConfigScore | None: """Fast fallback: use autosuggest's memory-min (CP,TP,PP) then score it.""" sug = auto_suggest(model, machine, s_kv, mode) @@ -272,7 +275,8 @@ def _best_parallelism_two_stage( cp_ring_variant="qoml" if mode == "decode" and sug.cp > 1 else "kv", ) cfg = FullConfig(model=model, topo=topo, machine=machine) - s = score_config(cfg, include_ffn=include_ffn) + s = score_config(cfg, include_attention=include_attention, + include_ffn=include_ffn) return s if s.fits_memory and s.placement_valid else None @@ -310,7 +314,7 @@ def compute_sensitivity( baseline_hw: HardwareCandidate, parallelism: ConfigScore, model: ModelConfig, s_kv: int, mode: str, - include_ffn: bool = True, + include_attention: bool = True, include_ffn: bool = True, ) -> list[SensitivityRow]: """For each HW knob, double it (holding others at baseline) and measure the latency change. Same parallelism used throughout so we isolate the @@ -321,7 +325,9 @@ def compute_sensitivity( baseline_cfg = FullConfig( model=model, topo=baseline_topo, machine=baseline_machine, ) - baseline_score = score_config(baseline_cfg, include_ffn=include_ffn) + baseline_score = score_config( + baseline_cfg, include_attention=include_attention, include_ffn=include_ffn, + ) baseline_latency = baseline_score.total_latency_ns for knob in _SENSITIVITY_KNOBS: @@ -330,7 +336,9 @@ def compute_sensitivity( doubled_cfg = FullConfig( model=model, topo=baseline_topo, machine=doubled_hw.as_machine(), ) - doubled_score = score_config(doubled_cfg, include_ffn=include_ffn) + doubled_score = score_config( + doubled_cfg, include_attention=include_attention, include_ffn=include_ffn, + ) rows.append(SensitivityRow( knob=knob, baseline_value=getattr(baseline_hw, knob), @@ -351,6 +359,7 @@ def joint_explore( s_kv: int, mode: str, depth: str = "balanced", + include_attention: bool = True, include_ffn: bool = True, ) -> JointExploreResult: """Sweep HW candidates × parallelism, return joint Pareto + sensitivity. @@ -358,7 +367,7 @@ def joint_explore( Sensitivity is computed around the LATENCY-MINIMUM joint point (the "best fast" config), doubling each HW knob one at a time. - ``include_ffn=False`` restricts scoring to attention stages only — + Scope flags select which stages contribute to the summed latency — both the per-HW parallelism search and the sensitivity ranking use the same restriction. """ @@ -367,11 +376,13 @@ def joint_explore( machine = hw.as_machine() if depth == "two_stage": par = _best_parallelism_two_stage( - model, machine, s_kv, mode, include_ffn=include_ffn, + model, machine, s_kv, mode, + include_attention=include_attention, include_ffn=include_ffn, ) else: par = _best_parallelism_for_hw( - model, machine, s_kv, mode, include_ffn=include_ffn, + model, machine, s_kv, mode, + include_attention=include_attention, include_ffn=include_ffn, ) if par is None: continue @@ -391,7 +402,7 @@ def joint_explore( best = all_scores[0] sensitivity = compute_sensitivity( best.hardware, best.parallelism, model, s_kv, mode, - include_ffn=include_ffn, + include_attention=include_attention, include_ffn=include_ffn, ) return JointExploreResult( diff --git a/tests/analytical_visualization/test_auto_explore.py b/tests/analytical_visualization/test_auto_explore.py index 28c70d6..06cf1cf 100644 --- a/tests/analytical_visualization/test_auto_explore.py +++ b/tests/analytical_visualization/test_auto_explore.py @@ -183,11 +183,31 @@ def test_attention_only_pareto_non_empty(): model = PRESETS["Llama 3.1 70B"].model machine = MachineParams() res = run_auto_explore(model, machine, s_kv=8192, mode="decode", - include_ffn=False) + include_attention=True, include_ffn=False) assert res.total_feasible > 0 assert len(res.pareto_scores) > 0 +def test_ffn_only_latency_less_than_full_and_different_from_attn(): + """FFN-only latency must be > 0, < full-model latency, and != attn-only.""" + model = PRESETS["Llama 3.1 70B"].model + machine = MachineParams() + r_full = run_auto_explore(model, machine, s_kv=8192, mode="decode", + include_attention=True, include_ffn=True) + r_attn = run_auto_explore(model, machine, s_kv=8192, mode="decode", + include_attention=True, include_ffn=False) + r_ffn = run_auto_explore(model, machine, s_kv=8192, mode="decode", + include_attention=False, include_ffn=True) + b_full = min(r_full.pareto_scores, key=lambda s: s.total_latency_ns) + b_attn = min(r_attn.pareto_scores, key=lambda s: s.total_latency_ns) + b_ffn = min(r_ffn.pareto_scores, key=lambda s: s.total_latency_ns) + assert b_ffn.total_latency_ns > 0 + assert b_ffn.total_latency_ns < b_full.total_latency_ns + assert abs(b_ffn.total_latency_ns - b_attn.total_latency_ns) > 1.0, ( + "FFN-only and attention-only latencies should not be identical" + ) + + def test_parallelism_sensitivity_includes_baseline_value(): """Each row's values contain the baseline_value.""" from tests.analytical_visualization.auto_explore import (