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>

tests/analytical_visualization/app.py

@@ -419,9 +419,14 @@ if _warnings:
 
 
 # ── Tabs ─────────────────────────────────────────────────────────
-tab_layout, tab_memory, tab_stages, tab_compare, tab_auto, tab_hw = st.tabs([
+# One tab per auto feature; each has TWO sweep buttons (Attention only /
+# Attn + FFN/MoE) so the user can toggle scope without switching tabs.
+# Auto Suggest is renamed "Parallelism" since it only varies parallelism
+# knobs (hardware is held fixed at the sidebar values).
+tab_auto, tab_layout, tab_memory, tab_stages, tab_compare, tab_hw = st.tabs([
+    "Auto Suggest Parallelism",
     "Physical layout", "Memory breakdown", "Per-stage latency",
-    "Save & compare", "Auto Explore", "Auto Hardware",
+    "Save & compare", "Auto Hardware",
 ])
 
 
@@ -1284,8 +1289,17 @@ with tab_compare:
         )
 
 
-# ── TAB 5: Auto Explore ─────────────────────────────────────────
-with tab_auto:
+# ── TAB (Auto Suggest Parallelism) — one tab, two sweep buttons ──
+
+
+def _render_auto_explore_tab():
+    """Render the Auto Suggest Parallelism tab body.
+
+    Two sweep buttons — Attention-only and Attn+FFN/MoE. Each caches its
+    result under its own session_state key. Whichever button was last
+    clicked drives the display below; both caches persist so re-clicking
+    either without changing model/workload just re-shows the cached view.
+    """
     from tests.analytical_visualization.auto_explore import (
         compute_parallelism_sensitivity,
         run_auto_explore,
@@ -1298,33 +1312,79 @@ with tab_auto:
         "(latency ↓, PEs ↓, efficiency ↑). Throughput = 1/latency for a "
         "single request so it collapses with latency and is shown for info only."
     )
+    st.caption(
+        "**Attention only**: drop FFN / MoE from the summed latency — "
+        "useful when isolating attention-kernel tuning (matches what our "
+        "kernbench attention sim measures).  \n"
+        "**Attn + FFN/MoE**: sum every stage — the full per-token cost."
+    )
 
-    ax_l, ax_r = st.columns([3, 1])
-    with ax_l:
     st.markdown(
         f"**Context:** {model.name}  |  S_kv=**{s_kv:,}**  |  mode=**{mode}**  |  "
         f"per-PE HBM = **{machine.pe_hbm_gb:.1f} GB**"
     )
-    with ax_r:
-        _run_sweep = st.button("Run sweep", type="primary", width='stretch')
 
-    if _run_sweep or st.session_state.get("_auto_explore_result") is not None:
-        if _run_sweep:
-            with st.spinner("Sweeping ~28k configs..."):
-                _res = run_auto_explore(model, machine, s_kv=s_kv, mode=mode)
-            st.session_state["_auto_explore_result"] = _res
-            st.session_state["_auto_explore_ctx"] = (
+    _b1, _b2, _b3 = st.columns([1, 1, 2])
+    with _b1:
+        _run_attn = st.button("Run sweep — Attention",
+                               type="primary", width='stretch',
+                               key="_auto_run_attn")
+    with _b2:
+        _run_full = st.button("Run sweep — Attn + FFN/MoE",
+                               type="primary", width='stretch',
+                               key="_auto_run_full")
+
+    if _run_attn:
+        with st.spinner("Sweeping ~28k configs (Attention only)..."):
+            _r = run_auto_explore(
+                model, machine, s_kv=s_kv, mode=mode, include_ffn=False,
+            )
+        st.session_state["_auto_explore_result_attn"] = _r
+        st.session_state["_auto_explore_ctx_attn"] = (
             model.name, s_kv, mode, machine.pe_hbm_gb,
         )
-        _res = st.session_state["_auto_explore_result"]
+        st.session_state["_auto_explore_active"] = "attn"
+    if _run_full:
+        with st.spinner("Sweeping ~28k configs (Attn + FFN/MoE)..."):
+            _r = run_auto_explore(
+                model, machine, s_kv=s_kv, mode=mode, include_ffn=True,
+            )
+        st.session_state["_auto_explore_result_full"] = _r
+        st.session_state["_auto_explore_ctx_full"] = (
+            model.name, s_kv, mode, machine.pe_hbm_gb,
+        )
+        st.session_state["_auto_explore_active"] = "full"
+
+    _active = st.session_state.get("_auto_explore_active")
+    if _active is None:
+        st.info(
+            "Click one of the run buttons above to enumerate valid "
+            "parallelism configurations and rank them on the Pareto "
+            "frontier. Takes ~5–10 s. Both button results are cached so "
+            "you can flip between the two scopes without re-running."
+        )
+        return
+
+    # Downstream code expects _suffix, _label, include_ffn, _res, _ctx_key.
+    include_ffn = (_active == "full")
+    _suffix = "_full" if include_ffn else "_attn"
+    _label = "Attn + FFN/MoE" if include_ffn else "Attention only"
+    _result_key = f"_auto_explore_result{_suffix}"
+    _ctx_key = f"_auto_explore_ctx{_suffix}"
+    _res = st.session_state.get(_result_key)
+    if _res is None:
+        st.info(f"Click **Run sweep — {_label.split(' ')[0]}** to compute this view.")
+        return
+
+    st.markdown(f"### Currently viewing: **{_label}**")
 
     # Warn if the cached result is stale relative to the current config.
-        _cached_ctx = st.session_state.get("_auto_explore_ctx")
+    _cached_ctx = st.session_state.get(_ctx_key)
     _cur_ctx = (model.name, s_kv, mode, machine.pe_hbm_gb)
     if _cached_ctx != _cur_ctx:
         st.warning(
-                "Sweep result is from a different model/workload. "
-                "Click **Run sweep** to refresh."
+            f"Cached {_label} result is from a different model/workload. "
+            f"Click **Run sweep — {_label.split(' ')[0]}** to refresh."
         )
 
         _m1, _m2, _m3, _m4 = st.columns(4)
@@ -1431,12 +1491,13 @@ with tab_auto:
                 _sens_row = st.number_input(
                     "Baseline row #",
                     min_value=0, max_value=len(_pareto_by_lat) - 1,
-                    value=0, step=1, key="_auto_sens_row",
+                    value=0, step=1, key=f"_auto_sens_row{_suffix}",
                     help="Which Pareto row is the sensitivity computed around?",
                 )
             _baseline_for_sens = _pareto_by_lat[int(_sens_row)]
             _sens_rows = compute_parallelism_sensitivity(
                 _baseline_for_sens, model, machine, s_kv=s_kv, mode=mode,
+                include_ffn=include_ffn,
             )
 
             import matplotlib.pyplot as _plt
@@ -1487,10 +1548,11 @@ with tab_auto:
             with _lc1:
                 _pick = st.number_input(
                     "Row #", min_value=0, max_value=len(_pareto_by_lat) - 1,
-                    value=0, step=1, key="_auto_pick_row",
+                    value=0, step=1, key=f"_auto_pick_row{_suffix}",
                 )
             with _lc2:
-                if st.button("Load into sidebar", type="secondary"):
+                if st.button("Load into sidebar", type="secondary",
+                              key=f"_auto_load_btn{_suffix}"):
                     _s = _pareto_by_lat[int(_pick)]
                     # Set the sidebar's session_state keys directly. The
                     # sidebar reads these on the next rerun.
@@ -1527,8 +1589,17 @@ with tab_auto:
         )
 
 
-# ── TAB 6: Auto Hardware ─────────────────────────────────────────
-with tab_hw:
+with tab_auto:
+    _render_auto_explore_tab()
+
+
+# ── TAB (Auto Hardware) — one tab, two sweep buttons ─────────────
+def _render_auto_hardware_tab():
+    """Render the Auto Hardware tab body.
+
+    Two sweep buttons — Attention-only and Attn+FFN/MoE. Each caches its
+    result. Whichever button was last clicked drives the display; both
+    caches persist for fast toggling."""
     from tests.analytical_visualization.auto_hardware import (
         _HW_KNOB_DEFAULTS, joint_explore,
     )
@@ -1538,11 +1609,16 @@ with tab_hw:
         "interconnect BW) **together** with parallelism. For each hardware "
         "candidate, we pick the latency-minimum parallelism that fits, then "
         "Pareto-rank the joint (latency ↓, hardware cost ↓) space. Also "
-        "computes per-knob sensitivity: which HW investment gives the biggest "
-        "speedup?"
+        "computes per-knob sensitivity: which HW investment gives the "
+        "biggest speedup?"
+    )
+    st.caption(
+        "**Attention only**: FFN / MoE stages are excluded from the summed "
+        "latency (both in the joint search and sensitivity).  \n"
+        "**Attn + FFN/MoE**: full per-token cost."
     )
 
