diff --git a/src/kernbench/benches/_gqa_attention_decode_long.py b/src/kernbench/benches/_gqa_attention_decode_long.py deleted file mode 100644 index 77bd132..0000000 --- a/src/kernbench/benches/_gqa_attention_decode_long.py +++ /dev/null @@ -1,353 +0,0 @@ -"""GQA fused-attention decode kernel — long context (ADR-0060). - -Each rank holds an ``S_local = S_kv / (C·P)`` slice of K, V and the full -Q (replicated). The local attention is computed via an S_kv-axis tile -sweep (ADR-0063 §A.2) so per-rank scratch is bounded by ``TILE_S_KV`` -regardless of ``S_local``. The partial ``(m, ℓ, O)`` is then reduced -via a 2-level pattern (intra-CUBE row+col, then inter-CUBE), and the -root writes the final output. - -Inter-CUBE reduce (selected by the ``sub_w`` launch arg): - - ``sub_w == 0`` (default): 1D chain along ``W``, root at CUBE 0. - Used for ``C ∈ {1, 4}`` panels. - - ``sub_w >= 2``: ADR-0060 §4.2 prescribed **lrab-adapted center-root - mesh reduce** over a ``sub_w × sub_h`` sub-mesh (``sub_h = C // - sub_w``). Phase 1 row reduce converges at ``root_col = sub_w//2``; - Phase 2 col reduce on ``root_col`` converges at - ``root_row = sub_h//2``. Root cube id: - ``root_row * sub_w + root_col``. Used for the C=8 single-KV-group - LLaMA-3.1-70B target (sub_w=4, sub_h=2, root=cube 6). - Requires ``sub_w >= 2 and sub_h >= 2`` — degenerate 1×N / N×1 - layouts must use the 1D-chain path with ``sub_w=0``. - -Topology / SFR: - - Requires ``configure_sfr_intercube_multisip`` when ``P > 1`` or - ``C > 1`` (provides disjoint ``intra_*`` and ``E/W/N/S`` namespaces). - - Intra-CUBE PEs are arranged as a 2×4 grid (no wrap). - - Inter-CUBE CUBEs: 1D row (``sub_w == 0``) or rectangular sub-mesh - of the SIP's 4×4 CUBE mesh starting at origin (0, 0) with - ``sub_w == mesh_w`` (``sub_w >= 2``). - -Layout caveats: - - GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``). - - K is loaded as ``(d_head, S_local)`` via byte-conserving reshape of - the deployed ``(S_local, h_kv·d_head)`` slice — correct for zero / - symmetric inputs (ADR-0060 §3 reshape-not-transpose caveat). -""" -from __future__ import annotations - - -TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). - - -def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): - """Online-softmax merge of two partial ``(m, ℓ, O)`` triples.""" - m_new = tl.maximum(m_local, m_other) - scale_old = tl.exp(m_local - m_new) - scale_new = tl.exp(m_other - m_new) - l_new = l_local * scale_old + l_other * scale_new - O_new = O_local * scale_old + O_other * scale_new - return m_new, l_new, O_new - - -def gqa_attention_decode_long_kernel( - q_ptr: int, - k_ptr: int, - v_ptr: int, - o_ptr: int, - T_q: int, - S_kv: int, - h_q: int, - h_kv: int, - d_head: int, - C: int, - P: int, - sub_w: int = 0, - *, - tl, -) -> None: - """GQA decode with M-fold + S_kv tile sweep + 2-level reduce-to-root. - - Tensor layout: - Q : (T_q, h_q · d_head) replicated on every rank; loaded as - (G·T_q, d_head) — byte-conserving and math-correct for T_q=1. - K : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank - loads its (d_head, S_local) slice via byte-conserving reshape. - V : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank - loads its (S_local, d_head) slice. - O : (T_q, h_q · d_head) — only PE 0 of the root CUBE stores - (CUBE 0 for ``sub_w=0``; geometric center cube for ``sub_w>=2``). - """ - if sub_w > 0: - sub_h = C // sub_w - if sub_w < 2 or sub_h < 2 or sub_w * sub_h != C: - raise ValueError( - f"sub_w={sub_w} requires sub_w>=2 and sub_h>=2 and " - f"sub_w*sub_h==C; got sub_h={sub_h}, C={C}. " - "Use sub_w=0 for the 1D-chain path." - ) - root_col = sub_w // 2 - root_row = sub_h // 2 - root_cube = root_row * sub_w + root_col - else: - sub_h = 0 - root_col = 0 - root_row = 0 - root_cube = 0 - - G = h_q // h_kv - n_ranks = C * P - S_local = S_kv // n_ranks - pe_id = tl.program_id(axis=0) - cube_id = tl.program_id(axis=1) - - # ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ── - Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16") - n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV - KV_ROW_BYTES = d_head * 2 # f16 - - # Tile 0: establishes persistent (m_local, l_local, O_local). - # - # Cannot be folded into the Tiles 1..N loop (kernbench-only limitation): - # - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``, - # otherwise scope teardown discards them before the next tile's - # merge can read them; - # - kernbench has no scratch-backed initializer — ``tl.zeros`` / - # ``tl.full`` return addr=0 handles with no backing storage, so - # they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0). - # So Tile 0 computes the initial running state directly; Tiles 1..N - # fold into it. Triton port: limitation does not apply (SSA tensors - # stay live across iterations) — a single unified loop suffices. - tile_s0 = min(TILE_S_KV, S_local) - K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16") - V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16") - scores = tl.dot(Q, K_T) - m_local = tl.max(scores, axis=-1) - centered = scores - m_local - exp_scores = tl.exp(centered) - l_local = tl.sum(exp_scores, axis=-1) - O_local = tl.dot(exp_scores, V) - - # Tiles 1..n_tiles-1: fold into running state via online-softmax merge. - # Triton port: drop the ``with tl.scratch_scope():`` line and replace - # each ``copy_to`` with a Python rebind. - for tile_idx in range(1, n_tiles): - tile_start = tile_idx * TILE_S_KV - tile_s = min(TILE_S_KV, S_local - tile_start) - with tl.scratch_scope(): - K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES, - shape=(d_head, tile_s), dtype="f16") - V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES, - shape=(tile_s, d_head), dtype="f16") - scores_t = tl.dot(Q, K_T_t) - m_tile = tl.max(scores_t, axis=-1) - centered_t = scores_t - m_tile - exp_scores_t = tl.exp(centered_t) - l_tile = tl.sum(exp_scores_t, axis=-1) - O_tile = tl.dot(exp_scores_t, V_t) - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - - # ── Communication: chain reduce-to-root at PE 0 of CUBE 0 ── - PE_GRID_COLS = 4 - pe_col = pe_id % PE_GRID_COLS - pe_row = pe_id // PE_GRID_COLS - pe_cols_used = min(PE_GRID_COLS, P) - pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS - - # Level-2 row chain (intra-CUBE, along intra_W, leftward). - if pe_cols_used > 1: - if pe_col < pe_cols_used - 1: - with tl.scratch_scope(): - m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16") - l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16") - O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16") - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - if pe_col > 0: - tl.send(dir="intra_W", src=m_local) - tl.send(dir="intra_W", src=l_local) - tl.send(dir="intra_W", src=O_local) - - # Level-2 col bridge (intra-CUBE, along intra_N, row-1 → row-0). - if pe_col == 0 and pe_rows_used > 1: - if pe_row < pe_rows_used - 1: - with tl.scratch_scope(): - m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16") - l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16") - O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16") - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - if pe_row > 0: - tl.send(dir="intra_N", src=m_local) - tl.send(dir="intra_N", src=l_local) - tl.send(dir="intra_N", src=O_local) - - # Level-1 inter-CUBE reduce (only PE 0 of each CUBE participates). - if pe_id == 0 and C > 1 and sub_w == 0: - # 1D chain along W, leftward; root at CUBE 0. - if cube_id < C - 1: - with tl.scratch_scope(): - m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16") - l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16") - O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16") - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - if cube_id > 0: - tl.send(dir="W", src=m_local) - tl.send(dir="W", src=l_local) - tl.send(dir="W", src=O_local) - elif pe_id == 0 and sub_w > 0: - # ── ADR-0060 §4.2 lrab-adapted center-root mesh reduce ── - # Adapts Phases 1-2 of lrab_hierarchical_allreduce.py: - # bidirectional row reduce converges at root_col, bidirectional - # col reduce on root_col converges at root_row. Reduce-only - # (Phases 3-5 dropped: no inter-SIP, no broadcast — attention - # needs the answer at one rank). Plain ``+`` replaced with the - # log-sum-exp ``_merge_running`` for (m, ℓ, O). - row = cube_id // sub_w - col = cube_id % sub_w - - # Phase 1: row reduce — converge at col == root_col. - if col == 0 and root_col > 0: - tl.send(dir="E", src=m_local) - tl.send(dir="E", src=l_local) - tl.send(dir="E", src=O_local) - elif 0 < col < root_col: - with tl.scratch_scope(): - m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16") - l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16") - O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16") - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - tl.send(dir="E", src=m_local) - tl.send(dir="E", src=l_local) - tl.send(dir="E", src=O_local) - elif col == root_col: - if root_col > 0: - with tl.scratch_scope(): - m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16") - l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16") - O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16") - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - if sub_w - 1 > root_col: - with tl.scratch_scope(): - m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16") - l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16") - O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16") - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - elif root_col < col < sub_w - 1: - with tl.scratch_scope(): - m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16") - l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16") - O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16") - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - tl.send(dir="W", src=m_local) - tl.send(dir="W", src=l_local) - tl.send(dir="W", src=O_local) - elif col == sub_w - 1 and sub_w - 1 > root_col: - tl.send(dir="W", src=m_local) - tl.send(dir="W", src=l_local) - tl.send(dir="W", src=O_local) - - # Phase 2: col reduce on col == root_col — converge at row == root_row. - if col == root_col: - if row == 0 and root_row > 0: - tl.send(dir="S", src=m_local) - tl.send(dir="S", src=l_local) - tl.send(dir="S", src=O_local) - elif 0 < row < root_row: - with tl.scratch_scope(): - m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16") - l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16") - O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16") - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - tl.send(dir="S", src=m_local) - tl.send(dir="S", src=l_local) - tl.send(dir="S", src=O_local) - elif row == root_row: - if root_row > 0: - with tl.scratch_scope(): - m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16") - l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16") - O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16") - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - if sub_h - 1 > root_row: - with tl.scratch_scope(): - m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16") - l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16") - O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16") - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - elif root_row < row < sub_h - 1: - with tl.scratch_scope(): - m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16") - l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16") - O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16") - m_new, l_new, O_new = _merge_running( - m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, - ) - tl.copy_to(m_local, m_new) - tl.copy_to(l_local, l_new) - tl.copy_to(O_local, O_new) - tl.send(dir="N", src=m_local) - tl.send(dir="N", src=l_local) - tl.send(dir="N", src=O_local) - elif row == sub_h - 1 and sub_h - 1 > root_row: - tl.send(dir="N", src=m_local) - tl.send(dir="N", src=l_local) - tl.send(dir="N", src=O_local) - - # ── Final normalise + store (root only) ── - if pe_id == 0 and cube_id == root_cube: - O_final = O_local / l_local - tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_repl_pe_sp.py b/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_repl_pe_sp.py index 7eb8957..224dc1c 100644 --- a/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_repl_pe_sp.py +++ b/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_repl_pe_sp.py @@ -25,12 +25,20 @@ Topology / SFR: """ from __future__ import annotations -from kernbench.benches._gqa_attention_decode_long import _merge_running - TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). +def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): + """Online-softmax merge of two partial ``(m, ℓ, O)`` triples.""" + m_new = tl.maximum(m_local, m_other) + scale_old = tl.exp(m_local - m_new) + scale_new = tl.exp(m_other - m_new) + l_new = l_local * scale_old + l_other * scale_new + O_new = O_local * scale_old + O_other * scale_new + return m_new, l_new, O_new + + def gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel( q_ptr: int, k_ptr: int, diff --git a/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py b/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py index e9201a7..4a44f11 100644 --- a/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py +++ b/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py @@ -1,26 +1,60 @@ """GQA decode kernel — Case 4 (Cube-SP × PE-SP). ★ slide-11 optimal. -This is a thin entry-point wrapper around the underlying decode_long -lrab-adapted center-root kernel with ``sub_w`` baked to 4 (the C=8 -single-KV-head-group geometry: 4×2 cube sub-mesh, root cube 6). The -math lives in ``_gqa_attention_decode_long.py`` (which still serves -non-4-cases panels via ``sub_w=0`` and the headline GQA bench); this -file exists so the 4-cases bench refers to all four cases via -parallel ``_gqa_attention_decode_.py`` kernel-file names. - Per GQA_full_deck.pptx slide 11: - K, V split 64-way (cube=row_wise, pe=row_wise); each rank owns - S_local = S_kv / (C·P). - - Two-level reduce on the partial (m, ℓ, O): intra-CUBE 8-way - (row chain + col bridge over the 2×4 PE grid) then inter-CUBE - 2-phase lrab over the 4×2 cube sub-mesh — converges at cube 6. - - PE 0 of CUBE 6 stores O. + ``S_local = S_kv / (C·P)``. + - Local attention via per-rank S_kv-axis tile sweep (ADR-0063 §A.2) + keeps scratch bounded by ``TILE_S_KV``. + - Two-level reduce on the partial ``(m, ℓ, O)``: + • intra-CUBE 8-way (row chain along intra_W + col bridge along + intra_N) over the 2×4 PE grid; + • inter-CUBE 2-phase lrab over the 4×2 CUBE sub-mesh (ADR-0060 + §4.2 lrab-adapted center-root reduce): Phase 1 row reduce + converges at root_col=2; Phase 2 col reduce on root_col + converges at root_row=1; root cube id = 6. + - PE 0 of CUBE 6 stores ``O``. + +Deviation from slide 13: slide prescribes AllReduce (every rank has +the answer); the kernel does reduce-to-root (only cube 6 has it) +per ADR-0060 §4. Treated as the kernbench Case-4 baseline. + +Tensor layout: + Q : (T_q, h_q · d_head) replicated; loaded as (G·T_q, d_head). + K : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank + loads its (d_head, S_local) slice via byte-conserving reshape + of (S_local, h_kv·d_head) — correct for zero / symmetric + inputs (ADR-0060 §3 reshape-not-transpose caveat). + V : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank + loads its (S_local, d_head) slice. + O : (T_q, h_q · d_head) — only PE 0 of CUBE 6 stores. + +Topology / SFR: + - Requires ``configure_sfr_intercube_multisip`` (provides disjoint + ``intra_*`` and ``E/W/N/S`` namespaces). + - Intra-CUBE PEs arranged as a 2×4 grid (no wrap). + - Inter-CUBE: 4×2 CUBE sub-mesh starting at origin (0, 0). """ from __future__ import annotations -from kernbench.benches._gqa_attention_decode_long import ( - gqa_attention_decode_long_kernel, -) + +TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). + +# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh). +_SUB_W = 4 +_SUB_H = 2 +_ROOT_COL = _SUB_W // 2 # 2 +_ROOT_ROW = _SUB_H // 2 # 1 +_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL # 6 + + +def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): + """Online-softmax merge of two partial ``(m, ℓ, O)`` triples.""" + m_new = tl.maximum(m_local, m_other) + scale_old = tl.exp(m_local - m_new) + scale_new = tl.exp(m_other - m_new) + l_new = l_local * scale_old + l_other * scale_new + O_new = O_local * scale_old + O_other * scale_new + return m_new, l_new, O_new def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel( @@ -39,8 +73,223 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel( tl, ) -> None: """Case-4 decode: Cube-SP × PE-SP, lrab center-root at cube 6.""" - gqa_attention_decode_long_kernel( - q_ptr, k_ptr, v_ptr, o_ptr, - T_q, S_kv, h_q, h_kv, d_head, C, P, 4, # sub_w=4 baked in - tl=tl, - ) + G = h_q // h_kv + n_ranks = C * P + S_local = S_kv // n_ranks + pe_id = tl.program_id(axis=0) + cube_id = tl.program_id(axis=1) + + # ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ── + Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16") + n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV + KV_ROW_BYTES = d_head * 2 # f16 + + # Tile 0: establishes persistent (m_local, l_local, O_local). Cannot + # fold into Tiles 1..N loop (kernbench-only): persistent tensors + # must live outside tl.scratch_scope or scope teardown discards + # them; there's no scratch-backed (-inf, 0, 0) initializer. + tile_s0 = min(TILE_S_KV, S_local) + K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16") + V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16") + scores = tl.dot(Q, K_T) + m_local = tl.max(scores, axis=-1) + centered = scores - m_local + exp_scores = tl.exp(centered) + l_local = tl.sum(exp_scores, axis=-1) + O_local = tl.dot(exp_scores, V) + + for tile_idx in range(1, n_tiles): + tile_start = tile_idx * TILE_S_KV + tile_s = min(TILE_S_KV, S_local - tile_start) + with tl.scratch_scope(): + K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES, + shape=(d_head, tile_s), dtype="f16") + V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES, + shape=(tile_s, d_head), dtype="f16") + scores_t = tl.dot(Q, K_T_t) + m_tile = tl.max(scores_t, axis=-1) + centered_t = scores_t - m_tile + exp_scores_t = tl.exp(centered_t) + l_tile = tl.sum(exp_scores_t, axis=-1) + O_tile = tl.dot(exp_scores_t, V_t) + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + + # ── Intra-CUBE reduce: row chain (intra_W) + col bridge (intra_N) ── + PE_GRID_COLS = 4 + pe_col = pe_id % PE_GRID_COLS + pe_row = pe_id // PE_GRID_COLS + pe_cols_used = min(PE_GRID_COLS, P) + pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS + + if pe_cols_used > 1: + if pe_col < pe_cols_used - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_col > 0: + tl.send(dir="intra_W", src=m_local) + tl.send(dir="intra_W", src=l_local) + tl.send(dir="intra_W", src=O_local) + + if pe_col == 0 and pe_rows_used > 1: + if pe_row < pe_rows_used - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_row > 0: + tl.send(dir="intra_N", src=m_local) + tl.send(dir="intra_N", src=l_local) + tl.send(dir="intra_N", src=O_local) + + # ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ── + # Only PE 0 of each CUBE participates. Adapts Phases 1-2 of + # lrab_hierarchical_allreduce.py: bidirectional row reduce converges + # at root_col; bidirectional col reduce on root_col converges at + # root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``. + if pe_id == 0: + row = cube_id // _SUB_W + col = cube_id % _SUB_W + + # Phase 1: row reduce — converge at col == _ROOT_COL. + if col == 0: + tl.send(dir="E", src=m_local) + tl.send(dir="E", src=l_local) + tl.send(dir="E", src=O_local) + elif 0 < col < _ROOT_COL: + with tl.scratch_scope(): + m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + tl.send(dir="E", src=m_local) + tl.send(dir="E", src=l_local) + tl.send(dir="E", src=O_local) + elif col == _ROOT_COL: + with tl.scratch_scope(): + m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + with tl.scratch_scope(): + m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + elif _ROOT_COL < col < _SUB_W - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + tl.send(dir="W", src=m_local) + tl.send(dir="W", src=l_local) + tl.send(dir="W", src=O_local) + elif col == _SUB_W - 1: + tl.send(dir="W", src=m_local) + tl.send(dir="W", src=l_local) + tl.send(dir="W", src=O_local) + + # Phase 2: col reduce on col == _ROOT_COL — converge at row == _ROOT_ROW. + if col == _ROOT_COL: + if row == 0: + tl.send(dir="S", src=m_local) + tl.send(dir="S", src=l_local) + tl.send(dir="S", src=O_local) + elif 0 < row < _ROOT_ROW: + with tl.scratch_scope(): + m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + tl.send(dir="S", src=m_local) + tl.send(dir="S", src=l_local) + tl.send(dir="S", src=O_local) + elif row == _ROOT_ROW: + with tl.scratch_scope(): + m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if _SUB_H - 1 > _ROOT_ROW: + with tl.scratch_scope(): + m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + elif _ROOT_ROW < row < _SUB_H - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + tl.send(dir="N", src=m_local) + tl.send(dir="N", src=l_local) + tl.send(dir="N", src=O_local) + elif row == _SUB_H - 1 and _SUB_H - 1 > _ROOT_ROW: + tl.send(dir="N", src=m_local) + tl.send(dir="N", src=l_local) + tl.send(dir="N", src=O_local) + + # ── Final normalise + store (root only: PE 0 of CUBE 6) ── + if pe_id == 0 and cube_id == _ROOT_CUBE: + O_final = O_local / l_local + tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_tp.py b/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_tp.py index 5e96ddb..eb0ef22 100644 --- a/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_tp.py +++ b/src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_tp.py @@ -27,12 +27,20 @@ Topology / SFR: """ from __future__ import annotations -from kernbench.benches._gqa_attention_decode_long import _merge_running - TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). +def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): + """Online-softmax merge of two partial ``(m, ℓ, O)`` triples.""" + m_new = tl.maximum(m_local, m_other) + scale_old = tl.exp(m_local - m_new) + scale_new = tl.exp(m_other - m_new) + l_new = l_local * scale_old + l_other * scale_new + O_new = O_local * scale_old + O_other * scale_new + return m_new, l_new, O_new + + def gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel( q_ptr: int, k_ptr: int, diff --git a/src/kernbench/benches/_gqa_attention_decode_opt2.py b/src/kernbench/benches/_gqa_attention_decode_opt2.py index 669564d..05de3b3 100644 --- a/src/kernbench/benches/_gqa_attention_decode_opt2.py +++ b/src/kernbench/benches/_gqa_attention_decode_opt2.py @@ -37,9 +37,10 @@ def gqa_attention_decode_opt2_kernel( ) -> None: """Single-rank (C=P=1) GQA decode using the opt2 two-composite form. - Layout mirrors ``gqa_attention_decode_long_kernel`` (M-fold Q, K loaded - as ``(d_head, S_local)``). The KV slice is split into two sub-tiles so - the second one drives the ``softmax_merge`` recipe composite. + Layout mirrors ``gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel`` + (M-fold Q, K loaded as ``(d_head, S_local)``). The KV slice is split + into two sub-tiles so the second one drives the ``softmax_merge`` + recipe composite. """ G = h_q // h_kv S_local = S_kv // (C * P) diff --git a/src/kernbench/benches/milestone_gqa_headline.py b/src/kernbench/benches/milestone_gqa_headline.py index bd27bff..d213838 100644 --- a/src/kernbench/benches/milestone_gqa_headline.py +++ b/src/kernbench/benches/milestone_gqa_headline.py @@ -1,29 +1,26 @@ -"""milestone-gqa-headline bench: real GQA + 2-level SP + Ring KV. +"""milestone-gqa-headline bench: real GQA prefill — single-KV-group target. -Wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels through 4 -panels with real GQA (h_q = G·h_kv, G > 1 on the decode side), writing -per-panel ``op_log_summary`` into ``sweep.json``. Independent from the -existing ``milestone-gqa-llama70b`` validation-scale bench (which stays -on the legacy baseline kernels). +Drives the LLaMA-3.1-70B single-KV-head-group prefill panel through +``_gqa_attention_prefill_long`` (C=8 snake Ring KV + intra-CUBE PE-SP, +all 64 ranks), writing ``op_log_summary`` into ``sweep.json``. +Independent from the existing ``milestone-gqa-llama70b`` validation-scale +bench (which stays on the legacy baseline kernels). + +Decode panels live in ``milestone_gqa_decode_long_ctx_4cases.py`` (the +4-cases long-context decode comparative study). Restrictions: - - Decode side capped at C ≤ 4 (1D chain reduce; 2D mesh wiring on - the decode panels deferred to the 4-cases decode comparative study) - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP headline deferred per ADR-0060 §B-item-1) + - T_q = S_kv = 1024 (scratch-limited; 32K headline awaits Q-axis + kernel tiling) - No figure renderers (defer to a separate cycle) Panels: single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV + - intra-CUBE PE-SP (all 64 ranks), - T_q=S_kv=1K, d_head=128 — the - LLaMA-3.1-70B single-KV-group target - (scratch-limited; 32K headline awaits - Q-axis kernel tiling) - -Decode panels live in ``milestone_gqa_decode_4cases.py`` (the 4-cases -comparative study). Legacy C=1 / C=4 panels were dropped — pytest -regression already covers those configurations. + intra-CUBE PE-SP, T_q=S_kv=1K, + d_head=128 — the LLaMA-3.1-70B + single-KV-group prefill target. Gated by ``GQA_HEADLINE_RUN=1``. """ @@ -33,14 +30,10 @@ import json import os from pathlib import Path -from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel from kernbench.benches.registry import bench from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import ( - configure_sfr_intercube_multisip, - configure_sfr_intercube_ring, -) +from kernbench.ccl.sfr_config import configure_sfr_intercube_ring from kernbench.policy.placement.dp import DPPolicy _OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline" @@ -51,10 +44,7 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep.json" _DTYPE = "f16" _D_HEAD = 64 _T_Q_PREFILL = 4 -_T_Q_DECODE = 1 _S_KV_PREFILL = 16 -_H_Q_DECODE = 8 # real GQA: G = H_Q_DECODE / H_KV_DECODE = 8 -_H_KV_DECODE = 1 _PANELS = ( "single_kv_group_prefill_gqa_c8_p8", @@ -133,43 +123,14 @@ def _run_prefill_panel( ) -def _run_decode_panel( - ctx, *, panel: str, C: int, P: int, S_kv: int, - sub_w: int = 0, - T_q: int = _T_Q_DECODE, - d_head: int = _D_HEAD, - h_q: int = _H_Q_DECODE, - h_kv: int = _H_KV_DECODE, -) -> None: - configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) - dp_full = DPPolicy(cube="replicate", pe="replicate", - num_cubes=C, num_pes=P) - dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="row_wise", num_cubes=C, num_pes=P) - q = ctx.zeros((T_q, h_q * d_head), - dtype=_DTYPE, dp=dp_full, name=f"{panel}_q") - k = ctx.zeros((S_kv, h_kv * d_head), - dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k") - v = ctx.zeros((S_kv, h_kv * d_head), - dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v") - o = ctx.empty((T_q, h_q * d_head), - dtype=_DTYPE, dp=dp_full, name=f"{panel}_o") - ctx.launch( - panel, gqa_attention_decode_long_kernel, - q, k, v, o, - T_q, S_kv, h_q, h_kv, d_head, C, P, sub_w, - _auto_dim_remap=False, - ) - - def _make_bench_fn(panel: str): kind, params = _PANEL_DISPATCH[panel] + assert kind == "prefill", ( + f"milestone-gqa-headline only registers prefill panels; got {kind!r}" + ) def _bench_fn(ctx): - if kind == "prefill": - _run_prefill_panel(ctx, panel=panel, **params) - else: - _run_decode_panel(ctx, panel=panel, **params) + _run_prefill_panel(ctx, panel=panel, **params) return _bench_fn @@ -234,10 +195,10 @@ def _run_panel(panel: str, topology: str) -> dict: @bench( name="milestone-gqa-headline", - description="Headline GQA milestone — real GQA h_q=8/h_kv=1 + 2-level SP (decode) + Ring KV (prefill).", + description="Headline GQA prefill milestone — LLaMA-3.1-70B single-KV-group target (C=8, P=8).", ) def run(torch) -> None: - """Drive 4 headline panels through the new GQA kernels; write sweep.json. + """Drive the headline prefill panel; write sweep.json. Gated by GQA_HEADLINE_RUN=1. """ @@ -254,10 +215,7 @@ def run(torch) -> None: "panels": list(_PANELS), "config": { "T_q_prefill": _T_Q_PREFILL, - "T_q_decode": _T_Q_DECODE, "S_kv_prefill": _S_KV_PREFILL, - "h_q_decode": _H_Q_DECODE, - "h_kv_decode": _H_KV_DECODE, "d_head": _D_HEAD, }, "rows": rows, diff --git a/tests/attention/test_gqa_decode.py b/tests/attention/test_gqa_decode.py deleted file mode 100644 index 5f8895f..0000000 --- a/tests/attention/test_gqa_decode.py +++ /dev/null @@ -1,143 +0,0 @@ -"""Phase 1 spec test for P1a GQA decode kernel (real GQA via M-fold). - -P1a is the first phase of the DDD-0060 plan, split out of the original -P1 (the composite-hybrid swap is P1b, deferred until the tl.composite -output-handle question is decided). P1a is the *correctness* unlock: -the kernel processes ONE KV head at a time and folds the G query heads -into the matmul M (row) dimension so a single Q·Kᵀ serves all G heads -sharing one K (ADR-0060 §5.2). This lifts the baseline's -``h_q == h_kv == 1`` cap pinned at -``tests/attention/test_milestone_gqa_llama70b.py:137-142``. - -P1a stays inside the existing ``tl`` API: the two attention GEMMs use the -blocking ``tl.dot`` so the chain ``Q·Kᵀ → softmax → P·V → store`` fits -without any composite-output chaining. The composite swap (P1b) will -revisit this once the API for feeding a composite's output into a -downstream MATH op is settled. - -Phase 1 (this commit): tests only — production code lands in Phase 2. -Tests fail at import in Phase 1 with ModuleNotFoundError; Phase 2 makes -them pass. -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 (Phase 2 deliverable) -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" - -# Decode shapes — P1a is one-shot, no tiling, single rank. P3 will tile. -T_Q = 1 -D_HEAD = 64 -S_KV = 16 -DTYPE = "f16" - - -def _engine_factory(t, d): - return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) - - -def _run_decode_p1(*, h_q: int, h_kv: int): - """One-shot GQA decode on a single PE in a single CUBE (P=1, no SP). - - Tensor layout (natural K — same as the P2a SP tests; the kernel - reshapes K to (d_head, S_local) via byte-conserving load): - Q : (T_q, h_q · d_head) — natural Q layout; kernel reshapes to - (G·T_q, d_head) — byte-conserving and math-correct because - T_q=1 collapses axis ordering. - K : (S_kv, h_kv · d_head) — natural K layout; kernel loads as - (d_head, S_local) — reshape-as-transpose caveat (ADR-0060 §3 - / §B item 2), correct for zero inputs used here. - V : (S_kv, h_kv · d_head) — natural V layout. - O : (T_q, h_q · d_head) — same shape as Q. - """ - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - dp = DPPolicy(cube="replicate", pe="replicate", - num_cubes=1, num_pes=1) - q = ctx.zeros((T_Q, h_q * D_HEAD), - dtype=DTYPE, dp=dp, name=f"q_h{h_q}_kv{h_kv}") - k = ctx.zeros((S_KV, h_kv * D_HEAD), - dtype=DTYPE, dp=dp, name=f"k_h{h_q}_kv{h_kv}") - v = ctx.zeros((S_KV, h_kv * D_HEAD), - dtype=DTYPE, dp=dp, name=f"v_h{h_q}_kv{h_kv}") - o = ctx.empty((T_Q, h_q * D_HEAD), - dtype=DTYPE, dp=dp, name=f"o_h{h_q}_kv{h_kv}") - ctx.launch( - f"gqa_decode_p1_h{h_q}_kv{h_kv}", - gqa_attention_decode_long_kernel, - q, k, v, o, - T_Q, S_KV, h_q, h_kv, D_HEAD, - 1, 1, # C=1, P=1 (no SP, degenerate) - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, - bench_fn=_bench_fn, - device=resolve_device(None), - engine_factory=_engine_factory, - ) - - -def _dma_read_count(op_log) -> int: - return sum(1 for r in op_log if r.op_name == "dma_read") - - -# ── Headline unlock: real GQA (h_q = G·h_kv) runs in data mode ────────── - - -def test_real_gqa_h_q_eight_h_kv_one_completes_in_data_mode(): - """ADR-0060 §A.1 headline unlock — the baseline raises - ``ValueError: Shape mismatch …`` in MemoryStore at h_q=8, h_kv=1 - because ``_view(K, (h_q·d, S_kv))`` only byte-conserves when h_q==h_kv. - M-fold processes one KV head at a time using only byte-conserving - reshapes, so this completes. - """ - result = _run_decode_p1(h_q=8, h_kv=1) - assert result.completion.ok, ( - f"real GQA (h_q=8, h_kv=1) decode failed: {result.completion}" - ) - - -# ── M-fold property: K/V loads do not scale with G ───────────────────── - - -def test_kv_dma_read_count_independent_of_g(): - """ADR-0060 TL;DR / §5.2: M-fold loads K and V once per KV head and - folds the G query heads into the GEMM M dim. The dma_read_count must - therefore be identical between (G=1, h_kv=1) and (G=8, h_kv=1) — both - issue exactly 3 reads (Q + K + V). This pins the GQA-reuse property - that the rest of the plan (composite streaming in P4, etc.) builds on. - """ - g1 = _run_decode_p1(h_q=1, h_kv=1) - g8 = _run_decode_p1(h_q=8, h_kv=1) - n_g1 = _dma_read_count(g1.engine.op_log) - n_g8 = _dma_read_count(g8.engine.op_log) - assert n_g1 == 3, f"G=1 dma_read_count must be 3 (Q+K+V); got {n_g1}" - assert n_g8 == 3, f"G=8 dma_read_count must be 3 (Q+K+V); got {n_g8}" - assert n_g1 == n_g8, ( - f"K/V dma_read_count must be independent of G; " - f"got G=1 -> {n_g1}, G=8 -> {n_g8}" - ) - - -# ── Backward-compat: degenerate G=1 still works ──────────────────────── - - -def test_degenerate_g_equals_one_still_works(): - """G=1 (h_q == h_kv == 1) is the baseline-compatible config. M-fold - degenerates to (T_q, d) = (1, 64) — the same shape the baseline - already exercises — so this proves no regression on that path. - """ - result = _run_decode_p1(h_q=1, h_kv=1) - assert result.completion.ok, ( - f"degenerate G=1 decode failed: {result.completion}" - ) diff --git a/tests/attention/test_gqa_decode_long_2d_reduce.py b/tests/attention/test_gqa_decode_long_2d_reduce.py deleted file mode 100644 index 107aa4f..0000000 --- a/tests/attention/test_gqa_decode_long_2d_reduce.py +++ /dev/null @@ -1,191 +0,0 @@ -"""Tests for lrab-adapted center-root inter-CUBE reduce in decode_long. - -Verifies the ADR-0060 §4.2 prescribed Level-1 collective for the -single-KV-group LLaMA-3.1-70B target (C=8 over a 2×4 sub-mesh, root at -the geometric center CUBE). - -Activation contract: - - ``sub_w == 0`` (default; omitted from launch args) → existing 1D - inter-CUBE chain that converges at CUBE 0. Byte-for-byte unchanged. - - ``sub_w >= 2`` with ``sub_h = C // sub_w >= 2`` → lrab-adapted - center-root mesh reduce. Root cube is - ``(sub_h//2)*sub_w + (sub_w//2)``. - -Degenerate (``sub_h < 2`` or ``sub_w < 2``) combinations are rejected at -launch time per ADR-0060 §4.2 + CLAUDE.md "Simplicity First": those -configurations belong to the 1D-chain code path, not lrab. - -Phase 1 (this commit): tests only — production code lands in Phase 2. -T1, T2, T4 fail today (kernel signature does not accept ``sub_w``). -T3 passes today as the backward-compat anchor for the existing -1D-chain path. -""" -from __future__ import annotations - -import re -from pathlib import Path - -import pytest - -from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 -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" -_CUBE_RE = re.compile(r"\bcube(\d+)\b") - - -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 _dma_write_cubes(op_log) -> list[int]: - """Return the list of CUBE ids that emitted a ``dma_write``. - - Decode's final-store path uses one ``tl.store`` from the root rank. - Multiple HBM-channel-level write records may correspond to the same - logical store; this just reports the source CUBE for each. - """ - cubes: list[int] = [] - for r in op_log: - if r.op_name != "dma_write": - continue - m = _CUBE_RE.search(r.component_id) - if m is not None: - cubes.append(int(m.group(1))) - return cubes - - -def _run_decode_long( - *, C: int, P: int, S_kv: int, sub_w: int | None, -): - """Drive a decode_long launch. ``sub_w=None`` means omit the arg - entirely (exercises the kernel's current default behaviour).""" - 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) - dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="row_wise", num_cubes=C, num_pes=P) - T_q = 1 - h_q = 8 - h_kv = 1 - ctx.zeros((T_q, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"q_dec_2d_c{C}") - k = ctx.zeros((S_kv, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"k_dec_2d_c{C}") - v = ctx.zeros((S_kv, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"v_dec_2d_c{C}") - o = ctx.empty((T_q, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"o_dec_2d_c{C}") - q = ctx.zeros((T_q, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"q_dec_2d_c{C}_2") - launch_args = [q, k, v, o, T_q, S_kv, h_q, h_kv, D_HEAD, C, P] - if sub_w is not None: - launch_args.append(sub_w) - ctx.launch( - f"gqa_decode_long_2d_c{C}_sw{sub_w}", - gqa_attention_decode_long_kernel, - *launch_args, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), engine_factory=_engine_factory, - ) - - -# ── T1: C=8 lrab-adapted center-root completes ─────────────────────── - - -def test_decode_2d_sub_w4_sub_h2_completes(): - """ADR-0060 §4.2 prescribed center-root mesh reduce at the - milestone target: C=8 over a 2×4 sub-mesh (sub_w=4, sub_h=2), - P=8 PEs per CUBE. With zero inputs, output is trivially zero; - the test guards that the kernel reaches completion without - deadlock or scratch overflow. - """ - result = _run_decode_long(C=8, P=8, S_kv=2048, sub_w=4) - assert result.completion.ok, ( - f"decode at C=8, sub_w=4 (lrab-adapted) must complete; " - f"got {result.completion}" - ) - - -# ── T2: root lives at the geometric center CUBE (cube 6) ───────────── - - -def test_decode_2d_sub_w4_sub_h2_root_at_center_cube_6(): - """For sub_w=4, sub_h=2: root_col=2, root_row=1, root_cube=6. - - Only the root CUBE issues the final ``tl.store`` — every other - rank short-circuits. The dma_write records must come exclusively - from CUBE 6. - """ - result = _run_decode_long(C=8, P=8, S_kv=2048, sub_w=4) - assert result.completion.ok - cubes = _dma_write_cubes(result.engine.op_log) - assert cubes, "expected at least one dma_write for the final O store" - distinct = set(cubes) - assert distinct == {6}, ( - f"final dma_write must come exclusively from CUBE 6 (sub_w=4, " - f"sub_h=2 root); got cubes={sorted(distinct)}" - ) - - -# ── T3: backward-compat — sub_w omitted → existing 1D chain at C=4 ─── - - -def test_decode_2d_backward_compat_sub_w0_default(): - """When ``sub_w`` is omitted from the launch, the kernel must - behave identically to today: 1D inter-CUBE chain along W, - converging at CUBE 0. This serves as a regression anchor — Phase 2 - must not change the existing C=4 panel behaviour. - - Passes today as well as after Phase 2. - """ - result = _run_decode_long(C=4, P=8, S_kv=2048, sub_w=None) - assert result.completion.ok, ( - f"decode at C=4 with default sub_w (1D chain) must complete; " - f"got {result.completion}" - ) - cubes = _dma_write_cubes(result.engine.op_log) - assert cubes, "expected at least one dma_write for the final O store" - distinct = set(cubes) - assert distinct == {0}, ( - f"1D-chain root must be CUBE 0; got cubes={sorted(distinct)}" - ) - - -# ── T4: degenerate sub_w configurations are rejected ───────────────── - - -def test_decode_2d_invalid_sub_w_rejected_when_sub_h_lt_2(): - """ADR-0060 §4.2 + CLAUDE.md "Simplicity First": only sub-meshes - with both sub_w >= 2 and sub_h >= 2 use lrab. C=4 with sub_w=4 - would give sub_h=1 (degenerate; Phase 2 of lrab becomes a no-op). - Callers must use the 1D-chain path (sub_w=0/omitted) for that case. - - The kernel must reject the degenerate combination with a clear - error rather than silently producing a 1D-chain result at the - wrong root location. - """ - with pytest.raises((ValueError, AssertionError), match=r"sub_[wh]"): - _run_decode_long(C=4, P=8, S_kv=2048, sub_w=4) diff --git a/tests/attention/test_gqa_decode_mc.py b/tests/attention/test_gqa_decode_mc.py deleted file mode 100644 index ce2f03d..0000000 --- a/tests/attention/test_gqa_decode_mc.py +++ /dev/null @@ -1,182 +0,0 @@ -"""Phase 1 spec test for P2b GQA decode multi-cube SP (both Level-2 + Level-1). - -P2b extends P2a to multiple CUBEs in one CUBE Group. The kernel uses the -canonical full SFR install (``configure_sfr_intercube_multisip``) which -provides disjoint direction namespaces: - - - ``intra_E / intra_W / intra_N / intra_S`` — PE↔PE within a CUBE - (logical 2×4 grid, no wrap) - - ``E / W / N / S`` — CUBE↔CUBE inter-CUBE - (mesh, no wrap) - -Reduce pattern (chain reduce-to-root, ADR-0060 §A.2 spirit, §4 chain -deviation noted): - - Level-2 (intra-CUBE, 2×4 grid): - row-then-col chain — each row reduces leftward along ``intra_W`` to - its col-0 PE, then PE 4 sends to PE 0 along ``intra_N``. 7 chain - steps per CUBE × 3 handles each = 21 ``ipcq_copy`` per CUBE. - - Level-1 (inter-CUBE): - only PE 0 of each CUBE participates. Chain leftward along ``W``. - (C-1) chain steps × 3 handles each. - - Final store at PE 0 of CUBE 0 only. - -Phase 1 (this commit): tests only — production code lands in Phase 2. -Phase 2 also updates ``test_gqa_decode.py`` (add C=1) and -``test_gqa_decode_sp.py`` (switch SFR + add C=1). -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 -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" - -T_Q = 1 -D_HEAD = 64 -DTYPE = "f16" - - -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_decode_mc(*, h_q: int, h_kv: int, C: int, P: int, S_kv: int): - """Multi-CUBE SP decode: C cubes × P PEs each share the work.""" - 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) - dp_kv = DPPolicy(cube="row_wise", pe="row_wise", - num_cubes=C, num_pes=P) - # Total KV split across C×P ranks; each rank sees S_kv/(C·P) rows. - q = ctx.zeros((T_Q, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, - name=f"q_h{h_q}_kv{h_kv}_c{C}_p{P}") - k = ctx.zeros((S_kv, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_kv, - name=f"k_h{h_q}_kv{h_kv}_c{C}_p{P}") - v = ctx.zeros((S_kv, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_kv, - name=f"v_h{h_q}_kv{h_kv}_c{C}_p{P}") - o = ctx.empty((T_Q, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, - name=f"o_h{h_q}_kv{h_kv}_c{C}_p{P}") - ctx.launch( - f"gqa_decode_mc_h{h_q}_kv{h_kv}_c{C}_p{P}", - gqa_attention_decode_long_kernel, - q, k, v, o, - T_Q, S_kv, h_q, h_kv, D_HEAD, - C, P, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, - bench_fn=_bench_fn, - device=resolve_device(None), - engine_factory=_engine_factory, - ) - - -def _count(op_log, name: str) -> int: - return sum(1 for r in op_log if r.op_name == name) - - -# ── Degenerate C=1 P=1 ──────────────────────────────────────────────── - - -def test_mc_c_one_p_one_degenerate(): - """C=1, P=1: single rank, no SP. No IPCQ traffic; one dma_write.""" - result = _run_decode_mc(h_q=1, h_kv=1, C=1, P=1, S_kv=16) - assert result.completion.ok, ( - f"C=1 P=1 degenerate failed: {result.completion}" - ) - assert _count(result.engine.op_log, "ipcq_copy") == 0 - assert _count(result.engine.op_log, "dma_write") == 1 - - -# ── Intra-CUBE only (single CUBE, P=8 on 2×4 grid) ──────────────────── - - -def test_mc_c_one_p_eight_intracube_grid(): - """C=1, P=8: intra-cube row-then-col chain on the 2×4 grid. - 7 chain steps × 3 handles (m, ℓ, O) = 21 ipcq_copy.""" - result = _run_decode_mc(h_q=1, h_kv=1, C=1, P=8, S_kv=64) - assert result.completion.ok, ( - f"C=1 P=8 intra-cube failed: {result.completion}" - ) - assert _count(result.engine.op_log, "dma_write") == 1 - n_copy = _count(result.engine.op_log, "ipcq_copy") - assert n_copy == 21, ( - f"C=1 P=8: expected 21 ipcq_copy (7 chain × 3 handles); got {n_copy}" - ) - - -# ── Multi-CUBE root-only write ──────────────────────────────────────── - - -def test_mc_c_two_p_eight_root_only_writes_o(): - """C=2, P=8: 16 ranks total. Only PE 0 of CUBE 0 writes O.""" - result = _run_decode_mc(h_q=1, h_kv=1, C=2, P=8, S_kv=128) - assert result.completion.ok, ( - f"C=2 P=8 multi-cube failed: {result.completion}" - ) - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 1, ( - f"root-only write must hold for C=2 P=8 (16 ranks); got {n_writes}" - ) - - -# ── Multi-CUBE total IPCQ chain count ───────────────────────────────── - - -def test_mc_c_two_p_eight_total_ipcq_count(): - """C=2, P=8: 21 intra-cube ipcq_copy per CUBE × 2 CUBEs + 3 inter-cube - chain ipcq_copy (C-1=1 step × 3 handles) = 45 total.""" - result = _run_decode_mc(h_q=1, h_kv=1, C=2, P=8, S_kv=128) - assert result.completion.ok, ( - f"C=2 P=8 multi-cube failed: {result.completion}" - ) - n_copy = _count(result.engine.op_log, "ipcq_copy") - expected = 21 * 2 + (2 - 1) * 3 - assert n_copy == expected, ( - f"C=2 P=8: expected {expected} ipcq_copy (21 intra × 2 CUBEs + 3 " - f"inter); got {n_copy}" - ) - - -# ── Real GQA × multi-CUBE SP combined ───────────────────────────────── - - -def test_mc_real_gqa_c_two_p_eight(): - """Headline: real GQA (h_q = G·h_kv with G=8) AND multi-CUBE SP - (C=2, P=8) together — the case the original baseline can express - neither part of.""" - result = _run_decode_mc(h_q=8, h_kv=1, C=2, P=8, S_kv=128) - assert result.completion.ok, ( - f"real GQA + multi-CUBE SP combined failed: {result.completion}" - ) - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 1, ( - f"root-only write must hold under M-fold + multi-CUBE; " - f"got {n_writes}" - ) diff --git a/tests/attention/test_gqa_decode_opt2.py b/tests/attention/test_gqa_decode_opt2.py index b824e08..ecfbc2d 100644 --- a/tests/attention/test_gqa_decode_opt2.py +++ b/tests/attention/test_gqa_decode_opt2.py @@ -181,55 +181,3 @@ def test_opt2_bench_completes_oplog_mode(): assert result.completion.ok, f"opt2 decode failed: {result.completion}" -# ── opt2 ↔ opt3 numeric parity in data mode (ADR-0065 N4) ───────────── - - -def _run_decode_data(kernel, name, q_d, k_d, v_d, *, S_kv=16): - """Run a decode kernel (C=P=1) in data mode with seeded Q/K/V; return the - HBM output array.""" - from pathlib import Path - - import numpy as np - 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.data_executor import DataExecutor - from kernbench.sim_engine.engine import GraphEngine - from kernbench.topology.builder import resolve_topology - - topo = resolve_topology(str(Path(__file__).resolve().parents[2] / "topology.yaml")) - cap: dict = {} - - def bench_fn(ctx): - dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1) - q = ctx.zeros((1, 8 * D), dtype="f16", dp=dp, name=name + "q") - k = ctx.zeros((S_kv, D), dtype="f16", dp=dp, name=name + "k") - v = ctx.zeros((S_kv, D), dtype="f16", dp=dp, name=name + "v") - o = ctx.empty((1, 8 * D), dtype="f16", dp=dp, name=name + "o") - q.copy_(ctx.from_numpy(q_d)); k.copy_(ctx.from_numpy(k_d)); v.copy_(ctx.from_numpy(v_d)) - cap["o"] = o - ctx.launch(name, kernel, q, k, v, o, 1, S_kv, 8, 1, D, 1, 1, _auto_dim_remap=False) - - r = run_bench(topology=topo, bench_fn=bench_fn, device=resolve_device(None), - engine_factory=lambda t, d: GraphEngine(getattr(t, "topology_obj", t), - enable_data=True)) - assert r.completion.ok - DataExecutor(r.engine.op_log, r.engine.memory_store).run() - o = cap["o"] - return r.engine.memory_store.read("hbm", o._handle.va_base, shape=(1, 8 * D), dtype="f16") - - -def test_opt2_matches_opt3_data_mode(): - import numpy as np - from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel - from kernbench.benches._gqa_attention_decode_opt2 import gqa_attention_decode_opt2_kernel - - rng = np.random.default_rng(0) - q_d = (0.1 * rng.standard_normal((1, 8 * D))).astype(np.float16) - k_d = (0.1 * rng.standard_normal((16, D))).astype(np.float16) - v_d = (0.1 * rng.standard_normal((16, D))).astype(np.float16) - - o3 = _run_decode_data(gqa_attention_decode_long_kernel, "p3", q_d, k_d, v_d) - o2 = _run_decode_data(gqa_attention_decode_opt2_kernel, "p2", q_d, k_d, v_d) - assert np.allclose(np.asarray(o2, np.float32), np.asarray(o3, np.float32), - rtol=5e-2, atol=5e-2), f"opt2 {np.asarray(o2).ravel()[:4]} vs opt3 {np.asarray(o3).ravel()[:4]}" diff --git a/tests/attention/test_gqa_decode_sp.py b/tests/attention/test_gqa_decode_sp.py deleted file mode 100644 index 5ecbafd..0000000 --- a/tests/attention/test_gqa_decode_sp.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Phase 1 spec test for P2a GQA decode SP (chain reduce-to-root, Level-2 only). - -P2a is the first half of DDD-0060 P2: the kernel becomes multi-PE within -one CUBE and reduces to root (PE 0) using a chain over the 1D intra-cube -ring (W direction). This **replaces the baseline's bidirectional O(N) -fan-out** where every rank ends with O — ADR-0060 §A.2's headline. - -Deviation from DDD-0060 §7 P2 gate: the gate text asks for -``⌈log₂ P⌉`` reduce rounds. The intra-cube SFR install -(``configure_sfr_intracube_pe_ring``) wires only a 1D E/W ring, so a -true tree would require either multi-hop forwarding or a new SFR install -(future ADR). P2a uses a **chain reduce-to-root**: ``P-1`` rounds along -W. The architectural property the ADR cares about -(root-only output vs every-rank-has-O) is preserved; the logarithmic -collective is deferred. - -P2b (deferred) covers Level-1 inter-CUBE center-mesh reduce (C>1). - -Phase 1 (this commit): tests only — production code lands in Phase 2. -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 -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" - -T_Q = 1 -D_HEAD = 64 -DTYPE = "f16" - - -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_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int): - """Single-CUBE SP decode: P PEs share the work along the intra-cube ring.""" - 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=1, num_pes=P) - dp_kv = DPPolicy(cube="replicate", pe="row_wise", - num_cubes=1, num_pes=P) - q = ctx.zeros((T_Q, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"q_h{h_q}_kv{h_kv}_p{P}") - # KV: total S_kv split across P PEs along axis 0 (row_wise sharding). - # Each PE sees (S_kv/P, h_kv·D_HEAD). - k = ctx.zeros((S_kv, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"k_h{h_q}_kv{h_kv}_p{P}") - v = ctx.zeros((S_kv, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"v_h{h_q}_kv{h_kv}_p{P}") - o = ctx.empty((T_Q, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"o_h{h_q}_kv{h_kv}_p{P}") - ctx.launch( - f"gqa_decode_sp_h{h_q}_kv{h_kv}_p{P}", - gqa_attention_decode_long_kernel, - q, k, v, o, - T_Q, S_kv, h_q, h_kv, D_HEAD, - 1, P, # C=1, P=P (single-CUBE SP) - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, - bench_fn=_bench_fn, - device=resolve_device(None), - engine_factory=_engine_factory, - ) - - -def _count(op_log, name: str) -> int: - return sum(1 for r in op_log if r.op_name == name) - - -# ── Root-only write ──────────────────────────────────────────────────── - - -def test_sp_chain_reduce_root_only_writes_o(): - """ADR-0060 §A.2: only the root rank (PE 0) writes O. Baseline today - has every rank write the full final O (bidirectional fan-out).""" - result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64) - assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}" - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 1, ( - f"reduce-to-root must produce exactly 1 dma_write (PE 0); " - f"got {n_writes}" - ) - - -# ── Chain step count ─────────────────────────────────────────────────── - - -def test_sp_chain_reduce_p_minus_one_ipcq_pairs(): - """Chain reduce-to-root has P-1 send→recv pairs along the W chain; - each pair logs one ``ipcq_copy`` (inbound DMA, per - ``milestone_gqa_llama70b._summarize_op_log``). Each chain step ships - the triplet (m, ℓ, O) → 3 handles per step → 7 steps × 3 = 21.""" - result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64) - assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}" - n_copy = _count(result.engine.op_log, "ipcq_copy") - expected = (8 - 1) * 3 - assert n_copy == expected, ( - f"chain reduce: expected {expected} ipcq_copy (P-1=7 steps × " - f"3 handles m/ℓ/O); got {n_copy}" - ) - - -# ── Real GQA × SP combined ───────────────────────────────────────────── - - -def test_sp_real_gqa_h_q_eight_h_kv_one_p_eight(): - """The combined unlock: real GQA (h_q=G·h_kv with G=8) AND SP - (P=8) together — neither expressible by the baseline.""" - result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64) - assert result.completion.ok, ( - f"real GQA + SP combined run failed: {result.completion}" - ) - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 1, ( - f"root-only write must hold under M-fold too; got {n_writes}" - ) - - -# ── Degenerate P=1 ──────────────────────────────────────────────────── - - -def test_sp_p_one_degenerate_no_ipcq_traffic(): - """P=1: SP degenerates to a single rank. No IPCQ traffic; one dma_write.""" - result = _run_decode_sp(h_q=8, h_kv=1, P=1, S_kv=16) - assert result.completion.ok, f"P=1 degenerate failed: {result.completion}" - n_send = _count(result.engine.op_log, "ipcq_send") - n_recv = _count(result.engine.op_log, "ipcq_recv") - assert n_send == 0, f"P=1 must have no ipcq_send; got {n_send}" - assert n_recv == 0, f"P=1 must have no ipcq_recv; got {n_recv}" - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 1, f"P=1: one dma_write; got {n_writes}" diff --git a/tests/attention/test_gqa_long_context.py b/tests/attention/test_gqa_long_context.py deleted file mode 100644 index d5354c7..0000000 --- a/tests/attention/test_gqa_long_context.py +++ /dev/null @@ -1,157 +0,0 @@ -"""Phase 1 spec test for Phase C: long-context regression for the GQA -kernels (ADR-0060 §B "long/short context split" + ADR-0063 §A.2). - -The headline panel today caps prefill at S_kv=16 / decode at S_kv≤128 — -NOT because the algorithm fails, but because the bump allocator never -recycles. ADR-0063 §A.2 (test req 3) requires a sweep at an S that -would overflow 1 MiB without recycling to complete after scope -discipline lands. - -Scratch-budget estimate at C=4, P=8 (current 1 MiB pool): - decode: per-rank S_local = S_kv / 32; intermediates ≲ 500 KB at - S_kv=32K (fits today — used as regression guard). - prefill: per-rank S_local = S_kv / 4; ~16·S_local bytes per ring - step × 4 steps ≈ 64·S_local bytes. Overflows at - ~64 K tokens (64 × 16K > 1 MiB). - -Phase 1 (this commit): tests only — production code lands in Phase 2. -The prefill test fails today with a RuntimeError("TLContext scratch -overflow"). After Phase 2 (scratch_scope + copy_to discipline) it -completes. -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 -from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 -from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import ( - configure_sfr_intercube_multisip, - configure_sfr_intercube_ring, -) -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" - - -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) - - -# ── Decode at moderate-long context (regression guard) ─────────────── - - -def test_decode_long_context_32k_completes(): - """Decode at S_kv=32K (C=1, P=8) — per-rank S_local=4K. Should - complete with current scratch usage (~few KB intermediates) and - must continue to complete after the rewrite. - - This is a regression guard: scratch discipline shouldn't break - decode's existing long-context capability. - """ - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) - P = 8 - S_kv = 32_768 - dp_full = DPPolicy(cube="replicate", pe="replicate", - num_cubes=1, num_pes=P) - dp_kv = DPPolicy(cube="replicate", pe="row_wise", - num_cubes=1, num_pes=P) - ctx.zeros((1, 8 * D_HEAD), dtype=DTYPE, dp=dp_full, name="q_long_dec") - k = ctx.zeros((S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name="k_long_dec") - v = ctx.zeros((S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name="v_long_dec") - o = ctx.empty((1, 8 * D_HEAD), - dtype=DTYPE, dp=dp_full, name="o_long_dec") - q = ctx.zeros((1, 8 * D_HEAD), - dtype=DTYPE, dp=dp_full, name="q_long_dec_2") - ctx.launch( - "gqa_decode_long_32k", - gqa_attention_decode_long_kernel, - q, k, v, o, - 1, S_kv, 8, 1, D_HEAD, - 1, P, - _auto_dim_remap=False, - ) - - result = run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), engine_factory=_engine_factory, - ) - assert result.completion.ok, ( - f"decode at S_kv=32K must complete; got {result.completion}" - ) - - -# ── Prefill at long context overflows without scope discipline ─────── - - -def test_prefill_long_context_completes_after_scope_discipline(): - """ADR-0063 §A.2 Test Req 3: a sweep at an ``S`` that would overflow - 1 MiB without recycling must complete after scope discipline lands. - - With C=4 and S_kv chosen so per-rank S_local·8·4 > 1 MiB - (~16K tokens per rank ⇒ S_kv ≥ 64K), the current kernel overflows - the 1 MiB scratch pool. After Phase 2 (scratch_scope wraps each - ring step's intermediates; copy_to persists running state), it - completes. - - This is the headline test that proves the S-ceiling is gone for - prefill. - """ - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - configure_sfr_intercube_ring( - ctx.engine, ctx.spec, _ccl_cfg(), ring_size=4, - ) - T_q = 4 - S_kv = 65_536 # 64 K — per-rank 16 K, ~2 MB scratch w/o scope - C = 4 - dp_q = DPPolicy(cube="replicate", pe="replicate", - num_cubes=C, num_pes=1) - dp_kv = DPPolicy(cube="row_wise", pe="replicate", - num_cubes=C, num_pes=1) - dp_o = DPPolicy(cube="row_wise", pe="replicate", - num_cubes=C, num_pes=1) - q = ctx.zeros((T_q, D_HEAD), - dtype=DTYPE, dp=dp_q, name="q_long_pre") - k = ctx.zeros((S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name="k_long_pre") - v = ctx.zeros((S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name="v_long_pre") - o = ctx.empty((T_q * C, D_HEAD), - dtype=DTYPE, dp=dp_o, name="o_long_pre") - ctx.launch( - "gqa_prefill_long_64k", - gqa_attention_prefill_long_kernel, - q, k, v, o, - T_q, S_kv, D_HEAD, C, - _auto_dim_remap=False, - ) - - result = run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), engine_factory=_engine_factory, - ) - assert result.completion.ok, ( - f"prefill at S_kv=64K must complete after scope discipline lands; " - f"got {result.completion}" - ) diff --git a/tests/attention/test_gqa_scoped_writeback.py b/tests/attention/test_gqa_scoped_writeback.py deleted file mode 100644 index d6541d0..0000000 --- a/tests/attention/test_gqa_scoped_writeback.py +++ /dev/null @@ -1,198 +0,0 @@ -"""Phase 1 spec test for Phase C: scratch_scope + tl.copy_to discipline -in the GQA kernels (ADR-0060 §5.2 / §5.5 + ADR-0063 §D3 / §D3.1). - -ADR-0060 §5.2 (decode pseudocode line 75 / §5.5 (prefill pseudocode line -96) both wrap per-tile / per-ring-step intermediates in -``with tl.scratch_scope():`` and persist the merged running ``(m, ℓ, O)`` -to a persistent arena allocated outside the scope. The original ADR-0063 -§D3 specifies the two-arena pattern; §D3.1 specifies the -``tl.copy_to(dst, src)`` writeback primitive used to persist scoped -results. - -Currently neither kernel uses ``scratch_scope`` or ``copy_to``; their -chain-merge / ring-merge bodies allocate every intermediate from the -bump cursor and never recycle. Result: op_log has 0 ``copy`` entries. - -After Phase 2: each per-tile / per-step merge writes the new running -``(m, ℓ, O)`` via ``copy_to`` to the persistent arena. Per merge step -→ 3 copy ops (m, ℓ, O). - -Phase 1 (this commit): tests only — production code lands in Phase 2. -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 -from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 -from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import ( - configure_sfr_intercube_multisip, - configure_sfr_intercube_ring, -) -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" - - -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 _count(op_log, name: str) -> int: - return sum(1 for r in op_log if r.op_name == name) - - -# ── Decode chain merges must use scratch_scope + copy_to ───────────── - - -def _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int): - """Single-CUBE SP decode (C=1, P PEs along intra-cube chain).""" - 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=1, num_pes=P) - dp_kv = DPPolicy(cube="replicate", pe="row_wise", - num_cubes=1, num_pes=P) - q = ctx.zeros((1, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"q_sc_{P}") - k = ctx.zeros((S_kv, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"k_sc_{P}") - v = ctx.zeros((S_kv, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"v_sc_{P}") - o = ctx.empty((1, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"o_sc_{P}") - ctx.launch( - f"gqa_decode_scoped_{P}", - gqa_attention_decode_long_kernel, - q, k, v, o, - 1, S_kv, h_q, h_kv, D_HEAD, - 1, P, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), engine_factory=_engine_factory, - ) - - -def test_decode_chain_merges_emit_copy_to_writeback(): - """ADR-0060 §5.2 + ADR-0063 §D3.1: each chain-merge step must wrap - its intermediates in ``scratch_scope`` and persist the new running - ``(m, ℓ, O)`` via ``tl.copy_to``. - - For (C=1, P=8): 7 intra-cube chain merges × 3 handles (m, ℓ, O) per - merge ⇒ 21 ``copy`` entries. - - Currently 0 because the kernel never calls ``tl.copy_to``. - """ - result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64) - assert result.completion.ok, f"decode SP failed: {result.completion}" - n_copy = _count(result.engine.op_log, "copy") - assert n_copy > 0, ( - f"decode kernel must emit copy_to writeback per merge step " - f"(ADR-0060 §5.2 + ADR-0063 §D3.1); got 0 ``copy`` entries" - ) - - -# ── Prefill Ring KV merges must use scratch_scope + copy_to ────────── - - -def _run_prefill_ring(*, T_q: int, S_kv: int, C: int): - """Head-parallel prefill with Ring KV across C CUBEs.""" - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - configure_sfr_intercube_ring( - ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C, - ) - dp_q = DPPolicy(cube="replicate", pe="replicate", - num_cubes=C, num_pes=1) - dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="replicate", num_cubes=C, num_pes=1) - dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="replicate", num_cubes=C, num_pes=1) - q = ctx.zeros((T_q, D_HEAD), - dtype=DTYPE, dp=dp_q, name=f"q_ring_{C}") - k = ctx.zeros((S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"k_ring_{C}") - v = ctx.zeros((S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"v_ring_{C}") - o = ctx.empty((T_q * C, D_HEAD), - dtype=DTYPE, dp=dp_o, name=f"o_ring_{C}") - ctx.launch( - f"gqa_prefill_scoped_{C}", - gqa_attention_prefill_long_kernel, - q, k, v, o, - T_q, S_kv, D_HEAD, C, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), engine_factory=_engine_factory, - ) - - -def test_prefill_ring_step_merges_emit_copy_to_writeback(): - """ADR-0060 §5.5 + ADR-0063 §D3.1: each Ring KV step's online-softmax - merge must wrap its intermediates in ``scratch_scope`` and persist - the new running ``(m, ℓ, O)`` via ``tl.copy_to``. - - For C=4: 3 ring-step merges (steps 1..C-1) × 3 handles (m, ℓ, O) per - merge ⇒ 9 ``copy`` entries per participating CUBE. Aggregated across - C CUBEs: ⇒ 36 ``copy`` entries. - - Currently 0 because the kernel never calls ``tl.copy_to``. - """ - result = _run_prefill_ring(T_q=4, S_kv=16, C=4) - assert result.completion.ok, f"prefill ring failed: {result.completion}" - n_copy = _count(result.engine.op_log, "copy") - assert n_copy > 0, ( - f"prefill ring kernel must emit copy_to writeback per merge step " - f"(ADR-0060 §5.5 + ADR-0063 §D3.1); got 0 ``copy`` entries" - ) - - -# ── Scoped kernels still produce the same op_log shape (regression) ── - - -def test_decode_scoped_still_has_root_only_write(): - """ADR-0060 §A.2 root-only output must hold under the rewrite: - adding scratch_scope + copy_to should not change the reduce - topology; only the per-PE scratch usage. Single PE 0 writes O.""" - result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64) - assert result.completion.ok - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 1, ( - f"scoped decode must still produce 1 dma_write (root-only); " - f"got {n_writes}" - ) - - -def test_prefill_scoped_still_has_per_cube_distributed_output(): - """ADR-0060 §5.5 per-CUBE distributed output must hold under the - rewrite: scoped prefill still writes one O slice per CUBE.""" - result = _run_prefill_ring(T_q=4, S_kv=16, C=4) - assert result.completion.ok - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 4, ( - f"scoped prefill must write one O per CUBE (C=4 → 4 dma_write); " - f"got {n_writes}" - ) diff --git a/tests/attention/test_gqa_tile_sweep.py b/tests/attention/test_gqa_tile_sweep.py deleted file mode 100644 index 322cd49..0000000 --- a/tests/attention/test_gqa_tile_sweep.py +++ /dev/null @@ -1,290 +0,0 @@ -"""Phase 1 spec test for P3b: S_kv tile sweep in the GQA kernels -(ADR-0063 §A.2 + ADR-0060 §B "long/short context split"). - -The local one-shot partial in three kernels — ``_gqa_attention_decode_long.py``, -``_gqa_attention_decode_short.py``, ``_gqa_attention_prefill_short.py`` — loads its rank's -entire ``(d_head, S_local)`` / ``(S_local, d_head)`` KV slice in one -shot, so per-rank scratch grows linearly with ``S_local``. The -``_attention_*`` baselines hit a TCM ceiling once ``S_local`` exceeds -the 1 MiB pool divided by the per-rank intermediate footprint. - -P3b replaces the one-shot partial with a tile sweep: - - Module-level ``TILE_S_KV = 1024`` per kernel file. - - Tile 0 establishes the persistent ``(m_local, l_local, O_local)``. - - Tiles 1..n_tiles-1 wrap their intermediates (K_T tile, V tile, - scores, centered, exp_scores, partials) in ``tl.scratch_scope()`` - and persist the merged ``(m, ℓ, O)`` to the persistent handles via - ``tl.copy_to`` (ADR-0063 §D3 / §D3.1). - -When ``S_local ≤ TILE_S_KV`` the loop body never runs and the op_log -is structurally identical to today's one-shot path — every existing -validation-scale test continues to pass. - -``_gqa_attention_prefill_long.py`` is **out of scope** for P3b. Its inner step -is IPCQ partial-recv (Ring KV rotation), not an HBM load; tiling it -intersects the ring-step structure and is a separate phase. - -Phase 1 (this commit): tests only — production code lands in Phase 2. -The four multi-tile tests fail today because the kernels never call -``copy_to`` outside the chain-reduce merges (so isolating to a -chain-free config gives ``copy_to == 0`` today). The 128K test fails -today with a ``TLContext`` scratch overflow. -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 -from kernbench.benches._gqa_attention_decode_short import gqa_attention_decode_short_kernel # noqa: F401 -from kernbench.benches._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel # noqa: F401 -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" - - -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 _count(op_log, name: str) -> int: - return sum(1 for r in op_log if r.op_name == name) - - -# ── decode_long runner (chain-free configs isolate tile-sweep behavior) ── - - -def _run_decode_long(*, C: int, P: int, S_kv: int, h_q: int = 1, h_kv: int = 1): - """Run the long-context decode kernel. - - For tile-sweep isolation, callers pass C=1, P=1 — this leaves both - chain-reduce levels inactive so any ``copy_to`` in op_log comes - solely from the tile-merge body. - """ - 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) - dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="row_wise" if P > 1 else "replicate", - num_cubes=C, num_pes=P) - q = ctx.zeros((1, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"q_tl_c{C}_p{P}_s{S_kv}") - k = ctx.zeros((S_kv, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"k_tl_c{C}_p{P}_s{S_kv}") - v = ctx.zeros((S_kv, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"v_tl_c{C}_p{P}_s{S_kv}") - o = ctx.empty((1, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"o_tl_c{C}_p{P}_s{S_kv}") - ctx.launch( - f"gqa_decode_long_tile_c{C}_p{P}_s{S_kv}", - gqa_attention_decode_long_kernel, - q, k, v, o, - 1, S_kv, h_q, h_kv, D_HEAD, C, P, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), engine_factory=_engine_factory, - ) - - -# ── decode_short / prefill_short runners ───────────────────────────── - - -def _run_decode_short(*, kv_per_cube: int, C: int, P: int, S_kv: int, - h_q: int = 8, h_kv: int = 8): - """For tile-sweep isolation: kv_per_cube=P → group_size=1 (no chain).""" - 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) - dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", - 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_dsh_s{S_kv}") - k = ctx.zeros((h_kv * S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"k_dsh_s{S_kv}") - v = ctx.zeros((h_kv * S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"v_dsh_s{S_kv}") - o = ctx.empty((1, h_q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"o_dsh_s{S_kv}") - ctx.launch( - f"gqa_decode_short_tile_s{S_kv}", - 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 _run_prefill_short(*, kv_per_cube: int, C: int, P: int, - T_q: int, S_kv: int, h_kv: int = 8): - 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) - dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="row_wise", num_cubes=C, num_pes=P) - q = ctx.zeros((T_q, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"q_psh_s{S_kv}") - k = ctx.zeros((h_kv * S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"k_psh_s{S_kv}") - v = ctx.zeros((h_kv * S_kv, D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"v_psh_s{S_kv}") - o = ctx.empty((T_q, h_kv * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"o_psh_s{S_kv}") - ctx.launch( - f"gqa_prefill_short_tile_s{S_kv}", - gqa_attention_prefill_short_kernel, - q, k, v, o, - T_q, S_kv, 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, - ) - - -# ── T1: single-tile path is op_log-stable (regression guard) ───────── - - -def test_decode_long_single_tile_path_unchanged(): - """ADR-0063 §A.2: when S_local <= TILE_S_KV the tile-sweep loop - body never runs and op_log is structurally identical to today's - one-shot path. - - Passes today (one-shot path) AND after Phase 2 (n_tiles=1 skips - the merge loop). The witness: with chain reduce inactive (C=1, - P=1), there must be zero ``copy_to`` entries — neither today nor - after Phase 2 should the kernel emit tile-merge writebacks here. - """ - result = _run_decode_long(C=1, P=1, S_kv=64) - assert result.completion.ok, ( - f"single-tile decode_long must complete; got {result.completion}" - ) - n_copy = _count(result.engine.op_log, "copy") - assert n_copy == 0, ( - f"S_local <= TILE_S_KV must not emit tile-merge copy_to; got {n_copy}" - ) - - -# ── T2: decode_long multi-tile emits tile-merge copy_to ────────────── - - -def test_decode_long_multi_tile_emits_copy_to_merges(): - """ADR-0063 §A.2 + §D3.1: when S_local > TILE_S_KV the kernel must - wrap each subsequent tile's intermediates in ``scratch_scope`` and - persist the merged ``(m, ℓ, O)`` via ``tl.copy_to``. - - Chain-free config (C=1, P=1) isolates the tile-merge body — any - ``copy_to`` in op_log comes from the tile sweep, not chain reduce. - - S_kv=2048, C=1, P=1 → S_local=2048 → 2 tiles → 1 merge × 3 handles - (m, ℓ, O) ⇒ 3 ``copy`` entries. - - Currently 0 because the kernel performs a one-shot partial. - """ - result = _run_decode_long(C=1, P=1, S_kv=2048) - assert result.completion.ok, ( - f"multi-tile decode_long must complete; got {result.completion}" - ) - n_copy = _count(result.engine.op_log, "copy") - assert n_copy >= 3, ( - f"decode_long multi-tile must emit >= 3 copy_to entries " - f"(1 merge × 3 handles); got {n_copy}" - ) - - -# ── T3: decode_short multi-tile emits tile-merge copy_to ───────────── - - -def test_decode_short_multi_tile_emits_copy_to_merges(): - """Same property as T2 for the short decode kernel. - - kv_per_cube=8, P=8, C=1 → group_size=1 (no chain reduce); S_kv=2048 - → S_local=2048 → 2 tiles per PE → 3 ``copy`` entries per PE × 8 PEs - = 24 total. - - Currently 0 because the kernel performs a one-shot partial. - """ - result = _run_decode_short(kv_per_cube=8, C=1, P=8, S_kv=2048) - assert result.completion.ok, ( - f"multi-tile decode_short must complete; got {result.completion}" - ) - n_copy = _count(result.engine.op_log, "copy") - assert n_copy >= 3 * 8, ( - f"decode_short multi-tile must emit >= 24 copy_to entries " - f"(1 merge × 3 handles × 8 PEs); got {n_copy}" - ) - - -# ── T4: prefill_short multi-tile emits tile-merge copy_to ──────────── - - -def test_prefill_short_multi_tile_emits_copy_to_merges(): - """Same property as T3 for the short prefill kernel. - - kv_per_cube=8, P=8, C=1, T_q=4, S_kv=2048 → group_size=1 (no chain), - S_local=2048 → 2 tiles per PE → 3 ``copy`` entries per PE × 8 PEs - = 24 total. - - Currently 0 because the kernel performs a one-shot partial. - """ - result = _run_prefill_short( - kv_per_cube=8, C=1, P=8, T_q=4, S_kv=2048, - ) - assert result.completion.ok, ( - f"multi-tile prefill_short must complete; got {result.completion}" - ) - n_copy = _count(result.engine.op_log, "copy") - assert n_copy >= 3 * 8, ( - f"prefill_short multi-tile must emit >= 24 copy_to entries " - f"(1 merge × 3 handles × 8 PEs); got {n_copy}" - ) - - -# ── T5: decode_long at S_kv=256K completes (TCM ceiling lifted) ────── - - -def test_decode_long_context_256k_completes(): - """ADR-0063 §A.2 (test req 3): headline ceiling-lift. C=1 P=8 - decode at S_kv=256K → S_local=32K. Today the per-rank score stack - (3 ×M·S_local·2 bytes, M=G·T_q=8) is ~1.5 MB → exceeds the 1 MiB - scratch pool → TLContext scratch overflow. After Phase 2 the tile - sweep bounds per-tile score stack at 3·M·TILE_S_KV·2 ≈ 48 KB and - the run completes regardless of S_kv. - """ - result = _run_decode_long(C=1, P=8, S_kv=262144, h_q=8, h_kv=1) - assert result.completion.ok, ( - f"decode_long at S_kv=256K must complete after tile sweep; " - f"got {result.completion}" - )