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

gqa-decode: topology-agnostic reduce + multi-model attention bench

Browse Source 
Make reduce_mlo derive its submesh dimensions (sub_w, sub_h, root_col,
root_row, root_cube) from C at call time, with peer-existence guards on
every inter-cube send/recv so any C ≥ 1 completes cleanly (previously
hardcoded to C=8's 4×2 mesh; non-rectangular C like 7 or 12 hit
IpcqDeadlock at the root cube's east receive). Byte-equal at C=8 —
existing composite digest tests still pass 8/8.

Callsites in the four Case-6 kernels (primitive / primitive-tiled /
composite / composite_extended) updated to pass C into reduce_mlo and
use root_cube_for(C) for the final tl.store gate.

Add the multi-model attention bench: sweeps the composite kernel at
S_kv=128K across six GQA models (Gemma-2 27B, LLaMA-3 8B, Qwen 2.5 7B,
LLaMA-3 70B, Qwen 2.5 72B, Command R+), with per-model topology
C = h_q (cubes per KV group = query heads per KV group). Captures
per-op-kind occupancy (matmul GEMM+MATH, comm DMA) alongside latency
and PE_CPU dispatch. Three-panel plot writes to bench-output and
paper-figures dir.

Headline result: same-h_q-different-family models land within 0.2 µs
(model-agnostic given fixed h_kv=1 per KV group + d_head=128); latency
climbs 263 → 492 µs as C climbs 2 → 12, driven by reduce depth
(comm/matmul ratio 0.4 → 0.6 across the sweep).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-07-22 09:45:51 -07:00
parent 23a2bc3fb3
commit d0db9ff40e
10 changed files with 840 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_models.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 169 KiB

src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_models.json
+173
View File
@@ -0,0 +1,173 @@
{
"version": 1,
"variant": "composite",
"S_kv": 131072,
"P": 8,
"note": "C = G (G-matched topology per model)",
"models": [
"llama3-8b",
"llama3-70b",
"qwen2.5-7b",
"qwen2.5-72b",
"gemma2-27b",
"command-r-plus"
],
"rows": [
{
"model": "llama3-8b",
"family": "Meta",
"params_b": 8,
"h_q": 4,
"h_kv": 1,
"G": 4,
"d_head": 128,
"full_h_q": 32,
"full_h_kv": 8,
"hidden": 4096,
"layers": 32,
"C": 4,
"P": 8,
"N": 32,
"T_q": 1,
"S_kv": 131072,
"S_local": 4096,
"variant": "composite",
"pe_cpu_cmd_count": 44,
"pe_cpu_dispatch_cycles": 494,
"latency_ns": 394539.4071250012,
"matmul_ns": 4783.8393749997485,
"comm_ns": 1812.7050000057789,
"other_ns": 1684807746.5613055
},
{
"model": "llama3-70b",
"family": "Meta",
"params_b": 70,
"h_q": 8,
"h_kv": 1,
"G": 8,
"d_head": 128,
"full_h_q": 64,
"full_h_kv": 8,
"hidden": 8192,
"layers": 80,
"C": 8,
"P": 8,
"N": 64,
"T_q": 1,
"S_kv": 131072,
"S_local": 2048,
"variant": "composite",
"pe_cpu_cmd_count": 94,
"pe_cpu_dispatch_cycles": 978,
"latency_ns": 460651.3170000027,
"matmul_ns": 10069.639499999117,
"comm_ns": 4287.750000016997,
"other_ns": 2079259822.423066
},
{
"model": "qwen2.5-7b",
"family": "Alibaba",
"params_b": 7,
"h_q": 7,
"h_kv": 1,
"G": 7,
"d_head": 128,
"full_h_q": 28,
"full_h_kv": 4,
"hidden": 3584,
"layers": 28,
"C": 7,
"P": 8,
"N": 56,
"T_q": 1,
"S_kv": 131072,
"S_local": 2340,
"variant": "composite",
"pe_cpu_cmd_count": 77,
"pe_cpu_dispatch_cycles": 813,
"latency_ns": 456482.3358750015,
"matmul_ns": 8809.787499998813,
"comm_ns": 3486.4500000112457,
"other_ns": 2032890493.459298
},
{
"model": "qwen2.5-72b",
"family": "Alibaba",
"params_b": 72,
"h_q": 8,
"h_kv": 1,
"G": 8,
"d_head": 128,
"full_h_q": 64,
"full_h_kv": 8,
"hidden": 8192,
"layers": 80,
"C": 8,
"P": 8,
"N": 64,
"T_q": 1,
"S_kv": 131072,
"S_local": 2048,
"variant": "composite",
"pe_cpu_cmd_count": 94,
"pe_cpu_dispatch_cycles": 978,
"latency_ns": 460651.3170000027,
"matmul_ns": 10069.639499999117,
"comm_ns": 4287.750000016997,
"other_ns": 2079259822.423066
},
{
"model": "gemma2-27b",
"family": "Google",
"params_b": 27,
"h_q": 2,
"h_kv": 1,
"G": 2,
"d_head": 128,
"full_h_q": 32,
"full_h_kv": 16,
"hidden": 4608,
"layers": 46,
"C": 2,
"P": 8,
"N": 16,
"T_q": 1,
"S_kv": 131072,
"S_local": 8192,
"variant": "composite",
"pe_cpu_cmd_count": 60,
"pe_cpu_dispatch_cycles": 648,
"latency_ns": 263522.86775000195,
"matmul_ns": 2326.0945000193315,
"comm_ns": 907.5250000030501,
"other_ns": 1235149167.507364
},
{
"model": "command-r-plus",
"family": "Cohere",
"params_b": 104,
"h_q": 12,
"h_kv": 1,
"G": 12,
"d_head": 128,
"full_h_q": 96,
"full_h_kv": 8,
"hidden": 12288,
"layers": 64,
"C": 12,
"P": 8,
"N": 96,
"T_q": 1,
"S_kv": 131072,
"S_local": 1365,
"variant": "composite",
"pe_cpu_cmd_count": 111,
"pe_cpu_dispatch_cycles": 1143,
"latency_ns": 491928.65900000196,
"matmul_ns": 15894.392999998527,
"comm_ns": 9904.505000027013,
"other_ns": 2372909985.3171477
}
]
}
src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py
+4 -4
View File
@@ -37,9 +37,9 @@ Topology / SFR:
from __future__ import annotations
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
_ROOT_CUBE,
_merge_running,
reduce_mlo,
root_cube_for,
)
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
@@ -108,9 +108,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
tl.copy_to(O_local, O_new)
# ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
# ── Final normalise + store (root only: PE 0 of CUBE 6) ──
if pe_id == 0 and cube_id == _ROOT_CUBE:
# ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
if pe_id == 0 and cube_id == root_cube_for(C):
O_final = O_local / l_local
tl.store(o_ptr, O_final)
src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite.py
+4 -4
View File
@@ -18,8 +18,8 @@ under the ADR-0064 Rev2 structural model (the CPU-offload win).
from __future__ import annotations
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
_ROOT_CUBE,
reduce_mlo,
root_cube_for,
)
@@ -63,9 +63,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel(
tl.composite(op="gemm", a=exp_scores, b=V, out=O_local) # P·V, one command
# ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
# ── Final normalise + store (root only: PE 0 of CUBE 6) ──
if pe_id == 0 and cube_id == _ROOT_CUBE:
# ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
if pe_id == 0 and cube_id == root_cube_for(C):
O_final = O_local / l_local
tl.store(o_ptr, O_final)
src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext.py
+4 -4
View File
@@ -27,8 +27,8 @@ numeric parity of the recipe's 8 MATH ops is a separate follow-up
from __future__ import annotations
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
_ROOT_CUBE,
reduce_mlo,
root_cube_for,
)
# Small primitive seed slice that establishes the running (m, ℓ, O) the
@@ -86,9 +86,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel(
)
# ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
# ── Final normalise + store (root only: PE 0 of CUBE 6) ──
if pe_id == 0 and cube_id == _ROOT_CUBE:
# ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
if pe_id == 0 and cube_id == root_cube_for(C):
O_final = O_local / l_local
tl.store(o_ptr, O_final)
src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16.py
+3 -3
View File
@@ -37,9 +37,9 @@ from math import ceil
from kernbench.common.pe_commands import GemmCmd
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
_ROOT_CUBE,
_merge_running,
reduce_mlo,
root_cube_for,
)
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
@@ -208,8 +208,8 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel(
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
if pe_id == 0 and cube_id == _ROOT_CUBE:
if pe_id == 0 and cube_id == root_cube_for(C):
O_final = O_local / l_local
tl.store(o_ptr, O_final)
src/kernbench/benches/gqa_helpers/long_ctx/_gqa_mlo_reduce.py
+198 -132
View File
@@ -1,33 +1,66 @@
"""Shared (m, ℓ, O) reduce for the Case-6 long-context decode kernels.
Extracted from the Cube-SP × PE-SP decode kernel so the three
command-form variants (primitive / composite / composite_extended)
Extracted from the Cube-SP × PE-SP decode kernel so the four command-form
variants (primitive / primitive-tiled / composite / composite_extended)
share one identical reduce. The local attention differs per variant;
the reduce does not. Behavior is byte-equal to the pre-extraction inline
reduce (guarded by tests/attention/test_gqa_decode_long_ctx_composite.py).
the reduce does not. Byte-equal behavior at C=8 is guarded by
``tests/attention/test_gqa_decode_long_ctx_composite.py``.
The reduce is two-level:
• intra-CUBE 8-way (row chain along intra_W + col bridge along
intra_N) over the 2×4 PE grid;
• inter-CUBE 2-phase lrab over the 4×2 CUBE sub-mesh (ADR-0060 §4.2
lrab-adapted center-root reduce): Phase 1 row reduce converges at
root_col=2; Phase 2 col reduce on root_col converges at root_row=1;
root cube id = 6.
Plain ``+`` is replaced by log-sum-exp ``_merge_running``.
• inter-CUBE 2-phase lrab-adapted center-root reduce over a
``sub_w × sub_h`` CUBE sub-mesh (ADR-0060 §4.2): Phase 1 does a
bidirectional row reduce converging at ``root_col = sub_w // 2``;
Phase 2 does a bidirectional col reduce on ``root_col`` converging
at ``root_row = sub_h // 2``.
The submesh dimensions ``(sub_w, sub_h)`` and root are **derived from
C** (number of cubes launched per KV group), not hardcoded, so the
same reduce works for any ``C ≥ 1``. Peer-existence guards on every
send/recv elide operations that would target un-launched neighbors
(needed for non-rectangular C like 7 or 12; safe no-ops when the
neighbor exists). For ``C = 8`` this reproduces the previous
hardcoded ``{sub_w=4, sub_h=2, root_col=2, root_row=1, root_cube=6}``
exactly.
Plain ``+`` in the tree is replaced by log-sum-exp ``_merge_running``.
"""
from __future__ import annotations
# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh).
_SUB_W = 4
_SUB_H = 2
_ROOT_COL = _SUB_W // 2 # 2
_ROOT_ROW = _SUB_H // 2 # 1
_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL # 6
PE_GRID_COLS = 4
def _submesh_for(C):
"""Derive ``(sub_w, sub_h, root_col, root_row, root_cube)`` for a
C-cube inter-cube reduce.
Submesh is at most 4 columns wide (matching the SIP's 4×4 cube
mesh); rows fill left-to-right, top-to-bottom. For the Case-6
baseline C=8, this returns ``(4, 2, 2, 1, 6)`` — identical to
the previously hardcoded values.
"""
sub_w = min(C, 4)
sub_h = (C + sub_w - 1) // sub_w
root_col = sub_w // 2
root_row = sub_h // 2
root_cube = root_row * sub_w + root_col
return sub_w, sub_h, root_col, root_row, root_cube
def root_cube_for(C):
"""Root cube id for a C-cube reduce. Callers use this to gate the
final ``tl.store`` in the kernel epilogue."""
_, _, _, _, root_cube = _submesh_for(C)
return root_cube
# Backward-compat alias: preserves the Case-6 baseline value (6) for
# any external code that still imports the module-level constant.
_ROOT_CUBE = root_cube_for(8)
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
m_new = tl.maximum(m_local, m_other)
@@ -38,13 +71,23 @@ def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
return m_new, l_new, O_new
def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, *, tl):
"""Two-level (m, ℓ, O) reduce-to-root (PE 0 of CUBE 6). In place.
def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, *, tl):
"""Two-level (m, ℓ, O) reduce-to-root (PE 0 of ``root_cube_for(C)``).
In place.
Updates ``m_local``/``l_local``/``O_local`` in place via ``copy_to``;
after this call only PE 0 of CUBE ``_ROOT_CUBE`` holds the full result.
Updates ``m_local`` / ``l_local`` / ``O_local`` in place via
``copy_to``; after this call only PE 0 of ``root_cube_for(C)``
holds the fully-merged result.
Peer-existence guards on every inter-cube send/recv ensure the
reduce completes for any launched C (rectangular or not); a
send/recv toward a non-launched neighbor is skipped. For C = 8
every neighbor exists → guards are all True → byte-equal to the
previously hardcoded Case-6 reduce.
"""
# ── Intra-CUBE reduce: row chain (intra_W) + col bridge (intra_N) ──
sub_w, sub_h, root_col, root_row, _root_cube = _submesh_for(C)
# ── Intra-CUBE reduce (unchanged — depends only on P) ─────────────
pe_col = pe_id % PE_GRID_COLS
pe_row = pe_id // PE_GRID_COLS
pe_cols_used = min(PE_GRID_COLS, P)
@@ -84,106 +127,126 @@ def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, *, tl):
tl.send(dir="intra_N", src=l_local)
tl.send(dir="intra_N", src=O_local)
# ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ──
# ── Inter-CUBE lrab-adapted center-root reduce ────────────────────
# Only PE 0 of each CUBE participates. Adapts Phases 1-2 of
# lrab_hierarchical_allreduce.py: bidirectional row reduce converges
# at root_col; bidirectional col reduce on root_col converges at
# root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``.
# at ``root_col``; bidirectional col reduce on ``root_col`` converges
# at ``root_row``. Plain ``+`` replaced by log-sum-exp
# ``_merge_running``. Peer-existence guards elide sends/recvs whose
# neighbor cube was not launched.
if pe_id == 0:
row = cube_id // _SUB_W
col = cube_id % _SUB_W
row = cube_id // sub_w
col = cube_id % sub_w
# Phase 1: row reduce — converge at col == _ROOT_COL.
# Peer existence within the launched set of C cubes.
east_exists = (col + 1 < sub_w) and (cube_id + 1 < C)
west_exists = col > 0
south_exists = cube_id + sub_w < C
north_exists = row > 0
# Phase 1: row reduce — converge at col == root_col.
if col == 0:
tl.send(dir="E", src=m_local)
tl.send(dir="E", src=l_local)
tl.send(dir="E", src=O_local)
elif 0 < col < _ROOT_COL:
with tl.scratch_scope():
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
tl.send(dir="E", src=m_local)
tl.send(dir="E", src=l_local)
tl.send(dir="E", src=O_local)
elif col == _ROOT_COL:
with tl.scratch_scope():
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
with tl.scratch_scope():
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
elif _ROOT_COL < col < _SUB_W - 1:
with tl.scratch_scope():
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
tl.send(dir="W", src=m_local)
tl.send(dir="W", src=l_local)
tl.send(dir="W", src=O_local)
elif col == _SUB_W - 1:
tl.send(dir="W", src=m_local)
tl.send(dir="W", src=l_local)
tl.send(dir="W", src=O_local)
if east_exists:
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:
if west_exists:
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 east_exists:
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 west_exists:
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 east_exists:
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:
if east_exists:
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 west_exists:
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:
if west_exists:
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:
# Phase 2: col reduce on col == root_col — converge at row == root_row.
if col == root_col:
if row == 0:
tl.send(dir="S", src=m_local)
tl.send(dir="S", src=l_local)
tl.send(dir="S", src=O_local)
elif 0 < row < _ROOT_ROW:
with tl.scratch_scope():
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
tl.send(dir="S", src=m_local)
tl.send(dir="S", src=l_local)
tl.send(dir="S", src=O_local)
elif row == _ROOT_ROW:
with tl.scratch_scope():
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
if _SUB_H - 1 > _ROOT_ROW:
if south_exists:
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:
if north_exists:
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 south_exists:
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 north_exists:
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 and south_exists:
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")
@@ -194,21 +257,24 @@ def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, *, 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)
elif root_row < row < sub_h - 1:
if south_exists:
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)
if north_exists:
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:
if north_exists:
tl.send(dir="N", src=m_local)
tl.send(dir="N", src=l_local)
tl.send(dir="N", src=O_local)
src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_models.py
+262
View File
@@ -0,0 +1,262 @@
"""GQA composite kernel across models — per-KV-group multi-model comparison.
The topology per model uses C = h_q (cubes per KV group = query heads
per KV group in the model). P = 8 PEs / cube is fixed to the physical
layout. So Gemma-2 27B (G=2) runs on a 2-cube slice; Command R+ (G=12)
runs on 12 cubes; LLaMA-3-70B (G=8) matches the Case-6 baseline
exactly.
Companion to ``gqa_decode_long_ctx_models.py``. That bench sweeps a
fixed set of topology sizes (N ∈ {16, 32, 64}) across six GQA models
and showed that latency is model-agnostic when only G varies.
This bench takes a different slice: it fixes the topology to match
each model's G ratio — one cube per Q head in the KV group:
C = G (cubes per KV group = h_q of the KV group)
P = 8 (fixed, physical PEs per cube)
N = C · P = 8·G
Six models, one topology per model, one context length:
Gemma 2 27B G=2 → C=2 → N=16
LLaMA-3 8B G=4 → C=4 → N=32
Qwen 2.5 7B G=7 → C=7 → N=56 (2 KV groups per SIP, 2 cubes idle)
LLaMA-3 70B G=8 → C=8 → N=64
Qwen 2.5 72B G=8 → C=8 → N=64
Command R+ G=12 → C=12 → N=96 (1 KV group per SIP, 4 cubes idle)
Physical SIP is 4×4 = 16 cubes (topology.yaml sip.cube_mesh: {w:4, h:4}),
so all six configurations fit in one SIP.
Total: 6 engine runs at S_kv = 131 072 (~40 min wall-clock).
"""
from __future__ import annotations
import json
from pathlib import Path
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite import ( # noqa: E501
gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
)
from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
from kernbench.policy.placement.dp import DPPolicy
_OUTPUT_DIR = (
Path(__file__).resolve().parents[2]
/ "1H_milestone_output" / "gqa" / "long_ctx"
)
_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_models.json"
# ── Model registry (per-KV-group shapes) ─────────────────────────────
# Kernel operates on ONE KV group at a time. h_q / h_kv here are the
# per-group counts (not the model's totals); G = h_q / h_kv.
# Full-model context (full_h_q, full_h_kv, hidden, layers) is metadata
# for annotation only — it does not enter the kernel launch.
_MODELS = {
"llama3-8b": dict(
family="Meta", params_b=8,
h_q=4, h_kv=1, d_head=128,
full_h_q=32, full_h_kv=8, hidden=4096, layers=32,
),
"llama3-70b": dict(
family="Meta", params_b=70,
h_q=8, h_kv=1, d_head=128,
full_h_q=64, full_h_kv=8, hidden=8192, layers=80,
),
"qwen2.5-7b": dict(
family="Alibaba", params_b=7,
h_q=7, h_kv=1, d_head=128,
full_h_q=28, full_h_kv=4, hidden=3584, layers=28,
),
"qwen2.5-72b": dict(
family="Alibaba", params_b=72,
h_q=8, h_kv=1, d_head=128,
full_h_q=64, full_h_kv=8, hidden=8192, layers=80,
),
"gemma2-27b": dict(
family="Google", params_b=27,
h_q=2, h_kv=1, d_head=128,
full_h_q=32, full_h_kv=16, hidden=4608, layers=46,
),
"command-r-plus": dict(
family="Cohere", params_b=104,
h_q=12, h_kv=1, d_head=128,
full_h_q=96, full_h_kv=8, hidden=12288, layers=64,
),
}
_MODEL_ORDER = tuple(_MODELS.keys())
_S_KV = 131_072 # 128 K context
_T_Q = 1
_P = 8 # fixed physical PEs per cube
def _c_for(model_key: str) -> int:
m = _MODELS[model_key]
return m["h_q"] // m["h_kv"]
def _end_to_end_ns(op_log) -> float:
if not op_log:
return 0.0
return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
def _op_breakdown_ns(op_log) -> dict[str, float]:
"""Sum occupancy per op_kind across all components (not
critical-path — ops overlap on independent components). First-order
view of where time is spent inside the engine."""
totals = {"matmul": 0.0, "comm": 0.0, "other": 0.0}
for r in op_log:
dur = r.t_end - r.t_start
if r.op_kind in ("gemm", "math"):
totals["matmul"] += dur
elif r.op_kind == "memory":
totals["comm"] += dur
else:
totals["other"] += dur
return totals
def _run_panel_fn(model_key: str, C: int, P: int, S_kv: int):
m = _MODELS[model_key]
panel = f"decode_models_{model_key}_c{C}p{P}_s{S_kv}"
def _bench_fn(ctx):
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_full = DPPolicy(cube="replicate", pe="replicate",
num_cubes=C, num_pes=P)
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
num_cubes=C, num_pes=P)
q = ctx.zeros((_T_Q, m["h_q"] * m["d_head"]),
dtype="f16", dp=dp_full, name=f"{panel}_q")
k = ctx.zeros((S_kv, m["h_kv"] * m["d_head"]),
dtype="f16", dp=dp_kv, name=f"{panel}_k")
v = ctx.zeros((S_kv, m["h_kv"] * m["d_head"]),
dtype="f16", dp=dp_kv, name=f"{panel}_v")
o = ctx.empty((_T_Q, m["h_q"] * m["d_head"]),
dtype="f16", dp=dp_full, name=f"{panel}_o")
ctx.launch(
panel,
gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
q, k, v, o,
_T_Q, S_kv, m["h_q"], m["h_kv"], m["d_head"], C, P,
_auto_dim_remap=False,
)
return _bench_fn
def _emit_dispatch(model_key: str, C: int, P: int,
S_kv: int) -> tuple[int, float]:
from kernbench.common.pe_commands import PeCpuOverheadCmd
from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL
from kernbench.triton_emu.tl_context import TLContext, run_kernel
m = _MODELS[model_key]
cube_id = min(6, C - 1)
tl = TLContext(
pe_id=0, num_programs=P, cost_model=DEFAULT_PE_COST_MODEL,
cube_id=cube_id, num_cubes=C, scratch_base=1 << 61,
scratch_size=1 << 20,
)
run_kernel(
gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel, tl,
0x1000, 0x2000, 0x3000, 0x4000,
_T_Q, S_kv, m["h_q"], m["h_kv"], m["d_head"], C, P,
)
cmds = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)]
return len(cmds), sum(c.cycles for c in cmds)
def _engine_run(model_key: str, C: int, P: int,
S_kv: int, topology: str):
"""Run the engine sim; return (latency_ns, op_kind_breakdown)."""
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=_run_panel_fn(model_key, C, P, S_kv),
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"gqa-decode-models {model_key} C={C} P={P} S={S_kv} "
f"failed: {result.completion}"
)
op_log = result.engine.op_log
latency = _end_to_end_ns(op_log)
breakdown = _op_breakdown_ns(op_log)
return latency, breakdown
def run_sweep(topology: str = "topology.yaml") -> int:
"""One row per model at its per-KV-group topology (C = h_q). Writes
sweep_decode_models.json."""
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
rows = []
for model_key in _MODEL_ORDER:
m = _MODELS[model_key]
C = _c_for(model_key)
N = C * _P
n_cmds, cycles = _emit_dispatch(model_key, C, _P, _S_KV)
latency, breakdown = _engine_run(model_key, C, _P, _S_KV, topology)
rows.append({
"model": model_key,
"family": m["family"],
"params_b": m["params_b"],
"h_q": m["h_q"],
"h_kv": m["h_kv"],
"G": C,
"d_head": m["d_head"],
"full_h_q": m["full_h_q"],
"full_h_kv": m["full_h_kv"],
"hidden": m["hidden"],
"layers": m["layers"],
"C": C, "P": _P, "N": N,
"T_q": _T_Q,
"S_kv": _S_KV,
"S_local": _S_KV // N,
"variant": "composite",
"pe_cpu_cmd_count": n_cmds,
"pe_cpu_dispatch_cycles": cycles,
"latency_ns": latency,
# Op-kind occupancy (summed across all components — not
# critical-path; first-order view of where engine time goes)
"matmul_ns": breakdown["matmul"],
"comm_ns": breakdown["comm"],
"other_ns": breakdown["other"],
})
print(
f" {model_key:<18} G={C:>2} C={C:>2} P={_P} N={N:>3} "
f"latency={latency/1e3:>7.1f} µs cmds={n_cmds:>4} "
f"matmul={breakdown['matmul']/1e3:>7.1f} "
f"comm={breakdown['comm']/1e3:>7.1f}"
)
sweep = {
"version": 1,
"variant": "composite",
"S_kv": _S_KV,
"P": _P,
"note": "C = h_q per model (cubes per KV group = query heads per KV group)",
"models": list(_MODEL_ORDER),
"rows": rows,
}
_SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
print(f" gqa-decode-models: {len(rows)} rows -> {_SWEEP_JSON}")
return len(rows)
if __name__ == "__main__":
run_sweep()