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,
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 short context the two mappings are effectively tied, and they separate
as context grows: the $\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube}$
latency ratio rises monotonically with $S_{kv}$ (decode: $0.98\times$ at
$8$K---within measurement noise---to $1.06\times$ at $64$K, where
\textsf{1-kv-per-cube} is the faster; 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; the
two mappings are effectively tied at $8$K and separate as context grows,
the \textsf{8-kv-per-cube}/\textsf{1-kv-per-cube} ratio widening from
$0.98\times$ at $8$K to $1.06\times$ at $64$K (\textsf{1-kv-per-cube}
faster at the large end). 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}

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