eusonice
486 lines
27 KiB
TeX
486 lines
27 KiB
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 splits the query tile across the PEs
|
|
of a group and broadcasts each KV tile from the group's root PE over
|
|
intra-CUBE IPCQ, so the HBM K/V read is paid once per group instead of
|
|
once per PE. The \emph{decode} kernel sequence-shards the KV cache across
|
|
the same group of PEs, runs per-PE local attention, then chain-reduces the
|
|
partial $(m,\ell,O)$ triples back up to the group 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 KV load overlaps score computation; and per-tile \emph{scratch
|
|
recycling} that keeps the running accumulators 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
|
|
restructures the decode step so the per-tile matrix products and the
|
|
online-softmax merge are issued as \emph{composite} commands rather than
|
|
hand-tiled primitives; \S\ref{sec:gqa-composite} measures it against the
|
|
primitive baseline at long context.
|
|
|
|
The design question for short context is therefore not whether to fuse
|
|
softmax---that is settled---but how to distribute the eight KV heads
|
|
across the eight CUBEs of a SIP. Four mappings cover the spectrum, named
|
|
by KV-head count per CUBE: \textsf{1-kv-per-cube} dedicates one whole
|
|
CUBE to each head and uses all eight PEs of that CUBE on the head's
|
|
sequence shard; \textsf{2-kv-per-cube} packs two heads per CUBE with
|
|
group size four; \textsf{4-kv-per-cube} packs four heads with group size
|
|
two; \textsf{8-kv-per-cube} loads all heads onto a single CUBE with one
|
|
PE per head. The mapping fixes a trade-off: as KV heads per CUBE grows,
|
|
HBM bandwidth pressure per CUBE grows linearly, the per-CUBE KV cache
|
|
grows linearly, but the intra-CUBE IPCQ broadcast/reduce cost shrinks to
|
|
zero (8-kv-per-cube uses a single PE per head and has no group
|
|
communication).
|
|
|
|
A second axis is how each GEMM tile is issued. We isolate this with three
|
|
\emph{composite tiers} applied to the same mapping: \textsf{without
|
|
composite} uses primitive \textsf{tl.dot} and unfused softmax;
|
|
\textsf{with composite (GEMM-only)} issues the GEMMs as
|
|
\textsf{tl.composite(op="gemm")} commands so the scheduler can pipeline
|
|
the tile stages but leaves softmax as primitives; \textsf{with composite
|
|
+ softmax\_merge} additionally folds the softmax fold into the $P\!\cdot\!V$
|
|
composite through a named \textsf{softmax\_merge} prologue recipe.
|
|
|
|
We sweep all four mappings $\times$ three composite tiers $\times$
|
|
$S_{kv}\in\{8\text{K},16\text{K},32\text{K},64\text{K}\}$ for both
|
|
phases (prefill uses $T_q{=}8$ sliced tiles; decode uses $T_q{=}1$). The
|
|
headline latency comparison is in Figure~\ref{fig:gqa-short-wall-baseline}.
|
|
At every context length the \textsf{1-kv-per-cube} mapping is the
|
|
fastest, and the $\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube}$ ratio
|
|
grows monotonically with $S_{kv}$ (decode: $0.98\times$ at $8$K,
|
|
$1.06\times$ at $64$K; prefill follows the same trend). The wall gap is
|
|
modest because short-context decode is dominated by setup overhead and
|
|
its $M{=}G{=}8$ skinny shape leaves the MAC array idle most of the time;
|
|
where the mappings really separate is per-CUBE pressure.
|
|
|
|
\begin{figure}[t]
|
|
\centering
|
|
\includegraphics[width=\linewidth]{gqa_short_context/wall_variant1_mode_compare.png}
|
|
\caption{Wall-clock latency of the four short-context mappings, baseline
|
|
kernel (no composite). Prefill and decode follow the same ordering at
|
|
every context length and the \textsf{1-kv-per-cube} mapping is the
|
|
fastest throughout; the gap to \textsf{8-kv-per-cube} grows from
|
|
$0.98\times$ at $8$K to $1.06\times$ at $64$K. The composite and
|
|
softmax\_merge tiers track the baseline within a few percent at this
|
|
shape (see Figure~\ref{fig:gqa-short-a1-variant}).}
|
|
\label{fig:gqa-short-wall-baseline}
|
|
\end{figure}
|
|
|
|
Figure~\ref{fig:gqa-short-tradeoff} makes the per-CUBE trade-off
|
|
explicit. Per-PE HBM bandwidth utilization is $\sim$\,$8\times$ higher on
|
|
\textsf{8-kv-per-cube} than on \textsf{1-kv-per-cube} at every context
|
|
length: putting all heads on a single CUBE saturates that CUBE's HBM
|
|
channel while leaving seven CUBEs idle, whereas spreading one head per
|
|
CUBE keeps each CUBE's read bandwidth well below the
|
|
\SI{256}{\giga\byte\per\second} per-PE ceiling. IPCQ traffic moves in the
|
|
opposite direction---\textsf{1-kv-per-cube} pays the broadcast or
|
|
chain-reduce cost across all eight PEs, while \textsf{8-kv-per-cube} pays
|
|
nothing---but absolute IPCQ volume stays under \SI{120}{\kibi\byte} per
|
|
run, two orders of magnitude below the HBM traffic. Because per-PE HBM
|
|
bandwidth and per-CUBE KV cache (which grows as $\text{kv\_per\_cube}\times
|
|
S_{kv}\times d_{\text{head}}$, so $8\times$ more on
|
|
\textsf{8-kv-per-cube}) both scale with the number of heads on a CUBE,
|
|
the small wall advantage of \textsf{1-kv-per-cube} compounds with much
|
|
larger headroom for longer contexts and for batched serving---which is
|
|
why it is the choice carried into long context (\S\ref{sec:gqa-long}).
|
|
|
|
\begin{figure}[t]
|
|
\centering
|
|
\includegraphics[width=\linewidth]{gqa_short_context/per_cube_tradeoff_mode_compare.png}
|
|
\caption{Per-CUBE trade-off across the four mappings (baseline kernel;
|
|
HBM and IPCQ totals are variant-invariant). Top row: per-PE HBM
|
|
bandwidth utilization---\textsf{8-kv-per-cube} concentrates roughly
|
|
$8\times$ more bandwidth on a single CUBE than \textsf{1-kv-per-cube}.
|
|
Bottom row: IPCQ traffic per run on a log scale---\textsf{1-kv-per-cube}
|
|
pays the broadcast/reduce cost across the largest group while
|
|
\textsf{8-kv-per-cube} (single PE per head) pays zero.}
|
|
\label{fig:gqa-short-tradeoff}
|
|
\end{figure}
|
|
|
|
The composite ablation in Figure~\ref{fig:gqa-short-a1-variant} shows
|
|
the three tiers within \textsf{1-kv-per-cube}: wall-clock is essentially
|
|
unchanged across tiers ($< 0.5\%$ in either direction). This is the
|
|
expected outcome for the decode-skinny shape---$M{=}G{=}8$ is well below
|
|
the scheduler's $\textsf{TILE\_M}{=}32$ supertile (\S\ref{sec:gemm}), so
|
|
the composite path pads $M$ by $4\times$ with zeros and the fusion has
|
|
no slack to win back. The padding shows up in
|
|
Figure~\ref{fig:gqa-short-gemm-util} as inflated GEMM engine utilization
|
|
on the composite and softmax\_merge tiers (roughly $3$ and $6\times$ the
|
|
baseline), which is bookkeeping cost rather than added useful work.
|
|
Demonstrating fusion benefit therefore requires lifting $M$---either by
|
|
batching multiple users (a separate batched-$M$ kernel, deferred to
|
|
\S\ref{sec:future}) or by the long-context path, where Q-tile splitting
|
|
yields a larger $M$ per PE. In the short-context regime the contribution
|
|
of composite/softmax\_merge is a no-regression dispatch-cost isolation;
|
|
the substantive design lever is the mapping, and
|
|
\textsf{1-kv-per-cube} wins it.
|
|
|
|
\begin{figure}[t]
|
|
\centering
|
|
\includegraphics[width=\linewidth]{gqa_short_context/wall_a1_variant_compare.png}
|
|
\caption{Composite-tier ablation on \textsf{1-kv-per-cube}: wall-clock
|
|
latency is flat across the three tiers within $0.5\%$. With
|
|
$M{=}G{=}8 < \textsf{TILE\_M}{=}32$, the composite scheduler pads the
|
|
$M$ dimension $4\times$ and any fusion overlap is absorbed by the
|
|
padding overhead. Fusion benefit is exercised only when $M$ fills the
|
|
supertile---batched-$M$ inference or long-context Q-tile splitting.}
|
|
\label{fig:gqa-short-a1-variant}
|
|
\end{figure}
|
|
|
|
\begin{figure}[t]
|
|
\centering
|
|
\includegraphics[width=\linewidth]{gqa_short_context/gemm_util_a1_variant_ablation.png}
|
|
\caption{GEMM engine utilization on \textsf{1-kv-per-cube} across the
|
|
three composite tiers. The $\sim 3$ and $\sim 6\times$ rise on the
|
|
composite and softmax\_merge tiers reflects supertile padding (the
|
|
$8\!\to\!32$ pad in $M$) being accounted as GEMM time, not added useful
|
|
work. With actual $M=8$, the engine is below \SI{2}{\percent} busy in
|
|
all tiers---the kernel is firmly setup- and KV-bandwidth-bound at this
|
|
shape.}
|
|
\label{fig:gqa-short-gemm-util}
|
|
\end{figure}
|
|
|
|
\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{Use of Composite Commands}
|
|
\label{sec:gqa-composite}
|
|
|
|
The decode kernel of \S\ref{sec:gqa-long} issues its local attention as
|
|
primitive operations: it walks each PE's $S_{\text{local}}$ token slice in
|
|
\SI{1024}{}-token tiles, and for every tile issues a $Q\!\cdot\!K^{\top}$
|
|
\textsf{dot}, the online-softmax primitives, and a $P\!\cdot\!V$
|
|
\textsf{dot}, merging the running $(m,\ell,O)$ state by hand. The number
|
|
of PE\_CPU commands this costs grows with the context. At a production
|
|
context of $S_{kv}{=}1\,\text{M}$ tokens the Case-6 64-way split gives
|
|
each PE $S_{\text{local}}{=}16384$ tokens, so its local attention is a
|
|
$Q\!\cdot\!K^{\top}$ of $(8,128)\!\cdot\!(128,16384)$ and a
|
|
$P\!\cdot\!V$ of $(8,16384)\!\cdot\!(16384,128)$---sixteen hand-issued
|
|
tiles, each a fresh batch of CPU commands.
|
|
|
|
The composite command lets the kernel hand that tiling to PE\_SCHEDULER.
|
|
We compare three command forms of the \emph{same} Case-6 kernel---identical
|
|
placement and identical $(m,\ell,O)$ reduce, differing only in how the
|
|
local attention is issued:
|
|
\begin{itemize}
|
|
\item \textbf{primitive}---the hand-tiled \textsf{dot}/softmax kernel
|
|
above (the baseline of \S\ref{sec:gqa-long}).
|
|
\item \textbf{composite}---each matrix product is one coarse
|
|
\textsf{composite} GEMM over the \emph{whole} $S_{\text{local}}$, with
|
|
$K$ and $V$ passed as HBM references so PE\_SCHEDULER streams and
|
|
tiles them on the fixed $32\!\times\!64\!\times\!32$ MAC tile; the
|
|
softmax stays primitive.
|
|
\item \textbf{composite\,+\,softmax\_merge}---additionally folds the
|
|
online-softmax merge and $P\!\cdot\!V$ into a single stateful
|
|
\textsf{softmax\_merge} recipe composite.
|
|
\end{itemize}
|
|
|
|
\begin{figure}[t]
|
|
\centering
|
|
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_composite.png}
|
|
\caption{Three command forms of the Case-6 decode kernel, swept over
|
|
context length ($S_{\text{local}}{=}S_{kv}/64$ per PE). \emph{Right:}
|
|
PE\_CPU commands issued. The hand-tiled primitive kernel rises
|
|
$O(n_{\text{tiles}})$---from 96 commands at one tile to 426 at the
|
|
1\,M-token, sixteen-tile production point---while both composite forms
|
|
issue a context-\emph{independent} $O(1)$ count (94 and 98) that
|
|
\emph{saturates}: one coarse descriptor offloads the entire per-tile
|
|
fan-out. \emph{Left:} the consequence for wall-clock latency is none---all
|
|
three land on the same curve (\SI{30.6}{}, \SI{231}{},
|
|
\SI{461}{\micro\second} at 8\,K\,/\,64\,K\,/\,128\,K), because decode is
|
|
bound by streaming the KV cache, not by issue. Command-count is measured
|
|
at emit time (exact, to 1\,M); latency on the data-mode engine over the
|
|
tractable range.}
|
|
\label{fig:gqa-composite}
|
|
\end{figure}
|
|
|
|
Figure~\ref{fig:gqa-composite} reads off the two quantities that matter,
|
|
and they point in opposite directions. The PE\_CPU command count (right)
|
|
collapses from a context-growing $O(n_{\text{tiles}})$ to a flat $O(1)$:
|
|
at 1\,M tokens the composite form issues \num{94} commands against the
|
|
primitive kernel's \num{426}, a $4.5\times$ reduction
|
|
($4.2\times$ in modeled dispatch cost), and---crucially---that number no
|
|
longer grows with context. The wall-clock latency (left), by contrast,
|
|
is unchanged across all three forms: decode is bound by streaming the KV
|
|
cache out of HBM, so the command form does not move the critical path.
|
|
|
|
That juxtaposition is the point. The composite command is not a latency
|
|
optimization for this memory-bound decode; it is a \emph{CPU-issue}
|
|
optimization. Its value is removing the per-tile dispatch work that would
|
|
otherwise grow without bound as context grows, freeing PE\_CPU to run
|
|
ahead and keep the engines fed---which is exactly what lets the
|
|
data-movement cost analyzed next show through as the true bottleneck
|
|
rather than being masked by issue overhead. The \textsf{softmax\_merge}
|
|
recipe folds the online merge into the same descriptor; on this
|
|
memory-bound path its marginal cost over the plain GEMM composite is small
|
|
(98 vs.\ 94 commands), and like the plain composite it keeps the issued
|
|
count flat as context scales.
|
|
|
|
\paragraph{The compute-bound mirror: prefill.} Decode's verdict---command
|
|
form is latency-neutral---is a property of its regime, not of the
|
|
composite command. A decode step has $T_q{=}1$, so its score and context
|
|
products are skinny ($M{=}G\,T_q{=}8$): the MAC array is barely fed and
|
|
the kernel is bound by streaming the KV cache. Prefill is the opposite
|
|
corner. It processes a block of query positions at once, so $M{=}G\,T_q$
|
|
is large and tile-filling, the GEMMs carry real arithmetic intensity
|
|
($\sim$$M$ flops/byte, well above the roofline ridge), and the kernel is
|
|
\emph{compute-bound}. This is the regime the composite command was built
|
|
for (\S\ref{sec:gemm}): it streams the per-HW-tile
|
|
DMA$\rightleftarrows$compute pipeline so the MAC array stays fed, whereas
|
|
the primitive kernel's blocking \textsf{tl.dot} serializes each tile's
|
|
load and compute and starves the array between tiles. We run the same
|
|
three command forms on a single-rank compute-bound prefill (FlashAttention
|
|
$Q$-block $\times$ $S_{kv}$-tile, online softmax) and sweep the context
|
|
length (Figure~\ref{fig:gqa-prefill-cb}).
|
|
|
|
\begin{figure}[t]
|
|
\centering
|
|
\includegraphics[width=\linewidth]{gqa_prefill_compute_bound.png}
|
|
\caption{Compute-bound prefill, three command forms, swept over context
|
|
length ($M{=}8\,T_q$ tile-filling). \emph{Left:} end-to-end latency.
|
|
\emph{Right:} MAC utilization (achieved $\div$ the
|
|
\SI{8}{\tera\flop\per\second} per-PE peak). The hand-tiled primitive sits
|
|
flat at $\sim$\SI{68}{\percent}---its serial load$\to$dot path leaves the
|
|
MAC array idle between tiles regardless of context. The composite forms
|
|
climb with context (\SI{67}{}$\to$\SI{80}{\percent} plain,
|
|
\SI{67}{}$\to$\SI{83}{\percent} with the recipe) because a deeper $P\!\cdot
|
|
\!V$ reduction gives more HW tiles to pipeline, and they convert that into
|
|
wall-clock: at \num{1024} the recipe form is \SI{646.9}{} vs.\
|
|
\SI{794.1}{\micro\second} (\SI{19}{\percent} faster). The margin
|
|
\emph{grows} with context---the compute-bound mirror of the GEMM result of
|
|
\S\ref{sec:gemm}.}
|
|
\label{fig:gqa-prefill-cb}
|
|
\end{figure}
|
|
|
|
The two studies together state the composite command's value precisely. It
|
|
has two distinct benefits, and which one matters is set by the workload's
|
|
roofline position. The first is \emph{host-issue offload}: one macro
|
|
command in place of $O(n_{\text{tiles}})$ fine ones, which removes
|
|
PE\_CPU dispatch work and is regime-independent (it shows in the decode
|
|
command count). The second is \emph{MAC-array feeding}: the
|
|
scheduler-internal per-tile DMA$\rightleftarrows$compute pipeline, which
|
|
only converts to latency when the workload is compute-bound enough to have
|
|
a MAC array worth keeping busy (it shows in the prefill utilization). A
|
|
memory-bound decode exercises only the first; a compute-bound prefill
|
|
exercises both. The composite command is the single mechanism that
|
|
delivers each where it applies.
|
|
|
|
\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).
|