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

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>

src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json

@@ -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_cmd_count": 414,
-      "pe_cpu_dispatch_cycles": 5011,
+      "pe_cpu_dispatch_cycles": 3918,
-      "latency_ns": 34727.00300000079,
+      "latency_ns": 63285.01999999998
-      "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
     },
     {
       "variant": "composite",
@@ -83,23 +68,9 @@
       "d_head": 128,
       "h_q": 8,
       "h_kv": 1,
-      "pe_cpu_cmd_count": 3603,
+      "pe_cpu_cmd_count": 2654,
-      "pe_cpu_dispatch_cycles": 34467,
+      "pe_cpu_dispatch_cycles": 24974,
-      "latency_ns": 265116.37900000165,
+      "latency_ns": 491824.02999999997
-      "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
     },
     {
       "variant": "composite",
@@ -136,23 +107,9 @@
       "d_head": 128,
       "h_q": 8,
       "h_kv": 1,
-      "pe_cpu_cmd_count": 7133,
+      "pe_cpu_cmd_count": 5235,
-      "pe_cpu_dispatch_cycles": 68228,
+      "pe_cpu_dispatch_cycles": 49242,
-      "latency_ns": 528202.5190000018,
+      "latency_ns": 959456.0549999998
-      "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
     },
     {
       "variant": "composite",
@@ -189,21 +146,8 @@
       "d_head": 128,
       "h_q": 8,
       "h_kv": 1,
-      "pe_cpu_cmd_count": 14193,
+      "pe_cpu_cmd_count": 10397,
-      "pe_cpu_dispatch_cycles": 135750,
+      "pe_cpu_dispatch_cycles": 97778,
-      "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,
       "latency_ns": null
     },
     {
@@ -241,21 +185,8 @@
       "d_head": 128,
       "h_q": 8,
       "h_kv": 1,
-      "pe_cpu_cmd_count": 28313,
+      "pe_cpu_cmd_count": 20721,
-      "pe_cpu_dispatch_cycles": 270794,
+      "pe_cpu_dispatch_cycles": 194850,
-      "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,
       "latency_ns": null
     },
     {
@@ -293,21 +224,8 @@
       "d_head": 128,
       "h_q": 8,
       "h_kv": 1,
-      "pe_cpu_cmd_count": 56553,
+      "pe_cpu_cmd_count": 41369,
-      "pe_cpu_dispatch_cycles": 540882,
+      "pe_cpu_dispatch_cycles": 388994,
-      "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,
       "latency_ns": null
     },
     {

src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16.py

@@ -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
+(no operand cache), and each (mi, ni) output tile updated **implicitly**
-**deferred sum outside the K-inner loop**.
+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
+V slice for P·V) → one ``GemmCmd`` emitted directly via ``tl._emit`` with
-for a given (mi, ni), the K/16 partial 16×16 results are summed via
+``out`` bound to the shared (M, N) handle and ``m/k/n`` overridden to
-sequential ``+`` (MathCmd add) into a single per-tile accumulator; the
+the 16³ block dims. All K-inner iterations at a given (mi, ni) write
-sum is discarded (the shared M×N output handle stays zero-initialised,
+into the same output tile; no compiler-emitted ``MathCmd`` accumulator
-which softmax reads under the zero-input decode-bench convention).
+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
+D8 single-op FIXED overhead on DMA (per operand slice) AND on GEMM,
-on the per-tile add chain, exposing the full PE_CPU dispatch pressure
+exposing the full PE_CPU dispatch pressure that composite forms
-that composite forms (ADR-0065) absorb into PE_SCHEDULER.
+(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
+convention), so the "overwrite instead of accumulate" is semantically
-shared handle are identically zero — the value flow is a placeholder;
+equivalent to sum (0 = 0 + 0), and ``GemmCmd`` writes populate the
-this study measures dispatch and timing.
+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
+    emit a ``GemmCmd`` directly with ``out`` bound to the shared (M, N)
-    ends, reduce the K/16 partial 16×16 results via sequential adds
+    handle and ``m/k/n`` overridden to the 16³ block dims. All K-inner
-    (MathCmd, no ``tl.copy_to``). The per-tile accumulator is discarded;
+    iterations at a given (mi, ni) write into the same shared ``out`` tile
-    the shared (M, N) output handle is primed from HBM (one small extra
+    — accumulation is implicit (MAC-side latching, mirroring how real
-    DMA per matmul) so the DataExecutor path can populate it under the
+    accelerators feed a per-tile accumulator register instead of running
-    zero-input decode-bench convention — downstream softmax reads a
+    a compiler-emitted MathCmd add chain).
-    valid region. Distinct nominal per-block address offsets keep DMA
+
-    descriptors logically-separable.
+    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
+    # Full-shape operand handles for data-mode correctness. Under the
-    # zeros) materialises a TCM-resident (M, N) handle so the
+    # simulator's DataExecutor, GemmCmd computes np.matmul over these
-    # DataExecutor's read of scores succeeds and the tile-loop's
+    # shapes and writes an (M, N) result to out.addr — populating the
-    # tl.copy_to on the TCM-resident O_local also validates.
+    # shared handle with a valid region so the downstream softmax's
-    x_hbm = tl.load(b_ptr, shape=(M, N), dtype="f16")
+    # strict MathCmd reads succeed. In engine timing, the m/k/n
-    out = tl.abs(x_hbm)
+    # 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
+            # Recycle per (mi, ni) — the per-block slice DMAs live only
-            # temporaries live only during this tile's processing.
+            # 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).
+                    # Per-block DMAs model the streaming architecture
-                    A_slice = tl.load(
+                    # (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))
+                    tl._emit(GemmCmd(
-                # Deferred K-inner sum (out of the K loop).
+                        a=A_full, b=B_full, out=out,
-                acc = partials[0]
+                        m=bm, k=bk, n=bn,
-                for p in partials[1:]:
+                    ))
-                    acc = acc + p
     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)
