paper(gqa): wire the 4-cases long-context decode comparison into §6

The collaborator commit 7c346de added the comparative figures
(gqa_decode_long_ctx_4cases_{latency,memory,traffic}.png) but the
fused-GQA section still only referenced the four headline panels.
This commit closes that loop:

- New §6 subsection "Long-context decode: parallelism strategies"
  added between Results and Analysis. It lays out the four
  parallelism strategies (Cube-SP/Repl x PE-TP/SP) and pulls in the
  three new figures.
- The discussion highlights the central trade: the fastest strategy
  (Case 3, Cube-Repl x PE-SP at 20.2 us) requires replicating the
  full KV cache to every CUBE, while Case 4 (Cube-SP x PE-SP, the
  chosen design marked *) gives back ~14 us in exchange for an 8x
  KV-memory reduction.  Case 4's ~190 IPCQ copies + ~190 DMA reads
  are precisely the on-device collective traffic PE_IPCQ and the
  torus links of §5 are provisioned to absorb -- a direct payoff
  of the communication-side codesign work.
- Connects back to §5 (PE_IPCQ / all-reduce) so the reader sees the
  capstone arc: the GEMM enabler exposes the data-movement bound,
  the communication enabler attacks it, and the long-context
  parallelism study shows how the choice between the two extremes is
  framed by KV memory vs. on-device collective traffic.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-16 10:44:32 -07:00
parent e3f08972da
commit 22f3968b51
2 changed files with 75 additions and 0 deletions
Binary file not shown.
@@ -104,6 +104,81 @@ shards are streamed. The PE control-processor dispatch cost registered as
zero in this configuration---command issue is simply not on the critical zero in this configuration---command issue is simply not on the critical
path when data movement is this dominant. 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} \subsection{Analysis and meaning}
These panels are the clearest statement of the codesign thesis in the These panels are the clearest statement of the codesign thesis in the