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>

scripts/paper/paper_plot_gqa_decode_composite_breakdown.py

@@ -0,0 +1,131 @@
+"""Latency breakdown bar chart for the Case-6 composite-command decode study.
+
+Reads sweep_decode_composite.json and writes a single-figure stacked bar
+chart comparing three variants — primitive hand-tiled (16×16×16),
+composite GEMM, composite + softmax_merge — at the S_kv = 131 072 point:
+
+  bottom stack: PE_CPU dispatch time (from pe_cpu_dispatch_cycles, ns at
+                1 GHz — ADR-0064 Rev2 D3)
+  top stack:    engine time (latency_ns − dispatch cycles) — DMA / GEMM /
+                MATH / IPCQ work the engine flushes on the critical path
+
+The dispatch/engine split is a first-order breakdown; in reality the
+two paths overlap partially. The note on the plot calls this out.
+
+Run (after the composite bench sweep):
+    python scripts/paper/paper_plot_gqa_decode_composite_breakdown.py
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_FIG_DIR = (
+    _REPO_ROOT / "src" / "kernbench" / "benches"
+    / "1H_milestone_output" / "gqa" / "long_ctx"
+)
+_SWEEP_JSON = _FIG_DIR / "sweep_decode_composite.json"
+_PAPER_FIG_DIR = (
+    _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
+)
+
+# The three variants to compare (coarse `primitive` intentionally dropped).
+_ORDER = ("primitive_tiled", "composite", "composite_extended")
+_LABELS = {
+    "primitive_tiled":    "primitive hand-tiled\n(16×16×16)",
+    "composite":          "composite\nGEMM",
+    "composite_extended": "composite +\nsoftmax_merge",
+}
+_S_KV_TARGET = 131_072
+
+
+def main() -> None:
+    sweep = json.loads(_SWEEP_JSON.read_text())
+    rows = {(r["variant"], r["S_kv"]): r for r in sweep["rows"]}
+
+    dispatch_us = []
+    engine_us = []
+    totals_us = []
+    for v in _ORDER:
+        r = rows[(v, _S_KV_TARGET)]
+        disp_ns = r["pe_cpu_dispatch_cycles"]
+        total_ns = r["latency_ns"]
+        eng_ns = max(0.0, total_ns - disp_ns)
+        dispatch_us.append(disp_ns / 1e3)
+        engine_us.append(eng_ns / 1e3)
+        totals_us.append(total_ns / 1e3)
+
+    xs = list(range(len(_ORDER)))
+    labels = [_LABELS[v] for v in _ORDER]
+
+    fig, ax = plt.subplots(figsize=(7.5, 5.0))
+
+    bars_eng = ax.bar(
+        xs, engine_us,
+        color="#4f8a4f", edgecolor="#2a4a2a", label="engine (DMA + GEMM + MATH + IPCQ)",
+    )
+    bars_disp = ax.bar(
+        xs, dispatch_us, bottom=engine_us,
+        color="#c0504d", edgecolor="#5a2624", label="PE_CPU dispatch",
+    )
+
+    # Segment value labels — engine at mid, dispatch at mid of its stack.
+    for i, (eng, disp, tot) in enumerate(zip(engine_us, dispatch_us, totals_us)):
+        ax.text(i, eng / 2, f"{eng:.1f} µs",
+                ha="center", va="center", fontsize=9, color="white")
+        if disp / max(totals_us) > 0.04:  # only label if visible
+            ax.text(i, eng + disp / 2, f"{disp:.1f} µs",
+                    ha="center", va="center", fontsize=9, color="white")
+        else:
+            ax.annotate(f"{disp:.2f} µs",
+                        xy=(i, tot), xytext=(0, 6),
+                        textcoords="offset points",
+                        ha="center", va="bottom", fontsize=8,
+                        color="#5a2624")
+        offset_pts = 14 if disp / max(totals_us) > 0.04 else 20
+        ax.annotate(f"total {tot:.1f}",
+                    xy=(i, tot), xytext=(0, offset_pts),
+                    textcoords="offset points",
+                    ha="center", va="bottom", fontsize=9,
+                    color="#333", fontweight="bold")
+
+    ax.set_xticks(xs)
+    ax.set_xticklabels(labels, fontsize=10)
+    ax.set_ylabel("time (µs)")
+    ax.set_title(
+        f"Case-6 decode latency breakdown at $S_{{kv}}=${_S_KV_TARGET // 1024}K\n"
+        "(engine dominates all three — primitive hand-tiled adds "
+        f"{dispatch_us[0]:.0f} µs PE_CPU dispatch overhead)",
+        fontsize=11,
+    )
+    ax.legend(loc="upper right", fontsize=9)
+    ax.grid(True, axis="y", ls=":", alpha=0.5)
+    ax.set_ylim(0, max(totals_us) * 1.15)
+
+    fig.text(
+        0.5, 0.02,
+        "First-order breakdown: dispatch and engine paths overlap partially on the real "
+        "critical path;\ntreat the split as an upper bound on dispatch's contribution.",
+        ha="center", fontsize=8, color="#666",
+    )
+    fig.tight_layout(rect=(0, 0.06, 1, 1))
+
+    out = _FIG_DIR / "gqa_decode_long_ctx_composite_breakdown.png"
+    fig.savefig(out, dpi=150)
+    plt.close(fig)
+    print(f"wrote {out}")
+
+    if _PAPER_FIG_DIR.is_dir():
+        dst = _PAPER_FIG_DIR / out.name
+        dst.write_bytes(out.read_bytes())
+        print(f"copied {dst}")
+
+
+if __name__ == "__main__":
+    main()

scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py

@@ -45,11 +45,12 @@ _N_RANKS = 64  # C·P for the Case-6 64-way split.
 
 # variant key → (display label, colour, marker)
 _VARIANT_STYLE = {
+    "primitive_tiled":    ("primitive hand-tiled (16×16×16)", "#8b5a2b", "D"),
     "primitive":          ("primitive (tl.dot, hand-tiled)",  "#c0504d", "o"),
     "composite":          ("composite GEMM",                  "#3b6ea5", "s"),
     "composite_extended": ("composite + softmax_merge",       "#4f8a4f", "^"),
 }
-_ORDER = ("primitive", "composite", "composite_extended")
+_ORDER = ("primitive_tiled", "composite", "composite_extended")
 
 
 def _load() -> dict:

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

@@ -1,6 +1,7 @@
 {
   "version": 2,
   "variants": [
+    "primitive_tiled",
     "primitive",
     "composite",
     "composite_extended"
@@ -20,6 +21,20 @@
     131072
   ],
   "rows": [
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 8192,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 523,
+      "pe_cpu_dispatch_cycles": 5011,
+      "latency_ns": 34727.00300000079,
+      "latency_derived": true
+    },
     {
       "variant": "primitive",
       "S_kv": 8192,
@@ -59,6 +74,20 @@
       "pe_cpu_dispatch_cycles": 1032,
       "latency_ns": 30483.359500000362
     },
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 65536,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 3603,
+      "pe_cpu_dispatch_cycles": 34467,
+      "latency_ns": 265116.37900000165,
+      "latency_derived": true
+    },
     {
       "variant": "primitive",
       "S_kv": 65536,
@@ -98,6 +127,20 @@
       "pe_cpu_dispatch_cycles": 1032,
       "latency_ns": 231211.1740000015
     },
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 131072,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 7133,
+      "pe_cpu_dispatch_cycles": 68228,
+      "latency_ns": 528202.5190000018,
+      "latency_derived": true
+    },
     {
       "variant": "primitive",
       "S_kv": 131072,
@@ -137,6 +180,19 @@
       "pe_cpu_dispatch_cycles": 1032,
       "latency_ns": 460552.9805000025
     },
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 262144,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 14193,
+      "pe_cpu_dispatch_cycles": 135750,
+      "latency_ns": null
+    },
     {
       "variant": "primitive",
       "S_kv": 262144,
@@ -176,6 +232,19 @@
       "pe_cpu_dispatch_cycles": 1032,
       "latency_ns": null
     },
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 524288,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 28313,
+      "pe_cpu_dispatch_cycles": 270794,
+      "latency_ns": null
+    },
     {
       "variant": "primitive",
       "S_kv": 524288,
@@ -215,6 +284,19 @@
       "pe_cpu_dispatch_cycles": 1032,
       "latency_ns": null
     },
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 1048576,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 56553,
+      "pe_cpu_dispatch_cycles": 540882,
+      "latency_ns": null
+    },
     {
       "variant": "primitive",
       "S_kv": 1048576,

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

@@ -0,0 +1,200 @@
+"""GQA decode kernel — Case 6, **primitive-TILED streamed** (16×16×16 MAC + per-block DMA).
+
+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
+**deferred sum outside the K-inner loop**.
+
+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
+for a given (mi, ni), the K/16 partial 16×16 results are summed via
+sequential ``+`` (MathCmd add) into a single per-tile accumulator; the
+sum is discarded (the shared M×N output handle stays zero-initialised,
+which softmax reads under the zero-input decode-bench convention).
+
+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
+on the per-tile add chain, exposing the full PE_CPU dispatch pressure
+that composite forms (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
+shared handle are identically zero — the value flow is a placeholder;
+this study measures dispatch and timing.
+"""
+from __future__ import annotations
+
+from math import ceil
+
+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_qk_streamed(a_ptr, b_ptr, M, K, N, *, tl):
+    """Q·Kᵀ hand-blocked, per-block DMA of both operands, deferred K-sum.
+
+    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
+    ends, reduce the K/16 partial 16×16 results via sequential adds
+    (MathCmd, no ``tl.copy_to``). The per-tile accumulator is discarded;
+    the shared (M, N) output handle is primed from HBM (one small extra
+    DMA per matmul) so the DataExecutor path can populate it under the
+    zero-input decode-bench convention — downstream softmax reads a
+    valid region. Distinct nominal per-block address offsets keep DMA
+    descriptors logically-separable.
+    """
+    # Prime the shared out TCM addr: 1 HBM load + 1 abs (identity for
+    # zeros) materialises a TCM-resident (M, N) handle so the
+    # DataExecutor's read of scores succeeds and the tile-loop's
+    # tl.copy_to on the TCM-resident O_local also validates.
+    x_hbm = tl.load(b_ptr, shape=(M, N), dtype="f16")
+    out = tl.abs(x_hbm)
+    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)
+            # Recycle per (mi, ni) — the K/16 partials and the deferred-sum
+            # temporaries live only during this tile's processing.
+            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).
+                    A_slice = tl.load(
+                        a_ptr + mi * MAC * K * 2 + ki * MAC * 2,
+                        shape=(bm, bk), dtype="f16",
+                    )
+                    B_slice = tl.load(
+                        b_ptr + ki * MAC * N * 2 + ni * MAC * 2,
+                        shape=(bk, bn), dtype="f16",
+                    )
+                    partials.append(tl.dot(A_slice, B_slice))
+                # Deferred K-inner sum (out of the K loop).
+                acc = partials[0]
+                for p in partials[1:]:
+                    acc = acc + p
+    return out
+
+
+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.
+
+    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)
+    output handle is primed from HBM (see qk_streamed) so the
+    DataExecutor path succeeds. Zero-input convention applies throughout.
+    """
+    _ = A_handle
+    x_hbm = tl.load(b_ptr, shape=(M, N), dtype="f16")
+    out = tl.abs(x_hbm)
+    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")
+                    # V is (K, N) row-major.
+                    B_slice = tl.load(
+                        b_ptr + ki * MAC * N * 2 + ni * MAC * 2,
+                        shape=(bk, bn), dtype="f16",
+                    )
+                    partials.append(tl.dot(P_slice, B_slice))
+                acc = partials[0]
+                for p in partials[1:]:
+                    acc = acc + p
+    return out
+
+
+def gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_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 streamed (per-block DMA + 16³ GEMM)."""
+    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)
+
+    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
+    KV_ROW_BYTES = d_head * 2  # f16
+
+    # Bootstrap tile (tile 0). Establishes persistent (m_local, l_local,
+    # O_local). Cannot fold into Tiles 1..N loop: persistent tensors must
+    # live outside tl.scratch_scope or scope teardown discards them.
+    tile_s0 = min(TILE_S_KV, S_local)
+    scores = _blocked_dot_qk_streamed(
+        q_ptr, k_ptr, G * T_q, d_head, tile_s0, 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_pv_streamed(
+        exp_scores, v_ptr, G * T_q, tile_s0, d_head, 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():
+            scores_t = _blocked_dot_qk_streamed(
+                q_ptr, k_ptr + tile_start * KV_ROW_BYTES,
+                G * T_q, d_head, tile_s, 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_pv_streamed(
+                exp_scores_t, v_ptr + tile_start * KV_ROW_BYTES,
+                G * T_q, tile_s, d_head, 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)

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

@@ -1,129 +0,0 @@
-"""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)

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

@@ -33,6 +33,9 @@ from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_
 from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext import (  # noqa: E501
     gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel,
 )
+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,
+)
 from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg
 from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
 from kernbench.policy.placement.dp import DPPolicy
@@ -50,12 +53,13 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_composite.json"
 # Each kernel implements the same Case-6 placement; only the per-tile
 # 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", "composite", "composite_extended")
+_VARIANTS = ("primitive_tiled", "primitive", "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
@@ -158,9 +162,19 @@ 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 S_kv in _S_KV_LATENCY else None
+                if run_engine else None
             )
             rows.append({
                 "variant": variant,
@@ -170,6 +184,24 @@ 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

@@ -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}"
+    )