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>

tests/analytical_visualization/app.py

@@ -1324,51 +1324,56 @@ def _render_auto_explore_tab():
         f"per-PE HBM = **{machine.pe_hbm_gb:.1f} GB**"
     )
 
-    _b1, _b2, _b3 = st.columns([1, 1, 2])
+    _b1, _b2, _b3, _b4 = st.columns([1, 1, 1, 1])
     with _b1:
         _run_attn = st.button("Run sweep — Attention",
                                type="primary", width='stretch',
                                key="_auto_run_attn")
     with _b2:
+        _run_ffn = st.button("Run sweep — FFN/MoE",
+                              type="primary", width='stretch',
+                              key="_auto_run_ffn")
+    with _b3:
         _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)..."):
+    _scope_meta = {
+        "attn": (True,  False, "Attention only"),
+        "ffn":  (False, True,  "FFN / MoE only"),
+        "full": (True,  True,  "Attn + FFN/MoE"),
+    }
+    _button_click = {"attn": _run_attn, "ffn": _run_ffn, "full": _run_full}
+
+    for _scope, _clicked in _button_click.items():
+        if not _clicked:
+            continue
+        _ia, _ff, _lbl = _scope_meta[_scope]
+        with st.spinner(f"Sweeping ~28k configs ({_lbl})..."):
             _r = run_auto_explore(
-                model, machine, s_kv=s_kv, mode=mode, include_ffn=False,
+                model, machine, s_kv=s_kv, mode=mode,
+                include_attention=_ia, include_ffn=_ff,
             )
-        st.session_state["_auto_explore_result_attn"] = _r
-        st.session_state["_auto_explore_ctx_attn"] = (
+        st.session_state[f"_auto_explore_result_{_scope}"] = _r
+        st.session_state[f"_auto_explore_ctx_{_scope}"] = (
             model.name, s_kv, mode, machine.pe_hbm_gb,
         )
-        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"
+        st.session_state["_auto_explore_active"] = _scope
 
     _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."
+            "frontier. Takes ~5–10 s per scope. All button results are "
+            "cached so you can flip between 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"
+    # Downstream code expects _suffix, _label, include_attention, include_ffn,
+    # _res, _ctx_key.
+    include_attention, include_ffn, _label = _scope_meta[_active]
+    _suffix = f"_{_active}"
     _result_key = f"_auto_explore_result{_suffix}"
     _ctx_key = f"_auto_explore_ctx{_suffix}"
     _res = st.session_state.get(_result_key)
@@ -1401,12 +1406,12 @@ def _render_auto_explore_tab():
             "per-PE HBM. Try raising `pe_hbm_gb` in the sidebar or "
             "reducing `s_kv`."
         )
-        else:
+        return
+
     # ── Pareto scatter: latency vs PEs, coloured by efficiency ─
     import matplotlib.pyplot as _plt
     _fig, _axs = _plt.subplots(1, 2, figsize=(10, 3.5))
 
-            # All feasible = grey points; Pareto = coloured by efficiency.
     _feas = [s for s in _res.all_scores
              if s.fits_memory and s.placement_valid]
     _axs[0].scatter(
@@ -1423,7 +1428,6 @@ def _render_auto_explore_tab():
         cmap="viridis", s=60, edgecolors="black", linewidths=0.8,
         label="Pareto",
     )
