mukesh
6ef09dd6f5
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>
117 lines
4.5 KiB
Python
117 lines
4.5 KiB
Python
"""Interface + invariant tests for auto_hardware."""
|
||
from __future__ import annotations
|
||
|
||
from tests.analytical_visualization.auto_hardware import (
|
||
HardwareCandidate,
|
||
JointScore,
|
||
_HW_KNOB_DEFAULTS,
|
||
enumerate_hardware,
|
||
joint_explore,
|
||
)
|
||
from tests.analytical_visualization.model_presets import PRESETS
|
||
|
||
|
||
# ── Enumeration ─────────────────────────────────────────────────────
|
||
|
||
|
||
def test_enumerate_two_stage_yields_one():
|
||
"""Two-stage depth yields exactly 1 HW candidate (defaults)."""
|
||
hws = list(enumerate_hardware("two_stage"))
|
||
assert len(hws) == 1
|
||
for knob, default_val in _HW_KNOB_DEFAULTS.items():
|
||
assert getattr(hws[0], knob) == default_val
|
||
|
||
|
||
def test_enumerate_balanced_yields_64():
|
||
"""Balanced = 2 values × 6 knobs = 2^6 = 64 candidates."""
|
||
hws = list(enumerate_hardware("balanced"))
|
||
assert len(hws) == 64
|
||
|
||
|
||
def test_enumerate_coarse_yields_729():
|
||
"""Coarse = 3 values × 6 knobs = 3^6 = 729 candidates."""
|
||
hws = list(enumerate_hardware("coarse"))
|
||
assert len(hws) == 729
|
||
|
||
|
||
def test_default_cost_score_is_6():
|
||
"""At every knob = its default, cost_score = 6 (one per knob)."""
|
||
hw = HardwareCandidate(
|
||
pe_hbm_gb=_HW_KNOB_DEFAULTS["pe_hbm_gb"],
|
||
bw_hbm_gbs=_HW_KNOB_DEFAULTS["bw_hbm_gbs"],
|
||
peak_tflops_f16=_HW_KNOB_DEFAULTS["peak_tflops_f16"],
|
||
bw_intra_gbs=_HW_KNOB_DEFAULTS["bw_intra_gbs"],
|
||
bw_inter_gbs=_HW_KNOB_DEFAULTS["bw_inter_gbs"],
|
||
bw_intersip_gbs=_HW_KNOB_DEFAULTS["bw_intersip_gbs"],
|
||
)
|
||
assert hw.cost_score == 6.0
|
||
|
||
|
||
# ── Joint explore + Pareto ──────────────────────────────────────────
|
||
|
||
|
||
def test_joint_explore_returns_pareto_non_empty():
|
||
"""Given a reasonable model+workload, at least one HW+parallelism fits."""
|
||
model = PRESETS["Llama 3.1 70B"].model
|
||
res = joint_explore(model, s_kv=131072, mode="decode", depth="two_stage")
|
||
assert res.total_joint >= 1
|
||
assert len(res.pareto_scores) >= 1
|
||
|
||
|
||
def test_pareto_subset_of_all_scores():
|
||
"""Every Pareto entry is present in all_scores."""
|
||
model = PRESETS["Llama 3.1 70B"].model
|
||
res = joint_explore(model, s_kv=131072, mode="decode", depth="two_stage")
|
||
pareto_ids = {id(s) for s in res.pareto_scores}
|
||
all_ids = {id(s) for s in res.all_scores}
|
||
assert pareto_ids.issubset(all_ids)
|
||
|
||
|
||
def test_pareto_non_dominated():
|
||
"""No Pareto entry is dominated by another Pareto entry on (lat, cost)."""
|
||
model = PRESETS["Llama 3.1 70B"].model
|
||
res = joint_explore(model, s_kv=131072, mode="decode", depth="balanced")
|
||
for i, a in enumerate(res.pareto_scores):
|
||
for j, b in enumerate(res.pareto_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
|
||
)
|
||
assert not (no_worse and strictly_better), (
|
||
f"Pareto entry {i} is dominated by entry {j}"
|
||
)
|
||
|
||
|
||
# ── Sensitivity ─────────────────────────────────────────────────────
|
||
|
||
|
||
def test_sensitivity_all_knobs_monotone_non_worsening():
|
||
"""Doubling any HW knob should not slow the sim down (rel_speedup ≥ 0
|
||
within floating-point tolerance)."""
|
||
model = PRESETS["Llama 3.1 70B"].model
|
||
res = joint_explore(model, s_kv=131072, mode="decode", depth="two_stage")
|
||
assert len(res.sensitivity) == 6
|
||
for row in res.sensitivity:
|
||
# Allow a tiny slop for floating-point but disallow real regressions.
|
||
assert row.rel_speedup >= -1e-9, (
|
||
f"knob {row.knob}: doubling slowed latency from "
|
||
f"{row.baseline_latency_ns} to {row.doubled_latency_ns}"
|
||
)
|
||
|
||
|
||
def test_sensitivity_hbm_bw_dominant_for_llama_decode():
|
||
"""For Llama 70B decode (memory-bound), HBM BW should be the top
|
||
sensitivity knob — doubling it gives more speedup than any other knob."""
|
||
model = PRESETS["Llama 3.1 70B"].model
|
||
res = joint_explore(model, s_kv=131072, mode="decode", depth="balanced")
|
||
assert res.sensitivity[0].knob == "bw_hbm_gbs", (
|
||
f"expected bw_hbm_gbs to top the sensitivity ranking for Llama 70B "
|
||
f"decode; got {res.sensitivity[0].knob}"
|
||
)
|