diff --git a/docs/report/1H-codesign-paper/sections/05-gqa.tex b/docs/report/1H-codesign-paper/sections/05-gqa.tex index db6731f..1a71ab5 100644 --- a/docs/report/1H-codesign-paper/sections/05-gqa.tex +++ b/docs/report/1H-codesign-paper/sections/05-gqa.tex @@ -111,29 +111,144 @@ redundant compute each one induces. 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 +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. -% 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/ +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). -% TODO: prefill performance figure (latency, stage breakdown). -% TODO: decode performance figure (latency, stage breakdown). -% Bench output for short_ctx to be generated. +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}