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:
+    _scope_meta = {
-        with st.spinner("Sweeping ~28k configs (Attention only)..."):
+        "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[f"_auto_explore_result_{_scope}"] = _r
-        st.session_state["_auto_explore_ctx_attn"] = (
+        st.session_state[f"_auto_explore_ctx_{_scope}"] = (
             model.name, s_kv, mode, machine.pe_hbm_gb,
         )
-        st.session_state["_auto_explore_active"] = "attn"
+        st.session_state["_auto_explore_active"] = _scope
-    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 "
+            "frontier. Takes ~5–10 s per scope. All button results are "
-            "you can flip between the two scopes without re-running."
+            "cached so you can flip between scopes without re-running."
         )
         return
 
-    # Downstream code expects _suffix, _label, include_ffn, _res, _ctx_key.
+    # Downstream code expects _suffix, _label, include_attention, include_ffn,
-    include_ffn = (_active == "full")
+    # _res, _ctx_key.
-    _suffix = "_full" if include_ffn else "_attn"
+    include_attention, include_ffn, _label = _scope_meta[_active]
-    _label = "Attn + FFN/MoE" if include_ffn else "Attention only"
+    _suffix = f"_{_active}"
     _result_key = f"_auto_explore_result{_suffix}"
     _ctx_key = f"_auto_explore_ctx{_suffix}"
     _res = st.session_state.get(_result_key)
@@ -1387,206 +1392,191 @@ def _render_auto_explore_tab():
             f"Click **Run sweep — {_label.split(' ')[0]}** to refresh."
         )
 
-        _m1, _m2, _m3, _m4 = st.columns(4)
+    _m1, _m2, _m3, _m4 = st.columns(4)
-        _m1.metric("Enumerated", f"{_res.total_enumerated:,}")
+    _m1.metric("Enumerated", f"{_res.total_enumerated:,}")
-        _m2.metric("Feasible", f"{_res.total_feasible:,}")
+    _m2.metric("Feasible", f"{_res.total_feasible:,}")
-        _m3.metric("Pareto configs", len(_res.pareto_scores))
+    _m3.metric("Pareto configs", len(_res.pareto_scores))
-        if _res.pareto_scores:
+    if _res.pareto_scores:
-            _best = min(_res.pareto_scores, key=lambda s: s.total_latency_ns)
+        _best = min(_res.pareto_scores, key=lambda s: s.total_latency_ns)
-            _m4.metric("Best latency", f"{_best.latency_ms:.2f} ms")
+        _m4.metric("Best latency", f"{_best.latency_ms:.2f} ms")
 
