ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

analytical-viz: add FFN-only scope + 3rd sweep button

Browse Source 
Both auto tabs now offer three sweep scopes via three buttons instead
of two:
  - Run sweep — Attention          → include_attention=T, include_ffn=F
  - Run sweep — FFN/MoE            → include_attention=F, include_ffn=T
  - Run sweep — Attn + FFN/MoE     → include_attention=T, include_ffn=T

Each button caches its result under its own session_state key; the most
recently clicked button drives the display. All three caches persist so
users can flip between scopes without re-running.

Core changes:
  auto_explore.py + auto_hardware.py:
    - New include_attention: bool = True param alongside include_ffn
    - _sum_visible_latency, _efficiency, score_config, run_auto_explore,
      compute_parallelism_sensitivity, joint_explore, compute_sensitivity,
      _best_parallelism_for_hw, _best_parallelism_two_stage all wired.
    - Attention and FFN are additive with no overlap: measured 7.35 ms
      (attn) + 5.50 ms (ffn) = 12.85 ms (full) for Llama 70B decode 128K.

Bug fix (drive-by): the display block in both tab render functions was
incorrectly nested inside the "cached ctx is stale" warning branch, so
metrics/scatter/table/sensitivity/load only rendered when the cache
was stale. Un-indented and removed the dead trailing else-info block.

Verified:
  - 24 pytest tests pass (added 1 new for FFN-only scope invariants)
  - Smoke: Llama 70B decode 128K all three scopes produce sensible Pareto:
      * Attention only: 7.35 ms (14 pareto configs)
      * FFN / MoE only: 5.50 ms (34 pareto configs)
      * Attn + FFN/MoE: 12.85 ms (6 pareto configs)
    Sums add up exactly, confirming no overlap in stage summation.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-07-28 14:09:24 -07:00
parent b8e1a3322f
commit bf5b659a3d
4 changed files with 339 additions and 294 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/analytical_visualization/app.py
+254 -262
View File
@@ -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: