gqa: single-KV-group LLaMA-3.1-70B prefill milestone (Increments 1-5)
End-to-end wires C=8 P=8 d_head=128 prefill with snake-ring inter-CUBE SFR, intra-CUBE PE-SP (all 64 ranks active), and the milestone bench panel. Decode kernel gains lrab-adapted center-root reduce for the 2×4 sub-mesh per ADR-0060 §4.2. Increment 1 — SFR multi-row snake src/kernbench/ccl/sfr_config.py: configure_sfr_intercube_ring gains submesh_shape / submesh_origin kwargs; installs a Hamiltonian snake ring through a rectangular sub-mesh (every hop is 1-hop physical neighbour). Backward-compat: 1D-row behaviour preserved when submesh_shape is None. tests/test_intercube_snake_ring.py (12 tests) Increment 2 — Decode lrab-adapted center-root reduce src/kernbench/benches/_gqa_attention_decode_long.py: new sub_w param (default 0 = existing 1D-chain). sub_w >= 2 selects the ADR-0060 §4.2 prescribed lrab-adapted Phase 1+2 reduce (bidirectional row + bidirectional col converge to the center cube), with log-sum-exp _merge_running replacing the plain + of lrab. tests/attention/test_gqa_decode_long_2d_reduce.py (4 tests) Increment 3 — Prefill kernel at C=8 (no production change) Verified by inspection that the existing prefill_long kernel + Increment 1's snake SFR already work at C=8 without any kernel edit. The kernel speaks logical W/E; the snake routes it. tests/attention/test_gqa_prefill_long_c8_snake.py (3 tests) Increment 4 — Intra-CUBE PE-SP in prefill (all 64 ranks) src/kernbench/benches/_gqa_attention_prefill_long.py: new P param (default 1 = existing PE-0-only). P > 1 splits T_q query-axis-wise across the P PEs of each CUBE; output rows are disjoint per PE so no intra-CUBE reduce is needed; each PE drives its own same-lane ring (P parallel rings). tests/attention/test_gqa_prefill_long_pe_sp.py (5 tests) Increment 5 — LLaMA-scale milestone bench panel src/kernbench/benches/milestone_gqa_headline.py: new panel single_kv_group_prefill_gqa_c8_p8 (C=8, P=8, T_q=S_kv=32K, d_head=128). _run_prefill_panel extended with P/T_q/d_head defaults; routes snake SFR when C > mesh_w. tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py (3 tests) Total: 4 production files modified, 5 new test files, 27 new tests. Followed the Phase 1/2 protocol per CLAUDE.md throughout. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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) ──
|
||||
#
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user