paper(gqa): two-regime "use of composite commands" — decode + compute-bound prefill

Complete the composite-command study across both attention regimes and
write up the result, resolving when the composite command helps latency.

Decode (memory-bound, T_q=1, M=8): the command form is latency-neutral —
the kernel is bound by KV streaming and the MAC array is near-idle, so the
composite's only benefit here is host-issue offload (O(n_tiles) -> O(1)
PE_CPU commands). Figure: gqa_decode_long_ctx_composite (latency flat,
command count saturates).

Prefill (compute-bound, large M=G*T_q tile-filling): add three command-form
variants of a single-rank FlashAttention prefill kernel
(_gqa_prefill_compute_bound: primitive / composite / composite_extended).
Here the composite's scheduler-internal per-tile DMA<->compute pipeline
keeps the MAC array fed while the primitive's blocking tl.dot starves it,
so composite wins on MAC utilization (68% flat -> 83%) and wall-clock
(19% faster at ctx=1024), with the margin growing with context (deeper
P.V reduction = more tiles to pipeline) — the compute-bound mirror of the
GEMM result in section 3. Figure: gqa_prefill_compute_bound (latency +
MAC util). Sweep wired as GQA_1H_SWEEPS=prefill_cb; tests cover structure,
e2e, and composite<primitive in the compute-bound regime.

The synthesis: composite has two benefits set by roofline position —
host-issue offload (always) and MAC-array feeding (compute-bound only).
Decode exercises the first, prefill both.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

docs/report/1H-codesign-paper/sections/05-gqa.tex

@@ -122,9 +122,10 @@ overlaps score computation; and per-tile \emph{scratch recycling} that
 keeps the running softmax accumulators ($m,\ell,O$) in a persistent arena
 while freeing per-tile temporaries, so the kernel fits the
 \SI{1}{\mebi\byte} scratch budget across many tiles. A further refinement
-that restructures the decode step into two stateful composites (a named
+restructures the decode step so the per-tile matrix products and the
-\textsf{softmax\_merge} recipe) is designed but not yet wired into the
+online-softmax merge are issued as \emph{composite} commands rather than
-measured path; results below reflect the implemented kernel only.
+hand-tiled primitives; \S\ref{sec:gqa-composite} measures it against the
+primitive baseline at long context.
 
 % TODO: CUBE <-> KV-head mapping diagram for the short-context regime
 % (h_kv=8 KV heads -> 8 CUBEs, 1:1; intra-CUBE PE usage).
@@ -221,6 +222,128 @@ memory-feasible family, and the cross-PE softmax reduction it does pay is
 precisely the traffic the communication-side codesign of this report is
 built to move quickly.
 