-    hx_l, hx_m, hx_r = st.columns([2, 2, 1])
+    hx_l, hx_m = st.columns([1, 1])
     with hx_l:
         st.markdown(
             f"**Context:** {model.name}  |  S_kv=**{s_kv:,}**  |  mode=**{mode}**"
@@ -1557,27 +1633,63 @@ with tab_hw:
                   "balanced (default): 64 HW × ~2k parallelism, ~10-20s.\n"
                   "coarse: 729 HW × ~2k parallelism, ~2-5 min."),
         )
-    with hx_r:
-        _run_hw = st.button("Run joint sweep", type="primary",
-                             width='stretch', key="_hw_run_btn")
 
-    if _run_hw or st.session_state.get("_hw_result") is not None:
-        if _run_hw:
-            with st.spinner(f"Sweeping HW × parallelism ({_depth})..."):
-                _hw_res = joint_explore(model, s_kv=s_kv,
-                                          mode=mode, depth=_depth)
-            st.session_state["_hw_result"] = _hw_res
-            st.session_state["_hw_ctx"] = (
+    _bh1, _bh2, _bh3 = st.columns([1, 1, 2])
+    with _bh1:
+        _run_hw_attn = st.button("Run joint sweep — Attention",
+                                   type="primary", width='stretch',
+                                   key="_hw_run_attn")
+    with _bh2:
+        _run_hw_full = st.button("Run joint sweep — Attn + FFN/MoE",
+                                   type="primary", width='stretch',
+                                   key="_hw_run_full")
+
+    if _run_hw_attn:
+        with st.spinner(f"Sweeping HW × parallelism ({_depth}, Attention only)..."):
+            _r = joint_explore(model, s_kv=s_kv, mode=mode,
+                                depth=_depth, include_ffn=False)
+        st.session_state["_hw_result_attn"] = _r
+        st.session_state["_hw_ctx_attn"] = (
             model.name, s_kv, mode, _depth,
         )
-        _hw_res = st.session_state["_hw_result"]
+        st.session_state["_hw_active"] = "attn"
+    if _run_hw_full:
+        with st.spinner(f"Sweeping HW × parallelism ({_depth}, Attn + FFN/MoE)..."):
+            _r = joint_explore(model, s_kv=s_kv, mode=mode,
+                                depth=_depth, include_ffn=True)
+        st.session_state["_hw_result_full"] = _r
+        st.session_state["_hw_ctx_full"] = (
+            model.name, s_kv, mode, _depth,
+        )
+        st.session_state["_hw_active"] = "full"
 
-        _cached_ctx = st.session_state.get("_hw_ctx")
+    _active = st.session_state.get("_hw_active")
+    if _active is None:
+        st.info(
+            "Click one of the run buttons to explore hardware × parallelism "
+            "co-design for this model + workload. Both button results are "
+            "cached so you can flip between the two scopes without re-running."
+        )
+        return
+
+    include_ffn = (_active == "full")
+    _suffix = "_full" if include_ffn else "_attn"
+    _label = "Attn + FFN/MoE" if include_ffn else "Attention only"
+    _hw_result_key = f"_hw_result{_suffix}"
+    _hw_ctx_key = f"_hw_ctx{_suffix}"
+    _hw_res = st.session_state.get(_hw_result_key)
+    if _hw_res is None:
+        st.info(f"Click **Run joint sweep — {_label.split(' ')[0]}** to compute this view.")
+        return
+
+    st.markdown(f"### Currently viewing: **{_label}**")
+
+    _cached_ctx = st.session_state.get(_hw_ctx_key)
     _cur_ctx = (model.name, s_kv, mode, _depth)
     if _cached_ctx != _cur_ctx:
         st.warning(
-                "Sweep result is from a different model/workload/depth. "
-                "Click **Run joint sweep** to refresh."
+            f"Cached {_label} result is from a different model/workload/depth. "
+            f"Click **Run joint sweep — {_label.split(' ')[0]}** to refresh."
         )
 
         _hm1, _hm2, _hm3, _hm4 = st.columns(4)
@@ -1691,11 +1803,11 @@ with tab_hw:
                 _pick_hw = st.number_input(
                     "Row #", min_value=0,
                     max_value=len(_pareto_by_lat) - 1,
-                    value=0, step=1, key="_hw_pick_row",
+                    value=0, step=1, key=f"_hw_pick_row{_suffix}",
                 )
             with _lch2:
                 if st.button("Load into sidebar", type="secondary",
-                              key="_hw_load_btn"):
+                              key=f"_hw_load_btn{_suffix}"):
                     _s = _pareto_by_lat[int(_pick_hw)]
                     _h = _s.hardware
                     _p = _s.parallelism
@@ -1748,3 +1860,7 @@ with tab_hw:
             "Click **Run joint sweep** to explore hardware and parallelism "
             "co-design space for this model + workload."
         )
+
+
+with tab_hw:
+    _render_auto_hardware_tab()

tests/analytical_visualization/auto_explore.py

@@ -164,7 +164,7 @@ def enumerate_configs(
 # ── Scoring ──────────────────────────────────────────────────────────
 
 
-def _sum_visible_latency(cfg: FullConfig) -> float:
+def _sum_visible_latency(cfg: FullConfig, include_ffn: bool = True) -> float:
     """Total single-request latency (seconds) across all model layers.
 
     A single request traverses every layer sequentially, whether the layers
@@ -176,23 +176,30 @@ def _sum_visible_latency(cfg: FullConfig) -> float:
     PP therefore does NOT reduce single-request latency; it only improves
     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.
     """
     attn = sum(s.visible_s for s in all_stages(cfg))
-    ffn = sum(s.visible_s for s in all_ffn_stages(cfg))
+    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) -> float:
+def _efficiency(cfg: FullConfig, latency_s: float,
+                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.
     """
     if latency_s <= 0:
         return 0.0
     attn = all_stages(cfg)
-    ffn = all_ffn_stages(cfg)
+    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)
     total_bytes = layers * sum(s.mem_bytes for s in attn + ffn)
@@ -209,19 +216,24 @@ def _efficiency(cfg: FullConfig, latency_s: float) -> float:
     return math.sqrt(compute_util * bw_util)
 
 
-def score_config(cfg: FullConfig) -> ConfigScore:
+def score_config(cfg: FullConfig, 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.
     """
     mem = compute_memory(cfg)
     placement_ok = cfg.topo.placement_valid
 
-    latency_s = _sum_visible_latency(cfg)
+    latency_s = _sum_visible_latency(cfg, include_ffn=include_ffn)
     throughput = 1.0 / latency_s if latency_s > 0 else 0.0
-    efficiency = _efficiency(cfg, 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)
     fits = not mem.over_budget
 
     reason = ""
@@ -333,6 +345,7 @@ def compute_parallelism_sensitivity(
     machine: MachineParams,
     s_kv: int,
     mode: str,
+    include_ffn: bool = True,
 ) -> list[ParallelismSensitivityRow]:
     """For each parallelism knob (CP, TP, PP, DP, EP), sweep its values
     holding the OTHER knobs fixed at the baseline. Reports latency + memory
@@ -364,7 +377,7 @@ 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)
+            score = score_config(cfg, include_ffn=include_ffn)
             latencies.append(score.total_latency_ns)
             fits.append(score.fits_memory and score.placement_valid)
         rows.append(ParallelismSensitivityRow(
@@ -386,19 +399,23 @@ def run_auto_explore(
     machine: MachineParams,
     s_kv: int,
     mode: str = "decode",
+    include_ffn: bool = True,
 ) -> AutoExploreResult:
     """Enumerate all configs, score each, extract Pareto frontier.
 
     Returns both the full ``all_scores`` list (for the table view) and the
     ``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.
     """
     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))
+        all_scores.append(score_config(cfg, include_ffn=include_ffn))
 
     all_scores.sort(key=lambda s: s.total_latency_ns)
     pareto = pareto_frontier(all_scores)

tests/analytical_visualization/auto_hardware.py

@@ -235,16 +235,18 @@ def _iter_reduced_parallelism(
 
 def _best_parallelism_for_hw(
     model: ModelConfig, machine: MachineParams,
-    s_kv: int, mode: str,
+    s_kv: int, mode: str, 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.
     """
     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)
