\section{Fused Grouped-Query Attention} \label{sec:gqa} \subsection{Why it is needed} Attention is the 1H focus, and it is where the two preceding optimizations have to come together. Grouped-Query Attention (GQA) shrinks the KV cache by sharing each KV head across a group of query heads (here $h_q=8$ query heads to $h_{kv}=1$ KV head, a group factor $G=8$), which makes decoding feasible at long context but also makes it acutely memory-bound: a decode step processes a single query position ($T_q=1$) against the entire KV history, so its arithmetic intensity is low and its time is dominated by streaming the KV cache out of HBM. FlashAttention-style tiling with an online-softmax merge avoids ever materializing the full score matrix, but realizing it as a fast \emph{fused} kernel needs both building blocks from this report: efficient GEMM issue (\S\ref{sec:gemm}) for the $Q\!\cdot\!K^{\top}$ and $P\!\cdot\!V$ products, and an efficient on-device reduction (\S\ref{sec:allreduce}) for the multi-user and sequence-parallel KV reductions. This section is the capstone: the fused kernel that uses the composite command and PE\_IPCQ at the same time. Multi-head attention (MHA) was studied in prior work and serves here as the established baseline rather than being re-derived. \subsection{Design} The fused GQA kernel issues its matrix products as scheduler-managed composite commands and keeps the online-softmax merge and the cross-device KV reduction inside the kernel, on PE\_IPCQ. Two kernel families cover the two phases. The \emph{prefill} kernel is head-parallel and rotates the KV shards around an inter-CUBE ring (``Ring KV''). The \emph{decode} kernel is head-replicated with a statically sharded KV cache and reduces partial attention outputs through an M-fold intra-CUBE chain and, for multiple users, a two-level reduce-to-root. Two further primitives make long context practical: a \emph{lazy load} that issues the KV \textsf{DMA\_READ} and returns immediately, auto-waiting only at first use so that KV load overlaps score computation; and per-tile \emph{scratch recycling} that keeps the running softmax accumulators ($m,\ell,O$) in a persistent arena while freeing per-tile temporaries, so the kernel fits the \SI{1}{\mebi\byte} scratch budget across many tiles. A further refinement that restructures the decode step into two stateful composites (a named \textsf{softmax\_merge} recipe) is designed but not yet wired into the measured path; results below reflect the implemented kernel only. \subsection{Results} We measure four headline panels that vary the user count $C$ and the phase: single- and multi-user prefill ($T_q=4$, $S_{kv}=16$), and single- and multi-user decode ($P=8$ PEs, $S_{kv}=64$ and $128$), all at $d_{\text{head}}=64$ and $G=8$. For each panel we harvest end-to-end latency (max event end minus min event start, the same window convention as the GEMM study) together with the per-engine busy time and the operation mix. Figure~\ref{fig:gqa-lat} and Figure~\ref{fig:gqa-break} report the result; the underlying numbers are in Table~\ref{tab:gqa}. \begin{table}[t] \centering \caption{Fused GQA per-panel latency and operation mix. Compute (GEMM, MATH) is a tiny fraction of DMA occupancy; IPCQ copies grow with users and PEs.} \label{tab:gqa} \small \begin{tabular}{@{}lrrrr@{}} \toprule \textbf{Panel} & \textbf{Lat.\ (ns)} & \textbf{GEMM} & \textbf{IPCQ} & \textbf{DMA rd} \\ \midrule prefill C=1 & 445 & 2 & 0 & 3 \\ prefill C=4 (Ring) & 4630 & 32 & 24 & 12 \\ decode C=1, P=8 & 3632 & 16 & 21 & 24 \\ decode C=4, P=8 & 6693 & 64 & 93 & 96 \\ \bottomrule \end{tabular} \end{table} \begin{figure}[t] \centering \includegraphics[width=\linewidth]{gqa_latency_by_panel.png} \caption{Fused GQA end-to-end latency. Latency grows from \SI{445}{\nano\second} (single-user prefill) to \SI{6693}{\nano\second} (four-user decode) as the KV history and the number of participating devices grow.} \label{fig:gqa-lat} \end{figure} \begin{figure}[t] \centering \includegraphics[width=\linewidth]{gqa_op_engine_breakdown.png} \caption{Where the work goes. Left: operation counts---GEMM and IPCQ-copy volume both scale with users and PEs. Right: summed engine occupancy on a log scale---the DMA engine dominates by two to three orders of magnitude over the GEMM and MATH engines in every panel.} \label{fig:gqa-break} \end{figure} The dominant observation is in Figure~\ref{fig:gqa-break}: the compute engines are almost idle. The GEMM engine accumulates only \SIrange{2}{33}{\nano\second} of busy time across the panels and the vector-math engine \SIrange{5}{688}{\nano\second}, while the DMA engine accumulates \SIrange{72}{15920}{\nano\second}. Fused GQA, as modeled here, is overwhelmingly data-movement bound. The operation mix shows why the collective machinery matters: IPCQ-copy count rises from zero (single-user prefill) to 93 (four-user decode) as the kernel reduces partial outputs across more PEs and CUBEs, and DMA-read count rises in step as more KV shards are streamed. The PE control-processor dispatch cost registered as zero in this configuration---command issue is simply not on the critical path when data movement is this dominant. \subsection{Long-context decode: parallelism strategies} The four headline panels above stress the kernel at moderate context lengths. Long-context decode---the regime where KV cache size, not attention compute, sets serving cost---turns the choice of how to parallelize across cubes and PEs into a first-order design knob. We compare four strategies on the LLaMA-3.1-70B single-KV-head-group target (8 CUBEs $\times$ 8 PEs, one KV-head group): \begin{itemize}\setlength\itemsep{1pt} \item \textbf{Case 1} (Cube-SP $\times$ PE-TP): KV split by $S_{kv}$ across CUBEs; PEs tensor-parallel on the batch dimension (wastes PE-TP work at $B{=}1$). \item \textbf{Case 2} (Cube-Repl $\times$ PE-TP): full KV replicated to every CUBE; PEs tensor-parallel on batch. \item \textbf{Case 3} (Cube-Repl $\times$ PE-SP): full KV replicated; PEs sequence-parallel on $S_{kv}$ with an intra-CUBE all-reduce. \item \textbf{Case 4} ($\star$, Cube-SP $\times$ PE-SP): KV split 64-way (across both CUBEs and PEs) with a two-phase all-reduce on the running softmax state $(m, \ell, O)$. \end{itemize} \begin{figure}[t] \centering \includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_latency.png} \caption{End-to-end decode latency per parallelism strategy (LLaMA-3.1-70B single-KV-head group, 8 CUBEs $\times$ 8 PEs). Replication into CUBEs (Cases 2/3) wins the latency race (\SI{20.2}{\micro\second} for Case 3), but Case~4 ($\star$, KV split 64-way) finishes within \SI{14}{\micro\second} of the leader while paying a different cost---visible in Figure~\ref{fig:gqa-4cases-mem}.} \label{fig:gqa-4cases-lat} \end{figure} \begin{figure}[t] \centering \includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_memory.png} \caption{Per-CUBE KV memory footprint for the four cases. Cases 1 and 4---both with the KV cache split across cubes (Cube-SP)---hold only \SI{0.5}{\mebi\byte} of KV state per CUBE; Cases 2 and 3, which replicate the full KV, hold \SI{4}{\mebi\byte} per CUBE, an \textbf{8$\times$} blowup at this configuration that scales linearly with context length.} \label{fig:gqa-4cases-mem} \end{figure} \begin{figure}[t] \centering \includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_traffic.png} \caption{Per-case op-count breakdown. The replicated-KV PE-TP design (Case 2) avoids almost all on-device communication (\textasciitilde0 IPCQ copies), but at the cost of KV memory. Case~4's two-phase reduce charges \textasciitilde190 IPCQ copies and \textasciitilde190 DMA reads---this is the traffic that PE\_IPCQ (\S\ref{sec:allreduce}) is built to absorb at on-device speed.} \label{fig:gqa-4cases-traffic} \end{figure} Three things stand out. First, the fastest case in pure latency (Case~3, \SI{20.2}{\micro\second}) is also the most memory-hungry, requiring the full KV state on every CUBE---an option that fails to scale once context length blows past the per-CUBE budget. Second, Case~4's KV-split design gives back roughly \SI{14}{\micro\second} versus Case~3 in exchange for an \textbf{8$\times$} KV-memory reduction; for practical long-context serving where KV capacity is the binding constraint, this is the trade the design chooses (marked $\star$). Third, Case~4 pays its way in \emph{communication}: the op-count panel shows \textasciitilde190 IPCQ copies and \textasciitilde190 DMA reads, precisely the on-device collective traffic that PE\_IPCQ and the torus links of \S\ref{sec:allreduce} are provisioned to move quickly---so the ``slower'' strategy is in fact the one that fully cashes in the communication-side codesign work of this report. \subsection{Analysis and meaning} These panels are the clearest statement of the codesign thesis in the report. Because the composite command keeps GEMM issue cheap and the MAC array barely occupied, the fused attention kernel's latency is set almost entirely by data movement: streaming the KV cache and reducing partials across devices. That is precisely the cost that the communication-side work targets---PE\_IPCQ for the on-device reduction, the lazy load for load/compute overlap, fast TCM staging and torus links for the reduction itself. In other words, the two enablers are not independent features that happen to appear in the same kernel; the GEMM optimization is what \emph{exposes} the data-movement bottleneck (by removing the compute and issue overhead that would otherwise hide it), and the communication optimization is what \emph{attacks} it. For an attention-dominated decoder the meaningful hardware investments are therefore the ones that move data faster and reduce it on-device---not additional MAC throughput, which this workload cannot use.