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

\section{Fused Grouped-Query Attention}
\label{sec:gqa}

Attention is the 1H focus, and it is where the two preceding optimizations
have to come together. Grouped-Query Attention (GQA) shrinks the KV cache
by sharing each KV head across a group of query heads (here $h_q=8$ query
heads to $h_{kv}=1$ KV head, a group factor $G=8$), which makes decoding
feasible at long context but also makes it acutely memory-bound: a decode
step processes a single query position ($T_q=1$) against the entire KV
history, so its arithmetic intensity is low and its time is dominated by
streaming the KV cache out of HBM. FlashAttention-style tiling with an
online-softmax merge avoids ever materializing the full score matrix, but
realizing it as a fast \emph{fused} kernel needs both building blocks from
this report: efficient GEMM issue (\S\ref{sec:gemm}) for the
$Q\!\cdot\!K^{\top}$ and $P\!\cdot\!V$ products, and an efficient on-device
reduction (\S\ref{sec:allreduce}) for the multi-user and
sequence-parallel KV reductions. \emph{Fused} here is meant in the
FlashAttention sense---$Q\!\cdot\!K^{\top}$, the online softmax, and
$P\!\cdot\!V$ collapse into a single kernel that never materializes the
score matrix---and, beyond that, the cross-device KV reduction is absorbed
into the same kernel (on PE\_IPCQ) rather than issued as a separate
all-reduce. This section is the capstone: the fused
kernel that uses the composite command and PE\_IPCQ at the same time.
Multi-head attention (MHA) was studied in prior work and serves here as
the established baseline rather than being re-derived.

\subsection{Data Placement Policy}
\label{sec:gqa-placement}

Long-context decode is bound by the KV cache, so the first-order design
question is how to place that cache---and the running softmax state it
feeds---across the two hardware axes the machine exposes: the CUBEs and,
within each CUBE, the PEs (here $C{=}8$ CUBEs $\times$ $P{=}8$ PEs, for
$C\!\cdot\!P{=}64$ attention engines over one KV-head group on the
LLaMA-3.1-70B target). Each axis can \emph{replicate} the KV cache or
\emph{shard} it, and a shard can run along the sequence dimension
$S_{kv}$ or the head dimension $d_{\text{head}}$. The cross product is a
small, enumerable taxonomy; six placements span its meaningful corners
(Figure~\ref{fig:gqa-kv-sharding}).

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_kv_sharding_diagram.png}
\caption{The six KV-placement strategies, drawn on the
$S_{kv}\!\times\!d_{\text{head}}$ KV tensor (rows = sequence, columns =
head dimension). Cube colour bands and dashed PE dividers show which
axis each level shards. Cases~1--3 either replicate the cache or shard
it on a single axis (8-way at most); Cases~4--6 reach a full 64-way
split, three different ways: Case~4 splits $S_{kv}$ across CUBEs and
$d_{\text{head}}$ across PEs, Case~5 the mirror, and Case~6~$\star$
splits $S_{kv}$ on \emph{both} axes.}
\label{fig:gqa-kv-sharding}
\end{figure}

