"""Joint hardware × parallelism exploration. For a fixed model + workload, sweep both: - hardware parameters (MachineParams: pe_hbm_gb, bw_hbm_gbs, peak_tflops_f16, bw_intra_gbs, bw_inter_gbs, bw_intersip_gbs), and - parallelism knobs (a reduced subset of auto_explore's 9 knobs) Then rank on: - latency ↓ - hardware cost proxy ↓ (normalized sum of knob-over-default) Also computes per-knob sensitivity — how much latency drops when each hardware knob is doubled from its baseline value. Useful for HW co-design ("which knob to invest in first?"). The default 'balanced' depth: 64 HW candidates × ~2k parallelism configs = ~130k joint evaluations, ~1-5 s at analytical speeds. Coarse and two-stage variants trade coverage for speed. """ from __future__ import annotations import math from collections.abc import Iterator from dataclasses import dataclass, field, replace from .auto_explore import ConfigScore, score_config from .autosuggest import auto_suggest from .memory_layout import compute_memory from .model_config import FullConfig, MachineParams, ModelConfig, TopologyConfig # ── Hardware search space ──────────────────────────────────────────── # # For each depth level, values per knob. Defaults are always included so the # baseline configuration always appears in the sweep. _HW_KNOB_DEFAULTS = { "pe_hbm_gb": 6.0, "bw_hbm_gbs": 256.0, "peak_tflops_f16": 8.0, "bw_intra_gbs": 512.0, "bw_inter_gbs": 128.0, "bw_intersip_gbs": 50.0, } _HW_VALUES_BALANCED = { "pe_hbm_gb": [6.0, 12.0], "bw_hbm_gbs": [256.0, 512.0], "peak_tflops_f16": [8.0, 16.0], "bw_intra_gbs": [512.0, 1024.0], "bw_inter_gbs": [128.0, 256.0], "bw_intersip_gbs": [50.0, 100.0], } _HW_VALUES_COARSE = { "pe_hbm_gb": [6.0, 12.0, 24.0], "bw_hbm_gbs": [256.0, 512.0, 1024.0], "peak_tflops_f16": [8.0, 16.0, 32.0], "bw_intra_gbs": [512.0, 1024.0, 2048.0], "bw_inter_gbs": [128.0, 256.0, 512.0], "bw_intersip_gbs": [50.0, 100.0, 200.0], } # For sensitivity: "double each knob independently from baseline." _SENSITIVITY_KNOBS = list(_HW_KNOB_DEFAULTS.keys()) # ── Reduced parallelism search space (for per-HW inner loop) ───────── # # The full auto_explore has 28,800 configs; using it inside a HW loop is # too slow. Reduce to CP × TP × PP × DP × kv_shard_mode with common # defaults for the remaining 4 knobs (~2k configs). _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") # ── Result types ───────────────────────────────────────────────────── @dataclass class HardwareCandidate: pe_hbm_gb: float bw_hbm_gbs: float peak_tflops_f16: float bw_intra_gbs: float bw_inter_gbs: float bw_intersip_gbs: float def as_machine(self) -> MachineParams: """Build a MachineParams from this HW candidate, holding alphas + compute_util at their MachineParams defaults.""" return MachineParams( pe_hbm_gb=self.pe_hbm_gb, bw_hbm_gbs=self.bw_hbm_gbs, peak_tflops_f16=self.peak_tflops_f16, bw_intra_gbs=self.bw_intra_gbs, bw_inter_gbs=self.bw_inter_gbs, bw_intersip_gbs=self.bw_intersip_gbs, ) @property def cost_score(self) -> float: """Normalized cost proxy: sum of (knob / default). 6.0 at defaults. This is NOT dollars — it's a rough silicon-area / capability proxy. A higher cost_score means "more capable hardware" (more HBM, faster BW, etc.). Under identical latency, lower cost_score wins. """ return sum( getattr(self, k) / _HW_KNOB_DEFAULTS[k] for k in _HW_KNOB_DEFAULTS ) @dataclass class JointScore: """One (hardware, parallelism) pair with its computed latency + cost.""" hardware: HardwareCandidate parallelism: ConfigScore total_latency_ns: float cost_score: float @property def latency_ms(self) -> float: return self.total_latency_ns / 1e6 @dataclass class SensitivityRow: knob: str baseline_value: float doubled_value: float baseline_latency_ns: float doubled_latency_ns: float @property def rel_speedup(self) -> float: """1 - (doubled/baseline). Positive = doubling this knob speeds things up.""" if self.baseline_latency_ns <= 0: return 0.0 return 1.0 - (self.doubled_latency_ns / self.baseline_latency_ns) @dataclass class JointExploreResult: model_name: str s_kv: int mode: str depth: str total_hw: int total_joint: int all_scores: list[JointScore] = field(default_factory=list) pareto_scores: list[JointScore] = field(default_factory=list) sensitivity: list[SensitivityRow] = field(default_factory=list) # ── Hardware enumeration ───────────────────────────────────────────── def enumerate_hardware(depth: str = "balanced") -> Iterator[HardwareCandidate]: """Yield HardwareCandidates for the given sweep depth. depth ∈ {"two_stage", "balanced", "coarse"}: - two_stage: only the default HW (1 candidate; parallelism sweep dominates) - balanced: 2 values per knob → 2^6 = 64 candidates - coarse: 3 values per knob → 3^6 = 729 candidates """ if depth == "two_stage": vals = {k: [v] for k, v in _HW_KNOB_DEFAULTS.items()} elif depth == "coarse": vals = _HW_VALUES_COARSE else: # balanced (default) vals = _HW_VALUES_BALANCED for pe_hbm in vals["pe_hbm_gb"]: for bw_hbm in vals["bw_hbm_gbs"]: for tflops in vals["peak_tflops_f16"]: for bw_intra in vals["bw_intra_gbs"]: for bw_inter in vals["bw_inter_gbs"]: for bw_intersip in vals["bw_intersip_gbs"]: yield HardwareCandidate( pe_hbm_gb=pe_hbm, bw_hbm_gbs=bw_hbm, peak_tflops_f16=tflops, bw_intra_gbs=bw_intra, bw_inter_gbs=bw_inter, bw_intersip_gbs=bw_intersip, ) # ── Reduced parallelism search per HW ──────────────────────────────── def _iter_reduced_parallelism( model: ModelConfig, s_kv: int, mode: str, ) -> Iterator[TopologyConfig]: """Reduced sweep: CP × TP × PP × DP × kv_shard_mode. Other 4 knobs held at defaults chosen to be latency-friendly for the decode/prefill common case: - ffn_shard_scope = "TP+CP" (common good choice; trades a small AR for lower per-PE weight bytes) - tp_placement = "cube" (packs TP across cubes; UCIe D2D) - cp_placement = "pe" (packs CP inside cubes; intra NoC) - cp_ring_variant = "qoml" for decode, "kv" for prefill Total: 8 × 6 × 5 × 4 × 2 = 1,920 configs per HW candidate. """ cp_ring = "qoml" if mode == "decode" else "kv" 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: # cp_ring=qoml with cp=1 is a no-op; fall back to kv. ring = "kv" if cp == 1 else cp_ring yield TopologyConfig( cp=cp, tp=tp, pp=pp, dp=dp, s_kv=s_kv, mode=mode, kv_shard_mode=kv_mode, ffn_shard_scope="TP+CP", tp_placement="cube", cp_placement="pe", cp_ring_variant=ring, ) def _best_parallelism_for_hw( model: ModelConfig, machine: MachineParams, s_kv: int, mode: str, include_ffn: bool = True, ) -> ConfigScore | None: """Return the latency-minimum feasible parallelism for this HW. Feasibility: fits memory + placement_valid. Returns None if nothing fits. ``include_ffn`` is forwarded to score_config so the "latency" ranked on matches the caller's attention-only vs full-model choice. """ best: ConfigScore | None = None for topo in _iter_reduced_parallelism(model, s_kv, mode): cfg = FullConfig(model=model, topo=topo, machine=machine) s = score_config(cfg, include_ffn=include_ffn) if not (s.fits_memory and s.placement_valid): continue if best is None or s.total_latency_ns < best.total_latency_ns: best = s return best def _best_parallelism_two_stage( model: ModelConfig, machine: MachineParams, s_kv: int, mode: str, include_ffn: bool = True, ) -> ConfigScore | None: """Fast fallback: use autosuggest's memory-min (CP,TP,PP) then score it.""" sug = auto_suggest(model, machine, s_kv, mode) if not sug.fits: return None topo = TopologyConfig( cp=sug.cp, tp=sug.tp, pp=sug.pp, dp=1, s_kv=s_kv, mode=mode, kv_shard_mode="split", ffn_shard_scope="TP+CP", tp_placement="cube", cp_placement="pe", cp_ring_variant="qoml" if mode == "decode" and sug.cp > 1 else "kv", ) cfg = FullConfig(model=model, topo=topo, machine=machine) s = score_config(cfg, include_ffn=include_ffn) return s if s.fits_memory and s.placement_valid else None # ── Pareto over joint (latency, cost) ──────────────────────────────── def _pareto_2d(scores: list[JointScore]) -> list[JointScore]: """Non-dominated set on (total_latency_ns ↓, cost_score ↓). O(N²).""" frontier: list[JointScore] = [] for i, a in enumerate(scores): dominated = False for j, b in enumerate(scores): if i == j: continue no_worse = ( b.total_latency_ns <= a.total_latency_ns and b.cost_score <= a.cost_score ) strictly_better = ( b.total_latency_ns < a.total_latency_ns or b.cost_score < a.cost_score ) if no_worse and strictly_better: dominated = True break if not dominated: frontier.append(a) return frontier # ── Per-knob sensitivity ───────────────────────────────────────────── def compute_sensitivity( baseline_hw: HardwareCandidate, parallelism: ConfigScore, model: ModelConfig, s_kv: int, mode: str, include_ffn: bool = True, ) -> list[SensitivityRow]: """For each HW knob, double it (holding others at baseline) and measure the latency change. Same parallelism used throughout so we isolate the HW knob's effect.""" rows: list[SensitivityRow] = [] baseline_machine = baseline_hw.as_machine() baseline_topo = parallelism.as_topology(s_kv, mode) baseline_cfg = FullConfig( model=model, topo=baseline_topo, machine=baseline_machine, ) baseline_score = score_config(baseline_cfg, include_ffn=include_ffn) baseline_latency = baseline_score.total_latency_ns for knob in _SENSITIVITY_KNOBS: doubled_val = 2.0 * getattr(baseline_hw, knob) doubled_hw = replace(baseline_hw, **{knob: doubled_val}) doubled_cfg = FullConfig( model=model, topo=baseline_topo, machine=doubled_hw.as_machine(), ) doubled_score = score_config(doubled_cfg, include_ffn=include_ffn) rows.append(SensitivityRow( knob=knob, baseline_value=getattr(baseline_hw, knob), doubled_value=doubled_val, baseline_latency_ns=baseline_latency, doubled_latency_ns=doubled_score.total_latency_ns, )) # Sort by biggest speedup first (most sensitive knob at the top). rows.sort(key=lambda r: -r.rel_speedup) return rows # ── Top-level driver ───────────────────────────────────────────────── def joint_explore( model: ModelConfig, s_kv: int, mode: str, depth: str = "balanced", include_ffn: bool = True, ) -> JointExploreResult: """Sweep HW candidates × parallelism, return joint Pareto + sensitivity. Sensitivity is computed around the LATENCY-MINIMUM joint point (the "best fast" config), doubling each HW knob one at a time. ``include_ffn=False`` restricts scoring to attention stages only — both the per-HW parallelism search and the sensitivity ranking use the same restriction. """ all_scores: list[JointScore] = [] for hw in enumerate_hardware(depth): machine = hw.as_machine() if depth == "two_stage": par = _best_parallelism_two_stage( model, machine, s_kv, mode, include_ffn=include_ffn, ) else: par = _best_parallelism_for_hw( model, machine, s_kv, mode, include_ffn=include_ffn, ) if par is None: continue all_scores.append(JointScore( hardware=hw, parallelism=par, total_latency_ns=par.total_latency_ns, cost_score=hw.cost_score, )) all_scores.sort(key=lambda s: s.total_latency_ns) pareto = _pareto_2d(all_scores) pareto.sort(key=lambda s: s.total_latency_ns) sensitivity: list[SensitivityRow] = [] if all_scores: best = all_scores[0] sensitivity = compute_sensitivity( best.hardware, best.parallelism, model, s_kv, mode, include_ffn=include_ffn, ) return JointExploreResult( model_name=model.name, s_kv=s_kv, mode=mode, depth=depth, total_hw=sum(1 for _ in enumerate_hardware(depth)), total_joint=len(all_scores), all_scores=all_scores, pareto_scores=pareto, sensitivity=sensitivity, )