diff --git a/docs/report/1H-codesign-paper/build/main.pdf b/docs/report/1H-codesign-paper/build/main.pdf index 9d341a0..b748e92 100644 Binary files a/docs/report/1H-codesign-paper/build/main.pdf and b/docs/report/1H-codesign-paper/build/main.pdf differ diff --git a/docs/report/1H-codesign-paper/figures/roofline_long_context.png b/docs/report/1H-codesign-paper/figures/roofline_long_context.png new file mode 100644 index 0000000..2cda3ee Binary files /dev/null and b/docs/report/1H-codesign-paper/figures/roofline_long_context.png differ diff --git a/docs/report/1H-codesign-paper/figures/roofline_short_context.png b/docs/report/1H-codesign-paper/figures/roofline_short_context.png new file mode 100644 index 0000000..61c8c72 Binary files /dev/null and b/docs/report/1H-codesign-paper/figures/roofline_short_context.png differ diff --git a/docs/report/1H-codesign-paper/main.tex b/docs/report/1H-codesign-paper/main.tex index ded3150..7186d81 100644 --- a/docs/report/1H-codesign-paper/main.tex +++ b/docs/report/1H-codesign-paper/main.tex @@ -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} diff --git a/docs/report/1H-codesign-paper/sections/05-gqa.tex b/docs/report/1H-codesign-paper/sections/05-gqa.tex index 13fd776..8e97760 100644 --- a/docs/report/1H-codesign-paper/sections/05-gqa.tex +++ b/docs/report/1H-codesign-paper/sections/05-gqa.tex @@ -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). diff --git a/docs/report/1H-codesign-paper/sections/05a-roofline.tex b/docs/report/1H-codesign-paper/sections/05a-roofline.tex new file mode 100644 index 0000000..cd2e3af --- /dev/null +++ b/docs/report/1H-codesign-paper/sections/05a-roofline.tex @@ -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. diff --git a/docs/report/1H-codesign-paper/sections/05b-capacity-planning.tex b/docs/report/1H-codesign-paper/sections/05b-capacity-planning.tex new file mode 100644 index 0000000..5ca3b9c --- /dev/null +++ b/docs/report/1H-codesign-paper/sections/05b-capacity-planning.tex @@ -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*} diff --git a/docs/report/1H-codesign-paper/sections/05c-parallelism-selection.tex b/docs/report/1H-codesign-paper/sections/05c-parallelism-selection.tex new file mode 100644 index 0000000..3ac415c --- /dev/null +++ b/docs/report/1H-codesign-paper/sections/05c-parallelism-selection.tex @@ -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. diff --git a/docs/report/1H-codesign-paper/sections/05x-fused-kernel.tex b/docs/report/1H-codesign-paper/sections/05x-fused-kernel.tex new file mode 100644 index 0000000..52c809d --- /dev/null +++ b/docs/report/1H-codesign-paper/sections/05x-fused-kernel.tex @@ -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). diff --git a/docs/report/1H-codesign-paper/sections/05z-summary.tex b/docs/report/1H-codesign-paper/sections/05z-summary.tex new file mode 100644 index 0000000..a4c050b --- /dev/null +++ b/docs/report/1H-codesign-paper/sections/05z-summary.tex @@ -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. diff --git a/tests/analytical_visualization/_gen_roofline_paper_figs.py b/tests/analytical_visualization/_gen_roofline_paper_figs.py new file mode 100644 index 0000000..14febe3 --- /dev/null +++ b/tests/analytical_visualization/_gen_roofline_paper_figs.py @@ -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()