+    signature for API symmetry with the Q·Kᵀ variant. Zero-input
-    output handle is primed from HBM (see qk_streamed) so the
+    convention applies throughout.
-    DataExecutor path succeeds. Zero-input convention applies throughout.
     """
-    _ = A_handle
+    # A_handle (= exp_scores from softmax) is TCM-resident with
-    x_hbm = tl.load(b_ptr, shape=(M, N), dtype="f16")
+    # shape (M, K), already populated by tl.exp's DataExecutor. Only V
-    out = tl.abs(x_hbm)
+    # 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")
+                    # Per-block V load — streaming-architecture dispatch
-                    # V is (K, N) row-major.
+                    # cost. GemmCmd below uses B_full for compute.
-                    B_slice = tl.load(
+                    _ = tl.load(
                         b_ptr + ki * MAC * N * 2 + ni * MAC * 2,
                         shape=(bk, bn), dtype="f16",
                     )
-                    partials.append(tl.dot(P_slice, B_slice))
+                    bm = min(MAC, M - _mi * MAC)
-                acc = partials[0]
+                    tl._emit(GemmCmd(
-                for p in partials[1:]:
+                        a=A_handle, b=B_full, out=out,
-                    acc = acc + p
+                        m=bm, k=bk, n=bn,
+                    ))
     return out
 
 

src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_composite.py

@@ -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),

tests/attention/test_gqa_decode_long_ctx_tiled.py

@@ -64,12 +64,11 @@ def test_tiled_gemm_count_matches_block_formula():
 
 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
+    blocks pay 1 DMA (V slice only; the P side is TCM-resident post-softmax).
-    _blocked_dot_* call also emits one small prime DMA at the top so the
+    Each matmul also carries full-shape coarse loads for the GemmCmd
-    DataExecutor's read of the shared M×N output succeeds — 2 primes per
+    operand handles (2 for Q·Kᵀ: Q + K_T, 1 for P·V: V only)."""
-    outer S_kv tile (one QK + one PV)."""
     n_tiles, qk, pv = _block_counts()
-    expected = n_tiles * (2 * qk + 1 * pv + 2)
+    expected = n_tiles * (2 * qk + 1 * pv + 3)
 
     cmds = _run(_tiled, _S_KV)
     got = sum(1 for c in cmds if isinstance(c, DmaReadCmd))
@@ -77,12 +76,12 @@ def test_tiled_dma_count_matches_streamed_formula():
 
 
 def test_tiled_amplifies_dispatch_vs_primitive():
-    """The streamed tiled kernel emits ≫50× more PE_CPU dispatches than
+    """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 > 50 * n_prim, (
+    assert n_tiled > 40 * n_prim, (
-        f"expected >50× dispatch amplification, got {n_tiled} vs {n_prim}"
+        f"expected >40× dispatch amplification, got {n_tiled} vs {n_prim}"
     )