analytical-viz: parallelism sensitivity chart in Auto Explore tab

Adds compute_parallelism_sensitivity() + ParallelismSensitivityRow to
auto_explore.py. For a baseline ConfigScore (picked from the Pareto set),
sweeps each parallelism knob (CP, TP, PP, DP, EP) individually while
holding others fixed. Reports latency + memory-fit per value.

Sweep values start at 1 and step by 2 (multiples of 2 rather than only
powers of 2), giving finer granularity for the visual than the enumerator's
sparser set. CP goes up to 256 (per real-deployment scale), TP to 64,
PP to 32.

New UI panel in Auto Explore tab: 5 subplots (log/log), one per knob:
  - Solid line = latency where the config fits memory
  - Red X markers = infeasible (out of budget)
  - Dotted vertical line = baseline value
User picks which Pareto row is the "baseline" via a number input; the
sensitivity chart re-computes around it.

Behaviour caveat noted during verification: for large PP the FFN AR can
cross a SIP boundary (uses stage_latencies.py:730 sips_used tier
selection), producing a step-up in latency. This is the existing model's
choice, honestly reflected in the chart. Whether the FFN AR should span
only one PP stage's ranks (thus not cross SIPs) is a separate discussion
about stage_latencies.py.

Verified:
- 11 pytest tests pass (added 2 new for parallelism sensitivity)
- Smoke: at Llama 70B decode 128K with baseline CP=8/TP=16/PP=1/DP=1:
  * CP sweep: 4 fits, 8 = baseline optimum, 16+ overshoots memory
  * TP sweep: 8/16 fit, 16 = baseline optimum, 32 slower (spans SIPs)
  * PP sweep: 1 = baseline, 2+ slower due to sips_used tier drop for FFN AR
  * DP sweep: same shape as PP
  * EP sweep: monotone decreasing (bigger EP = smaller per-PE FFN)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

tests/analytical_visualization/app.py

@@ -1286,7 +1286,10 @@ with tab_compare:
 
 # ── TAB 5: Auto Explore ─────────────────────────────────────────
 with tab_auto:
-    from tests.analytical_visualization.auto_explore import run_auto_explore
+    from tests.analytical_visualization.auto_explore import (
+        compute_parallelism_sensitivity,
+        run_auto_explore,
+    )
 
     st.markdown(
         "Sweep every valid parallelism config (CP, TP, PP, DP, kv_shard, "
@@ -1416,6 +1419,68 @@ with tab_auto:
             _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="_auto_sens_row",
+                    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,
+            )
+
+            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])

tests/analytical_visualization/auto_explore.py

@@ -25,6 +25,17 @@ from .memory_layout import compute_memory
 from .model_config import FullConfig, MachineParams, ModelConfig, TopologyConfig
 from .stage_latencies import all_ffn_stages, all_stages
 
+# ── Parallelism sensitivity sweep values (see compute_parallelism_sensitivity)
+# Multiples of 2 (finer grid than powers of 2 alone), capped at values that
+# are physically meaningful for the modelled topology.
+_PARALLELISM_SWEEP_VALUES = {
+    "cp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64, 96, 128, 192, 256),
+    "tp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64),
+    "pp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 32),
+    "dp": (1, 2, 4, 6, 8, 12, 16),
+    "ep": (1, 2, 4, 6, 8, 12, 16, 32),
+}
+
 
 # ── Search space ─────────────────────────────────────────────────────
 
@@ -298,6 +309,75 @@ def pareto_frontier(scores: list[ConfigScore]) -> list[ConfigScore]:
     return frontier
 
 
