feat(gqa): Case-6 composite-command decode variants + on-chip-operand DMA fix

Add two command-form variants of the Case-6 (Cube-SP x PE-SP) long-context decode kernel alongside the primitive baseline: - composite: per-tile GEMMs issued as one coarse tl.composite over the full S_local (PE_SCHEDULER tiles/streams K,V) -- O(1) PE_CPU commands vs the primitive kernel's O(n_tiles). - composite_extended: Q.K^T composite + softmax_merge recipe composite (ADR-0065), folding the per-tile online merge + P.V. Extract the shared two-level (m,l,O) reduce into _gqa_mlo_reduce and refactor the baseline to use it (byte-equal, guarded by a 64-rank command-stream digest test). Engine fix: _make_compute_out omitted pinned=True, so an on-chip TCM compute result consumed as a composite GEMM operand was scheduled as an HBM DMA_READ of its bit-61 scratch address -> PhysAddrError in data mode. Mark compute outputs pinned (resident), matching the composite auto-output and recipe scratch handles. .pinned is read only by the tiling DMA gate, so this is byte-equal for existing benches. Add the composite sweep (emit-level PE_CPU dispatch to 1M + data-mode latency to 128K), its plot, umbrella wiring (GQA_1H_SWEEPS=composite), and tests (refactor guard, dispatch saturation, engine-fix, e2e). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-23 16:07:41 -07:00
parent e9a5c438e3
commit cd6f0ed91d
10 changed files with 977 additions and 190 deletions
@@ -36,26 +36,14 @@ Topology / SFR:
 """
 from __future__ import annotations

+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).

-# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh).
-_SUB_W = 4
-_SUB_H = 2
-_ROOT_COL = _SUB_W // 2   # 2
-_ROOT_ROW = _SUB_H // 2   # 1
-_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL  # 6
-
-
-def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
-    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
-    m_new = tl.maximum(m_local, m_other)
-    scale_old = tl.exp(m_local - m_new)
-    scale_new = tl.exp(m_other - m_new)
-    l_new = l_local * scale_old + l_other * scale_new
-    O_new = O_local * scale_old + O_other * scale_new
-    return m_new, l_new, O_new
-

 def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
    q_ptr: int,
@@ -119,175 +107,8 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
            tl.copy_to(l_local, l_new)
            tl.copy_to(O_local, O_new)

-    # ── Intra-CUBE reduce: row chain (intra_W) + col bridge (intra_N) ──
-    PE_GRID_COLS = 4
-    pe_col = pe_id % PE_GRID_COLS
-    pe_row = pe_id // PE_GRID_COLS
-    pe_cols_used = min(PE_GRID_COLS, P)
-    pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
-
-    if pe_cols_used > 1:
-        if pe_col < pe_cols_used - 1:
-            with tl.scratch_scope():
-                m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
-                l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
-                O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
-                m_new, l_new, O_new = _merge_running(
-                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
-                )
-                tl.copy_to(m_local, m_new)
-                tl.copy_to(l_local, l_new)
-                tl.copy_to(O_local, O_new)
-        if pe_col > 0:
-            tl.send(dir="intra_W", src=m_local)
-            tl.send(dir="intra_W", src=l_local)
-            tl.send(dir="intra_W", src=O_local)
-
-    if pe_col == 0 and pe_rows_used > 1:
-        if pe_row < pe_rows_used - 1:
-            with tl.scratch_scope():
-                m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
-                l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
-                O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
-                m_new, l_new, O_new = _merge_running(
-                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
-                )
-                tl.copy_to(m_local, m_new)
-                tl.copy_to(l_local, l_new)
-                tl.copy_to(O_local, O_new)
-        if pe_row > 0:
-            tl.send(dir="intra_N", src=m_local)
-            tl.send(dir="intra_N", src=l_local)
-            tl.send(dir="intra_N", src=O_local)
-
-    # ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ──
-    # Only PE 0 of each CUBE participates. Adapts Phases 1-2 of
-    # lrab_hierarchical_allreduce.py: bidirectional row reduce converges
-    # at root_col; bidirectional col reduce on root_col converges at
-    # root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``.
-    if pe_id == 0:
-        row = cube_id // _SUB_W
-        col = cube_id % _SUB_W
-
-        # Phase 1: row reduce — converge at col == _ROOT_COL.
-        if col == 0:
-            tl.send(dir="E", src=m_local)
-            tl.send(dir="E", src=l_local)
-            tl.send(dir="E", src=O_local)
-        elif 0 < col < _ROOT_COL:
-            with tl.scratch_scope():
-                m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
-                l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
-                O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
-                m_new, l_new, O_new = _merge_running(
-                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
-                )
-                tl.copy_to(m_local, m_new)
-                tl.copy_to(l_local, l_new)
-                tl.copy_to(O_local, O_new)
-            tl.send(dir="E", src=m_local)
-            tl.send(dir="E", src=l_local)
-            tl.send(dir="E", src=O_local)
-        elif col == _ROOT_COL:
-            with tl.scratch_scope():
-                m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
-                l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
-                O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
-                m_new, l_new, O_new = _merge_running(
-                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
-                )
-                tl.copy_to(m_local, m_new)
-                tl.copy_to(l_local, l_new)
-                tl.copy_to(O_local, O_new)
-            with tl.scratch_scope():
-                m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
-                l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
-                O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
-                m_new, l_new, O_new = _merge_running(
-                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
-                )
-                tl.copy_to(m_local, m_new)
-                tl.copy_to(l_local, l_new)
-                tl.copy_to(O_local, O_new)
-        elif _ROOT_COL < col < _SUB_W - 1:
-            with tl.scratch_scope():
-                m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
-                l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
-                O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
-                m_new, l_new, O_new = _merge_running(
-                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
-                )
-                tl.copy_to(m_local, m_new)
-                tl.copy_to(l_local, l_new)
-                tl.copy_to(O_local, O_new)
-            tl.send(dir="W", src=m_local)
-            tl.send(dir="W", src=l_local)
-            tl.send(dir="W", src=O_local)
-        elif col == _SUB_W - 1:
-            tl.send(dir="W", src=m_local)
-            tl.send(dir="W", src=l_local)
-            tl.send(dir="W", src=O_local)
-
-        # Phase 2: col reduce on col == _ROOT_COL — converge at row == _ROOT_ROW.
-        if col == _ROOT_COL:
-            if row == 0:
-                tl.send(dir="S", src=m_local)
-                tl.send(dir="S", src=l_local)
-                tl.send(dir="S", src=O_local)
-            elif 0 < row < _ROOT_ROW:
-                with tl.scratch_scope():
-                    m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
-                    l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
-                    O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
-                    m_new, l_new, O_new = _merge_running(
-                        m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
-                    )
-                    tl.copy_to(m_local, m_new)
-                    tl.copy_to(l_local, l_new)
-                    tl.copy_to(O_local, O_new)
-                tl.send(dir="S", src=m_local)
-                tl.send(dir="S", src=l_local)
-                tl.send(dir="S", src=O_local)
-            elif row == _ROOT_ROW:
-                with tl.scratch_scope():
-                    m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
-                    l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
-                    O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
-                    m_new, l_new, O_new = _merge_running(
-                        m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
-                    )
-                    tl.copy_to(m_local, m_new)
-                    tl.copy_to(l_local, l_new)
-                    tl.copy_to(O_local, O_new)
-                if _SUB_H - 1 > _ROOT_ROW:
-                    with tl.scratch_scope():
-                        m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
-                        l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
-                        O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
-                        m_new, l_new, O_new = _merge_running(
-                            m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
-                        )
-                        tl.copy_to(m_local, m_new)
-                        tl.copy_to(l_local, l_new)
-                        tl.copy_to(O_local, O_new)
-            elif _ROOT_ROW < row < _SUB_H - 1:
-                with tl.scratch_scope():
-                    m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
-                    l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
-                    O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
-                    m_new, l_new, O_new = _merge_running(
-                        m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
-                    )
-                    tl.copy_to(m_local, m_new)
-                    tl.copy_to(l_local, l_new)
-                    tl.copy_to(O_local, O_new)
-                tl.send(dir="N", src=m_local)
-                tl.send(dir="N", src=l_local)
-                tl.send(dir="N", src=O_local)
-            elif row == _SUB_H - 1 and _SUB_H - 1 > _ROOT_ROW:
-                tl.send(dir="N", src=m_local)
-                tl.send(dir="N", src=l_local)
-                tl.send(dir="N", src=O_local)
+    # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
+    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)

    # ── Final normalise + store (root only: PE 0 of CUBE 6) ──
    if pe_id == 0 and cube_id == _ROOT_CUBE:
@@ -0,0 +1,71 @@
+"""GQA decode kernel — Case 6, **composite-GEMM** command form.
+
+Identical placement and (m, ℓ, O) reduce as the primitive Case-6 kernel
+(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the difference is the
+command *granularity* of the local attention. The primitive kernel walks
+its KV slice in ``TILE_S_KV``-wide tiles, issuing per-tile loads + GEMMs
+ a manual online-softmax merge — O(n_tiles) PE_CPU commands. This kernel
+instead issues **one coarse composite GEMM over the whole ``S_local``**
+for each matrix product, passing K and V as HBM refs so PE_SCHEDULER
+streams and tiles them (the DMA fan-out moves off PE_CPU). The
+online-softmax stays primitive but now runs once over the full score row.
+
+So the kernel issues O(1) coarse commands; the scheduler expands each
+composite into the same per-tile DMA + MAC work the primitive kernel
+issued by hand. Fewer, coarser PE_CPU commands ⇒ lower dispatch cost
+under the ADR-0064 Rev2 structural model (the CPU-offload win).
+"""
+from __future__ import annotations
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
+    _ROOT_CUBE,
+    reduce_mlo,
+)
+
+
+def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_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 (Cube-SP × PE-SP) — coarse composite-GEMM command form."""
+    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)
+
+    # ── Local attention as two coarse composite GEMMs ──
+    # K, V are HBM refs: PE_SCHEDULER streams + tiles them over S_local
+    # (the per-tile DMA fan-out the primitive kernel issued by hand).
+    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
+    K_T = tl.ref(k_ptr, shape=(d_head, S_local), dtype="f16")
+    V = tl.ref(v_ptr, shape=(S_local, d_head), dtype="f16")
+
+    scores = tl.composite(op="gemm", a=Q, b=K_T)          # Q·Kᵀ, one command
+    m_local = tl.max(scores, axis=-1)
+    centered = scores - m_local
+    exp_scores = tl.exp(centered)
+    l_local = tl.sum(exp_scores, axis=-1)
+    # P·V into a pre-allocated TCM output (explicit out — the composite
+    # output must be a real TCM handle, not auto-allocated scratch).
+    O_local = tl.zeros((G * T_q, d_head), dtype="f16")
+    tl.composite(op="gemm", a=exp_scores, b=V, out=O_local)  # P·V, one command
+
+    # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
+    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
+
+    # ── Final normalise + store (root only: PE 0 of CUBE 6) ──
+    if pe_id == 0 and cube_id == _ROOT_CUBE:
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)
@@ -0,0 +1,94 @@
+"""GQA decode kernel — Case 6, **composite_extended** command form.
+
+Same placement and (m, ℓ, O) reduce as the primitive / composite-GEMM
+Case-6 kernels; the local attention is the opt2 **two-composite** form
+(ADR-0060 §8 item 4 / ADR-0065), issued coarsely:
+
+  establish  a small first KV slice computes the running (m, ℓ, O) with
+             primitives (kernbench has no scratch-backed ``-inf``
+             initializer, so the recipe — which *merges* into an existing
+             accumulator — needs a seed).
+  #1  Q·Kᵀ   one coarse composite GEMM over the remaining S_local
+             (K passed as an HBM ref → PE_SCHEDULER streams + tiles it).
+  #2  softmax → one ``softmax_merge`` recipe composite (online-softmax
+      + P·V        merge of (m, ℓ, O)) whose head GEMM is P·V over the same
+                   remaining slice (V as an HBM ref), with an ``add``
+                   epilogue folding the P·V contribution into ``O``.
+
+So the bulk of the KV slice is two coarse PE_CPU commands, and the recipe
+additionally folds the per-tile online merge + P·V into one descriptor —
+the fewest / cheapest commands of the three forms (the headline
+CPU-offload win, ADR-0064 Rev2).
+
+Scope: runnable in op_log mode (latency / dispatch only). Full data-mode
+numeric parity of the recipe's 8 MATH ops is a separate follow-up
+(DDD-0065 / the P5-numerics note).
+"""
+from __future__ import annotations
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
+    _ROOT_CUBE,
+    reduce_mlo,
+)
+
+# Small primitive seed slice that establishes the running (m, ℓ, O) the
+# recipe then merges into. One scheduler K-tile wide (ADR-0064 TILE_K).
+_SEED_S_KV = 64
+
+
+def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_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 (Cube-SP × PE-SP) — softmax_merge recipe command form."""
+    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)
+    KV_ROW_BYTES = d_head * 2  # f16
+
+    # ── Establish running (m, ℓ, O) on a small seed slice (primitive) ──
+    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
+    seed = min(_SEED_S_KV, S_local)
+    K_T0 = tl.load(k_ptr, shape=(d_head, seed), dtype="f16")
+    V0 = tl.load(v_ptr, shape=(seed, d_head), dtype="f16")
+    scores0 = tl.dot(Q, K_T0)
+    m_local = tl.max(scores0, axis=-1)
+    exp0 = tl.exp(scores0 - m_local)
+    l_local = tl.sum(exp0, axis=-1)
+    O_local = tl.dot(exp0, V0)
+
+    # ── Remaining slice: one Q·Kᵀ composite + one softmax_merge recipe ──
+    rest = S_local - seed
+    if rest > 0:
+        K_T1 = tl.ref(k_ptr + seed * KV_ROW_BYTES,
+                      shape=(d_head, rest), dtype="f16")
+        V1 = tl.ref(v_ptr + seed * KV_ROW_BYTES,
+                    shape=(rest, d_head), dtype="f16")
+        scores1 = tl.composite(op="gemm", a=Q, b=K_T1)
+        tl.composite(
+            prologue=[{"op": "softmax_merge", "s": scores1,
+                       "m": m_local, "l": l_local, "O": O_local}],
+            op="gemm", b=V1, out=O_local,
+            epilogue=[{"op": "add", "other": O_local}],
+        )
+
+    # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
+    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
+
+    # ── Final normalise + store (root only: PE 0 of CUBE 6) ──
+    if pe_id == 0 and cube_id == _ROOT_CUBE:
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)
@@ -0,0 +1,214 @@
+"""Shared (m, ℓ, O) reduce for the Case-6 long-context decode kernels.
+
+Extracted from the Cube-SP × PE-SP decode kernel so the three
+command-form variants (primitive / composite / composite_extended)
+share one identical reduce. The local attention differs per variant;
+the reduce does not. Behavior is byte-equal to the pre-extraction inline
+reduce (guarded by tests/attention/test_gqa_decode_long_ctx_composite.py).
+
+The reduce is two-level:
+  • intra-CUBE 8-way (row chain along intra_W + col bridge along
+    intra_N) over the 2×4 PE grid;
+  • inter-CUBE 2-phase lrab over the 4×2 CUBE sub-mesh (ADR-0060 §4.2
+    lrab-adapted center-root reduce): Phase 1 row reduce converges at
+    root_col=2; Phase 2 col reduce on root_col converges at root_row=1;
+    root cube id = 6.
+Plain ``+`` is replaced by log-sum-exp ``_merge_running``.
+"""
+from __future__ import annotations
+
+
+# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh).
+_SUB_W = 4
+_SUB_H = 2
+_ROOT_COL = _SUB_W // 2   # 2
+_ROOT_ROW = _SUB_H // 2   # 1
+_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL  # 6
+
+PE_GRID_COLS = 4
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
+    m_new = tl.maximum(m_local, m_other)
+    scale_old = tl.exp(m_local - m_new)
+    scale_new = tl.exp(m_other - m_new)
+    l_new = l_local * scale_old + l_other * scale_new
+    O_new = O_local * scale_old + O_other * scale_new
+    return m_new, l_new, O_new
+
+
+def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, *, tl):
+    """Two-level (m, ℓ, O) reduce-to-root (PE 0 of CUBE 6). In place.
+
+    Updates ``m_local``/``l_local``/``O_local`` in place via ``copy_to``;
+    after this call only PE 0 of CUBE ``_ROOT_CUBE`` holds the full result.
+    """
+    # ── Intra-CUBE reduce: row chain (intra_W) + col bridge (intra_N) ──
+    pe_col = pe_id % PE_GRID_COLS
+    pe_row = pe_id // PE_GRID_COLS
+    pe_cols_used = min(PE_GRID_COLS, P)
+    pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
+
+    if pe_cols_used > 1:
+        if pe_col < pe_cols_used - 1:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        if pe_col > 0:
+            tl.send(dir="intra_W", src=m_local)
+            tl.send(dir="intra_W", src=l_local)
+            tl.send(dir="intra_W", src=O_local)
+
+    if pe_col == 0 and pe_rows_used > 1:
+        if pe_row < pe_rows_used - 1:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        if pe_row > 0:
+            tl.send(dir="intra_N", src=m_local)
+            tl.send(dir="intra_N", src=l_local)
+            tl.send(dir="intra_N", src=O_local)
+
+    # ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ──
+    # Only PE 0 of each CUBE participates. Adapts Phases 1-2 of
+    # lrab_hierarchical_allreduce.py: bidirectional row reduce converges
+    # at root_col; bidirectional col reduce on root_col converges at
+    # root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``.
+    if pe_id == 0:
+        row = cube_id // _SUB_W
+        col = cube_id % _SUB_W
+
+        # Phase 1: row reduce — converge at col == _ROOT_COL.
+        if col == 0:
+            tl.send(dir="E", src=m_local)
+            tl.send(dir="E", src=l_local)
+            tl.send(dir="E", src=O_local)
+        elif 0 < col < _ROOT_COL:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+            tl.send(dir="E", src=m_local)
+            tl.send(dir="E", src=l_local)
+            tl.send(dir="E", src=O_local)
+        elif col == _ROOT_COL:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        elif _ROOT_COL < col < _SUB_W - 1:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+            tl.send(dir="W", src=m_local)
+            tl.send(dir="W", src=l_local)
+            tl.send(dir="W", src=O_local)
+        elif col == _SUB_W - 1:
+            tl.send(dir="W", src=m_local)
+            tl.send(dir="W", src=l_local)
+            tl.send(dir="W", src=O_local)
+
+        # Phase 2: col reduce on col == _ROOT_COL — converge at row == _ROOT_ROW.
+        if col == _ROOT_COL:
+            if row == 0:
+                tl.send(dir="S", src=m_local)
+                tl.send(dir="S", src=l_local)
+                tl.send(dir="S", src=O_local)
+            elif 0 < row < _ROOT_ROW:
+                with tl.scratch_scope():
+                    m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
+                    l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
+                    O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
+                    m_new, l_new, O_new = _merge_running(
+                        m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                    )
+                    tl.copy_to(m_local, m_new)
+                    tl.copy_to(l_local, l_new)
+                    tl.copy_to(O_local, O_new)
+                tl.send(dir="S", src=m_local)
+                tl.send(dir="S", src=l_local)
+                tl.send(dir="S", src=O_local)
+            elif row == _ROOT_ROW:
+                with tl.scratch_scope():
+                    m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
+                    l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
+                    O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
+                    m_new, l_new, O_new = _merge_running(
+                        m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                    )
+                    tl.copy_to(m_local, m_new)
+                    tl.copy_to(l_local, l_new)
+                    tl.copy_to(O_local, O_new)
+                if _SUB_H - 1 > _ROOT_ROW:
+                    with tl.scratch_scope():
+                        m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
+                        l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
+                        O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
+                        m_new, l_new, O_new = _merge_running(
+                            m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                        )
+                        tl.copy_to(m_local, m_new)
+                        tl.copy_to(l_local, l_new)
+                        tl.copy_to(O_local, O_new)
+            elif _ROOT_ROW < row < _SUB_H - 1:
+                with tl.scratch_scope():
+                    m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
+                    l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
+                    O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
+                    m_new, l_new, O_new = _merge_running(
+                        m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                    )
+                    tl.copy_to(m_local, m_new)
+                    tl.copy_to(l_local, l_new)
+                    tl.copy_to(O_local, O_new)
+                tl.send(dir="N", src=m_local)
+                tl.send(dir="N", src=l_local)
+                tl.send(dir="N", src=O_local)
+            elif row == _SUB_H - 1 and _SUB_H - 1 > _ROOT_ROW:
+                tl.send(dir="N", src=m_local)
+                tl.send(dir="N", src=l_local)
+                tl.send(dir="N", src=O_local)
@@ -0,0 +1,182 @@
+"""milestone-1h-gqa decode: Case-6 composite-command study.
+
+Three command-form variants of the Case-6 (Cube-SP × PE-SP) long-context
+decode kernel, swept over S_kv so the local-attention tile loop runs
+multiple tiles (S_local = S_kv/(C·P) > TILE_S_KV = 1024):
+
+    primitive           tl.dot + primitive online-softmax merge (baseline).
+    composite           per-tile GEMMs as tl.composite GEMM commands.
+    composite_extended  per-tile attention as a Q·Kᵀ composite + a
+                        softmax_merge recipe composite (ADR-0065).
+
+The placement, DP policy, and (m, ℓ, O) reduce are identical across
+variants; only the per-tile command form differs. The sweep writes
+per-(variant, S_kv) latency, engine occupancy, and PE_CPU command count
+to sweep_decode_composite.json so the comparative plot
+(paper_plot_gqa_decode_long_ctx_composite.py) can read off the
+CPU-offload win.
+
+Run in op_log mode (enable_data=False): latency / dispatch only — recipe
+data-mode numeric parity is a separate follow-up (DDD-0065).
+"""
+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,
+)
+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.shared._gqa_panel_helpers import _ccl_cfg
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+
+_OUTPUT_DIR = (
+    Path(__file__).resolve().parents[2]
+    / "1H_milestone_output" / "gqa" / "long_ctx"
+)
+_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_composite.json"
+
+
+# ── Variant + S_kv registry ──────────────────────────────────────────
+
+
+# Each kernel implements the same Case-6 placement; only the per-tile
+# command form differs (see module docstring).
+_VARIANT_KERNELS = {
+    "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")
+
+# 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
+# tiles): the per-PE work is Q·Kᵀ (8,128)·(128,16384) and P·V
+# (8,16384)·(16384,128). End-to-end latency needs the data-mode engine,
+# whose cost scales with S_kv, so it is swept only over a tractable range.
+_S_KV_OPCOUNT = (8192, 65_536, 131_072, 262_144, 524_288, 1_048_576)
+_S_KV_LATENCY = (8192, 32_768, 65_536, 131_072)
+
+# LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs), one decode step.
+_BASE_PARAMS = dict(C=8, P=8, T_q=1, d_head=128, h_q=8, h_kv=1)
+
+
+# ── Per-(variant, S_kv) runner ───────────────────────────────────────
+
+
+def _run_panel_fn(variant: str, S_kv: int):
+    kernel = _VARIANT_KERNELS[variant]
+    p = _BASE_PARAMS
+    panel = f"decode_long_ctx_composite_{variant}_s{S_kv}"
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_full = DPPolicy(cube="replicate", pe="replicate",
+                           num_cubes=p["C"], num_pes=p["P"])
+        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                         num_cubes=p["C"], num_pes=p["P"])
+        q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                      dtype="f16", dp=dp_full, name=f"{panel}_q")
+        k = ctx.zeros((S_kv, p["h_kv"] * p["d_head"]),
+                      dtype="f16", dp=dp_kv, name=f"{panel}_k")
+        v = ctx.zeros((S_kv, p["h_kv"] * p["d_head"]),
+                      dtype="f16", dp=dp_kv, name=f"{panel}_v")
+        o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                      dtype="f16", dp=dp_full, name=f"{panel}_o")
+        ctx.launch(panel, kernel, q, k, v, o,
+                   p["T_q"], S_kv, p["h_q"], p["h_kv"], p["d_head"],
+                   p["C"], p["P"], _auto_dim_remap=False)
+
+    return _bench_fn
+
+
+def _end_to_end_ns(op_log) -> float:
+    if not op_log:
+        return 0.0
+    return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
+
+
+def _emit_dispatch(variant: str, S_kv: int) -> tuple[int, float]:
+    """PE_CPU dispatch (# commands, summed cycles) the kernel emits at the
+    lrab center rank (cube 6, pe 0) — computed at command-emit time, so it
+    is exact and S_kv-independent for the composite forms (no engine)."""
+    from kernbench.common.pe_commands import PeCpuOverheadCmd
+    from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL
+    from kernbench.triton_emu.tl_context import TLContext, run_kernel
+
+    p = _BASE_PARAMS
+    tl = TLContext(
+        pe_id=0, num_programs=p["P"], cost_model=DEFAULT_PE_COST_MODEL,
+        cube_id=6, num_cubes=p["C"], scratch_base=1 << 61, scratch_size=1 << 20,
+    )
+    run_kernel(
+        _VARIANT_KERNELS[variant], tl,
+        0x1000, 0x2000, 0x3000, 0x4000,
+        p["T_q"], S_kv, p["h_q"], p["h_kv"], p["d_head"], p["C"], p["P"],
+    )
+    cmds = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)]
+    return len(cmds), sum(c.cycles for c in cmds)
+
+
+def _engine_latency_ns(variant: str, S_kv: int, topology: str) -> float:
+    from kernbench.runtime_api.bench_runner import run_bench
+    from kernbench.runtime_api.types import resolve_device
+    from kernbench.sim_engine.engine import GraphEngine
+    from kernbench.topology.builder import resolve_topology
+
+    topo = resolve_topology(topology)
+    result = run_bench(
+        topology=topo, bench_fn=_run_panel_fn(variant, S_kv),
+        device=resolve_device(None),
+        engine_factory=lambda t, d: GraphEngine(
+            getattr(t, "topology_obj", t), enable_data=True,
+        ),
+    )
+    if not result.completion.ok:
+        raise RuntimeError(
+            f"gqa-decode-composite {variant}@{S_kv} failed: {result.completion}"
+        )
+    return _end_to_end_ns(result.engine.op_log)
+
+
+# ── Sweep entry (called by the umbrella milestone_1h_gqa) ────────────
+
+
+def run_sweep(topology: str = "topology.yaml") -> int:
+    """Emit-level dispatch over the full S_kv range (to 1M) + engine
+    latency over the tractable range; write sweep.json. Returns row count."""
+    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    rows = []
+    for S_kv in _S_KV_OPCOUNT:
+        for variant in _VARIANTS:
+            n_cmds, cycles = _emit_dispatch(variant, S_kv)
+            latency = (
+                _engine_latency_ns(variant, S_kv, topology)
+                if S_kv in _S_KV_LATENCY else None
+            )
+            rows.append({
+                "variant": variant,
+                "S_kv": S_kv,
+                **_BASE_PARAMS,
+                "pe_cpu_cmd_count": n_cmds,
+                "pe_cpu_dispatch_cycles": cycles,
+                "latency_ns": latency,
+            })
+    sweep = {
+        "version": 2,
+        "variants": list(_VARIANTS),
+        "s_kv_opcount": list(_S_KV_OPCOUNT),
+        "s_kv_latency": list(_S_KV_LATENCY),
+        "rows": rows,
+    }
+    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
+    print(f"  gqa-decode-composite: {len(rows)} rows -> {_SWEEP_JSON}")
+    return len(rows)
@@ -8,12 +8,13 @@ Currently exercises (long-context only — short-context panels are
 future work):
  - 4-cases prefill comparative study (gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases)
  - 4-cases decode comparative study  (gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases)
+  - Case-6 composite-command study    (gqa_helpers.long_ctx.gqa_decode_long_ctx_composite)

-Each sub-sweep writes its own ``sweep_{prefill,decode}.json`` into the
-shared output dir ``benches/1H_milestone_output/gqa/gqa_long_ctx/``.
+Each sub-sweep writes its own ``sweep_*.json`` into the shared output
+dir ``benches/1H_milestone_output/gqa/gqa_long_ctx/``.
 Selection via the env var ``GQA_1H_SWEEPS=prefill,decode`` (default
-runs both). Toggle individual sweeps with ``GQA_1H_SWEEPS=prefill``
-or ``GQA_1H_SWEEPS=decode``.
+runs prefill+decode). The composite study is opt-in (it sweeps the
+data-mode engine over several S_kv points): ``GQA_1H_SWEEPS=composite``.

 Gated by ``GQA_1H_RUN=1`` to keep CI fast.
 """
@@ -24,6 +25,9 @@ import os
 from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import (
    run_sweep as _run_decode_sweep,
 )
+from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (
+    run_sweep as _run_composite_sweep,
+)
 from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases import (
    run_sweep as _run_prefill_sweep,
 )
@@ -57,6 +61,7 @@ def run(torch) -> None:
    runners = {
        "prefill": _run_prefill_sweep,
        "decode": _run_decode_sweep,
+        "composite": _run_composite_sweep,
    }
    unknown = [s for s in sweeps if s not in runners]
    if unknown:
@@ -264,6 +264,10 @@ class TLContext:
            id=self._next_handle_id(),
            addr=addr, shape=shape, dtype=dtype,
            nbytes=nbytes, space="tcm",
+            # TCM-resident: a downstream composite operand reads it on-chip,
+            # not via a DMA_READ of its (bit-61) scratch address — same as
+            # the composite auto-output and recipe scratch handles below.
+            pinned=True,
        )

    # ── Reference (no DMA, metadata only) ────────────────────────