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

analytical-viz: attention-only vs attn+FFN via two sweep buttons

Browse Source 
Both auto tabs now accept a scope choice at run time:
  - "Run sweep — Attention"        → include_ffn=False
  - "Run sweep — Attn + FFN/MoE"   → include_ffn=True

Each button runs an independent sweep and caches its result under its
own session_state key. The most recently clicked button determines the
displayed view; both caches persist so users can flip between the two
scopes without re-running.

Tabs:
  - Renamed "Auto Explore" → "Auto Suggest Parallelism"
    (accurately reflects that it only varies parallelism knobs; HW is
    held at the sidebar values).
  - "Auto Hardware" tab unchanged.
  - Still 6 top-level tabs; no additional tabs added.

Core changes:
  auto_explore.py:
    - New include_ffn: bool = True parameter on _sum_visible_latency,
      _efficiency, score_config, run_auto_explore, compute_parallelism_
      sensitivity. False drops all FFN stages from the summed latency.

  auto_hardware.py:
    - New include_ffn: bool = True parameter on joint_explore,
      compute_sensitivity, _best_parallelism_for_hw, _best_parallelism_
      two_stage. Forwards to score_config.

Both defaults keep existing tests byte-identical.

Verified:
  - 23 pytest tests pass (added 3 new: attn-only latency lower, attn-only
    Pareto non-empty, joint HW attn-only faster than full).
  - Smoke: Llama 70B decode 128K:
      * Attn+FFN best latency: 12.85 ms (unchanged)
      * Attention-only best:    7.35 ms (~57% of full)
      * Both sensitivities top-rank bw_hbm_gbs (physics preserved).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-07-28 13:58:29 -07:00
