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{Long-context decode: parallelism strategies}

The four headline panels above stress the kernel at moderate context
lengths. Long-context decode---the regime where KV cache size, not
attention compute, sets serving cost---turns the choice of how to
parallelize across cubes and PEs into a first-order design knob. We
compare four strategies on the LLaMA-3.1-70B single-KV-head-group
target (8 CUBEs $\times$ 8 PEs, one KV-head group):

\begin{itemize}\setlength\itemsep{1pt}
  \item \textbf{Case 1} (Cube-SP $\times$ PE-TP): KV split by $S_{kv}$
        across CUBEs; PEs tensor-parallel on the batch dimension
        (wastes PE-TP work at $B{=}1$).
  \item \textbf{Case 2} (Cube-Repl $\times$ PE-TP): full KV
        replicated to every CUBE; PEs tensor-parallel on batch.
  \item \textbf{Case 3} (Cube-Repl $\times$ PE-SP): full KV
        replicated; PEs sequence-parallel on $S_{kv}$ with an
        intra-CUBE all-reduce.
  \item \textbf{Case 4} ($\star$, Cube-SP $\times$ PE-SP): KV split
        64-way (across both CUBEs and PEs) with a two-phase all-reduce
        on the running softmax state $(m, \ell, O)$.
\end{itemize}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_latency.png}
\caption{End-to-end decode latency per parallelism strategy
(LLaMA-3.1-70B single-KV-head group, 8 CUBEs $\times$ 8 PEs).
Replication into CUBEs (Cases 2/3) wins the latency race
(\SI{20.2}{\micro\second} for Case 3), but Case~4 ($\star$, KV split
64-way) finishes within \SI{14}{\micro\second} of the leader while
paying a different cost---visible in Figure~\ref{fig:gqa-4cases-mem}.}
\label{fig:gqa-4cases-lat}
\end{figure}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_memory.png}
\caption{Per-CUBE KV memory footprint for the four cases. Cases 1
and 4---both with the KV cache split across cubes (Cube-SP)---hold
only \SI{0.5}{\mebi\byte} of KV state per CUBE; Cases 2 and 3, which
replicate the full KV, hold \SI{4}{\mebi\byte} per CUBE, an
\textbf{8$\times$} blowup at this configuration that scales linearly
with context length.}
\label{fig:gqa-4cases-mem}
\end{figure}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_traffic.png}
\caption{Per-case op-count breakdown. The replicated-KV PE-TP design
(Case 2) avoids almost all on-device communication
(\textasciitilde0 IPCQ copies), but at the cost of KV memory.
Case~4's two-phase reduce charges \textasciitilde190 IPCQ copies and
\textasciitilde190 DMA reads---this is the traffic that PE\_IPCQ
(\S\ref{sec:allreduce}) is built to absorb at on-device speed.}
\label{fig:gqa-4cases-traffic}
\end{figure}

Three things stand out. First, the fastest case in pure latency
(Case~3, \SI{20.2}{\micro\second}) is also the most memory-hungry,
requiring the full KV state on every CUBE---an option that fails to
scale once context length blows past the per-CUBE budget. Second,
Case~4's KV-split design gives back roughly \SI{14}{\micro\second}
versus Case~3 in exchange for an \textbf{8$\times$} KV-memory
reduction; for practical long-context serving where KV capacity is
the binding constraint, this is the trade the design chooses
(marked $\star$). Third, Case~4 pays its way in
\emph{communication}: the op-count panel shows \textasciitilde190
IPCQ copies and \textasciitilde190 DMA reads, precisely the
on-device collective traffic that PE\_IPCQ and the torus links of
\S\ref{sec:allreduce} are provisioned to move quickly---so the
``slower'' strategy is in fact the one that fully cashes in the
communication-side codesign work of this report.

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