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

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:
2026-07-22 09:45:51 -07:00
parent 23a2bc3fb3
commit d0db9ff40e
10 changed files with 840 additions and 147 deletions
Binary file not shown.

After

Width:  |  Height:  |  Size: 169 KiB

@@ -0,0 +1,192 @@
"""Bar plot for the multi-model GQA composite bench.
Topology per model: cubes per KV group C = h_q, P = 8 PEs / cube.
Reads ``sweep_decode_models.json`` (produced by
``gqa_decode_long_ctx_models.py``) and writes a two-panel PNG
comparing end-to-end latency and PE_CPU dispatch across six GQA models
where each model uses a topology sized to match its G value
(cubes per KV group = h_q per KV group; PEs per cube = 8):
Left panel — end-to-end decode latency (µs), one bar per model, sorted
by G. Bar labels show C, N, and latency.
Right panel — PE_CPU command count per model.
Run (after the bench):
python scripts/paper/paper_plot_gqa_decode_models.py
"""
from __future__ import annotations
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt # noqa: E402
import numpy as np # noqa: E402
_REPO_ROOT = Path(__file__).resolve().parents[2]
_FIG_DIR = (
_REPO_ROOT / "src" / "kernbench" / "benches"
/ "1H_milestone_output" / "gqa" / "long_ctx"
)
_SWEEP_JSON = _FIG_DIR / "sweep_decode_models.json"
_PAPER_FIG_DIR = (
_REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
)
# Model display labels
_MODEL_LABELS = {
"gemma2-27b": "Gemma 2 27B\n(G=2)",
"llama3-8b": "LLaMA-3 8B\n(G=4)",
"qwen2.5-7b": "Qwen 2.5 7B\n(G=7)",
"llama3-70b": "LLaMA-3 70B\n(G=8)",
"qwen2.5-72b": "Qwen 2.5 72B\n(G=8)",
"command-r-plus": "Command R+\n(G=12)",
}
# Bar color per model family
_FAMILY_COLOUR = {
"Google": "#4f8a4f",
"Meta": "#3b6ea5",
"Alibaba": "#c86432",
"Cohere": "#8b5a2b",
}
def main() -> None:
sweep = json.loads(_SWEEP_JSON.read_text())
rows_by_model = {r["model"]: r for r in sweep["rows"]}
# Sort models by G (ascending) so the trend is left-to-right.
ordered = sorted(_MODEL_LABELS.keys(), key=lambda m: rows_by_model[m]["G"])
fig, (ax_lat, ax_cmd, ax_brk) = plt.subplots(1, 3, figsize=(18.5, 5.0))
xs = np.arange(len(ordered))
# ── Left: latency ─────────────────────────────────────────────
latencies = [rows_by_model[m]["latency_ns"] / 1e3 for m in ordered]
colours = [_FAMILY_COLOUR[rows_by_model[m]["family"]] for m in ordered]
bars = ax_lat.bar(xs, latencies, 0.62,
color=colours, edgecolor="#222", linewidth=0.7)
for i, m in enumerate(ordered):
r = rows_by_model[m]
# Latency label above bar
ax_lat.text(i, latencies[i] + max(latencies) * 0.015,
f"{latencies[i]:.1f} µs",
ha="center", va="bottom", fontsize=9,
fontweight="bold", color="#222")
# Topology label inside bar
ax_lat.text(i, latencies[i] / 2,
f"cubes = {r['C']}",
ha="center", va="center", fontsize=9, color="white")
ax_lat.set_xticks(xs)
ax_lat.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
ax_lat.set_ylabel("end-to-end decode latency (µs)")
ax_lat.set_title(
f"Composite Case-6 decode latency at $S_{{kv}}=128$K\n"
"topology per model: cubes per KV group = h_q of the KV group",
fontsize=11,
)
ax_lat.grid(True, axis="y", ls=":", alpha=0.5)
ax_lat.set_ylim(0, max(latencies) * 1.18)
# Legend (family)
from matplotlib.patches import Patch
seen_families = []
handles = []
for m in ordered:
fam = rows_by_model[m]["family"]
if fam not in seen_families:
seen_families.append(fam)
handles.append(Patch(facecolor=_FAMILY_COLOUR[fam],
edgecolor="#222", label=fam))
ax_lat.legend(handles=handles, loc="upper left", fontsize=9,
title="family")
# ── Right: PE_CPU dispatch ─────────────────────────────────────
cmds = [rows_by_model[m]["pe_cpu_cmd_count"] for m in ordered]
bars2 = ax_cmd.bar(xs, cmds, 0.62,
color=colours, edgecolor="#222", linewidth=0.7)
for i, m in enumerate(ordered):
ax_cmd.text(i, cmds[i] + max(cmds) * 0.02,
f"{cmds[i]}",
ha="center", va="bottom", fontsize=9,
fontweight="bold", color="#222")
r = rows_by_model[m]
ax_cmd.text(i, cmds[i] / 2,
f"cubes = {r['C']}",
ha="center", va="center", fontsize=9, color="white")
ax_cmd.set_xticks(xs)
ax_cmd.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
ax_cmd.set_ylabel("PE_CPU commands issued")
ax_cmd.set_title(
"PE_CPU dispatch (composite, per KV group)\n"
"growing sub-mesh width or height → more reduce hops",
fontsize=11,
)
ax_cmd.grid(True, axis="y", ls=":", alpha=0.5)
ax_cmd.set_ylim(0, max(cmds) * 1.18)
# ── Third: matmul vs comm (op-kind occupancy) ─────────────────────
# Only the two kinds that carry useful work — matmul (GEMM + MATH
# occupancy summed across all engines) and comm (DMA occupancy).
# These are op-log sums across components, not critical-path
# attribution; parallelism means they don't sum to wall-clock
# latency. Bars are absolute µs so the reduce-cost growth with C
# is visible.
matmul_us = [rows_by_model[m].get("matmul_ns", 0.0) / 1e3 for m in ordered]
comm_us = [rows_by_model[m].get("comm_ns", 0.0) / 1e3 for m in ordered]
max_v = max(max(matmul_us), max(comm_us))
width = 0.35
ax_brk.bar(xs - width / 2, matmul_us, width,
color="#3b6ea5", edgecolor="#222", linewidth=0.6,
label="matmul (GEMM + MATH)")
ax_brk.bar(xs + width / 2, comm_us, width,
color="#c0504d", edgecolor="#222", linewidth=0.6,
label="communication (DMA)")
for i in range(len(ordered)):
ax_brk.text(xs[i] - width / 2, matmul_us[i] + max_v * 0.015,
f"{matmul_us[i]:.1f}",
ha="center", va="bottom", fontsize=8,
color="#222", fontweight="bold")
ax_brk.text(xs[i] + width / 2, comm_us[i] + max_v * 0.015,
f"{comm_us[i]:.1f}",
ha="center", va="bottom", fontsize=8,
color="#222", fontweight="bold")
ax_brk.set_xticks(xs)
ax_brk.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
ax_brk.set_ylabel("op-log occupancy (µs, summed across engines)")
ax_brk.set_title(
"Matmul vs communication engine work per model\n"
"sum of GEMM/MATH and DMA occupancy — not wall-clock (overlap)",
fontsize=11,
)
ax_brk.grid(True, axis="y", ls=":", alpha=0.5)
ax_brk.set_ylim(0, max_v * 1.20)
ax_brk.legend(loc="upper left", fontsize=9)
fig.suptitle(
"Multi-model attention comparison — Case-6 composite kernel"
" ($S_{kv}=128$K, $T_q=1$, P=8 PEs / cube)",
fontsize=12, fontweight="bold",
)
fig.tight_layout(rect=(0, 0, 1, 0.93))
out = _FIG_DIR / "gqa_decode_models.png"
fig.savefig(out, dpi=150)
plt.close(fig)
print(f"wrote {out}")
if _PAPER_FIG_DIR.is_dir():
dst = _PAPER_FIG_DIR / out.name
dst.write_bytes(out.read_bytes())
print(f"copied {dst}")
if __name__ == "__main__":
main()
Binary file not shown.

