kernbench2/tests/analytical_visualization/autosuggest.py

"""Auto-suggest (CP, TP, PP) that fits the model in the per-PE budget.

Strategy: iterate over candidate (CP, TP, PP) triples in order of
increasing total PE count and return the first one that satisfies
memory + kernel-support constraints with >10% slack. If none fits, return
the best-effort configuration and flag over-budget.
"""
from __future__ import annotations

from dataclasses import dataclass, replace
from typing import Iterable

from .model_config import FullConfig, ModelConfig, TopologyConfig, MachineParams
from .memory_layout import compute_memory


# Candidate parallelism dimensions (powers of 2 mostly, up to sensible limits).
_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
_PP_OPTIONS = (1, 2, 4, 8, 16, 32)


@dataclass
class Suggestion:
    cp: int
    tp: int
    pp: int
    weights_gb: float
    kv_gb: float
    transient_gb: float
    slack_gb: float
    fits: bool
    pes_used: int
    sips_used: int
    cubes_used: int = 0            # picked to minimize this (see _score_candidate)
    cp_placement: str = "cube"     # "pe" packs CP into intra-cube PEs when it fits
    reason: str = ""


def _iter_candidates(model: ModelConfig) -> Iterable[tuple[int, int, int]]:
    """Yield (CP, TP, PP) in order of increasing PE count."""
    triples: list[tuple[int, int, int]] = []
    for tp in _TP_OPTIONS:
        for cp in _CP_OPTIONS:
            for pp in _PP_OPTIONS:
                # PP must not exceed layer count.
                if pp > model.layers:
                    continue
                # Skip TP > H_q * some factor (unrealistic).
                if tp > model.h_q * 4:
                    continue
                triples.append((cp, tp, pp))
    # Sort by pe_count then by (pp, tp, cp) — prefer smaller PP first
    # (avoids pipeline bubbles), then smaller TP (avoids head-dim split).
    triples.sort(key=lambda t: (t[0] * t[1] * t[2], t[2], t[1], t[0]))
    for t in triples:
        yield t


def _score_candidate(cp: int, tp: int, pp: int,
                     model: ModelConfig, machine: MachineParams,
                     s_kv: int, mode: str,
                     slack_frac: float = 0.10,
                     b: int = 1,
                     include_attention: bool = True,
                     include_ffn: bool = True) -> Suggestion:
    # For each triple, try both cp_placement options and keep the one
    # with fewer cubes (breaks the historical "CP always across cubes"
    # default when a smaller pack is possible). The pe placement is only
    # valid when CP·TP fits within a single cube's PE count.
    best_topo: TopologyConfig | None = None
    best_placement = "cube"
    _pe_per_cube_hw = TopologyConfig().pes_per_cube_hw  # instance-invariant HW const
    _placements_to_try = ["cube"]
    if cp * tp <= _pe_per_cube_hw:
        _placements_to_try.append("pe")
    for _place in _placements_to_try:
        _t = TopologyConfig(
            cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode, b=b,
            cp_placement=_place,
        )
        if best_topo is None or _t.cubes_used < best_topo.cubes_used:
            best_topo = _t
            best_placement = _place
    topo = best_topo
    cfg = FullConfig(model=model, topo=topo, machine=machine)
    mem = compute_memory(
        cfg, include_attention=include_attention, include_ffn=include_ffn,
    )

    fits = (not mem.over_budget
            and mem.slack_bytes >= slack_frac * mem.budget_bytes)
    reason = ""
    if mem.over_budget:
        reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
                  f"exceeds budget ({mem.budget_bytes/1e9:.2f} GB)")
    elif not fits:
        reason = f"slack ({mem.slack_bytes/1e9:.2f} GB) below 10% of budget"
    return Suggestion(
        cp=cp, tp=tp, pp=pp,
        weights_gb=mem.weights_bytes / 1e9,
        kv_gb=mem.kv_cache_bytes / 1e9,
        transient_gb=mem.transient_bytes / 1e9,
        slack_gb=mem.slack_bytes / 1e9,
        fits=fits,
        pes_used=topo.total_pes,
        sips_used=topo.sips_used,
        cubes_used=topo.cubes_used,
        cp_placement=best_placement,
        reason=reason,
    )


def auto_suggest(model: ModelConfig, machine: MachineParams,
                 s_kv: int, mode: str = "decode",
                 slack_frac: float = 0.10,
                 b: int = 1,
                 include_attention: bool = True,
                 include_ffn: bool = True) -> Suggestion:
    """Return the smallest deployment (fewer cubes, then fewer PEs)
    that fits.

    Sort key: (cubes_used ↑, pes_used ↑, pp ↑, tp ↑, cp ↑).
    Fewer cubes wins first because a cube is the physical die-level
    hardware unit; PE count is the tiebreaker. Preferring fewer PP then
    TP then CP keeps the earlier historical bias (avoid pipeline
    bubbles / head-dim splits) among equal-cube-and-PE ties.

    If no candidate fits, returns the best-effort one (highest slack,
    even if negative) with fits=False.
    """
    scored: list[Suggestion] = []
    for cp, tp, pp in _iter_candidates(model):
        scored.append(
            _score_candidate(cp, tp, pp, model, machine, s_kv, mode,
                              slack_frac, b=b,
                              include_attention=include_attention,
                              include_ffn=include_ffn)
        )
    scored.sort(key=lambda s: (s.cubes_used, s.pes_used, s.pp, s.tp, s.cp))

    for s in scored:
        if s.fits:
            return s
    # No fit — return the best-effort (largest slack, closest to fitting).
    return max(scored, key=lambda s: s.slack_gb)