Why a fast AI chip spends most of its time waiting

"""Roofline arithmetic and tiled matrix multiplication with traffic counting.

Model used throughout the lesson: two levels of memory. Slow memory is large and has a
limited bandwidth. Fast memory is small, and traffic inside it costs nothing. We count the
elements that cross between the two levels, and the floating-point operations (FLOPs) done.
One multiply-add is 2 FLOPs, the same convention the vendor uses for its peak numbers.

Run `python3 tiled_roofline.py` for the numbers quoted in the lesson, and
`python3 tiled_roofline.py --time` to time numpy and a pure-Python loop on your own machine.
"""
from __future__ import annotations

import sys
import time
from dataclasses import dataclass
from math import ceil, isqrt

import numpy as np

# NVIDIA A100 Tensor Core GPU data sheet (revision "Jun21"), 80GB SXM column.
# BFLOAT16 / FP16 Tensor Core: "312 TFLOPS | 624 TFLOPS*" (* with sparsity). We use the dense value.
# GPU memory bandwidth: 2,039 GB/s (HBM2e). We read GB as 10**9 bytes.
A100_PEAK_FLOPS = 312e12
A100_BANDWIDTH = 2039e9
FP16_BYTES = 2


# --- FLOPs, bytes, intensity, roofline ---------------------------------------------------

def matmul_flops(m: int, n: int, k: int) -> int:
    """C[m, n] = A[m, k] @ B[k, n]: m*n*k multiply-adds, 2 FLOPs each."""
    return 2 * m * n * k


def matmul_min_bytes(m: int, n: int, k: int, elem_bytes: int) -> int:
    """Least traffic possible: read A once, read B once, write C once."""
    return (m * k + k * n + m * n) * elem_bytes


def gemm_intensity(m: int, n: int, k: int, elem_bytes: int) -> float:
    """FLOP per byte if every element moves exactly once (the best case)."""
    return matmul_flops(m, n, k) / matmul_min_bytes(m, n, k, elem_bytes)


def ridge(peak_flops: float, bandwidth: float) -> float:
    """Intensity (FLOP/byte) where the memory slope meets the flat compute roof."""
    return peak_flops / bandwidth


def attainable(intensity: float, peak_flops: float, bandwidth: float) -> float:
    """The roofline: attainable FLOP/s = min(peak, bandwidth * intensity)."""
    return min(peak_flops, bandwidth * intensity)


def classify(intensity: float, peak_flops: float, bandwidth: float) -> str:
    """A tie at the ridge counts as compute-bound."""
    return "memory-bound" if intensity < ridge(peak_flops, bandwidth) else "compute-bound"


def time_floors(flops: float, nbytes: float, peak_flops: float, bandwidth: float) -> tuple[float, float]:
    """The two lower bounds on run time, in seconds: (math alone, traffic alone).

    Neither the math nor the traffic can finish sooner than its own floor, and the two overlap,
    so the run takes at least the larger of the two.
    """
    return flops / peak_flops, nbytes / bandwidth


# --- traffic of a tiled matrix multiply --------------------------------------------------

@dataclass
class Traffic:
    loads: int = 0    # elements read from slow memory
    stores: int = 0   # elements written to slow memory
    flops: int = 0

    def nbytes(self, elem_bytes: int) -> int:
        return (self.loads + self.stores) * elem_bytes

    def intensity(self, elem_bytes: int) -> float:
        return self.flops / self.nbytes(elem_bytes)


def tiled_matmul(A: np.ndarray, B: np.ndarray, b: int) -> tuple[np.ndarray, Traffic]:
    """Multiply with b x b tiles and count what crosses between slow and fast memory.

    For each output tile, the C tile stays in fast memory while we walk along k. Every
    k-step loads one A tile and one B tile. The C tile is written once at the end.
    Edge tiles are smaller when a size is not a multiple of b.
    """
    m, k = A.shape
    k2, n = B.shape
    if k != k2 or b < 1:
        raise ValueError("shapes must agree and the tile size must be at least 1")
    C = np.zeros((m, n), dtype=np.result_type(A, B))
    t = Traffic()
    for i in range(0, m, b):
        for j in range(0, n, b):
            acc = np.zeros((min(b, m - i), min(b, n - j)), dtype=C.dtype)
            for p in range(0, k, b):
                a_tile = A[i:i + b, p:p + b]
                b_tile = B[p:p + b, j:j + b]
                t.loads += a_tile.size + b_tile.size
                t.flops += 2 * a_tile.shape[0] * a_tile.shape[1] * b_tile.shape[1]
                acc += a_tile @ b_tile
            C[i:i + b, j:j + b] = acc
            t.stores += acc.size
    return C, t