Two quantities decide which placement is viable, and they pull against
each other. The first is \textbf{per-PE KV memory}. With a per-PE HBM
budget of \SI{6.0}{\giga\byte} and \SI{1.76}{\giga\byte} of attention
weights resident, the KV headroom is \SI{4.24}{\giga\byte} per PE. At a
production context of $S_{kv}{=}1\,\text{M}$ tokens the unsharded cache
is \SI{40}{\giga\byte}/PE (Case~1), an 8-way shard is \SI{5}{\giga\byte}
(Cases~2--3)---both \emph{over} the headroom---while only the 64-way
placements bring it to \SI{640}{\mega\byte}/PE (Cases~4--6), comfortably
inside budget. Memory alone therefore eliminates Cases~1--3 at long
context. The second quantity is \textbf{communication per token}, and it
is what separates the three survivors. Sharding $d_{\text{head}}$
(Cases~4--5) makes each PE hold only a slice of every head, so the
$Q\!\cdot\!K^{\top}$ score is \emph{partial} and must be all-reduced
across the slice owners on every token---a reduction whose volume scales
with $S_{kv}$ ($\sim$\SI{166}{\mega\byte}/token analytically, intra-CUBE
on the NoC for Case~4, inter-CUBE on UCIe for Case~5). Case~6~$\star$
instead shards $S_{kv}$ on both axes, so every PE computes a
\emph{complete} score over its own token range and only the small running
softmax state $(m,\ell,O)$ is merged across PEs
($\sim$\SI{6.2}{\mega\byte}/token)---a $\sim$27$\times$ lighter collective
than the $d_{\text{head}}$-split designs.

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_summary.png}
\caption{Long-context placement analysis at $S_{kv}{=}1\,\text{M}$ tokens.
\emph{Left:} the per-PE HBM budget---\SI{1.76}{\giga\byte} of attention
weights leave \SI{4.24}{\giga\byte} of KV headroom (red line).
\emph{Middle:} per-PE KV memory per case (log scale); only the 64-way
placements (Cases~4--6, \SI{640}{\mega\byte}) clear the headroom, while
the unsharded (\SI{40}{\giga\byte}) and 8-way (\SI{5}{\giga\byte}) cases
overflow. \emph{Right:} analytical communication per token (log scale);
the $d_{\text{head}}$-split Cases~4--5 pay a partial-score all-reduce
($\sim$\SI{166}{\mega\byte}/token) that the both-axes-$S_{kv}$ split of
Case~6~$\star$ avoids ($\sim$\SI{6.2}{\mega\byte}/token, merging only the
softmax state). On these two axes Case~6 (marked $\star$ in the figure)
is the single placement that lands both inside the memory budget and at
low per-token communication; whether that combination is the right one to
pick is a regime-specific question taken up next.}
\label{fig:gqa-budget}
\end{figure}

These two costs---per-PE KV memory and per-token communication---are
intrinsic properties of each placement, fixed by how it shards the cache
and independent of the workload regime. They do not by themselves name a
winner: a short prompt where the whole cache fits on one PE values low
communication and tolerates replication, whereas a million-token decode
is bound by the memory wall and will pay communication to escape it.
Which placement is appropriate is therefore a per-regime question, which
the short- and long-context subsections that follow answer by running the
options on the simulator and reading off latency, traffic, and the
redundant compute each one induces.

\subsection{Inference with Short-Context Length}
\label{sec:gqa-short}

The fused GQA kernel issues its matrix products as scheduler-managed
composite commands and keeps the online-softmax merge and the cross-device
KV reduction inside the kernel, on PE\_IPCQ. Two kernel families cover the
two phases. The \emph{prefill} kernel is head-parallel and rotates the KV
shards around an inter-CUBE ring (``Ring KV''). The \emph{decode} kernel
is head-replicated with a statically sharded KV cache and reduces partial
attention outputs through an M-fold intra-CUBE chain and, for multiple
users, a two-level reduce-to-root. Two further primitives make long
context practical: a \emph{lazy load} that issues the KV \textsf{DMA\_READ}
and returns immediately, auto-waiting only at first use so that KV load
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
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).
% Bench code: src/kernbench/benches/gqa_helpers/short_ctx/

% TODO: prefill performance figure (latency, stage breakdown).
% TODO: decode performance figure (latency, stage breakdown).
%       Bench output for short_ctx to be generated.

\subsection{Inference with Long-Context Length}
\label{sec:gqa-long}

% TODO: prefill long-context kernel implementation description
% (Sequence-Parallel partition of S_kv, per-case mechanics).
% Bench code: src/kernbench/benches/gqa_helpers/long_ctx/

% TODO: prefill long-context performance figure.

