kernbench2/tests/analytical_visualization/auto_hardware.py

"""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,
    )