analytical-viz: attention-only vs attn+FFN via two sweep buttons
Both auto tabs now accept a scope choice at run time:
- "Run sweep — Attention" → include_ffn=False
- "Run sweep — Attn + FFN/MoE" → include_ffn=True
Each button runs an independent sweep and caches its result under its
own session_state key. The most recently clicked button determines the
displayed view; both caches persist so users can flip between the two
scopes without re-running.
Tabs:
- Renamed "Auto Explore" → "Auto Suggest Parallelism"
(accurately reflects that it only varies parallelism knobs; HW is
held at the sidebar values).
- "Auto Hardware" tab unchanged.
- Still 6 top-level tabs; no additional tabs added.
Core changes:
auto_explore.py:
- New include_ffn: bool = True parameter on _sum_visible_latency,
_efficiency, score_config, run_auto_explore, compute_parallelism_
sensitivity. False drops all FFN stages from the summed latency.
auto_hardware.py:
- New include_ffn: bool = True parameter on joint_explore,
compute_sensitivity, _best_parallelism_for_hw, _best_parallelism_
two_stage. Forwards to score_config.
Both defaults keep existing tests byte-identical.
Verified:
- 23 pytest tests pass (added 3 new: attn-only latency lower, attn-only
Pareto non-empty, joint HW attn-only faster than full).
- Smoke: Llama 70B decode 128K:
* Attn+FFN best latency: 12.85 ms (unchanged)
* Attention-only best: 7.35 ms (~57% of full)
* Both sensitivities top-rank bw_hbm_gbs (physics preserved).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -161,6 +161,33 @@ def test_parallelism_sensitivity_has_all_five_knobs():
|
||||
assert knobs == {"cp", "tp", "pp", "dp", "ep"}
|
||||
|
||||
|
||||
def test_attention_only_latency_is_lower_than_full():
|
||||
"""include_ffn=False must produce lower latency than include_ffn=True
|
||||
for the same model+workload (FFN cost is dropped)."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
r_full = run_auto_explore(model, machine, s_kv=8192, mode="decode",
|
||||
include_ffn=True)
|
||||
r_attn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
|
||||
include_ffn=False)
|
||||
best_full = min(r_full.pareto_scores, key=lambda s: s.total_latency_ns)
|
||||
best_attn = min(r_attn.pareto_scores, key=lambda s: s.total_latency_ns)
|
||||
assert best_attn.total_latency_ns < best_full.total_latency_ns, (
|
||||
f"attn-only ({best_attn.latency_us:.2f} us) should be less than "
|
||||
f"full ({best_full.latency_us:.2f} us)"
|
||||
)
|
||||
|
||||
|
||||
def test_attention_only_pareto_non_empty():
|
||||
"""The attention-only sweep must still produce a non-empty Pareto set."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
res = run_auto_explore(model, machine, s_kv=8192, mode="decode",
|
||||
include_ffn=False)
|
||||
assert res.total_feasible > 0
|
||||
assert len(res.pareto_scores) > 0
|
||||
|
||||
|
||||
def test_parallelism_sensitivity_includes_baseline_value():
|
||||
"""Each row's values contain the baseline_value."""
|
||||
from tests.analytical_visualization.auto_explore import (
|
||||
|
||||
Reference in New Issue
Block a user