mukesh
fc4747668e
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>
27 lines
1.5 KiB
TeX
27 lines
1.5 KiB
TeX
\section{Fused Grouped-Query Attention}
|
|
\label{sec:gqa}
|
|
|
|
Attention is the 1H focus, and it is where the two preceding optimizations
|
|
have to come together. Grouped-Query Attention (GQA) shrinks the KV cache
|
|
by sharing each KV head across a group of query heads (here $h_q=8$ query
|
|
heads to $h_{kv}=1$ KV head, a group factor $G=8$), which makes decoding
|
|
feasible at long context but also makes it acutely memory-bound: a decode
|
|
step processes a single query position ($T_q=1$) against the entire KV
|
|
history, so its arithmetic intensity is low and its time is dominated by
|
|
streaming the KV cache out of HBM. FlashAttention-style tiling with an
|
|
online-softmax merge avoids ever materializing the full score matrix, but
|
|
realizing it as a fast \emph{fused} kernel needs both building blocks from
|
|
this report: efficient GEMM issue (\S\ref{sec:gemm}) for the
|
|
$Q\!\cdot\!K^{\top}$ and $P\!\cdot\!V$ products, and an efficient on-device
|
|
reduction (\S\ref{sec:allreduce}) for the multi-user and
|
|
sequence-parallel KV reductions. \emph{Fused} here is meant in the
|
|
FlashAttention sense---$Q\!\cdot\!K^{\top}$, the online softmax, and
|
|
$P\!\cdot\!V$ collapse into a single kernel that never materializes the
|
|
score matrix---and, beyond that, the cross-device KV reduction is absorbed
|
|
into the same kernel (on PE\_IPCQ) rather than issued as a separate
|
|
all-reduce. This section is the capstone: the fused
|
|
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.
|
|
|