"""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 # ── 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 ───────────────────────────────────────────────────── _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, 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 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. Scope selection (both default True — full per-token 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)) 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) Scope flags mirror :func:`_sum_visible_latency`. """ if latency_s <= 0: return 0.0 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) 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, 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. Scope flags (default: full transformer) restrict *latency + efficiency* to attention only, FFN only, or both. Memory feasibility is unchanged — 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_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_attention=include_attention, include_ffn=include_ffn) 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 # ── 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, include_attention: bool = True, include_ffn: bool = True, ) -> 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, 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( 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 ───────────────────────────────────────────────── def run_auto_explore( model: ModelConfig, 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. 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. Scope: at least one of ``include_attention`` / ``include_ffn`` must be 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_attention=include_attention, include_ffn=include_ffn, )) 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, )