From 85f6716fc10a9410cda9f8b97d1571c8f0601c1c Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Tue, 28 Jul 2026 12:58:10 -0700 Subject: [PATCH] =?UTF-8?q?analytical-viz:=20auto=5Fexplore.py=20=E2=80=94?= =?UTF-8?q?=209-knob=20Pareto=20search?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../analytical_visualization/auto_explore.py | 337 ++++++++++++++++++ .../test_auto_explore.py | 142 ++++++++ 2 files changed, 479 insertions(+) create mode 100644 tests/analytical_visualization/auto_explore.py create mode 100644 tests/analytical_visualization/test_auto_explore.py diff --git a/tests/analytical_visualization/auto_explore.py b/tests/analytical_visualization/auto_explore.py new file mode 100644 index 0000000..a9bac4b --- /dev/null +++ b/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, + ) diff --git a/tests/analytical_visualization/test_auto_explore.py b/tests/analytical_visualization/test_auto_explore.py new file mode 100644 index 0000000..163c203 --- /dev/null +++ b/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}" + )