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:
2026-07-28 13:58:29 -07:00
parent 91b63eb9f6
commit b8e1a3322f
5 changed files with 264 additions and 75 deletions
@@ -105,6 +105,22 @@ def test_sensitivity_all_knobs_monotone_non_worsening():
)
def test_joint_explore_attention_only_faster_than_full():
"""include_ffn=False produces smaller best-latency than include_ffn=True
for the same HW+model."""
model = PRESETS["Llama 3.1 70B"].model
r_full = joint_explore(model, s_kv=131072, mode="decode",
depth="two_stage", include_ffn=True)
r_attn = joint_explore(model, s_kv=131072, mode="decode",
depth="two_stage", include_ffn=False)
b_full = min(r_full.pareto_scores, key=lambda s: s.total_latency_ns)
b_attn = min(r_attn.pareto_scores, key=lambda s: s.total_latency_ns)
assert b_attn.total_latency_ns < b_full.total_latency_ns, (
f"attn-only ({b_attn.latency_ms:.2f}ms) should be less than "
f"full ({b_full.latency_ms:.2f}ms)"
)
def test_sensitivity_hbm_bw_dominant_for_llama_decode():
"""For Llama 70B decode (memory-bound), HBM BW should be the top
sensitivity knob — doubling it gives more speedup than any other knob."""