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:
@@ -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))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user