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
|
hand-tiled primitives; \S\ref{sec:gqa-composite} measures it against the
|
||||||
primitive baseline at long context.
|
primitive baseline at long context.
|
||||||
|
|
||||||
The design question for short context is therefore not whether to fuse
|
This is a different decomposition from the six-case placement taxonomy of
|
||||||
softmax---that is settled---but how to distribute the eight KV heads
|
\S\ref{sec:gqa-placement}: that taxonomy spreads a single KV-head group
|
||||||
across the eight CUBEs of a SIP. Four mappings cover the spectrum, named
|
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
|
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
|
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
|
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
|
$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
|
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}.
|
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
|
At short context the two mappings are effectively tied, and they separate
|
||||||
fastest, and the $\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube}$ ratio
|
as context grows: the $\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube}$
|
||||||
grows monotonically with $S_{kv}$ (decode: $0.98\times$ at $8$K,
|
latency ratio rises monotonically with $S_{kv}$ (decode: $0.98\times$ at
|
||||||
$1.06\times$ at $64$K; prefill follows the same trend). The wall gap is
|
$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
|
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;
|
its $M{=}G{=}8$ skinny shape leaves the MAC array idle most of the time;
|
||||||
where the mappings really separate is per-CUBE pressure.
|
where the mappings really separate is per-CUBE pressure.
|
||||||
@@ -167,10 +174,11 @@ where the mappings really separate is per-CUBE pressure.
|
|||||||
\centering
|
\centering
|
||||||
\includegraphics[width=\linewidth]{gqa_short_context/wall_variant1_mode_compare.png}
|
\includegraphics[width=\linewidth]{gqa_short_context/wall_variant1_mode_compare.png}
|
||||||
\caption{Wall-clock latency of the four short-context mappings, baseline
|
\caption{Wall-clock latency of the four short-context mappings, baseline
|
||||||
kernel (no composite). Prefill and decode follow the same ordering at
|
kernel (no composite). Prefill and decode follow the same ordering; the
|
||||||
every context length and the \textsf{1-kv-per-cube} mapping is the
|
two mappings are effectively tied at $8$K and separate as context grows,
|
||||||
fastest throughout; the gap to \textsf{8-kv-per-cube} grows from
|
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. The composite and
|
$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
|
softmax\_merge tiers track the baseline within a few percent at this
|
||||||
shape (see Figure~\ref{fig:gqa-short-a1-variant}).}
|
shape (see Figure~\ref{fig:gqa-short-a1-variant}).}
|
||||||
\label{fig:gqa-short-wall-baseline}
|
\label{fig:gqa-short-wall-baseline}
|
||||||
@@ -320,22 +328,17 @@ absorb at link speed.}
|
|||||||
\label{fig:gqa-6cases-traffic}
|
\label{fig:gqa-6cases-traffic}
|
||||||
\end{figure}
|
\end{figure}
|
||||||
|
|
||||||
The measurements select the right placement for this regime. The fastest
|
Taken together the three panels select the placement. The only
|
||||||
raw latency (Case~3, \SI{17.8}{\micro\second}) comes from replicating the
|
memory-feasible family at production context is the 64-way splits
|
||||||
full KV cache into every CUBE---which is exactly the placement the
|
(Cases~4--6), and within it Case~6~$\star$ is both the fastest and the
|
||||||
long-context memory budget forbids, and which the parallelism panel shows
|
lightest-communicating, because it merges only the running softmax state
|
||||||
wastes 8$\times$ the compute. Restricting attention to the placements that
|
rather than the partial scores that the $d_{\text{head}}$-split
|
||||||
fit production-context memory (the 64-way splits, Cases~4--6), the choice
|
Cases~4--5 must all-reduce. Case~3's lower raw latency is beside the
|
||||||
is Case~6~$\star$: it is the fastest of the three
|
point---it replicates the full cache into every CUBE, overflowing the
|
||||||
(\SI{30.6}{\micro\second}), and the op-count panel shows why---its
|
per-PE budget and wasting $8\times$ the compute. For long-context decode
|
||||||
softmax-state-only reduction charges 189 IPCQ copies and one DMA write
|
the placement of choice is therefore the both-axes sequence shard, and the
|
||||||
against the 280 copies and 8 DMA writes the $d_{\text{head}}$-split
|
cross-PE softmax reduction it does pay is precisely the traffic the
|
||||||
Cases~4--5 pay for their partial-score all-reduce. For long-context
|
communication-side codesign of this report is built to move quickly.
|
||||||
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{Use of Composite Commands}
|
\subsection{Use of Composite Commands}
|
||||||
\label{sec:gqa-composite}
|
\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
|
exercises both. The composite command is the single mechanism that
|
||||||
delivers each where it applies.
|
delivers each where it applies.
|
||||||
|
|
||||||
\subsection{Comprehensive Analysis}
|
\subsection{Summary}
|
||||||
\label{sec:gqa-analysis}
|
\label{sec:gqa-analysis}
|
||||||
|
|
||||||
These panels are the clearest statement of the codesign thesis in the
|
The section's results line up into one picture of the fused kernel.
|
||||||
report. Because the composite command keeps GEMM issue cheap and the MAC
|
Placement is decided by memory first and communication second: only the
|
||||||
array barely occupied, the fused attention kernel's latency is set almost
|
64-way splits fit production context, and among them the both-axes
|
||||||
entirely by data movement: streaming the KV cache and reducing partials
|
sequence shard (Case~6) minimizes the per-token collective by merging
|
||||||
across devices. That is precisely the cost that the communication-side
|
softmax state rather than partial scores. Command form is decided by
|
||||||
work targets---PE\_IPCQ for the on-device reduction, the lazy load for
|
roofline position: the composite command removes the per-tile issue work
|
||||||
load/compute overlap, fast TCM staging and torus links for the reduction
|
in both regimes, but converts to wall-clock only in compute-bound prefill,
|
||||||
itself. In other words, the two enablers are not independent features that
|
while memory-bound decode stays bound by KV streaming regardless of command
|
||||||
happen to appear in the same kernel; the GEMM optimization is what
|
form. The thread connecting the two is that, once composite issue makes
|
||||||
\emph{exposes} the data-movement bottleneck (by removing the compute and
|
GEMM cheap and leaves the MAC array idle, the fused kernel's latency is set
|
||||||
issue overhead that would otherwise hide it), and the communication
|
almost entirely by data movement---streaming the KV cache and reducing
|
||||||
optimization is what \emph{attacks} it. For an attention-dominated decoder
|
partials across PEs---which is exactly the cost the communication-side
|
||||||
the meaningful hardware investments are therefore the ones that move data
|
work (on-device reduction, lazy load/compute overlap, fast TCM staging and
|
||||||
faster and reduce it on-device---not additional MAC throughput, which this
|
torus links) is built to attack. What this implies for hardware investment
|
||||||
workload cannot use.
|
across the report as a whole is taken up in the discussion.
|
||||||
|
|
||||||
% TODO: cross-regime DP (data parallelism) applicability:
|
% TODO: cross-regime DP (data parallelism) applicability:
|
||||||
% - Does Case-4 long-context placement compose with batch-level DP
|
% - Does Case-4 long-context placement compose with batch-level DP
|
||||||
|
|||||||
Reference in New Issue
Block a user