mukesh
23a2bc3fb3
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>
88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
"""Command-count checks for the Case-6 primitive-TILED (streamed) decode kernel.
|
||
|
||
The tiled kernel hand-blocks each Q·Kᵀ / P·V matmul into 16×16×16
|
||
GEMMs, with per-block DMAs of the HBM operand slices and a deferred
|
||
K-inner sum outside the K loop (ADR-0064 Rev2 D8 single-op FIXED path
|
||
exercised across DMA, GEMM, and MATH engines).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from math import ceil
|
||
|
||
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 as _primitive,
|
||
)
|
||
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 as _tiled,
|
||
)
|
||
from kernbench.common.pe_commands import (
|
||
DmaReadCmd, GemmCmd, PeCpuOverheadCmd,
|
||
)
|
||
from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL
|
||
from kernbench.triton_emu.tl_context import TLContext, run_kernel
|
||
|
||
_C, _P = 8, 8
|
||
_S_KV = 131_072
|
||
_D_HEAD, _H_Q, _H_KV, _T_Q = 128, 8, 1, 1
|
||
_TILE_S_KV = 1024
|
||
_MAC = 16
|
||
|
||
|
||
def _run(kernel, S_kv: int):
|
||
tl = TLContext(
|
||
pe_id=0, num_programs=_P, cost_model=DEFAULT_PE_COST_MODEL,
|
||
cube_id=6, num_cubes=_C, scratch_base=0x200000, scratch_size=1 << 20,
|
||
)
|
||
run_kernel(
|
||
kernel, tl, 0x1000, 0x2000, 0x3000, 0x4000,
|
||
_T_Q, S_kv, _H_Q, _H_KV, _D_HEAD, _C, _P,
|
||
)
|
||
return tl.commands
|
||
|
||
|
||
def _block_counts():
|
||
"""Return (n_tiles, qk_blocks, pv_blocks) for the fixed test params."""
|
||
G = _H_Q // _H_KV
|
||
S_local = _S_KV // (_C * _P)
|
||
n_tiles = (S_local + _TILE_S_KV - 1) // _TILE_S_KV
|
||
tile_s = min(_TILE_S_KV, S_local)
|
||
M = G * _T_Q
|
||
qk = ceil(M / _MAC) * ceil(tile_s / _MAC) * ceil(_D_HEAD / _MAC)
|
||
pv = ceil(M / _MAC) * ceil(_D_HEAD / _MAC) * ceil(tile_s / _MAC)
|
||
return n_tiles, qk, pv
|
||
|
||
|
||
def test_tiled_gemm_count_matches_block_formula():
|
||
"""One GemmCmd per 16³ block across both matmuls."""
|
||
n_tiles, qk, pv = _block_counts()
|
||
expected = n_tiles * (qk + pv)
|
||
|
||
cmds = _run(_tiled, _S_KV)
|
||
got = sum(1 for c in cmds if isinstance(c, GemmCmd))
|
||
assert got == expected, f"GemmCmd count: got {got}, expected {expected}"
|
||
|
||
|
||
def test_tiled_dma_count_matches_streamed_formula():
|
||
"""Streamed operand DMAs — Q·Kᵀ blocks pay 2 DMAs (A+B slice), P·V
|
||
blocks pay 1 DMA (V slice only; the P side is TCM-resident post-softmax).
|
||
Each matmul also carries full-shape coarse loads for the GemmCmd
|
||
operand handles (2 for Q·Kᵀ: Q + K_T, 1 for P·V: V only)."""
|
||
n_tiles, qk, pv = _block_counts()
|
||
expected = n_tiles * (2 * qk + 1 * pv + 3)
|
||
|
||
cmds = _run(_tiled, _S_KV)
|
||
got = sum(1 for c in cmds if isinstance(c, DmaReadCmd))
|
||
assert got == expected, f"DmaReadCmd count: got {got}, expected {expected}"
|
||
|
||
|
||
def test_tiled_amplifies_dispatch_vs_primitive():
|
||
"""The streamed tiled kernel emits ≫40× more PE_CPU dispatches than
|
||
the coarse primitive — the whole point of adding this variant."""
|
||
n_tiled = sum(1 for c in _run(_tiled, _S_KV)
|
||
if isinstance(c, PeCpuOverheadCmd))
|
||
n_prim = sum(1 for c in _run(_primitive, _S_KV)
|
||
if isinstance(c, PeCpuOverheadCmd))
|
||
assert n_tiled > 40 * n_prim, (
|
||
f"expected >40× dispatch amplification, got {n_tiled} vs {n_prim}"
|
||
)
|