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:
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
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user