-            # Connect Pareto in PE order to show the trade-off curve.
     _axs[0].plot(
         [s.pes_used for s in _pareto_sorted],
         [s.latency_ms for s in _pareto_sorted],
@@ -1437,7 +1441,6 @@ def _render_auto_explore_tab():
     _axs[0].legend(loc="upper right", fontsize=8)
     _fig.colorbar(_sc, ax=_axs[0], label="efficiency")
 
-            # HBM utilization vs latency for Pareto — where does the config land?
     _pareto_by_lat = sorted(_res.pareto_scores,
                              key=lambda s: s.total_latency_ns)
     _axs[1].scatter(
@@ -1458,7 +1461,7 @@ def _render_auto_explore_tab():
     st.pyplot(_fig, width='stretch')
     _plt.close(_fig)
 
-            # ── Pareto table + Load button ─────────────────────────────
+    # ── Pareto table ────────────────────────────────────────────
     st.markdown("**Pareto configurations** (sorted by latency)")
     _rows = []
     for i, s in enumerate(_pareto_by_lat):
@@ -1481,10 +1484,9 @@ def _render_auto_explore_tab():
 
     # ── Parallelism sensitivity subplots ──────────────────────
     st.markdown(
-                "**Parallelism sensitivity** — for the *baseline* config "
-                "picked below, vary each knob one at a time (holding others "
-                "fixed). Answers 'if I only tweak CP / TP / PP / DP / EP, "
-                "what happens?' Solid line = fits memory; red X = infeasible."
+        "**Parallelism sensitivity** — for the *baseline* config picked "
+        "below, vary each knob one at a time (holding others fixed). "
+        "Solid line = fits memory; red X = infeasible."
     )
     _sens_c1, _sens_c2 = st.columns([1, 4])
     with _sens_c1:
@@ -1497,21 +1499,20 @@ def _render_auto_explore_tab():
     _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,
+        include_attention=include_attention, include_ffn=include_ffn,
     )
 
-            import matplotlib.pyplot as _plt
     _fig_sens, _axes_sens = _plt.subplots(1, 5, figsize=(14, 3.2))
     for _ax, _row in zip(_axes_sens, _sens_rows):
         _fit_vals, _fit_lats = [], []
         _bad_vals = []
         for _v, _lat_ns, _fit in zip(_row.values, _row.latencies_ns,
                                       _row.fits_flags):
-                    if _lat_ns != _lat_ns:  # NaN — skip out-of-bounds
+            if _lat_ns != _lat_ns:
                 continue
             if _fit:
                 _fit_vals.append(_v)
-                        _fit_lats.append(_lat_ns / 1e6)  # ms
+                _fit_lats.append(_lat_ns / 1e6)
             else:
                 _bad_vals.append((_v, _lat_ns / 1e6))
         if _fit_vals:
@@ -1524,7 +1525,6 @@ def _render_auto_explore_tab():
                 marker="x", color="#c0392b", s=40, alpha=0.6,
                 label="no fit",
             )
-                # Baseline marker
         _ax.axvline(_row.baseline_value, color="#3498db",
                      linestyle=":", linewidth=1.0)
         _ax.set_xscale("log")
@@ -1542,7 +1542,7 @@ def _render_auto_explore_tab():
     st.pyplot(_fig_sens, width='stretch')
     _plt.close(_fig_sens)
 
-            # ── Load-into-sidebar ─────────────────────────────────────
+    # ── Load-into-sidebar ───────────────────────────────────────
     st.markdown("**Load a Pareto config into the main sidebar sliders**")
     _lc1, _lc2 = st.columns([1, 3])
     with _lc1:
@@ -1554,8 +1554,6 @@ def _render_auto_explore_tab():
         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.
             st.session_state["cp"] = _s.cp
             st.session_state["tp"] = _s.tp
             st.session_state["pp"] = _s.pp
@@ -1564,8 +1562,6 @@ def _render_auto_explore_tab():
             st.session_state["cp_placement"] = _s.cp_placement
             st.session_state["cp_ring_variant"] = _s.cp_ring_variant
             st.session_state["kv_mode"] = _s.kv_shard_mode
-                    # ffn_scope label is dynamic: "TP+CP (div=…)" etc.
-                    # Reconstruct with the target cp/tp/dp values.
             _tp, _cp, _dp = _s.tp, _s.cp, _s.dp
             _ffn_label_map = {
                 "TP":       f"TP only (div={max(1, _tp)})",
@@ -1581,12 +1577,6 @@ def _render_auto_explore_tab():
                 f"full breakdown."
             )
             st.rerun()
-    else:
-        st.info(
-            "Click **Run sweep** to enumerate all valid parallelism "
-            "configurations and rank them on the Pareto frontier. "
-            "Takes ~5–10 s for the full search."
-        )
 
 
 with tab_auto:
@@ -1634,47 +1624,53 @@ def _render_auto_hardware_tab():
                   "coarse: 729 HW × ~2k parallelism, ~2-5 min."),
         )
 
-    _bh1, _bh2, _bh3 = st.columns([1, 1, 2])
+    _bh1, _bh2, _bh3, _bh4 = st.columns([1, 1, 1, 1])
     with _bh1:
         _run_hw_attn = st.button("Run joint sweep — Attention",
                                    type="primary", width='stretch',
                                    key="_hw_run_attn")
     with _bh2:
