analytical-viz: auto_explore.py — 9-knob Pareto search

Extends the memory-only autosuggest to search the full 9-dimensional
parallelism space (CP, TP, PP, DP, kv_shard_mode, ffn_shard_scope,
tp_placement, cp_placement, cp_ring_variant) and rank feasible configs
on a 3D Pareto frontier: (latency ↓, pes_used ↓, efficiency ↑).

Throughput is stored on ConfigScore for display but is deliberately NOT
a Pareto axis because for a single-request analysis it collapses to
1 / latency, which would collapse the frontier.

Reuses existing physics (stage_latencies.all_stages + all_ffn_stages,
memory_layout.compute_memory) — no new formulas.

Single-request latency formula fix: PP does NOT reduce single-request
decode/prefill latency because the request has to traverse every layer
sequentially regardless of pipeline depth. The initial version had
latency ~ per_layer × layers_per_stage, which incorrectly rewarded high
PP. Corrected to latency ~ per_layer × model.layers.

Enumerator prunes:
  - PP > model.layers
  - TP > 4 × h_q
  - ffn_shard_scope contains 'DP' when dp=1 (redundant)
  - cp_ring_variant='qoml' when cp=1 (no-op)

Full sweep on Llama 3.1 70B: ~28,800 configs enumerated in ~7s, ~7k-10k
feasible (varies with S_kv), 2-7 unique Pareto configs. Faster context
lengths produce richer frontiers; at 1M, memory forces a single
dominant config (128 PEs, HBM 85%).

Verified:
- 9 pytest tests pass (enumerate, score, Pareto, subset invariants)
- Manual: Llama 70B decode at 8K/64K/128K/1M produces physically
  sensible Pareto (CP=8/TP=16 wins latency; smaller-PE options
  appear at longer context up to memory limits)

Next: Streamlit tab UI in app.py + verification against more presets.

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

tests/analytical_visualization/auto_explore.py

