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/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,
)