\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 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 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. % TODO: CUBE <-> KV-head mapping diagram for the short-context regime % (h_kv=8 KV heads -> 8 CUBEs, 1:1; intra-CUBE PE usage). % Bench code: src/kernbench/benches/gqa_helpers/short_ctx/ % TODO: prefill performance figure (latency, stage breakdown). % TODO: decode performance figure (latency, stage breakdown). % Bench output for short_ctx to be generated. \subsection{Inference with Long-Context Length} \label{sec:gqa-long} % TODO: prefill long-context kernel implementation description % (Sequence-Parallel partition of S_kv, per-case mechanics). % Bench code: src/kernbench/benches/gqa_helpers/long_ctx/ % TODO: prefill long-context performance figure. 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 \emph{at the 64-way production scale} (a single-rank caveat follows, Figure~\ref{fig:gqa-decode-stream}); 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{Isolating the rank: the masked streaming win.} That neutrality is a property of the \emph{full 64-way} critical path, not of the local attention: at production scale the inter-CUBE $(m,\ell,O)$ reduce tail and shared-HBM contention set the wall clock, so a faster local attention does not surface. Stripping those away---a single rank, no cross-CUBE reduce, swept over the per-rank context $S_{kv}$ (so $S_{kv}{=}16$\,K here is the per-PE load of a 1M-token, 64-way-sharded decode)---exposes the local attention directly (Figure~\ref{fig:gqa-decode-stream}), and the composite \emph{does} win, by \SI{25}{}--\SI{28}{\percent}. The reason is the memory-bound mirror of prefill: its scheduler-streamed concurrent per-tile DMAs keep the HBM pipeline full and reach \SI{233}{\giga\byte\per\second} ---\SI{91}{\percent} of the per-rank \SI{256}{\giga\byte\per\second} roofline---whereas the primitive kernel's blocking \textsf{tl.dot} serializes one tile DMA at a time and plateaus at \SI{166}{\giga\byte\per\second}. So even for memory-bound decode the composite is not \emph{only} a CPU-issue optimization---it also extracts bandwidth---but that latency benefit materializes only when the local attention is on the critical path, which at 64-way production scale it is not. \begin{figure}[t] \centering \includegraphics[width=\linewidth]{gqa_decode_streaming.png} \caption{Single-rank memory-bound decode ($T_q{=}1$, $M{=}8$), three command forms, swept over per-rank context. \emph{Left:} end-to-end latency---the composite forms run \SI{25}{}--\SI{28}{\percent} below the primitive, a gap that widens with context. \emph{Right:} achieved HBM bandwidth against the per-rank \SI{256}{\giga\byte\per\second} roofline. The primitive's blocking load$\rightarrow$dot serializes the KV stream and plateaus at \SI{166}{\giga\byte\per\second}; the composite forms pipeline concurrent per-tile DMAs through the scheduler and reach \SI{233}{\giga\byte\per\second}. This is the memory-bound mirror of the prefill result (Figure~\ref{fig:gqa-prefill-cb}): there the composite approaches the MAC roofline, here the bandwidth roofline. Capped at $16$\,K---the plain composite materializes the full $(M,S_{kv})$ scores in TCM, so beyond that only the \textsf{softmax\_merge} recipe, which tiles the softmax, stays within scratch.} \label{fig:gqa-decode-stream} \end{figure} \paragraph{The compute-bound mirror: prefill.} Decode's \emph{production} verdict---command form is latency-neutral at 64-way scale---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).