@@ -0,0 +1,337 @@
+"""Auto-explore parallelism configuration space and rank by Pareto frontier.
+
+Extends the existing ``autosuggest`` (memory-only, single winner) to the
+full 9-knob search (CP, TP, PP, DP, kv_shard_mode, ffn_shard_scope,
+tp_placement, cp_placement, cp_ring_variant) with four objectives:
+
+    - min total_latency_ns   (single-request decode/prefill step)
+    - max throughput_tok_s   (naive tokens/sec = 1 / latency for single-req)
+    - max efficiency_score   (geo-mean of compute + BW utilization)
+    - min pes_used
+
+Reuses ``stage_latencies`` and ``memory_layout`` as the physics; adds
+enumeration + 4D Pareto sort on top.
+
+Analytical model runs in ~microseconds per config, so the full sweep
+(~5-10k feasible configs after pruning) completes in a few seconds.
+"""
+from __future__ import annotations
+
+import math
+from collections.abc import Iterator
+from dataclasses import dataclass, field, replace
+
+from .memory_layout import compute_memory
+from .model_config import FullConfig, MachineParams, ModelConfig, TopologyConfig
+from .stage_latencies import all_ffn_stages, all_stages
+
+
+# ── Search space ─────────────────────────────────────────────────────
+
+_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
+_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
+_PP_OPTIONS = (1, 2, 4, 8, 16)
+_DP_OPTIONS = (1, 2, 4)
+_KV_SHARD_MODES = ("split", "replicate")
+_FFN_SHARD_SCOPES = ("TP", "TP+CP", "TP+CP+DP")
+_TP_PLACEMENTS = ("pe", "cube")
+_CP_PLACEMENTS = ("cube", "pe")
+_CP_RING_VARIANTS = ("kv", "qoml")
+
+
+# ── Result types ─────────────────────────────────────────────────────
+
+
+@dataclass
+class ConfigScore:
+    """A single (config, computed-metrics) tuple. All floats in SI units
+    (seconds, tokens/sec, dimensionless) unless suffixed otherwise."""
+
+    # ── Config (the 9 knobs) ─────
+    cp: int
+    tp: int
+    pp: int
+    dp: int
+    kv_shard_mode: str
+    ffn_shard_scope: str
+    tp_placement: str
+    cp_placement: str
+    cp_ring_variant: str
+
+    # ── Objectives ─────
+    total_latency_ns: float      # ↓ minimize
+    throughput_tok_s: float       # ↑ maximize
+    efficiency_score: float       # ↑ maximize (0..1)
+    pes_used: int                 # ↓ minimize
+
+    # ── Info-only (not ranked on) ─────
+    hbm_utilization: float        # bytes_used / budget (0..1+; may exceed 1 if over-budget)
+    weights_gb: float
+    kv_gb: float
+    transient_gb: float
+    sips_used: int
+    fits_memory: bool
+    placement_valid: bool
+    reason: str = ""
+
+    @property
+    def latency_us(self) -> float:
+        return self.total_latency_ns / 1e3
+
+    @property
+    def latency_ms(self) -> float:
+        return self.total_latency_ns / 1e6
+
+    def as_topology(self, s_kv: int, mode: str) -> TopologyConfig:
+        """Reconstruct the TopologyConfig this score was computed for."""
+        return TopologyConfig(
+            cp=self.cp, tp=self.tp, pp=self.pp, dp=self.dp,
+            s_kv=s_kv, mode=mode,
+            kv_shard_mode=self.kv_shard_mode,
+            ffn_shard_scope=self.ffn_shard_scope,
+            tp_placement=self.tp_placement,
+            cp_placement=self.cp_placement,
+            cp_ring_variant=self.cp_ring_variant,
+        )
+
+
+@dataclass
+class AutoExploreResult:
+    model_name: str
+    s_kv: int
+    mode: str
+    total_enumerated: int         # after basic domain pruning
+    total_feasible: int           # after memory + placement checks
+    all_scores: list[ConfigScore] = field(default_factory=list)  # every feasible config
+    pareto_scores: list[ConfigScore] = field(default_factory=list)  # non-dominated set
+
+
+# ── Enumeration + pruning ────────────────────────────────────────────
+
+
+def enumerate_configs(
+    model: ModelConfig,
+    s_kv: int,
+    mode: str,
+) -> Iterator[TopologyConfig]:
+    """Yield every domain-valid TopologyConfig.
+
+    Domain rules (fast pruning; no memory check yet):
+      - PP ≤ model.layers                            (can't have more stages than layers)
+      - TP ≤ 4 × model.h_q                            (unrealistic head-dim splits above this)
+      - Skip ffn_shard_scope containing 'DP' when DP=1 (redundant with plain 'TP+CP')
+      - Skip cp_ring_variant='qoml' when CP=1        (no ring, variant is a no-op)
+    """
+    for cp in _CP_OPTIONS:
+        for tp in _TP_OPTIONS:
+            if tp > 4 * model.h_q:
+                continue
+            for pp in _PP_OPTIONS:
+                if pp > model.layers:
+                    continue
+                for dp in _DP_OPTIONS:
+                    for kv_mode in _KV_SHARD_MODES:
+                        for ffn_scope in _FFN_SHARD_SCOPES:
+                            if "DP" in ffn_scope and dp == 1:
+                                continue
+                            for tp_place in _TP_PLACEMENTS:
+                                for cp_place in _CP_PLACEMENTS:
+                                    for cp_ring in _CP_RING_VARIANTS:
+                                        if cp == 1 and cp_ring == "qoml":
+                                            continue
+                                        yield TopologyConfig(
+                                            cp=cp, tp=tp, pp=pp, dp=dp,
+                                            s_kv=s_kv, mode=mode,
+                                            kv_shard_mode=kv_mode,
+                                            ffn_shard_scope=ffn_scope,
+                                            tp_placement=tp_place,
+                                            cp_placement=cp_place,
+                                            cp_ring_variant=cp_ring,
+                                        )
+
+
+# ── Scoring ──────────────────────────────────────────────────────────
+
+
+def _sum_visible_latency(cfg: FullConfig) -> float:
+    """Total single-request latency (seconds) across all model layers.
+
+    A single request traverses every layer sequentially, whether the layers
+    sit on one PP stage or are spread across many:
+      - PP=1  : one rank holds all L layers, latency = L × per_layer
+      - PP=K  : K ranks each hold L/K layers, but request crosses all K
+                stages sequentially → still L × per_layer
+
+    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.
+    """
+    attn = sum(s.visible_s for s in all_stages(cfg))
+    ffn = sum(s.visible_s for s in all_ffn_stages(cfg))
+    per_layer = attn + ffn
+    return per_layer * cfg.model.layers
+
+
+def _efficiency(cfg: FullConfig, latency_s: float) -> 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)
+    """
+    if latency_s <= 0:
+        return 0.0
+    attn = all_stages(cfg)
+    ffn = all_ffn_stages(cfg)
+    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)
+
+    pes = cfg.topo.total_pes
+    peak_flops = cfg.machine.peak_flops * pes
+    peak_bw = cfg.machine.bw_hbm * pes
+
+    compute_util = total_flops / (peak_flops * latency_s) if peak_flops > 0 else 0.0
+    bw_util = total_bytes / (peak_bw * latency_s) if peak_bw > 0 else 0.0
+    compute_util = min(1.0, max(0.0, compute_util))
+    bw_util = min(1.0, max(0.0, bw_util))
+    # Geo-mean; if either is 0 the score is 0 (avoids overrewarding lopsided configs).
+    return math.sqrt(compute_util * bw_util)
+
+
+def score_config(cfg: FullConfig) -> 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.
+    """
+    mem = compute_memory(cfg)
+    placement_ok = cfg.topo.placement_valid
+
+    latency_s = _sum_visible_latency(cfg)
+    throughput = 1.0 / latency_s if latency_s > 0 else 0.0
+    efficiency = _efficiency(cfg, latency_s) if latency_s > 0 else 0.0
+    fits = not mem.over_budget
+
+    reason = ""
+    if not fits:
+        reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
+                  f"exceeds per-PE budget ({mem.budget_bytes/1e9:.2f} GB)")
+    elif not placement_ok:
+        reason = (f"intra-cube demand ({cfg.topo.intra_cube_dims}) "
+                  f"exceeds PEs/cube ({cfg.topo.pes_per_cube_hw})")
+
+    return ConfigScore(
+        cp=cfg.topo.cp, tp=cfg.topo.tp, pp=cfg.topo.pp, dp=cfg.topo.dp,
+        kv_shard_mode=cfg.topo.kv_shard_mode,
+        ffn_shard_scope=cfg.topo.ffn_shard_scope,
+        tp_placement=cfg.topo.tp_placement,
+        cp_placement=cfg.topo.cp_placement,
+        cp_ring_variant=cfg.topo.cp_ring_variant,
+        total_latency_ns=latency_s * 1e9,
+        throughput_tok_s=throughput,
+        efficiency_score=efficiency,
+        pes_used=cfg.topo.total_pes,
+        hbm_utilization=mem.used_bytes / mem.budget_bytes if mem.budget_bytes > 0 else 0.0,
+        weights_gb=mem.weights_bytes / 1e9,
+        kv_gb=mem.kv_cache_bytes / 1e9,
+        transient_gb=mem.transient_bytes / 1e9,
+        sips_used=cfg.topo.sips_used,
+        fits_memory=fits,
+        placement_valid=placement_ok,
+        reason=reason,
+    )
+
+
+# ── Pareto sort ──────────────────────────────────────────────────────
+
+
+def _dominates(a: ConfigScore, b: ConfigScore) -> bool:
+    """True iff a is no worse than b on every axis AND strictly better on ≥ 1.
+
+    Axes (with direction) — 3D Pareto:
+      total_latency_ns    ↓  (a ≤ b)
+      pes_used            ↓  (a ≤ b)  — proxy for cost / deployment size
+      efficiency_score    ↑  (a ≥ b)  — geo-mean of compute+BW utilization
+
+    Note: ``throughput_tok_s`` is deliberately NOT a Pareto axis. For a
+    single-request analysis, throughput = 1 / latency, so it's collinear
+    with latency and would collapse the frontier. It stays on ConfigScore
+    for display but doesn't participate in domination.
+    """
+    no_worse = (
+        a.total_latency_ns <= b.total_latency_ns
+        and a.pes_used <= b.pes_used
+        and a.efficiency_score >= b.efficiency_score
+    )
+    if not no_worse:
+        return False
+    strictly_better = (
+        a.total_latency_ns < b.total_latency_ns
+        or a.pes_used < b.pes_used
+        or a.efficiency_score > b.efficiency_score
+    )
+    return strictly_better
+
+
+def pareto_frontier(scores: list[ConfigScore]) -> list[ConfigScore]:
+    """Extract non-dominated set on (latency↓, pe↓, throughput↑, efficiency↑).
+
+    Only feasible configs (fits_memory AND placement_valid) participate; the
+    non-feasible are excluded from the frontier (they can still be returned in
+    ``all_scores`` for informational display).
+
+    O(N²) — fine for N up to ~50k with early exits.
+    """
+    feasible = [s for s in scores if s.fits_memory and s.placement_valid]
+    frontier: list[ConfigScore] = []
+    for i, a in enumerate(feasible):
+        dominated = False
+        for j, b in enumerate(feasible):
+            if i == j:
+                continue
+            if _dominates(b, a):
+                dominated = True
+                break
+        if not dominated:
+            frontier.append(a)
+    return frontier
+
+
+# ── Top-level driver ─────────────────────────────────────────────────
+
+
+def run_auto_explore(
+    model: ModelConfig,
+    machine: MachineParams,
+    s_kv: int,
+    mode: str = "decode",
+) -> 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.
+    """
+    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.sort(key=lambda s: s.total_latency_ns)
+    pareto = pareto_frontier(all_scores)
+    pareto.sort(key=lambda s: s.total_latency_ns)
+
+    feasible_count = sum(1 for s in all_scores if s.fits_memory and s.placement_valid)
+
+    return AutoExploreResult(
+        model_name=model.name,
+        s_kv=s_kv,
+        mode=mode,
+        total_enumerated=total_enumerated,
+        total_feasible=feasible_count,
+        all_scores=all_scores,
+        pareto_scores=pareto,
+    )

tests/analytical_visualization/test_auto_explore.py

@@ -0,0 +1,142 @@
+"""Interface + smoke tests for auto_explore.
+
+Covers:
+  - enumerate_configs yields valid TopologyConfigs with expected pruning
+  - score_config returns a ConfigScore with all fields populated
+  - pareto_frontier is a subset of feasible and is non-empty for a
+    reasonable model+workload
+  - run_auto_explore returns a coherent AutoExploreResult (all_scores
+    sorted by latency, pareto ⊆ feasible ⊆ all_scores)
+"""
+from __future__ import annotations
+
+import pytest
+
+from tests.analytical_visualization.auto_explore import (
+    ConfigScore,
+    enumerate_configs,
+    pareto_frontier,
+    run_auto_explore,
+    score_config,
+)
+from tests.analytical_visualization.model_config import (
+    FullConfig, MachineParams,
+)
+from tests.analytical_visualization.model_presets import PRESETS
+
+
+# ── Enumeration ─────────────────────────────────────────────────────
+
+
+def test_enumerate_yields_valid_configs():
+    """Enumerator produces TopologyConfigs with all 9 knobs set."""
+    model = PRESETS["Llama 3.1 70B"].model
+    it = enumerate_configs(model, s_kv=8192, mode="decode")
+    first = next(it)
+    for f in ("cp", "tp", "pp", "dp", "kv_shard_mode",
+              "ffn_shard_scope", "tp_placement", "cp_placement",
+              "cp_ring_variant"):
+        assert getattr(first, f) is not None, f"knob {f} not set"
+    assert first.s_kv == 8192
+    assert first.mode == "decode"
+
+
+def test_enumerate_prunes_pp_beyond_layers():
+    """PP > model.layers is skipped by the enumerator."""
+    model = PRESETS["Llama 3.1 70B"].model  # layers=80
+    for cfg in enumerate_configs(model, s_kv=8192, mode="decode"):
+        assert cfg.pp <= model.layers
+
+
+def test_enumerate_prunes_ffn_dp_when_dp_is_1():
+    """ffn_shard_scope containing 'DP' is skipped when dp=1 (redundant)."""
+    model = PRESETS["Llama 3.1 70B"].model
+    for cfg in enumerate_configs(model, s_kv=8192, mode="decode"):
+        if cfg.dp == 1:
+            assert "DP" not in cfg.ffn_shard_scope
+
+
+# ── Scoring ─────────────────────────────────────────────────────────
+
+
+def test_score_returns_populated_config_score():
+    """score_config returns a fully-populated ConfigScore."""
+    model = PRESETS["Llama 3.1 70B"].model
+    topo = next(enumerate_configs(model, s_kv=8192, mode="decode"))
+    machine = MachineParams()
+    cfg = FullConfig(model=model, topo=topo, machine=machine)
+    score = score_config(cfg)
+    assert isinstance(score, ConfigScore)
+    assert score.total_latency_ns > 0
+    assert score.pes_used == topo.total_pes
+    assert 0.0 <= score.efficiency_score <= 1.0
+    assert score.hbm_utilization >= 0.0
+
+
+# ── Pareto ──────────────────────────────────────────────────────────
+
+
+def test_pareto_subset_of_feasible():
+    """Pareto scores are always feasible (memory + placement)."""
+    model = PRESETS["Llama 3.1 70B"].model
+    machine = MachineParams()
+    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
+    for p in res.pareto_scores:
+        assert p.fits_memory
+        assert p.placement_valid
+
+
+def test_pareto_non_empty_when_feasible_configs_exist():
+    """When at least one config fits, Pareto must have ≥ 1 entry."""
+    model = PRESETS["Llama 3.1 70B"].model
+    machine = MachineParams()
+    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
+    assert res.total_feasible > 0
+    assert len(res.pareto_scores) > 0
+
+
+def test_pareto_frontier_non_dominated():
+    """No Pareto entry is dominated by another."""
+    from tests.analytical_visualization.auto_explore import _dominates
+    model = PRESETS["Llama 3.1 70B"].model
+    machine = MachineParams()
+    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
+    for i, a in enumerate(res.pareto_scores):
+        for j, b in enumerate(res.pareto_scores):
+            if i == j:
+                continue
+            assert not _dominates(b, a), (
+                f"Pareto entry {i} is dominated by entry {j}"
+            )
+
+
+# ── End-to-end sanity ───────────────────────────────────────────────
+
+
+def test_run_auto_explore_shape():
+    """all_scores sorted asc by latency; pareto ⊆ feasible ⊆ all_scores."""
+    model = PRESETS["Llama 3.1 70B"].model
+    machine = MachineParams()
+    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
+
+    assert res.total_enumerated == len(res.all_scores)
+    assert res.total_feasible == sum(
+        1 for s in res.all_scores if s.fits_memory and s.placement_valid
+    )
+    assert len(res.pareto_scores) <= res.total_feasible
+
+    latencies = [s.total_latency_ns for s in res.all_scores]
+    assert latencies == sorted(latencies)
+
+
+def test_pp_does_not_reduce_single_request_latency():
+    """A single request traverses all layers regardless of PP.
+    Latency-optimal Pareto configs should NOT prefer PP>1 for decode."""
+    model = PRESETS["Llama 3.1 70B"].model
+    machine = MachineParams()
+    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
+    # Sort Pareto by latency; the fastest should not need PP>1 to win.
+    fastest = res.pareto_scores[0]
+    assert fastest.pp == 1, (
+        f"expected PP=1 for the fastest decode config; got PP={fastest.pp}"
+    )