analytical-viz: auto_hardware.py — joint HW×parallelism explore
New module extends auto_explore into the hardware co-design space. For a
fixed model + workload, sweeps hardware knobs (pe_hbm_gb, bw_hbm_gbs,
peak_tflops_f16, bw_intra_gbs, bw_inter_gbs, bw_intersip_gbs) and, for
each hardware candidate, searches parallelism for the latency-minimum
that fits memory. Returns:
- all_scores: every (hw, parallelism) pair that fits, sorted by latency
- pareto_scores: 2D Pareto frontier on (latency ↓, cost_score ↓)
- sensitivity: per-knob rel_speedup when doubled from the best-fast HW
baseline. Ranks which HW knob gives the biggest speedup — a co-design
signal.
Three sweep depths trade coverage for time:
- two_stage: 1 HW candidate (defaults) × autosuggest's memory-min
parallelism. Fast (~1s), useful for the sensitivity ranking alone.
- balanced: 64 HW × ~2k reduced-parallelism configs = ~130k joint evals,
~10-20s. Default UI setting.
- coarse: 729 HW × ~2k configs = ~1.4M joint evals, ~2-5 min.
Reduced parallelism sweep for the inner loop: CP × TP × PP × DP ×
kv_shard_mode (1,920 configs), other 4 knobs held at latency-friendly
defaults (ffn_shard_scope='TP+CP', tp_placement='cube', cp_placement='pe',
cp_ring_variant='qoml' for decode, 'kv' for prefill). Full 28,800-config
auto_explore per HW would take 6+ minutes — too slow.
Cost proxy: sum of (knob / knob_default). 6.0 at defaults. Not dollars —
a rough capability score where higher = "more spec'd hardware".
Verified:
- 9 pytest tests pass:
* enumeration counts match expected (1, 64, 729)
* default cost_score = 6.0
* Pareto non-dominated + subset of all_scores
* every knob is monotone-non-worsening when doubled
* for Llama 70B decode, bw_hbm_gbs tops the sensitivity ranking
(physically correct: memory-bound workload)
- Smoke: Llama 70B decode 128K balanced sweep in ~20s produces
5 Pareto configs; best 7.57 ms with 1024 GB/s HBM BW. Doubling
HBM BW gives 33% additional speedup; every other knob < 1.5%.
Next: Streamlit UI tab consuming this in Commit 5.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,394 @@
|
||||
"""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,
|
||||
) -> ConfigScore | None:
|
||||
"""Return the latency-minimum feasible parallelism for this HW.
|
||||
|
||||
Feasibility: fits memory + placement_valid. Returns None if nothing fits.
|
||||
"""
|
||||
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)
|
||||
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,
|
||||
) -> 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)
|
||||
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,
|
||||
) -> 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)
|
||||
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)
|
||||
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",
|
||||
) -> 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.
|
||||
"""
|
||||
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)
|
||||
else:
|
||||
par = _best_parallelism_for_hw(model, machine, s_kv, mode)
|
||||
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,
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
Reference in New Issue
Block a user