mukesh
b8e1a3322f
Both auto tabs now accept a scope choice at run time:
- "Run sweep — Attention" → include_ffn=False
- "Run sweep — Attn + FFN/MoE" → include_ffn=True
Each button runs an independent sweep and caches its result under its
own session_state key. The most recently clicked button determines the
displayed view; both caches persist so users can flip between the two
scopes without re-running.
Tabs:
- Renamed "Auto Explore" → "Auto Suggest Parallelism"
(accurately reflects that it only varies parallelism knobs; HW is
held at the sidebar values).
- "Auto Hardware" tab unchanged.
- Still 6 top-level tabs; no additional tabs added.
Core changes:
auto_explore.py:
- New include_ffn: bool = True parameter on _sum_visible_latency,
_efficiency, score_config, run_auto_explore, compute_parallelism_
sensitivity. False drops all FFN stages from the summed latency.
auto_hardware.py:
- New include_ffn: bool = True parameter on joint_explore,
compute_sensitivity, _best_parallelism_for_hw, _best_parallelism_
two_stage. Forwards to score_config.
Both defaults keep existing tests byte-identical.
Verified:
- 23 pytest tests pass (added 3 new: attn-only latency lower, attn-only
Pareto non-empty, joint HW attn-only faster than full).
- Smoke: Llama 70B decode 128K:
* Attn+FFN best latency: 12.85 ms (unchanged)
* Attention-only best: 7.35 ms (~57% of full)
* Both sensitivities top-rank bw_hbm_gbs (physics preserved).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
435 lines
17 KiB
Python
435 lines
17 KiB
Python
"""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_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.
|
||
|
||
``include_ffn=False`` restricts to attention stages only — useful for
|
||
isolating attention-kernel tuning from FFN cost.
|
||
"""
|
||
attn = sum(s.visible_s for s in all_stages(cfg))
|
||
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_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)
|
||
|
||
``include_ffn=False`` mirrors the same restriction as
|
||
:func:`_sum_visible_latency` — attention stages only.
|
||
"""
|
||
if latency_s <= 0:
|
||
return 0.0
|
||
attn = all_stages(cfg)
|
||
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_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.
|
||
|
||
``include_ffn=False`` restricts latency + efficiency to attention stages
|
||
only. Memory feasibility is unchanged — still checks weights+KV+transient
|
||
fit in per-PE HBM since the model still exists physically.
|
||
"""
|
||
mem = compute_memory(cfg)
|
||
placement_ok = cfg.topo.placement_valid
|
||
|
||
latency_s = _sum_visible_latency(cfg, include_ffn=include_ffn)
|
||
throughput = 1.0 / latency_s if latency_s > 0 else 0.0
|
||
efficiency = (_efficiency(cfg, latency_s, 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_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_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_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.
|
||
|
||
``include_ffn=False`` restricts to attention stages only — useful for
|
||
isolating attention-kernel tuning independent of FFN cost.
|
||
"""
|
||
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_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,
|
||
)
|