mukesh
eec61838b5
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>
89 lines
3.3 KiB
Python
89 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). Each
|
||
_blocked_dot_* call also emits one small prime DMA at the top so the
|
||
DataExecutor's read of the shared M×N output succeeds — 2 primes per
|
||
outer S_kv tile (one QK + one PV)."""
|
||
n_tiles, qk, pv = _block_counts()
|
||
expected = n_tiles * (2 * qk + 1 * pv + 2)
|
||
|
||
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 ≫50× 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 > 50 * n_prim, (
|
||
f"expected >50× dispatch amplification, got {n_tiled} vs {n_prim}"
|
||
)
|