analytical-viz: add FFN-only scope + 3rd sweep button
Both auto tabs now offer three sweep scopes via three buttons instead
of two:
- Run sweep — Attention → include_attention=T, include_ffn=F
- Run sweep — FFN/MoE → include_attention=F, include_ffn=T
- Run sweep — Attn + FFN/MoE → include_attention=T, include_ffn=T
Each button caches its result under its own session_state key; the most
recently clicked button drives the display. All three caches persist so
users can flip between scopes without re-running.
Core changes:
auto_explore.py + auto_hardware.py:
- New include_attention: bool = True param alongside include_ffn
- _sum_visible_latency, _efficiency, score_config, run_auto_explore,
compute_parallelism_sensitivity, joint_explore, compute_sensitivity,
_best_parallelism_for_hw, _best_parallelism_two_stage all wired.
- Attention and FFN are additive with no overlap: measured 7.35 ms
(attn) + 5.50 ms (ffn) = 12.85 ms (full) for Llama 70B decode 128K.
Bug fix (drive-by): the display block in both tab render functions was
incorrectly nested inside the "cached ctx is stale" warning branch, so
metrics/scatter/table/sensitivity/load only rendered when the cache
was stale. Un-indented and removed the dead trailing else-info block.
Verified:
- 24 pytest tests pass (added 1 new for FFN-only scope invariants)
- Smoke: Llama 70B decode 128K all three scopes produce sensible Pareto:
* Attention only: 7.35 ms (14 pareto configs)
* FFN / MoE only: 5.50 ms (34 pareto configs)
* Attn + FFN/MoE: 12.85 ms (6 pareto configs)
Sums add up exactly, confirming no overlap in stage summation.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -164,7 +164,11 @@ def enumerate_configs(
|
||||
# ── Scoring ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _sum_visible_latency(cfg: FullConfig, include_ffn: bool = True) -> float:
|
||||
def _sum_visible_latency(
|
||||
cfg: FullConfig,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True,
|
||||
) -> float:
|
||||
"""Total single-request latency (seconds) across all model layers.
|
||||
|
||||
A single request traverses every layer sequentially, whether the layers
|
||||
@@ -177,28 +181,31 @@ def _sum_visible_latency(cfg: FullConfig, include_ffn: bool = True) -> float:
|
||||
throughput under batching. This function is the single-request cost, so
|
||||
we multiply by full model.layers regardless of PP.
|
||||
|
||||
``include_ffn=False`` restricts to attention stages only — useful for
|
||||
isolating attention-kernel tuning from FFN cost.
|
||||
Scope selection (both default True — full per-token cost):
|
||||
- ``include_attention=True, include_ffn=True`` → full transformer
|
||||
- ``include_attention=True, include_ffn=False`` → attention only
|
||||
- ``include_attention=False, include_ffn=True`` → FFN / MoE only
|
||||
- ``include_attention=False, include_ffn=False`` → zero (rejected caller-side)
|
||||
"""
|
||||
attn = sum(s.visible_s for s in all_stages(cfg))
|
||||
attn = sum(s.visible_s for s in all_stages(cfg)) if include_attention else 0.0
|
||||
ffn = sum(s.visible_s for s in all_ffn_stages(cfg)) if include_ffn else 0.0
|
||||
per_layer = attn + ffn
|
||||
return per_layer * cfg.model.layers
|
||||
|
||||
|
||||
def _efficiency(cfg: FullConfig, latency_s: float,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True) -> float:
|
||||
"""Geo-mean of compute-util and BW-util. Range ~ (0, 1].
|
||||
|
||||
- compute_util = achieved_flops / (peak_flops × pes × latency)
|
||||
- bw_util = achieved_bytes / (peak_bw × pes × latency)
|
||||
|
||||
``include_ffn=False`` mirrors the same restriction as
|
||||
:func:`_sum_visible_latency` — attention stages only.
|
||||
Scope flags mirror :func:`_sum_visible_latency`.
|
||||
"""
|
||||
if latency_s <= 0:
|
||||
return 0.0
|
||||
attn = all_stages(cfg)
|
||||
attn = all_stages(cfg) if include_attention else []
|
||||
ffn = all_ffn_stages(cfg) if include_ffn else []
|
||||
layers = math.ceil(cfg.model.layers / cfg.topo.pp)
|
||||
total_flops = layers * sum(s.flops for s in attn + ffn)
|
||||
@@ -216,24 +223,32 @@ def _efficiency(cfg: FullConfig, latency_s: float,
|
||||
return math.sqrt(compute_util * bw_util)
|
||||
|
||||
|
||||
def score_config(cfg: FullConfig, include_ffn: bool = True) -> ConfigScore:
|
||||
def score_config(cfg: FullConfig,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True) -> ConfigScore:
|
||||
"""Compute all 4 objectives + info fields for one config.
|
||||
|
||||
Feasibility (memory + placement) is stored but does NOT gate scoring —
|
||||
infeasible configs get returned with fits_memory=False so callers can
|
||||
filter or display them.
|
||||
|
||||
``include_ffn=False`` restricts latency + efficiency to attention stages
|
||||
only. Memory feasibility is unchanged — still checks weights+KV+transient
|
||||
fit in per-PE HBM since the model still exists physically.
|
||||
Scope flags (default: full transformer) restrict *latency + efficiency*
|
||||
to attention only, FFN only, or both. Memory feasibility is unchanged —
|
||||
still checks weights+KV+transient fit in per-PE HBM since the model
|
||||
still exists physically regardless of what the caller is scoring.
|
||||
"""
|
||||
mem = compute_memory(cfg)
|
||||
placement_ok = cfg.topo.placement_valid
|
||||
|
||||
latency_s = _sum_visible_latency(cfg, include_ffn=include_ffn)
|
||||
latency_s = _sum_visible_latency(
|
||||
cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
throughput = 1.0 / latency_s if latency_s > 0 else 0.0
|
||||
efficiency = (_efficiency(cfg, latency_s, include_ffn=include_ffn)
|
||||
if latency_s > 0 else 0.0)
|
||||
efficiency = (
|
||||
_efficiency(cfg, latency_s,
|
||||
include_attention=include_attention, include_ffn=include_ffn)
|
||||
if latency_s > 0 else 0.0
|
||||
)
|
||||
fits = not mem.over_budget
|
||||
|
||||
reason = ""
|
||||
@@ -345,6 +360,7 @@ def compute_parallelism_sensitivity(
|
||||
machine: MachineParams,
|
||||
s_kv: int,
|
||||
mode: str,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True,
|
||||
) -> list[ParallelismSensitivityRow]:
|
||||
"""For each parallelism knob (CP, TP, PP, DP, EP), sweep its values
|
||||
@@ -377,7 +393,9 @@ def compute_parallelism_sensitivity(
|
||||
# Build a variant TopologyConfig with just this one knob changed.
|
||||
trial = replace(baseline_topo, **{knob: v})
|
||||
cfg = FullConfig(model=model, topo=trial, machine=machine)
|
||||
score = score_config(cfg, include_ffn=include_ffn)
|
||||
score = score_config(
|
||||
cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
latencies.append(score.total_latency_ns)
|
||||
fits.append(score.fits_memory and score.placement_valid)
|
||||
rows.append(ParallelismSensitivityRow(
|
||||
@@ -399,6 +417,7 @@ def run_auto_explore(
|
||||
machine: MachineParams,
|
||||
s_kv: int,
|
||||
mode: str = "decode",
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True,
|
||||
) -> AutoExploreResult:
|
||||
"""Enumerate all configs, score each, extract Pareto frontier.
|
||||
@@ -407,15 +426,18 @@ def run_auto_explore(
|
||||
``pareto_scores`` subset (for the scatter/highlight view). Both are sorted
|
||||
by total_latency_ns ascending.
|
||||
|
||||
``include_ffn=False`` restricts to attention stages only — useful for
|
||||
isolating attention-kernel tuning independent of FFN cost.
|
||||
Scope: at least one of ``include_attention`` / ``include_ffn`` must be
|
||||
True. Setting both False would give zero latency for every config —
|
||||
the caller should not do that.
|
||||
"""
|
||||
all_scores: list[ConfigScore] = []
|
||||
total_enumerated = 0
|
||||
for topo in enumerate_configs(model, s_kv, mode):
|
||||
total_enumerated += 1
|
||||
cfg = FullConfig(model=model, topo=topo, machine=machine)
|
||||
all_scores.append(score_config(cfg, include_ffn=include_ffn))
|
||||
all_scores.append(score_config(
|
||||
cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
))
|
||||
|
||||
all_scores.sort(key=lambda s: s.total_latency_ns)
|
||||
pareto = pareto_frontier(all_scores)
|
||||
|
||||
Reference in New Issue
Block a user