+\subsection{Use of Composite Commands}
+\label{sec:gqa-composite}
+
+The decode kernel of \S\ref{sec:gqa-long} issues its local attention as
+primitive operations: it walks each PE's $S_{\text{local}}$ token slice in
+\SI{1024}{}-token tiles, and for every tile issues a $Q\!\cdot\!K^{\top}$
+\textsf{dot}, the online-softmax primitives, and a $P\!\cdot\!V$
+\textsf{dot}, merging the running $(m,\ell,O)$ state by hand. The number
+of PE\_CPU commands this costs grows with the context. At a production
+context of $S_{kv}{=}1\,\text{M}$ tokens the Case-6 64-way split gives
+each PE $S_{\text{local}}{=}16384$ tokens, so its local attention is a
+$Q\!\cdot\!K^{\top}$ of $(8,128)\!\cdot\!(128,16384)$ and a
+$P\!\cdot\!V$ of $(8,16384)\!\cdot\!(16384,128)$---sixteen hand-issued
+tiles, each a fresh batch of CPU commands.
+
+The composite command lets the kernel hand that tiling to PE\_SCHEDULER.
+We compare three command forms of the \emph{same} Case-6 kernel---identical
+placement and identical $(m,\ell,O)$ reduce, differing only in how the
+local attention is issued:
+\begin{itemize}
+  \item \textbf{primitive}---the hand-tiled \textsf{dot}/softmax kernel
+    above (the baseline of \S\ref{sec:gqa-long}).
+  \item \textbf{composite}---each matrix product is one coarse
+    \textsf{composite} GEMM over the \emph{whole} $S_{\text{local}}$, with
+    $K$ and $V$ passed as HBM references so PE\_SCHEDULER streams and
+    tiles them on the fixed $32\!\times\!64\!\times\!32$ MAC tile; the
+    softmax stays primitive.
+  \item \textbf{composite\,+\,softmax\_merge}---additionally folds the
+    online-softmax merge and $P\!\cdot\!V$ into a single stateful
+    \textsf{softmax\_merge} recipe composite.
+\end{itemize}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_decode_long_ctx_composite.png}
+\caption{Three command forms of the Case-6 decode kernel, swept over
+context length ($S_{\text{local}}{=}S_{kv}/64$ per PE). \emph{Right:}
+PE\_CPU commands issued. The hand-tiled primitive kernel rises
+$O(n_{\text{tiles}})$---from 96 commands at one tile to 426 at the
+1\,M-token, sixteen-tile production point---while both composite forms
+issue a context-\emph{independent} $O(1)$ count (94 and 98) that
+\emph{saturates}: one coarse descriptor offloads the entire per-tile
+fan-out. \emph{Left:} the consequence for wall-clock latency is none---all
+three land on the same curve (\SI{30.6}{}, \SI{231}{},
+\SI{461}{\micro\second} at 8\,K\,/\,64\,K\,/\,128\,K), because decode is
+bound by streaming the KV cache, not by issue. Command-count is measured
+at emit time (exact, to 1\,M); latency on the data-mode engine over the
+tractable range.}
+\label{fig:gqa-composite}
+\end{figure}
+
+Figure~\ref{fig:gqa-composite} reads off the two quantities that matter,
+and they point in opposite directions. The PE\_CPU command count (right)
+collapses from a context-growing $O(n_{\text{tiles}})$ to a flat $O(1)$:
+at 1\,M tokens the composite form issues \num{94} commands against the
+primitive kernel's \num{426}, a $4.5\times$ reduction
+($4.2\times$ in modeled dispatch cost), and---crucially---that number no
+longer grows with context. The wall-clock latency (left), by contrast,
+is unchanged across all three forms: decode is bound by streaming the KV
+cache out of HBM, so the command form does not move the critical path.
+
+That juxtaposition is the point. The composite command is not a latency
+optimization for this memory-bound decode; it is a \emph{CPU-issue}
+optimization. Its value is removing the per-tile dispatch work that would
+otherwise grow without bound as context grows, freeing PE\_CPU to run
+ahead and keep the engines fed---which is exactly what lets the
+data-movement cost analyzed next show through as the true bottleneck
+rather than being masked by issue overhead. The \textsf{softmax\_merge}
+recipe folds the online merge into the same descriptor; on this
+memory-bound path its marginal cost over the plain GEMM composite is small
+(98 vs.\ 94 commands), and like the plain composite it keeps the issued
+count flat as context scales.
+
+\paragraph{The compute-bound mirror: prefill.} Decode's verdict---command
+form is latency-neutral---is a property of its regime, not of the
+composite command. A decode step has $T_q{=}1$, so its score and context
+products are skinny ($M{=}G\,T_q{=}8$): the MAC array is barely fed and
+the kernel is bound by streaming the KV cache. Prefill is the opposite
+corner. It processes a block of query positions at once, so $M{=}G\,T_q$
+is large and tile-filling, the GEMMs carry real arithmetic intensity
+($\sim$$M$ flops/byte, well above the roofline ridge), and the kernel is
+\emph{compute-bound}. This is the regime the composite command was built
+for (\S\ref{sec:gemm}): it streams the per-HW-tile
+DMA$\rightleftarrows$compute pipeline so the MAC array stays fed, whereas
+the primitive kernel's blocking \textsf{tl.dot} serializes each tile's
+load and compute and starves the array between tiles. We run the same
+three command forms on a single-rank compute-bound prefill (FlashAttention
+$Q$-block $\times$ $S_{kv}$-tile, online softmax) and sweep the context
+length (Figure~\ref{fig:gqa-prefill-cb}).
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_prefill_compute_bound.png}
+\caption{Compute-bound prefill, three command forms, swept over context
+length ($M{=}8\,T_q$ tile-filling). \emph{Left:} end-to-end latency.
+\emph{Right:} MAC utilization (achieved $\div$ the
+\SI{8}{\tera\flop\per\second} per-PE peak). The hand-tiled primitive sits
+flat at $\sim$\SI{68}{\percent}---its serial load$\to$dot path leaves the
+MAC array idle between tiles regardless of context. The composite forms
+climb with context (\SI{67}{}$\to$\SI{80}{\percent} plain,
+\SI{67}{}$\to$\SI{83}{\percent} with the recipe) because a deeper $P\!\cdot
+\!V$ reduction gives more HW tiles to pipeline, and they convert that into
+wall-clock: at \num{1024} the recipe form is \SI{646.9}{} vs.\
+\SI{794.1}{\micro\second} (\SI{19}{\percent} faster). The margin
+\emph{grows} with context---the compute-bound mirror of the GEMM result of
+\S\ref{sec:gemm}.}
+\label{fig:gqa-prefill-cb}
+\end{figure}
+
+The two studies together state the composite command's value precisely. It
+has two distinct benefits, and which one matters is set by the workload's
+roofline position. The first is \emph{host-issue offload}: one macro
+command in place of $O(n_{\text{tiles}})$ fine ones, which removes
+PE\_CPU dispatch work and is regime-independent (it shows in the decode
+command count). The second is \emph{MAC-array feeding}: the
+scheduler-internal per-tile DMA$\rightleftarrows$compute pipeline, which
+only converts to latency when the workload is compute-bound enough to have
+a MAC array worth keeping busy (it shows in the prefill utilization). A
+memory-bound decode exercises only the first; a compute-bound prefill
+exercises both. The composite command is the single mechanism that
+delivers each where it applies.
+
 \subsection{Comprehensive Analysis}
 \label{sec:gqa-analysis}
 

scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py

@@ -113,8 +113,8 @@ def main() -> None:
     )
     ax_cmd.set_ylabel("PE_CPU commands issued")
     ax_cmd.set_title(
-        "PE_CPU command count — primitive rises O(n$_\\mathrm{tiles}$); "
+        "PE_CPU command count per command form\n"
-        "composite saturates O(1)"
+        "(primitive O(n$_\\mathrm{tiles}$) rises; composite O(1) saturates)"
     )
     ax_cmd.axvline(1 << 20, color="#888", ls="--", lw=1, alpha=0.7)
     ax_cmd.annotate("1M production\ncontext", xy=(1 << 20, 0),

scripts/paper/paper_plot_gqa_prefill_compute_bound.py

@@ -0,0 +1,111 @@
+"""Comparative figure for the compute-bound prefill composite study.
+
+Reads sweep_prefill_compute_bound.json (emitted by milestone-1h-gqa,
+sweep ``prefill_cb``) and writes one two-panel PNG:
+
+  gqa_prefill_compute_bound.png
+    Left  — end-to-end prefill latency (µs) vs context length.
+    Right — MAC utilization (achieved / 8 TFLOP·s⁻¹ per-PE peak) vs context.
+
+Unlike memory-bound decode (where command form is latency-neutral), in
+compute-bound prefill the composite command keeps the MAC array fed by
+streaming DMA↔compute per HW tile, so it wins on both latency and
+utilization — and the margin grows with context (deeper P·V reduction =
+more tiles to pipeline).
+
+Run (after the bench):
+    GQA_1H_RUN=1 GQA_1H_SWEEPS=prefill_cb python -m kernbench.cli.main run \\
+        --bench milestone-1h-gqa --topology topology.yaml
+    python scripts/paper/paper_plot_gqa_prefill_compute_bound.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_prefill_compute_bound.json"
+_PAPER_FIG_DIR = (
+    _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
+)
+
+_VARIANT_STYLE = {
+    "primitive":          ("primitive (tl.dot, hand-tiled)", "#c0504d", "o"),
+    "composite":          ("composite GEMM",                 "#3b6ea5", "s"),
+    "composite_extended": ("composite + softmax_merge",      "#4f8a4f", "^"),
+}
+_ORDER = ("primitive", "composite", "composite_extended")
+
+
+def _ctx_label(c: int) -> str:
+    return f"{c // 1024}K" if c >= 1024 else str(c)
+
+
+def _series(rows, variant, key):
+    pts = sorted(((r["ctx_len"], r[key]) for r in rows
+                  if r["variant"] == variant), key=lambda t: t[0])
+    return [p[0] for p in pts], [p[1] for p in pts]
+
+
+def main() -> None:
+    sweep = json.loads(_SWEEP_JSON.read_text())
+    rows = sweep["rows"]
+    ctxs = sweep["ctx_points"]
+
+    fig, (ax_lat, ax_util) = plt.subplots(1, 2, figsize=(13.0, 4.8))
+
+    for v in _ORDER:
+        label, color, marker = _VARIANT_STYLE[v]
+        xs, lat = _series(rows, v, "latency_ns")
+        ax_lat.plot(xs, [y / 1e3 for y in lat], marker=marker,
+                    color=color, label=label, lw=2)
+        xs, util = _series(rows, v, "mac_util")
+        ax_util.plot(xs, [u * 100 for u in util], marker=marker,
+                     color=color, label=label, lw=2)
+
+    for ax in (ax_lat, ax_util):
+        ax.set_xscale("log", base=2)
+        ax.set_xticks(ctxs)
+        ax.set_xticklabels([_ctx_label(c) for c in ctxs])
+        ax.set_xlabel(r"context length (= $T_q$ = $S_{kv}$)")
+        ax.grid(True, ls=":", alpha=0.5)
+        ax.legend(fontsize=9)
+
+    ax_lat.set_ylabel("end-to-end prefill latency (µs)")
+    ax_lat.set_title("Compute-bound prefill latency per command form")
+    ax_util.set_ylabel("MAC utilization (% of 8 TFLOP·s⁻¹ peak)")
+    ax_util.set_title(
+        "MAC utilization — composite keeps the array fed; primitive starves"
+    )
+    ax_util.axhline(100, color="#888", ls="--", lw=1, alpha=0.6)
+
+    fig.suptitle(
+        "Compute-bound prefill attention — use of composite commands\n"
+        "single-rank, GQA single-KV-head group ($h_q{=}8$, $d_{\\text{head}}"
+        "{=}128$); $M{=}8T_q$ tile-filling",
+        fontsize=11,
+    )
+    fig.tight_layout(rect=(0, 0, 1, 0.92))
+
+    out = _FIG_DIR / "gqa_prefill_compute_bound.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()

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

@@ -0,0 +1,258 @@
+{
+  "version": 2,
+  "variants": [
+    "primitive",
+    "composite",
+    "composite_extended"
+  ],
+  "s_kv_opcount": [
+    8192,
+    65536,
+    131072,
+    262144,
+    524288,
+    1048576
+  ],
+  "s_kv_latency": [
+    8192,
+    32768,
+    65536,
+    131072
+  ],
+  "rows": [
+    {
+      "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",
+      "S_kv": 8192,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": 30566.1840000007
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 8192,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "latency_ns": 30483.359500000362
+    },
+    {
+      "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",
+      "S_kv": 65536,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": 231270.18400000152
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 65536,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "latency_ns": 231211.1740000015
+    },
+    {
+      "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",
+      "S_kv": 131072,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": 460651.3170000027
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 131072,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "latency_ns": 460552.9805000025
+    },
+    {
+      "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
+    },
+    {
+      "variant": "composite",
+      "S_kv": 262144,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": null
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 262144,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "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
+    },
+    {
+      "variant": "composite",
+      "S_kv": 524288,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": null
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 524288,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "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
+    },
+    {
+      "variant": "composite",
+      "S_kv": 1048576,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": null
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 1048576,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "latency_ns": null
+    }
+  ]
+}

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

@@ -0,0 +1,105 @@
+{
+  "version": 1,
+  "variants": [
+    "primitive",
+    "composite",
+    "composite_extended"
+  ],
+  "ctx_points": [
+    256,
+    512,
+    1024
+  ],
+  "rows": [
+    {
+      "variant": "primitive",
+      "ctx_len": 256,
+      "M": 2048,
+      "latency_ns": 48857.361999999725,
+      "gemm_busy_ns": 33554.43200000003,
+      "dma_busy_ns": 20880.000000000004,
+      "achieved_tflops": 5.4942683151825005,
+      "mac_util": 0.6867835393978126
+    },
+    {
+      "variant": "composite",
+      "ctx_len": 256,
+      "M": 2048,
+      "latency_ns": 49887.505999999594,
+      "gemm_busy_ns": 34531.32799999314,
+      "dma_busy_ns": 1095773.440000072,
+      "achieved_tflops": 5.380815308746887,
+      "mac_util": 0.6726019135933609
+    },
+    {
+      "variant": "composite_extended",
+      "ctx_len": 256,
+      "M": 2048,
+      "latency_ns": 50280.91399999841,
+      "gemm_busy_ns": 129096.19199996963,
+      "dma_busy_ns": 633563.5200000411,
+      "achieved_tflops": 5.338714725830331,
+      "mac_util": 0.6673393407287914
+    },
+    {
+      "variant": "primitive",
+      "ctx_len": 512,
+      "M": 4096,
+      "latency_ns": 197496.81800000605,
+      "gemm_busy_ns": 134217.72800000012,
+      "dma_busy_ns": 66880.0,
+      "achieved_tflops": 5.436755056985105,
+      "mac_util": 0.6795943821231382
+    },
+    {
+      "variant": "composite",
+      "ctx_len": 512,
+      "M": 4096,
+      "latency_ns": 177341.23400000148,
+      "gemm_busy_ns": 141168.63999995415,
+      "dma_busy_ns": 8567063.039998509,
+      "achieved_tflops": 6.054665346469796,
+      "mac_util": 0.7568331683087245
+    },
+    {
+      "variant": "composite_extended",
+      "ctx_len": 512,
+      "M": 4096,
+      "latency_ns": 174702.7699999963,
+      "gemm_busy_ns": 1076198.3999996716,
+      "dma_busy_ns": 6594394.87999886,
+      "achieved_tflops": 6.146106464139193,
+      "mac_util": 0.7682633080173992
+    },
+    {
+      "variant": "primitive",
+      "ctx_len": 1024,
+      "M": 8192,
+      "latency_ns": 794090.4820000405,
+      "gemm_busy_ns": 536870.9120000004,
+      "dma_busy_ns": 234240.0,
+      "achieved_tflops": 5.4086623544239645,
+      "mac_util": 0.6760827943029956
+    },
+    {
+      "variant": "composite",
+      "ctx_len": 1024,
+      "M": 8192,
+      "latency_ns": 668110.7059999453,
+      "gemm_busy_ns": 873011.1999923651,
+      "dma_busy_ns": 67794150.3999821,
+      "achieved_tflops": 6.428526376585786,
+      "mac_util": 0.8035657970732233
+    },
+    {
+      "variant": "composite_extended",
+      "ctx_len": 1024,
+      "M": 8192,
+      "latency_ns": 646937.2019999702,
+      "gemm_busy_ns": 8870907.903995967,
+      "dma_busy_ns": 59655820.79998447,
+      "achieved_tflops": 6.638924586068552,
+      "mac_util": 0.829865573258569
+    }
+  ]
+}

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

@@ -0,0 +1,161 @@
+"""Compute-bound prefill attention — 3 command-form variants (single-rank).
+
+Companion to the memory-bound decode study
+(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp*``). Prefill processes a
+block of T_q query positions at once, so the score / context GEMMs have a
+large M = G·T_q and high arithmetic intensity (~M flops/byte) — the
+workload is **compute-bound** (above the roofline ridge), unlike T_q=1
+decode (M=8, memory-bound). This is the regime where the composite
+command's value shows: it streams DMA↔compute per HW tile to keep the MAC
+array fed, while the primitive kernel serializes load→dot and starves it.
+
+Single-rank (C=P=1): no cross-device reduce — the focus is the per-PE
+GEMM-issue mechanism. FlashAttention 2-D tiling (Q-block × S_kv-tile,
+online softmax) bounds the TCM scratch.
+
+Three forms, differing only in how each Q-block's local attention is
+issued (placement/softmax identical):
+  primitive           tl.dot per S_kv-tile + primitive online merge.
+  composite           one coarse composite GEMM per Q-block over the whole
+                      S_kv (K, V as HBM refs → PE_SCHEDULER tiles/streams).
+  composite_extended  Q·Kᵀ composite + softmax_merge recipe composite.
+"""
+from __future__ import annotations
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import _merge_running
+
+M_BLOCK = 128       # query rows per FlashAttention Q-block (tile-filling).
+TILE_S_KV = 256     # primitive per-tile S_kv width.
+_SEED_S_KV = 64     # composite_extended recipe seed slice.
+
+
+def _qblock_bounds(T_q: int, h_q: int, h_kv: int):
+    G = h_q // h_kv
+    M_total = G * T_q
+    n_qblocks = (M_total + M_BLOCK - 1) // M_BLOCK
+    return G, M_total, n_qblocks
+
+
+# ── primitive: hand-tiled monolithic dots + online merge ─────────────
+
+
+def gqa_prefill_primitive_kernel(
+    q_ptr, k_ptr, v_ptr, o_ptr,
+    T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
+) -> None:
+    G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
+    ROW = d_head * 2  # f16
+    n_tiles = (S_kv + TILE_S_KV - 1) // TILE_S_KV
+
+    for qb in range(n_qblocks):
+        # Per-Q-block scratch_scope: this block's running (m,ℓ,O) plus its
+        # tile-0 transients are freed before the next block, so scratch
+        # stays O(one block) rather than accumulating across all blocks.
+        with tl.scratch_scope():
+            m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
+            q_off = qb * M_BLOCK * ROW
+            Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
+
+            tile_s0 = min(TILE_S_KV, S_kv)
+            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 = tl.dot(Q, K_T)
+            m_local = tl.max(scores, axis=-1)
+            exp_s = tl.exp(scores - m_local)
+            l_local = tl.sum(exp_s, axis=-1)
+            O_local = tl.dot(exp_s, V)
+
+            for ti in range(1, n_tiles):
+                start = ti * TILE_S_KV
+                tile_s = min(TILE_S_KV, S_kv - start)
+                with tl.scratch_scope():
+                    K_T_t = tl.load(k_ptr + start * ROW,
+                                    shape=(d_head, tile_s), dtype="f16")
+                    V_t = tl.load(v_ptr + start * ROW,
+                                  shape=(tile_s, d_head), dtype="f16")
+                    scores_t = tl.dot(Q, K_T_t)
+                    m_tile = tl.max(scores_t, axis=-1)
+                    exp_t = tl.exp(scores_t - m_tile)
+                    l_tile = tl.sum(exp_t, axis=-1)
+                    O_tile = tl.dot(exp_t, V_t)
+                    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)
+
+            O_final = O_local / l_local
+            tl.store(o_ptr + q_off, O_final)
+
+
+# ── composite: one coarse composite GEMM per Q-block ─────────────────
+
+
+def gqa_prefill_composite_kernel(
+    q_ptr, k_ptr, v_ptr, o_ptr,
+    T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
+) -> None:
+    G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
+    ROW = d_head * 2
+
+    for qb in range(n_qblocks):
+        with tl.scratch_scope():
+            m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
+            q_off = qb * M_BLOCK * ROW
+            Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
+            K_T = tl.ref(k_ptr, shape=(d_head, S_kv), dtype="f16")
+            V = tl.ref(v_ptr, shape=(S_kv, 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)
+            exp_s = tl.exp(scores - m_local)
+            l_local = tl.sum(exp_s, axis=-1)
+            O_local = tl.zeros((m_blk, d_head), dtype="f16")
+            tl.composite(op="gemm", a=exp_s, b=V, out=O_local)  # P·V, one command
+
+            O_final = O_local / l_local
+            tl.store(o_ptr + q_off, O_final)
+
+
+# ── composite_extended: Q·Kᵀ composite + softmax_merge recipe ────────
+
+
+def gqa_prefill_composite_ext_kernel(
+    q_ptr, k_ptr, v_ptr, o_ptr,
+    T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
+) -> None:
+    G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
+    ROW = d_head * 2
+
+    for qb in range(n_qblocks):
+        with tl.scratch_scope():
+            m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
+            q_off = qb * M_BLOCK * ROW
+            Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
+
+            seed = min(_SEED_S_KV, S_kv)
+            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)
+
+            rest = S_kv - seed
+            if rest > 0:
+                K_T1 = tl.ref(k_ptr + seed * ROW,
+                              shape=(d_head, rest), dtype="f16")
+                V1 = tl.ref(v_ptr + seed * ROW,
+                            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}],
+                )
+
+            O_final = O_local / l_local
+            tl.store(o_ptr + q_off, O_final)

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

@@ -0,0 +1,142 @@
+"""milestone-1h-gqa: compute-bound prefill composite-command study.
+
+Three command-form variants of a single-rank compute-bound prefill
+attention kernel (``_gqa_prefill_compute_bound``), swept over context
+length S_kv = T_q. Unlike the memory-bound decode study, prefill has a
+large M = G·T_q, so the score / context GEMMs are compute-bound — the
+regime where the composite command keeps the MAC array fed (DMA↔compute
+pipelining) and the hand-tiled primitive starves it on the serial
+load→dot path.
+
+Records per (variant, context) the end-to-end latency, the GEMM-engine
+busy time, and the MAC occupancy (gemm_busy / e2e) so the comparative
+plot can show composite winning on both latency and utilization.
+
+Runs in data mode (engine latency). Gated via the umbrella
+``GQA_1H_SWEEPS=prefill_cb``.
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
+    gqa_prefill_composite_ext_kernel,
+    gqa_prefill_composite_kernel,
+    gqa_prefill_primitive_kernel,
+)
+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_prefill_compute_bound.json"
+
+_VARIANT_KERNELS = {
+    "primitive": gqa_prefill_primitive_kernel,
+    "composite": gqa_prefill_composite_kernel,
+    "composite_extended": gqa_prefill_composite_ext_kernel,
+}
+_VARIANTS = ("primitive", "composite", "composite_extended")
+
+# Context length S_kv = T_q (prefill processes T_q tokens against S_kv=T_q
+# keys). M = G·T_q = 8·T_q is tile-filling/compute-bound at every point.
+_CTX_POINTS = (256, 512, 1024)
+
+_H_Q, _H_KV, _D_HEAD = 8, 1, 128
+_PEAK_TFLOPS = 8.0  # per-PE f16 GEMM peak (topology.yaml pe_gemm.peak_tflops_f16)
+
+
+def _run_panel_fn(variant: str, ctx_len: int):
+    kernel = _VARIANT_KERNELS[variant]
+    panel = f"prefill_cb_{variant}_c{ctx_len}"
+
+    def _bench_fn(ctx):
+        dp = DPPolicy(cube="replicate", pe="replicate",
+                      num_cubes=1, num_pes=1)
+        q = ctx.zeros((ctx_len, _H_Q * _D_HEAD),
+                      dtype="f16", dp=dp, name=f"{panel}_q")
+        k = ctx.zeros((ctx_len, _H_KV * _D_HEAD),
+                      dtype="f16", dp=dp, name=f"{panel}_k")
+        v = ctx.zeros((ctx_len, _H_KV * _D_HEAD),
+                      dtype="f16", dp=dp, name=f"{panel}_v")
+        o = ctx.empty((ctx_len, _H_Q * _D_HEAD),
+                      dtype="f16", dp=dp, name=f"{panel}_o")
+        ctx.launch(panel, kernel, q, k, v, o,
+                   ctx_len, ctx_len, _H_Q, _H_KV, _D_HEAD, 1, 1,
+                   _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 _engine_busy_ns(op_log, suffix: str) -> float:
+    return sum(r.t_end - r.t_start
+               for r in op_log if r.component_id.endswith("." + suffix))
+
+
+def _run_panel(variant: str, ctx_len: int, topology: str) -> dict:
+    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, ctx_len),
+        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-prefill-cb {variant}@{ctx_len} failed: {result.completion}"
+        )
+    op_log = result.engine.op_log
+    e2e = _end_to_end_ns(op_log)
+    gemm = _engine_busy_ns(op_log, "pe_gemm")
+    dma = _engine_busy_ns(op_log, "pe_dma")
+    # Useful attention flops (Q·Kᵀ + P·V), single rank.
+    G = _H_Q // _H_KV
+    M = G * ctx_len
+    useful_flops = 4.0 * M * _D_HEAD * ctx_len
+    # MAC utilization = achieved / peak. ``achieved_tflops`` uses wall-clock
+    # (useful_flops / e2e) so it is bounded by peak even when the composite
+    # path overlaps many tile GEMMs (gemm_busy is a sum over overlapping ops
+    # and is kept only for diagnostics).
+    return {
+        "variant": variant,
+        "ctx_len": ctx_len,
+        "M": M,
+        "latency_ns": e2e,
+        "gemm_busy_ns": gemm,
+        "dma_busy_ns": dma,
+        "achieved_tflops": (useful_flops / e2e / 1e3) if e2e > 0 else 0.0,
+        "mac_util": (useful_flops / e2e / 1e3 / _PEAK_TFLOPS) if e2e > 0 else 0.0,
+    }
+
+
+def run_sweep(topology: str = "topology.yaml") -> int:
+    """Drive all (variant, context) prefill panels; write sweep.json."""
+    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    rows = [
+        _run_panel(variant, ctx_len, topology)
+        for ctx_len in _CTX_POINTS
+        for variant in _VARIANTS
+    ]
+    sweep = {
+        "version": 1,
+        "variants": list(_VARIANTS),
+        "ctx_points": list(_CTX_POINTS),
+        "rows": rows,
+    }
+    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
+    print(f"  gqa-prefill-cb: {len(rows)} rows -> {_SWEEP_JSON}")
+    return len(rows)

