experiment(gqa-decode): 16x16x16 MAC-block dispatch model
Research artifact (NOT wired into the production sweep) for the decode composite-vs-primitive investigation. Models a primitive decode kernel that hand-blocks each tl.dot into 16x16x16 GemmCmds, charging per-MAC-block PE_CPU dispatch -- testing whether faithfully charging the dispatch that the composite form offloads to PE_SCHEDULER flips the 'composite gives no decode latency benefit' result.
Finding: it does NOT. Even with up to 512 serialized blocking GemmCmds per matmul (vs 2 coarse tl.dot), end-to-end decode latency is unchanged (8192: 28.9us, 32768: 114.9us) -- the inter-cube (m,l,O) reduce DMA tail dominates and the local-attention GEMM/dispatch sits in critical-path slack. Confirms the two-regime conclusion (24c7054): composite helps compute-bound prefill, not memory/reduce-bound decode.
Caveats: measured at S_kv 8K/32K (reduce-dominated); large-S_kv streaming regime extrapolated only. The -5.5% tiled<untiled is reduce-tail ordering noise, not fully nailed down.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
+129
@@ -0,0 +1,129 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user