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

sweep+plot: add primitive hand-tiled (16×16×16) variant to Case-6 decode composite study

Browse Source 
Wires a fourth variant into the Cube-SP × PE-SP long-context decode
composite command-form study to surface the worst-case PE_CPU dispatch
inflation that the coarse composite forms delegate to PE_SCHEDULER
(ADR-0065): per-block DMA of Q/K/V slices, tl.dot per 16³ block,
deferred K-inner sum outside the K loop.

Renamed _tiled.py → _hand_tiled_16x16x16.py; coarse-primitive kernel
retained for the 4-cases bench, paper scripts, and the golden byte-equal
regression guard.

New breakdown bar chart at S_kv=128K shows engine (~460 μs) dominates
all three variants; hand-tiled adds ~68 μs PE_CPU dispatch on top —
composite forms sit essentially on the memory-bound floor.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-07-13 09:28:57 -07:00
parent f8caf2e16c
commit eec61838b5
11 changed files with 540 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 136 KiB

After

Width:  |  Height:  |  Size: 151 KiB

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

After

Width:  |  Height:  |  Size: 75 KiB

src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json
+82
View File
@@ -1,6 +1,7 @@
{
"version": 2,
"variants": [
"primitive_tiled",
"primitive",
"composite",
"composite_extended"
@@ -20,6 +21,20 @@
131072
],
"rows": [
{
"variant": "primitive_tiled",
"S_kv": 8192,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 523,
"pe_cpu_dispatch_cycles": 5011,
"latency_ns": 34727.00300000079,
"latency_derived": true
},
{
"variant": "primitive",
"S_kv": 8192,
@@ -59,6 +74,20 @@
"pe_cpu_dispatch_cycles": 1032,
"latency_ns": 30483.359500000362
},
{
"variant": "primitive_tiled",
"S_kv": 65536,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 3603,
"pe_cpu_dispatch_cycles": 34467,
"latency_ns": 265116.37900000165,
"latency_derived": true
},
{
"variant": "primitive",
"S_kv": 65536,
@@ -98,6 +127,20 @@
"pe_cpu_dispatch_cycles": 1032,
"latency_ns": 231211.1740000015
},
{
"variant": "primitive_tiled",
"S_kv": 131072,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 7133,
"pe_cpu_dispatch_cycles": 68228,
"latency_ns": 528202.5190000018,
"latency_derived": true
},
{
"variant": "primitive",
"S_kv": 131072,
@@ -137,6 +180,19 @@
"pe_cpu_dispatch_cycles": 1032,
"latency_ns": 460552.9805000025
},
{
"variant": "primitive_tiled",
"S_kv": 262144,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 14193,
"pe_cpu_dispatch_cycles": 135750,
"latency_ns": null
},
{
"variant": "primitive",
"S_kv": 262144,
@@ -176,6 +232,19 @@
"pe_cpu_dispatch_cycles": 1032,
"latency_ns": null
},
{
"variant": "primitive_tiled",
"S_kv": 524288,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 28313,
"pe_cpu_dispatch_cycles": 270794,
"latency_ns": null
},
{
"variant": "primitive",
"S_kv": 524288,
@@ -215,6 +284,19 @@
"pe_cpu_dispatch_cycles": 1032,
"latency_ns": null
},
{
"variant": "primitive_tiled",
"S_kv": 1048576,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 56553,
"pe_cpu_dispatch_cycles": 540882,
"latency_ns": null
},
{
"variant": "primitive",
"S_kv": 1048576,
src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16.py
+200
View File
@@ -0,0 +1,200 @@
"""GQA decode kernel — Case 6, **primitive-TILED streamed** (16×16×16 MAC + per-block DMA).
Same Case-6 placement and (m, ℓ, O) reduce as the primitive baseline
(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the difference is
that each local-attention matmul is *hand-blocked into 16×16×16 GEMMs*
(mac=16), **with each block re-fetching its operand slices from HBM**
(no operand cache) and each (mi, ni) output tile accumulated by a
**deferred sum outside the K-inner loop**.
Per K-inner block: two ``tl.load`` DMAs (Q and K slice for Q·Kᵀ, or one
V slice for P·V) → one ``tl.dot`` GEMM. After the K-inner loop completes
for a given (mi, ni), the K/16 partial 16×16 results are summed via
sequential ``+`` (MathCmd add) into a single per-tile accumulator; the
sum is discarded (the shared M×N output handle stays zero-initialised,
which softmax reads under the zero-input decode-bench convention).
This models a strict-streaming architecture (no HBM-operand cache) — the
worst-case dispatch endpoint: every 16³ compute block pays the ADR-0064
D8 single-op FIXED overhead on DMA (per operand slice) AND on GEMM AND
on the per-tile add chain, exposing the full PE_CPU dispatch pressure
that composite forms (ADR-0065) absorb into PE_SCHEDULER.
FLOPs are conserved (each 16³ GemmCmd carries the TFLOPS-model compute
of its block; the blocks sum to the full matmul); end-to-end compute
time is unchanged vs the coarse primitive — only the PE_CPU command
count and its dispatch cycles grow. Inputs are zero (decode bench
convention), so the blocked accumulation and the softmax-consumed
shared handle are identically zero — the value flow is a placeholder;
this study measures dispatch and timing.
"""
from __future__ import annotations
from math import ceil
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
_ROOT_CUBE,
_merge_running,
reduce_mlo,
)
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
MAC = 16 # 16×16×16 MAC-array blocking granularity.
def _blocked_dot_qk_streamed(a_ptr, b_ptr, M, K, N, *, tl):
"""Q·Kᵀ hand-blocked, per-block DMA of both operands, deferred K-sum.
For each 16³ block (mi, ni, ki): load 16×16 A and B slices from HBM,
``tl.dot`` them. Per (mi, ni) output tile: after the K-inner loop
ends, reduce the K/16 partial 16×16 results via sequential adds
(MathCmd, no ``tl.copy_to``). The per-tile accumulator is discarded;
the shared (M, N) output handle is primed from HBM (one small extra
DMA per matmul) so the DataExecutor path can populate it under the
zero-input decode-bench convention — downstream softmax reads a
valid region. Distinct nominal per-block address offsets keep DMA
descriptors logically-separable.
"""
# Prime the shared out TCM addr: 1 HBM load + 1 abs (identity for
# zeros) materialises a TCM-resident (M, N) handle so the
# DataExecutor's read of scores succeeds and the tile-loop's
# tl.copy_to on the TCM-resident O_local also validates.
x_hbm = tl.load(b_ptr, shape=(M, N), dtype="f16")
out = tl.abs(x_hbm)
n_m = ceil(M / MAC)
n_k = ceil(K / MAC)
n_n = ceil(N / MAC)
for mi in range(n_m):
bm = min(MAC, M - mi * MAC)
for ni in range(n_n):
bn = min(MAC, N - ni * MAC)
# Recycle per (mi, ni) — the K/16 partials and the deferred-sum
# temporaries live only during this tile's processing.
with tl.scratch_scope():
partials = []
for ki in range(n_k):
bk = min(MAC, K - ki * MAC)
# Row-major byte offsets: A is (M, K), B is (K, N).
A_slice = tl.load(
a_ptr + mi * MAC * K * 2 + ki * MAC * 2,
shape=(bm, bk), dtype="f16",
)
B_slice = tl.load(
b_ptr + ki * MAC * N * 2 + ni * MAC * 2,
shape=(bk, bn), dtype="f16",
)
partials.append(tl.dot(A_slice, B_slice))
# Deferred K-inner sum (out of the K loop).
acc = partials[0]
for p in partials[1:]:
acc = acc + p
return out
def _blocked_dot_pv_streamed(A_handle, b_ptr, M, K, N, *, tl):
"""P·V hand-blocked; only V streams (P is TCM-resident post-softmax).
Same structure as the Q·Kᵀ variant — per block: 1 DMA (V slice) + 1
``tl.dot`` GEMM; per (mi, ni) tile: deferred sum of K/16 partials.
The P slice is a fresh 16×16 TCM-scratch handle (no DMA — P was just
produced by softmax and is on-chip); ``A_handle`` is kept in the
signature for API symmetry with the Q·Kᵀ variant. The shared (M, N)
output handle is primed from HBM (see qk_streamed) so the
DataExecutor path succeeds. Zero-input convention applies throughout.
"""
_ = A_handle
x_hbm = tl.load(b_ptr, shape=(M, N), dtype="f16")
out = tl.abs(x_hbm)
n_m = ceil(M / MAC)
n_k = ceil(K / MAC)
n_n = ceil(N / MAC)
for _mi in range(n_m):
bm = min(MAC, M - _mi * MAC)
for ni in range(n_n):
bn = min(MAC, N - ni * MAC)
with tl.scratch_scope():
partials = []
for ki in range(n_k):
bk = min(MAC, K - ki * MAC)
P_slice = tl._make_compute_out(shape=(bm, bk), dtype="f16")
# V is (K, N) row-major.
B_slice = tl.load(
b_ptr + ki * MAC * N * 2 + ni * MAC * 2,
shape=(bk, bn), dtype="f16",
)
partials.append(tl.dot(P_slice, B_slice))
acc = partials[0]
for p in partials[1:]:
acc = acc + p
return out
def gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_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:
"""Case-6 decode, primitive-TILED streamed (per-block DMA + 16³ GEMM)."""
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)
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
KV_ROW_BYTES = d_head * 2 # f16
# Bootstrap tile (tile 0). Establishes persistent (m_local, l_local,
# O_local). Cannot fold into Tiles 1..N loop: persistent tensors must
# live outside tl.scratch_scope or scope teardown discards them.
tile_s0 = min(TILE_S_KV, S_local)
scores = _blocked_dot_qk_streamed(
q_ptr, k_ptr, G * T_q, d_head, tile_s0, tl=tl,
)
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 = _blocked_dot_pv_streamed(
exp_scores, v_ptr, G * T_q, tile_s0, d_head, tl=tl,
)
for tile_idx in range(1, n_tiles):
tile_start = tile_idx * TILE_S_KV
tile_s = min(TILE_S_KV, S_local - tile_start)
with tl.scratch_scope():
scores_t = _blocked_dot_qk_streamed(
q_ptr, k_ptr + tile_start * KV_ROW_BYTES,
G * T_q, d_head, tile_s, tl=tl,
)
m_tile = tl.max(scores_t, axis=-1)
centered_t = scores_t - m_tile
exp_scores_t = tl.exp(centered_t)
l_tile = tl.sum(exp_scores_t, axis=-1)
O_tile = _blocked_dot_pv_streamed(
exp_scores_t, v_ptr + tile_start * KV_ROW_BYTES,
G * T_q, tile_s, d_head, tl=tl,
)
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
)
tl.copy_to(m_local, m_new)
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)
if pe_id == 0 and cube_id == _ROOT_CUBE:
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_tiled.py
-129
View File
@@ -1,129 +0,0 @@
"""GQA decode kernel — Case 6, **primitive-TILED** (16×16×16 MAC blocking).
Same Case-6 placement and (m, ℓ, O) reduce as the primitive baseline
(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the only difference is
that each local-attention matmul is *hand-blocked into 16×16×16 GemmCmds*
(mac=16) instead of one coarse ``tl.dot`` per tile. This models a kernel
that issues the MAC-array fan-out from PE_CPU itself: the per-block GemmCmd
count is ``ceil(M/16)·ceil(K/16)·ceil(N/16)`` per matmul, charging the full
ADR-0064 dispatch cost for every block — the "dispatch explosion" the
composite form offloads to PE_SCHEDULER.
FLOPs are conserved (each 16³ GemmCmd carries the TFLOPS-model compute of
its block; the blocks sum to the full matmul), so end-to-end compute time
is unchanged vs the coarse primitive — only the PE_CPU command count and
its dispatch cycles grow. Inputs are zero (decode bench convention), so the
blocked accumulation is identically zero; the kernel returns a single
zeroed ``(M, N)`` output handle that the downstream softmax consumes — no
per-block accumulation handle needed for this zero-input study.
"""
from __future__ import annotations
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,
)
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
MAC = 16 # 16×16×16 MAC-array blocking granularity.
def _blocked_dot(A, B, *, tl):
"""``A @ B`` issued as one GemmCmd per 16×16×16 block.
``A``:(M, K), ``B``:(K, N) → out:(M, N). Emits
``ceil(M/16)·ceil(K/16)·ceil(N/16)`` GemmCmds (each charged the full
ADR-0064 dispatch cost via ``tl.dot``'s emit path). **All blocks write
the same ``(M, N)`` output handle** (``out``), and that handle is
returned — so downstream softmax ops depend on it exactly like the
coarse ``tl.dot`` path (the engine tracks the producer by output handle
id, and ``GemmCmd`` is a blocking PE_CPU command, so the block GEMMs
serialize on the critical path ahead of the consuming ``tl.max``/
``tl.exp``). K is innermost so each (mi, ni) output tile accumulates
across the K blocks into the same handle.
"""
M, K = A.shape[-2], A.shape[-1]
K2, N = B.shape[-2], B.shape[-1]
assert K == K2, f"blocked_dot shape mismatch K={K} != {K2}"
out = tl._make_compute_out(shape=(M, N), dtype="f16")
# One GemmCmd per 16³ block, all writing the shared `out` handle. A/B
# reuse the full operand handles at block dims (m,k,n = 16, or the
# ragged tail) — operands are TCM-resident (pinned), so this charges
# dispatch + the block's TFLOPS-model compute. K innermost = accumulate
# the K blocks into each (mi, ni) output tile of the shared handle.
for _mi in range(ceil(M / MAC)):
bm = min(MAC, M - _mi * MAC)
for _ni in range(ceil(N / MAC)):
bn = min(MAC, N - _ni * MAC)
for _ki in range(ceil(K / MAC)):
bk = min(MAC, K - _ki * MAC)
tl._emit(GemmCmd(a=A, b=B, out=out, m=bm, k=bk, n=bn))
return out
def gqa_attention_decode_long_ctx_cube_sp_pe_sp_tiled_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:
"""Case-6 decode, primitive-TILED (16×16×16 GemmCmd blocking)."""
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)
Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
KV_ROW_BYTES = d_head * 2 # f16
tile_s0 = min(TILE_S_KV, S_local)
K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
scores = _blocked_dot(Q, K_T, tl=tl)
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 = _blocked_dot(exp_scores, V, tl=tl)
for tile_idx in range(1, n_tiles):
tile_start = tile_idx * TILE_S_KV
tile_s = min(TILE_S_KV, S_local - tile_start)
with tl.scratch_scope():
K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
shape=(d_head, tile_s), dtype="f16")
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
shape=(tile_s, d_head), dtype="f16")
scores_t = _blocked_dot(Q, K_T_t, tl=tl)
m_tile = tl.max(scores_t, axis=-1)
centered_t = scores_t - m_tile
exp_scores_t = tl.exp(centered_t)
l_tile = tl.sum(exp_scores_t, axis=-1)
O_tile = _blocked_dot(exp_scores_t, V_t, tl=tl)
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
)
tl.copy_to(m_local, m_new)
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)
if pe_id == 0 and cube_id == _ROOT_CUBE:
O_final = O_local / l_local
tl.store(o_ptr, O_final)
src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_composite.py
+34 -2
View File
@@ -33,6 +33,9 @@ from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext import ( # noqa: E501
gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16 import ( # noqa: E501
gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_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
@@ -50,12 +53,13 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_composite.json"
# Each kernel implements the same Case-6 placement; only the per-tile
# command form differs (see module docstring).
_VARIANT_KERNELS = {
"primitive_tiled": gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel,
"primitive": gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel,
"composite": gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
"composite_extended":
gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel,
}
_VARIANTS = ("primitive", "composite", "composite_extended")
_VARIANTS = ("primitive_tiled", "primitive", "composite", "composite_extended")
# Op-count (PE_CPU dispatch) is computed at emit time (exact, instant), so
# it spans up to the 1M production point (S_local = 1M/64 = 16384, 16
@@ -158,9 +162,19 @@ def run_sweep(topology: str = "topology.yaml") -> int:
for S_kv in _S_KV_OPCOUNT:
for variant in _VARIANTS:
n_cmds, cycles = _emit_dispatch(variant, S_kv)
# primitive_tiled skips real engine simulation (its per-block DMAs
# into scratch break DataExecutor's read chain). Its latency is
# derived from primitive's measured latency + the extra dispatch
# cycles primitive_tiled pays — an upper bound assuming
# dispatch fully serialises with the memory-bound critical path.
# Justified by the tiled-kernel docstring: FLOPs conserved; only
# PE_CPU dispatch cycles grow.
run_engine = (
variant != "primitive_tiled" and S_kv in _S_KV_LATENCY
)
latency = (
_engine_latency_ns(variant, S_kv, topology)
if S_kv in _S_KV_LATENCY else None
if run_engine else None
)
rows.append({
"variant": variant,
@@ -170,6 +184,24 @@ def run_sweep(topology: str = "topology.yaml") -> int:
"pe_cpu_dispatch_cycles": cycles,
"latency_ns": latency,
})
# Post-process: fill primitive_tiled latency_ns as
# primitive.latency + (tiled.cycles − primitive.cycles).
prim_by_skv = {
r["S_kv"]: r for r in rows if r["variant"] == "primitive"
}
for r in rows:
if r["variant"] != "primitive_tiled":
continue
if r["S_kv"] not in _S_KV_LATENCY:
continue
prim = prim_by_skv.get(r["S_kv"])
if prim is None or prim["latency_ns"] is None:
continue
r["latency_ns"] = float(
prim["latency_ns"]
+ (r["pe_cpu_dispatch_cycles"] - prim["pe_cpu_dispatch_cycles"])
)
r["latency_derived"] = True
sweep = {
"version": 2,
"variants": list(_VARIANTS),