ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

paper(1H): 6-case GQA long-ctx + IPCQ design alternatives + GEMM terminology

Browse Source 
- §6 GQA: rewrite long-context decode from 4-case to 6-case. Data
  Placement Policy now presents the six placement options and their
  intrinsic per-PE memory / per-token comm costs (no winner predicted);
  the long-context subsection selects the best placement for that regime
  (Case 6 ★, both-axes S_kv shard) from the measured 6-case sweep. Add
  KV-sharding diagram + analytical budget/summary figures.
- Regenerate the 6-row decode sweep (milestone-1h-gqa) and the
  sweep-dependent decode panels (latency/traffic/parallelism/memory) so
  figures, prose, and sweep_decode.json are mutually consistent.
- §5 All-Reduce: add IPCQ design alternatives (architecture + decision
  matrix) as design-rationale schematics (illustrative step-counts, not
  measured) and the bench-generated topology diagram.
- §4 GEMM: rename "user-orchestrated/user-level" -> "kernel-orchestrated/
  kernel-level" (orchestration runs in the kernel program vs the
  scheduler-orchestrated composite); minor accuracy fixes (7.18 TFLOP/s
  ~10% below peak; ~781 ns DMA).
- Recompile build/main.pdf.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook
2026-06-22 00:55:54 -07:00
parent bb668a3ac3
commit e9a5c438e3
20 changed files with 351 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/report/1H-codesign-paper/sections/03-gemm.tex
+11 -11
View File
@@ -172,16 +172,16 @@ kernel of \S\ref{sec:gqa} reaches for next — keeping the right
working set on-chip so the composite pipeline lands in the
compute-rich regime rather than the BW-bound one.
\subsection{Why composite, and not user-orchestrated async loading?}
\subsection{Why composite, and not kernel-orchestrated async loading?}
\label{sec:gemm-vs-async}
A reader familiar with double-buffered GEMM kernels on conventional
hardware may ask: why a hardware-side composite command at all? Why
isn't the obvious user-level pattern --- async-load each operand,
isn't the obvious kernel-level pattern --- async-load each operand,
overlap with compute, accumulate --- sufficient?
To answer this concretely we contrast composite against two
user-orchestrated baselines that have access to the same single-op
kernel-orchestrated baselines that have access to the same single-op
primitives the platform exposes (\textsf{tl.load} for async DMA into
TCM, \textsf{tl.dot} for a single-op GEMM command on TCM-resident
operands, \textsf{tl.store} for a DMA write-back). Both baselines
@@ -199,7 +199,7 @@ way at the runtime API surface to express ``start computing on
tile 0 of $B$ while tile 1 is still in flight.'' Load-of-$B$ and
GEMM therefore serialize.
\paragraph{Async-tiled (chunked prefetch).} The user-level workaround is to
\paragraph{Async-tiled (chunked prefetch).} The kernel-level workaround is to
split $B$ along $K$ into \textsf{TILE\_K}-sized chunks, issue async
\textsf{tl.load}s for those chunks, issue one \textsf{tl.dot} per
chunk (each blocking only on its own $b_i$), and accumulate via
@@ -227,7 +227,7 @@ $K_{\text{KV}}=4096, d_{\text{head}}=128$ already needs
--- past the cap. async-full and queue-all async-tiled are therefore
not just slower than composite but \emph{architecturally infeasible}
at LLM context length. The depth-2 async-tiled kernel is the only
user-level
kernel-level
variant whose peak TCM footprint stays
$O(2 \cdot \textsf{TILE\_K} \cdot N)$ regardless of $K$, the same
order as composite's per-tile streaming buffer. It is the apples-to-apples
@@ -249,7 +249,7 @@ onto $K / \textsf{TILE\_K} = 48$ hardware tiles, so per-tile costs
amplify into the largest measurable gap. The work content is
identical for all four kernels: $\sim$6.3 M f16 MACs and
$\sim$386 KiB of $B$ traffic from HBM, which together require
$\sim$786 ns of GEMM-engine compute and $\sim$750 ns of DMA on a
$\sim$786 ns of GEMM-engine compute and $\sim$781 ns of DMA on a
saturated per-PE link. What differs is the number of host commands
the same work is decomposed into --- 2 for composite (one
\textsf{tl.load(A)} plus one composite), 4 for async-full, and 192
@@ -257,7 +257,7 @@ for either async-tiled variant (48 $A$-loads + 48 $B$-loads + 48
\textsf{tl.dot}s + 47 elementwise adds + 1 store). The
engine-pipeline-window throughput tracks that decomposition closely:
composite reaches \SI{7.18}{\tera\flop\per\second} (post-overlap
limit, only \SI{12}{\percent} below the \SI{8}{\tera\flop\per\second}
limit, only \SI{10}{\percent} below the \SI{8}{\tera\flop\per\second}
per-PE GEMM peak), async-full \SI{3.91}{\tera\flop\per\second}
(DMA and compute serialize on a single big dot), and both async-tiled
variants $\sim$\SI{2.53}{\tera\flop\per\second} (192 commands' worth of
@@ -272,7 +272,7 @@ specific simulator mechanisms.
four issuance patterns: composite (one command, scheduler streams
per-tile internally), async-full (one \textsf{tl.dot} on
fully-loaded $B$), async-tiled with depth-2 double-buffer
(TCM-bounded; the only user-level variant that scales to LLM
(TCM-bounded; the only kernel-level variant that scales to LLM
context length), and async-tiled with depth-$\infty$
(all B-tiles queued up front; included as a sanity check that the
prefetch depth is \emph{not} what separates the async-tiled kernel from
@@ -293,7 +293,7 @@ and all three async kernels beat it.}
\end{figure*}
\paragraph{Decomposing the gap.} Three structural mechanisms separate
composite from the user-level baselines, and they layer.
composite from the kernel-level baselines, and they layer.
\emph{1. Inter-engine token routing happens below the host-side
dispatch path.} The composite encodes the full
@@ -309,7 +309,7 @@ hand-offs total --- behind a single command from the host's point of
view.
\emph{2. \textsf{tl.dot} cannot replicate that per-tile pipeline at the
user level.} A single-op GEMM command is handled on the GEMM engine as
kernel level.} A single-op GEMM command is handled on the GEMM engine as
a single monolithic compute timeout for the supplied $M{\times}K{\times}N$;
there is no internal token loop that would let a streaming DMA of
$B[i{+}1]$ overlap with the GEMM of $B[i]$ inside one
@@ -388,7 +388,7 @@ the four kernels to combine (a) macro-command dispatch at the host
boundary (amortizing the structural CPU cost across all the work a
single GEMM does), (b) scheduler-internal per-HW-tile streaming of
DMA$\rightleftarrows$compute, and (c) TCM-bounded streaming buffer.
User-orchestrated async kernels can have any two of those, not all
Kernel-orchestrated async kernels can have any two of those, not all
three: async-full pays one host dispatch (a) but forfeits per-tile
overlap (b) and pins all of $B$ in TCM (c); depth-$\infty$ async-tiled
achieves inter-chunk overlap but at $N_K$ host dispatches and
docs/report/1H-codesign-paper/sections/04-allreduce.tex
+72 -2
View File
@@ -127,6 +127,63 @@ across whatever inter-device topology the configuration specifies,
with the IPCQ ring buffer placed in on-PE TCM, PE-local HBM, or
cube-shared SRAM---the third knob the results section sweeps.
\subsection{Design alternatives}
\label{sec:ipcq-alternatives}
PE\_IPCQ is one point in a small space of hardware mechanisms for
moving a short message from one PE to a neighbor and signalling its
arrival. Three established alternatives anchor the space, each the
HW realization of a familiar host-networking idea: a \emph{doorbell +
polling} scheme (the classic MMIO doorbell---write the payload by DMA,
write a doorbell, let the peer poll or take an interrupt); a
\emph{hardware message queue} (HMQ, the NVLink-style descriptor engine
that pushes a queue entry to the peer, with large payloads still
riding a second DMA); and a \emph{completion-queue} design (RDMA-CQ,
the InfiniBand/RoCE pattern where a DMA write auto-posts a completion
entry the peer's CQ polls). PE\_IPCQ is the fourth: a hardware ring
with credit return, splitting the control plane into PE\_IPCQ and the
data plane into PE\_DMA, with head updates riding the payload and tail
updates riding a 16\,B side-channel credit (\S\ref{sec:allreduce}).
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{ipcq_alternatives_architecture_stacked.png}
\caption{Per-send data and control flow for the four PE-to-PE
signalling mechanisms (sender\,$\rightarrow$\,NoC\,$\rightarrow$\,receiver).
Doorbell and RDMA-CQ each issue two fabric transactions (payload then
doorbell / completion) and leave the peer polling or taking an
interrupt; HMQ adds a dedicated descriptor engine but still moves large
payloads on a second DMA; PE\_IPCQ folds head-pointer signalling into
the payload flit train and returns the tail credit on a side channel,
so a send is one MMIO write and a receive is a flip-flop read. This is
a \emph{design schematic}, not a measured comparison.}
\label{fig:ipcq-arch}
\end{figure}
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{ipcq_alternatives_decision_matrix.png}
\caption{Why the ring+credit design was chosen, across five criteria:
single-send latency, whether the host CPU sits on the critical path,
whether the receiver must poll or take a wake-up interrupt, whether the
control and data datapaths are duplicated, and whether the mechanism is
right-sized for single-owner PE-to-PE traffic (rather than a
multi-tenant fabric). PE\_IPCQ is the only design that clears every
criterion. The accompanying per-send step-count tally
($\sim$28 control events for IPCQ versus $\sim$38 for HMQ, $\sim$53 for
RDMA-CQ, and $\sim$56 for doorbell+polling) is an \emph{illustrative}
order-of-magnitude comparator over hand-counted pipeline steps---not a
simulator measurement. The measured, simulator-grounded results follow
in the next subsection.}
\label{fig:ipcq-decision}
\end{figure}
The qualitative comparison motivates the design but is not a
quantitative claim: the cycle-step tallies above are hand-counted
control events, deliberately separated from the measured latencies that
follow. Everything in the results subsection runs on the PE\_IPCQ
substrate and is simulator-grounded.
\subsection{Results}
All measurements in this section run on the PE\_IPCQ substrate
@@ -138,8 +195,21 @@ the collective sweep builds its own six-device (six-SIP, $2\times3$)
configurations---distinct from the two-SIP default of
Table~\ref{tab:hw}---and measures all-reduce latency as a function of
payload size for three inter-device topologies: a 1D ring, a 2D mesh
(no wrap), and a 2D torus. Table~\ref{tab:allreduce} and
Figure~\ref{fig:allreduce-cmp} report the result.
(no wrap), and a 2D torus (Figure~\ref{fig:allreduce-topo}).
Table~\ref{tab:allreduce} and Figure~\ref{fig:allreduce-cmp} report the
result.
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{allreduce_topology.png}
\caption{The three six-device ($2\times3$) inter-device topologies the
collective sweep runs over, and the hierarchical local-reduce /
global all-reduce-broadcast schedule mapped onto each: a 1D ring, a 2D
mesh (no wrap-around), and a 2D torus (wrap-around links on both axes).
The torus's wrap links shorten the worst-case reduction path, which is
what the latency sweep below rewards.}
\label{fig:allreduce-topo}
\end{figure}
\begin{table}[t]
\centering
docs/report/1H-codesign-paper/sections/05-gqa.tex
+141 -68
View File
@@ -27,15 +27,83 @@ the established baseline rather than being re-derived.
\subsection{Data Placement Policy}
\label{sec:gqa-placement}
% TODO: compare data-placement options that apply across both the
% short- and long-context regimes. Candidate axes:
% - KV cache: per-CUBE shard vs. replicate; per-PE shard vs. replicate
% - Q / W_qkv / W_o weights: static partition across CUBEs and PEs
% - Workspace (m, l, O softmax state): scratch arena placement
% The 4-case taxonomy used for long-context decode in
% \S\ref{sec:gqa-long} (Cube-{SP,Repl} x PE-{TP,SP}) is one instantiation
% of this framework; the short-context mapping in \S\ref{sec:gqa-short}
% is another.
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}
@@ -75,78 +143,83 @@ measured path; results below reflect the implemented kernel only.
% TODO: prefill long-context performance figure.
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}
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_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}
\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_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}
\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_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}
\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}
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.
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{Comprehensive Analysis}
\label{sec:gqa-analysis}