From f08dda3bd4d3785e669e24ace9f6c84e91f30737 Mon Sep 17 00:00:00 2001 From: eusonice <98355549+eusonice@users.noreply.github.com> Date: Tue, 23 Jun 2026 16:29:10 -0700 Subject: [PATCH] test short context length attention kernel --- tests/attention/test_gqa_short_context.py | 613 ++++++++++++------ .../test_gqa_short_context_sweep_decode.py | 131 ++++ .../test_gqa_short_context_sweep_prefill.py | 145 +++++ 3 files changed, 695 insertions(+), 194 deletions(-) create mode 100644 tests/attention/test_gqa_short_context_sweep_decode.py create mode 100644 tests/attention/test_gqa_short_context_sweep_prefill.py diff --git a/tests/attention/test_gqa_short_context.py b/tests/attention/test_gqa_short_context.py index bdc85a2..30c1094 100644 --- a/tests/attention/test_gqa_short_context.py +++ b/tests/attention/test_gqa_short_context.py @@ -1,31 +1,17 @@ -"""Phase 1 spec test for Phase D: short-context GQA kernels -(ADR-0060 §B "Items from the long/short context split", item B.split.2). +"""Short-context GQA kernel tests — unified A1/A2/A4/B (ADR-0070). -The long-context kernels (§5.2 decode / §5.5 prefill) shard each KV -head row-wise across all CUBEs. At short context (S_kv < 256K), that -shard is too thin to feed the engines and the cube-level collective -overhead dominates. The short-context kernels drop cube-SP entirely: -each CUBE owns ``kv_per_cube`` *whole* KV heads, no S_kv sharding across -CUBEs. +Both ``_gqa_attention_prefill_short.py`` and +``_gqa_attention_decode_short.py`` are single unified kernels selected +at launch via ``kv_per_cube ∈ {1, 2, 4, 8}``: -Design (per AskUserQuestion answers in this session): - - PE-parallel heads: P PEs split into ``kv_per_cube`` groups, each - group does PE-SP across (P/kv_per_cube) PEs for one owned head. - - scratch_scope + tl.copy_to discipline mirrors the long kernels. - - One short kernel handles kv_per_cube ∈ {1, 2, 4, 8} via a parameter. + Mode kv_per_cube C group_size + ---- ----------- ------- ---------- + A1 1 h_kv P (=8) + A2 2 h_kv/2 P/2 (=4) + A4 4 h_kv/4 P/4 (=2) + B 8 1 1 -Group layout on the 2×4 PE grid: - kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge. - kv_per_cube=2, group=4 PEs (one row): row chain only. - kv_per_cube=4, group=2 PEs (adj cols): 1-step chain. - kv_per_cube=8, group=1 PE: no chain. - -After chain reduce-to-group-root, the group's root PE writes its -owned head's output. No inter-CUBE reduce. - -Phase 1 (this commit): tests only — production code lands in Phase 2. -All tests fail because the short kernels (``_gqa_attention_decode_short.py`` and -``_gqa_attention_prefill_short.py``) do not exist yet. +Tests are mode-parametrized over the four (kv_per_cube, C) tuples. """ from __future__ import annotations @@ -45,6 +31,17 @@ TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" D_HEAD = 64 DTYPE = "f16" +TILE_S_KV = 1024 +P = 8 +H_KV = 8 + +# (kv_per_cube, C) for the four modes. +MODES = [ + pytest.param(1, 8, id="A1"), + pytest.param(2, 4, id="A2"), + pytest.param(4, 2, id="A4"), + pytest.param(8, 1, id="B"), +] def _ccl_cfg(): @@ -61,166 +58,41 @@ def _count(op_log, name: str) -> int: return sum(1 for r in op_log if r.op_name == name) -# ── Decode short-context kernel ────────────────────────────────────── +# ── Prefill ────────────────────────────────────────────────────────── -def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int, - C: int, P: int, S_kv: int): - """Run the short-context decode kernel with PE-parallel heads. - - Layout (after design iteration — see Phase D failure-recovery): - Q: (T_q, h_q·D_HEAD) replicated; kernel reshapes byte-conservingly. - K, V: (h_kv·S_kv, D_HEAD) head-stacked, ``cube=row_wise, pe=row_wise`` - so each PE gets a contiguous (S_local, D_HEAD) chunk at its own - addressable shard (no partial reads needed). - O: replicated; each group root writes the full byte-conserving - (h_q·T_q, D_HEAD) result. - """ - from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import gqa_attention_decode_short_kernel # Phase 2 - +def _run_prefill(*, kv_per_cube: int, C: int, T_q: int, S_kv: int, + h_q: int = 8): + """Run the unified prefill kernel in the given mode.""" + from kernbench.benches._gqa_attention_prefill_short import ( + _validate_config as _validate_prefill_config, + gqa_attention_prefill_short_kernel, + ) + _validate_prefill_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) + n_tiles = S_kv // TILE_S_KV + Q_ROWS = kv_per_cube * T_q + Q_COLS = (h_q * D_HEAD) // kv_per_cube topo = resolve_topology(str(TOPOLOGY_DEFAULT)) def _bench_fn(ctx): configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) - dp_full = DPPolicy(cube="replicate", pe="replicate", - num_cubes=C, num_pes=P) - # Head-stacked KV with row_wise sharding so each PE's chunk is - # contiguous and exactly (S_local, D_HEAD) addressable. - dp_kv = DPPolicy(cube="row_wise", pe="row_wise", - num_cubes=C, num_pes=P) - q = ctx.zeros((1, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"q_short_{kv_per_cube}_{C}") - k = ctx.zeros((h_kv * S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"k_short_{kv_per_cube}_{C}") - v = ctx.zeros((h_kv * S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"v_short_{kv_per_cube}_{C}") - o = ctx.empty((1, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"o_short_{kv_per_cube}_{C}") + dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P) + dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P) + dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P) + q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q, + name=f"q_pre_kv{kv_per_cube}") + k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV), + dtype=DTYPE, dp=dp_kv, name=f"k_pre_kv{kv_per_cube}") + v = ctx.zeros((H_KV * S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"v_pre_kv{kv_per_cube}") + o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o, + name=f"o_pre_kv{kv_per_cube}") ctx.launch( - f"gqa_decode_short_{kv_per_cube}_{C}", - gqa_attention_decode_short_kernel, - q, k, v, o, - 1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), engine_factory=_engine_factory, - ) - - -def test_short_decode_smoke_kv_per_cube_2_C_4(): - """ADR-0060 §B.split.2 headline: kv_per_cube=2, C=4 — each CUBE - owns 2 heads (8 heads / 4 CUBEs). PE-parallel heads splits the 8 - PEs into 2 groups of 4, each group does PE-SP for one owned head. - Smoke: kernel completes.""" - result = _run_decode_short( - h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64, - ) - assert result.completion.ok, ( - f"short decode kv_per_cube=2 must complete; got {result.completion}" - ) - - -def test_short_decode_smoke_kv_per_cube_4_C_2(): - """kv_per_cube=4, C=2 — half the CUBEs participate, each owns 4 - heads, 4 PE groups of 2 PEs each.""" - result = _run_decode_short( - h_q=8, h_kv=8, kv_per_cube=4, C=2, P=8, S_kv=64, - ) - assert result.completion.ok, ( - f"short decode kv_per_cube=4 must complete; got {result.completion}" - ) - - -def test_short_decode_smoke_kv_per_cube_8_C_1(): - """kv_per_cube=8, C=1 — all heads on one CUBE. 8 PE groups of 1 PE - each → no chain reduce, each PE writes its head's output.""" - result = _run_decode_short( - h_q=8, h_kv=8, kv_per_cube=8, C=1, P=8, S_kv=64, - ) - assert result.completion.ok, ( - f"short decode kv_per_cube=8 must complete; got {result.completion}" - ) - - -def test_short_decode_dma_write_count_equals_h_kv(): - """ADR-0060 §B.split.2: each owned head produces exactly one - output (no inter-CUBE reduce). Total dma_writes = h_kv across all - participating CUBEs and groups. - - For h_kv=8, kv_per_cube=2, C=4: each CUBE writes 2 outputs → - 4 × 2 = 8 dma_writes total. - """ - result = _run_decode_short( - h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64, - ) - assert result.completion.ok - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 8, ( - f"short decode: expected 8 dma_writes (h_kv); got {n_writes}" - ) - - -def test_short_decode_no_inter_cube_traffic(): - """ADR-0060 §B.split.2: each head is fully owned by one CUBE → no - inter-CUBE reduce. The kernel must not invoke CUBE-level E/W IPCQ. - - Today's long-context kernel at C=4 emits ~12 inter-CUBE ipcq_copy - via direction "E"/"W". The short kernel must emit zero of those, - keeping all IPCQ traffic on the ``intra_*`` directions. - """ - result = _run_decode_short( - h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64, - ) - assert result.completion.ok - # Count IPCQ traffic that targeted the CUBE-level "E"/"W" directions. - inter_cube_ipcq = sum( - 1 for r in result.engine.op_log - if r.op_name == "ipcq_copy" - and r.params.get("direction") in ("E", "W") - ) - assert inter_cube_ipcq == 0, ( - f"short decode must have no inter-CUBE E/W IPCQ; got {inter_cube_ipcq}" - ) - - -# ── Prefill short-context kernel ───────────────────────────────────── - - -def _run_prefill_short(*, h_kv: int, kv_per_cube: int, - C: int, P: int, T_q: int, S_kv: int): - """Run the short-context prefill kernel with PE-parallel heads. - - Layout: same head-stacked K/V scheme as decode short, with Q/O - replicated and byte-conserving reshape inside the kernel. - """ - from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel # Phase 2 - - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) - dp_q = DPPolicy(cube="replicate", pe="replicate", - num_cubes=C, num_pes=P) - dp_kv = DPPolicy(cube="row_wise", pe="row_wise", - num_cubes=C, num_pes=P) - dp_o = DPPolicy(cube="replicate", pe="replicate", - num_cubes=C, num_pes=P) - q = ctx.zeros((T_q, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_q, name=f"q_pre_short_{kv_per_cube}_{C}") - k = ctx.zeros((h_kv * S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"k_pre_short_{kv_per_cube}_{C}") - v = ctx.zeros((h_kv * S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"v_pre_short_{kv_per_cube}_{C}") - o = ctx.empty((T_q, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_o, name=f"o_pre_short_{kv_per_cube}_{C}") - ctx.launch( - f"gqa_prefill_short_{kv_per_cube}_{C}", + f"gqa_prefill_kv{kv_per_cube}", gqa_attention_prefill_short_kernel, q, k, v, o, - T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube, + T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube, _auto_dim_remap=False, ) @@ -230,31 +102,384 @@ def _run_prefill_short(*, h_kv: int, kv_per_cube: int, ) -def test_short_prefill_smoke_kv_per_cube_2_C_4(): - """Short prefill kv_per_cube=2, C=4. Smoke: completes.""" - result = _run_prefill_short( - h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64, - ) - assert result.completion.ok, ( - f"short prefill kv_per_cube=2 must complete; got {result.completion}" +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_prefill_smoke(kv_per_cube, C): + """All four prefill modes complete on a single-tile mini config.""" + r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024) + assert r.completion.ok, ( + f"prefill kv_per_cube={kv_per_cube}: {r.completion}" ) -def test_short_prefill_no_ring_KV_traffic(): - """ADR-0060 §B.split.2: short prefill DOES NOT use Ring KV — each - CUBE owns its KV heads fully, no rotation. The kernel must not - emit any KV-rotation IPCQ traffic at the CUBE level. +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_prefill_qtile_split_dma_writes(kv_per_cube, C): + """Every PE in every active group stores its q-tile output: + dma_writes = group_size · kv_per_cube · C.""" + r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024) + assert r.completion.ok + group_size = P // kv_per_cube + expected = group_size * kv_per_cube * C + n_writes = _count(r.engine.op_log, "dma_write") + assert n_writes == expected, ( + f"prefill kv_per_cube={kv_per_cube}: expected {expected} " + f"dma_writes; got {n_writes}" + ) + + +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_prefill_kv_hbm_read_collapsed(kv_per_cube, C): + """IPCQ KV broadcast collapses HBM K/V reads to one per group per tile: + dma_reads = (P · C) Q-reads + (kv_per_cube · 2 · n_tiles · C) KV-reads. """ - result = _run_prefill_short( - h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64, + r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024) + assert r.completion.ok + n_tiles = 1024 // TILE_S_KV + expected = (P * C) + (kv_per_cube * 2 * n_tiles * C) + n_reads = _count(r.engine.op_log, "dma_read") + assert n_reads == expected, ( + f"prefill kv_per_cube={kv_per_cube}: expected {expected} dma_reads " + f"(Q + IPCQ-broadcasted KV); got {n_reads}" ) - assert result.completion.ok + + +def test_prefill_no_ring_KV_traffic(): + """Short prefill never uses Ring KV — no inter-CUBE E/W IPCQ.""" + r = _run_prefill(kv_per_cube=2, C=4, T_q=8, S_kv=1024) + assert r.completion.ok inter_cube_ipcq = sum( - 1 for r in result.engine.op_log - if r.op_name == "ipcq_copy" - and r.params.get("direction") in ("E", "W") + 1 for rec in r.engine.op_log + if rec.op_name == "ipcq_copy" + and rec.params.get("direction") in ("E", "W") ) assert inter_cube_ipcq == 0, ( - f"short prefill must have no Ring KV (no inter-CUBE E/W IPCQ); " + f"short prefill must have no inter-CUBE E/W IPCQ; " f"got {inter_cube_ipcq}" ) + + +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_prefill_multitile_scaling(kv_per_cube, C): + """Prefill dma_read scales linearly with n_tiles in every mode. + + Per-mode formula (IPCQ broadcast: cube's group-root PE reads K/V): + dma_reads = P·C + 2·kv_per_cube·C·n_tiles + = (Q: 1 per PE) + (K + V: 1 per group per tile) + """ + for S_kv in (1024, 2048, 4096): + r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=S_kv) + assert r.completion.ok, f"kv={kv_per_cube} S_kv={S_kv}" + n_tiles = S_kv // TILE_S_KV + expected = (P * C) + (2 * kv_per_cube * C * n_tiles) + n_reads = _count(r.engine.op_log, "dma_read") + assert n_reads == expected, ( + f"prefill kv={kv_per_cube} n_tiles={n_tiles}: expected " + f"{expected} dma_reads; got {n_reads}" + ) + + +# ── Decode ─────────────────────────────────────────────────────────── + + +def _run_decode(*, kv_per_cube: int, C: int, S_kv: int, h_q: int = 8): + """Run the unified decode kernel in the given mode.""" + from kernbench.benches._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 total elements = h_kv · S_kv · d_head; mode-invariant deploy. + 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()) + dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P) + dp_kv = DPPolicy(cube="row_wise", pe="row_wise", num_cubes=C, num_pes=P) + dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P) + q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q, + name=f"q_dec_kv{kv_per_cube}") + k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv, + name=f"k_dec_kv{kv_per_cube}") + v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, + name=f"v_dec_kv{kv_per_cube}") + o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o, + name=f"o_dec_kv{kv_per_cube}") + ctx.launch( + f"gqa_decode_kv{kv_per_cube}", + gqa_attention_decode_short_kernel, + q, k, v, o, + T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube, + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), engine_factory=_engine_factory, + ) + + +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_decode_smoke(kv_per_cube, C): + """All four decode modes complete at S_kv = 8K (min multi-tile-aligned + config that satisfies every mode's group-size constraint).""" + r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192) + assert r.completion.ok, ( + f"decode kv_per_cube={kv_per_cube}: {r.completion}" + ) + + +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_decode_chain_reduce_store_count(kv_per_cube, C): + """Decode chain-reduce: only group root (PE 0 per group) stores — + dma_writes = kv_per_cube · C (one per group root × cubes).""" + r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192) + assert r.completion.ok + expected = kv_per_cube * C + n_writes = _count(r.engine.op_log, "dma_write") + assert n_writes == expected, ( + f"decode kv_per_cube={kv_per_cube}: expected {expected} dma_writes " + f"(1 per group root); got {n_writes}" + ) + + +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_decode_multitile_per_pe(kv_per_cube, C): + """Decode multi-tile sweep path exercised in every mode. + + S_kv is chosen so each PE owns 2 tiles (n_tiles_per_pe == 2), + forcing the post-tile-0 sweep loop to run. + """ + group_size = P // kv_per_cube + S_kv = group_size * 2 * TILE_S_KV + r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv) + assert r.completion.ok, ( + f"decode kv={kv_per_cube} S_kv={S_kv}: {r.completion}" + ) + + +# ── ADR-0011 D-VA1 contract regression tests ─────────────────────── + + +def _per_cube_dma_busy(op_log): + """Sum of dma_read t_end-t_start grouped by cube index.""" + from collections import defaultdict + busy = defaultdict(float) + for rec in op_log: + if rec.op_name != "dma_read": + continue + cid = rec.component_id or "" + for part in cid.split("."): + if part.startswith("cube"): + try: + busy[int(part[4:])] += rec.t_end - rec.t_start + except ValueError: + pass + break + return dict(busy) + + +def _per_cube_large_read_addrs(op_log, *, min_nbytes: int = 1024): + """Group dma_read src_addrs (>= min_nbytes) by issuing cube.""" + from collections import defaultdict + per_cube: dict[int, set[int]] = defaultdict(set) + for rec in op_log: + if rec.op_name != "dma_read": + continue + if rec.params.get("nbytes", 0) < min_nbytes: + continue + cid = rec.component_id or "" + for part in cid.split("."): + if part.startswith("cube"): + try: + per_cube[int(part[4:])].add(rec.params["src_addr"]) + except ValueError: + pass + break + return per_cube + + +def _assert_cubes_disjoint(per_cube_addrs, label): + seen = sorted(per_cube_addrs.items()) + for i, (ci, ai) in enumerate(seen): + for cj, aj in seen[i + 1:]: + assert ai.isdisjoint(aj), ( + f"{label}: cube{ci} and cube{cj} share {len(ai & aj)} " + f"dma_read src_addrs; cubes must target disjoint HBM regions." + ) + + +def _assert_cube_balanced(busy, label, max_ratio=1.1): + if len(busy) <= 1: + return # single-cube modes auto-pass + vals = list(busy.values()) + ratio = max(vals) / min(vals) if min(vals) > 0 else 0 + assert ratio < max_ratio, ( + f"{label}: per-cube DMA_READ unbalanced: max/min = {ratio:.2f}× " + f"(expected < {max_ratio}×). Per-cube busy (μs): " + f"{[f'cube{c}={v/1000:.1f}' for c,v in sorted(busy.items())]}" + ) + + +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_decode_per_cube_disjoint_src_addrs(kv_per_cube, C): + """Decode: cubes must read from disjoint HBM regions (per ADR-0011 + D-VA1, kernel computes cube/PE shard offset from program_id). + + Regression guard: pre-fix all 64 PEs read offset 0 → all cubes + share the same src_addrs. + """ + r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192) + assert r.completion.ok + per_cube = _per_cube_large_read_addrs(r.engine.op_log) + assert len(per_cube) == C, ( + f"decode kv={kv_per_cube}: expected {C} cubes; got {len(per_cube)}" + ) + _assert_cubes_disjoint(per_cube, f"decode kv={kv_per_cube}") + + +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_decode_per_cube_dma_balanced(kv_per_cube, C): + """Decode: per-cube DMA_READ busy must be ~equal (cubes operate on + independent local HBM). max/min < 1.1× for multi-cube modes. + + Regression guard: pre-fix 11.5× ratio (cross-cube traffic to cube0). + """ + r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192) + assert r.completion.ok + busy = _per_cube_dma_busy(r.engine.op_log) + assert len(busy) == C, ( + f"decode kv={kv_per_cube}: expected {C} cubes; got {sorted(busy.keys())}" + ) + _assert_cube_balanced(busy, f"decode kv={kv_per_cube}") + + +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_prefill_per_cube_disjoint_src_addrs(kv_per_cube, C): + """Prefill: cubes must read from disjoint HBM regions.""" + r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024) + assert r.completion.ok + per_cube = _per_cube_large_read_addrs(r.engine.op_log) + assert len(per_cube) == C, ( + f"prefill kv={kv_per_cube}: expected {C} cubes; got {len(per_cube)}" + ) + _assert_cubes_disjoint(per_cube, f"prefill kv={kv_per_cube}") + + +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_prefill_per_cube_dma_balanced(kv_per_cube, C): + """Prefill: per-cube DMA_READ busy must be ~equal (cube-parallel).""" + r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024) + assert r.completion.ok + busy = _per_cube_dma_busy(r.engine.op_log) + assert len(busy) == C, ( + f"prefill kv={kv_per_cube}: expected {C} cubes; got {sorted(busy.keys())}" + ) + _assert_cube_balanced(busy, f"prefill kv={kv_per_cube}") + + +# ── Composite (second-level) variant smoke tests ──────────────────── + + +def _run_prefill_composite(*, kv_per_cube: int, C: int, T_q: int, + S_kv: int, h_q: int = 8): + """Same helper as _run_prefill but for the composite-GEMM kernel.""" + from kernbench.benches._gqa_attention_prefill_short_composite import ( + _validate_config as _validate_prefill_cmp_config, + gqa_attention_prefill_short_composite_kernel, + ) + _validate_prefill_cmp_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) + n_tiles = S_kv // TILE_S_KV + Q_ROWS = kv_per_cube * T_q + Q_COLS = (h_q * D_HEAD) // kv_per_cube + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P) + dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P) + dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P) + q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q, + name=f"q_pre_cmp_kv{kv_per_cube}") + k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV), + dtype=DTYPE, dp=dp_kv, name=f"k_pre_cmp_kv{kv_per_cube}") + v = ctx.zeros((H_KV * S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"v_pre_cmp_kv{kv_per_cube}") + o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o, + name=f"o_pre_cmp_kv{kv_per_cube}") + ctx.launch( + f"gqa_prefill_cmp_kv{kv_per_cube}", + gqa_attention_prefill_short_composite_kernel, + q, k, v, o, + T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube, + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), engine_factory=_engine_factory, + ) + + +def _run_decode_composite(*, kv_per_cube: int, C: int, S_kv: int, + h_q: int = 8): + """Same helper as _run_decode but for the composite-GEMM kernel.""" + from kernbench.benches._gqa_attention_decode_short_composite import ( + _validate_config as _validate_decode_cmp_config, + gqa_attention_decode_short_composite_kernel, + ) + T_q = 1 + _validate_decode_cmp_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()) + dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P) + dp_kv = DPPolicy(cube="row_wise", pe="row_wise", num_cubes=C, num_pes=P) + dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P) + q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q, + name=f"q_dec_cmp_kv{kv_per_cube}") + k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv, + name=f"k_dec_cmp_kv{kv_per_cube}") + v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, + name=f"v_dec_cmp_kv{kv_per_cube}") + o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o, + name=f"o_dec_cmp_kv{kv_per_cube}") + ctx.launch( + f"gqa_decode_cmp_kv{kv_per_cube}", + gqa_attention_decode_short_composite_kernel, + q, k, v, o, + T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube, + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), engine_factory=_engine_factory, + ) + + +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_prefill_composite_smoke(kv_per_cube, C): + """Composite-GEMM prefill completes in all 4 modes.""" + r = _run_prefill_composite(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024) + assert r.completion.ok, ( + f"prefill-composite kv_per_cube={kv_per_cube}: {r.completion}" + ) + + +@pytest.mark.parametrize("kv_per_cube,C", MODES) +def test_decode_composite_smoke(kv_per_cube, C): + """Composite-GEMM decode completes in all 4 modes.""" + r = _run_decode_composite(kv_per_cube=kv_per_cube, C=C, S_kv=8192) + assert r.completion.ok, ( + f"decode-composite kv_per_cube={kv_per_cube}: {r.completion}" + ) diff --git a/tests/attention/test_gqa_short_context_sweep_decode.py b/tests/attention/test_gqa_short_context_sweep_decode.py new file mode 100644 index 0000000..25e8fc9 --- /dev/null +++ b/tests/attention/test_gqa_short_context_sweep_decode.py @@ -0,0 +1,131 @@ +"""Decode 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 — chain reduce) + 5. KV cache space cost (per-cube bytes) + +Output: ``docs/sweeps/short_context_decode_sweep.csv`` (16 rows). + +Run: ``pytest tests/attention/test_gqa_short_context_sweep_decode.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_decode, +) + +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] + +CSV_OUT = (Path(__file__).resolve().parents[2] + / "docs" / "sweeps" / "short_context_decode_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) + + 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) + + 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_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_decode_mode_context_sweep(): + """Run 4 modes × 4 context lengths, dump 5 metrics to CSV.""" + rows = [] + for mode, kv_per_cube, C in MODES: + for S_kv in CONTEXT_LENGTHS: + print(f">>> DECODE {mode} S_kv={S_kv//1024}K " + f"(kv_per_cube={kv_per_cube}, C={C})", flush=True) + r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv, h_q=64) + assert r.completion.ok, ( + f"DECODE {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}") diff --git a/tests/attention/test_gqa_short_context_sweep_prefill.py b/tests/attention/test_gqa_short_context_sweep_prefill.py new file mode 100644 index 0000000..93ccd85 --- /dev/null +++ b/tests/attention/test_gqa_short_context_sweep_prefill.py @@ -0,0 +1,145 @@ +"""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}")