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
-\textsf{softmax\_merge} recipe) is designed but not yet wired into the
-measured path; results below reflect the implemented kernel only.
+restructures the decode step so the per-tile matrix products and the
+online-softmax merge are issued as \emph{composite} commands rather than
+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}