paper(§5): add roofline, capacity planning, parallelism selection subsections
Three new subsections at the top of §5, motivated by the deployment-sizing questions the analytical visualization tool exposes: - 5.1 Roofline Analysis — B*, L*, short vs long context batch effect (measured: 19× per-token cost reduction at short context, 28% at long) - 5.2 Capacity Planning — 3-axis sizing (weights / KV / SLO), SLO targets, regime rules, deployment templates, per-axis playbook - 5.3 Parallelism Selection — TP × CP × PP × DP × EP comparison, symptom-driven axis selection, add/stop criteria, misconceptions Restructure: - 05-gqa.tex trimmed to section header + intro - Existing 5.4-5.7 content (placement, short/long ctx, composite) moved to 05x-fused-kernel.tex - Summary extracted to 05z-summary.tex, cross-refs updated - main.tex \input order sets the requested subsection sequence - lmodern loaded to satisfy microtype font-expansion Two new figures generated from tests/analytical_visualization/chip_roofline: - roofline_short_context.png (S_kv=8K, batch wins) - roofline_long_context.png (S_kv=1M, batch stalls at KV floor) - Generator: tests/analytical_visualization/_gen_roofline_paper_figs.py 4 tables (regime rules, deployment templates, playbook, parallelism criteria) use table* so they span both columns without overflow. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -24,533 +24,3 @@ kernel that uses the composite command and PE\_IPCQ at the same time.
|
||||
Multi-head attention (MHA) was studied in prior work and serves here as
|
||||
the established baseline rather than being re-derived.
|
||||
|
||||
\subsection{Data Placement Policy}
|
||||
\label{sec:gqa-placement}
|
||||
|
||||
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}
|
||||
|
||||
The fused GQA kernel issues its matrix products as scheduler-managed
|
||||
composite commands and keeps the online-softmax merge and the cross-device
|
||||
KV reduction inside the kernel, on PE\_IPCQ. Two kernel families cover the
|
||||
two phases. The \emph{prefill} kernel splits the query tile across the PEs
|
||||
of a group and broadcasts each KV tile from the group's root PE over
|
||||
intra-CUBE IPCQ, so the HBM K/V read is paid once per group instead of
|
||||
once per PE. The \emph{decode} kernel sequence-shards the KV cache across
|
||||
the same group of PEs, runs per-PE local attention, then chain-reduces the
|
||||
partial $(m,\ell,O)$ triples back up to the group root. Two further
|
||||
primitives make long context practical: a \emph{lazy load} that issues the
|
||||
KV \textsf{DMA\_READ} and returns immediately, auto-waiting only at first
|
||||
use so KV load overlaps score computation; and per-tile \emph{scratch
|
||||
recycling} that keeps the running accumulators in a persistent arena while
|
||||
freeing per-tile temporaries, so the kernel fits the
|
||||
\SI{1}{\mebi\byte} scratch budget across many tiles. A further refinement
|
||||
restructures the decode step so the per-tile matrix products and the
|
||||
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.
|
||||
|
||||
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
|
||||
group size four; \textsf{4-kv-per-cube} packs four heads with group size
|
||||
two; \textsf{8-kv-per-cube} loads all heads onto a single CUBE with one
|
||||
PE per head. The mapping fixes a trade-off: as KV heads per CUBE grows, the per-CUBE KV
|
||||
footprint grows linearly and the mapping uses proportionally fewer CUBEs
|
||||
(eight down to one), but the intra-CUBE IPCQ broadcast/reduce cost shrinks
|
||||
to zero (8-kv-per-cube uses a single PE per head and has no group
|
||||
communication).
|
||||
|
||||
A second axis is how each GEMM tile is issued. We isolate this with three
|
||||
\emph{composite tiers} applied to the same mapping: \textsf{without
|
||||
composite} uses primitive \textsf{tl.dot} and unfused softmax;
|
||||
\textsf{with composite (GEMM-only)} issues the GEMMs as
|
||||
\textsf{tl.composite(op="gemm")} commands so the scheduler can pipeline
|
||||
the tile stages but leaves softmax as primitives; \textsf{with composite
|
||||
+ softmax\_merge} additionally folds the softmax fold into the $P\!\cdot\!V$
|
||||
composite through a named \textsf{softmax\_merge} prologue recipe.
|
||||
|
||||
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}.
|
||||
The mappings separate cleanly, and in proportion to how many PEs each
|
||||
dedicates to a head: \textsf{1-kv-per-cube} spreads one head across all
|
||||
eight PEs of its CUBE, \textsf{8-kv-per-cube} gives each head a single PE,
|
||||
and the four mappings form a $\{64,32,16,8\}$-active-PE ladder. Decode
|
||||
latency tracks that ladder inversely---at $S_{kv}{=}64$K the
|
||||
\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube} ratio is $7.4\times$
|
||||
($\SI{87.0}{}$ vs.\ $\SI{11.8}{\micro\second}$), and even at $8$K it is
|
||||
$4.8\times$ ($\SI{10.8}{}$ vs.\ $\SI{2.2}{\micro\second}$); prefill shows
|
||||
the same ordering more gently ($2.3\times$ at $64$K). The reason is that
|
||||
decode here is \emph{bandwidth-bound}: each active PE streams its KV shard
|
||||
at \SI{46}{}--\SI{76}{\percent} of the \SI{256}{\giga\byte\per\second}
|
||||
per-PE ceiling (rising with context), so a mapping's speed is set by how
|
||||
many PEs it puts to work---more PEs, more aggregate HBM bandwidth, lower
|
||||
latency. The $M{=}G{=}8$ score GEMM is skinny and leaves the MAC array
|
||||
lightly loaded throughout; the lever here is data movement, not compute.
|
||||
|
||||
\begin{figure}[t]
|
||||
\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; log scale). The mappings separate along the
|
||||
$\{64,32,16,8\}$-active-PE ladder set by heads-per-CUBE:
|
||||
\textsf{1-kv-per-cube} (64 PEs) is fastest and \textsf{8-kv-per-cube}
|
||||
(8 PEs) slowest, by $4.8\times$ at $8$K growing to $7.4\times$ at $64$K in
|
||||
decode ($2.3\times$ at $64$K in prefill). Decode is bandwidth-bound, so
|
||||
latency scales inversely with the number of active PEs. Composite tiers
|
||||
track the baseline at this skinny shape
|
||||
(Figure~\ref{fig:gqa-short-a1-variant}).}
|
||||
\label{fig:gqa-short-wall-baseline}
|
||||
\end{figure}
|
||||
|
||||
Figure~\ref{fig:gqa-short-tradeoff} shows what the mappings trade for that
|
||||
latency. The differentiator is \emph{density}: \textsf{8-kv-per-cube}
|
||||
concentrates all eight heads onto one CUBE, so its per-CUBE KV footprint is
|
||||
$8\times$ that of \textsf{1-kv-per-cube} (one CUBE per head), and it
|
||||
occupies a single one of the SIP's sixteen CUBEs against
|
||||
\textsf{1-kv-per-cube}'s eight. Per-PE HBM \emph{utilization}, by contrast,
|
||||
is similar across mappings---every active PE runs near the same
|
||||
\SI{46}{}--\SI{76}{\percent} of the per-PE ceiling---so the latency ladder
|
||||
comes from the \emph{count} of active PEs, not from any per-PE
|
||||
concentration. IPCQ traffic runs opposite to density:
|
||||
\textsf{1-kv-per-cube} pays the eight-PE chain-reduce while
|
||||
\textsf{8-kv-per-cube} (one PE per head) pays none, but absolute IPCQ volume
|
||||
stays under \SI{120}{\kibi\byte} per run, two orders of magnitude below the
|
||||
HBM traffic. The result is a genuine trade-off rather than a single winner:
|
||||
\textsf{1-kv-per-cube} minimizes per-request latency by spending the most
|
||||
hardware (eight CUBEs, 64 PEs) on one request, while denser mappings leave
|
||||
CUBEs free for other requests---the batched-serving axis taken up below.
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_short_context/per_cube_tradeoff_mode_compare.png}
|
||||
\caption{Per-CUBE trade-off across the four mappings (baseline kernel;
|
||||
log scale). Top row: per-CUBE KV footprint---\textsf{8-kv-per-cube} packs
|
||||
all eight heads onto one CUBE, an $8\times$ larger per-CUBE KV than
|
||||
\textsf{1-kv-per-cube} and one CUBE used against eight (the $8\times$ ratio
|
||||
is context-invariant; the absolute footprint grows with $S_{kv}$). Bottom
|
||||
row: IPCQ traffic per run---\textsf{1-kv-per-cube} pays the eight-PE
|
||||
chain-reduce while \textsf{8-kv-per-cube} (single PE per head) pays zero.
|
||||
The density (top) is what sets how many concurrent requests a SIP can hold.}
|
||||
\label{fig:gqa-short-tradeoff}
|
||||
\end{figure}
|
||||
|
||||
The composite ablation in Figure~\ref{fig:gqa-short-a1-variant} shows the
|
||||
three tiers within \textsf{1-kv-per-cube}: wall-clock is nearly flat, with
|
||||
the GEMM-only composite tier a few percent \emph{slower}
|
||||
(\SI{14.0}{} vs.\ \SI{11.8}{\micro\second} at $64$K decode) and the
|
||||
\textsf{softmax\_merge} tier matching the baseline. This is the expected
|
||||
outcome for the decode-skinny shape---$M{=}G{=}8$ is well below the
|
||||
scheduler's $\textsf{TILE\_M}{=}32$ supertile (\S\ref{sec:gemm}), so the
|
||||
composite path pads $M$ by $4\times$ with zeros and the fusion has no slack
|
||||
to win back; the padding is pure overhead here. It shows up in
|
||||
Figure~\ref{fig:gqa-short-gemm-util} as inflated GEMM-engine busy time on
|
||||
the composite tiers---the padded-$M$ GEMM is charged in full against a
|
||||
sub-\SI{15}{\micro\second} decode wall, driving the accounted utilization
|
||||
well above the baseline's \SI{12}{}--\SI{18}{\percent}. Fusion benefit
|
||||
requires lifting $M$, either by the long-context Q-tile split
|
||||
(\S\ref{sec:gqa-composite}) or by batching multiple users---which is also
|
||||
the lever that converts a mapping's density into throughput.
|
||||
|
||||
\paragraph{Latency versus batched throughput.} The latency ladder and the
|
||||
density trade-off pull in opposite directions once a SIP serves more than
|
||||
one request. Decode users are independent---each attends its own KV
|
||||
cache---and these mappings generate no cross-CUBE traffic, so a SIP runs
|
||||
several users at once on disjoint CUBE groups, up to $16/C$ users: two for
|
||||
\textsf{1-kv-per-cube}, sixteen for \textsf{8-kv-per-cube}. We measure this
|
||||
directly, launching $B$ concurrent users and reading aggregate latency off
|
||||
the shared clock (Figure~\ref{fig:gqa-batch}). The overlap is nearly
|
||||
perfect: aggregate latency stays within \SI{3}{}--\SI{5}{\percent} of the
|
||||
single-user latency all the way to a full SIP (\textsf{8-kv-per-cube} at
|
||||
$16$ users is \SI{11.4}{} vs.\ \SI{10.8}{\micro\second}), so that small
|
||||
residual is the only cross-user interference the shared fabric adds.
|
||||
Aggregate throughput at $8$K therefore rises almost linearly with $B$, from
|
||||
\SI{0.86}{}~requests\,/\,\si{\micro\second} (\textsf{1-kv-per-cube}, capped
|
||||
at two users) to \SI{1.41}{} (\textsf{8-kv-per-cube}, sixteen users): the
|
||||
\emph{dense} mapping wins on throughput even though it is the \emph{slowest}
|
||||
per request.
|
||||
|
||||
The reason is worth spelling out, because both mappings fill the same
|
||||
hardware: at capacity \textsf{1-kv-per-cube} runs two users
|
||||
$\times$ 64 PEs and \textsf{8-kv-per-cube} sixteen users $\times$ 8 PEs, so
|
||||
each puts all 128 PEs of the SIP to work. What differs is per-PE
|
||||
\emph{efficiency}. Splitting one head across eight PEs
|
||||
(\textsf{1-kv-per-cube}) leaves each PE only $S_{kv}/8$ tokens---a single
|
||||
tile---so its fixed overheads (pipeline fill on the lazy KV load, and the
|
||||
eight-PE online-softmax chain-reduce) dominate the little streaming work it
|
||||
does, and it sustains just \SI{46}{\percent} of the per-PE HBM ceiling.
|
||||
Giving each head a whole PE (\textsf{8-kv-per-cube}) streams the full
|
||||
$S_{kv}$-token cache contiguously with no reduce, amortizing those overheads
|
||||
over $8\times$ more work and reaching \SI{76}{\percent}. The measured
|
||||
throughput ratio ($1.41/0.86 = 1.6\times$) matches the utilization ratio
|
||||
($0.76/0.46 = 1.7\times$) almost exactly: because decode is
|
||||
bandwidth-bound, aggregate throughput is set by total HBM efficiency, not
|
||||
by PE count. \textsf{1-kv-per-cube} is really spending hardware on
|
||||
\emph{latency}---eight PEs per head buy only a $4.8\times$ speedup, a $60\%$
|
||||
strong-scaling efficiency---and that same $40\%$ overhead is what caps its
|
||||
throughput once the SIP is full. The design conclusion is regime-dependent: \textsf{1-kv-per-cube} is the
|
||||
latency-optimal choice for single-stream or low-batch decode, while denser
|
||||
mappings are the throughput-optimal choice for high-batch short-context
|
||||
serving. At long context the per-request latencies already scale as $1/C$,
|
||||
so the same accounting predicts the throughputs converge; a measured
|
||||
long-context batch sweep is 2H work (\S\ref{sec:future}).
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_short_context/batch_scaling.png}
|
||||
\caption{Batch scaling at $S_{kv}{=}8$K: aggregate throughput (requests per
|
||||
\si{\micro\second}) versus the number of concurrent decode users on one SIP,
|
||||
each user on a disjoint CUBE group. Each mapping scales almost linearly with
|
||||
$B$ up to its SIP capacity $16/C$---\textsf{1-kv-per-cube} saturates at two
|
||||
users, \textsf{8-kv-per-cube} at sixteen. Solid segments are measured
|
||||
concurrent runs; the dashed extensions are the saturated-throughput ceiling
|
||||
beyond capacity, where further users run in waves at that sustained rate
|
||||
(so the SIP is full, not idle). The dense mapping reaches the highest
|
||||
ceiling; the spread mapping is latency-optimal but capacity-limited.
|
||||
Concurrent users overlap to within \SI{3}{}--\SI{5}{\percent} of the
|
||||
single-user latency.}
|
||||
\label{fig:gqa-batch}
|
||||
\end{figure}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_short_context/wall_a1_variant_compare.png}
|
||||
\caption{Composite-tier ablation on \textsf{1-kv-per-cube} (log scale):
|
||||
wall-clock is nearly flat across tiers---the GEMM-only composite runs a few
|
||||
percent slower (padding overhead) and \textsf{softmax\_merge} matches the
|
||||
baseline. With $M{=}G{=}8 < \textsf{TILE\_M}{=}32$ the composite scheduler
|
||||
pads $M$ by $4\times$, leaving no fusion slack to recover at this shape.
|
||||
Fusion pays off only when $M$ fills the supertile---long-context Q-tile
|
||||
splitting or batched-$M$ inference.}
|
||||
\label{fig:gqa-short-a1-variant}
|
||||
\end{figure}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_short_context/gemm_util_a1_variant_ablation.png}
|
||||
\caption{GEMM-engine busy fraction on \textsf{1-kv-per-cube} across the
|
||||
three composite tiers. The baseline runs at \SI{12}{}--\SI{18}{\percent};
|
||||
the composite tiers read much higher---up to and past
|
||||
\SI{100}{\percent} of the short decode wall---because the padded
|
||||
$8\!\to\!32$ $M$ dimension is charged as GEMM time, bookkeeping overhead
|
||||
rather than useful work. The engine is never the bottleneck at this shape;
|
||||
decode is KV-bandwidth-bound.}
|
||||
\label{fig:gqa-short-gemm-util}
|
||||
\end{figure}
|
||||
|
||||
\subsection{Inference with Long-Context Length}
|
||||
\label{sec:gqa-long}
|
||||
|
||||
% TODO: prefill long-context kernel implementation description
|
||||
% (Sequence-Parallel partition of S_kv, per-case mechanics).
|
||||
% Bench code: src/kernbench/benches/gqa_helpers/long_ctx/
|
||||
|
||||
% TODO: prefill long-context performance figure.
|
||||
|
||||
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_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_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_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}
|
||||
|
||||
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}
|
||||
|
||||
The decode kernel of \S\ref{sec:gqa-long} issues its local attention as
|
||||
primitive operations: it walks each PE's $S_{\text{local}}$ token slice in
|
||||
\SI{1024}{}-token tiles, and for every tile issues a $Q\!\cdot\!K^{\top}$
|
||||
\textsf{dot}, the online-softmax primitives, and a $P\!\cdot\!V$
|
||||
\textsf{dot}, merging the running $(m,\ell,O)$ state by hand. The number
|
||||
of PE\_CPU commands this costs grows with the context. At a production
|
||||
context of $S_{kv}{=}1\,\text{M}$ tokens the Case-6 64-way split gives
|
||||
each PE $S_{\text{local}}{=}16384$ tokens, so its local attention is a
|
||||
$Q\!\cdot\!K^{\top}$ of $(8,128)\!\cdot\!(128,16384)$ and a
|
||||
$P\!\cdot\!V$ of $(8,16384)\!\cdot\!(16384,128)$---sixteen hand-issued
|
||||
tiles, each a fresh batch of CPU commands.
|
||||
|
||||
The composite command lets the kernel hand that tiling to PE\_SCHEDULER.
|
||||
We compare three command forms of the \emph{same} Case-6 kernel---identical
|
||||
placement and identical $(m,\ell,O)$ reduce, differing only in how the
|
||||
local attention is issued:
|
||||
\begin{itemize}
|
||||
\item \textbf{primitive}---the hand-tiled \textsf{dot}/softmax kernel
|
||||
above (the baseline of \S\ref{sec:gqa-long}).
|
||||
\item \textbf{composite}---each matrix product is one coarse
|
||||
\textsf{composite} GEMM over the \emph{whole} $S_{\text{local}}$, with
|
||||
$K$ and $V$ passed as HBM references so PE\_SCHEDULER streams and
|
||||
tiles them on the fixed $32\!\times\!64\!\times\!32$ MAC tile; the
|
||||
softmax stays primitive.
|
||||
\item \textbf{composite\,+\,softmax\_merge}---additionally folds the
|
||||
online-softmax merge and $P\!\cdot\!V$ into a single stateful
|
||||
\textsf{softmax\_merge} recipe composite.
|
||||
\end{itemize}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_composite.png}
|
||||
\caption{Three command forms of the Case-6 decode kernel, swept over
|
||||
context length ($S_{\text{local}}{=}S_{kv}/64$ per PE). \emph{Right:}
|
||||
PE\_CPU commands issued. The hand-tiled primitive kernel rises
|
||||
$O(n_{\text{tiles}})$---from 96 commands at one tile to 426 at the
|
||||
1\,M-token, sixteen-tile production point---while both composite forms
|
||||
issue a context-\emph{independent} $O(1)$ count (94 and 98) that
|
||||
\emph{saturates}: one coarse descriptor offloads the entire per-tile
|
||||
fan-out. \emph{Left:} the consequence for wall-clock latency is none---all
|
||||
three land on the same curve (\SI{30.6}{}, \SI{231}{},
|
||||
\SI{461}{\micro\second} at 8\,K\,/\,64\,K\,/\,128\,K), because decode is
|
||||
bound by streaming the KV cache, not by issue. Command-count is measured
|
||||
at emit time (exact, to 1\,M); latency on the data-mode engine over the
|
||||
tractable range.}
|
||||
\label{fig:gqa-composite}
|
||||
\end{figure}
|
||||
|
||||
Figure~\ref{fig:gqa-composite} reads off the two quantities that matter,
|
||||
and they point in opposite directions. The PE\_CPU command count (right)
|
||||
collapses from a context-growing $O(n_{\text{tiles}})$ to a flat $O(1)$:
|
||||
at 1\,M tokens the composite form issues \num{94} commands against the
|
||||
primitive kernel's \num{426}, a $4.5\times$ reduction
|
||||
($4.2\times$ in modeled dispatch cost), and---crucially---that number no
|
||||
longer grows with context. The wall-clock latency (left), by contrast,
|
||||
is unchanged across all three forms: decode is bound by streaming the KV
|
||||
cache out of HBM, so the command form does not move the critical path.
|
||||
|
||||
That juxtaposition is the point. The composite command is not a latency
|
||||
optimization for this memory-bound decode; it is a \emph{CPU-issue}
|
||||
optimization. Its value is removing the per-tile dispatch work that would
|
||||
otherwise grow without bound as context grows, freeing PE\_CPU to run
|
||||
ahead and keep the engines fed---which is exactly what lets the
|
||||
data-movement cost analyzed next show through as the true bottleneck
|
||||
rather than being masked by issue overhead. The \textsf{softmax\_merge}
|
||||
recipe folds the online merge into the same descriptor; on this
|
||||
memory-bound path its marginal cost over the plain GEMM composite is small
|
||||
(98 vs.\ 94 commands), and like the plain composite it keeps the issued
|
||||
count flat as context scales.
|
||||
|
||||
\paragraph{The compute-bound mirror: prefill.} Decode's verdict---command
|
||||
form is latency-neutral---is a property of its regime, not of the
|
||||
composite command. A decode step has $T_q{=}1$, so its score and context
|
||||
products are skinny ($M{=}G\,T_q{=}8$): the MAC array is barely fed and
|
||||
the kernel is bound by streaming the KV cache. Prefill is the opposite
|
||||
corner. It processes a block of query positions at once, so $M{=}G\,T_q$
|
||||
is large and tile-filling, the GEMMs carry real arithmetic intensity
|
||||
($\sim$$M$ flops/byte, well above the roofline ridge), and the kernel is
|
||||
\emph{compute-bound}. This is the regime the composite command was built
|
||||
for (\S\ref{sec:gemm}): it streams the per-HW-tile
|
||||
DMA$\rightleftarrows$compute pipeline so the MAC array stays fed, whereas
|
||||
the primitive kernel's blocking \textsf{tl.dot} serializes each tile's
|
||||
load and compute and starves the array between tiles. We run the same
|
||||
three command forms on a single-rank compute-bound prefill (FlashAttention
|
||||
$Q$-block $\times$ $S_{kv}$-tile, online softmax) and sweep the context
|
||||
length (Figure~\ref{fig:gqa-prefill-cb}).
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_prefill_compute_bound.png}
|
||||
\caption{Compute-bound prefill, three command forms, swept over context
|
||||
length ($M{=}8\,T_q$ tile-filling). \emph{Left:} end-to-end latency.
|
||||
\emph{Right:} MAC utilization (achieved $\div$ the
|
||||
\SI{8}{\tera\flop\per\second} per-PE peak). The hand-tiled primitive sits
|
||||
flat at $\sim$\SI{68}{\percent}---its serial load$\to$dot path leaves the
|
||||
MAC array idle between tiles regardless of context. The composite forms
|
||||
climb with context (\SI{67}{}$\to$\SI{80}{\percent} plain,
|
||||
\SI{67}{}$\to$\SI{83}{\percent} with the recipe) because a deeper $P\!\cdot
|
||||
\!V$ reduction gives more HW tiles to pipeline, and they convert that into
|
||||
wall-clock: at \num{1024} the recipe form is \SI{646.9}{} vs.\
|
||||
\SI{794.1}{\micro\second} (\SI{19}{\percent} faster). The margin
|
||||
\emph{grows} with context---the compute-bound mirror of the GEMM result of
|
||||
\S\ref{sec:gemm}.}
|
||||
\label{fig:gqa-prefill-cb}
|
||||
\end{figure}
|
||||
|
||||
The two studies together state the composite command's value precisely. It
|
||||
has two distinct benefits, and which one matters is set by the workload's
|
||||
roofline position. The first is \emph{host-issue offload}: one macro
|
||||
command in place of $O(n_{\text{tiles}})$ fine ones, which removes
|
||||
PE\_CPU dispatch work and is regime-independent (it shows in the decode
|
||||
command count). The second is \emph{MAC-array feeding}: the
|
||||
scheduler-internal per-tile DMA$\rightleftarrows$compute pipeline, which
|
||||
only converts to latency when the workload is compute-bound enough to have
|
||||
a MAC array worth keeping busy (it shows in the prefill utilization). A
|
||||
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{Summary}
|
||||
\label{sec:gqa-analysis}
|
||||
|
||||
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
|
||||
% without further changes?
|
||||
% - Does the short-context placement compose the same way?
|
||||
% - Implications for multi-user serving (single vs. mixed regimes).
|
||||
|
||||
Reference in New Issue
Block a user