-        if not _res.pareto_scores:
+    if not _res.pareto_scores:
-            st.error(
+        st.error(
-                "No feasible configs — every parallelism choice exceeds "
+            "No feasible configs — every parallelism choice exceeds "
-                "per-PE HBM. Try raising `pe_hbm_gb` in the sidebar or "
+            "per-PE HBM. Try raising `pe_hbm_gb` in the sidebar or "
-                "reducing `s_kv`."
+            "reducing `s_kv`."
-            )
-        else:
-            # ── 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(
-                [s.pes_used for s in _feas],
-                [s.latency_ms for s in _feas],
-                c="lightgrey", s=8, alpha=0.4, label="feasible",
-            )
-            _pareto_sorted = sorted(_res.pareto_scores,
-                                    key=lambda s: s.pes_used)
-            _sc = _axs[0].scatter(
-                [s.pes_used for s in _pareto_sorted],
-                [s.latency_ms for s in _pareto_sorted],
-                c=[s.efficiency_score for s in _pareto_sorted],
-                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],
-                "k--", linewidth=0.8, alpha=0.5,
-            )
-            _axs[0].set_xlabel("PEs used"); _axs[0].set_ylabel("Latency (ms)")
-            _axs[0].set_xscale("log", base=2)
-            _axs[0].set_yscale("log")
-            _axs[0].set_title("Pareto: latency vs PEs")
-            _axs[0].grid(alpha=0.3)
-            _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(
-                [s.hbm_utilization * 100 for s in _pareto_by_lat],
-                [s.latency_ms for s in _pareto_by_lat],
-                c=[s.pes_used for s in _pareto_by_lat],
-                cmap="plasma", s=60, edgecolors="black", linewidths=0.8,
-            )
-            _axs[1].set_xlabel("HBM utilization (%)")
-            _axs[1].set_ylabel("Latency (ms)")
-            _axs[1].set_yscale("log")
-            _axs[1].set_title("HBM usage vs latency (Pareto)")
-            _axs[1].grid(alpha=0.3)
-            _sc2 = _axs[1].collections[0]
-            _fig.colorbar(_sc2, ax=_axs[1], label="PEs")
-
-            _fig.tight_layout()
-            st.pyplot(_fig, width='stretch')
-            _plt.close(_fig)
-
-            # ── Pareto table + Load button ─────────────────────────────
-            st.markdown("**Pareto configurations** (sorted by latency)")
-            _rows = []
-            for i, s in enumerate(_pareto_by_lat):
-                _rows.append({
-                    "#": i,
-                    "CP": s.cp, "TP": s.tp, "PP": s.pp, "DP": s.dp,
-                    "kv": s.kv_shard_mode,
-                    "ffn": s.ffn_shard_scope,
-                    "tp_place": s.tp_placement,
-                    "cp_place": s.cp_placement,
-                    "cp_ring": s.cp_ring_variant,
-                    "lat (ms)": round(s.latency_ms, 3),
-                    "eff": round(s.efficiency_score, 4),
-                    "PEs": s.pes_used,
-                    "SIPs": s.sips_used,
-                    "HBM %": round(s.hbm_utilization * 100, 1),
-                })
-            _df = pd.DataFrame(_rows)
-            st.dataframe(_df, width='stretch', hide_index=True)
-
-            # ── 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."
-            )
-            _sens_c1, _sens_c2 = st.columns([1, 4])
-            with _sens_c1:
-                _sens_row = st.number_input(
-                    "Baseline row #",
-                    min_value=0, max_value=len(_pareto_by_lat) - 1,
-                    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
-            _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
-                        continue
-                    if _fit:
-                        _fit_vals.append(_v)
-                        _fit_lats.append(_lat_ns / 1e6)  # ms
-                    else:
-                        _bad_vals.append((_v, _lat_ns / 1e6))
-                if _fit_vals:
-                    _ax.plot(_fit_vals, _fit_lats, "o-", color="#2c3e50",
-                              linewidth=1.4, markersize=4)
-                if _bad_vals:
-                    _ax.scatter(
-                        [v for v, _ in _bad_vals],
-                        [l for _, l in _bad_vals],
-                        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")
-                _ax.set_yscale("log")
-                _ax.set_title(
-                    f"{_row.knob.upper()}  (baseline={_row.baseline_value})",
-                    fontsize=10,
-                )
-                _ax.set_xlabel("value")
-                _ax.set_ylabel("lat (ms)")
-                _ax.grid(alpha=0.3)
-                if _bad_vals:
-                    _ax.legend(loc="best", fontsize=7)
-            _fig_sens.tight_layout()
-            st.pyplot(_fig_sens, width='stretch')
-            _plt.close(_fig_sens)
-
-            # ── Load-into-sidebar ─────────────────────────────────────
-            st.markdown("**Load a Pareto config into the main sidebar sliders**")
-            _lc1, _lc2 = st.columns([1, 3])
-            with _lc1:
-                _pick = st.number_input(
-                    "Row #", min_value=0, max_value=len(_pareto_by_lat) - 1,
-                    value=0, step=1, key=f"_auto_pick_row{_suffix}",
-                )
-            with _lc2:
-                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
-                    st.session_state["dp"] = _s.dp
-                    st.session_state["tp_placement"] = _s.tp_placement
-                    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)})",
-                        "TP+CP":    f"TP*CP (div={max(1, _tp * _cp)})",
-                        "TP+CP+DP": f"TP*CP*DP (div={max(1, _tp * _cp * _dp)})",
-                    }
-                    st.session_state["ffn_scope_label"] = \
-                        _ffn_label_map[_s.ffn_shard_scope]
-                    st.success(
-                        f"Loaded row {_pick}: CP={_s.cp} TP={_s.tp} PP={_s.pp} "
-                        f"DP={_s.dp} → latency {_s.latency_ms:.2f} ms, "
-                        f"{_s.pes_used} PEs. Flip to another tab to see the "
-                        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."
         )