parent 91b63eb9f6
commit b8e1a3322f
5 changed files with 264 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/analytical_visualization/app.py
+174 -58
View File
@@ -419,9 +419,14 @@ if _warnings:
# ── Tabs ─────────────────────────────────────────────────────────
tab_layout, tab_memory, tab_stages, tab_compare, tab_auto, tab_hw = st.tabs([
# One tab per auto feature; each has TWO sweep buttons (Attention only /
# Attn + FFN/MoE) so the user can toggle scope without switching tabs.
# Auto Suggest is renamed "Parallelism" since it only varies parallelism
# knobs (hardware is held fixed at the sidebar values).
tab_auto, tab_layout, tab_memory, tab_stages, tab_compare, tab_hw = st.tabs([
"Auto Suggest Parallelism",
"Physical layout", "Memory breakdown", "Per-stage latency",
"Save & compare", "Auto Explore", "Auto Hardware",
"Save & compare", "Auto Hardware",
])
@@ -1284,8 +1289,17 @@ with tab_compare:
)
# ── TAB 5: Auto Explore ─────────────────────────────────────────
with tab_auto:
# ── TAB (Auto Suggest Parallelism) — one tab, two sweep buttons ──
def _render_auto_explore_tab():
"""Render the Auto Suggest Parallelism tab body.
Two sweep buttons — Attention-only and Attn+FFN/MoE. Each caches its
result under its own session_state key. Whichever button was last
clicked drives the display below; both caches persist so re-clicking
either without changing model/workload just re-shows the cached view.
"""
from tests.analytical_visualization.auto_explore import (
compute_parallelism_sensitivity,
run_auto_explore,
@@ -1298,34 +1312,80 @@ with tab_auto:
"(latency ↓, PEs ↓, efficiency ↑). Throughput = 1/latency for a "
"single request so it collapses with latency and is shown for info only."
)
st.caption(
"**Attention only**: drop FFN / MoE from the summed latency — "
"useful when isolating attention-kernel tuning (matches what our "
"kernbench attention sim measures). \n"
"**Attn + FFN/MoE**: sum every stage — the full per-token cost."
)
ax_l, ax_r = st.columns([3, 1])
with ax_l:
st.markdown(
f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}** | "
f"per-PE HBM = **{machine.pe_hbm_gb:.1f} GB**"
st.markdown(
f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}** | "
f"per-PE HBM = **{machine.pe_hbm_gb:.1f} GB**"
)
_b1, _b2, _b3 = st.columns([1, 1, 2])
with _b1:
_run_attn = st.button("Run sweep — Attention",
type="primary", width='stretch',
key="_auto_run_attn")
with _b2:
_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)..."):
_r = run_auto_explore(
model, machine, s_kv=s_kv, mode=mode, include_ffn=False,
)
st.session_state["_auto_explore_result_attn"] = _r
st.session_state["_auto_explore_ctx_attn"] = (
model.name, s_kv, mode, machine.pe_hbm_gb,
)
with ax_r:
_run_sweep = st.button("Run sweep", type="primary", width='stretch')
if _run_sweep or st.session_state.get("_auto_explore_result") is not None:
if _run_sweep:
with st.spinner("Sweeping ~28k configs..."):
_res = run_auto_explore(model, machine, s_kv=s_kv, mode=mode)
st.session_state["_auto_explore_result"] = _res
st.session_state["_auto_explore_ctx"] = (
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,
)
_res = st.session_state["_auto_explore_result"]
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"
# Warn if the cached result is stale relative to the current config.
_cached_ctx = st.session_state.get("_auto_explore_ctx")
_cur_ctx = (model.name, s_kv, mode, machine.pe_hbm_gb)
if _cached_ctx != _cur_ctx:
st.warning(
"Sweep result is from a different model/workload. "
"Click **Run sweep** to refresh."
)
_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."
)
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"
_result_key = f"_auto_explore_result{_suffix}"
_ctx_key = f"_auto_explore_ctx{_suffix}"
_res = st.session_state.get(_result_key)
if _res is None:
st.info(f"Click **Run sweep — {_label.split(' ')[0]}** to compute this view.")
return
st.markdown(f"### Currently viewing: **{_label}**")
# Warn if the cached result is stale relative to the current config.
_cached_ctx = st.session_state.get(_ctx_key)
_cur_ctx = (model.name, s_kv, mode, machine.pe_hbm_gb)
if _cached_ctx != _cur_ctx:
st.warning(
f"Cached {_label} result is from a different model/workload. "
f"Click **Run sweep — {_label.split(' ')[0]}** to refresh."
)
_m1, _m2, _m3, _m4 = st.columns(4)
_m1.metric("Enumerated", f"{_res.total_enumerated:,}")
@@ -1431,12 +1491,13 @@ with tab_auto:
_sens_row = st.number_input(
"Baseline row #",
min_value=0, max_value=len(_pareto_by_lat) - 1,
value=0, step=1, key="_auto_sens_row",
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
@@ -1487,10 +1548,11 @@ with tab_auto:
with _lc1:
_pick = st.number_input(
"Row #", min_value=0, max_value=len(_pareto_by_lat) - 1,
value=0, step=1, key="_auto_pick_row",
value=0, step=1, key=f"_auto_pick_row{_suffix}",
)
with _lc2:
if st.button("Load into sidebar", type="secondary"):
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.
@@ -1527,8 +1589,17 @@ with tab_auto:
)
# ── TAB 6: Auto Hardware ─────────────────────────────────────────
with tab_hw:
with tab_auto:
_render_auto_explore_tab()
# ── TAB (Auto Hardware) — one tab, two sweep buttons ─────────────
def _render_auto_hardware_tab():
"""Render the Auto Hardware tab body.
Two sweep buttons — Attention-only and Attn+FFN/MoE. Each caches its
result. Whichever button was last clicked drives the display; both
caches persist for fast toggling."""
from tests.analytical_visualization.auto_hardware import (
_HW_KNOB_DEFAULTS, joint_explore,
)
@@ -1538,11 +1609,16 @@ with tab_hw:
"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?"
"computes per-knob sensitivity: which HW investment gives the "
"biggest speedup?"
)
st.caption(
"**Attention only**: FFN / MoE stages are excluded from the summed "
"latency (both in the joint search and sensitivity). \n"
"**Attn + FFN/MoE**: full per-token cost."
)
hx_l, hx_m, hx_r = st.columns([2, 2, 1])
hx_l, hx_m = st.columns([1, 1])
with hx_l:
st.markdown(
f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}**"
@@ -1557,28 +1633,64 @@ with tab_hw:
"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"]
_bh1, _bh2, _bh3 = st.columns([1, 1, 2])
with _bh1:
_run_hw_attn = st.button("Run joint sweep — Attention",
type="primary", width='stretch',
key="_hw_run_attn")
with _bh2:
_run_hw_full = st.button("Run joint sweep — Attn + FFN/MoE",
type="primary", width='stretch',
key="_hw_run_full")
_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."
)
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"] = (
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"
_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."
)
return
include_ffn = (_active == "full")
_suffix = "_full" if include_ffn else "_attn"
_label = "Attn + FFN/MoE" if include_ffn else "Attention only"
_hw_result_key = f"_hw_result{_suffix}"
_hw_ctx_key = f"_hw_ctx{_suffix}"
_hw_res = st.session_state.get(_hw_result_key)
if _hw_res is None:
st.info(f"Click **Run joint sweep — {_label.split(' ')[0]}** to compute this view.")
return
st.markdown(f"### Currently viewing: **{_label}**")
_cached_ctx = st.session_state.get(_hw_ctx_key)
_cur_ctx = (model.name, s_kv, mode, _depth)
if _cached_ctx != _cur_ctx:
st.warning(
f"Cached {_label} result is from a different model/workload/depth. "
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)
@@ -1691,11 +1803,11 @@ with tab_hw:
_pick_hw = st.number_input(
"Row #", min_value=0,
max_value=len(_pareto_by_lat) - 1,
value=0, step=1, key="_hw_pick_row",
value=0, step=1, key=f"_hw_pick_row{_suffix}",
)
with _lch2:
if st.button("Load into sidebar", type="secondary",
key="_hw_load_btn"):
key=f"_hw_load_btn{_suffix}"):
_s = _pareto_by_lat[int(_pick_hw)]
_h = _s.hardware
_p = _s.parallelism
@@ -1748,3 +1860,7 @@ with tab_hw:
"Click **Run joint sweep** to explore hardware and parallelism "
"co-design space for this model + workload."
)
with tab_hw:
_render_auto_hardware_tab()