ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

paper(§5): add roofline, capacity planning, parallelism selection subsections

Browse Source 
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:
Mukesh Garg
2026-07-30 15:56:41 -07:00
parent bfe3e0e8d1
commit fc4747668e
11 changed files with 942 additions and 531 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/report/1H-codesign-paper/build/main.pdf
BIN
View File
Binary file not shown.
docs/report/1H-codesign-paper/figures/roofline_long_context.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 98 KiB

docs/report/1H-codesign-paper/figures/roofline_short_context.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 106 KiB

docs/report/1H-codesign-paper/main.tex
+7 -1
View File
@@ -15,6 +15,7 @@
\usepackage{caption}
\usepackage{subcaption}
\captionsetup{font=small,labelfont=bf}
\usepackage{lmodern} % scalable Latin Modern fonts (required by microtype expansion)
\usepackage{microtype}
\usepackage{tikz}
\usetikzlibrary{arrows.meta,positioning,calc,fit}
@@ -39,7 +40,12 @@ AGI Computing Lab, System Technology Group\\
\input{sections/02-platform}
\input{sections/03-gemm}
\input{sections/04-allreduce}
\input{sections/05-gqa}
\input{sections/05-gqa} % section header + intro
\input{sections/05a-roofline} % 5.1 Roofline Analysis
\input{sections/05b-capacity-planning} % 5.2 Capacity Planning
\input{sections/05c-parallelism-selection} % 5.3 Parallelism Selection
\input{sections/05x-fused-kernel} % 5.4-5.7 Placement + short/long ctx + composite
\input{sections/05z-summary} % 5.8 Summary
\input{sections/06-agentic}
\input{sections/07-hw-spec-search}
\input{sections/08-discussion}
docs/report/1H-codesign-paper/sections/05-gqa.tex
-530
View File
@@ -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).
docs/report/1H-codesign-paper/sections/05a-roofline.tex
+85
View File
@@ -0,0 +1,85 @@
\subsection{Roofline Analysis: Batch and Context Regimes}
\label{sec:roofline}
Before deploying any sharding strategy, the workload's position on the
decode roofline sets what \emph{can} be improved and what \emph{cannot}.
Two quantities matter: the machine's \textbf{critical batch}
$B^\ast = C \cdot b / (2 W)$ (where $C$ is per-PE FLOPs, $W$ is per-PE
HBM bandwidth, $b$ is bytes per parameter), and its \textbf{balance
context} $L^\ast = 2 N / (\mathrm{AI} \cdot \mathrm{kv\_bpt})$ (where
$N$ is active parameters and $\mathrm{AI} = C/W$ is arithmetic
intensity). Above $B^\ast$ the deployment leaves the memory-bound
regime; above $L^\ast$ per-user KV streaming dominates and no amount
of batching hides it. For Llama-3.1-70B on the default AHBM machine
($C \!=\! \SI{8}{TFLOPS}$/PE, $W\!=\! \SI{256}{GB/s}$/PE), we get
$B^\ast \!\approx\! 31$ and $L^\ast \!\approx\! \SI{13.4}{K}$ tokens.
\paragraph{Short-context regime ($S_{kv} \!<\! L^\ast$).}
Figure~\ref{fig:roofline-short} shows one decode step at
$S_{kv}\!=\! \SI{8}{K}$. The step-latency panel (left) makes it
obvious that weight fetch dominates at low $B$: at $B\!=\!1$, one
step spends $\sim\!\SI{535}{ms}$ streaming weights and only
$\SI{28}{ms}$ on per-sequence compute and KV combined. The
cost-per-token panel (right) is where the batch story lives: dividing
step latency by $B$ shrinks the weight term as $1/B$, so per-token
cost drops from $\SI{562}{ms}$ at $B\!=\!1$ to $\SI{29.7}{ms}$ at
$B\!=\!256$ --- a $19\times$ reduction. This is the memory-bound-to-
compute-bound crossover: past $B^\ast$, weight cost is fully
amortized and per-token time asymptotes to the compute + KV floor
($\sim\!\SI{27}{ms}$). \textbf{At short context, batching directly
lowers cost per token.}
\begin{figure*}[h]
\centering
\includegraphics[width=\textwidth]{roofline_short_context.png}
\caption{Roofline decomposition for Llama-3.1-70B decode at
$S_{kv}\!=\! \SI{8}{K}$ (short, well below $L^\ast\!\approx\! \SI{13.4}{K}$).
\textbf{Left}: step latency vs.\ batch $B$; weight fetch is flat
(one HBM sweep per step), compute and KV grow linearly with $B$.
\textbf{Right}: per-token cost ($=$ step $\div B$); the $1/B$
weight-fetch term amortizes rapidly, dropping total per-token time
from \SI{562}{ms} at $B\!=\!1$ to \SI{29.7}{ms} at $B\!=\!256$
--- a $19\times$ throughput gain from batching alone.}
\label{fig:roofline-short}
\end{figure*}
\paragraph{Long-context regime ($S_{kv} \!\gg\! L^\ast$).}
Figure~\ref{fig:roofline-long} shows the same decomposition at
$S_{kv}\!=\! \SI{1}{M}$ ($\sim\!78 \times L^\ast$). The step-latency
panel shows KV fetch has swelled by two orders of magnitude:
per-token KV cost is now $\sim\!\SI{1342}{ms}$, dwarfing both
weight ($\SI{535}{ms}$ at $B\!=\!1$) and compute ($\SI{17}{ms}$).
The cost-per-token panel is the punchline: increasing $B$ still
shrinks the weight term but leaves the giant KV floor untouched,
because \textbf{KV fetch is per-sequence} --- adding another user
adds a full extra copy of their KV read. Total per-token cost falls
only from \SI{1894}{ms} at $B\!=\!1$ to \SI{1361}{ms} at $B\!=\!256$
--- a mere $28\%$ reduction despite $256\times$ the batch.
\textbf{At long context, batching stops paying because per-user KV
streaming dominates.}
\begin{figure*}[h]
\centering
\includegraphics[width=\textwidth]{roofline_long_context.png}
\caption{Same decomposition at $S_{kv}\!=\! \SI{1}{M}$ (long,
$\sim\!78\times L^\ast$). KV fetch has grown to \SI{1342}{ms}/token
and is now the dominant term at every $B$. Batching only shrinks
the weight component; the KV floor is unmovable because each new
user brings a full per-sequence KV read. Total per-token cost falls
just $28\%$ from $B\!=\!1$ to $B\!=\!256$, versus $19\times$ at
short context.}
\label{fig:roofline-long}
\end{figure*}
\paragraph{Implication for deployment.} The two regimes call for
opposite strategies. In the short-context regime, the operator packs
$B$ as high as HBM allows to sit on the compute floor --- this is
where cost-per-token is minimized and hardware utilization is
highest. In the long-context regime, per-user KV is the binding
resource; batching offers little benefit, so the operator instead
shrinks $\mathrm{kv\_bpt}$ (GQA / MQA / MLA, INT4 KV, sparse
attention) and shards the sequence dimension itself (CP), routing
long-context requests to a dedicated pool with more chips per user.
The next two subsections (\S\ref{sec:capacity-planning},
\S\ref{sec:parallelism-selection}) turn these regime observations
into concrete sizing and sharding rules.
docs/report/1H-codesign-paper/sections/05b-capacity-planning.tex
+118
View File
@@ -0,0 +1,118 @@
\subsection{Capacity Planning: LLM Serving on AHBM}
\label{sec:capacity-planning}
Once a fused kernel is chosen, the deployment question is orthogonal:
\emph{how many cubes (or SIPs) does a given model, context length, and
latency SLO require?} Hyperscalers decide this along three axes that
must all be satisfied simultaneously. The total PE count is
$\max(A, B, C) \times N_{\text{replicas}}$, where $A$ is the capacity
floor, $B$ the KV-cache headroom, and $C$ the throughput SLO
(Table~\ref{tab:capacity-axes}).
\begin{table}[h]
\centering
\small
\caption{Three-axis sizing. $N$ = active params, $b$ = bytes per param,
$\mathrm{HBM}_{\mathrm{PE}}$ = per-PE HBM budget, $S_{kv}$ = context
length, $\mathrm{kv\_bpt}$ = KV cache bytes per token per user
($=2 \cdot h_{kv} \cdot d_{\text{head}} \cdot b \cdot L$).}
\label{tab:capacity-axes}
\begin{tabular}{@{}l l@{}}
\toprule
Axis & Formula \\
\midrule
A. Capacity floor & $\lceil N b / \mathrm{HBM}_{\mathrm{PE}} \rceil$ \\
B. KV headroom & $\lceil (N b + u \cdot S_{kv} \cdot \mathrm{kv\_bpt}) / \mathrm{HBM}_{\mathrm{PE}} \rceil$ \\
C. Throughput SLO & $N_{\text{replicas}} = \lceil n_{\text{users}} / B_{\mathrm{SLO}} \rceil$ \\
\bottomrule
\end{tabular}
\end{table}
\paragraph{SLO targets.} The per-token latency budget $B_{\mathrm{SLO}}$
follows the use case. TTFT (Time To First Token) is dominated by
prefill; TPOT (Time Per Output Token) by one decode step
(Table~\ref{tab:slo-targets}).
\begin{table}[h]
\centering
\small
\caption{Typical SLO targets. Voice needs a tight per-token cadence;
chat balances both; batch is priced on throughput not latency.}
\label{tab:slo-targets}
\begin{tabular}{@{}l l l@{}}
\toprule
Workload & TTFT & TPOT \\
\midrule
Voice / real-time & 200--300 ms & 10--25 ms \\
Interactive chat & 300 ms -- 1 s & 20--50 ms \\
Batch / offline & seconds -- minutes & N/A \\
\bottomrule
\end{tabular}
\end{table}
\paragraph{Regime-aware rules of thumb.} Context length relative to the
machine's balance point $L^*=2N/(\mathrm{AI}\cdot\mathrm{kv\_bpt})$
determines which strategy applies. Below $L^*$ the deployment is
compute-bound and batch scales to $\sim\!2B^*$; above, KV streaming
dominates and $B$ must shrink (Table~\ref{tab:regime-rules}).
\begin{table*}[t]
\centering
\small
\caption{Deployment strategy by context regime.}
\label{tab:regime-rules}
\begin{tabular}{@{}l l l@{}}
\toprule
Regime & Batch & Cost/token \\
\midrule
Short ($S_{kv} < L^*$) & $B \approx 2 B^*$ & low \\
Long ($S_{kv} > L^*$) & smaller $B$, more CP & higher \\
Extreme ($S_{kv} \gg L^*$) & dedicated pool, disagg.\ prefill & much higher \\
\bottomrule
\end{tabular}
\end{table*}
\paragraph{Sample deployment templates.} A single base model is
typically routed to one of several sharding tiers depending on the
request's expected context length; the API gateway dispatches to the
appropriate replica pool (Table~\ref{tab:deploy-templates}).
\begin{table*}[t]
\centering
\small
\caption{Sample deployment tiers for one base model. TP is fixed to
match $h_{kv}$; CP grows with context; $B$ shrinks to hold TPOT.}
\label{tab:deploy-templates}
\begin{tabular}{@{}l c c c c@{}}
\toprule
Tier & CP$\times$TP & Cubes/repl. & Max $S_{kv}$ & Typical $B$ \\
\midrule
Small & $1 \times 8$ & 1 & 32k & 64 \\
Medium & $4 \times 8$ & 4 & 128k & 32 \\
Large & $32 \times 8$ & 32 & 1M & 4 \\
\bottomrule
\end{tabular}
\end{table*}
\paragraph{What to do when each axis binds.} The dominant axis dictates
the first lever to pull (Table~\ref{tab:axis-playbook}). Axes $A$ and
$C$ scale nearly linearly with hardware; axis $B$ is the one where
algorithmic techniques (GQA, MLA, INT4 KV, sparse attention) buy the
most because they attack $\mathrm{kv\_bpt}$ directly rather than adding
chips.
\begin{table*}[t]
\centering
\small
\caption{Playbook per binding axis.}
\label{tab:axis-playbook}
\begin{tabular}{@{}l l l@{}}
\toprule
Binding & First lever & Second lever \\
\midrule
A. Capacity & $\uparrow$ TP / PP & bigger-HBM chip; FP8 / INT4 weights \\
B. KV & $\uparrow$ CP; $\downarrow B$ & GQA / MQA / MLA; INT4 KV; sparse attn \\
C. Throughput & $\uparrow$ replicas (DP) & loosen SLO; smaller model; spec.\ decoding \\
\bottomrule
\end{tabular}
\end{table*}
docs/report/1H-codesign-paper/sections/05c-parallelism-selection.tex
+91
View File
@@ -0,0 +1,91 @@
\subsection{Parallelism Selection: TP $\times$ CP $\times$ PP $\times$ DP $\times$ EP}
\label{sec:parallelism-selection}
Given a capacity-feasible layout (\S\ref{sec:capacity-planning}), the
remaining decision is \emph{how to shard}: which axes to enable and at
what degrees. Memory is the feasibility filter; once you fit, the
design problem is minimising \emph{exposed} communication (communication
that could not be overlapped with compute). The core principle is to
\textbf{rank techniques by how much they communicate per unit of
compute, and assign the chattiest ones to the fastest links.}
\begin{table*}[h]
\centering
\small
\caption{Parallelism axes at a glance. TP AllReduce volume does not
shrink with degree — only compute does — which sets its practical
ceiling near a fast-link (NVLink / intra-SIP) domain.}
\label{tab:parallelism-compare}
\begin{tabular}{@{}l l l l l@{}}
\toprule
Axis & Shards & Comm pattern & Frequency & Hard ceiling \\
\midrule
DP & optimizer state (w/ ZeRO) & AllReduce (grads) & 1$\times$/step & critical batch \\
TP & weights + acts + KV heads & AllReduce & 2$\times$/layer & fast-link domain; $\lesssim h_{kv}$ \\
PP & weights + KV by layer & P2P activation hand-off & per stage boundary & $n_{\text{layers}}$ \\
CP & activations + KV by position & Ring P2P per hop & per layer (overlappable) & $\mathrm{seq\_len} / \mathrm{block}$ \\
EP & expert weights (MoE only) & All-to-all & 2$\times$/MoE layer & $n_{\text{experts}}$ \\
\bottomrule
\end{tabular}
\end{table*}
\paragraph{Symptom-driven axis selection.} The dominant problem dictates
the first knob to try (Table~\ref{tab:parallelism-problem}).
\begin{table}[h]
\centering
\small
\caption{Which axis to reach for first, by dominant problem.}
\label{tab:parallelism-problem}
\begin{tabular}{@{}l l@{}}
\toprule
Dominant problem & First axis \\
\midrule
Weight memory doesn't fit & TP within a fast domain \\
KV of ONE long sequence & CP \\
KV for MANY sequences & DP replicas \\
Single-request TPOT & TP ($\lesssim 8$) \\
Aggregate throughput & Outer DP \\
Model spans multiple nodes & TP intra-node, PP inter-node \\
MoE expert weight memory & EP \\
\bottomrule
\end{tabular}
\end{table}
\paragraph{When to add, when to stop.} Every axis has a natural
saturation point past which further degree hurts more than it helps
(Table~\ref{tab:parallelism-criteria}). Sizing decisions are almost
always constrained by two axes simultaneously (e.g.\ TP by NVLink
domain and $h_{kv}$; CP by ring-hop hiding and per-rank block size).
\begin{table*}[t]
\centering
\small
\caption{Add / stop criteria per axis.}
\label{tab:parallelism-criteria}
\begin{tabular}{@{}l l l@{}}
\toprule
Axis & Add when & Stop when \\
\midrule
DP & more independent requests & global batch $>$ critical \\
TP & weights need sharding or TPOT tight & collectives dominate; local GEMMs shrink \\
PP & depth too large; TP would cross slow links & bubble $>$ $\sim$10\% \\
CP & one sequence needs position sharding & ring hops can't overlap compute \\
EP & MoE expert memory & expert GEMMs too small; routing imbalance \\
\bottomrule
\end{tabular}
\end{table*}
\paragraph{Common misconceptions.} Three widely repeated rules survive
contact with real workloads only partially:
\emph{(i)} memory is not just a feasibility filter but a continuous
performance variable --- more free HBM enables larger $B$, more prefix
cache, higher throughput even after weights fit;
\emph{(ii)} the ``TP $\leq h_{kv}$'' ceiling is an efficiency
preference, not a correctness constraint --- vLLM and others support
KV-head replication and head-dim splitting for TP $>$ $h_{kv}$;
\emph{(iii)} ``always start at TP=8'' is disproven by public
disclosures (DeepSeek-V3 inference uses TP=4, some MLPs at TP=1;
DeepSeek-V3 training uses PP=16 + EP=64 with no TP), which shows the
right starting point is the smallest TP that fits memory and meets
latency, followed by measurement.
docs/report/1H-codesign-paper/sections/05x-fused-kernel.tex
+515
View File
@@ -0,0 +1,515 @@
\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.
% Summary moved to sections/05z-summary.tex so the three new
% subsections (roofline, capacity planning, parallelism selection)
% can appear between the technical results and the closing summary.
% 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).
docs/report/1H-codesign-paper/sections/05z-summary.tex
+22
View File
@@ -0,0 +1,22 @@
\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. Roofline framing
(\S\ref{sec:roofline}) turns this observation into a sizing and sharding
programme in the next two subsections: capacity planning
(\S\ref{sec:capacity-planning}) and parallelism selection
(\S\ref{sec:parallelism-selection}). What this implies for hardware
investment across the report as a whole is taken up in the discussion.
tests/analytical_visualization/_gen_roofline_paper_figs.py
+104
View File
@@ -0,0 +1,104 @@
"""Generate roofline figures for the paper's Roofline Analysis section.
Produces two 2-panel PNGs into docs/report/1H-codesign-paper/figures/:
- roofline_short_context.png (S_kv = 8K) — batch drops per-token cost
- roofline_long_context.png (S_kv = 1M) — batch effect vanishes
Each figure shows both the step-latency view (raw step time vs B, with
weight/compute/KV components) and the cost-per-token view (÷B), which
together make the batch-vs-context story visible at a glance.
Run: python tests/analytical_visualization/_gen_roofline_paper_figs.py
"""
from __future__ import annotations
import sys
from pathlib import Path
# Ensure the repo root is on sys.path so `tests.analytical_visualization`
# imports resolve when invoked as a script.
_REPO_ROOT = Path(__file__).resolve().parents[2]
if str(_REPO_ROOT) not in sys.path:
sys.path.insert(0, str(_REPO_ROOT))
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from tests.analytical_visualization.chip_roofline import (
critical_batch,
per_token_latency_curve,
step_latency_curve,
)
from tests.analytical_visualization.model_config import MachineParams
from tests.analytical_visualization.model_presets import PRESETS
FIGURES_DIR = Path(__file__).resolve().parents[2] / (
"docs/report/1H-codesign-paper/figures"
)
FIGURES_DIR.mkdir(parents=True, exist_ok=True)
MACHINE = MachineParams()
MODEL = PRESETS["Llama 3 70B"].model
B_RANGE = [1, 2, 4, 8, 16, 32, 64, 128, 256]
def _plot_pair(s_kv: int, label: str, out_path: Path) -> None:
step = step_latency_curve(MACHINE, MODEL, B_RANGE, s_kv=s_kv)
tok = per_token_latency_curve(MACHINE, MODEL, B_RANGE, s_kv=s_kv)
xs = [p.batch for p in step]
b_star = critical_batch(MACHINE, MODEL)
fig, (axA, axB) = plt.subplots(1, 2, figsize=(11, 4.2))
# ── Left: step latency (undivided) ────────────────────────────────
axA.plot(xs, [p.weight_s * 1e3 for p in step], "^-",
color="#ffbe0b", label="Weight fetch (flat)")
axA.plot(xs, [p.compute_s * 1e3 for p in step], "o-",
color="#3a86ff", label="Compute (linear)")
axA.plot(xs, [p.kv_s * 1e3 for p in step], "s-",
color="#d90429", label="KV fetch (linear)")
axA.plot(xs, [p.total_s * 1e3 for p in step], "-",
color="#212529", linewidth=2.5, label="Total")
axA.set_xscale("log", base=2)
axA.set_yscale("log")
axA.set_xlabel("Batch size B")
axA.set_ylabel("Step time (ms)")
axA.set_title(f"Step latency ({label} context, $S_{{kv}}={s_kv:,}$)")
axA.grid(True, which="both", alpha=0.3)
axA.legend(fontsize=8, loc="upper left")
# ── Right: cost per token (÷ B) ───────────────────────────────────
axB.plot(xs, [p.weight_s * 1e3 for p in tok], "^-",
color="#ffbe0b", label="Weight fetch ($\\propto 1/B$)")
axB.plot(xs, [p.compute_s * 1e3 for p in tok], "o--",
color="#3a86ff", label="Compute (flat)")
axB.plot(xs, [p.kv_s * 1e3 for p in tok], "s--",
color="#d90429", label="KV fetch (flat)")
axB.plot(xs, [p.total_s * 1e3 for p in tok], "-",
color="#212529", linewidth=2.5, label="Total")
axB.axvline(b_star, linestyle=":", color="#2e7d32",
label=f"$B^*={b_star:.0f}$")
axB.set_xscale("log", base=2)
axB.set_yscale("log")
axB.set_xlabel("Batch size B")
axB.set_ylabel("Per-token time (ms)")
axB.set_title(f"Cost per token ({label} context, $S_{{kv}}={s_kv:,}$)")
axB.grid(True, which="both", alpha=0.3)
axB.legend(fontsize=8, loc="upper right")
plt.tight_layout()
fig.savefig(out_path, dpi=150, bbox_inches="tight")
plt.close(fig)
print(f" wrote {out_path.name}")
def main() -> None:
print("Generating roofline figures for the paper ...")
_plot_pair(8_192, "short", FIGURES_DIR / "roofline_short_context.png")
_plot_pair(1_048_576, "long", FIGURES_DIR / "roofline_long_context.png")
print(f"Done. Output: {FIGURES_DIR}")
if __name__ == "__main__":
main()