src/kernbench/benches/milestone_1h_gqa.py

@@ -28,6 +28,9 @@ from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import (
 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_compute_bound import (
+    run_sweep as _run_prefill_cb_sweep,
+)
 from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases import (
     run_sweep as _run_prefill_sweep,
 )
@@ -62,6 +65,7 @@ def run(torch) -> None:
         "prefill": _run_prefill_sweep,
         "decode": _run_decode_sweep,
         "composite": _run_composite_sweep,
+        "prefill_cb": _run_prefill_cb_sweep,
     }
     unknown = [s for s in sweeps if s not in runners]
     if unknown:

tests/attention/test_gqa_prefill_compute_bound.py

@@ -0,0 +1,61 @@
+"""Tests for the compute-bound prefill composite-command study.
+
+Unlike memory-bound decode (where command form is latency-neutral), in
+compute-bound prefill the composite command keeps the MAC array fed by
+pipelining DMA↔compute per HW tile, so it wins on wall-clock latency. The
+rigorous claim here is therefore the *opposite* of the decode study:
+``composite_latency < primitive_latency`` at a compute-bound context.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.common.pe_commands import CompositeCmd
+from kernbench.triton_emu.tl_context import TLContext, run_kernel
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
+    gqa_prefill_composite_ext_kernel,
+    gqa_prefill_composite_kernel,
+    gqa_prefill_primitive_kernel,
+)
+
+_TOPOLOGY = Path(__file__).resolve().parents[2] / "topology.yaml"
+_H_Q, _H_KV, _D_HEAD = 8, 1, 128
+
+
+def _n_composites(kernel, T_q: int, S_kv: int) -> int:
+    tl = TLContext(pe_id=0, num_programs=1, scratch_base=1 << 61,
+                   scratch_size=1 << 20)
+    run_kernel(kernel, tl, 0x1000, 0x2000, 0x3000, 0x4000,
+               T_q, S_kv, _H_Q, _H_KV, _D_HEAD, 1, 1)
+    return sum(1 for c in tl.commands if isinstance(c, CompositeCmd))
+
+
+def test_command_form_composite_counts():
+    """One Q-block (T_q=16 → M=128): the primitive issues no composites;
+    each composite form issues exactly two (Q·Kᵀ and P·V / recipe)."""
+    assert _n_composites(gqa_prefill_primitive_kernel, 16, 256) == 0
+    assert _n_composites(gqa_prefill_composite_kernel, 16, 256) == 2
+    assert _n_composites(gqa_prefill_composite_ext_kernel, 16, 256) == 2
+
+
+def _latency_ns(variant: str, ctx_len: int) -> float:
+    from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
+        _run_panel,
+    )
+    return _run_panel(variant, ctx_len, str(_TOPOLOGY))["latency_ns"]
+
+
+def test_three_variants_complete_in_data_mode():
+    """All three prefill forms run end-to-end through the engine."""
+    for variant in ("primitive", "composite", "composite_extended"):
+        assert _latency_ns(variant, 256) > 0
+
+
+def test_composite_faster_in_compute_bound_prefill():
+    """At a compute-bound context (512), the composite command's DMA↔compute
+    pipelining beats the primitive's serial load→dot — the opposite of the
+    memory-bound decode study, where the two tie."""
+    prim = _latency_ns("primitive", 512)
+    comp = _latency_ns("composite", 512)
+    assert comp < prim, f"composite {comp} not < primitive {prim}"