After

Width:  |  Height:  |  Size: 169 KiB

@@ -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
}
]
}
@@ -37,9 +37,9 @@ Topology / SFR:
from __future__ import annotations from __future__ import annotations
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import ( from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
_ROOT_CUBE,
_merge_running, _merge_running,
reduce_mlo, reduce_mlo,
root_cube_for,
) )
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). 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) tl.copy_to(O_local, O_new)
# ── Two-level (m, , O) reduce-to-root (shared helper) ── # ── 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) ── # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
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 O_final = O_local / l_local
tl.store(o_ptr, O_final) tl.store(o_ptr, O_final)
@@ -18,8 +18,8 @@ under the ADR-0064 Rev2 structural model (the CPU-offload win).
from __future__ import annotations from __future__ import annotations
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import ( from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
_ROOT_CUBE,
reduce_mlo, 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 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) ── # ── 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) ── # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
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 O_final = O_local / l_local
tl.store(o_ptr, O_final) tl.store(o_ptr, O_final)
@@ -27,8 +27,8 @@ numeric parity of the recipe's 8 MATH ops is a separate follow-up
from __future__ import annotations from __future__ import annotations
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import ( from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
_ROOT_CUBE,
reduce_mlo, reduce_mlo,
root_cube_for,
) )
# Small primitive seed slice that establishes the running (m, , O) the # 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) ── # ── 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) ── # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
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 O_final = O_local / l_local
tl.store(o_ptr, O_final) tl.store(o_ptr, O_final)
@@ -37,9 +37,9 @@ from math import ceil
from kernbench.common.pe_commands import GemmCmd from kernbench.common.pe_commands import GemmCmd
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import ( from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
_ROOT_CUBE,
_merge_running, _merge_running,
reduce_mlo, reduce_mlo,
root_cube_for,
) )
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). 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(l_local, l_new)
tl.copy_to(O_local, O_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 O_final = O_local / l_local
tl.store(o_ptr, O_final) tl.store(o_ptr, O_final)
@@ -1,33 +1,66 @@
"""Shared (m, , O) reduce for the Case-6 long-context decode kernels. """Shared (m, , O) reduce for the Case-6 long-context decode kernels.
Extracted from the Cube-SP × PE-SP decode kernel so the three Extracted from the Cube-SP × PE-SP decode kernel so the four command-form
command-form variants (primitive / composite / composite_extended) variants (primitive / primitive-tiled / composite / composite_extended)
share one identical reduce. The local attention differs per variant; share one identical reduce. The local attention differs per variant;
the reduce does not. Behavior is byte-equal to the pre-extraction inline the reduce does not. Byte-equal behavior at C=8 is guarded by
reduce (guarded by tests/attention/test_gqa_decode_long_ctx_composite.py). ``tests/attention/test_gqa_decode_long_ctx_composite.py``.
The reduce is two-level: The reduce is two-level:
• intra-CUBE 8-way (row chain along intra_W + col bridge along • intra-CUBE 8-way (row chain along intra_W + col bridge along
intra_N) over the 2×4 PE grid; intra_N) over the 2×4 PE grid;
• inter-CUBE 2-phase lrab over the 4×2 CUBE sub-mesh (ADR-0060 §4.2 • inter-CUBE 2-phase lrab-adapted center-root reduce over a
lrab-adapted center-root reduce): Phase 1 row reduce converges at ``sub_w × sub_h`` CUBE sub-mesh (ADR-0060 §4.2): Phase 1 does a
root_col=2; Phase 2 col reduce on root_col converges at root_row=1; bidirectional row reduce converging at ``root_col = sub_w // 2``;
root cube id = 6. Phase 2 does a bidirectional col reduce on ``root_col`` converging
Plain ``+`` is replaced by log-sum-exp ``_merge_running``. 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 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 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): 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.""" """Online-softmax merge of two partial ``(m, , O)`` triples."""
m_new = tl.maximum(m_local, m_other) 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 return m_new, l_new, O_new
def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, *, tl): 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 CUBE 6). In place. """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``; Updates ``m_local`` / ``l_local`` / ``O_local`` in place via
after this call only PE 0 of CUBE ``_ROOT_CUBE`` holds the full result. ``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_col = pe_id % PE_GRID_COLS
pe_row = pe_id // PE_GRID_COLS pe_row = pe_id // PE_GRID_COLS
pe_cols_used = min(PE_GRID_COLS, P) 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=l_local)
tl.send(dir="intra_N", src=O_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 # Only PE 0 of each CUBE participates. Adapts Phases 1-2 of
# lrab_hierarchical_allreduce.py: bidirectional row reduce converges # lrab_hierarchical_allreduce.py: bidirectional row reduce converges
# at root_col; bidirectional col reduce on root_col converges at # at ``root_col``; bidirectional col reduce on ``root_col`` converges
# root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``. # 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: if pe_id == 0:
row = cube_id // _SUB_W row = cube_id // sub_w
col = 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: if col == 0:
tl.send(dir="E", src=m_local) if east_exists:
tl.send(dir="E", src=l_local) tl.send(dir="E", src=m_local)
tl.send(dir="E", src=O_local) tl.send(dir="E", src=l_local)
elif 0 < col < _ROOT_COL: tl.send(dir="E", src=O_local)
with tl.scratch_scope(): elif 0 < col < root_col:
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16") if west_exists:
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16") with tl.scratch_scope():
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16") m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running( l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
) m_new, l_new, O_new = _merge_running(
tl.copy_to(m_local, m_new) m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
tl.copy_to(l_local, l_new) )
tl.copy_to(O_local, O_new) tl.copy_to(m_local, m_new)
tl.send(dir="E", src=m_local) tl.copy_to(l_local, l_new)
tl.send(dir="E", src=l_local) tl.copy_to(O_local, O_new)
tl.send(dir="E", src=O_local) if east_exists:
elif col == _ROOT_COL: tl.send(dir="E", src=m_local)
with tl.scratch_scope(): tl.send(dir="E", src=l_local)
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16") tl.send(dir="E", src=O_local)
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16") elif col == root_col:
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16") if west_exists:
m_new, l_new, O_new = _merge_running( with tl.scratch_scope():
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
) l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
tl.copy_to(m_local, m_new) O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
tl.copy_to(l_local, l_new) m_new, l_new, O_new = _merge_running(
tl.copy_to(O_local, O_new) m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
with tl.scratch_scope(): )
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16") tl.copy_to(m_local, m_new)
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16") tl.copy_to(l_local, l_new)
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16") tl.copy_to(O_local, O_new)
m_new, l_new, O_new = _merge_running( if east_exists:
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, with tl.scratch_scope():
) m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
tl.copy_to(m_local, m_new) l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
tl.copy_to(l_local, l_new) O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
tl.copy_to(O_local, O_new) m_new, l_new, O_new = _merge_running(
elif _ROOT_COL < col < _SUB_W - 1: m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
with tl.scratch_scope(): )
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16") tl.copy_to(m_local, m_new)
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16") tl.copy_to(l_local, l_new)
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16") tl.copy_to(O_local, O_new)
m_new, l_new, O_new = _merge_running( elif root_col < col < sub_w - 1:
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, if east_exists:
) with tl.scratch_scope():
tl.copy_to(m_local, m_new) m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
tl.copy_to(l_local, l_new) l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
tl.copy_to(O_local, O_new) O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
tl.send(dir="W", src=m_local) m_new, l_new, O_new = _merge_running(
tl.send(dir="W", src=l_local) m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
tl.send(dir="W", src=O_local) )
elif col == _SUB_W - 1: tl.copy_to(m_local, m_new)
tl.send(dir="W", src=m_local) tl.copy_to(l_local, l_new)
tl.send(dir="W", src=l_local) tl.copy_to(O_local, O_new)
tl.send(dir="W", src=O_local) 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. # Phase 2: col reduce on col == root_col — converge at row == root_row.
if col == _ROOT_COL: if col == root_col:
if row == 0: if row == 0:
tl.send(dir="S", src=m_local) if south_exists:
tl.send(dir="S", src=l_local) tl.send(dir="S", src=m_local)
tl.send(dir="S", src=O_local) tl.send(dir="S", src=l_local)
elif 0 < row < _ROOT_ROW: tl.send(dir="S", src=O_local)
with tl.scratch_scope(): elif 0 < row < root_row:
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16") if north_exists:
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16") with tl.scratch_scope():
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16") m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running( l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
) m_new, l_new, O_new = _merge_running(
tl.copy_to(m_local, m_new) m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
tl.copy_to(l_local, l_new) )
tl.copy_to(O_local, O_new) tl.copy_to(m_local, m_new)
tl.send(dir="S", src=m_local) tl.copy_to(l_local, l_new)
tl.send(dir="S", src=l_local) tl.copy_to(O_local, O_new)
tl.send(dir="S", src=O_local) if south_exists:
elif row == _ROOT_ROW: tl.send(dir="S", src=m_local)
with tl.scratch_scope(): tl.send(dir="S", src=l_local)
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16") tl.send(dir="S", src=O_local)
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16") elif row == root_row:
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16") if north_exists:
m_new, l_new, O_new = _merge_running( with tl.scratch_scope():
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
) l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
tl.copy_to(m_local, m_new) O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
tl.copy_to(l_local, l_new) m_new, l_new, O_new = _merge_running(
tl.copy_to(O_local, O_new) m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
if _SUB_H - 1 > _ROOT_ROW: )
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(): with tl.scratch_scope():
m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16") m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="S", shape=l_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(m_local, m_new)
tl.copy_to(l_local, l_new) tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new) tl.copy_to(O_local, O_new)
elif _ROOT_ROW < row < _SUB_H - 1: elif root_row < row < sub_h - 1:
with tl.scratch_scope(): if south_exists:
m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16") with tl.scratch_scope():
l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16") m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16") l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running( O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, 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(m_local, m_new)
tl.copy_to(O_local, O_new) tl.copy_to(l_local, l_new)
tl.send(dir="N", src=m_local) tl.copy_to(O_local, O_new)
tl.send(dir="N", src=l_local) if north_exists:
tl.send(dir="N", src=O_local) tl.send(dir="N", src=m_local)
elif row == _SUB_H - 1 and _SUB_H - 1 > _ROOT_ROW: tl.send(dir="N", src=l_local)
tl.send(dir="N", src=m_local) tl.send(dir="N", src=O_local)
tl.send(dir="N", src=l_local) elif row == sub_h - 1 and sub_h - 1 > root_row:
tl.send(dir="N", src=O_local) if north_exists:
tl.send(dir="N", src=m_local)
tl.send(dir="N", src=l_local)
tl.send(dir="N", src=O_local)
@@ -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()