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

analytical-viz: sidebar memory-min buttons are scope-aware (Attn / FFN / Attn+FFN)

Browse Source 
Replace the three memory-optimal Pareto-search buttons in the Parallelism
sidebar with three scope-aware memory-min autosuggest buttons. Each
button sizes for a different per-PE memory footprint:

  - Attn      -> attention weights + KV cache (no FFN weights)
  - FFN/MoE   -> FFN weights only (no attention, no KV)
  - Attn+FFN  -> everything (traditional autosuggest)

memory_layout: per_pe_weight_bytes and compute_memory take
include_attention / include_ffn flags; KV cache is zeroed when attention
scope is excluded.

autosuggest: _score_candidate and auto_suggest forward the flags into
compute_memory, so the smallest-fit search now respects the chosen
scope.

app.py: single "Apply memory-min" caption + 3 column buttons. Each
button runs auto_suggest with its scope filter, snaps the sliders to
the picked (CP, TP, PP), and sets the Physical Layout scope filter to
match. Removes the standalone Apply memory-min button (was duplicating
the Attn+FFN case) and the Pareto-search import.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-07-28 23:34:05 -07:00
parent 47c4f40f4d
commit 4ded67a443
3 changed files with 79 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/analytical_visualization/app.py
+31 -60
View File
@@ -206,25 +206,18 @@ with st.sidebar.expander("Parallelism", expanded=True):
f"{_default_suggestion.pes_used} PEs "
f"(cp_placement={_default_suggestion.cp_placement})"
)
# Memory-min apply — snap sliders back to what the caption says.
if st.button("Apply memory-min", width='stretch',
key="_sb_apply_memmin",
help="Set sliders to the memory-min config shown above "
"(smallest deployment that fits)."):
st.session_state["cp"] = _snap(_default_suggestion.cp, _CP_OPTS)
st.session_state["tp"] = _snap(_default_suggestion.tp, _TP_OPTS)
st.session_state["pp"] = _snap(_default_suggestion.pp, _PP_OPTS)
st.session_state["dp"] = 1
st.session_state["cp_placement"] = _default_suggestion.cp_placement
st.session_state["_pl_active_scope"] = "full"
st.rerun()
# Three memory-optimal apply buttons — pick a scope, load the
# smallest-fit Pareto config into the sidebar sliders.
from tests.analytical_visualization.auto_explore import (
run_auto_explore as _run_auto_explore_sb,
# Three scope-aware memory-min buttons. Each runs auto_suggest with a
# different (include_attention, include_ffn) filter so the memory-fit
# constraint counts only the relevant block:
# - Attn → attention weights + KV cache (no FFN)
# - FFN/MoE → FFN weights only (no attn / no KV)
# - Attn+FFN → everything (the traditional autosuggest)
st.caption(
"**Apply memory-min** — smallest deployment that fits the chosen "
"block. Each button sizes for a different memory footprint."
)
_sb_scope_meta = {
"attn": (True, False, "Attention"),
"attn": (True, False, "Attn"),
"ffn": (False, True, "FFN/MoE"),
"full": (True, True, "Attn+FFN"),
}
@@ -234,56 +227,34 @@ with st.sidebar.expander("Parallelism", expanded=True):
with _sb_col:
_sb_clicks[_sb_scope] = st.button(
_sb_scope_meta[_sb_scope][2], width='stretch',
key=f"_sb_apply_{_sb_scope}",
help=f"Apply memory-optimal (smallest-fit) parallelism for "
f"{_sb_scope_meta[_sb_scope][2]} scope.",
key=f"_sb_apply_memmin_{_sb_scope}",
help=(f"Memory-min autosuggest — memory budget counts only "
f"{_sb_scope_meta[_sb_scope][2]} weights"
f"{' + KV cache' if _sb_scope == 'attn' else ''}"
f". Picks the smallest deployment that fits."),
)
for _sb_scope, _sb_click in _sb_clicks.items():
if not _sb_click:
continue
_sb_ia, _sb_ff, _sb_lbl = _sb_scope_meta[_sb_scope]
with st.spinner(f"Finding memory-optimal parallelism ({_sb_lbl})..."):
_sb_res = _run_auto_explore_sb(
model, _default_machine, s_kv=s_kv, mode=mode,
include_attention=_sb_ia, include_ffn=_sb_ff, b=b_batch,
with st.spinner(f"Sizing memory-min for {_sb_lbl}..."):
_sb_sug = auto_suggest(
model, _default_machine, s_kv=s_kv, mode=mode, b=b_batch,
include_attention=_sb_ia, include_ffn=_sb_ff,
)
if not _sb_res.pareto_scores:
st.error(
f"No feasible parallelism for {_sb_lbl}. Try raising "
"`pe_hbm_gb` or reducing `s_kv`."
if not _sb_sug.fits:
st.warning(
f"{_sb_lbl} memory-min: no config fits at 10% slack; "
f"best-effort {_sb_sug.pes_used} PEs, {_sb_sug.cubes_used} "
f"cubes ({_sb_sug.reason})."
)
else:
# Smallest-fit Pareto point: fewest PEs first, then lowest
# HBM utilization, then fastest latency as final tiebreaker.
# Gives the smallest deployment that's still Pareto-optimal
# for the chosen scope (matches how the memory-min autosuggest
# thinks — but restricted to this scope's Pareto set).
_sb_best = min(_sb_res.pareto_scores,
key=lambda s: (s.pes_used,
s.hbm_utilization,
s.total_latency_ns))
st.session_state["cp"] = _snap(_sb_best.cp, _CP_OPTS)
st.session_state["tp"] = _snap(_sb_best.tp, _TP_OPTS)
st.session_state["pp"] = _snap(_sb_best.pp, _PP_OPTS)
st.session_state["dp"] = _snap(_sb_best.dp, _DP_OPTS)
st.session_state["tp_placement"] = _sb_best.tp_placement
st.session_state["cp_placement"] = _sb_best.cp_placement
st.session_state["cp_ring_variant"] = _sb_best.cp_ring_variant
st.session_state["kv_mode"] = _sb_best.kv_shard_mode
_sb_tp2 = _snap(_sb_best.tp, _TP_OPTS)
_sb_cp2 = _snap(_sb_best.cp, _CP_OPTS)
_sb_dp2 = _snap(_sb_best.dp, _DP_OPTS)
_sb_ffn_map = {
"TP": f"TP only (div={max(1, _sb_tp2)})",
"TP+CP": f"TP*CP (div={max(1, _sb_tp2 * _sb_cp2)})",
"TP+CP+DP": f"TP*CP*DP (div={max(1, _sb_tp2 * _sb_cp2 * _sb_dp2)})",
}
st.session_state["ffn_scope_label"] = \
_sb_ffn_map[_sb_best.ffn_shard_scope]
# Track for the Physical Layout scope filter, so the tables
# below also filter to this scope.
st.session_state["_pl_active_scope"] = _sb_scope
st.rerun()
st.session_state["cp"] = _snap(_sb_sug.cp, _CP_OPTS)
st.session_state["tp"] = _snap(_sb_sug.tp, _TP_OPTS)
st.session_state["pp"] = _snap(_sb_sug.pp, _PP_OPTS)
st.session_state["dp"] = 1
st.session_state["cp_placement"] = _sb_sug.cp_placement
st.session_state["_pl_active_scope"] = _sb_scope
st.rerun()
# Sharding dims — arrange in 2x3 grid to save vertical space
p_row1 = st.columns(3)
tests/analytical_visualization/autosuggest.py
+12 -4
View File
@@ -61,7 +61,9 @@ def _score_candidate(cp: int, tp: int, pp: int,
model: ModelConfig, machine: MachineParams,
s_kv: int, mode: str,
slack_frac: float = 0.10,
b: int = 1) -> Suggestion:
b: int = 1,
include_attention: bool = True,
include_ffn: bool = True) -> Suggestion:
# For each triple, try both cp_placement options and keep the one
# with fewer cubes (breaks the historical "CP always across cubes"
# default when a smaller pack is possible). The pe placement is only
@@ -82,7 +84,9 @@ def _score_candidate(cp: int, tp: int, pp: int,
best_placement = _place
topo = best_topo
cfg = FullConfig(model=model, topo=topo, machine=machine)
mem = compute_memory(cfg)
mem = compute_memory(
cfg, include_attention=include_attention, include_ffn=include_ffn,
)
fits = (not mem.over_budget
and mem.slack_bytes >= slack_frac * mem.budget_bytes)
@@ -110,7 +114,9 @@ def _score_candidate(cp: int, tp: int, pp: int,
def auto_suggest(model: ModelConfig, machine: MachineParams,
s_kv: int, mode: str = "decode",
slack_frac: float = 0.10,
b: int = 1) -> Suggestion:
b: int = 1,
include_attention: bool = True,
include_ffn: bool = True) -> Suggestion:
"""Return the smallest deployment (fewer cubes, then fewer PEs)
that fits.
@@ -127,7 +133,9 @@ def auto_suggest(model: ModelConfig, machine: MachineParams,
for cp, tp, pp in _iter_candidates(model):
scored.append(
_score_candidate(cp, tp, pp, model, machine, s_kv, mode,
slack_frac, b=b)
slack_frac, b=b,
include_attention=include_attention,
include_ffn=include_ffn)
)
scored.sort(key=lambda s: (s.cubes_used, s.pes_used, s.pp, s.tp, s.cp))
tests/analytical_visualization/memory_layout.py
+36 -15
View File
@@ -30,12 +30,19 @@ class MemoryBreakdown:
return self.used_bytes > self.budget_bytes
def per_pe_weight_bytes(cfg: FullConfig) -> int:
def per_pe_weight_bytes(cfg: FullConfig,
include_attention: bool = True,
include_ffn: bool = True) -> int:
"""Attention + FFN weights per PE, bf16. Divided by TP, PP, and EP.
EP divides the FFN (experts) across ranks; attention weights are
unaffected. When kv_shard_mode='replicate' and TP > H_kv, each KV
head is replicated (per-PE W_K/W_V size doesn't drop below one head).
Scope flags let callers count only attention or only FFN weights —
useful for "smallest deployment for just this block" sizing. Both
True (default) matches the traditional full-model per-PE weight
footprint.
"""
m = cfg.model
tp = cfg.topo.tp
@@ -50,18 +57,23 @@ def per_pe_weight_bytes(cfg: FullConfig) -> int:
# Head-dim split: fractional head allowed (< 1 head bytes when TP > H_kv).
hkv_per_pe_bytes = m.h_kv / tp
per_layer_attn = (
m.hidden * hq_per_pe * m.d_head + # W_Q
m.hidden * hkv_per_pe_bytes * m.d_head + # W_K
m.hidden * hkv_per_pe_bytes * m.d_head + # W_V
hq_per_pe * m.d_head * m.hidden # W_O
)
# FFN divisor: TP (default) or TP*CP or TP*CP*DP if user opts in.
# EP further divides (MoE experts).
ffn_div = cfg.ffn_shard_divisor * ep
per_layer_ffn = 3 * m.hidden * (m.ffn_dim // ffn_div)
per_layer = (per_layer_attn + per_layer_ffn) * m.bytes_per_elem
per_layer_attn = 0
if include_attention:
per_layer_attn = (
m.hidden * hq_per_pe * m.d_head + # W_Q
m.hidden * hkv_per_pe_bytes * m.d_head + # W_K
m.hidden * hkv_per_pe_bytes * m.d_head + # W_V
hq_per_pe * m.d_head * m.hidden # W_O
)
per_layer_ffn = 0
if include_ffn:
# FFN divisor: TP (default) or TP*CP or TP*CP*DP if user opts in.
# EP further divides (MoE experts).
ffn_div = cfg.ffn_shard_divisor * ep
per_layer_ffn = 3 * m.hidden * (m.ffn_dim // ffn_div)
per_layer = (per_layer_attn + per_layer_ffn) * m.bytes_per_elem
layers_per_stage = (m.layers + pp - 1) // pp
return int(per_layer * layers_per_stage)
@@ -103,10 +115,19 @@ def per_pe_transient_bytes(cfg: FullConfig) -> int:
return int(2 * tile_score + T_q * m.hidden * m.bytes_per_elem)
def compute_memory(cfg: FullConfig) -> MemoryBreakdown:
def compute_memory(cfg: FullConfig,
include_attention: bool = True,
include_ffn: bool = True) -> MemoryBreakdown:
"""Per-PE memory breakdown. Scope flags let callers count only the
attention or only the FFN block (KV cache goes with the attention
block; transient activation buffer is small either way and left in).
"""
kv_bytes = per_pe_kv_cache_bytes(cfg) if include_attention else 0
return MemoryBreakdown(
weights_bytes=per_pe_weight_bytes(cfg),
kv_cache_bytes=per_pe_kv_cache_bytes(cfg),
weights_bytes=per_pe_weight_bytes(
cfg, include_attention=include_attention, include_ffn=include_ffn,
),
kv_cache_bytes=kv_bytes,
transient_bytes=per_pe_transient_bytes(cfg),
budget_bytes=cfg.machine.pe_budget_bytes,
)