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 <noreply@anthropic.com>
This commit is contained in:
2026-07-28 23:34:05 -07:00
parent 47c4f40f4d
commit 4ded67a443
3 changed files with 79 additions and 79 deletions
+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))