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
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.

% 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{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).