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:
tests/analytical_visualization/auto_explore.py
+40 -18
View File
@@ -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)
tests/analytical_visualization/auto_hardware.py
+24 -13
View File
@@ -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(
tests/analytical_visualization/test_auto_explore.py
+21 -1
View File
@@ -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 (