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>
2026-07-28 12:58:10 -07:00
parent 09e9331026
commit 85f6716fc1
2 changed files with 479 additions and 0 deletions
@@ -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}"
+    )