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}

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.

This is a different decomposition from the six-case placement taxonomy of
\S\ref{sec:gqa-placement}: that taxonomy spreads a single KV-head group
across all 64 PEs and is the right lens for long context
(\S\ref{sec:gqa-long}), whereas short context---where the whole cache fits
comfortably---turns instead on how the eight KV heads are distributed
across the eight CUBEs. 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, the per-CUBE KV
footprint grows linearly and the mapping uses proportionally fewer CUBEs
(eight down to one), 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}.
The mappings separate cleanly, and in proportion to how many PEs each
dedicates to a head: \textsf{1-kv-per-cube} spreads one head across all
eight PEs of its CUBE, \textsf{8-kv-per-cube} gives each head a single PE,
and the four mappings form a $\{64,32,16,8\}$-active-PE ladder. Decode
latency tracks that ladder inversely---at $S_{kv}{=}64$K the
\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube} ratio is $7.4\times$
($\SI{87.0}{}$ vs.\ $\SI{11.8}{\micro\second}$), and even at $8$K it is
$4.8\times$ ($\SI{10.8}{}$ vs.\ $\SI{2.2}{\micro\second}$); prefill shows
the same ordering more gently ($2.3\times$ at $64$K). The reason is that
decode here is \emph{bandwidth-bound}: each active PE streams its KV shard
at \SI{46}{}--\SI{76}{\percent} of the \SI{256}{\giga\byte\per\second}
per-PE ceiling (rising with context), so a mapping's speed is set by how
many PEs it puts to work---more PEs, more aggregate HBM bandwidth, lower
latency. The $M{=}G{=}8$ score GEMM is skinny and leaves the MAC array
lightly loaded throughout; the lever here is data movement, not compute.

\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; log scale). The mappings separate along the
$\{64,32,16,8\}$-active-PE ladder set by heads-per-CUBE:
\textsf{1-kv-per-cube} (64 PEs) is fastest and \textsf{8-kv-per-cube}
(8 PEs) slowest, by $4.8\times$ at $8$K growing to $7.4\times$ at $64$K in
decode ($2.3\times$ at $64$K in prefill). Decode is bandwidth-bound, so
latency scales inversely with the number of active PEs. Composite tiers
track the baseline at this skinny shape
(Figure~\ref{fig:gqa-short-a1-variant}).}
\label{fig:gqa-short-wall-baseline}
\end{figure}

Figure~\ref{fig:gqa-short-tradeoff} shows what the mappings trade for that
latency. The differentiator is \emph{density}: \textsf{8-kv-per-cube}
concentrates all eight heads onto one CUBE, so its per-CUBE KV footprint is
$8\times$ that of \textsf{1-kv-per-cube} (one CUBE per head), and it
occupies a single one of the SIP's sixteen CUBEs against
\textsf{1-kv-per-cube}'s eight. Per-PE HBM \emph{utilization}, by contrast,
is similar across mappings---every active PE runs near the same
\SI{46}{}--\SI{76}{\percent} of the per-PE ceiling---so the latency ladder
comes from the \emph{count} of active PEs, not from any per-PE
concentration. IPCQ traffic runs opposite to density:
\textsf{1-kv-per-cube} pays the eight-PE chain-reduce while
\textsf{8-kv-per-cube} (one PE per head) pays none, but absolute IPCQ volume
stays under \SI{120}{\kibi\byte} per run, two orders of magnitude below the
HBM traffic. The result is a genuine trade-off rather than a single winner:
\textsf{1-kv-per-cube} minimizes per-request latency by spending the most
hardware (eight CUBEs, 64 PEs) on one request, while denser mappings leave
CUBEs free for other requests---the batched-serving axis taken up below.

\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;
log scale). Top row: per-CUBE KV footprint---\textsf{8-kv-per-cube} packs
all eight heads onto one CUBE, an $8\times$ larger per-CUBE KV than
\textsf{1-kv-per-cube} and one CUBE used against eight (the $8\times$ ratio
is context-invariant; the absolute footprint grows with $S_{kv}$). Bottom
row: IPCQ traffic per run---\textsf{1-kv-per-cube} pays the eight-PE
chain-reduce while \textsf{8-kv-per-cube} (single PE per head) pays zero.
The density (top) is what sets how many concurrent requests a SIP can hold.}
\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 nearly flat, with
the GEMM-only composite tier a few percent \emph{slower}
(\SI{14.0}{} vs.\ \SI{11.8}{\micro\second} at $64$K decode) and the
\textsf{softmax\_merge} tier matching the baseline. 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 is pure overhead here. It shows up in
Figure~\ref{fig:gqa-short-gemm-util} as inflated GEMM-engine busy time on
the composite tiers---the padded-$M$ GEMM is charged in full against a
sub-\SI{15}{\micro\second} decode wall, driving the accounted utilization
well above the baseline's \SI{12}{}--\SI{18}{\percent}. Fusion benefit
requires lifting $M$, either by the long-context Q-tile split
(\S\ref{sec:gqa-composite}) or by batching multiple users---which is also
the lever that converts a mapping's density into throughput.

