report(1H): tighten §5 GQA (mapping tie, taxonomy bridge, dedup, Summary)
- Soften short-context claim: 1-kv-/8-kv-per-cube are tied at 8K and separate as context grows (1-kv faster at the large end). - Add a bridge noting the six-case placement taxonomy (§5.1) is the long-context lens; short context turns on head-to-CUBE distribution. - Remove caption-duplicating numbers from the long-context conclusion. - Reduce the closing subsection to a GQA-scoped Summary; defer the cross-cutting hardware-investment thesis to the Discussion. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Binary file not shown.
@@ -128,9 +128,14 @@ online-softmax merge are issued as \emph{composite} commands rather than
|
||||
hand-tiled primitives; \S\ref{sec:gqa-composite} measures it against the
|
||||
primitive baseline at long context.
|
||||
|
||||
The design question for short context is therefore not whether to fuse
|
||||
softmax---that is settled---but how to distribute the eight KV heads
|
||||
across the eight CUBEs of a SIP. Four mappings cover the spectrum, named
|
||||
This is a different decomposition from the six-case placement taxonomy of
|
||||
\S\ref{sec:gqa-placement}: that taxonomy spreads a single KV-head group
|
||||
across all 64 PEs and is the right lens for long context
|
||||
(\S\ref{sec:gqa-long}), whereas short context---where the whole cache fits
|
||||
comfortably---turns instead on how the eight KV heads are distributed
|
||||
across the eight CUBEs. The design question for short context is therefore
|
||||
not whether to fuse softmax---that is settled---but how to distribute the
|
||||
eight KV heads across the eight CUBEs of a SIP. Four mappings cover the spectrum, named
|
||||
by KV-head count per CUBE: \textsf{1-kv-per-cube} dedicates one whole
|
||||
CUBE to each head and uses all eight PEs of that CUBE on the head's
|
||||
sequence shard; \textsf{2-kv-per-cube} packs two heads per CUBE with
|
||||
@@ -155,10 +160,12 @@ We sweep all four mappings $\times$ three composite tiers $\times$
|
||||
$S_{kv}\in\{8\text{K},16\text{K},32\text{K},64\text{K}\}$ for both
|
||||
phases (prefill uses $T_q{=}8$ sliced tiles; decode uses $T_q{=}1$). The
|
||||
headline latency comparison is in Figure~\ref{fig:gqa-short-wall-baseline}.
|
||||
At every context length the \textsf{1-kv-per-cube} mapping is the
|
||||
fastest, and the $\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube}$ ratio
|
||||
grows monotonically with $S_{kv}$ (decode: $0.98\times$ at $8$K,
|
||||
$1.06\times$ at $64$K; prefill follows the same trend). The wall gap is
|
||||
At short context the two mappings are effectively tied, and they separate
|
||||
as context grows: the $\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube}$
|
||||
latency ratio rises monotonically with $S_{kv}$ (decode: $0.98\times$ at
|
||||
$8$K---within measurement noise---to $1.06\times$ at $64$K, where
|
||||
\textsf{1-kv-per-cube} is the faster; prefill follows the same trend). The
|
||||
wall gap is
|
||||
modest because short-context decode is dominated by setup overhead and
|
||||
its $M{=}G{=}8$ skinny shape leaves the MAC array idle most of the time;
|
||||
where the mappings really separate is per-CUBE pressure.
|
||||
@@ -167,10 +174,11 @@ where the mappings really separate is per-CUBE pressure.
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_short_context/wall_variant1_mode_compare.png}
|
||||
\caption{Wall-clock latency of the four short-context mappings, baseline
|
||||
kernel (no composite). Prefill and decode follow the same ordering at
|
||||
every context length and the \textsf{1-kv-per-cube} mapping is the
|
||||
fastest throughout; the gap to \textsf{8-kv-per-cube} grows from
|
||||
$0.98\times$ at $8$K to $1.06\times$ at $64$K. The composite and
|
||||
kernel (no composite). Prefill and decode follow the same ordering; the
|
||||
two mappings are effectively tied at $8$K and separate as context grows,
|
||||
the \textsf{8-kv-per-cube}/\textsf{1-kv-per-cube} ratio widening from
|
||||
$0.98\times$ at $8$K to $1.06\times$ at $64$K (\textsf{1-kv-per-cube}
|
||||
faster at the large end). The composite and
|
||||
softmax\_merge tiers track the baseline within a few percent at this
|
||||
shape (see Figure~\ref{fig:gqa-short-a1-variant}).}
|
||||
\label{fig:gqa-short-wall-baseline}
|
||||
@@ -320,22 +328,17 @@ absorb at link speed.}
|
||||
\label{fig:gqa-6cases-traffic}
|
||||
\end{figure}
|
||||
|
||||
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.
|
||||
Taken together the three panels select the placement. The only
|
||||
memory-feasible family at production context is the 64-way splits
|
||||
(Cases~4--6), and within it Case~6~$\star$ is both the fastest and the
|
||||
lightest-communicating, because it merges only the running softmax state
|
||||
rather than the partial scores that the $d_{\text{head}}$-split
|
||||
Cases~4--5 must all-reduce. Case~3's lower raw latency is beside the
|
||||
point---it replicates the full cache into every CUBE, overflowing the
|
||||
per-PE budget and wasting $8\times$ the compute. For long-context decode
|
||||
the placement of choice is therefore the both-axes sequence shard, 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{Use of Composite Commands}
|
||||
\label{sec:gqa-composite}
|
||||
@@ -459,24 +462,24 @@ memory-bound decode exercises only the first; a compute-bound prefill
|
||||
exercises both. The composite command is the single mechanism that
|
||||
delivers each where it applies.
|
||||
|
||||
\subsection{Comprehensive Analysis}
|
||||
\subsection{Summary}
|
||||
\label{sec:gqa-analysis}
|
||||
|
||||
These panels are the clearest statement of the codesign thesis in the
|
||||
report. Because the composite command keeps GEMM issue cheap and the MAC
|
||||
array barely occupied, the fused attention kernel's latency is set almost
|
||||
entirely by data movement: streaming the KV cache and reducing partials
|
||||
across devices. That is precisely the cost that the communication-side
|
||||
work targets---PE\_IPCQ for the on-device reduction, the lazy load for
|
||||
load/compute overlap, fast TCM staging and torus links for the reduction
|
||||
itself. In other words, the two enablers are not independent features that
|
||||
happen to appear in the same kernel; the GEMM optimization is what
|
||||
\emph{exposes} the data-movement bottleneck (by removing the compute and
|
||||
issue overhead that would otherwise hide it), and the communication
|
||||
optimization is what \emph{attacks} it. For an attention-dominated decoder
|
||||
the meaningful hardware investments are therefore the ones that move data
|
||||
faster and reduce it on-device---not additional MAC throughput, which this
|
||||
workload cannot use.
|
||||
The section's results line up into one picture of the fused kernel.
|
||||
Placement is decided by memory first and communication second: only the
|
||||
64-way splits fit production context, and among them the both-axes
|
||||
sequence shard (Case~6) minimizes the per-token collective by merging
|
||||
softmax state rather than partial scores. Command form is decided by
|
||||
roofline position: the composite command removes the per-tile issue work
|
||||
in both regimes, but converts to wall-clock only in compute-bound prefill,
|
||||
while memory-bound decode stays bound by KV streaming regardless of command
|
||||
form. The thread connecting the two is that, once composite issue makes
|
||||
GEMM cheap and leaves the MAC array idle, the fused kernel's latency is set
|
||||
almost entirely by data movement---streaming the KV cache and reducing
|
||||
partials across PEs---which is exactly the cost the communication-side
|
||||
work (on-device reduction, lazy load/compute overlap, fast TCM staging and
|
||||
torus links) is built to attack. What this implies for hardware investment
|
||||
across the report as a whole is taken up in the discussion.
|
||||
|
||||
% TODO: cross-regime DP (data parallelism) applicability:
|
||||
% - Does Case-4 long-context placement compose with batch-level DP
|
||||
|
||||
Reference in New Issue
Block a user