diff --git a/src/kernbench/benches/_gqa_attention_decode_long.py b/src/kernbench/benches/_gqa_attention_decode_long.py index 576a4d1..77bd132 100644 --- a/src/kernbench/benches/_gqa_attention_decode_long.py +++ b/src/kernbench/benches/_gqa_attention_decode_long.py @@ -4,14 +4,29 @@ 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 -to PE 0 of CUBE 0 via a 2-level chain (intra-CUBE row+col, then -inter-CUBE), and the root writes the final output. +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 are arranged as a 1D row (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``). @@ -47,10 +62,11 @@ def gqa_attention_decode_long_kernel( d_head: int, C: int, P: int, + sub_w: int = 0, *, tl, ) -> None: - """GQA decode with M-fold + S_kv tile sweep + 2-level chain reduce-to-root. + """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 @@ -59,8 +75,26 @@ def gqa_attention_decode_long_kernel( 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 CUBE 0 stores. + 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 @@ -161,8 +195,9 @@ def gqa_attention_decode_long_kernel( tl.send(dir="intra_N", src=l_local) tl.send(dir="intra_N", src=O_local) - # Level-1 inter-CUBE chain (along W, leftward; only PE 0 of each CUBE). - if pe_id == 0 and C > 1: + # 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") @@ -178,8 +213,141 @@ def gqa_attention_decode_long_kernel( 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 == 0: + 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_prefill_long.py b/src/kernbench/benches/_gqa_attention_prefill_long.py index 40a1ec9..77970b8 100644 --- a/src/kernbench/benches/_gqa_attention_prefill_long.py +++ b/src/kernbench/benches/_gqa_attention_prefill_long.py @@ -16,10 +16,17 @@ persistent scratch is ``(m, ℓ, O)`` only (~1 KB); per-tile in-scope scratch is bounded by ``TILE_S_KV`` regardless of ``S_local``. Topology / SFR: - - Requires ``configure_sfr_intercube_ring(ring_size=C)`` (1D ring of - C CUBEs with wrap at the CUBE level). - - Only PE 0 of each CUBE participates (head-parallel; intra-CUBE PE - parallelism is a separate phase). + - Requires ``configure_sfr_intercube_ring`` — either ``ring_size=C`` + (1D row, ``C ≤ mesh_w``) or ``submesh_shape=(rows, cols)`` (snake + Hamiltonian cycle through a rectangular sub-mesh, for ``C > mesh_w``). + - ``P == 1`` (default): only PE 0 of each CUBE participates; + intra-CUBE PE parallelism is disabled. + - ``P > 1``: all P PEs of each CUBE participate via query-axis + split (ADR-0060 §5.5 last bullet + §B-item-3) — each PE owns + ``T_q/P`` disjoint query rows. Output rows are disjoint per PE, + so no intra-CUBE reduce is needed. Each PE drives its own + same-lane ring via independent IPCQ channels (P parallel rings). + Requires ``T_q % P == 0``. Layout caveats: - GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``). @@ -55,36 +62,56 @@ def gqa_attention_prefill_long_kernel( S_kv: int, d_head: int, C: int, + P: int = 1, *, tl, ) -> None: """Head-parallel prefill attention with tile-granular Ring KV (ADR-0060 §5.5). - Tensor layout: - Q : (T_q, d_head) one head per CUBE; replicated. + Tensor layout (per-rank shapes): + Q : (T_q_local, d_head) one head per CUBE; replicated within + a CUBE for P=1, or sharded ``pe="row_wise"`` for P>1 so each + PE owns ``T_q_local = T_q // P`` query rows. K : (S_kv, d_head) sharded cube_row_wise → each CUBE owns (S_local, d_head). Tiles loaded as (d_head, tile_s) via byte-conserving reshape. V : (S_kv, d_head) sharded cube_row_wise → each CUBE owns (S_local, d_head). Tiles loaded as (tile_s, d_head). O : (T_q * C, d_head) sharded cube_row_wise → each CUBE - writes its own (T_q, d_head) slice. NO reduce. + writes its own ``T_q`` rows; for P>1 each PE writes a + disjoint ``(T_q_local, d_head)`` slice (no intra-CUBE + reduce — disjoint output rows). NO inter-CUBE reduce. Algorithm: nested loop over (ring_step k, tile_idx t). At k=0 each - CUBE loads its own block's tiles from HBM; at k > 0 it receives + rank loads its own block's tiles from HBM; at k > 0 it receives tiles from its E neighbour. Tiles are forwarded W to the next ring step. Each tile's partial is folded into the running (m, ℓ, O) via online-softmax merge. """ pe_id = tl.program_id(axis=0) - # Head-parallel: only PE 0 of each CUBE participates. - if pe_id != 0: - return + if P == 1: + # Head-parallel only: PE 0 of each CUBE participates. + if pe_id != 0: + return + T_q_local = T_q + else: + # Intra-CUBE PE-SP (ADR-0060 §5.5 last bullet + §B-item-3): + # query-axis split across P PEs of each CUBE. Output rows are + # disjoint per PE ⇒ no intra-CUBE reduce. Each PE drives its + # own same-lane ring via independent IPCQ channels. + if pe_id >= P: + return + if T_q % P != 0: + raise ValueError( + f"T_q={T_q} must be divisible by P={P} when P > 1; " + "use P=1 for the PE-0-only fallback path." + ) + T_q_local = T_q // P S_local = S_kv // C n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV KV_ROW_BYTES = d_head * 2 # f16 - Q = tl.load(q_ptr, shape=(T_q, d_head), dtype="f16") + Q = tl.load(q_ptr, shape=(T_q_local, d_head), dtype="f16") # ── Bootstrap: (t=0, k=0) — load my own tile 0, establish (m, ℓ, O) ── # diff --git a/src/kernbench/benches/milestone_gqa_headline.py b/src/kernbench/benches/milestone_gqa_headline.py index 94e55d1..d8023b3 100644 --- a/src/kernbench/benches/milestone_gqa_headline.py +++ b/src/kernbench/benches/milestone_gqa_headline.py @@ -6,19 +6,24 @@ 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). -Restrictions (P7 first cut): - - C ≤ 4 (single-row inter-CUBE ring SFR; multi-row deferred) +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) + headline deferred per ADR-0060 §B-item-1) - No figure renderers (defer to a separate cycle) Panels: - single_user_prefill_gqa : prefill C=1, T_q=4, S_kv=16 - multi_user_prefill_gqa : prefill C=4 Ring KV, T_q=4, S_kv=16 - single_user_decode_gqa : decode C=1, P=8, h_q=8, h_kv=1, S_kv=64 - (M-fold + intra-cube row-then-col chain) - multi_user_decode_gqa : decode C=4, P=8, h_q=8, h_kv=1, S_kv=128 - (M-fold + 2-level chain reduce-to-root) + single_user_prefill_gqa : prefill C=1, T_q=4, S_kv=16 + multi_user_prefill_gqa : prefill C=4 Ring KV, T_q=4, S_kv=16 + single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV + + intra-CUBE PE-SP (all 64 ranks), + T_q=S_kv=32K, d_head=128 — the + LLaMA-3.1-70B single-KV-group target + single_user_decode_gqa : decode C=1, P=8, h_q=8, h_kv=1, S_kv=64 + (M-fold + intra-cube row-then-col chain) + multi_user_decode_gqa : decode C=4, P=8, h_q=8, h_kv=1, S_kv=128 + (M-fold + 2-level chain reduce-to-root) Gated by ``GQA_HEADLINE_RUN=1``. """ @@ -54,6 +59,7 @@ _H_KV_DECODE = 1 _PANELS = ( "single_user_prefill_gqa", "multi_user_prefill_gqa", + "single_kv_group_prefill_gqa_c8_p8", "single_user_decode_gqa", "multi_user_decode_gqa", ) @@ -62,6 +68,11 @@ _PANELS = ( _PANEL_DISPATCH: dict[str, tuple[str, dict]] = { "single_user_prefill_gqa": ("prefill", {"C": 1, "S_kv": _S_KV_PREFILL}), "multi_user_prefill_gqa": ("prefill", {"C": 4, "S_kv": _S_KV_PREFILL}), + "single_kv_group_prefill_gqa_c8_p8": ("prefill", { + "C": 8, "P": 8, + "T_q": 32_768, "S_kv": 32_768, + "d_head": 128, + }), "single_user_decode_gqa": ("decode", {"C": 1, "P": 8, "S_kv": 64}), "multi_user_decode_gqa": ("decode", {"C": 4, "P": 8, "S_kv": 128}), } @@ -76,29 +87,49 @@ def _ccl_cfg(): # ── Per-kind launch helpers ────────────────────────────────────────── -def _run_prefill_panel(ctx, *, panel: str, C: int, S_kv: int) -> None: +def _run_prefill_panel( + ctx, *, panel: str, C: int, S_kv: int, + P: int = 1, + T_q: int = _T_Q_PREFILL, + d_head: int = _D_HEAD, +) -> None: if C > 1: - 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) + mesh_w = int(ctx.spec["sip"]["cube_mesh"]["w"]) + if C <= mesh_w: + configure_sfr_intercube_ring( + ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C, + ) + else: + if C % mesh_w != 0: + raise ValueError( + f"C={C} > mesh_w={mesh_w} requires C divisible by mesh_w" + ) + configure_sfr_intercube_ring( + ctx.engine, ctx.spec, _ccl_cfg(), + submesh_shape=(C // mesh_w, mesh_w), + ) + # Q and O switch to pe="row_wise" when intra-CUBE PE-SP is active + # (ADR-0060 §5.5 + §B-item-3: disjoint query-row split across PEs). + q_pe = "row_wise" if P > 1 else "replicate" + o_pe = "row_wise" if P > 1 else "replicate" + dp_q = DPPolicy(cube="replicate", pe=q_pe, + num_cubes=C, num_pes=P) dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="replicate", num_cubes=C, num_pes=1) + pe="replicate", num_cubes=C, num_pes=P) dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="replicate", num_cubes=C, num_pes=1) - q = ctx.zeros((_T_Q_PREFILL, _D_HEAD), + pe=o_pe, num_cubes=C, num_pes=P) + q = ctx.zeros((T_q, d_head), dtype=_DTYPE, dp=dp_q, name=f"{panel}_q") - k = ctx.zeros((S_kv, _D_HEAD), + k = ctx.zeros((S_kv, d_head), dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k") - v = ctx.zeros((S_kv, _D_HEAD), + v = ctx.zeros((S_kv, d_head), dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v") - o = ctx.empty((_T_Q_PREFILL * C, _D_HEAD), + o = ctx.empty((T_q * C, d_head), dtype=_DTYPE, dp=dp_o, name=f"{panel}_o") ctx.launch( panel, gqa_attention_prefill_long_kernel, q, k, v, o, - _T_Q_PREFILL, S_kv, _D_HEAD, C, + T_q, S_kv, d_head, C, P, _auto_dim_remap=False, ) diff --git a/src/kernbench/ccl/sfr_config.py b/src/kernbench/ccl/sfr_config.py index 81e47b5..15eba1b 100644 --- a/src/kernbench/ccl/sfr_config.py +++ b/src/kernbench/ccl/sfr_config.py @@ -248,14 +248,29 @@ def configure_sfr_intercube_ring( cfg: dict, *, ring_size: int | None = None, + submesh_shape: tuple[int, int] | None = None, + submesh_origin: tuple[int, int] = (0, 0), ) -> dict[str, Any]: - """Install intra-cube PE grid + a 1D CUBE-level ring with wrap. + """Install intra-cube PE grid + a CUBE-level ring with wrap. + + Two ring layouts: + + - **1D row** (default; ``submesh_shape=None``): cubes + ``0..ring_size-1`` form a 1D ring with wrap. ``ring_size`` must + be ≤ ``mesh_w`` so every hop is a 1-hop CUBE NOC neighbour. + - **Snake/serpentine** (``submesh_shape=(rows, cols)``, + ADR-0060 §5.5 prefill Ring KV at C=G=8 on a 2×4 sub-mesh): + a boustrophedon Hamiltonian cycle through the ``rows × cols`` + sub-mesh rooted at ``submesh_origin``. Every consecutive pair + on the snake (including the wrap) is a 1-hop CUBE NOC + neighbour, so the kernel sees a 1D logical E/W ring without + being aware of the underlying 2D layout. Direction namespaces (disjoint, same as ``configure_sfr_intercube_multisip``): - ``intra_N/S/E/W`` : 2×4 PE grid within each cube (no wrap) - - ``E/W`` : 1D ring of cubes 0..ring_size-1 WITH WRAP + - ``E/W`` : ring along the resolved path WITH WRAP (symmetric to ``configure_sfr_intracube_pe_ring`` at PE level — wrap applied at CUBE level here) - ``global_*`` : SIP topology (same as multisip) @@ -264,10 +279,15 @@ def configure_sfr_intercube_ring( ``configure_sfr_intercube_multisip`` for the full 4×4 cube mesh. Args: - ring_size: number of CUBEs in the ring (wrap applies to cubes - 0..ring_size-1). Defaults to the full cube_mesh count. - Must be ≤ mesh_w (single row); multi-row rings span - non-neighbour boundaries. + ring_size: number of CUBEs in the ring. Defaults to the full + cube_mesh count for the 1D-row case, or ``rows*cols`` for + the snake case. If passed alongside ``submesh_shape``, must + equal ``rows*cols``. + submesh_shape: ``(rows, cols)`` of the snake sub-mesh. If + given, ring follows a boustrophedon path through that + rectangle. If ``None``, falls back to the 1D-row layout. + submesh_origin: ``(row, col)`` top-left of the sub-mesh inside + the cube mesh. Defaults to ``(0, 0)``. """ cm = spec["sip"]["cube_mesh"] mesh_w = int(cm["w"]) @@ -281,13 +301,42 @@ def configure_sfr_intercube_ring( sip_w = int(sip_w) if sip_w is not None else None sip_h = int(sip_h) if sip_h is not None else None - if ring_size is None: - ring_size = n_cubes - if ring_size > mesh_w: - raise ValueError( - f"intercube_ring ring_size={ring_size} > mesh_w={mesh_w}; " - "multi-row rings cross non-neighbour boundaries" - ) + if submesh_shape is not None: + sub_h, sub_w = submesh_shape + origin_row, origin_col = submesh_origin + if (sub_h <= 0 or sub_w <= 0 + or origin_row < 0 or origin_col < 0 + or origin_row + sub_h > mesh_h + or origin_col + sub_w > mesh_w): + raise ValueError( + f"submesh_shape={submesh_shape} at origin={submesh_origin} " + f"does not fit cube_mesh ({mesh_h}x{mesh_w})" + ) + expected_size = sub_h * sub_w + if ring_size is not None and ring_size != expected_size: + raise ValueError( + f"ring_size={ring_size} inconsistent with " + f"submesh_shape={submesh_shape} (expected {expected_size})" + ) + ring_size = expected_size + # Boustrophedon: even rows L→R, odd rows R→L. + ring_path: list[int] = [] + for r in range(sub_h): + cols = range(sub_w) if r % 2 == 0 else range(sub_w - 1, -1, -1) + for c in cols: + ring_path.append((origin_row + r) * mesh_w + (origin_col + c)) + else: + if ring_size is None: + ring_size = n_cubes + if ring_size > mesh_w: + raise ValueError( + f"intercube_ring ring_size={ring_size} > mesh_w={mesh_w}; " + "multi-row rings cross non-neighbour boundaries" + ) + ring_path = list(range(ring_size)) + + ring_pos: dict[int, int] = {c: i for i, c in enumerate(ring_path)} + ring_len = len(ring_path) if sip_topology not in _TOPO_BUILTINS: raise ValueError( @@ -329,10 +378,11 @@ def configure_sfr_intercube_ring( for d, peer_pe in _intra_cube_neighbors(pe).items(): nbrs[d] = _pe_idx(sip, cube, peer_pe) - # ── Cube ring (E/W with wrap for cubes 0..ring_size-1) ── - if cube < ring_size: - nbrs["E"] = _pe_idx(sip, (cube + 1) % ring_size, pe) - nbrs["W"] = _pe_idx(sip, (cube - 1) % ring_size, pe) + # ── Cube ring (E/W along resolved ring_path, with wrap) ── + pos = ring_pos.get(cube) + if pos is not None: + nbrs["E"] = _pe_idx(sip, ring_path[(pos + 1) % ring_len], pe) + nbrs["W"] = _pe_idx(sip, ring_path[(pos - 1) % ring_len], pe) # ── Inter-SIP same-(cube, pe) (global_*) ── if n_sips > 1: diff --git a/tests/attention/test_gqa_decode_long_2d_reduce.py b/tests/attention/test_gqa_decode_long_2d_reduce.py new file mode 100644 index 0000000..107aa4f --- /dev/null +++ b/tests/attention/test_gqa_decode_long_2d_reduce.py @@ -0,0 +1,191 @@ +"""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_prefill_long_c8_snake.py b/tests/attention/test_gqa_prefill_long_c8_snake.py new file mode 100644 index 0000000..58ca62e --- /dev/null +++ b/tests/attention/test_gqa_prefill_long_c8_snake.py @@ -0,0 +1,154 @@ +"""Tests for prefill_long Ring KV at C=8 over a 2×4 snake sub-mesh. + +ADR-0060 §5.5 design target for the single-KV-group LLaMA-3.1-70B +configuration: ``C = G = 8`` (one Q head per CUBE), Ring KV rotates +the 8 KV slices around all 8 CUBEs. + +Increment 1 wired ``configure_sfr_intercube_ring(submesh_shape=(2, 4))`` +to map the kernel's logical E/W to a snake/serpentine 1-hop path +through the top 2×4 sub-mesh of the 4×4 SIP CUBE mesh. The kernel +itself uses only logical ``dir="W"`` / ``dir="E"`` — the snake is +transparent at kernel level. + +This file verifies the assembly works end-to-end at C=8: + T1 kernel completes (no scratch overflow, no deadlock) + T2 per-CUBE head-parallel output (8 dma_writes, one per CUBE) + T3 tile-granular ring traffic at C=8 follows the same + ``(C-1)·n_tiles·2·C`` formula as the existing tile-ring tests +""" +from __future__ import annotations + +from pathlib import Path + +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_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) + + +def _run_prefill_c8_snake(*, T_q: int, S_kv: int): + """Head-parallel prefill at C=8 with snake-mapped Ring KV over the + top 2×4 sub-mesh of the 4×4 SIP CUBE mesh.""" + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + C = 8 + + def _bench_fn(ctx): + configure_sfr_intercube_ring( + ctx.engine, ctx.spec, _ccl_cfg(), + submesh_shape=(2, 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=f"q_c8_snake_t{T_q}_s{S_kv}") + k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, + name=f"k_c8_snake_t{T_q}_s{S_kv}") + v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, + name=f"v_c8_snake_t{T_q}_s{S_kv}") + o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o, + name=f"o_c8_snake_t{T_q}_s{S_kv}") + ctx.launch( + f"gqa_prefill_long_c8_snake_t{T_q}_s{S_kv}", + 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, + ) + + +# ── T1: kernel completes end-to-end at C=8 ─────────────────────────── + + +def test_prefill_long_c8_snake_completes(): + """ADR-0060 §5.5 design target: prefill Ring KV at C=G=8 over the + 2×4 snake sub-mesh. With zero inputs, output is trivially zero; + this test guards that the kernel reaches completion without + deadlock, scratch overflow, or routing failure. + + Configuration: T_q=4, S_kv=8192 → S_local=1024, n_tiles=1 + (TILE_S_KV=1024). Bounded scratch. + """ + result = _run_prefill_c8_snake(T_q=4, S_kv=8192) + assert result.completion.ok, ( + f"prefill at C=8 with snake ring must complete; " + f"got {result.completion}" + ) + + +# ── T2: 8 dma_writes (head-parallel, one head per CUBE) ────────────── + + +def test_prefill_long_c8_snake_distributed_output_count(): + """Head-parallel: each CUBE owns one Q head and writes its own + head's output (T_q, d_head) rows. No inter-CUBE reduce on the + output side — so ``dma_write_count == C == 8``. + + This is the C=8 analogue of + ``test_prefill_long_tile_ring_dma_write_count`` at C=4. + """ + result = _run_prefill_c8_snake(T_q=4, S_kv=8192) + assert result.completion.ok + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 8, ( + f"head-parallel C=8: expected 8 dma_writes (one per CUBE); " + f"got {n_writes}" + ) + + +# ── T3: tile-granular ring ipcq_copy count scales as (C-1)·n_tiles·2·C + + +def test_prefill_long_c8_snake_tile_ipcq_count(): + """ADR-0060 §5.5.1 tile-granular ring: per-CUBE send count is + ``2·n_tiles·(C-1)`` (K + V per tile per ring step). Aggregated + across C CUBEs, total ipcq_copy = ``(C-1)·n_tiles·2·C``. + + Configuration: T_q=4, S_kv=8192, C=8 → S_local=1024, n_tiles=1 + (TILE_S_KV=1024). Expected total = ``7·1·2·8 = 112``. + + Verifies that the snake-mapped 1-hop physical links carry the + same logical ring traffic as the existing C=4 1D-row ring tests. + """ + C = 8 + n_tiles = 1 # S_local=1024 / TILE_S_KV=1024 + result = _run_prefill_c8_snake(T_q=4, S_kv=8192) + assert result.completion.ok + n_copy = _count(result.engine.op_log, "ipcq_copy") + expected = (C - 1) * n_tiles * 2 * C + assert n_copy == expected, ( + f"tile-granular ring at C=8: expected {expected} ipcq_copy " + f"((C-1)·n_tiles·2·C = {C - 1}·{n_tiles}·2·{C}); got {n_copy}" + ) diff --git a/tests/attention/test_gqa_prefill_long_pe_sp.py b/tests/attention/test_gqa_prefill_long_pe_sp.py new file mode 100644 index 0000000..8c8b775 --- /dev/null +++ b/tests/attention/test_gqa_prefill_long_pe_sp.py @@ -0,0 +1,233 @@ +"""Tests for intra-CUBE PE-SP in prefill_long (query-axis split). + +ADR-0060 §5.5 last bullet + §B-item-3: split the head's query rows +``[T_q, d_head]`` across the ``P`` PEs of each CUBE so all P PEs work +in parallel. Output rows are disjoint across PEs ⇒ no intra-CUBE +reduce needed. + +Same-lane SFR wiring: PE ``i`` in CUBE A has its own E/W ring link to +PE ``i`` in CUBE B (the snake's prev/next). All P PEs of a CUBE see +the same K, V (HBM-resident, ``pe="replicate"``) ⇒ P parallel rings +run in lockstep, each PE rotating its own K/V copies via its own IPCQ +channels. + +Activation contract: + - ``P == 1`` (default; omitted from launch args) → existing + PE-0-only behavior. Byte-for-byte unchanged. + - ``P > 1`` → all P PEs active; each handles ``T_q // P`` query + rows. Requires ``T_q % P == 0`` (degenerate T_q < P is rejected + — caller must use ``P=1`` for that workload, ADR-0060 §B-item-3). + +Phase 1: tests only — production code lands in Phase 2. +T1, T2, T3, T5 fail today (TypeError: kernel signature has no P). +T4 passes today as the backward-compat anchor. +""" +from __future__ import annotations + +from pathlib import Path + +import pytest + +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_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) + + +def _run_prefill_pe_sp( + *, T_q: int, S_kv: int, C: int, P: int | None, + snake: bool, +): + """Drive a prefill_long launch with optional intra-CUBE PE-SP. + + ``P=None`` → omit P from the launch args (exercises the kernel's + default behaviour). + ``snake=True`` → install the 2×4 snake ring SFR (Increment 1); + else install the 1D-row ring at ``ring_size=C``. + """ + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + if snake: + configure_sfr_intercube_ring( + ctx.engine, ctx.spec, _ccl_cfg(), + submesh_shape=(2, 4), + ) + else: + configure_sfr_intercube_ring( + ctx.engine, ctx.spec, _ccl_cfg(), + ring_size=C, + ) + num_pes = P if (P is not None and P > 1) else 1 + # When PE-SP is active, Q and O are split row-wise across PEs; + # K, V remain replicated within a CUBE. + q_pe = "row_wise" if num_pes > 1 else "replicate" + o_pe = "row_wise" if num_pes > 1 else "replicate" + dp_q = DPPolicy(cube="replicate", pe=q_pe, + num_cubes=C, num_pes=num_pes) + dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", + pe="replicate", num_cubes=C, num_pes=num_pes) + dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate", + pe=o_pe, num_cubes=C, num_pes=num_pes) + suffix = f"c{C}_p{P}_t{T_q}_s{S_kv}" + q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q, + name=f"q_pe_sp_{suffix}") + k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, + name=f"k_pe_sp_{suffix}") + v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, + name=f"v_pe_sp_{suffix}") + o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o, + name=f"o_pe_sp_{suffix}") + launch_args = [q, k, v, o, T_q, S_kv, D_HEAD, C] + if P is not None: + launch_args.append(P) + ctx.launch( + f"gqa_prefill_long_pe_sp_{suffix}", + gqa_attention_prefill_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 P=8 completes end-to-end ───────────────────────────────── + + +def test_prefill_c8_p8_pe_sp_completes(): + """ADR-0060 §5.5 + §B-item-3 design target: 64 ranks active + (8 CUBEs × 8 PEs), query-axis split across PEs, snake-mapped + Ring KV. Guards no scratch overflow, no deadlock across the + P parallel same-lane rings. + """ + result = _run_prefill_pe_sp( + T_q=8, S_kv=8192, C=8, P=8, snake=True, + ) + assert result.completion.ok, ( + f"prefill at C=8, P=8 (PE-SP) must complete; " + f"got {result.completion}" + ) + + +# ── T2: 64 dma_writes (one per PE, disjoint query rows) ────────────── + + +def test_prefill_c8_p8_pe_sp_64_dma_writes(): + """With query-axis split, each PE owns ``T_q/P = 1`` query row + and stores its own ``(1, d_head)`` slice. Across all 64 ranks + (8 CUBEs × 8 PEs), ``dma_write_count == 64``. + + This is the structural signal that PE-SP is actually wired: + PE-0-only would give 8 dma_writes (one per CUBE). + """ + result = _run_prefill_pe_sp( + T_q=8, S_kv=8192, C=8, P=8, snake=True, + ) + assert result.completion.ok + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 64, ( + f"PE-SP at C=8, P=8: expected 64 dma_writes " + f"(8 cubes × 8 PEs, each writes its T_q/P=1 query row); " + f"got {n_writes}" + ) + + +# ── T3: P parallel rings — ipcq_copy scales by P ───────────────────── + + +def test_prefill_c8_p8_pe_sp_ring_ipcq_count(): + """Per-PE same-lane rings: each PE runs its own ring traffic via + its own IPCQ channels. Total inter-CUBE ipcq_copy: + ``(C-1) · n_tiles · 2 · C · P`` (= existing C=8 formula × P). + + Configuration: T_q=8, S_kv=8192, C=8 → S_local=1024, n_tiles=1 + (TILE_S_KV=1024). Expected = ``7·1·2·8·8`` = **896**. + """ + C = 8 + P = 8 + n_tiles = 1 + result = _run_prefill_pe_sp( + T_q=8, S_kv=8192, C=C, P=P, snake=True, + ) + assert result.completion.ok + n_copy = _count(result.engine.op_log, "ipcq_copy") + expected = (C - 1) * n_tiles * 2 * C * P + assert n_copy == expected, ( + f"PE-SP parallel rings at C=8, P=8: expected {expected} " + f"ipcq_copy ((C-1)·n_tiles·2·C·P = " + f"{C - 1}·{n_tiles}·2·{C}·{P}); got {n_copy}" + ) + + +# ── T4: backward-compat — P omitted → existing PE-0-only behaviour ─── + + +def test_prefill_p_default_1_backward_compat(): + """When ``P`` is omitted from the launch args, the kernel must + behave identically to today: only PE 0 of each CUBE participates; + one head per CUBE; one dma_write per CUBE. + + At C=4 this gives 4 dma_writes (matches the existing + ``test_prefill_long_tile_ring_dma_write_count``). Phase 2 must + not regress this path. + + Passes today AND after Phase 2. + """ + C = 4 + result = _run_prefill_pe_sp( + T_q=4, S_kv=8192, C=C, P=None, snake=False, + ) + assert result.completion.ok, ( + f"prefill at C=4 with default P (PE-0-only) must complete; " + f"got {result.completion}" + ) + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == C, ( + f"default P=1 (PE-0-only): expected {C} dma_writes " + f"(one per CUBE); got {n_writes}" + ) + + +# ── T5: validation — T_q must be divisible by P when P > 1 ─────────── + + +def test_prefill_pe_sp_rejects_non_divisible_t_q(): + """ADR-0060 §B-item-3 fallback (KV-block split + intra-CUBE + reduce for T_q < P) is deferred. Callers must request ``P=1`` for + workloads where T_q < P or T_q % P != 0. + + C=4, P=8, T_q=4 violates ``T_q % P == 0`` (and also T_q < P). + The kernel must raise ValueError with a clear error message + rather than silently producing wrong results. + """ + with pytest.raises((ValueError, AssertionError), match=r"T_q"): + _run_prefill_pe_sp( + T_q=4, S_kv=8192, C=4, P=8, snake=False, + ) diff --git a/tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py b/tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py new file mode 100644 index 0000000..0ba67d5 --- /dev/null +++ b/tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py @@ -0,0 +1,146 @@ +"""Tests for the single-KV-group prefill panel in milestone_gqa_headline. + +The single-KV-group (LLaMA-3.1-70B target) prefill panel wires together +all Increment 1–4 work in a single bench panel: + - C=8 head-parallel + Ring KV via snake-mapped 2×4 sub-mesh (Inc 1, 3) + - P=8 intra-CUBE PE-SP (Inc 4) + - d_head=128 (LLaMA-3.1-70B), one-shot prefill T_q = S_kv = 32K + +This file verifies the bench-config wiring: + T1 the new ``single_kv_group_prefill_gqa_c8_p8`` panel is registered + with the expected dims in ``_PANELS`` + ``_PANEL_DISPATCH``. + T2 ``_run_prefill_panel`` accepts the new ``P``, ``T_q``, ``d_head`` + kwargs and drives the prefill kernel to completion at the + milestone-target ``(C, P) = (8, 8)``. Uses smaller T_q/S_kv to + keep test time bounded — full 32K runs come from + ``kernbench run milestone-gqa-headline``. + T3 Existing prefill panels still work when ``_run_prefill_panel`` + is called without the new kwargs (backward compat anchor). + +Phase 1: tests only — production code lands in Phase 2. +T1 fails today (panel not registered). +T2 fails today (helper doesn't accept the new kwargs). +T3 passes today as the backward-compat anchor. +""" +from __future__ import annotations + +from pathlib import Path + +from kernbench.benches.milestone_gqa_headline import ( # noqa: F401 + _PANEL_DISPATCH, + _PANELS, + _run_prefill_panel, +) +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" + + +def _engine_factory(t, d): + return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) + + +# ── T1: new panel registered ───────────────────────────────────────── + + +def test_single_kv_group_prefill_panel_registered(): + """The ``single_kv_group_prefill_gqa_c8_p8`` panel must be in ``_PANELS`` and + its dispatch entry must have the expected LLaMA-3.1-70B params. + + Headline config: + C = 8 (head-parallel, snake sub-mesh) + P = 8 (intra-CUBE PE-SP) + T_q = 32_768 (one-shot long-context prefill) + S_kv = 32_768 + d_head = 128 (LLaMA-3.1-70B) + """ + panel_name = "single_kv_group_prefill_gqa_c8_p8" + assert panel_name in _PANELS, ( + f"{panel_name!r} not in _PANELS; got {_PANELS}" + ) + assert panel_name in _PANEL_DISPATCH, ( + f"{panel_name!r} not in _PANEL_DISPATCH" + ) + kind, params = _PANEL_DISPATCH[panel_name] + assert kind == "prefill", f"kind={kind!r}, expected 'prefill'" + assert params.get("C") == 8, f"C={params.get('C')}, expected 8" + assert params.get("P") == 8, f"P={params.get('P')}, expected 8" + assert params.get("T_q") == 32_768, ( + f"T_q={params.get('T_q')}, expected 32_768" + ) + assert params.get("S_kv") == 32_768, ( + f"S_kv={params.get('S_kv')}, expected 32_768" + ) + assert params.get("d_head") == 128, ( + f"d_head={params.get('d_head')}, expected 128" + ) + + +# ── T2: helper drives the C=8, P=8 prefill kernel to completion ────── + + +def test_single_kv_group_prefill_panel_runner_smoke(): + """``_run_prefill_panel`` must accept the new ``P``, ``T_q``, + ``d_head`` kwargs and successfully launch the prefill kernel at + ``(C, P) = (8, 8)`` with the snake-mapped 2×4 SFR. + + Uses **smaller** T_q/S_kv than the headline panel so the test + completes quickly. The headline 32K dims are exercised via + ``kernbench run milestone-gqa-headline``, not pytest. + """ + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + _run_prefill_panel( + ctx, + panel="single_kv_group_prefill_gqa_c8_p8", + C=8, P=8, + T_q=8, S_kv=8192, # smoke dims, not the 32K headline + d_head=128, + ) + + result = run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), + engine_factory=_engine_factory, + ) + assert result.completion.ok, ( + f"single_kv_group prefill panel smoke at C=8 P=8 must complete; " + f"got {result.completion}" + ) + + +# ── T3: existing panel calls still work (backward compat) ──────────── + + +def test_existing_prefill_panel_runner_backward_compat(): + """Existing callers of ``_run_prefill_panel`` (no ``P``/``T_q``/ + ``d_head`` overrides) must continue to work after the signature + extension. Exercises the ``multi_user_prefill_gqa`` panel dims + from ``_PANEL_DISPATCH``. + + Passes today AND after Phase 2. + """ + panel_name = "multi_user_prefill_gqa" + _kind, params = _PANEL_DISPATCH[panel_name] + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + _run_prefill_panel( + ctx, + panel=panel_name, + C=params["C"], S_kv=params["S_kv"], + ) + + result = run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), + engine_factory=_engine_factory, + ) + assert result.completion.ok, ( + f"existing prefill panel {panel_name!r} must still run with " + f"the original kwargs; got {result.completion}" + ) diff --git a/tests/test_intercube_snake_ring.py b/tests/test_intercube_snake_ring.py new file mode 100644 index 0000000..130f28e --- /dev/null +++ b/tests/test_intercube_snake_ring.py @@ -0,0 +1,273 @@ +"""Tests for snake/serpentine ring extension of configure_sfr_intercube_ring. + +Verifies the multi-row snake ring SFR wiring used for ADR-0060 §5.5 +prefill Ring KV at C=8 on a 2×4 sub-mesh of the 4×4 CUBE mesh. + +The snake path for ``submesh_shape=(2, 4), origin=(0, 0)`` is:: + + (0,0)→(0,1)→(0,2)→(0,3)→(1,3)→(1,2)→(1,1)→(1,0)→wrap to (0,0) + +In cube indices on a mesh_w=4 grid: ``[0, 1, 2, 3, 7, 6, 5, 4]``. + +Every consecutive pair (including the wrap) is a 1-hop CUBE NOC neighbour. +The kernel sees a 1D logical E/W ring; the snake is invisible to it. +""" +from __future__ import annotations + +from pathlib import Path + +import pytest + +from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config +from kernbench.ccl.sfr_config import configure_sfr_intercube_ring +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml" + +N_CUBES = 16 # 4×4 SIP CUBE mesh +PES_PER_CUBE = 8 +MESH_W = 4 + +# Canonical snake path for (2,4) sub-mesh at origin (0,0). +SNAKE_2X4_O00 = [0, 1, 2, 3, 7, 6, 5, 4] + +# Canonical snake path for (2,4) sub-mesh at origin (1,0) (rows 1-2). +SNAKE_2X4_O10 = [4, 5, 6, 7, 11, 10, 9, 8] + + +def _engine_and_spec(): + topo = resolve_topology(str(TOPOLOGY_PATH)) + engine = GraphEngine(topo.topology_obj, enable_data=True) + return engine, topo.topology_obj.spec + + +def _merged_cfg(): + cfg = load_ccl_config() + return resolve_algorithm_config(cfg, name="lrab_hierarchical_allreduce") + + +def _qp(engine, sip: int, cube: int, pe: int): + return engine._components[f"sip{sip}.cube{cube}.pe{pe}.pe_ipcq"].queue_pairs + + +def _row_col(cube: int) -> tuple[int, int]: + return cube // MESH_W, cube % MESH_W + + +def _l1(c1: int, c2: int) -> int: + r1, col1 = _row_col(c1) + r2, col2 = _row_col(c2) + return abs(r1 - r2) + abs(col1 - col2) + + +class TestSnakeRing2x4Origin00: + """Snake ring through the top 2×4 sub-mesh of the 4×4 SIP.""" + + def test_snake_ring_2x4_world_size(self): + """All PEs on all CUBEs of all SIPs still enumerated (snake only + restricts which CUBEs have ring links, not the world).""" + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + plan = configure_sfr_intercube_ring( + engine, spec, cfg, submesh_shape=(2, 4), + ) + n_sips = int(spec["system"]["sips"]["count"]) + assert plan["world_size"] == n_sips * N_CUBES * PES_PER_CUBE + assert len(plan["rank_to_pe"]) == plan["world_size"] + + def test_snake_ring_2x4_path_corners_E(self): + """The four 'interesting' E hops on the snake path: + + cube0.E → cube1 (start of row 0) + cube3.E → cube7 (row-bridge: top-right corner down) + cube7.E → cube6 (row 1, going leftward) + cube4.E → cube0 (wrap: bottom-left to top-left) + """ + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + configure_sfr_intercube_ring( + engine, spec, cfg, submesh_shape=(2, 4), + ) + assert _qp(engine, 0, 0, 0)["E"]["peer"].cube == 1 + assert _qp(engine, 0, 3, 0)["E"]["peer"].cube == 7 + assert _qp(engine, 0, 7, 0)["E"]["peer"].cube == 6 + assert _qp(engine, 0, 4, 0)["E"]["peer"].cube == 0 + + def test_snake_ring_2x4_path_corners_W(self): + """Symmetric W hops (W is reverse of E along the snake): + + cube0.W → cube4 (wrap reverse) + cube7.W → cube3 (row-bridge reverse: up) + cube6.W → cube7 (row 1 reverse: rightward) + cube4.W → cube5 (row-1 leftmost, W goes right within row 1) + + Note: under the snake path [0,1,2,3,7,6,5,4], W(cube_i) is the + predecessor on that path. So W(4)=5 (previous on path). + """ + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + configure_sfr_intercube_ring( + engine, spec, cfg, submesh_shape=(2, 4), + ) + assert _qp(engine, 0, 0, 0)["W"]["peer"].cube == 4 + assert _qp(engine, 0, 7, 0)["W"]["peer"].cube == 3 + assert _qp(engine, 0, 6, 0)["W"]["peer"].cube == 7 + assert _qp(engine, 0, 4, 0)["W"]["peer"].cube == 5 + + def test_snake_ring_2x4_interior(self): + """Interior hops along each row (not at row-bridge or wrap).""" + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + configure_sfr_intercube_ring( + engine, spec, cfg, submesh_shape=(2, 4), + ) + # Row 0 interior: cube1 between cube0 and cube2. + assert _qp(engine, 0, 1, 0)["E"]["peer"].cube == 2 + assert _qp(engine, 0, 1, 0)["W"]["peer"].cube == 0 + # Row 1 interior: cube5 between cube6 and cube4 along the snake. + assert _qp(engine, 0, 5, 0)["E"]["peer"].cube == 4 + assert _qp(engine, 0, 5, 0)["W"]["peer"].cube == 6 + + def test_snake_ring_2x4_all_hops_are_1_hop(self): + """Every E hop on the snake is L1-distance 1 in the cube mesh. + + Defensive — holds by construction (snake = Hamiltonian cycle + on the 2×4 grid graph), but explicit guards against future + path-builder regressions. + """ + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + configure_sfr_intercube_ring( + engine, spec, cfg, submesh_shape=(2, 4), + ) + for cube in SNAKE_2X4_O00: + qp = _qp(engine, 0, cube, 0) + e_peer = qp["E"]["peer"].cube + w_peer = qp["W"]["peer"].cube + assert _l1(cube, e_peer) == 1, ( + f"snake E hop cube{cube}→cube{e_peer} is not 1-hop" + ) + assert _l1(cube, w_peer) == 1, ( + f"snake W hop cube{cube}→cube{w_peer} is not 1-hop" + ) + + def test_snake_ring_2x4_off_path_cubes_have_no_ring_links(self): + """CUBEs not on the snake path (rows 2-3, i.e. cubes 8..15) + get no E/W ring entries — consistent with the current + ``if cube < ring_size`` behaviour.""" + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + configure_sfr_intercube_ring( + engine, spec, cfg, submesh_shape=(2, 4), + ) + for cube in range(8, 16): + qp = _qp(engine, 0, cube, 0) + assert "E" not in qp, ( + f"sip0.cube{cube}.pe0 has E but is off the snake path" + ) + assert "W" not in qp + + def test_snake_ring_2x4_intra_namespace_unchanged(self): + """Snake doesn't disturb the intra-cube 2×4 PE grid wiring + (intra_N/S/E/W) — it only changes cube-level E/W.""" + from kernbench.ccl.sfr_config import _intra_cube_neighbors + + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + configure_sfr_intercube_ring( + engine, spec, cfg, submesh_shape=(2, 4), + ) + # Spot-check a few cubes (including off-path) and all 8 PEs. + for cube in (0, 3, 4, 7, 10, 15): + for pe in range(PES_PER_CUBE): + qp = _qp(engine, 0, cube, pe) + expected = _intra_cube_neighbors(pe) + for d, expected_pe in expected.items(): + assert d in qp, f"cube{cube}.pe{pe} missing {d}" + assert qp[d]["peer"].pe == expected_pe + assert qp[d]["peer"].cube == cube + + +class TestSnakeRing2x4OriginShift: + """Snake ring through rows 1-2 (origin=(1, 0)) — proves the + origin parameter shifts the sub-mesh.""" + + def test_snake_ring_origin_shift_path(self): + """For origin=(1, 0), snake path is [4,5,6,7,11,10,9,8]: + + row 1 left→right: cube4 → cube5 → cube6 → cube7 + ↓ + row 2 right→left: cube11 → cube10 → cube9 → cube8 + ↺ wraps to cube4 + """ + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + configure_sfr_intercube_ring( + engine, spec, cfg, submesh_shape=(2, 4), submesh_origin=(1, 0), + ) + # Path corners. + assert _qp(engine, 0, 4, 0)["E"]["peer"].cube == 5 + assert _qp(engine, 0, 7, 0)["E"]["peer"].cube == 11 + assert _qp(engine, 0, 11, 0)["E"]["peer"].cube == 10 + assert _qp(engine, 0, 8, 0)["E"]["peer"].cube == 4 # wrap + + def test_snake_ring_origin_shift_off_path_no_links(self): + """For origin=(1, 0), cubes 0..3 (row 0) and 12..15 (row 3) + get no ring links.""" + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + configure_sfr_intercube_ring( + engine, spec, cfg, submesh_shape=(2, 4), submesh_origin=(1, 0), + ) + for cube in list(range(0, 4)) + list(range(12, 16)): + qp = _qp(engine, 0, cube, 0) + assert "E" not in qp + assert "W" not in qp + + +class TestSnakeRingBackwardCompat: + """The existing 1D-row API (no submesh_shape) must behave identically + to before the snake extension lands.""" + + def test_snake_ring_backward_compat_1d_default(self): + """ring_size=4 (single-row) without submesh_shape — identical + behaviour to today's 1D-row ring: cube0..3 with E/W wrap; + cubes 4..15 have no ring links.""" + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + configure_sfr_intercube_ring( + engine, spec, cfg, ring_size=4, + ) + # Row-0 1D ring with wrap. + assert _qp(engine, 0, 0, 0)["E"]["peer"].cube == 1 + assert _qp(engine, 0, 0, 0)["W"]["peer"].cube == 3 # wrap + assert _qp(engine, 0, 3, 0)["E"]["peer"].cube == 0 # wrap + # Off-row cubes get no ring links. + for cube in range(4, 16): + qp = _qp(engine, 0, cube, 0) + assert "E" not in qp + assert "W" not in qp + + +class TestSnakeRingValidation: + """Input validation of the submesh_shape / ring_size combinations.""" + + def test_snake_ring_invalid_submesh_overflow(self): + """submesh_shape that doesn't fit the cube mesh is rejected.""" + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + with pytest.raises(ValueError, match=r"sub.?mesh"): + configure_sfr_intercube_ring( + engine, spec, cfg, submesh_shape=(2, 5), # col overflow + ) + + def test_snake_ring_ring_size_mismatch(self): + """submesh_shape and ring_size disagree → ValueError.""" + engine, spec = _engine_and_spec() + cfg = _merged_cfg() + with pytest.raises(ValueError, match=r"ring_size"): + configure_sfr_intercube_ring( + engine, spec, cfg, + submesh_shape=(2, 4), ring_size=4, + )