+        return
+
+    # ── Pareto scatter: latency vs PEs, coloured by efficiency ─
+    import matplotlib.pyplot as _plt
+    _fig, _axs = _plt.subplots(1, 2, figsize=(10, 3.5))
+
+    _feas = [s for s in _res.all_scores
+             if s.fits_memory and s.placement_valid]
+    _axs[0].scatter(
+        [s.pes_used for s in _feas],
+        [s.latency_ms for s in _feas],
+        c="lightgrey", s=8, alpha=0.4, label="feasible",
+    )
+    _pareto_sorted = sorted(_res.pareto_scores,
+                             key=lambda s: s.pes_used)
+    _sc = _axs[0].scatter(
+        [s.pes_used for s in _pareto_sorted],
+        [s.latency_ms for s in _pareto_sorted],
+        c=[s.efficiency_score for s in _pareto_sorted],
+        cmap="viridis", s=60, edgecolors="black", linewidths=0.8,
+        label="Pareto",
+    )
+    _axs[0].plot(
+        [s.pes_used for s in _pareto_sorted],
+        [s.latency_ms for s in _pareto_sorted],
+        "k--", linewidth=0.8, alpha=0.5,
+    )
+    _axs[0].set_xlabel("PEs used"); _axs[0].set_ylabel("Latency (ms)")
+    _axs[0].set_xscale("log", base=2)
+    _axs[0].set_yscale("log")
+    _axs[0].set_title("Pareto: latency vs PEs")
+    _axs[0].grid(alpha=0.3)
+    _axs[0].legend(loc="upper right", fontsize=8)
+    _fig.colorbar(_sc, ax=_axs[0], label="efficiency")
+
+    _pareto_by_lat = sorted(_res.pareto_scores,
+                             key=lambda s: s.total_latency_ns)
+    _axs[1].scatter(
+        [s.hbm_utilization * 100 for s in _pareto_by_lat],
+        [s.latency_ms for s in _pareto_by_lat],
+        c=[s.pes_used for s in _pareto_by_lat],
+        cmap="plasma", s=60, edgecolors="black", linewidths=0.8,
+    )
+    _axs[1].set_xlabel("HBM utilization (%)")
+    _axs[1].set_ylabel("Latency (ms)")
+    _axs[1].set_yscale("log")
+    _axs[1].set_title("HBM usage vs latency (Pareto)")
+    _axs[1].grid(alpha=0.3)
+    _sc2 = _axs[1].collections[0]
+    _fig.colorbar(_sc2, ax=_axs[1], label="PEs")
+
+    _fig.tight_layout()
+    st.pyplot(_fig, width='stretch')
+    _plt.close(_fig)
+
+    # ── Pareto table ────────────────────────────────────────────
+    st.markdown("**Pareto configurations** (sorted by latency)")
+    _rows = []
+    for i, s in enumerate(_pareto_by_lat):
+        _rows.append({
+            "#": i,
+            "CP": s.cp, "TP": s.tp, "PP": s.pp, "DP": s.dp,
+            "kv": s.kv_shard_mode,
+            "ffn": s.ffn_shard_scope,
+            "tp_place": s.tp_placement,
+            "cp_place": s.cp_placement,
+            "cp_ring": s.cp_ring_variant,
+            "lat (ms)": round(s.latency_ms, 3),
+            "eff": round(s.efficiency_score, 4),
+            "PEs": s.pes_used,
+            "SIPs": s.sips_used,
+            "HBM %": round(s.hbm_utilization * 100, 1),
+        })
+    _df = pd.DataFrame(_rows)
+    st.dataframe(_df, width='stretch', hide_index=True)
+
+    # ── Parallelism sensitivity subplots ──────────────────────
+    st.markdown(
+        "**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:
+        _sens_row = st.number_input(
+            "Baseline row #",
+            min_value=0, max_value=len(_pareto_by_lat) - 1,
+            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_attention=include_attention, include_ffn=include_ffn,
+    )
+
+    _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:
+                continue
+            if _fit:
+                _fit_vals.append(_v)
+                _fit_lats.append(_lat_ns / 1e6)
+            else:
+                _bad_vals.append((_v, _lat_ns / 1e6))
+        if _fit_vals:
+            _ax.plot(_fit_vals, _fit_lats, "o-", color="#2c3e50",
+                      linewidth=1.4, markersize=4)
+        if _bad_vals:
+            _ax.scatter(
+                [v for v, _ in _bad_vals],
+                [l for _, l in _bad_vals],
+                marker="x", color="#c0392b", s=40, alpha=0.6,
+                label="no fit",
+            )
+        _ax.axvline(_row.baseline_value, color="#3498db",
+                     linestyle=":", linewidth=1.0)
+        _ax.set_xscale("log")
+        _ax.set_yscale("log")
+        _ax.set_title(
+            f"{_row.knob.upper()}  (baseline={_row.baseline_value})",
+            fontsize=10,
+        )
+        _ax.set_xlabel("value")
+        _ax.set_ylabel("lat (ms)")
+        _ax.grid(alpha=0.3)
+        if _bad_vals:
+            _ax.legend(loc="best", fontsize=7)
+    _fig_sens.tight_layout()
+    st.pyplot(_fig_sens, width='stretch')
+    _plt.close(_fig_sens)
+
+    # ── Load-into-sidebar ───────────────────────────────────────
+    st.markdown("**Load a Pareto config into the main sidebar sliders**")
+    _lc1, _lc2 = st.columns([1, 3])
+    with _lc1:
+        _pick = st.number_input(
+            "Row #", min_value=0, max_value=len(_pareto_by_lat) - 1,
+            value=0, step=1, key=f"_auto_pick_row{_suffix}",
+        )
+    with _lc2:
+        if st.button("Load into sidebar", type="secondary",
+                      key=f"_auto_load_btn{_suffix}"):
+            _s = _pareto_by_lat[int(_pick)]
+            st.session_state["cp"] = _s.cp
+            st.session_state["tp"] = _s.tp
+            st.session_state["pp"] = _s.pp
+            st.session_state["dp"] = _s.dp
+            st.session_state["tp_placement"] = _s.tp_placement
+            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
+            _tp, _cp, _dp = _s.tp, _s.cp, _s.dp
+            _ffn_label_map = {
+                "TP":       f"TP only (div={max(1, _tp)})",
+                "TP+CP":    f"TP*CP (div={max(1, _tp * _cp)})",
+                "TP+CP+DP": f"TP*CP*DP (div={max(1, _tp * _cp * _dp)})",
+            }
+            st.session_state["ffn_scope_label"] = \
+                _ffn_label_map[_s.ffn_shard_scope]
+            st.success(
+                f"Loaded row {_pick}: CP={_s.cp} TP={_s.tp} PP={_s.pp} "
+                f"DP={_s.dp} → latency {_s.latency_ms:.2f} ms, "
+                f"{_s.pes_used} PEs. Flip to another tab to see the "
+                f"full breakdown."
+            )
+            st.rerun()
 
 
 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:
+    _scope_meta_hw = {
-        with st.spinner(f"Sweeping HW × parallelism ({_depth}, Attention only)..."):
+        "attn": (True,  False, "Attention only"),
-            _r = joint_explore(model, s_kv=s_kv, mode=mode,
+        "ffn":  (False, True,  "FFN / MoE only"),
-                                depth=_depth, include_ffn=False)
+        "full": (True,  True,  "Attn + FFN/MoE"),
-        st.session_state["_hw_result_attn"] = _r
+    }
-        st.session_state["_hw_ctx_attn"] = (
+    _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"
+        st.session_state["_hw_active"] = _scope
-    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"
 
     _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 "
+            "co-design for this model + workload. All button results are "
-            "cached so you can flip between the two scopes without re-running."
+            "cached so you can flip between scopes without re-running."
         )
         return
 
-    include_ffn = (_active == "full")
+    include_attention, include_ffn, _label = _scope_meta_hw[_active]
-    _suffix = "_full" if include_ffn else "_attn"
+    _suffix = f"_{_active}"
-    _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)
@@ -1692,21 +1688,22 @@ def _render_auto_hardware_tab():
             f"Click **Run joint sweep — {_label.split(' ')[0]}** to refresh."
         )
 
