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:
@@ -235,18 +235,20 @@ def _iter_reduced_parallelism(
|
||||
|
||||
def _best_parallelism_for_hw(
|
||||
model: ModelConfig, machine: MachineParams,
|
||||
s_kv: int, mode: str, include_ffn: bool = True,
|
||||
s_kv: int, mode: str,
|
||||
include_attention: bool = True, include_ffn: bool = True,
|
||||
) -> ConfigScore | None:
|
||||
"""Return the latency-minimum feasible parallelism for this HW.
|
||||
|
||||
Feasibility: fits memory + placement_valid. Returns None if nothing fits.
|
||||
``include_ffn`` is forwarded to score_config so the "latency" ranked on
|
||||
matches the caller's attention-only vs full-model choice.
|
||||
Scope flags are forwarded to score_config so the "latency" ranked on
|
||||
matches the caller's attention/FFN/full choice.
|
||||
"""
|
||||
best: ConfigScore | None = None
|
||||
for topo in _iter_reduced_parallelism(model, s_kv, mode):
|
||||
cfg = FullConfig(model=model, topo=topo, machine=machine)
|
||||
s = score_config(cfg, include_ffn=include_ffn)
|
||||
s = score_config(cfg, include_attention=include_attention,
|
||||
include_ffn=include_ffn)
|
||||
if not (s.fits_memory and s.placement_valid):
|
||||
continue
|
||||
if best is None or s.total_latency_ns < best.total_latency_ns:
|
||||
@@ -256,7 +258,8 @@ def _best_parallelism_for_hw(
|
||||
|
||||
def _best_parallelism_two_stage(
|
||||
model: ModelConfig, machine: MachineParams,
|
||||
s_kv: int, mode: str, include_ffn: bool = True,
|
||||
s_kv: int, mode: str,
|
||||
include_attention: bool = True, include_ffn: bool = True,
|
||||
) -> ConfigScore | None:
|
||||
"""Fast fallback: use autosuggest's memory-min (CP,TP,PP) then score it."""
|
||||
sug = auto_suggest(model, machine, s_kv, mode)
|
||||
@@ -272,7 +275,8 @@ def _best_parallelism_two_stage(
|
||||
cp_ring_variant="qoml" if mode == "decode" and sug.cp > 1 else "kv",
|
||||
)
|
||||
cfg = FullConfig(model=model, topo=topo, machine=machine)
|
||||
s = score_config(cfg, include_ffn=include_ffn)
|
||||
s = score_config(cfg, include_attention=include_attention,
|
||||
include_ffn=include_ffn)
|
||||
return s if s.fits_memory and s.placement_valid else None
|
||||
|
||||
|
||||
@@ -310,7 +314,7 @@ def compute_sensitivity(
|
||||
baseline_hw: HardwareCandidate,
|
||||
parallelism: ConfigScore,
|
||||
model: ModelConfig, s_kv: int, mode: str,
|
||||
include_ffn: bool = True,
|
||||
include_attention: bool = True, include_ffn: bool = True,
|
||||
) -> list[SensitivityRow]:
|
||||
"""For each HW knob, double it (holding others at baseline) and measure
|
||||
the latency change. Same parallelism used throughout so we isolate the
|
||||
@@ -321,7 +325,9 @@ def compute_sensitivity(
|
||||
baseline_cfg = FullConfig(
|
||||
model=model, topo=baseline_topo, machine=baseline_machine,
|
||||
)
|
||||
baseline_score = score_config(baseline_cfg, include_ffn=include_ffn)
|
||||
baseline_score = score_config(
|
||||
baseline_cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
baseline_latency = baseline_score.total_latency_ns
|
||||
|
||||
for knob in _SENSITIVITY_KNOBS:
|
||||
@@ -330,7 +336,9 @@ def compute_sensitivity(
|
||||
doubled_cfg = FullConfig(
|
||||
model=model, topo=baseline_topo, machine=doubled_hw.as_machine(),
|
||||
)
|
||||
doubled_score = score_config(doubled_cfg, include_ffn=include_ffn)
|
||||
doubled_score = score_config(
|
||||
doubled_cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
rows.append(SensitivityRow(
|
||||
knob=knob,
|
||||
baseline_value=getattr(baseline_hw, knob),
|
||||
@@ -351,6 +359,7 @@ def joint_explore(
|
||||
s_kv: int,
|
||||
mode: str,
|
||||
depth: str = "balanced",
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True,
|
||||
) -> JointExploreResult:
|
||||
"""Sweep HW candidates × parallelism, return joint Pareto + sensitivity.
|
||||
@@ -358,7 +367,7 @@ def joint_explore(
|
||||
Sensitivity is computed around the LATENCY-MINIMUM joint point (the
|
||||
"best fast" config), doubling each HW knob one at a time.
|
||||
|
||||
``include_ffn=False`` restricts scoring to attention stages only —
|
||||
Scope flags select which stages contribute to the summed latency —
|
||||
both the per-HW parallelism search and the sensitivity ranking use
|
||||
the same restriction.
|
||||
"""
|
||||
@@ -367,11 +376,13 @@ def joint_explore(
|
||||
machine = hw.as_machine()
|
||||
if depth == "two_stage":
|
||||
par = _best_parallelism_two_stage(
|
||||
model, machine, s_kv, mode, include_ffn=include_ffn,
|
||||
model, machine, s_kv, mode,
|
||||
include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
else:
|
||||
par = _best_parallelism_for_hw(
|
||||
model, machine, s_kv, mode, include_ffn=include_ffn,
|
||||
model, machine, s_kv, mode,
|
||||
include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
if par is None:
|
||||
continue
|
||||
@@ -391,7 +402,7 @@ def joint_explore(
|
||||
best = all_scores[0]
|
||||
sensitivity = compute_sensitivity(
|
||||
best.hardware, best.parallelism, model, s_kv, mode,
|
||||
include_ffn=include_ffn,
|
||||
include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
|
||||
return JointExploreResult(
|
||||
|
||||
Reference in New Issue
Block a user