analytical-viz: add Auto Hardware Streamlit tab
Consumes auto_hardware.joint_explore(). UI:
- Sweep-depth radio: two_stage / balanced (default) / coarse
- Run joint sweep button + spinner
- 4 metric cards: HW candidates, feasible joint, Pareto count,
best latency
- Panel 1 (Pareto scatter): latency vs hardware cost proxy, feasible
in grey, Pareto colored by PE count, dashed line connects Pareto in
cost order — this is the "optimal HW config for the model" plot
- Panel 2 (sensitivity bar chart): per-knob relative speedup when
doubled from the baseline. Green bars, sorted biggest first;
annotated with the baseline → doubled values. Answers "where to
invest next?"
- Panel 3 (table): sortable Pareto joint configs with parallelism +
HW spec columns (PE HBM, HBM BW, TFLOPs, PE↔PE, D2D, C2C)
- Load into sidebar: picks a Pareto row and syncs both the HW
selectboxes (Per-PE + Interconnect sub-tabs) AND the parallelism
sliders. Values only get loaded when the target is one of the
selectbox options; otherwise the sidebar keeps its current value.
Session-state cache under _hw_result keyed by (model, s_kv, mode, depth);
if any of these drift, a warning suggests refreshing.
Verified: app.py parses; auto_hardware smoke run on Llama 70B decode
128K in ~20s produces sensible HW co-design signal (HBM BW 33% speedup,
everything else <1.5%).
Next: verification across Qwen 3 8B, Mixtral 8x7B (Commit 6).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -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."
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user