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

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>
2026-07-13 09:28:57 -07:00
parent f8caf2e16c
commit eec61838b5
11 changed files with 540 additions and 135 deletions
@@ -0,0 +1,88 @@
+"""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}"
+    )