ywkang
b715002b5a
§2 platform:
- Accuracy promoted from \paragraph to \subsection (sec:accuracy);
it sits at the same heading level as Why / Device / Latency / HW
config and so reads as a first-class component of the platform
description rather than a tail-end footnote.
- Three architecture diagrams (SIP, CUBE, PE) collapsed into one
figure* using subcaption: Fig.~\ref{fig:hw-arch}(a) SIP and (b)
CUBE side-by-side on row 1, (c) PE wide on row 2, total height
capped at ~40% of the page. Subfigure cross-references rewritten
to Fig.~ref{fig:hw-arch}\subref{...} in the body text. main.tex
now pulls in the subcaption package.
- Table 1 (modeled hardware configuration) promoted to table*
(two-column / full page width) and the row layout rewritten as a
4-column tabular so Hierarchy + PE + Command-issue sit on the
left and Memory + Interconnect sit on the right. The previous
single-column rendering was getting cut at the right edge of the
printed column.
§6 GQA:
- Removed the historical four-panel headline table (tab:gqa), the
two figures (fig:gqa-lat, fig:gqa-break), and the prose paragraph
that cited their per-panel numbers. Reason: the underlying
milestone-gqa-headline bench has been simplified in collaborator
commit 65c365f ("drop misleading single_user_/multi_user_
panels") and no longer reproduces that data, so the section was
left referencing a dataset the current bench cannot regenerate.
- The previously-added 4-cases long-context decode comparison is
now §6's only Results subsection, retitled to
"Results: long-context decode and parallelism strategies".
The headline GQA result that survives into the paper is therefore
the parallelism trade-off study (Case 4 / Cube-SP × PE-SP at 34 us
buying an 8x KV-memory reduction over the latency leaders).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
136 lines
7.1 KiB
TeX
136 lines
7.1 KiB
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: long-context decode and 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.
|