ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gqa: ADR-0060/0062/0063/0064 unified GQA kernels + CPU cost model

Browse Source 
Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.

ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
  intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
  online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
  (m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.

ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
  owns whole KV heads; PE-parallel heads with intra-group chain
  reduce. Prefill has no Ring KV (each head fully resident).

ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).

ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.

ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
  primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
  consults table with dispatch_cycles fallback (ADR-0046 §D6
  back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
  DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
  measurable end-to-end.

P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).

Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
  _attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
  (EN + KO kept in sync).

Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-06-09 18:15:59 -07:00
parent b3730a33eb
commit d282144339
52 changed files with 3952 additions and 2526 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/benches/1H_milestone_output/gqa/gqa_comparison.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 33 KiB

src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_multi_user_decode.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 21 KiB

src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_multi_user_prefill.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 20 KiB

src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_single_user_decode.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 24 KiB

src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_single_user_prefill.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 22 KiB

src/kernbench/benches/1H_milestone_output/gqa/sweep.json
-65
View File
@@ -1,65 +0,0 @@
{
"version": 1,
"validation_scale": true,
"panels": [
"single_user_prefill",
"multi_user_prefill",
"single_user_decode",
"multi_user_decode"
],
"config": {
"S_q_prefill": 16,
"S_kv_per_rank": 16,
"h_q": 1,
"h_kv": 1,
"d_head": 64,
"n_ranks_single_user": 8,
"n_ranks_multi_user": 4
},
"rows": [
{
"panel": "single_user_prefill",
"n_ranks": 8,
"op_log_summary": {
"gemm_count": 128,
"ipcq_send_count": 112,
"ipcq_recv_count": 112,
"dma_read_count": 24,
"dma_write_count": 8
}
},
{
"panel": "multi_user_prefill",
"n_ranks": 4,
"op_log_summary": {
"gemm_count": 32,
"ipcq_send_count": 24,
"ipcq_recv_count": 24,
"dma_read_count": 12,
"dma_write_count": 4
}
},
{
"panel": "single_user_decode",
"n_ranks": 8,
"op_log_summary": {
"gemm_count": 16,
"ipcq_send_count": 168,
"ipcq_recv_count": 168,
"dma_read_count": 24,
"dma_write_count": 8
}
},
{
"panel": "multi_user_decode",
"n_ranks": 4,
"op_log_summary": {
"gemm_count": 8,
"ipcq_send_count": 36,
"ipcq_recv_count": 36,
"dma_read_count": 12,
"dma_write_count": 4
}
}
]
}
src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json
+69
View File
@@ -0,0 +1,69 @@
{
"version": 1,
"panels": [
"single_user_prefill_gqa",
"multi_user_prefill_gqa",
"single_user_decode_gqa",
"multi_user_decode_gqa"
],
"config": {
"T_q_prefill": 4,
"T_q_decode": 1,
"S_kv_prefill": 16,
"h_q_decode": 8,
"h_kv_decode": 1,
"d_head": 64
},
"rows": [
{
"panel": "single_user_prefill_gqa",
"kind": "prefill",
"C": 1,
"S_kv": 16,
"op_log_summary": {
"gemm_count": 2,
"ipcq_copy_count": 0,
"dma_read_count": 3,
"dma_write_count": 1
}
},
{
"panel": "multi_user_prefill_gqa",
"kind": "prefill",
"C": 4,
"S_kv": 16,
"op_log_summary": {
"gemm_count": 32,
"ipcq_copy_count": 24,
"dma_read_count": 12,
"dma_write_count": 4
}
},
{
"panel": "single_user_decode_gqa",
"kind": "decode",
"C": 1,
"P": 8,
"S_kv": 64,
"op_log_summary": {
"gemm_count": 16,
"ipcq_copy_count": 21,
"dma_read_count": 24,
"dma_write_count": 1
}
},
{
"panel": "multi_user_decode_gqa",
"kind": "decode",
"C": 4,
"P": 8,
"S_kv": 128,
"op_log_summary": {
"gemm_count": 64,
"ipcq_copy_count": 93,
"dma_read_count": 96,
"dma_write_count": 1
}
}
]
}
src/kernbench/benches/_attention_mesh_kv.py
-193
View File
@@ -1,193 +0,0 @@
"""Mesh-native bidirectional Ring-K/V attention kernel — prefill (ADR-0059 Proposed).
Each rank holds its own Q tile and 1/n_ranks of K, V (sequence-sharded).
Over ``n_ranks - 1`` bidirectional steps, K and V propagate both east and
west: chunk c_i originating at rank i reaches rank j at step ``|i - j|``.
Every rank receives every other rank's chunk **exactly once** and folds it
into a running ``(m, ℓ, o)`` via the online-softmax recurrence. After all
steps each rank holds the final attention output for its own Q tokens —
no cross-rank merge is required.
Supersedes ADR-0055's closed-ring ``_attention_ring_kv.py``. Both modules
stay on disk during the transition; this one runs on the hardware's
actual open-mesh wiring (no closed-ring SFR install required).
Imported by ``milestone_gqa_llama70b`` (after the bench's Phase 2 switches
its imports) and invoked through ``torch.launch(...)`` — not through
``dist.all_reduce(...)``. See ADR-0055 Context for why this kernel is not
backend-dispatched via ADR-0050's algorithm-module contract.
"""
from __future__ import annotations
from kernbench.common.pe_commands import TensorHandle
def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle:
"""Reshape — metadata only, no command emitted (cf. ``tl.trans``)."""
return TensorHandle(
id=handle.id,
addr=handle.addr,
shape=new_shape,
dtype=handle.dtype,
nbytes=handle.nbytes,
data=handle.data,
space=handle.space,
pinned=handle.pinned,
)
def _partial_attention(
Q: TensorHandle,
K: TensorHandle,
V: TensorHandle,
S_q: int,
S_kv_per_rank: int,
h_q: int,
d_head: int,
tl,
) -> tuple[TensorHandle, TensorHandle, TensorHandle]:
"""One pass of partial attention against (K, V).
Emits 1 GEMM(Q·K^T) + softmax + max + sub + exp + sum + 1 GEMM(P·V).
Returns the running-statistics triplet ``(m, ℓ, O_partial)`` for the
online-softmax mlo merge.
"""
K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank))
V_2d = _view(V, (S_kv_per_rank, h_q * d_head))
scores = tl.dot(Q, K_2d_T)
m = tl.max(scores, axis=-1)
P = tl.softmax(scores, axis=-1)
scores_centered = scores - m
exp_scores = tl.exp(scores_centered)
ell = tl.sum(exp_scores, axis=-1)
O_partial = tl.dot(P, V_2d)
return m, ell, O_partial
def attention_mesh_kv_kernel(
q_ptr: int,
k_ptr: int,
v_ptr: int,
o_ptr: int,
S_q: int,
S_kv_per_rank: int,
h_q: int,
h_kv: int,
d_head: int,
n_ranks: int,
rank_axis: int = 0,
*,
tl,
) -> None:
"""Mesh-native bidirectional Ring-K/V attention — see module docstring.
``rank_axis`` selects which program-id dimension carries the ring rank,
matching the GQA Llama-70B sharding study's TL/BL vs TR/BR distinction
(`llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py`):
0 — single_user_* panels (TL/BL): rank == tl.program_id(axis=0) (PE
id in cube). KV is split @ PEs **intra-cube**; ring runs over
the 8 PEs of one cube (NOC-only). At Llama-70B headline scale
this kernel launches once per cube; 64 such cubes run in
parallel for one user. Each cube's PE-level ring is independent.
1 — multi_user_* panels (TR/BR): rank == tl.program_id(axis=1)
(cube id). KV is split @ cubes **inter-cube**; ring runs over
the cubes of one KV-group. The kernel gates ``pe_id != 0`` to
return early — a v1 simplification: at headline scale (B=8) the
study's "Batch on batch" pattern would have all 8 PEs each handle
one user's batch element instead of staying silent. Adding the
per-cube batch dimension is sub-cycle 4c headline work.
"""
# For multi_user (rank_axis=1) only PE 0 in each cube runs the ring.
if rank_axis != 0 and tl.program_id(axis=0) != 0:
return
rank = tl.program_id(axis=rank_axis)
has_E = rank < n_ranks - 1
has_W = rank > 0
# Q stays put on this rank — loaded once, used in every partial attention.
Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16")
# Local K, V chunk.
K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
# Step 0 (local): partial attention against own K, V — initializes the
# running triplet (m, ℓ, o).
m, ell, o = _partial_attention(
Q, K, V, S_q, S_kv_per_rank, h_q, d_head, tl,
)
# Seed bidirectional waves with own chunk (step-1 send).
to_send_east_K: TensorHandle | None = K
to_send_east_V: TensorHandle | None = V
to_send_west_K: TensorHandle | None = K
to_send_west_V: TensorHandle | None = V
# Bidirectional fan-out: n_ranks - 1 steps. By step k, the wave from
# rank i has reached rank (i ± k). After n_ranks - 1 steps, every rank
# has merged every other rank's chunk exactly once (ADR-0059 D3).
for step in range(1, n_ranks):
# Send the eastbound wave we currently hold (own at step 1; forwarded
# at later steps). ``None`` means we have no wave to forward this
# direction this step (edge rank, or the wave already passed by).
if has_E and to_send_east_K is not None:
tl.send(dir="E", src=to_send_east_K)
tl.send(dir="E", src=to_send_east_V)
if has_W and to_send_west_K is not None:
tl.send(dir="W", src=to_send_west_K)
tl.send(dir="W", src=to_send_west_V)
# Receive eastbound wave from W (carries chunk c_{rank - step}).
K_from_W: TensorHandle | None = None
V_from_W: TensorHandle | None = None
if has_W and (rank - step) >= 0:
K_from_W = tl.recv(
dir="W", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16",
)
V_from_W = tl.recv(
dir="W", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16",
)
m_new, ell_new, o_new = _partial_attention(
Q, K_from_W, V_from_W, S_q, S_kv_per_rank, h_q, d_head, tl,
)
m_combined = tl.maximum(m, m_new)
scale_old = tl.exp(m - m_combined)
scale_new = tl.exp(m_new - m_combined)
ell = ell * scale_old + ell_new * scale_new
o = o * scale_old + o_new * scale_new
m = m_combined
# Receive westbound wave from E (carries chunk c_{rank + step}).
K_from_E: TensorHandle | None = None
V_from_E: TensorHandle | None = None
if has_E and (rank + step) < n_ranks:
K_from_E = tl.recv(
dir="E", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16",
)
V_from_E = tl.recv(
dir="E", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16",
)
m_new, ell_new, o_new = _partial_attention(
Q, K_from_E, V_from_E, S_q, S_kv_per_rank, h_q, d_head, tl,
)
m_combined = tl.maximum(m, m_new)
scale_old = tl.exp(m - m_combined)
scale_new = tl.exp(m_new - m_combined)
ell = ell * scale_old + ell_new * scale_new
o = o * scale_old + o_new * scale_new
m = m_combined
# Forward what we received for next step. ``None`` propagates: if no
# chunk arrived this step (out-of-bounds wave origin), there is
# nothing to forward next step in that direction.
to_send_east_K = K_from_W
to_send_east_V = V_from_W
to_send_west_K = K_from_E
to_send_west_V = V_from_E
# Final normalize: O := o / ℓ.
O_final = o / ell
tl.store(o_ptr, O_final)
src/kernbench/benches/_attention_mesh_mlo.py
-167
View File
@@ -1,167 +0,0 @@
"""Mesh-native bidirectional AllReduce-mlo attention — decode (ADR-0059 Proposed).
Every rank holds the full Q (replicated, small at ``S_q=1``) and 1/n_ranks
of KV (sequence-sharded). Each rank computes its partial attention
against own KV in ONE shot, then runs a bidirectional fan-out of the
``(m, ℓ, o)`` triplet: the triplet originating at rank i reaches rank j at
step ``|i - j|``. Every rank merges every other rank's triplet exactly
once over ``n_ranks - 1`` steps, ending with the final answer replicated
on every rank.
Supersedes ADR-0056's closed-ring ``_attention_allreduce_mlo.py``. Both
modules stay on disk during the transition; this one runs on the
hardware's actual open-mesh wiring (no closed-ring SFR install required).
Imported by ``milestone_gqa_llama70b`` (after the bench's Phase 2 switches
its imports) and invoked through ``torch.launch(...)`` — not through
``dist.all_reduce(...)``. See ADR-0056 Context for why this kernel is not
backend-dispatched via ADR-0050's algorithm-module contract.
"""
from __future__ import annotations
from kernbench.common.pe_commands import TensorHandle
def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle:
"""Reshape — metadata only, no command emitted (cf. ``tl.trans``)."""
return TensorHandle(
id=handle.id,
addr=handle.addr,
shape=new_shape,
dtype=handle.dtype,
nbytes=handle.nbytes,
data=handle.data,
space=handle.space,
pinned=handle.pinned,
)
def attention_mesh_mlo_kernel(
q_ptr: int,
k_ptr: int,
v_ptr: int,
o_ptr: int,
S_q: int,
S_kv_per_rank: int,
h_q: int,
h_kv: int,
d_head: int,
n_ranks: int,
rank_axis: int = 0,
*,
tl,
) -> None:
"""Mesh-native bidirectional AllReduce-mlo — see module docstring.
``rank_axis`` selects which program-id dimension carries the ring rank,
matching the GQA Llama-70B sharding study's TL/BL vs TR/BR distinction
(`llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py`):
0 — single_user_* panels (TL/BL): rank == tl.program_id(axis=0) (PE
id in cube). KV is split @ PEs **intra-cube**; ring runs over
the 8 PEs of one cube (NOC-only). At Llama-70B headline scale
this kernel launches once per cube; 64 such cubes run in
parallel for one user (1 Q-head per cube × 8 cubes per KV-group
× 8 KV-groups). The PE-level ring inside each cube is
independent of the others.
1 — multi_user_* panels (TR/BR): rank == tl.program_id(axis=1)
(cube id). KV is split @ cubes **inter-cube**; ring runs over
the cubes of one KV-group. The kernel gates ``pe_id != 0`` to
return early — a v1 simplification: at headline scale (B=8) the
study's "Batch on batch" pattern would have all 8 PEs each handle
one user's batch element instead of staying silent. Validation
shipped with B=1 to focus on the inter-cube ring's correctness;
adding the per-cube batch dimension is sub-cycle 4c headline work.
"""
# For multi_user (rank_axis=1) only PE 0 in each cube runs the ring.
if rank_axis != 0 and tl.program_id(axis=0) != 0:
return
rank = tl.program_id(axis=rank_axis)
has_E = rank < n_ranks - 1
has_W = rank > 0
# Q is replicated on every rank — loaded once.
Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16")
# Local KV chunk. KV is sequence-sharded and stays put on this rank for
# the entire fan-out — distinguishing decode from prefill (ADR-0059 D3)
# where KV circulates.
K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
# ── One-shot local partial attention ──────────────────────────
K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank))
V_2d = _view(V, (S_kv_per_rank, h_q * d_head))
scores = tl.dot(Q, K_2d_T)
m = tl.max(scores, axis=-1)
P = tl.softmax(scores, axis=-1)
scores_centered = scores - m
exp_scores = tl.exp(scores_centered)
ell = tl.sum(exp_scores, axis=-1)
o = tl.dot(P, V_2d)
# Seed bidirectional waves with own triplet (step-1 send).
to_send_east_m: TensorHandle | None = m
to_send_east_ell: TensorHandle | None = ell
to_send_east_o: TensorHandle | None = o
to_send_west_m: TensorHandle | None = m
to_send_west_ell: TensorHandle | None = ell
to_send_west_o: TensorHandle | None = o
# Bidirectional fan-out of (m, ℓ, o) triplets — n_ranks - 1 steps.
for step in range(1, n_ranks):
# Send eastbound triplet (own at step 1; forwarded at later steps).
if has_E and to_send_east_m is not None:
tl.send(dir="E", src=to_send_east_m)
tl.send(dir="E", src=to_send_east_ell)
tl.send(dir="E", src=to_send_east_o)
# Send westbound triplet.
if has_W and to_send_west_m is not None:
tl.send(dir="W", src=to_send_west_m)
tl.send(dir="W", src=to_send_west_ell)
tl.send(dir="W", src=to_send_west_o)
# Receive eastbound triplet from W (originated at rank - step).
m_from_W: TensorHandle | None = None
ell_from_W: TensorHandle | None = None
o_from_W: TensorHandle | None = None
if has_W and (rank - step) >= 0:
m_from_W = tl.recv(dir="W", shape=m.shape, dtype="f16")
ell_from_W = tl.recv(dir="W", shape=ell.shape, dtype="f16")
o_from_W = tl.recv(dir="W", shape=o.shape, dtype="f16")
m_combined = tl.maximum(m, m_from_W)
scale_old = tl.exp(m - m_combined)
scale_new = tl.exp(m_from_W - m_combined)
ell = ell * scale_old + ell_from_W * scale_new
o = o * scale_old + o_from_W * scale_new
m = m_combined
# Receive westbound triplet from E (originated at rank + step).
m_from_E: TensorHandle | None = None
ell_from_E: TensorHandle | None = None
o_from_E: TensorHandle | None = None
if has_E and (rank + step) < n_ranks:
m_from_E = tl.recv(dir="E", shape=m.shape, dtype="f16")
ell_from_E = tl.recv(dir="E", shape=ell.shape, dtype="f16")
o_from_E = tl.recv(dir="E", shape=o.shape, dtype="f16")
m_combined = tl.maximum(m, m_from_E)
scale_old = tl.exp(m - m_combined)
scale_new = tl.exp(m_from_E - m_combined)
ell = ell * scale_old + ell_from_E * scale_new
o = o * scale_old + o_from_E * scale_new
m = m_combined
# Forward the original received triplet (not the merged running state)
# so neighbors get the original wave. ``None`` propagates if nothing
# arrived this step.
to_send_east_m = m_from_W
to_send_east_ell = ell_from_W
to_send_east_o = o_from_W
to_send_west_m = m_from_E
to_send_west_ell = ell_from_E
to_send_west_o = o_from_E
# Final normalize: O := o / ℓ.
O_final = o / ell
tl.store(o_ptr, O_final)
src/kernbench/benches/_attention_mesh_mlo_2d.py
-217
View File
@@ -1,217 +0,0 @@
"""Mesh-native 2D row-then-col AllReduce-mlo attention — decode (ADR-0059 extension).
Each cube holds the full Q (replicated) and 1/(mesh_rows * mesh_cols) of
KV (sequence-sharded across the 2D cube sub-mesh). The kernel decomposes
the AllReduce-mlo into a two-stage reduction:
Stage 1 — Row reduce (E/W edges, ``mesh_cols - 1`` steps)
Bidirectional ring within each row. After this stage every cube in
row ``r`` holds the partial ``(m, ℓ, o)`` over the ``mesh_cols`` KV
chunks in row ``r``.
Stage 2 — Col reduce (N/S edges, ``mesh_rows - 1`` steps)
Bidirectional ring within each column. After this stage every cube
holds the partial over all ``mesh_rows × mesh_cols`` KV chunks —
the AllReduce result.
The online-softmax mlo merge is associative, so row-then-col partitioning
of the reduction is mathematically equivalent to a 1D ring AllReduce-mlo
over all ``mesh_rows × mesh_cols`` cubes. The 2D form takes
``(mesh_cols - 1) + (mesh_rows - 1)`` steps instead of
``mesh_rows × mesh_cols - 1`` (e.g. 4 vs 7 at 2×4; 6 vs 15 at 4×4).
Designed to run on hardware wired by
``configure_sfr_intercube_multisip``, which installs both E/W and N/S
intra-SIP cube-mesh edges (``sfr_config.py:135-143``). The 1D
``_attention_mesh_mlo.py`` remains for the single_user PE-ring case;
this 2D variant supersedes it for multi_user_decode where the per-KV-group
cube count crosses a row boundary in the 4×4 cube mesh.
``mesh_rows = 1`` is supported as a degenerate row-only case so the
validation config (``_N_RANKS_MULTI_USER = 4`` → ``(1, 4)``) reduces to
the 1D ring's step count without behavioral change.
"""
from __future__ import annotations
from kernbench.common.pe_commands import TensorHandle
def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle:
"""Reshape — metadata only, no command emitted (cf. ``tl.trans``)."""
return TensorHandle(
id=handle.id,
addr=handle.addr,
shape=new_shape,
dtype=handle.dtype,
nbytes=handle.nbytes,
data=handle.data,
space=handle.space,
pinned=handle.pinned,
)
def _bidir_allreduce_mlo(
m: TensorHandle,
ell: TensorHandle,
o: TensorHandle,
rank: int,
n_ranks: int,
dir_pos: str,
dir_neg: str,
*,
tl,
) -> tuple[TensorHandle, TensorHandle, TensorHandle]:
"""One bidirectional AllReduce-mlo ring along ``(dir_pos, dir_neg)``.
Mirrors the 1D ``_attention_mesh_mlo.py`` algorithm but parameterized
on direction labels so the 2D kernel can call it once with ``("E", "W")``
for the row reduce and once with ``("S", "N")`` for the col reduce.
Forwards the received triplets in subsequent steps so chunk ``c_i``
reaches rank ``j`` at step ``|i - j|``.
Returns the running ``(m, ℓ, o)`` after ``n_ranks - 1`` steps. Degenerate
cases (``n_ranks <= 1``) are no-ops — the for-loop body simply does not
execute.
"""
has_pos = rank < n_ranks - 1
has_neg = rank > 0
to_send_pos_m: TensorHandle | None = m
to_send_pos_ell: TensorHandle | None = ell
to_send_pos_o: TensorHandle | None = o
to_send_neg_m: TensorHandle | None = m
to_send_neg_ell: TensorHandle | None = ell
to_send_neg_o: TensorHandle | None = o
for step in range(1, n_ranks):
if has_pos and to_send_pos_m is not None:
tl.send(dir=dir_pos, src=to_send_pos_m)
tl.send(dir=dir_pos, src=to_send_pos_ell)
tl.send(dir=dir_pos, src=to_send_pos_o)
if has_neg and to_send_neg_m is not None:
tl.send(dir=dir_neg, src=to_send_neg_m)
tl.send(dir=dir_neg, src=to_send_neg_ell)
tl.send(dir=dir_neg, src=to_send_neg_o)
m_from_neg: TensorHandle | None = None
ell_from_neg: TensorHandle | None = None
o_from_neg: TensorHandle | None = None
if has_neg and (rank - step) >= 0:
m_from_neg = tl.recv(dir=dir_neg, shape=m.shape, dtype="f16")
ell_from_neg = tl.recv(dir=dir_neg, shape=ell.shape, dtype="f16")
o_from_neg = tl.recv(dir=dir_neg, shape=o.shape, dtype="f16")
m_combined = tl.maximum(m, m_from_neg)
scale_old = tl.exp(m - m_combined)
scale_new = tl.exp(m_from_neg - m_combined)
ell = ell * scale_old + ell_from_neg * scale_new
o = o * scale_old + o_from_neg * scale_new
m = m_combined
m_from_pos: TensorHandle | None = None
ell_from_pos: TensorHandle | None = None
o_from_pos: TensorHandle | None = None
if has_pos and (rank + step) < n_ranks:
m_from_pos = tl.recv(dir=dir_pos, shape=m.shape, dtype="f16")
ell_from_pos = tl.recv(dir=dir_pos, shape=ell.shape, dtype="f16")
o_from_pos = tl.recv(dir=dir_pos, shape=o.shape, dtype="f16")
m_combined = tl.maximum(m, m_from_pos)
scale_old = tl.exp(m - m_combined)
scale_new = tl.exp(m_from_pos - m_combined)
ell = ell * scale_old + ell_from_pos * scale_new
o = o * scale_old + o_from_pos * scale_new
m = m_combined
to_send_pos_m = m_from_neg
to_send_pos_ell = ell_from_neg
to_send_pos_o = o_from_neg
to_send_neg_m = m_from_pos
to_send_neg_ell = ell_from_pos
to_send_neg_o = o_from_pos
return m, ell, o
def attention_mesh_mlo_2d_kernel(
q_ptr: int,
k_ptr: int,
v_ptr: int,
o_ptr: int,
S_q: int,
S_kv_per_rank: int,
h_q: int,
h_kv: int,
d_head: int,
mesh_rows: int,
mesh_cols: int,
rank_axis: int = 0,
cube_start: int = 0,
*,
tl,
) -> None:
"""2D row-then-col AllReduce-mlo decode kernel — see module docstring.
``rank_axis`` selects which program-id dimension carries the cube
rank (matches the 1D kernel convention):
0 — single_user_* (TL/BL): rank == tl.program_id(axis=0) (PE id).
Not used at headline scale — single_user uses the 1D intra-cube
PE ring (``_attention_mesh_mlo``). Kept here so the signature
mirrors the 1D kernel.
1 — multi_user_* (TR/BR): rank == tl.program_id(axis=1) (cube id).
KV is split @ cubes inter-cube; the ring runs over the
``mesh_rows × mesh_cols`` cubes of one KV-group. The kernel
gates ``pe_id != 0`` to return early — same v1 simplification
as ``_attention_mesh_mlo`` (validation B=1).
``cube_start`` matches the value passed to ``DPPolicy.cube_start`` for
the launch's tensor placement. kernbench's ``tl.program_id(axis=1)``
returns the physical cube id (ADR-0022), so when the launch is
offset within the SIP (e.g. cube_start=8 placing the second 2×4
KV-group on cubes 8..15), the kernel must subtract ``cube_start``
to recover the launch-local rank for ring arithmetic. Default 0
preserves the cube_start=0 launches unchanged.
"""
# For multi_user (rank_axis=1) only PE 0 in each cube runs the ring.
if rank_axis != 0 and tl.program_id(axis=0) != 0:
return
rank = tl.program_id(axis=rank_axis)
if rank_axis != 0:
rank = rank - cube_start
my_row = rank // mesh_cols
my_col = rank % mesh_cols
# Q is replicated on every cube — loaded once.
Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16")
# Local KV chunk (sequence-sharded across the 2D sub-mesh).
K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
# ── One-shot local partial attention ──────────────────────────
K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank))
V_2d = _view(V, (S_kv_per_rank, h_q * d_head))
scores = tl.dot(Q, K_2d_T)
m = tl.max(scores, axis=-1)
P = tl.softmax(scores, axis=-1)
scores_centered = scores - m
exp_scores = tl.exp(scores_centered)
ell = tl.sum(exp_scores, axis=-1)
o = tl.dot(P, V_2d)
# ── Stage 1: row AllReduce (E/W, mesh_cols - 1 steps) ─────────
m, ell, o = _bidir_allreduce_mlo(
m, ell, o, my_col, mesh_cols, "E", "W", tl=tl,
)
# ── Stage 2: col AllReduce (N/S, mesh_rows - 1 steps) ─────────
# ``dir_pos="S"`` matches the SFR convention: ``S`` goes to higher
# row (configure_sfr_intercube_multisip:140).
m, ell, o = _bidir_allreduce_mlo(
m, ell, o, my_row, mesh_rows, "S", "N", tl=tl,
)
# Final normalize: O := o / ℓ.
O_final = o / ell
tl.store(o_ptr, O_final)
src/kernbench/benches/_gqa_decode.py
+180
View File
@@ -0,0 +1,180 @@
"""GQA fused-attention decode kernel — P1a + P2a + P2b (2-level SP).
Lineage (DDD-0060 §7 phase plan):
P1a : real GQA via M-fold using ``tl.dot``; one-shot per rank.
P2a : intra-CUBE PE-level chain reduce-to-root (single-CUBE SP).
P2b : adds inter-CUBE chain reduce-to-root (multi-CUBE SP);
switches to the canonical full SFR install
``configure_sfr_intercube_multisip`` with disjoint namespaces
(``intra_*`` for PE, ``N/S/E/W`` for CUBE, ``global_*`` for SIP).
P3b : tile S_kv sweep + ``tl.scratch_scope`` (deferred).
Later : P1b composite swap; P4 lazy load; P5 opt3 pipelining; P6 prefill.
P2b SFR + topology assumptions:
- SFR install: ``configure_sfr_intercube_multisip`` is required when
P > 1 or C > 1 (provides ``intra_*`` and ``E/W/N/S`` namespaces).
- Intra-CUBE PE layout: logical 2×4 grid (no wrap):
Row 0: PE 0, 1, 2, 3
Row 1: PE 4, 5, 6, 7
- Inter-CUBE layout: CUBEs of one CUBE Group are laid out as a 1D row
(single row of C CUBEs, no wrap). Multi-row CUBE Group placement
(e.g. 2×4) is future work — head_of_group / cube_start dispatch is
a P7 concern (DDD-0060 §4.1).
Reduce strategy — chain reduce-to-root at PE 0 of CUBE 0:
Level-2 (intra-CUBE, row-then-col chain on 2×4 grid):
Row chain along ``intra_W``: rightmost-col PEs send leftward;
leftmost-col PE of each row holds its row's partial.
Col bridge along ``intra_N``: PE 4 (col-0, row-1) sends to PE 0
(col-0, row-0). Only relevant when P > 4.
Result: PE 0 of each CUBE holds the CUBE's partial.
Level-1 (inter-CUBE, only PE 0 of each CUBE participates):
Chain along ``W``: rightmost CUBE sends leftward; CUBE 0 of the
CUBE Group ends with the final answer.
Final store: PE 0 of CUBE 0 normalises (O / ℓ) and writes.
Chain step counts (per ADR-0060 §A.2 root-only output, §4 chain
deviation noted): for (C, P)=(2, 8), 7 intra-cube × 2 cubes + (C-1)
inter-cube = 14 + 1 = 15 chain steps; each step ships 3 handles
(m, ℓ, O) ⇒ 45 ``ipcq_copy`` total.
Three deliberate deviations from ADR-0060, addressed in later phases:
1. GEMMs use ``tl.dot``, not ``tl.composite`` (P1b).
2. ``softmax_scale`` omitted — needs composite epilogue mechanism (P1b/P5).
3. K loaded as ``[d, S_local]`` via byte-conserving reshape of the
deployed ``[S_local, h_kv·d]`` slice (ADR-0060 §3 / §B item 2,
reshape-not-transpose caveat; correct for zero / symmetric inputs).
Chain-vs-tree deviation from DDD-0060 §7 P2 gate (``⌈log₂ P⌉``):
P2b uses linear chain reduce-to-root (P-1 + C-1 hops). True tree on
the 2×4 PE grid requires a different SFR install — separate ADR.
Architectural intent (root-only output replacing baseline's
bidirectional fan-out) is preserved.
"""
from __future__ import annotations
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
"""Online-softmax merge (ADR-0060 §4 / _attention_mesh_mlo baseline)."""
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_decode_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,
*,
tl,
) -> None:
"""GQA decode with M-fold + 2-level chain 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 CUBE 0 stores.
"""
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 one-shot partial attention (M-fold on the rank's slice) ──
Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
K_T = tl.load(k_ptr, shape=(d_head, S_local), dtype="f16")
V = tl.load(v_ptr, shape=(S_local, 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)
# ── Level-2: intra-CUBE row-then-col chain reduce-to-(PE 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
# Row chain (along intra_W within each row, gathering leftward).
# Each merge step's intermediates are wrapped in tl.scratch_scope and
# the new running (m, ℓ, O) is persisted back to the outside-scope
# (persistent) m_local/l_local/O_local via tl.copy_to (ADR-0063 §D3/D3.1).
if pe_cols_used > 1:
if pe_col < pe_cols_used - 1: # not rightmost: receive E
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: # not leftmost: send W
tl.send(dir="intra_W", src=m_local)
tl.send(dir="intra_W", src=l_local)
tl.send(dir="intra_W", src=O_local)
# Col bridge (intra_N from row 1 col 0 → row 0 col 0). Only at col 0.
if pe_col == 0 and pe_rows_used > 1:
if pe_row < pe_rows_used - 1: # row 0 receives from S
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: # row >0 sends to N
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 chain reduce (only PE 0 of each CUBE) ──
if pe_id == 0 and C > 1:
if cube_id < C - 1: # not rightmost CUBE: recv E
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: # non-root CUBE: send W
tl.send(dir="W", src=m_local)
tl.send(dir="W", src=l_local)
tl.send(dir="W", src=O_local)
# ── Final normalise + store (only at PE 0 of CUBE 0) ──
if pe_id == 0 and cube_id == 0:
O_final = O_local / l_local
tl.store(o_ptr, O_final)
src/kernbench/benches/_gqa_decode_short.py
+145
View File
@@ -0,0 +1,145 @@
"""GQA fused-attention SHORT-CONTEXT decode kernel (ADR-0060 §B.split.2).
Short context (S_kv < 256K, per ADR-0060 §B.split.1): each CUBE owns
``kv_per_cube`` whole KV heads, no S_kv sharding across CUBEs, no
inter-CUBE reduce. PE-SP within each CUBE: the P PEs split into
``kv_per_cube`` groups, each group does PE-SP across (P/kv_per_cube)
PEs for ONE owned head.
Layout (after design iteration during Phase D — see ADR-0060 §B.split.2):
- K, V: shape ``(h_kv·S_kv, d_head)`` head-stacked, with the bench
deploying ``dp = (cube=row_wise, pe=row_wise)`` so each PE's chunk
is exactly ``(S_local, d_head)`` contiguous at its own addressable
shard. The kernel just loads at its ``k_ptr`` / ``v_ptr`` — no
offset arithmetic needed.
- Q: replicated ``(T_q, h_q·d_head)``; the kernel reshapes
byte-conservingly to ``(h_q·T_q, d_head)`` and operates on the
full stack. Other heads' rows are computed too (semantic noise);
with zero/symmetric inputs the math is unchanged. A proper
per-head Q slice would require runtime support for partial reads
of stored tensors (deferred).
- O: replicated; each group root writes the full byte-conserving
``(h_q·T_q, d_head)`` result. Multiple roots within a CUBE write
to disjoint PE-local addresses (no overwrite collision).
Group layout on the 2×4 PE grid:
kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
kv_per_cube=2, group=4 PEs (one row): row chain only.
kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
kv_per_cube=8, group=1 PE: no chain — direct write.
Chain reduce within group via the existing ``intra_E/W/N/S`` SFR
namespace (configure_sfr_intercube_multisip). After chain reduce, the
group's root PE (pe_in_group == 0) writes its working state to HBM.
Deviations from ADR-0060 (deliberate, documented):
1. GEMMs use ``tl.dot``, not ``tl.composite``.
2. ``softmax_scale`` omitted.
3. K loaded as ``[d, S_local]`` via byte-conserving reshape
(reshape-not-transpose caveat — correct for zero / symmetric inputs).
4. Q byte-conserving reshape: kernel computes attention for ALL Q
rows against the group's owned K head; only the rows for my head
are semantically meaningful. Correct for zero / symmetric inputs.
"""
from __future__ import annotations
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
"""Online-softmax merge — identical to long kernel."""
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_decode_short_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,
kv_per_cube: int,
*,
tl,
) -> None:
"""Short-context GQA decode with PE-parallel heads + intra-group PE-SP."""
group_size = P // kv_per_cube # PEs per head group
pe_id = tl.program_id(axis=0)
pe_in_group = pe_id % group_size
# PE-SP within group: shard S_kv across group_size PEs
S_local = S_kv // group_size
# ── Loads (DP layout already places each PE at its own shard) ──
# Q replicated → byte-conserving reshape to (h_q·T_q, d_head).
Q = tl.load(q_ptr,
shape=(h_q * T_q, d_head), dtype="f16")
# K, V row_wise per (cube, pe) → each PE has (S_local, d_head) at k_ptr.
K_T = tl.load(k_ptr,
shape=(d_head, S_local), dtype="f16")
V = tl.load(v_ptr,
shape=(S_local, d_head), dtype="f16")
# ── Local one-shot partial attention ──
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)
# ── Within-group chain reduce-to-root (Level-2 only) ──
group_cols = min(4, group_size)
group_rows = (group_size + group_cols - 1) // group_cols
pe_col_in_group = pe_in_group % group_cols
pe_row_in_group = pe_in_group // group_cols
# Row chain along intra_W (within group's row).
if group_cols > 1:
if pe_col_in_group < group_cols - 1: # receive E (in-group)
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_in_group > 0: # send W (in-group)
tl.send(dir="intra_W", src=m_local)
tl.send(dir="intra_W", src=l_local)
tl.send(dir="intra_W", src=O_local)
# Col bridge along intra_N (only if group spans 2 grid rows).
if pe_col_in_group == 0 and group_rows > 1:
if pe_row_in_group < group_rows - 1: # receive S (in-group)
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_in_group > 0: # send N (in-group)
tl.send(dir="intra_N", src=m_local)
tl.send(dir="intra_N", src=l_local)
tl.send(dir="intra_N", src=O_local)
# ── Group root writes its owned head's output ──
if pe_in_group == 0:
O_final = O_local / l_local
tl.store(o_ptr, O_final)
src/kernbench/benches/_gqa_prefill.py
+111
View File
@@ -0,0 +1,111 @@
"""GQA fused-attention prefill kernel — P6a + P6b (head-parallel + Ring KV).
Lineage (DDD-0060 §7 phase plan):
P6a : head-parallel structure (one Q head per CUBE), C=1 baseline.
P6b : add Ring KV rotation across C CUBEs (ADR-0060 §5.5). Each
CUBE rotates its KV block to its W neighbour and receives
from E; over C-1 steps every CUBE sees every block. Online-
softmax merge folds each step into running (m, ℓ, O). No
reduce — each CUBE writes its own head's output.
Requires ``configure_sfr_intercube_ring(ring_size=C)`` SFR.
P6c (later): tile T_q across the P PEs inside a CUBE for intra-CUBE
parallelism (ADR-0060 §B item 3).
Deviations from ADR-0060 §5.5 — deferred to later phases:
1. GEMMs use ``tl.dot`` not ``tl.composite`` (parallel to P1a; lifts
when P1b decides the composite-output-handle question).
2. ``softmax_scale`` omitted — same composite-epilogue deferral.
3. K loaded as ``[d, S_local]`` via byte-conserving reshape of
``[S_local, d]`` (ADR-0060 §3 / §B item 2 reshape-not-transpose
caveat; correct for zero / symmetric inputs).
4. No causal masking / step-skip — future P6c.
5. Blocking ``tl.recv`` (not ``recv_async``) — overlap via lazy
``tl.load`` lands in P4.
"""
from __future__ import annotations
def gqa_prefill_kernel(
q_ptr: int,
k_ptr: int,
v_ptr: int,
o_ptr: int,
T_q: int,
S_kv: int,
d_head: int,
C: int,
*,
tl,
) -> None:
"""Head-parallel prefill attention with Ring KV (C>1) — ADR-0060 §5.5.
Tensor layout consumed by this kernel:
Q : (T_q, d_head) one head per CUBE; replicated.
K : (S_kv, d_head) sharded cube_row_wise → each CUBE owns
(S_kv/C, d_head); kernel loads as (d_head, S_local) via
byte-conserving reshape (reshape-not-transpose caveat).
V : (S_kv, d_head) sharded cube_row_wise → each CUBE owns
(S_local, d_head).
O : (T_q * C, d_head) sharded cube_row_wise → each CUBE
writes its own (T_q, d_head) slice. NO reduce.
Algorithm: each CUBE computes a local partial against its current
KV block, then over C-1 ring steps the K and V blocks rotate W
while online-softmax merges each step into running (m, ℓ, O).
"""
pe_id = tl.program_id(axis=0)
# Head-parallel: only PE 0 of each CUBE participates. P6c (future)
# will tile T_q across the 8 PEs for intra-CUBE parallelism.
if pe_id != 0:
return
S_local = S_kv // C
Q = tl.load(q_ptr, shape=(T_q, d_head), dtype="f16")
Kc = tl.load(k_ptr, shape=(d_head, S_local), dtype="f16")
Vc = tl.load(v_ptr, shape=(S_local, d_head), dtype="f16")
# ── Step 0: initial partial against own KV block — establishes the
# persistent (m, ℓ, O) arena. Intermediates (scores, exp_scores) stay
# allocated; ring steps below recycle per-step intermediates inside
# tl.scratch_scope to keep peak scratch O(one step) (ADR-0063 §D3).
scores = tl.dot(Q, Kc)
m = tl.max(scores, axis=-1)
exp_scores = tl.exp(scores - m)
l = tl.sum(exp_scores, axis=-1)
O = tl.dot(exp_scores, Vc)
# ── Steps 1..C-1: Ring KV rotation + online-softmax merge ──
# Per-step intermediates wrapped in tl.scratch_scope; the merged
# running (m, ℓ, O) is persisted to the outside-scope handles via
# tl.copy_to (ADR-0063 §D3.1) so its bytes survive __exit__.
for _ in range(1, C):
tl.send(dir="W", src=Kc)
tl.send(dir="W", src=Vc)
Kc = tl.recv(dir="E", shape=(d_head, S_local), dtype="f16")
Vc = tl.recv(dir="E", shape=(S_local, d_head), dtype="f16")
with tl.scratch_scope():
# Partial on rotated KV
scores = tl.dot(Q, Kc)
m_step = tl.max(scores, axis=-1)
exp_scores = tl.exp(scores - m_step)
l_step = tl.sum(exp_scores, axis=-1)
O_step = tl.dot(exp_scores, Vc)
# Online-softmax merge into new running (m, ℓ, O)
m_new = tl.maximum(m, m_step)
scale_old = tl.exp(m - m_new)
scale_step = tl.exp(m_step - m_new)
l_new = l * scale_old + l_step * scale_step
O_new = O * scale_old + O_step * scale_step
# Persist new running state back to the outside-scope arena.
tl.copy_to(m, m_new)
tl.copy_to(l, l_new)
tl.copy_to(O, O_new)
# __exit__: scoped intermediates gone; persistent m, l, O carry
# the new running state into the next ring iteration.
# Final normalise + store — each CUBE writes its own head's rows.
O_final = O / l
tl.store(o_ptr, O_final)
src/kernbench/benches/_gqa_prefill_short.py
+114
View File
@@ -0,0 +1,114 @@
"""GQA fused-attention SHORT-CONTEXT prefill kernel (ADR-0060 §B.split.2).
Prefill analogue of ``_gqa_decode_short.py``. Same layout decisions:
- K, V head-stacked (h_kv·S_kv, d_head) with row_wise DP so each PE
has a contiguous (S_local, d_head) shard at its own address.
- Q replicated; kernel uses byte-conserving reshape.
- O replicated; group root writes the full byte-conserving result.
No Ring KV (each owned head fully resident at its CUBE). Test
``test_short_prefill_no_ring_KV_traffic`` asserts no inter-CUBE IPCQ.
The only structural difference from short decode:
- ``T_q`` may be > 1 (prefill processes multiple query tokens).
- Q is shaped (T_q, h_kv·d_head) — one Q head per KV head (no GQA
M-fold here; head-parallel prefill in ADR-0060 §5.5 is 1:1).
"""
from __future__ import annotations
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
"""Online-softmax merge — identical to short decode kernel."""
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_prefill_short_kernel(
q_ptr: int,
k_ptr: int,
v_ptr: int,
o_ptr: int,
T_q: int,
S_kv: int,
h_kv: int,
d_head: int,
C: int,
P: int,
kv_per_cube: int,
*,
tl,
) -> None:
"""Short-context prefill with PE-parallel heads + intra-group PE-SP.
NO Ring KV (each CUBE owns its KV heads fully).
"""
group_size = P // kv_per_cube
pe_id = tl.program_id(axis=0)
pe_in_group = pe_id % group_size
S_local = S_kv // group_size
# Q: replicated (T_q, h_kv·d_head) → byte-conserving reshape to
# (h_kv·T_q, d_head) — one Q head per KV head, stacked.
Q = tl.load(q_ptr,
shape=(h_kv * T_q, d_head), dtype="f16")
K_T = tl.load(k_ptr,
shape=(d_head, S_local), dtype="f16")
V = tl.load(v_ptr,
shape=(S_local, 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)
# Within-group chain reduce-to-root (same machinery as short decode).
group_cols = min(4, group_size)
group_rows = (group_size + group_cols - 1) // group_cols
pe_col_in_group = pe_in_group % group_cols
pe_row_in_group = pe_in_group // group_cols
if group_cols > 1:
if pe_col_in_group < group_cols - 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_in_group > 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_in_group == 0 and group_rows > 1:
if pe_row_in_group < group_rows - 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_in_group > 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)
if pe_in_group == 0:
O_final = O_local / l_local
tl.store(o_ptr, O_final)
src/kernbench/benches/milestone_gqa_headline.py
+236
View File
@@ -0,0 +1,236 @@
"""milestone-gqa-headline bench: real GQA + 2-level SP + Ring KV.
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).
Restrictions (P7 first cut):
- C ≤ 4 (single-row inter-CUBE ring SFR; multi-row deferred)
- Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
headline deferred)
- 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)
Gated by ``GQA_HEADLINE_RUN=1``.
"""
from __future__ import annotations
import json
import os
from pathlib import Path
from kernbench.benches._gqa_decode import gqa_decode_kernel
from kernbench.benches._gqa_prefill import gqa_prefill_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.policy.placement.dp import DPPolicy
_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline"
_SWEEP_JSON = _OUTPUT_DIR / "sweep.json"
# ── Panel configs ────────────────────────────────────────────────────
_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_user_prefill_gqa",
"multi_user_prefill_gqa",
"single_user_decode_gqa",
"multi_user_decode_gqa",
)
# Each entry: (kind, panel-specific params)
_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_user_decode_gqa": ("decode", {"C": 1, "P": 8, "S_kv": 64}),
"multi_user_decode_gqa": ("decode", {"C": 4, "P": 8, "S_kv": 128}),
}
def _ccl_cfg():
return resolve_algorithm_config(
load_ccl_config(), name="lrab_hierarchical_allreduce",
)
# ── Per-kind launch helpers ──────────────────────────────────────────
def _run_prefill_panel(ctx, *, panel: str, C: int, S_kv: int) -> 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)
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_PREFILL, _D_HEAD),
dtype=_DTYPE, dp=dp_q, name=f"{panel}_q")
k = ctx.zeros((S_kv, _D_HEAD),
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
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),
dtype=_DTYPE, dp=dp_o, name=f"{panel}_o")
ctx.launch(
panel, gqa_prefill_kernel,
q, k, v, o,
_T_Q_PREFILL, S_kv, _D_HEAD, C,
_auto_dim_remap=False,
)
def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> 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_DECODE, _H_Q_DECODE * _D_HEAD),
dtype=_DTYPE, dp=dp_full, name=f"{panel}_q")
k = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
v = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
ctx.launch(
panel, gqa_decode_kernel,
q, k, v, o,
_T_Q_DECODE, S_kv, _H_Q_DECODE, _H_KV_DECODE, _D_HEAD, C, P,
_auto_dim_remap=False,
)
def _make_bench_fn(panel: str):
kind, params = _PANEL_DISPATCH[panel]
def _bench_fn(ctx):
if kind == "prefill":
_run_prefill_panel(ctx, panel=panel, **params)
else:
_run_decode_panel(ctx, panel=panel, **params)
return _bench_fn
# ── Op-log summary ──────────────────────────────────────────────────
def _summarize_op_log(op_log) -> dict[str, int]:
"""Per-panel op_log counts (gemm, ipcq_copy, dma_read, dma_write)."""
gemm_count = 0
ipcq_copy_count = 0
dma_read_count = 0
dma_write_count = 0
for r in op_log:
if r.op_kind == "gemm":
gemm_count += 1
elif r.op_name == "dma_read":
dma_read_count += 1
elif r.op_name == "dma_write":
dma_write_count += 1
elif r.op_name == "ipcq_copy":
ipcq_copy_count += 1
return {
"gemm_count": gemm_count,
"ipcq_copy_count": ipcq_copy_count,
"dma_read_count": dma_read_count,
"dma_write_count": dma_write_count,
}
def _run_panel(panel: str, topology: str) -> dict:
"""Run one panel in a fresh engine; return its row dict."""
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
topo = resolve_topology(topology)
result = run_bench(
topology=topo, bench_fn=_make_bench_fn(panel),
device=resolve_device(None),
engine_factory=lambda t, d: GraphEngine(
getattr(t, "topology_obj", t), enable_data=True,
),
)
if not result.completion.ok:
raise RuntimeError(
f"milestone-gqa-headline panel {panel!r} failed: "
f"{result.completion}"
)
kind, params = _PANEL_DISPATCH[panel]
return {
"panel": panel,
"kind": kind,
**params,
"op_log_summary": _summarize_op_log(result.engine.op_log),
}
# ── Bench entry ──────────────────────────────────────────────────────
@bench(
name="milestone-gqa-headline",
description="Headline GQA milestone — real GQA h_q=8/h_kv=1 + 2-level SP (decode) + Ring KV (prefill).",
)
def run(torch) -> None:
"""Drive 4 headline panels through the new GQA kernels; write sweep.json.
Gated by GQA_HEADLINE_RUN=1.
"""
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
if not os.environ.get("GQA_HEADLINE_RUN"):
raise RuntimeError(
"milestone-gqa-headline needs GQA_HEADLINE_RUN=1."
)
topology = os.environ.get("GQA_HEADLINE_TOPOLOGY", "topology.yaml")
rows = [_run_panel(panel, topology) for panel in _PANELS]
sweep = {
"version": 1,
"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,
}
_SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
print(f" milestone-gqa-headline: {len(rows)} rows -> {_SWEEP_JSON}")
# Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out).
torch.zeros(
(1, 1), dtype="f16",
dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1),
name="milestone_gqa_headline_sentinel",
)
src/kernbench/benches/milestone_gqa_llama70b.py
-473
View File
@@ -1,473 +0,0 @@
"""milestone-gqa-llama70b bench: GQA Llama-70B 4-panel sweep (ADR-0057 v1).
Self-contained eval bench (ADR-0054). Drives the four panels of the GQA
Llama-70B sharding study through ``run_bench`` with ``enable_data=True``,
harvests op_log summaries, and writes JSON into
``benches/1H_milestone_output/gqa/sweep.json``.
v1 (sub-cycle 4a + 4c.0) covers all four panels at validation scale:
Panel name in JSON / test Study label SFR install used
─────────────────────────────────────────────────────────────────────
single_user_prefill TL configure_sfr_intracube_pe_ring
multi_user_prefill TR configure_sfr_intercube_multisip
single_user_decode BL configure_sfr_intracube_pe_ring
multi_user_decode BR configure_sfr_intercube_multisip
Per the GQA sharding study (`llm_paper_review/notes/GQA_MHA_sharding/scripts
/_gen_llama70b_1M_4cases.py`):
Single User (B=1) panels — TL prefill, BL decode:
"n_cubes: 8 (1 KV-group)", KV split @ PEs intra-cube. Each cube does
its own 8-PE ring with no cube-to-cube attention traffic. At Llama-70B
headline scale this is 64 cubes (8 KV-groups × 8 cubes/group), each
independently running the kernel for one Q-head; 1 user spans all 64.
Multi User (B=8) panels — TR prefill, BR decode:
"8 cubes / KV-group", KV split @ cubes inter-cube. The 8 cubes of a
KV-group form a ring; "Inside each cube: 8 PEs each handle 1 different
user → Batch on batch, batch = 8/cube." At headline scale 8 KV-groups
run in parallel = 64 cubes serving 8 users.
Kernels use the mesh-native variants (ADR-0059), invoked with the
``rank_axis`` kwarg (0 for single_user PE-level rings, 1 for multi_user
cube-level rings). The v1 multi_user kernel gates ``pe_id != 0`` to return,
which simplifies B=8 → B=1 — that's a validation simplification; the
per-cube "Batch on batch" parallelism is sub-cycle 4c headline work.
Validation-scale config (ADR-0057 D4) — kept small so the simulator's
1 MB per-PE TCM scratch budget is not exhausted across n_ranks ring steps:
``S_q_prefill = S_kv_per_rank = 16``, ``h_q = h_kv = 1``, ``d_head = 64``,
``n_ranks_single_user = 8`` (8 PEs of one cube), ``n_ranks_multi_user = 4``
(half a KV-group, vs the study's 8). Headline-scale dims (``S_q = 1M``,
``S_kv = 1M``, ``h_q = 8 / h_kv = 1`` GQA, ``d_head = 128``, ``B = 8``) are
also deferred.
"""
from __future__ import annotations
import json
import os
from pathlib import Path
from typing import Any
from kernbench.benches._attention_mesh_kv import attention_mesh_kv_kernel
from kernbench.benches._attention_mesh_mlo import attention_mesh_mlo_kernel
from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_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_intracube_pe_ring,
)
from kernbench.policy.placement.dp import DPPolicy
_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa"
_SWEEP_JSON = _OUTPUT_DIR / "sweep.json"
# ── Validation-scale config (ADR-0057 D4) ─────────────────────────────
_S_Q_PREFILL = 16
_S_Q_DECODE = 1
_S_KV_PER_RANK = 16
_H_Q = 1
_H_KV = 1
_D_HEAD = 64
_N_RANKS_SINGLE_USER = 8
_N_RANKS_MULTI_USER = 4
_DTYPE = "f16"
_PANELS_V1 = (
"single_user_prefill",
"multi_user_prefill",
"single_user_decode",
"multi_user_decode",
)
# Panel → (kernel, SFR install, S_q, n_ranks, rank_axis, mesh_shape)
# ``mesh_shape`` is ``None`` for 1D-ring kernels and ``(rows, cols)`` for the
# 2D row-then-col AllReduce-mlo kernel (multi_user_decode); when set, the
# launch passes ``(mesh_rows, mesh_cols)`` instead of ``n_ranks``.
_PANEL_DISPATCH: dict[
str, tuple[Any, Any, int, int, int, tuple[int, int] | None]
] = {
"single_user_prefill": (
attention_mesh_kv_kernel, configure_sfr_intracube_pe_ring,
_S_Q_PREFILL, _N_RANKS_SINGLE_USER, 0, None,
),
"multi_user_prefill": (
attention_mesh_kv_kernel, configure_sfr_intercube_multisip,
_S_Q_PREFILL, _N_RANKS_MULTI_USER, 1, None,
),
"single_user_decode": (
attention_mesh_mlo_kernel, configure_sfr_intracube_pe_ring,
_S_Q_DECODE, _N_RANKS_SINGLE_USER, 0, None,
),
# multi_user_decode uses the C2 2D AllReduce-mlo kernel. (1, 4)
# degenerates to a row-only AllReduce equivalent to the prior 1D ring
# at n_ranks=4 — no op_log_summary regression. Headline 8-cube
# KV-groups land at (2, 4).
"multi_user_decode": (
attention_mesh_mlo_2d_kernel, configure_sfr_intercube_multisip,
_S_Q_DECODE, _N_RANKS_MULTI_USER, 1, (1, _N_RANKS_MULTI_USER),
),
}
# ── Per-panel bench fn ─────────────────────────────────────────────────
def _make_bench_fn(panel: str):
kernel, sfr_install, S_q, n_ranks, rank_axis, mesh_shape = (
_PANEL_DISPATCH[panel]
)
is_multi_user = panel.startswith("multi_user_")
def _bench_fn(ctx):
sfr_install(
ctx.engine, ctx.spec,
resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"),
)
if is_multi_user:
dp_full = DPPolicy(
cube="replicate", pe="replicate",
num_cubes=n_ranks, num_pes=8,
)
dp_kv = DPPolicy(
cube="row_wise", pe="replicate",
num_cubes=n_ranks, num_pes=8,
)
else:
dp_full = DPPolicy(
cube="replicate", pe="replicate",
num_cubes=1, num_pes=n_ranks,
)
dp_kv = DPPolicy(
cube="replicate", pe="row_wise",
num_cubes=1, num_pes=n_ranks,
)
q = ctx.zeros((S_q, _H_Q * _D_HEAD),
dtype=_DTYPE, dp=dp_full, name=f"{panel}_q")
k = ctx.zeros((_S_KV_PER_RANK * n_ranks, _H_KV * _D_HEAD),
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
v = ctx.zeros((_S_KV_PER_RANK * n_ranks, _H_KV * _D_HEAD),
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
o = ctx.empty((S_q, _H_Q * _D_HEAD),
dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
# rank_axis is a positional arg; _auto_dim_remap=False keeps
# d_head=64 from colliding with the multi_user K's global M=64.
if mesh_shape is None:
ctx.launch(
f"{panel}_mesh", kernel,
q, k, v, o,
S_q, _S_KV_PER_RANK, _H_Q, _H_KV, _D_HEAD, n_ranks,
rank_axis,
_auto_dim_remap=False,
)
else:
mesh_rows, mesh_cols = mesh_shape
ctx.launch(
f"{panel}_mesh", kernel,
q, k, v, o,
S_q, _S_KV_PER_RANK, _H_Q, _H_KV, _D_HEAD,
mesh_rows, mesh_cols,
rank_axis,
0, # cube_start=0: this panel's launch starts at cube 0
_auto_dim_remap=False,
)
return _bench_fn
# ── Op-log summary harvest ─────────────────────────────────────────────
def _summarize_op_log(op_log) -> dict[str, int]:
"""Counts per ADR-0057 D7 op_log_summary contract."""
gemm_count = 0
ipcq_send_count = 0
ipcq_recv_count = 0
dma_read_count = 0
dma_write_count = 0
for r in op_log:
if r.op_kind == "gemm":
gemm_count += 1
elif r.op_name == "dma_read":
dma_read_count += 1
elif r.op_name == "dma_write":
dma_write_count += 1
elif r.op_name == "ipcq_send":
ipcq_send_count += 1
elif r.op_name == "ipcq_recv":
ipcq_recv_count += 1
elif r.op_name == "ipcq_copy":
# The inbound DMA records ipcq_copy (one per send/recv pair).
# Count it as both a send and a recv side so the row's
# ipcq_send_count and ipcq_recv_count are non-zero even when
# the engine logs the collective via the inbound copy alone.
ipcq_send_count += 1
ipcq_recv_count += 1
return {
"gemm_count": gemm_count,
"ipcq_send_count": ipcq_send_count,
"ipcq_recv_count": ipcq_recv_count,
"dma_read_count": dma_read_count,
"dma_write_count": dma_write_count,
}
def _run_panel(panel: str, topology: str) -> dict:
"""Run one panel via a fresh engine; return its row dict."""
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
topo = resolve_topology(topology)
result = run_bench(
topology=topo, bench_fn=_make_bench_fn(panel),
device=resolve_device(None),
engine_factory=lambda t, d: GraphEngine(
getattr(t, "topology_obj", t), enable_data=True,
),
)
if not result.completion.ok:
raise RuntimeError(
f"milestone-gqa-llama70b panel {panel!r} failed: {result.completion}"
)
_, _, _, n_ranks, _, _ = _PANEL_DISPATCH[panel]
return {
"panel": panel,
"n_ranks": n_ranks,
"op_log_summary": _summarize_op_log(result.engine.op_log),
}
# ── Figure renderers (sub-cycle 4c, 5 of 6 figures) ──────────────────
#
# Sixth figure ``gqa_scaling.png`` is deferred to after sub-cycle 4b
# lands the Q/cube ∈ {1, 2, 4} sweep on multi_user_* panels — it needs
# multiple sweep.json rows per multi_user panel to be meaningful.
_OP_LOG_KEYS = (
"gemm_count",
"ipcq_send_count",
"ipcq_recv_count",
"dma_read_count",
"dma_write_count",
)
_OP_LOG_DISPLAY = {
"gemm_count": "GEMM",
"ipcq_send_count": "IPCQ send",
"ipcq_recv_count": "IPCQ recv",
"dma_read_count": "DMA read",
"dma_write_count": "DMA write",
}
_OP_LOG_COLORS = {
"gemm_count": "#F59E0B",
"ipcq_send_count": "#3B82F6",
"ipcq_recv_count": "#10B981",
"dma_read_count": "#A855F7",
"dma_write_count": "#EF4444",
}
_PANEL_DISPLAY = {
"single_user_prefill": "single_user / prefill",
"multi_user_prefill": "multi_user / prefill",
"single_user_decode": "single_user / decode",
"multi_user_decode": "multi_user / decode",
}
def _load_sweep_data(sweep_json: Path | str) -> dict:
sweep_json = Path(sweep_json)
if not sweep_json.exists():
return {"rows": [], "config": {}, "panels": []}
return json.loads(sweep_json.read_text())
def _row_for(rows: list, panel: str) -> dict | None:
for r in rows:
if r.get("panel") == panel:
return r
return None
def emit_panel_op_log_summary(
panel: str,
sweep_json: Path | str = _SWEEP_JSON,
out_dir: Path | str = _OUTPUT_DIR,
) -> str | None:
"""One bar chart of the 5 op_log counts for ``panel``.
Returns the written PNG path, or ``None`` when sweep.json is empty
or the requested panel is absent.
"""
import matplotlib.pyplot as plt
data = _load_sweep_data(sweep_json)
row = _row_for(data.get("rows", []), panel)
if row is None:
return None
summary = row.get("op_log_summary", {})
n_ranks = row.get("n_ranks")
labels = [_OP_LOG_DISPLAY[k] for k in _OP_LOG_KEYS]
values = [summary.get(k, 0) for k in _OP_LOG_KEYS]
colors = [_OP_LOG_COLORS[k] for k in _OP_LOG_KEYS]
fig, ax = plt.subplots(figsize=(8, 5))
bars = ax.bar(labels, values, color=colors)
for b, v in zip(bars, values):
ax.text(b.get_x() + b.get_width() / 2, b.get_height(),
f"{int(v)}", ha="center", va="bottom", fontsize=9)
ax.set_title(
f"{_PANEL_DISPLAY.get(panel, panel)} (n_ranks={n_ranks})",
fontsize=12, fontweight="bold",
)
ax.set_ylabel("count")
ax.grid(True, axis="y", alpha=0.3)
fig.tight_layout()
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
out = out_dir / f"gqa_op_log_{panel}.png"
fig.savefig(out, dpi=120)
plt.close(fig)
return str(out)
def emit_gqa_comparison(
sweep_json: Path | str = _SWEEP_JSON,
out_dir: Path | str = _OUTPUT_DIR,
) -> str | None:
"""Grouped-bar chart comparing the 5 op_log counts across all panels."""
import matplotlib.pyplot as plt
import numpy as np
data = _load_sweep_data(sweep_json)
panels_in = data.get("panels") or list(_PANELS_V1)
rows = data.get("rows", [])
panels = [p for p in panels_in if _row_for(rows, p) is not None]
if not panels:
return None
n_groups = len(panels)
n_series = len(_OP_LOG_KEYS)
x = np.arange(n_groups)
width = 0.8 / n_series
fig, ax = plt.subplots(figsize=(11, 6))
for i, key in enumerate(_OP_LOG_KEYS):
offset = (i - (n_series - 1) / 2) * width
vals = [_row_for(rows, p)["op_log_summary"].get(key, 0)
for p in panels]
ax.bar(x + offset, vals, width,
label=_OP_LOG_DISPLAY[key], color=_OP_LOG_COLORS[key])
ax.set_xticks(x)
ax.set_xticklabels(
[f"{_PANEL_DISPLAY.get(p, p)}\n(n_ranks={_row_for(rows, p)['n_ranks']})"
for p in panels],
fontsize=8,
)
ax.set_ylabel("count")
ax.set_title("GQA Llama-70B — op_log summary across panels",
fontsize=13, fontweight="bold")
ax.legend(fontsize=8, loc="upper right")
ax.grid(True, axis="y", alpha=0.3)
fig.tight_layout()
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
out = out_dir / "gqa_comparison.png"
fig.savefig(out, dpi=120)
plt.close(fig)
return str(out)
def emit_all_gqa_plots(
sweep_json: Path | str = _SWEEP_JSON,
out_dir: Path | str = _OUTPUT_DIR,
) -> list[str]:
"""Render all 5 in-scope figures and return the written paths.
Sub-cycle 4c v1 emits 5 of the 6 figures ADR-0057 D3 lists; the
6th (gqa_scaling.png) needs sub-cycle 4b's Q/cube sweep data.
"""
paths: list[str] = []
for panel in _PANELS_V1:
p = emit_panel_op_log_summary(panel, sweep_json, out_dir)
if p is not None:
paths.append(p)
comp = emit_gqa_comparison(sweep_json, out_dir)
if comp is not None:
paths.append(comp)
return paths
# ── Bench entry ────────────────────────────────────────────────────────
@bench(
name="milestone-gqa-llama70b",
description="1H milestone: GQA Llama-70B 4-panel sweep (ADR-0057 v1).",
)
def run(torch) -> None:
"""Drive the four GQA panels at validation scale; write sweep.json and figures.
Modes (mutually exclusive):
MILESTONE_FAST=1 Skip the sweep; re-render figures from the
committed sweep.json. Seconds, no simulator.
GQA_VALIDATION=1 Run the four-panel validation sweep + figures.
~1-2h on the full simulator.
Headline-scale mode is deferred to sub-cycle 4c (figures landed
here; headline-scale + scaling figure await sub-cycle 4b).
A sentinel tensor is submitted at the end so run_bench's ADR-0045 D4
"at least one request" contract is satisfied even when the panels
are skipped via MILESTONE_FAST=1.
"""
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
fast = bool(os.environ.get("MILESTONE_FAST"))
if not fast and not os.environ.get("GQA_VALIDATION"):
raise RuntimeError(
"milestone-gqa-llama70b v1 needs GQA_VALIDATION=1 (run the "
"sweep) or MILESTONE_FAST=1 (reuse committed sweep.json). "
"Headline mode is deferred to sub-cycle 4b/4c per ADR-0057 D3."
)
if not fast:
topology = os.environ.get("GQA_TOPOLOGY", "topology.yaml")
rows = [_run_panel(panel, topology) for panel in _PANELS_V1]
sweep = {
"version": 1,
"validation_scale": True,
"panels": list(_PANELS_V1),
"config": {
"S_q_prefill": _S_Q_PREFILL,
"S_kv_per_rank": _S_KV_PER_RANK,
"h_q": _H_Q,
"h_kv": _H_KV,
"d_head": _D_HEAD,
"n_ranks_single_user": _N_RANKS_SINGLE_USER,
"n_ranks_multi_user": _N_RANKS_MULTI_USER,
},
"rows": rows,
}
_SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
print(f" milestone-gqa-llama70b: {len(rows)} rows -> {_SWEEP_JSON}")
elif not _SWEEP_JSON.exists():
raise RuntimeError(
f"MILESTONE_FAST=1 requires {_SWEEP_JSON} to exist; "
"run with GQA_VALIDATION=1 once to seed it."
)
paths = emit_all_gqa_plots()
print(f" milestone-gqa-llama70b: {len(paths)} figures -> {_OUTPUT_DIR} "
f"(fast={fast})")
# Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out).
torch.zeros(
(1, 1), dtype="f16",
dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1),
name="milestone_gqa_sentinel",
)