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

gqa short-context: exclude KV deploy from wall; correct HBM peak; batch scaffold

Browse Source 
Metric fixes (test harnesses, deploy artifact + wrong constant):
- wall = max(t_end) - min(t_start): exclude the one-time KV-cache deploy
  from the measured decode/prefill step (it was 92-99% of wall, mapping-
  invariant, and masked the real per-mapping separation).
- PEAK_PE_HBM_BPS 128 -> 256 B/ns (8 channels/PE x 32 GB/s); the old value
  was half the modeled per-PE HBM BW, so hbm_bw_util read >1.0 once wall
  was corrected. All six short-context sweep CSVs regenerated/repatched.

Result: the four KV mappings now separate along a {64,32,16,8}-active-PE
ladder (decode 8-kv/1-kv = 4.8x at 8K, 7.4x at 64K), not "modest/tied" as
before; decode is bandwidth-bound at 46-76% of the 256 GB/s per-PE ceiling.

Report (S5.2 rewrite):
- Replace the tied-wall / 8x-per-PE-util claims (both deploy artifacts)
  with the corrected separation and a density trade-off (per-CUBE KV
  footprint, Fig 16 top panel switched to a wall-invariant metric).
- Add a projected latency-vs-batched-throughput analysis (marked
  projected, not measured): dense mappings win throughput, 1-kv wins
  latency; converges at long context.
- Regenerate Fig 15/16/17/18; fix plot script hardcoded ROOT path.

Batch experiment (Part 2, cube_base):
- Add backward-compatible cube_base=0 scalar to the decode kernel so a
  batched user placed at DPPolicy.cube_start addresses 0-based shards.
  Default preserves single-user behavior (32 decode tests pass unchanged).
- New batch harness (skipped): concurrent B>=2 launches hit a sim routing
  issue (sip0.cube0.pe8); single-user path verified. Concurrency fix next.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook Kang