+# ── Parallelism sensitivity ─────────────────────────────────────────
+
+
+@dataclass
+class ParallelismSensitivityRow:
+    """One knob's sweep result: latency + fit at every tested value.
+
+    All *other* parallelism knobs (including HW) are held at their baseline
+    values, so this isolates the effect of just this one knob.
+    """
+    knob: str                       # "cp" / "tp" / "pp" / "dp" / "ep"
+    values: list[int]               # tested values (from _PARALLELISM_SWEEP_VALUES)
+    latencies_ns: list[float]       # one per value; NaN when infeasible
+    fits_flags: list[bool]          # per-value memory+placement feasibility
+    baseline_value: int             # what the baseline config uses for this knob
+    baseline_latency_ns: float
+
+
+def compute_parallelism_sensitivity(
+    baseline: ConfigScore,
+    model: ModelConfig,
+    machine: MachineParams,
+    s_kv: int,
+    mode: str,
+) -> 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
+    fit per value.
+
+    Answers: 'if I only change this one knob, how does latency and memory
+    fit change?' — the complement to auto_hardware.compute_sensitivity
+    which does the same for hardware knobs.
+    """
+    baseline_topo = baseline.as_topology(s_kv, mode)
+
+    rows: list[ParallelismSensitivityRow] = []
+    for knob, values in _PARALLELISM_SWEEP_VALUES.items():
+        baseline_val = getattr(baseline_topo, knob, 1)
+        # EP isn't stored on ConfigScore's as_topology output today (dp
+        # attribute exists but ep does not always). Fall back to 1.
+        latencies: list[float] = []
+        fits: list[bool] = []
+        for v in values:
+            # Skip nonsensical values that would violate hard bounds.
+            if knob == "pp" and v > model.layers:
+                latencies.append(float("nan"))
+                fits.append(False)
+                continue
+            if knob == "tp" and v > 4 * model.h_q:
+                latencies.append(float("nan"))
+                fits.append(False)
+                continue
+            # 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)
+            latencies.append(score.total_latency_ns)
+            fits.append(score.fits_memory and score.placement_valid)
+        rows.append(ParallelismSensitivityRow(
+            knob=knob,
+            values=list(values),
+            latencies_ns=latencies,
+            fits_flags=fits,
+            baseline_value=int(baseline_val),
+            baseline_latency_ns=baseline.total_latency_ns,
+        ))
+    return rows
+
+
 # ── Top-level driver ─────────────────────────────────────────────────
 
 

tests/analytical_visualization/test_auto_explore.py

@@ -140,3 +140,41 @@ def test_pp_does_not_reduce_single_request_latency():
     assert fastest.pp == 1, (
         f"expected PP=1 for the fastest decode config; got PP={fastest.pp}"
     )
+
+
+# ── Parallelism sensitivity ─────────────────────────────────────────
+
+
+def test_parallelism_sensitivity_has_all_five_knobs():
+    """compute_parallelism_sensitivity returns one row per knob."""
+    from tests.analytical_visualization.auto_explore import (
+        compute_parallelism_sensitivity,
+    )
+    model = PRESETS["Llama 3.1 70B"].model
+    machine = MachineParams()
+    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
+    baseline = res.pareto_scores[0]
+    rows = compute_parallelism_sensitivity(
+        baseline, model, machine, s_kv=8192, mode="decode",
+    )
+    knobs = {r.knob for r in rows}
+    assert knobs == {"cp", "tp", "pp", "dp", "ep"}
+
+
+def test_parallelism_sensitivity_includes_baseline_value():
+    """Each row's values contain the baseline_value."""
+    from tests.analytical_visualization.auto_explore import (
+        compute_parallelism_sensitivity,
+    )
+    model = PRESETS["Llama 3.1 70B"].model
+    machine = MachineParams()
+    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
+    baseline = res.pareto_scores[0]
+    rows = compute_parallelism_sensitivity(
+        baseline, model, machine, s_kv=8192, mode="decode",
+    )
+    for r in rows:
+        # Baseline value must be sweepable OR mapped to something in values.
+        assert r.baseline_value in r.values or r.baseline_value == 1, (
+            f"knob {r.knob}: baseline_value={r.baseline_value} not in {r.values}"
+        )