paper(gqa): single-rank decode-streaming study — composite's masked bandwidth win

Adds the memory-bound mirror of the compute-bound prefill figure. New single-rank decode sweep (T_q=1, M=8) reuses the prefill_compute_bound kernels, sweeping per-rank S_kv with NO cross-CUBE reduce — isolating the local attention that the 64-way Case-6 decode masks under its reduce tail.

Finding: even memory-bound decode benefits from the composite command. Its scheduler-streamed concurrent per-tile DMAs reach ~233 GB/s (91% of the 256 GB/s per-rank roofline) while the primitive's blocking tl.dot serializes one tile DMA and plateaus at ~166 GB/s — a 25-28% latency win that widens with context. This refines the paper's 'decode is latency-neutral' claim: neutral at 64-way production scale (reduce-dominated), but the composite extracts bandwidth at the local level. The bandwidth roofline here mirrors the MAC roofline in prefill.

Adds: gqa_decode_streaming.py sweep (wired into the milestone umbrella as GQA_1H_SWEEPS=decode_streaming), paper_plot_gqa_decode_streaming.py, the figure, and a new fig:gqa-decode-stream + paragraph in 05-gqa.tex; rebuilt main.pdf. Prefill (Experiment B) re-verified byte-identical (794.1/668.1/646.9us).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-23 19:39:45 -07:00
parent faf011dae5
commit a4a2683aad
8 changed files with 514 additions and 3 deletions
Binary file not shown.
Binary file not shown.

After

Width:  |  Height:  |  Size: 136 KiB

@@ -284,7 +284,9 @@ 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 for this memory-bound decode \emph{at the 64-way production scale} (a
single-rank caveat follows, Figure~\ref{fig:gqa-decode-stream}); 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
@@ -295,8 +297,49 @@ 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
\paragraph{Isolating the rank: the masked streaming win.} That neutrality
is a property of the \emph{full 64-way} critical path, not of the local
attention: at production scale the inter-CUBE $(m,\ell,O)$ reduce tail and
shared-HBM contention set the wall clock, so a faster local attention does
not surface. Stripping those away---a single rank, no cross-CUBE reduce,
swept over the per-rank context $S_{kv}$ (so $S_{kv}{=}16$\,K here is the
per-PE load of a 1M-token, 64-way-sharded decode)---exposes the local
attention directly (Figure~\ref{fig:gqa-decode-stream}), and the composite
\emph{does} win, by \SI{25}{}--\SI{28}{\percent}. The reason is the
memory-bound mirror of prefill: its scheduler-streamed concurrent per-tile
DMAs keep the HBM pipeline full and reach \SI{233}{\giga\byte\per\second}
---\SI{91}{\percent} of the per-rank \SI{256}{\giga\byte\per\second}
roofline---whereas the primitive kernel's blocking \textsf{tl.dot}
serializes one tile DMA at a time and plateaus at
\SI{166}{\giga\byte\per\second}. So even for memory-bound decode the
composite is not \emph{only} a CPU-issue optimization---it also extracts
bandwidth---but that latency benefit materializes only when the local
attention is on the critical path, which at 64-way production scale it is
not.
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_decode_streaming.png}
\caption{Single-rank memory-bound decode ($T_q{=}1$, $M{=}8$), three
command forms, swept over per-rank context. \emph{Left:} end-to-end
latency---the composite forms run \SI{25}{}--\SI{28}{\percent} below the
primitive, a gap that widens with context. \emph{Right:} achieved HBM
bandwidth against the per-rank \SI{256}{\giga\byte\per\second} roofline.
The primitive's blocking load$\rightarrow$dot serializes the KV stream and
plateaus at \SI{166}{\giga\byte\per\second}; the composite forms pipeline
concurrent per-tile DMAs through the scheduler and reach
\SI{233}{\giga\byte\per\second}. This is the memory-bound mirror of the
prefill result (Figure~\ref{fig:gqa-prefill-cb}): there the composite
approaches the MAC roofline, here the bandwidth roofline. Capped at
$16$\,K---the plain composite materializes the full $(M,S_{kv})$ scores in
TCM, so beyond that only the \textsf{softmax\_merge} recipe, which tiles
the softmax, stays within scratch.}
\label{fig:gqa-decode-stream}
\end{figure}
\paragraph{The compute-bound mirror: prefill.} Decode's \emph{production}
verdict---command form is latency-neutral at 64-way scale---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