report(gqa-short): draft results subsection (mode + composite-tier comparison, 4 figures)

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
+two phases. The \emph{prefill} kernel splits the query tile across the PEs
-shards around an inter-CUBE ring (``Ring KV''). The \emph{decode} kernel
+of a group and broadcasts each KV tile from the group's root PE over
-is head-replicated with a statically sharded KV cache and reduces partial
+intra-CUBE IPCQ, so the HBM K/V read is paid once per group instead of
-attention outputs through an M-fold intra-CUBE chain and, for multiple
+once per PE. The \emph{decode} kernel sequence-shards the KV cache across
-users, a two-level reduce-to-root. Two further primitives make long
+the same group of PEs, runs per-PE local attention, then chain-reduces the
-context practical: a \emph{lazy load} that issues the KV \textsf{DMA\_READ}
+partial $(m,\ell,O)$ triples back up to the group root. Two further
-and returns immediately, auto-waiting only at first use so that KV load
+primitives make long context practical: a \emph{lazy load} that issues the
-overlaps score computation; and per-tile \emph{scratch recycling} that
+KV \textsf{DMA\_READ} and returns immediately, auto-waiting only at first
-keeps the running softmax accumulators ($m,\ell,O$) in a persistent arena
+use so KV load overlaps score computation; and per-tile \emph{scratch
-while freeing per-tile temporaries, so the kernel fits the
+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
+The design question for short context is therefore not whether to fuse
-% (h_kv=8 KV heads -> 8 CUBEs, 1:1; intra-CUBE PE usage).
+softmax---that is settled---but how to distribute the eight KV heads
-% Bench code: src/kernbench/benches/gqa_helpers/short_ctx/
+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).
+A second axis is how each GEMM tile is issued. We isolate this with three
-% TODO: decode performance figure (latency, stage breakdown).
+\emph{composite tiers} applied to the same mapping: \textsf{without
-%       Bench output for short_ctx to be generated.
+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}