kernbench2/tests/attention/test_gqa_decode_batch_sweep.py

"""Batch GQA decode sweep — concurrent independent users on ONE SIP.

Fixes the target to a single 16-cube SIP and scales the batch B = number
of concurrent decode users, each with its own (Q, K, V) placed on a
disjoint cube group ``[u*C, u*C + C)`` via ``DPPolicy.cube_start`` and
addressed by the kernel's ``cube_base`` local-index parameter. All B
launches are deferred (``_defer_wait=True``) so they overlap on the shared
discrete-event clock; ``run_bench`` drains them together.

Per mapping the SIP holds ``16 // C`` users at once (A1=2, A2=4, A4=8,
B=16). We sweep B up to that capacity and record aggregate latency,
throughput (users / latency), and mean per-user latency, to see which
mapping's throughput scales best with batch.

Metric convention matches the corrected short-context sweeps:
``wall = max(t_end) - min(t_start)`` over the op log (deploy excluded).
"""
from __future__ import annotations

import csv
from pathlib import Path

import pytest

from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
from kernbench.policy.placement.dp import DPPolicy
from kernbench.runtime_api.bench_runner import run_bench
from kernbench.runtime_api.types import resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology

TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"

D_HEAD = 64
DTYPE = "f16"
TILE_S_KV = 1024
P = 8
H_KV = 8
CUBES_PER_SIP = 16

# (mode, kv_per_cube, C); users-per-SIP capacity = CUBES_PER_SIP // C.
MODES = [("A1", 1, 8), ("A2", 2, 4), ("A4", 4, 2), ("B", 8, 1)]
CONTEXT_LENGTHS = [8 * 1024]

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


def _ccl_cfg():
    return resolve_algorithm_config(
        load_ccl_config(), name="lrab_hierarchical_allreduce",
    )


def _engine_factory(t, d):
    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)


def _run_batch(*, kv_per_cube: int, C: int, S_kv: int, B: int, h_q: int = 64):
    """Launch B concurrent decode users on disjoint cube groups of SIP 0."""
    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import (
        _validate_config as _validate_decode_config,
        gqa_attention_decode_short_kernel,
    )
    T_q = 1
    _validate_decode_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P, C=C,
                            h_q=h_q, h_kv=H_KV, S_kv=S_kv)
    Q_ROWS = kv_per_cube * T_q
    Q_COLS = (h_q * D_HEAD) // kv_per_cube
    k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
    topo = resolve_topology(str(TOPOLOGY_DEFAULT))

    def _bench_fn(ctx):
        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
        # Deploy ALL users' tensors first. Each deploy blocks (advances the
        # clock), so if a launch were already submitted a later deploy would
        # drive it to completion and serialize the batch. Creating every
        # tensor before any launch keeps the deploys from touching kernels.
        users = []
        for u in range(B):
            cs = u * C
            dp_q = DPPolicy(cube="column_wise", pe="replicate",
                            num_cubes=C, num_pes=P, cube_start=cs)
            dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
                             num_cubes=C, num_pes=P, cube_start=cs)
            dp_o = DPPolicy(cube="column_wise", pe="replicate",
                            num_cubes=C, num_pes=P, cube_start=cs)
            q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
                          name=f"q_u{u}")
            k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
                          name=f"k_u{u}")
            v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
                          name=f"v_u{u}")
            o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
                          name=f"o_u{u}")
            users.append((u, cs, q, k, v, o))
        # Now submit every launch deferred (no blocking wait between them),
        # then drain together so all B run concurrently on the shared clock.
        pending = []
        for u, cs, q, k, v, o in users:
            pending += ctx.launch(
                f"gqa_decode_u{u}",
                gqa_attention_decode_short_kernel,
                q, k, v, o,
                T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube, cs,
                _auto_dim_remap=False,
                _defer_wait=True,
            )
        for h, _sip_id, meta in pending:
            ctx.wait(h, _meta=meta)

    return run_bench(
        topology=topo, bench_fn=_bench_fn,
        device=resolve_device("sip:0"), engine_factory=_engine_factory,
    )


def _cube_of(rec) -> int:
    """Physical cube index of an op record, or -1 if not a PE op."""
    cid = rec.component_id or ""
    if ".cube" in cid:
        try:
            return int(cid.split(".cube")[1].split(".")[0])
        except (ValueError, IndexError):
            return -1
    return -1


def _metrics(r, *, mode, kv_per_cube, C, S_kv, B):
    op_log = r.engine.op_log
    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)

    # Per-user latency: group ops by which user's cube band they touch.
    per_user = []
    for u in range(B):
        lo, hi = u * C, u * C + C
        starts = [rec.t_start for rec in op_log if lo <= _cube_of(rec) < hi]
        ends = [rec.t_end for rec in op_log if lo <= _cube_of(rec) < hi]
        if starts and ends:
            per_user.append(max(ends) - min(starts))
    mean_user_ns = sum(per_user) / len(per_user) if per_user else 0.0

    agg_us = wall_ns / 1000
    return {
        "mode": mode,
        "kv_per_cube": kv_per_cube,
        "C": C,
        "S_kv": S_kv,
        "B": B,
        "capacity": CUBES_PER_SIP // C,
        "agg_latency_us": agg_us,
        "mean_user_latency_us": mean_user_ns / 1000,
        "throughput_users_per_us": (B / agg_us) if agg_us else 0.0,
    }


@pytest.mark.slow
def test_decode_batch_sweep():
    """Sweep 4 modes × contexts × B ∈ {1, 2, capacity}; dump batch CSV.

    Throughput is near-linear in B (aggregate latency stays within a few
    percent of the single-user latency across the batch), so the endpoints
    plus a low point capture the scaling; the high-B runs of the dense
    mappings are the expensive ones and full 1..capacity granularity is not
    needed for the trade-off.
    """
    rows = []
    for mode, kv_per_cube, C in MODES:
        capacity = CUBES_PER_SIP // C
        batch_sizes = sorted({1, 2, capacity})
        for S_kv in CONTEXT_LENGTHS:
            for B in batch_sizes:
                print(f">>> BATCH {mode} S_kv={S_kv//1024}K B={B}/{capacity}",
                      flush=True)
                r = _run_batch(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv, B=B)
                assert r.completion.ok, (
                    f"BATCH {mode} S_kv={S_kv} B={B}: {r.completion}"
                )
                m = _metrics(r, mode=mode, kv_per_cube=kv_per_cube,
                             C=C, S_kv=S_kv, B=B)
                rows.append(m)
                print(f"    agg={m['agg_latency_us']:.2f}us  "
                      f"user={m['mean_user_latency_us']:.2f}us  "
                      f"tput={m['throughput_users_per_us']:.3f} u/us",
                      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()
        w.writerows(rows)
    print(f"\nWrote {CSV_OUT} ({len(rows)} rows)")