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

gqa-decode-composite: measure primitive hand-tiled (16×16×16) latency

Browse Source 
Switch primitive hand-tiled Q·Kᵀ / P·V from a deferred-K-sum tl.dot
chain to per-block GemmCmd writes into a shared (M, N) out handle
(implicit MAC-side accumulation). Full-shape coarse Q / K_T / V loads
carry data-mode correctness; per-block DMAs remain for the streaming
architecture dispatch story. The kernel now runs in engine mode
end-to-end, so its 128K latency is measured instead of derived as
primitive + Δdispatch. Drop the coarse-primitive baseline from the
sweep and plots; keep the kernel file for reference.

At S_kv=128K: primitive hand-tiled = 959.5 µs vs composite = 460.7 µs
(2.08× from 5 235 vs 94 PE_CPU commands).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-07-13 21:49:57 -07:00
parent eec61838b5
commit 23a2bc3fb3
8 changed files with 95 additions and 195 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: 151 KiB

After

Width:  |  Height:  |  Size: 149 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 Swipe Overlay

Before

Width:  |  Height:  |  Size: 75 KiB

After

Width:  |  Height:  |  Size: 73 KiB

src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json
+15 -97
View File
@@ -2,7 +2,6 @@
"version": 2,
"variants": [
"primitive_tiled",
"primitive",
"composite",
"composite_extended"
],
@@ -30,23 +29,9 @@
"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,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 96,
"pe_cpu_dispatch_cycles": 930,
"latency_ns": 30646.00300000079
"pe_cpu_cmd_count": 414,
"pe_cpu_dispatch_cycles": 3918,
"latency_ns": 63285.01999999998
},
{
"variant": "composite",
@@ -83,23 +68,9 @@
"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,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 96,
"pe_cpu_dispatch_cycles": 930,
"latency_ns": 231579.37900000165
"pe_cpu_cmd_count": 2654,
"pe_cpu_dispatch_cycles": 24974,
"latency_ns": 491824.02999999997
},
{
"variant": "composite",
@@ -136,23 +107,9 @@
"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,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 118,
"pe_cpu_dispatch_cycles": 1145,
"latency_ns": 461119.51900000183
"pe_cpu_cmd_count": 5235,
"pe_cpu_dispatch_cycles": 49242,
"latency_ns": 959456.0549999998
},
{
"variant": "composite",
@@ -189,21 +146,8 @@
"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,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 162,
"pe_cpu_dispatch_cycles": 1575,
"pe_cpu_cmd_count": 10397,
"pe_cpu_dispatch_cycles": 97778,
"latency_ns": null
},
{
@@ -241,21 +185,8 @@
"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,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 250,
"pe_cpu_dispatch_cycles": 2435,
"pe_cpu_cmd_count": 20721,
"pe_cpu_dispatch_cycles": 194850,
"latency_ns": null
},
{
@@ -293,21 +224,8 @@
"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,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 426,
"pe_cpu_dispatch_cycles": 4155,
"pe_cpu_cmd_count": 41369,
"pe_cpu_dispatch_cycles": 388994,
"latency_ns": null
},
{
src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16.py
+71 -56
View File
@@ -1,37 +1,41 @@
"""GQA decode kernel — Case 6, **primitive-TILED streamed** (16×16×16 MAC + per-block DMA).
"""GQA decode kernel — Case 6, **primitive hand-tiled 16×16×16** (per-block DMA + MAC-side accumulate).
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**.
(no operand cache), and each (mi, ni) output tile updated **implicitly**
by successive K-inner ``GemmCmd`` writes into the shared (M, N) output —
mirroring a MAC-array accumulator register that latches across K-inner
cycles on real accelerators.
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).
V slice for P·V) → one ``GemmCmd`` emitted directly via ``tl._emit`` with
``out`` bound to the shared (M, N) handle and ``m/k/n`` overridden to
the 16³ block dims. All K-inner iterations at a given (mi, ni) write
into the same output tile; no compiler-emitted ``MathCmd`` accumulator
chain.
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.
D8 single-op FIXED overhead on DMA (per operand slice) AND on GEMM,
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.
convention), so the "overwrite instead of accumulate" is semantically
equivalent to sum (0 = 0 + 0), and ``GemmCmd`` writes populate the
shared ``out`` handle so the downstream softmax's strict ``MathCmd``
reads succeed — the kernel runs in engine (data) mode.
"""
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,
@@ -43,24 +47,31 @@ 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.
"""Q·Kᵀ hand-blocked, per-block DMA of both operands, MAC-side accumulate.
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.
emit a ``GemmCmd`` directly with ``out`` bound to the shared (M, N)
handle and ``m/k/n`` overridden to the 16³ block dims. All K-inner
iterations at a given (mi, ni) write into the same shared ``out`` tile
— accumulation is implicit (MAC-side latching, mirroring how real
accelerators feed a per-tile accumulator register instead of running
a compiler-emitted MathCmd add chain).
Under the zero-input decode-bench convention, the "overwrite instead
of accumulate" is semantically equivalent to sum (0 = 0 + 0), and
``GemmCmd`` writes populate ``out.addr`` in ``MemoryStore`` so the
downstream softmax's strict ``MathCmd`` reads succeed. 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)
# Full-shape operand handles for data-mode correctness. Under the
# simulator's DataExecutor, GemmCmd computes np.matmul over these
# shapes and writes an (M, N) result to out.addr — populating the
# shared handle with a valid region so the downstream softmax's
# strict MathCmd reads succeed. In engine timing, the m/k/n
# override on each GemmCmd charges only the 16³ block work.
A_full = tl.load(a_ptr, shape=(M, K), dtype="f16")
B_full = tl.load(b_ptr, shape=(K, N), dtype="f16")
out = tl._make_compute_out(shape=(M, N), dtype="f16")
n_m = ceil(M / MAC)
n_k = ceil(K / MAC)
n_n = ceil(N / MAC)
@@ -68,26 +79,29 @@ def _blocked_dot_qk_streamed(a_ptr, b_ptr, M, K, N, *, tl):
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.
# Recycle per (mi, ni) — the per-block slice DMAs live only
# during this tile's processing; the shared out survives.
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(
# Per-block DMAs model the streaming architecture
# (no HBM-operand cache). Row-major byte offsets:
# A is (M, K), B is (K, N). Nominal handles — the
# GemmCmd below uses the full-shape handles above
# for data-mode compute; these per-block loads pay
# their ADR-0064 dispatch cost as descriptor work.
_ = tl.load(
a_ptr + mi * MAC * K * 2 + ki * MAC * 2,
shape=(bm, bk), dtype="f16",
)
B_slice = tl.load(
_ = 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
tl._emit(GemmCmd(
a=A_full, b=B_full, out=out,
m=bm, k=bk, n=bn,
))
return out
@@ -95,38 +109,39 @@ 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.
``GemmCmd`` writing to the shared (M, N) ``out`` handle with implicit
K-inner accumulation via successive block writes.
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.
signature for API symmetry with the Q·Kᵀ variant. Zero-input
convention applies throughout.
"""
_ = A_handle
x_hbm = tl.load(b_ptr, shape=(M, N), dtype="f16")
out = tl.abs(x_hbm)
# A_handle (= exp_scores from softmax) is TCM-resident with
# shape (M, K), already populated by tl.exp's DataExecutor. Only V
# needs a full-shape coarse load from HBM for data-mode correctness.
B_full = tl.load(b_ptr, shape=(K, N), dtype="f16")
out = tl._make_compute_out(shape=(M, N), dtype="f16")
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(
# Per-block V load — streaming-architecture dispatch
# cost. GemmCmd below uses B_full for compute.
_ = 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
bm = min(MAC, M - _mi * MAC)
tl._emit(GemmCmd(
a=A_handle, b=B_full, out=out,
m=bm, k=bk, n=bn,
))
return out
src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_composite.py
+2 -34
View File
@@ -24,9 +24,6 @@ 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 import ( # noqa: E501
gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel,
)
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,
)
@@ -54,12 +51,11 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_composite.json"
# 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_tiled", "primitive", "composite", "composite_extended")
_VARIANTS = ("primitive_tiled", "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
@@ -162,19 +158,9 @@ 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 run_engine else None
if S_kv in _S_KV_LATENCY else None
)
rows.append({
"variant": variant,
@@ -184,24 +170,6 @@ 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),