kernbench2/tests/attention/test_gqa_short_context_sweep_prefill.py

"""Prefill mode-comparison sweep across context lengths.

Measures per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}):
  1. Wall clock latency (μs)
  2. Hardware utilization
       - GEMM engine utilization (Σ gemm time / wall × n_pe)
       - MATH op count
  3. HBM bandwidth utilization
       (Σ dma_read+write bytes / (wall × n_pe × peak_pe_hbm_bw))
  4. Communication cost (IPCQ bytes — broadcast)
  5. KV cache space cost (per-cube bytes)

Sliced-prefill convention: T_q = 8.

Output: ``docs/sweeps/prefill_sweep.csv`` (16 rows).

Run: ``pytest tests/attention/test_gqa_short_context_sweep_prefill.py -m slow``
"""
from __future__ import annotations

import csv
from pathlib import Path

import pytest

from tests.attention.test_gqa_short_context import (
    D_HEAD, H_KV, P, TILE_S_KV, _run_prefill,
)

# Peak HBM BW per PE — derived from topology (cube total / PE count).
# topology.yaml: hbm_total_bw_gbs = 1024.0 per cube, 8 PE per cube.
PEAK_PE_HBM_BPS = 1024.0 / 8  # 128 bytes/ns

MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
            "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
            "max_elem"}

MODES = [
    ("A1", 1, 8),
    ("A2", 2, 4),
    ("A4", 4, 2),
    ("B",  8, 1),
]

CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]

T_Q_PREFILL = 8

CSV_OUT = (Path(__file__).resolve().parents[2]
           / "docs" / "sweeps" / "short_context_prefill_sweep.csv")


def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
    op_log = r.engine.op_log
    wall_ns = max(rec.t_end for rec in op_log)

    # Per-PE detection (cube.pe)
    active_pes = set()
    for rec in op_log:
        cid = rec.component_id or ""
        if cid and ".cube" in cid and ".pe" in cid:
            parts = cid.split(".")
            active_pes.add(".".join(parts[:3]))
    n_pe = len(active_pes)

    # Stage time / count / bytes
    gemm_time_ns = 0.0
    math_count = 0
    dma_read_bytes = 0
    dma_write_bytes = 0
    ipcq_bytes = 0
    for rec in op_log:
        name = rec.op_name
        nbytes = rec.params.get("nbytes", 0) or 0
        if name == "gemm_f16":
            gemm_time_ns += rec.t_end - rec.t_start
        elif name in MATH_OPS:
            math_count += 1
        elif name == "dma_read":
            dma_read_bytes += nbytes
        elif name == "dma_write":
            dma_write_bytes += nbytes
        elif name == "ipcq_copy":
            ipcq_bytes += nbytes

    gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
    hbm_bw_util = (
        (dma_read_bytes + dma_write_bytes)
        / (wall_ns * n_pe * PEAK_PE_HBM_BPS)
    ) if n_pe else 0.0

    # KV cache space — per-cube bytes (K + V, f16)
    # Per cube owns kv_per_cube heads × S_kv tokens × d_head × 2 (K+V) × 2 (f16)
    kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2

    return {
        "mode": mode,
        "kv_per_cube": kv_per_cube,
        "C": C,
        "S_kv": S_kv,
        "wall_us": wall_ns / 1000,
        "n_pe": n_pe,
        "gemm_util": gemm_util,
        "math_count": math_count,
        "hbm_bw_util": hbm_bw_util,
        "hbm_read_mb": dma_read_bytes / (1024 * 1024),
        "hbm_write_kb": dma_write_bytes / 1024,
        "ipcq_kb": ipcq_bytes / 1024,
        "kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
    }


@pytest.mark.slow
def test_prefill_mode_context_sweep():
    """Run 4 modes × 4 context lengths, dump 5 metrics to CSV.

    Smoke-asserts each run completes; the measurement output is the CSV.
    """
    rows = []
    for mode, kv_per_cube, C in MODES:
        for S_kv in CONTEXT_LENGTHS:
            print(f">>> PREFILL {mode} S_kv={S_kv//1024}K "
                  f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
            r = _run_prefill(kv_per_cube=kv_per_cube, C=C,
                             T_q=T_Q_PREFILL, S_kv=S_kv, h_q=64)
            assert r.completion.ok, (
                f"PREFILL {mode} S_kv={S_kv}: {r.completion}"
            )
            m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
                                 C=C, S_kv=S_kv)
            rows.append(m)
            print(f"    wall={m['wall_us']:.1f}μs  "
                  f"gemm_util={m['gemm_util']:.3f}  "
                  f"hbm_bw_util={m['hbm_bw_util']:.3f}  "
                  f"ipcq={m['ipcq_kb']:.1f}KB  "
                  f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
                  flush=True)

    CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
    with CSV_OUT.open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        for row in rows:
            w.writerow(row)
    print(f"\nWrote {len(rows)} rows to {CSV_OUT}")