2026-07-22 16:25:04 -07:00
parent cb6d21f145
commit 711a9a257f
23 changed files with 403 additions and 187 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/report/1H-codesign-paper/sections/05-gqa.tex
+102 -73
View File
@@ -141,10 +141,10 @@ 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,
HBM bandwidth pressure per CUBE grows linearly, the per-CUBE KV cache
grows linearly, 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
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
@@ -160,101 +160,130 @@ We sweep all four mappings $\times$ three composite tiers $\times$
$S_{kv}\in\{8\text{K},16\text{K},32\text{K},64\text{K}\}$ for both
phases (prefill uses $T_q{=}8$ sliced tiles; decode uses $T_q{=}1$). The
headline latency comparison is in Figure~\ref{fig:gqa-short-wall-baseline}.
At short context the two mappings are effectively tied, and they separate
as context grows: the $\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube}$
latency ratio rises monotonically with $S_{kv}$ (decode: $0.98\times$ at
$8$K---within measurement noise---to $1.06\times$ at $64$K, where
\textsf{1-kv-per-cube} is the faster; prefill follows the same trend). The
wall gap is
modest because short-context decode is dominated by setup overhead and
its $M{=}G{=}8$ skinny shape leaves the MAC array idle most of the time;
where the mappings really separate is per-CUBE pressure.
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). Prefill and decode follow the same ordering; the
two mappings are effectively tied at $8$K and separate as context grows,
the \textsf{8-kv-per-cube}/\textsf{1-kv-per-cube} ratio widening from
$0.98\times$ at $8$K to $1.06\times$ at $64$K (\textsf{1-kv-per-cube}
faster at the large end). The composite and
softmax\_merge tiers track the baseline within a few percent at this
shape (see Figure~\ref{fig:gqa-short-a1-variant}).}
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} makes the per-CUBE trade-off
explicit. Per-PE HBM bandwidth utilization is $\sim$\,$8\times$ higher on
\textsf{8-kv-per-cube} than on \textsf{1-kv-per-cube} at every context
length: putting all heads on a single CUBE saturates that CUBE's HBM
channel while leaving seven CUBEs idle, whereas spreading one head per
CUBE keeps each CUBE's read bandwidth well below the
\SI{256}{\giga\byte\per\second} per-PE ceiling. IPCQ traffic moves in the
opposite direction---\textsf{1-kv-per-cube} pays the broadcast or
chain-reduce cost across all eight PEs, while \textsf{8-kv-per-cube} pays
nothing---but absolute IPCQ volume stays under \SI{120}{\kibi\byte} per
run, two orders of magnitude below the HBM traffic. Because per-PE HBM
bandwidth and per-CUBE KV cache (which grows as $\text{kv\_per\_cube}\times
S_{kv}\times d_{\text{head}}$, so $8\times$ more on
\textsf{8-kv-per-cube}) both scale with the number of heads on a CUBE,
the small wall advantage of \textsf{1-kv-per-cube} compounds with much
larger headroom for longer contexts and for batched serving---which is
why it is the choice carried into long context (\S\ref{sec:gqa-long}).
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;
HBM and IPCQ totals are variant-invariant). Top row: per-PE HBM
bandwidth utilization---\textsf{8-kv-per-cube} concentrates roughly
$8\times$ more bandwidth on a single CUBE than \textsf{1-kv-per-cube}.
Bottom row: IPCQ traffic per run on a log scale---\textsf{1-kv-per-cube}
pays the broadcast/reduce cost across the largest group while
\textsf{8-kv-per-cube} (single PE per head) pays zero.}
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 essentially
unchanged across tiers ($< 0.5\%$ in either direction). 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 shows up in
Figure~\ref{fig:gqa-short-gemm-util} as inflated GEMM engine utilization
on the composite and softmax\_merge tiers (roughly $3$ and $6\times$ the
baseline), which is bookkeeping cost rather than added useful work.
Demonstrating fusion benefit therefore requires lifting $M$---either by
batching multiple users (a separate batched-$M$ kernel, deferred to
\S\ref{sec:future}) or by the long-context path, where Q-tile splitting
yields a larger $M$ per PE. In the short-context regime the contribution
of composite/softmax\_merge is a no-regression dispatch-cost isolation;
the substantive design lever is the mapping, and
\textsf{1-kv-per-cube} wins it.
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 (projected).} 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 can run 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}.
Projecting from the measured single-user latencies under perfect overlap
(an idealization---no concurrent multi-user run has been measured yet; see
\S\ref{sec:future}), aggregate throughput at a full SIP is
$\text{users}/\text{latency}$: at $8$K it rises from \SI{0.89}{}
(\textsf{1-kv-per-cube}, 2 users) to \SI{1.48}{}~requests\,/\,\si{\micro\second}
(\textsf{8-kv-per-cube}, 16 users), so the \emph{dense} mapping wins on
throughput even though it is the \emph{slowest} per request; at $64$K the
projected throughputs converge (\SI{0.17}{}--\SI{0.18}{}) because latency
there already scales as $1/C$. The design conclusion is therefore
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. These are
\emph{projections} assuming zero cross-user interference; validating them
against a concurrent multi-user run is 2H work (\S\ref{sec:future}).
\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}: wall-clock
latency is flat across the three tiers within $0.5\%$. With
$M{=}G{=}8 < \textsf{TILE\_M}{=}32$, the composite scheduler pads the
$M$ dimension $4\times$ and any fusion overlap is absorbed by the
padding overhead. Fusion benefit is exercised only when $M$ fills the
supertile---batched-$M$ inference or long-context Q-tile splitting.}
\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 utilization on \textsf{1-kv-per-cube} across the
three composite tiers. The $\sim 3$ and $\sim 6\times$ rise on the
composite and softmax\_merge tiers reflects supertile padding (the
$8\!\to\!32$ pad in $M$) being accounted as GEMM time, not added useful
work. With actual $M=8$, the engine is below \SI{2}{\percent} busy in
all tiers---the kernel is firmly setup- and KV-bandwidth-bound at this
shape.}
\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}