paper(1H): 6-case GQA long-ctx + IPCQ design alternatives + GEMM terminology
- §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:
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user