Long-context decode is the regime where the KV cache, not attention
compute, sets serving cost, so the placement question of
\S\ref{sec:gqa-placement} becomes decisive here. To pick the right
placement for this regime we run each of the six options as a fused
decode kernel on the simulator (one decode step on the LLaMA-3.1-70B
single-KV-head-group target, $C{=}8$ CUBEs $\times$ $P{=}8$ PEs) at a
tractable $S_{kv}{=}8192$, and read off end-to-end latency, on-device op
traffic, and the redundant compute each one induces. The swept context is
small enough that all six fit in memory at $S_{kv}{=}8192$; the
placements the long-context memory budget rules out (Cases~1--3) are
drawn in red, run here only to expose their issue and communication
structure.

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_latency.png}
\caption{Measured end-to-end decode latency per placement
($S_{kv}{=}8192$; red = ruled out by the long-context memory budget,
blue = the predicted Pareto choice). The fastest raw latency belongs to
Case~3 (\SI{17.8}{\micro\second})---but Case~3 replicates the full KV
cache into every CUBE, so it overflows the per-PE budget at production
context and wastes 8$\times$ the compute
(Figure~\ref{fig:gqa-6cases-par}). Among the placements that actually fit
1\,M-token memory (Cases~4--6), Case~6~$\star$ is the fastest
(\SI{30.6}{\micro\second}, versus \SI{31.4}{} and \SI{34.5}{\micro\second}
for the $d_{\text{head}}$-split Cases~4 and~5)---making it the placement
of choice for long-context decode.}
\label{fig:gqa-6cases-lat}
\end{figure}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_parallelism.png}
\caption{Redundant compute per placement, measured as active-PE
$\times$ $S_{\text{local}}$ (PE-tokens; lower means less wasted work).
The minimum is \num{8192} PE-tokens---one pass over the sequence.
Case~3 inflates this 8$\times$ to \num{65536} by replicating the KV
cache across all eight CUBEs so every CUBE redundantly re-attends the
whole sequence; the $d_{\text{head}}$-split Cases~4--5 likewise carry
\num{65536} because each token is processed across eight head slices.
Case~6~$\star$ achieves the full 64-way split at the minimal
\num{8192} PE-tokens---fully parallel, no replication.}
\label{fig:gqa-6cases-par}
\end{figure}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_traffic.png}
\caption{Measured on-device op traffic per placement. The unsharded
Case~1 issues no IPCQ copies (each PE has the full cache, nothing to
reduce); the single-axis Case~3 charges 168. Among the 64-way splits,
the $d_{\text{head}}$-split Cases~4--5 charge the most---280 IPCQ copies
and 8 DMA writes each, the partial-score all-reduce that head-slicing
forces---while Case~6~$\star$ needs only 189 IPCQ copies and a single DMA
write, because it merges just the running softmax state $(m,\ell,O)$
rather than partial scores. This is the on-device collective traffic that
PE\_IPCQ and the torus links of \S\ref{sec:allreduce} are provisioned to
absorb at link speed.}
\label{fig:gqa-6cases-traffic}
\end{figure}

The measurements select the right placement for this regime. The fastest
raw latency (Case~3, \SI{17.8}{\micro\second}) comes from replicating the
full KV cache into every CUBE---which is exactly the placement the
long-context memory budget forbids, and which the parallelism panel shows
wastes 8$\times$ the compute. Restricting attention to the placements that
fit production-context memory (the 64-way splits, Cases~4--6), the choice
is Case~6~$\star$: it is the fastest of the three
(\SI{30.6}{\micro\second}), and the op-count panel shows why---its
softmax-state-only reduction charges 189 IPCQ copies and one DMA write
against the 280 copies and 8 DMA writes the $d_{\text{head}}$-split
Cases~4--5 pay for their partial-score all-reduce. For long-context
decode, then, the appropriate data placement is the both-axes sequence
shard (Case~6): it is the cheapest-communicating member of the only
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 \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
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{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
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}

These panels are the clearest statement of the codesign thesis in the
report. Because the composite command keeps GEMM issue cheap and the MAC
array barely occupied, the fused attention kernel's latency is set almost
entirely by data movement: streaming the KV cache and reducing partials
across devices. That is precisely the cost that the communication-side
work targets---PE\_IPCQ for the on-device reduction, the lazy load for
load/compute overlap, fast TCM staging and torus links for the reduction
itself. In other words, the two enablers are not independent features that
happen to appear in the same kernel; the GEMM optimization is what
\emph{exposes} the data-movement bottleneck (by removing the compute and
issue overhead that would otherwise hide it), and the communication
optimization is what \emph{attacks} it. For an attention-dominated decoder
the meaningful hardware investments are therefore the ones that move data
faster and reduce it on-device---not additional MAC throughput, which this
workload cannot use.

% TODO: cross-regime DP (data parallelism) applicability:
%   - Does Case-4 long-context placement compose with batch-level DP
%     without further changes?
%   - Does the short-context placement compose the same way?
%   - Implications for multi-user serving (single vs. mixed regimes).