def tiled_traffic(m: int, n: int, k: int, b: int) -> Traffic:
    """Closed form of what `tiled_matmul` counts, with no loops over tiles.

    Every A element is loaded once per column of output tiles, ceil(n / b) times.
    Every B element is loaded once per row of output tiles, ceil(m / b) times.
    """
    return Traffic(loads=ceil(n / b) * m * k + ceil(m / b) * k * n,
                   stores=m * n,
                   flops=matmul_flops(m, n, k))


def min_traffic(m: int, n: int, k: int) -> Traffic:
    """Best case: every element of A and B is loaded once, every element of C stored once."""
    return Traffic(loads=m * k + k * n, stores=m * n, flops=matmul_flops(m, n, k))


def naive_traffic(m: int, n: int, k: int) -> Traffic:
    """No reuse: each multiply-add fetches its two operands; C is written once."""
    return Traffic(loads=2 * m * n * k, stores=m * n, flops=matmul_flops(m, n, k))


def tile_steps(m: int, n: int, k: int, b: int) -> list[tuple[int, int, int, int]]:
    """The (output row tile, output column tile, k tile, elements loaded) of each step."""
    steps = []
    for i in range(0, m, b):
        for j in range(0, n, b):
            for p in range(0, k, b):
                bm, bn, bk = min(b, m - i), min(b, n - j), min(b, k - p)
                steps.append((i // b, j // b, p // b, bm * bk + bk * bn))
    return steps


def fast_bytes_needed(b: int, elem_bytes: int) -> int:
    """One A tile, one B tile and one C tile must sit in fast memory together."""
    return 3 * b * b * elem_bytes


def max_tile(fast_bytes: int, elem_bytes: int) -> int:
    """Largest b with 3 * b * b * elem_bytes <= fast_bytes."""
    return isqrt(fast_bytes // (3 * elem_bytes))


def tiled_intensity(n: int, b: int, elem_bytes: int) -> float:
    """FLOP per byte of an n x n x n multiply with b x b tiles: 2n^3 / ((2n^3 / b + n^2) * s)."""
    return tiled_traffic(n, n, n, b).intensity(elem_bytes)


def ridge_tile_if_tile_count_were_fractional(n: int, elem_bytes: int, peak_flops: float, bandwidth: float) -> float:
    """Smallest b with (b / s) / (1 + b / 2n) at or above the ridge, if n / b could be fractional.

    This keeps the term for writing C but drops the whole-tile effect. It is the middle step
    between the rule of thumb s * ridge (306 on the A100) and the exact answer (342 at n = 4096).
    """
    sr = elem_bytes * ridge(peak_flops, bandwidth)
    return sr / (1 - sr / (2 * n))


def smallest_tile_for_ridge(n: int, elem_bytes: int, peak_flops: float, bandwidth: float) -> int:
    """Smallest b whose n x n x n tiled multiply is no longer memory-bound.

    Not the same as the rule of thumb elem_bytes * ridge (306 on the A100), for two causes.
    Writing C adds traffic, which alone moves the answer to 318. The tile count per side,
    ceil(n / b), must also be a whole number: at n = 4096, 13 tiles per side give 151.7 and
    12 give 163.8 (the ridge is 153.0), so b must be at least 4096 / 12 = 341.3, which is 342.
    """
    r = ridge(peak_flops, bandwidth)
    return next(b for b in range(1, n + 1) if tiled_intensity(n, b, elem_bytes) >= r)


# --- decode versus prefill ---------------------------------------------------------------

def batch_to_reach_ridge(n: int, k: int, elem_bytes: int, peak_flops: float, bandwidth: float) -> int:
    """Smallest number of tokens m for which (m x k) @ (k x n) is no longer memory-bound."""
    r = ridge(peak_flops, bandwidth)
    m = 1
    while gemm_intensity(m, n, k, elem_bytes) < r:
        m += 1
    return m


def weight_read_seconds(params: float, elem_bytes: int, bandwidth: float) -> float:
    """Time to read every weight of a model once: the floor for one decoding step."""
    return params * elem_bytes / bandwidth


def decode_tokens_per_s(params: float, elem_bytes: int, bandwidth: float) -> float:
    """Upper bound on one decoding request if each token reads every weight once.

    Ignores the attention cache (stored keys and values) that each token also reads.
    """
    return 1 / weight_read_seconds(params, elem_bytes, bandwidth)


# --- the worked example and a slow reference ---------------------------------------------

EXAMPLE_A = np.array([[1, 2, 0, 1], [3, 1, 2, 0], [0, 1, 1, 2], [2, 0, 1, 1]])
EXAMPLE_B = np.array([[1, 2, 2, 1], [0, 1, 1, 0], [2, 1, 0, 1], [1, 2, 1, 0]])


def naive_matmul(A: list[list[float]], B: list[list[float]]) -> list[list[float]]:
    """Three plain Python loops. Slow on purpose: it shows what the library saves you."""
    m, k, n = len(A), len(B), len(B[0])
    C = [[0.0] * n for _ in range(m)]
    for i in range(m):
        for j in range(n):
            s = 0.0
            for p in range(k):
                s += A[i][p] * B[p][j]
            C[i][j] = s
    return C


def time_on_this_machine(n_small: int = 96, n_mm: int = 1024, n_mv: int = 4096) -> dict[str, float]:
    """Time four jobs and return rates (GFLOP/s, plus GB/s for the matrix-vector case). Your numbers depend on your machine."""
    rng = np.random.default_rng(0)

    def best(fn, repeats=3):
        times = []
        for _ in range(repeats):
            t0 = time.perf_counter()
            fn()
            times.append(time.perf_counter() - t0)
        return min(times)

    a_small = rng.standard_normal((n_small, n_small))
    a_list, b_list = a_small.tolist(), a_small.T.tolist()
    A = rng.standard_normal((n_mm, n_mm), dtype=np.float32)
    B = rng.standard_normal((n_mm, n_mm), dtype=np.float32)
    M = rng.standard_normal((n_mv, n_mv), dtype=np.float32)
    x = rng.standard_normal(n_mv, dtype=np.float32)
    t_py = best(lambda: naive_matmul(a_list, b_list), repeats=1)
    t_np = best(lambda: a_small @ a_small.T)
    t_mm = best(lambda: A @ B)
    t_mv = best(lambda: M @ x)
    return {
        f"pure Python loop, {n_small}x{n_small} matmul (GFLOP/s)": matmul_flops(n_small, n_small, n_small) / t_py / 1e9,
        f"numpy, {n_small}x{n_small} matmul (GFLOP/s)": matmul_flops(n_small, n_small, n_small) / t_np / 1e9,
        f"numpy, {n_mm}x{n_mm} matmul, float32 (GFLOP/s)": matmul_flops(n_mm, n_mm, n_mm) / t_mm / 1e9,
        f"numpy, {n_mv}x{n_mv} matrix-vector, float32 (GFLOP/s)": matmul_flops(n_mv, 1, n_mv) / t_mv / 1e9,
        f"numpy, {n_mv}x{n_mv} matrix-vector, matrix bytes read (GB/s)": n_mv * n_mv * 4 / t_mv / 1e9,
    }


# --- report: the numbers quoted in the lesson ---------------------------------------------

def report() -> None:
    s = FP16_BYTES
    r = ridge(A100_PEAK_FLOPS, A100_BANDWIDTH)
    print(f"A100 SXM 80GB: peak {A100_PEAK_FLOPS / 1e12:.0f} TFLOP/s, bandwidth {A100_BANDWIDTH / 1e9:.0f} GB/s")
    print(f"ridge point: {r:.1f} FLOP/byte")

    print("\n-- tiny example: 2x2 matmul, float32")
    print("flops", matmul_flops(2, 2, 2), "min bytes", matmul_min_bytes(2, 2, 2, 4),
          "intensity", round(gemm_intensity(2, 2, 2, 4), 4))
    for n in (2, 16, 256, 4096):
        print(f"best-case intensity, n={n}, float32: {gemm_intensity(n, n, n, 4):.3f}")

    n = 4096
    print(f"\n-- naive vs ideal, n={n}, FP16")
    nv = naive_traffic(n, n, n)
    best_t = min_traffic(n, n, n)
    print(f"flops {nv.flops:.4e}")
    print(f"naive bytes {nv.nbytes(s):.4e}, intensity {nv.intensity(s):.3f}")
    print(f"min bytes {best_t.nbytes(s):.4e}, intensity {best_t.intensity(s):.1f}")
    print(f"naive attainable {attainable(nv.intensity(s), A100_PEAK_FLOPS, A100_BANDWIDTH) / 1e12:.3f} TFLOP/s "
          f"= {100 * attainable(nv.intensity(s), A100_PEAK_FLOPS, A100_BANDWIDTH) / A100_PEAK_FLOPS:.2f}% of peak")
    t_math, t_naive = time_floors(nv.flops, nv.nbytes(s), A100_PEAK_FLOPS, A100_BANDWIDTH)
    _, t_min = time_floors(best_t.flops, best_t.nbytes(s), A100_PEAK_FLOPS, A100_BANDWIDTH)
    print(f"time floor from math {t_math * 1e3:.3f} ms; from naive traffic {t_naive * 1e3:.1f} ms "
          f"({t_naive / t_math:.0f} times the math floor); from minimum traffic {t_min * 1e3:.4f} ms")

    print("\n-- 4x4 with 2x2 tiles")
    C, t = tiled_matmul(EXAMPLE_A, EXAMPLE_B, 2)
    print("loads", t.loads, "stores", t.stores, "flops", t.flops, "naive", naive_traffic(4, 4, 4))
    print("C =", C.tolist())
    print(f"intensity per number moved: naive {naive_traffic(4, 4, 4).intensity(1):.2f}, tiled {t.intensity(1):.2f}")

    print("\n-- simulated loop, n=128, FP16: counted by tiled_matmul vs the formula 2 n^3 / b")
    rng = np.random.default_rng(0)
    A, B = rng.standard_normal((128, 128)), rng.standard_normal((128, 128))
    for b in (4, 8, 16, 32, 64, 128):
        C, t = tiled_matmul(A, B, b)
        assert np.allclose(C, A @ B)
        print(f"b={b}: loads {t.loads}, formula {2 * 128**3 // b}, stores {t.stores}, flops {t.flops}, "
              f"intensity {t.intensity(s):.2f} FLOP/byte")

    print(f"\n-- tile size sweep, n={n}, FP16")
    print("b, loads (elements), traffic/min, intensity, fast bytes needed")
    for b in (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 4096):
        t = tiled_traffic(n, n, n, b)
        print(b, t.loads, round(t.nbytes(s) / best_t.nbytes(s), 2), round(t.intensity(s), 2),
              fast_bytes_needed(b, s))

    print("\n-- capacity (example sizes, not a claim about any chip)")
    for kib in (16, 64, 256, 1024):
        b = max_tile(kib * 1024, s)
        print(f"{kib} KiB -> b <= {b}, intensity ~ {tiled_intensity(n, b, s):.1f}")
    rule = s * r
    b_need = smallest_tile_for_ridge(n, s, A100_PEAK_FLOPS, A100_BANDWIDTH)
    print(f"large-n rule: intensity ~ b / elem_bytes, so the ridge {r:.1f} needs b ~ {rule:.0f}")
    b_frac = ridge_tile_if_tile_count_were_fractional(n, s, A100_PEAK_FLOPS, A100_BANDWIDTH)
    print(f"with the C-store term but a fractional tile count: b >= {b_frac:.1f}, so {ceil(b_frac)}")
    print(f"exact at n={n}: smallest b at or above the ridge is {b_need}; "
          f"intensity there {tiled_intensity(n, b_need, s):.1f}, one smaller {tiled_intensity(n, b_need - 1, s):.1f}; "
          f"fast bytes {fast_bytes_needed(b_need, s)}")

    print("\n-- decode vs prefill: (m x 4096) @ (4096 x 4096), FP16")
    for m in (1, 2, 8, 16, 32, 64, 128, 165, 166, 256, 512, 1024, 4096):
        i = gemm_intensity(m, 4096, 4096, s)
        att = attainable(i, A100_PEAK_FLOPS, A100_BANDWIDTH)
        print(f"m={m}: intensity {i:.2f}, {classify(i, A100_PEAK_FLOPS, A100_BANDWIDTH)}, "
              f"{att / 1e12:.2f} TFLOP/s ({100 * att / A100_PEAK_FLOPS:.1f}% of peak)")
    print("batch to reach ridge:", batch_to_reach_ridge(4096, 4096, s, A100_PEAK_FLOPS, A100_BANDWIDTH))

    params = 7e9
    print(f"\n-- 7B params, FP16: weights {params * s / 1e9:.0f} GB; "
          f"read once = {weight_read_seconds(params, s, A100_BANDWIDTH) * 1e3:.2f} ms; "
          f"at most {decode_tokens_per_s(params, s, A100_BANDWIDTH):.0f} tokens/s at batch 1 (upper bound)")

    print("\n-- check-yourself: 48 KiB of fast memory, 4-byte numbers, n = 4096")
    b = max_tile(48 * 1024, 4)
    print(f"largest tile {b}, intensity {tiled_intensity(n, b, 4):.1f}")


if __name__ == "__main__":
    if "--time" in sys.argv:
        for name, value in time_on_this_machine().items():
            print(f"{name}: {value:.2f}")
    else:
        report()

import math

import numpy as np
import pytest

import tiled_roofline as tr


@pytest.mark.parametrize("shape, b", [((8, 8, 8), 2), ((37, 23, 41), 8), ((1, 33, 17), 8), ((7, 5, 3), 1), ((9, 11, 13), 64)])
def test_tiled_matmul_is_correct_when_sizes_do_not_divide(shape, b):
    """Guards an off-by-one on edge tiles: a wrong slice end gives a wrong C, not an error."""
    m, k, n = shape
    rng = np.random.default_rng(0)
    A, B = rng.standard_normal((m, k)), rng.standard_normal((k, n))
    C, _ = tr.tiled_matmul(A, B, b)
    assert np.allclose(C, A @ B, atol=1e-9)
    assert np.allclose(C, tr.naive_matmul(A.tolist(), B.tolist()), atol=1e-9)


@pytest.mark.parametrize("shape, b", [((64, 64, 64), 8), ((37, 23, 41), 8), ((1, 33, 17), 8), ((9, 11, 13), 64), ((5, 3, 7), 1)])
def test_counted_loop_matches_closed_form(shape, b):
    """The formula in the lesson and the loop must agree, ragged edges included.
    Guards a closed form that drifts from the loop it claims to describe."""
    m, k, n = shape
    A, B = np.ones((m, k)), np.ones((k, n))
    _, counted = tr.tiled_matmul(A, B, b)
    assert counted == tr.tiled_traffic(m, n, k, b)
    assert counted.flops == tr.matmul_flops(m, n, k)


def test_square_formula_and_tile_one_is_the_naive_count():
    """Loads = 2 n^3 / b for square sizes; tile 1 is exactly the no-reuse count."""
    n, b = 64, 16
    assert tr.tiled_traffic(n, n, n, b).loads == 2 * n**3 // b
    assert tr.tiled_traffic(n, n, n, 1) == tr.naive_traffic(n, n, n)
    assert tr.tiled_traffic(n, n, n, n) == tr.min_traffic(n, n, n)    # one tile: every element moves once


def test_worked_example_four_by_four_with_two_by_two_tiles():
    """The numbers the lesson prints by hand: 8 steps of 8 loads, 16 stores, 128 FLOPs."""
    C, t = tr.tiled_matmul(tr.EXAMPLE_A, tr.EXAMPLE_B, 2)
    assert np.array_equal(C, tr.EXAMPLE_A @ tr.EXAMPLE_B)
    assert C.tolist() == [[2, 6, 5, 1], [7, 9, 7, 5], [4, 6, 3, 1], [5, 7, 5, 3]]    # the matrix printed in the lesson
    steps = tr.tile_steps(4, 4, 4, 2)
    assert len(steps) == 8 and {s[3] for s in steps} == {8}
    assert sum(s[3] for s in tr.tile_steps(5, 3, 7, 2)) == tr.tiled_traffic(5, 3, 7, 2).loads    # ragged edges: guards a wrong edge-tile size in the per-step loads
    assert (t.loads, t.stores, t.flops) == (64, 16, 128)
    assert tr.naive_traffic(4, 4, 4).loads == 128
    assert tr.naive_traffic(4, 4, 4).intensity(1) == pytest.approx(0.89, abs=0.005)    # per number moved: 128 / 144
    assert t.intensity(1) == pytest.approx(1.6)                                          # 128 / 80
    # top-left output tile, the two partial products the lesson tabulates
    a, b = tr.EXAMPLE_A, tr.EXAMPLE_B
    step0 = a[0:2, 0:2] @ b[0:2, 0:2]
    step1 = a[0:2, 2:4] @ b[2:4, 0:2]
    assert step0.tolist() == [[1, 4], [3, 7]]
    assert step1.tolist() == [[1, 2], [4, 2]]
    assert (step0 + step1).tolist() == [[2, 6], [7, 9]]


def test_a100_ridge_from_the_datasheet_numbers():
    """312 TFLOP/s over 2,039 GB/s. Guards a unit slip (TFLOP vs GFLOP, GB vs GiB)."""
    assert tr.ridge(tr.A100_PEAK_FLOPS, tr.A100_BANDWIDTH) == pytest.approx(153.0, abs=0.05)


def test_roofline_min_and_classification_tie():
    peak, bw = 100.0, 10.0               # ridge = 10
    assert tr.attainable(2, peak, bw) == 20 and tr.attainable(50, peak, bw) == 100
    assert tr.classify(9.99, peak, bw) == "memory-bound"
    assert tr.classify(10.0, peak, bw) == "compute-bound"    # a tie sits on the roof


def test_intensity_rises_with_tile_and_is_about_b_over_element_size():
    n, s = 4096, tr.FP16_BYTES
    values = [tr.tiled_intensity(n, b, s) for b in (1, 4, 16, 64, 256)]
    assert values == sorted(values) and len(set(values)) == len(values)
    assert tr.tiled_intensity(n, 1, s) == pytest.approx(0.5, rel=1e-3)
    assert tr.tiled_intensity(n, 64, s) == pytest.approx(64 / s, rel=0.02)
    # one tile covering everything moves each element once: the best case
    assert tr.tiled_intensity(n, n, s) == pytest.approx(tr.gemm_intensity(n, n, n, s))


@pytest.mark.parametrize("kib, expected_b, expected_intensity", [(16, 52, 25.8), (64, 104, 50.6), (256, 209, 99.9), (1024, 418, 195.0)])
def test_tile_capacity_bound(kib, expected_b, expected_intensity):
    """Three b x b tiles must fit. Guards forgetting the C tile (2 b^2 gives 128 at 64 KiB) and a rounded-up
    root (105). The rows are the capacity table in the lesson."""
    s = tr.FP16_BYTES
    b = tr.max_tile(kib * 1024, s)
    assert tr.fast_bytes_needed(b, s) <= kib * 1024 < tr.fast_bytes_needed(b + 1, s)
    assert b == expected_b
    assert tr.tiled_intensity(4096, b, s) == pytest.approx(expected_intensity, abs=0.05)


def test_check_yourself_capacity_answer():
    """Check-yourself question 2: 48 KiB, 4-byte numbers -> b = 64, 15.9 FLOP/byte at n = 4096."""
    b = tr.max_tile(48 * 1024, 4)
    assert b == 64
    assert tr.tiled_intensity(4096, b, 4) == pytest.approx(15.9, abs=0.05)


@pytest.mark.parametrize("n, expected", [(2, 1 / 3), (16, 2.667), (256, 42.667), (4096, 682.667)])
def test_best_case_intensity_table_for_four_byte_floats(n, expected):
    """The lesson's best-case table, 2n / (3s) with s = 4. Guards a slip in the byte count or the 3n^2 numbers."""
    assert tr.gemm_intensity(n, n, n, 4) == pytest.approx(expected, abs=0.0005)


def test_smallest_tile_reaching_the_ridge_is_above_the_large_n_rule_for_two_causes():
    """Lesson quotes 342, not the rule of thumb 306. Guards quoting s * ridge as exact. Two causes:
    writing C alone moves 306 to 318; the whole-tile count (13 tiles per side give 151.7, 12 give 163.8)
    moves it to 342."""
    s, peak, bw = tr.FP16_BYTES, tr.A100_PEAK_FLOPS, tr.A100_BANDWIDTH
    ridge = tr.ridge(peak, bw)
    assert round(s * ridge) == 306
    assert math.ceil(tr.ridge_tile_if_tile_count_were_fractional(4096, s, peak, bw)) == 318
    b = tr.smallest_tile_for_ridge(4096, s, peak, bw)
    assert b == 342 and b > 318
    assert tr.tiled_intensity(4096, b, s) >= ridge > tr.tiled_intensity(4096, b - 1, s)
    assert tr.tiled_intensity(4096, 342, s) == pytest.approx(163.8, abs=0.05)    # 12 tiles per side
    assert tr.tiled_intensity(4096, 341, s) == pytest.approx(151.7, abs=0.05)    # 13 tiles per side
    assert tr.fast_bytes_needed(b, s) == 701_784


def test_decode_is_memory_bound_and_batching_crosses_the_ridge():
    s, peak, bw = tr.FP16_BYTES, tr.A100_PEAK_FLOPS, tr.A100_BANDWIDTH
    assert tr.gemm_intensity(1, 4096, 4096, s) == pytest.approx(1.0, abs=0.001)
    assert tr.classify(tr.gemm_intensity(1, 4096, 4096, s), peak, bw) == "memory-bound"
    assert tr.classify(tr.gemm_intensity(4096, 4096, 4096, s), peak, bw) == "compute-bound"
    m = tr.batch_to_reach_ridge(4096, 4096, s, peak, bw)
    assert m == 166
    assert tr.gemm_intensity(m - 1, 4096, 4096, s) < tr.ridge(peak, bw) <= tr.gemm_intensity(m, 4096, 4096, s)


@pytest.mark.parametrize("m, intensity, tflops, share", [
    (1, 1.00, 2.04, 0.7), (8, 7.97, 16.25, 5.2), (32, 31.51, 64.24, 20.6), (128, 120.47, 245.64, 78.7),
    (165, 152.70, 311.35, 99.8), (166, 153.55, 312.00, 100.0), (1024, 682.67, 312.00, 100.0)])
def test_decode_table_rows_match_the_lesson(m, intensity, tflops, share):
    """Guards the decode/prefill table drifting from the code that prints it."""
    peak, bw = tr.A100_PEAK_FLOPS, tr.A100_BANDWIDTH
    i = tr.gemm_intensity(m, 4096, 4096, tr.FP16_BYTES)
    att = tr.attainable(i, peak, bw)
    assert i == pytest.approx(intensity, abs=0.005)
    assert att / 1e12 == pytest.approx(tflops, abs=0.005)
    assert 100 * att / peak == pytest.approx(share, abs=0.05)


def test_decode_ceiling_for_a_seven_billion_parameter_model():
    """6.87 ms to read 14 GB of FP16 weights at 2,039 GB/s -> at most 146 tokens/s. Guards a bytes/params slip."""
    s, bw = tr.FP16_BYTES, tr.A100_BANDWIDTH
    assert tr.weight_read_seconds(7e9, s, bw) * 1e3 == pytest.approx(6.866, abs=0.001)
    assert tr.decode_tokens_per_s(7e9, s, bw) == pytest.approx(145.6, abs=0.1)


def test_naive_multiply_floors_on_the_a100():
    """The naive 4096 table: 1.019 TFLOP/s (0.327% of peak), floors 0.4405 ms and 134.83 ms (306 times),
    0.0494 ms at the minimum traffic. Guards time_floors and the bytes-per-element slip."""
    n, s, peak, bw = 4096, tr.FP16_BYTES, tr.A100_PEAK_FLOPS, tr.A100_BANDWIDTH
    naive, best = tr.naive_traffic(n, n, n), tr.min_traffic(n, n, n)
    assert tr.attainable(naive.intensity(s), peak, bw) == pytest.approx(1.019e12, abs=5e8)
    assert tr.attainable(naive.intensity(s), peak, bw) / peak * 100 == pytest.approx(0.327, abs=0.0005)
    t_math, t_naive = tr.time_floors(naive.flops, naive.nbytes(s), peak, bw)
    _, t_min = tr.time_floors(best.flops, best.nbytes(s), peak, bw)
    assert t_math * 1e3 == pytest.approx(0.4405, abs=0.00005)
    assert t_naive * 1e3 == pytest.approx(134.83, abs=0.005)
    assert t_naive / t_math == pytest.approx(306.07, abs=0.005)
    assert t_min * 1e3 == pytest.approx(0.0494, abs=0.00005) and t_min < t_math
    assert best.nbytes(s) == 100_663_296


def test_timing_helper_runs_and_reports_positive_rates():
    """Smoke test only: the reader runs this on their own machine, so no timing is asserted."""
    rates = tr.time_on_this_machine(n_small=8, n_mm=16, n_mv=32)
    assert len(rates) == 5 and all(v > 0 for v in rates.values())