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}

% TODO: compare data-placement options that apply across both the
% short- and long-context regimes. Candidate axes:
%   - KV cache:  per-CUBE shard vs. replicate; per-PE shard vs. replicate
%   - Q / W_qkv / W_o weights:  static partition across CUBEs and PEs
%   - Workspace (m, l, O softmax state):  scratch arena placement
% The 4-case taxonomy used for long-context decode in
% \S\ref{sec:gqa-long} (Cube-{SP,Repl} x PE-{TP,SP}) is one instantiation
% of this framework; the short-context mapping in \S\ref{sec:gqa-short}
% is another.

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

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