diff --git a/docs/report/1H-codesign-paper/build/main.pdf b/docs/report/1H-codesign-paper/build/main.pdf index ac74200..cae60c5 100644 Binary files a/docs/report/1H-codesign-paper/build/main.pdf and b/docs/report/1H-codesign-paper/build/main.pdf differ diff --git a/docs/report/1H-codesign-paper/sections/05-gqa.tex b/docs/report/1H-codesign-paper/sections/05-gqa.tex index 111ee2d..2a0710f 100644 --- a/docs/report/1H-codesign-paper/sections/05-gqa.tex +++ b/docs/report/1H-codesign-paper/sections/05-gqa.tex @@ -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