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

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

\subsection{Why it is needed}

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. 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{Design}

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.

\subsection{Results}

We measure four headline panels that vary the user count $C$ and the phase:
single- and multi-user prefill ($T_q=4$, $S_{kv}=16$), and single- and
multi-user decode ($P=8$ PEs, $S_{kv}=64$ and $128$), all at $d_{\text{head}}=64$
and $G=8$. For each panel we harvest end-to-end latency
(max event end minus min event start, the same window convention as the
GEMM study) together with the per-engine busy time and the operation mix.
Figure~\ref{fig:gqa-lat} and Figure~\ref{fig:gqa-break} report the result;
the underlying numbers are in Table~\ref{tab:gqa}.

\begin{table}[t]
\centering
\caption{Fused GQA per-panel latency and operation mix. Compute (GEMM,
MATH) is a tiny fraction of DMA occupancy; IPCQ copies grow with users and
PEs.}
\label{tab:gqa}
\small
\begin{tabular}{@{}lrrrr@{}}
\toprule
\textbf{Panel} & \textbf{Lat.\ (ns)} & \textbf{GEMM} & \textbf{IPCQ} & \textbf{DMA rd} \\
\midrule
prefill C=1        & 445   & 2  & 0  & 3  \\
prefill C=4 (Ring) & 4630  & 32 & 24 & 12 \\
decode C=1, P=8    & 3632  & 16 & 21 & 24 \\
decode C=4, P=8    & 6693  & 64 & 93 & 96 \\
\bottomrule
\end{tabular}
\end{table}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_latency_by_panel.png}
\caption{Fused GQA end-to-end latency. Latency grows from
\SI{445}{\nano\second} (single-user prefill) to \SI{6693}{\nano\second}
(four-user decode) as the KV history and the number of participating
devices grow.}
\label{fig:gqa-lat}
\end{figure}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_op_engine_breakdown.png}
\caption{Where the work goes. Left: operation counts---GEMM and IPCQ-copy
volume both scale with users and PEs. Right: summed engine occupancy on a
log scale---the DMA engine dominates by two to three orders of magnitude
over the GEMM and MATH engines in every panel.}
\label{fig:gqa-break}
\end{figure}

The dominant observation is in Figure~\ref{fig:gqa-break}: the compute
engines are almost idle. The GEMM engine accumulates only
\SIrange{2}{33}{\nano\second} of busy time across the panels and the
vector-math engine \SIrange{5}{688}{\nano\second}, while the DMA engine
accumulates \SIrange{72}{15920}{\nano\second}. Fused GQA, as modeled here,
is overwhelmingly data-movement bound. The operation mix shows why the
collective machinery matters: IPCQ-copy count rises from zero (single-user
prefill) to 93 (four-user decode) as the kernel reduces partial outputs
across more PEs and CUBEs, and DMA-read count rises in step as more KV
shards are streamed. The PE control-processor dispatch cost registered as
zero in this configuration---command issue is simply not on the critical
path when data movement is this dominant.

\subsection{Analysis and meaning}

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.