+        _run_hw_ffn = st.button("Run joint sweep — FFN/MoE",
+                                  type="primary", width='stretch',
+                                  key="_hw_run_ffn")
+    with _bh3:
         _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"] = (
+    _scope_meta_hw = {
+        "attn": (True,  False, "Attention only"),
+        "ffn":  (False, True,  "FFN / MoE only"),
+        "full": (True,  True,  "Attn + FFN/MoE"),
+    }
+    _button_click_hw = {
+        "attn": _run_hw_attn, "ffn": _run_hw_ffn, "full": _run_hw_full,
+    }
+
+    for _scope, _clicked in _button_click_hw.items():
+        if not _clicked:
+            continue
+        _ia, _ff, _lbl = _scope_meta_hw[_scope]
+        with st.spinner(f"Sweeping HW × parallelism ({_depth}, {_lbl})..."):
+            _r = joint_explore(model, s_kv=s_kv, mode=mode, depth=_depth,
+                                include_attention=_ia, include_ffn=_ff)
+        st.session_state[f"_hw_result_{_scope}"] = _r
+        st.session_state[f"_hw_ctx_{_scope}"] = (
             model.name, s_kv, mode, _depth,
         )
-        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"
+        st.session_state["_hw_active"] = _scope
 
     _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."
+            "co-design for this model + workload. All button results are "
+            "cached so you can flip between 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"
+    include_attention, include_ffn, _label = _scope_meta_hw[_active]
+    _suffix = f"_{_active}"
     _hw_result_key = f"_hw_result{_suffix}"
     _hw_ctx_key = f"_hw_ctx{_suffix}"
     _hw_res = st.session_state.get(_hw_result_key)
@@ -1706,7 +1702,8 @@ def _render_auto_hardware_tab():
             "No feasible joint configurations — every HW+parallelism "
             "combo exceeds per-PE HBM at this workload."
         )
-        else:
+        return
+    if True:
             # ── Panel 1: latency vs cost Pareto ────────────────────────
             import matplotlib.pyplot as _plt
             _fig1, _ax1 = _plt.subplots(figsize=(8, 4))
@@ -1855,11 +1852,6 @@ def _render_auto_hardware_tab():
                         f"full breakdown."
                     )
                     st.rerun()
-    else:
-        st.info(
-            "Click **Run joint sweep** to explore hardware and parallelism "
-            "co-design space for this model + workload."
-        )
 
 
 with tab_hw:

tests/analytical_visualization/auto_explore.py

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

tests/analytical_visualization/auto_hardware.py

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

tests/analytical_visualization/test_auto_explore.py

@@ -183,11 +183,31 @@ def test_attention_only_pareto_non_empty():
     model = PRESETS["Llama 3.1 70B"].model
     machine = MachineParams()
     res = run_auto_explore(model, machine, s_kv=8192, mode="decode",
-                             include_ffn=False)
+                             include_attention=True, include_ffn=False)
     assert res.total_feasible > 0
     assert len(res.pareto_scores) > 0
 
 
+def test_ffn_only_latency_less_than_full_and_different_from_attn():
+    """FFN-only latency must be > 0, < full-model latency, and != attn-only."""
+    model = PRESETS["Llama 3.1 70B"].model
+    machine = MachineParams()
+    r_full = run_auto_explore(model, machine, s_kv=8192, mode="decode",
+                                include_attention=True, include_ffn=True)
+    r_attn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
+                                include_attention=True, include_ffn=False)
+    r_ffn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
+                               include_attention=False, include_ffn=True)
+    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)
+    b_ffn = min(r_ffn.pareto_scores, key=lambda s: s.total_latency_ns)
+    assert b_ffn.total_latency_ns > 0
+    assert b_ffn.total_latency_ns < b_full.total_latency_ns
+    assert abs(b_ffn.total_latency_ns - b_attn.total_latency_ns) > 1.0, (
+        "FFN-only and attention-only latencies should not be identical"
+    )
+
+
 def test_parallelism_sensitivity_includes_baseline_value():
     """Each row's values contain the baseline_value."""
     from tests.analytical_visualization.auto_explore import (