\paragraph{Latency versus batched throughput (projected).} The latency
ladder and the density trade-off pull in opposite directions once a SIP
serves more than one request. Decode users are independent---each attends
its own KV cache---and these mappings generate no cross-CUBE traffic, so a
SIP can run several users at once on disjoint CUBE groups, up to $16/C$
users: two for \textsf{1-kv-per-cube}, sixteen for \textsf{8-kv-per-cube}.
Projecting from the measured single-user latencies under perfect overlap
(an idealization---no concurrent multi-user run has been measured yet; see
\S\ref{sec:future}), aggregate throughput at a full SIP is
$\text{users}/\text{latency}$: at $8$K it rises from \SI{0.89}{}
(\textsf{1-kv-per-cube}, 2 users) to \SI{1.48}{}~requests\,/\,\si{\micro\second}
(\textsf{8-kv-per-cube}, 16 users), so the \emph{dense} mapping wins on
throughput even though it is the \emph{slowest} per request; at $64$K the
projected throughputs converge (\SI{0.17}{}--\SI{0.18}{}) because latency
there already scales as $1/C$. The design conclusion is therefore
regime-dependent: \textsf{1-kv-per-cube} is the latency-optimal choice for
single-stream or low-batch decode, while denser mappings are the
throughput-optimal choice for high-batch short-context serving. These are
\emph{projections} assuming zero cross-user interference; validating them
against a concurrent multi-user run is 2H work (\S\ref{sec:future}).

\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} (log scale):
wall-clock is nearly flat across tiers---the GEMM-only composite runs a few
percent slower (padding overhead) and \textsf{softmax\_merge} matches the
baseline. With $M{=}G{=}8 < \textsf{TILE\_M}{=}32$ the composite scheduler
pads $M$ by $4\times$, leaving no fusion slack to recover at this shape.
Fusion pays off only when $M$ fills the supertile---long-context Q-tile
splitting or batched-$M$ inference.}
\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 busy fraction on \textsf{1-kv-per-cube} across the
three composite tiers. The baseline runs at \SI{12}{}--\SI{18}{\percent};
the composite tiers read much higher---up to and past
\SI{100}{\percent} of the short decode wall---because the padded
$8\!\to\!32$ $M$ dimension is charged as GEMM time, bookkeeping overhead
rather than useful work. The engine is never the bottleneck at this shape;
decode is KV-bandwidth-bound.}
\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}

Taken together the three panels select the placement. The only
memory-feasible family at production context is the 64-way splits
(Cases~4--6), and within it Case~6~$\star$ is both the fastest and the
lightest-communicating, because it merges only the running softmax state
rather than the partial scores that the $d_{\text{head}}$-split
Cases~4--5 must all-reduce. Case~3's lower raw latency is beside the
point---it replicates the full cache into every CUBE, overflowing the
per-PE budget and wasting $8\times$ the compute. For long-context decode
the placement of choice is therefore the both-axes sequence shard, 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{Summary}
\label{sec:gqa-analysis}

The section's results line up into one picture of the fused kernel.
Placement is decided by memory first and communication second: only the
64-way splits fit production context, and among them the both-axes
sequence shard (Case~6) minimizes the per-token collective by merging
softmax state rather than partial scores. Command form is decided by
roofline position: the composite command removes the per-tile issue work
in both regimes, but converts to wall-clock only in compute-bound prefill,
while memory-bound decode stays bound by KV streaming regardless of command
form. The thread connecting the two is that, once composite issue makes
GEMM cheap and leaves the MAC array idle, the fused kernel's latency is set
almost entirely by data movement---streaming the KV cache and reducing
partials across PEs---which is exactly the cost the communication-side
work (on-device reduction, lazy load/compute overlap, fast TCM staging and
torus links) is built to attack. What this implies for hardware investment
across the report as a whole is taken up in the discussion.

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