From 4ded67a44371b1e5f046658d9d87ff9d455aa6cd Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Tue, 28 Jul 2026 23:34:05 -0700 Subject: [PATCH] analytical-viz: sidebar memory-min buttons are scope-aware (Attn / FFN / Attn+FFN) 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 --- tests/analytical_visualization/app.py | 91 +++++++------------ tests/analytical_visualization/autosuggest.py | 16 +++- .../analytical_visualization/memory_layout.py | 51 ++++++++--- 3 files changed, 79 insertions(+), 79 deletions(-) diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 3d6877a..86e379a 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -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) diff --git a/tests/analytical_visualization/autosuggest.py b/tests/analytical_visualization/autosuggest.py index d6340ee..cac456a 100644 --- a/tests/analytical_visualization/autosuggest.py +++ b/tests/analytical_visualization/autosuggest.py @@ -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)) diff --git a/tests/analytical_visualization/memory_layout.py b/tests/analytical_visualization/memory_layout.py index 40c48c6..7162717 100644 --- a/tests/analytical_visualization/memory_layout.py +++ b/tests/analytical_visualization/memory_layout.py @@ -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, )