+        s = score_config(cfg, 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:
@@ -254,7 +256,7 @@ def _best_parallelism_for_hw(
 
 def _best_parallelism_two_stage(
     model: ModelConfig, machine: MachineParams,
-    s_kv: int, mode: str,
+    s_kv: int, mode: str, 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)
@@ -270,7 +272,7 @@ 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)
+    s = score_config(cfg, include_ffn=include_ffn)
     return s if s.fits_memory and s.placement_valid else None
 
 
@@ -308,6 +310,7 @@ def compute_sensitivity(
     baseline_hw: HardwareCandidate,
     parallelism: ConfigScore,
     model: ModelConfig, s_kv: int, mode: str,
+    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
@@ -318,7 +321,7 @@ def compute_sensitivity(
     baseline_cfg = FullConfig(
         model=model, topo=baseline_topo, machine=baseline_machine,
     )
-    baseline_score = score_config(baseline_cfg)
+    baseline_score = score_config(baseline_cfg, include_ffn=include_ffn)
     baseline_latency = baseline_score.total_latency_ns
 
     for knob in _SENSITIVITY_KNOBS:
@@ -327,7 +330,7 @@ def compute_sensitivity(
         doubled_cfg = FullConfig(
             model=model, topo=baseline_topo, machine=doubled_hw.as_machine(),
         )
-        doubled_score = score_config(doubled_cfg)
+        doubled_score = score_config(doubled_cfg, include_ffn=include_ffn)
         rows.append(SensitivityRow(
             knob=knob,
             baseline_value=getattr(baseline_hw, knob),
@@ -348,19 +351,28 @@ def joint_explore(
     s_kv: int,
     mode: str,
     depth: str = "balanced",
+    include_ffn: bool = True,
 ) -> JointExploreResult:
     """Sweep HW candidates × parallelism, return joint Pareto + sensitivity.
 
     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 —
+    both the per-HW parallelism search and the sensitivity ranking use
+    the same restriction.
     """
     all_scores: list[JointScore] = []
     for hw in enumerate_hardware(depth):
         machine = hw.as_machine()
         if depth == "two_stage":
-            par = _best_parallelism_two_stage(model, machine, s_kv, mode)
+            par = _best_parallelism_two_stage(
+                model, machine, s_kv, mode, include_ffn=include_ffn,
+            )
         else:
-            par = _best_parallelism_for_hw(model, machine, s_kv, mode)
+            par = _best_parallelism_for_hw(
+                model, machine, s_kv, mode, include_ffn=include_ffn,
+            )
         if par is None:
             continue
         all_scores.append(JointScore(
@@ -379,6 +391,7 @@ def joint_explore(
         best = all_scores[0]
         sensitivity = compute_sensitivity(
             best.hardware, best.parallelism, model, s_kv, mode,
+            include_ffn=include_ffn,
         )
 
     return JointExploreResult(

tests/analytical_visualization/test_auto_explore.py

@@ -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 (

tests/analytical_visualization/test_auto_hardware.py

@@ -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."""