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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user