kernbench2/tests/analytical_visualization/chip_roofline.py

"""Chip-level roofline math: AI, B*, L*, per-token latency curves.

Surfaces the arithmetic-intensity story from LLM-serving practice:

- **AI**  = C / W  (peak FLOPs per byte of HBM bandwidth).
- **B\***  = C * b / (2 * W) * sparsity  — the critical batch size at
  which weight-fetch time equals compute time for one decode step.
  Sparsity = N_total / N_active (MoE factor; 1 for dense).
- **L\***  = 2 * N_active / (AI * kv_bytes_per_token) — the balance
  context length at which KV-read time equals compute time.
- **B_knee(S_kv)** = B* / (1 - S_kv/L*) — the batch size where the
  cost curve bends (weight-fetch drops below the compute+KV floor).
  Diverges at S_kv = L* and no knee exists past it.

Per-token decode-step latency, per PE (dense-approx, no comm):

    t(B, S_kv)  =  N_active * b / (W * B)     <-- weight fetch, 1/B
                 + 2 * N_active / C            <-- compute (peak), flat
                 + S_kv * kv_bpt / W           <-- KV read, flat

All numbers per PE / per one forward pass. **Peak roofline — no
utilization factor.** Comm cost and TP/CP sharding are intentionally
NOT in the roofline — this is the back-of-envelope chip-vs-model view
the transcript talks about, not the full latency model that
stage_latencies.py builds. With this convention weight_s == compute_s
exactly at B*.
"""
from __future__ import annotations

from dataclasses import dataclass

from .model_config import MachineParams, ModelConfig


# Per-token bytes of BF16 MAC arithmetic: one multiply + one add = 2 FLOPs.
_FLOPS_PER_PARAM_PER_TOKEN = 2


# ── Chip / model derived quantities ────────────────────────────────


def arithmetic_intensity(machine: MachineParams) -> float:
    """FLOPs per byte of HBM bandwidth. Peak roofline; utilization
    is applied only in the compute-time formula, not here."""
    return machine.peak_flops / machine.bw_hbm


def total_active_params(model: ModelConfig) -> int:
    """Full-model parameter count (attention + FFN, all layers).

    Attention: 4 projections × hidden × H_q * d_head effective per layer.
    (W_Q hidden×H_q*d_h, W_O H_q*d_h×hidden, W_K/W_V hidden×H_kv*d_h.)
    FFN: 3 × hidden × ffn_dim per layer (gate, up, down).
    """
    m = model
    attn = (
        m.hidden * m.h_q * m.d_head            # W_Q
        + m.hidden * m.h_kv * m.d_head * 2     # W_K + W_V
        + m.h_q * m.d_head * m.hidden          # W_O
    )
    ffn = 3 * m.hidden * m.ffn_dim
    return (attn + ffn) * m.layers


def kv_bytes_per_token(model: ModelConfig) -> int:
    """Bytes of KV cache one new token adds across ALL layers, per one
    sequence, un-sharded (K + V, H_kv heads * d_h * bytes)."""
    m = model
    return 2 * m.h_kv * m.d_head * m.bytes_per_elem * m.layers


def critical_batch(machine: MachineParams, model: ModelConfig,
                    sparsity: float = 1.0) -> float:
    """B* = C * b / (2 * W) * sparsity.

    Sparsity = N_total / N_active (>= 1). Dense = 1. MoE 8-of-256 = 8.
    """
    b = model.bytes_per_elem
    ai = arithmetic_intensity(machine)
    return ai * b / _FLOPS_PER_PARAM_PER_TOKEN * sparsity


def balance_context(machine: MachineParams, model: ModelConfig) -> float:
    """L* = 2 * N_active / (AI * kv_bpt).

    Context length (in tokens) at which per-step KV read matches the
    per-step compute cost. Beyond L*, the KV term is dominant and no
    batch size gets you compute-bound.
    """
    n_active = total_active_params(model)
    ai = arithmetic_intensity(machine)
    kv_bpt = kv_bytes_per_token(model)
    return _FLOPS_PER_PARAM_PER_TOKEN * n_active / (ai * kv_bpt)


def knee_batch(machine: MachineParams, model: ModelConfig,
                s_kv: int) -> float | None:
    """B_knee(S_kv) = B* / (1 - S_kv/L*).

    Returns None when S_kv >= L* (no knee exists — the total-cost
    curve never touches the compute floor).
    """
    b_star = critical_batch(machine, model)
    l_star = balance_context(machine, model)
    r = s_kv / l_star
    if r >= 1.0:
        return None
    return b_star / (1.0 - r)


# ── Per-token latency curves ───────────────────────────────────────


@dataclass
class RooflinePoint:
    batch: int
    weight_s: float   # weight fetch time, 1/B
    compute_s: float  # compute time, flat
    kv_s: float       # KV read time, flat
    total_s: float


def per_token_latency_curve(machine: MachineParams, model: ModelConfig,
                             batch_range: list[int],
                             s_kv: int) -> list[RooflinePoint]:
    """Per-token decode-step latency curve across a range of batch sizes.

    Returns one point per batch. All times per PE, dense-approx,
    utilization from machine.compute_util. Comm and TP/CP sharding are
    excluded — this is the roofline model.
    """
    n_active = total_active_params(model)
    b = model.bytes_per_elem
    weight_bytes = n_active * b
    compute_flops = _FLOPS_PER_PARAM_PER_TOKEN * n_active
    kv_read_bytes = s_kv * kv_bytes_per_token(model)

    compute_s = compute_flops / machine.peak_flops
    kv_s = kv_read_bytes / machine.bw_hbm

    points: list[RooflinePoint] = []
    for bs in batch_range:
        weight_s = weight_bytes / machine.bw_hbm / max(1, bs)
        points.append(RooflinePoint(
            batch=bs,
            weight_s=weight_s,
            compute_s=compute_s,
            kv_s=kv_s,
            total_s=weight_s + compute_s + kv_s,
        ))
    return points


def bound_regime(machine: MachineParams, model: ModelConfig,
                  batch: int, s_kv: int) -> str:
    """Which term dominates at the current (batch, S_kv) point.

    Returns 'memory-bound' if weight_fetch is the largest term,
    'kv-bound' if KV read is largest, 'compute-bound' if compute.
    """
    pts = per_token_latency_curve(machine, model, [batch], s_kv)
    p = pts[0]
    parts = {"memory-bound": p.weight_s,
             "kv-bound": p.kv_s,
             "compute-bound": p.compute_s}
    return max(parts, key=parts.get)