-        _hm1, _hm2, _hm3, _hm4 = st.columns(4)
+    _hm1, _hm2, _hm3, _hm4 = st.columns(4)
-        _hm1.metric("HW candidates", _hw_res.total_hw)
+    _hm1.metric("HW candidates", _hw_res.total_hw)
-        _hm2.metric("Feasible joint", _hw_res.total_joint)
+    _hm2.metric("Feasible joint", _hw_res.total_joint)
-        _hm3.metric("Pareto configs", len(_hw_res.pareto_scores))
+    _hm3.metric("Pareto configs", len(_hw_res.pareto_scores))
-        if _hw_res.pareto_scores:
+    if _hw_res.pareto_scores:
-            _best = min(_hw_res.pareto_scores,
+        _best = min(_hw_res.pareto_scores,
-                         key=lambda s: s.total_latency_ns)
+                     key=lambda s: s.total_latency_ns)
-            _hm4.metric("Best latency", f"{_best.latency_ms:.2f} ms")
+        _hm4.metric("Best latency", f"{_best.latency_ms:.2f} ms")
 
-        if not _hw_res.pareto_scores:
+    if not _hw_res.pareto_scores:
-            st.error(
+        st.error(
-                "No feasible joint configurations — every HW+parallelism "
+            "No feasible joint configurations — every HW+parallelism "
-                "combo exceeds per-PE HBM at this workload."
+            "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
+    Scope selection (both default True — full per-token cost):
-    isolating attention-kernel tuning from FFN 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
+    Scope flags mirror :func:`_sum_visible_latency`.
-    :func:`_sum_visible_latency` — attention stages only.
     """
     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
+    Scope flags (default: full transformer) restrict *latency + efficiency*
-    only. Memory feasibility is unchanged — still checks weights+KV+transient
+    to attention only, FFN only, or both. Memory feasibility is unchanged —
-    fit in per-PE HBM since the model still exists physically.
+    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)
+    efficiency = (
-                  if latency_s > 0 else 0.0)
+        _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
+    Scope: at least one of ``include_attention`` / ``include_ffn`` must be
-    isolating attention-kernel tuning independent of FFN cost.
+    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
+    Scope flags are forwarded to score_config so the "latency" ranked on
-    matches the caller's attention-only vs full-model choice.
+    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 (