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

Three new subsections at the top of §5, motivated by the deployment-sizing
questions the analytical visualization tool exposes:

- 5.1 Roofline Analysis — B*, L*, short vs long context batch effect
  (measured: 19× per-token cost reduction at short context, 28% at long)
- 5.2 Capacity Planning — 3-axis sizing (weights / KV / SLO), SLO targets,
  regime rules, deployment templates, per-axis playbook
- 5.3 Parallelism Selection — TP × CP × PP × DP × EP comparison,
  symptom-driven axis selection, add/stop criteria, misconceptions

Restructure:
- 05-gqa.tex trimmed to section header + intro
- Existing 5.4-5.7 content (placement, short/long ctx, composite) moved
  to 05x-fused-kernel.tex
- Summary extracted to 05z-summary.tex, cross-refs updated
- main.tex \input order sets the requested subsection sequence
- lmodern loaded to satisfy microtype font-expansion

Two new figures generated from tests/analytical_visualization/chip_roofline:
- roofline_short_context.png (S_kv=8K, batch wins)
- roofline_long_context.png (S_kv=1M, batch stalls at KV floor)
- Generator: tests/analytical_visualization/_gen_roofline_paper_figs.py

4 tables (regime rules, deployment templates, playbook, parallelism
criteria) use table* so they span both columns without overflow.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

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}

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).

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.

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*}

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.

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).

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.

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()

tests/analytical_visualization/app.py

@@ -555,13 +555,14 @@ if _warnings:
 # Auto Suggest is renamed "Parallelism" since it only varies parallelism
 # knobs (hardware is held fixed at the sidebar values).
 (tab_layout, tab_auto, tab_memory, tab_stages, tab_compare, tab_hw,
- tab_roofline, tab_planning) = st.tabs([
+ tab_roofline, tab_planning, tab_prules) = st.tabs([
     "Physical layout",
     "Auto Suggest Parallelism",
     "Memory breakdown", "Per-stage latency",
     "Save & compare", "Auto Hardware",
     "Chip roofline & B*",
     "Capacity planning",
+    "Parallelism rules",
 ])
 
 
@@ -3122,3 +3123,508 @@ with tab_planning:
     ]
     st.dataframe(pd.DataFrame(_playbook_rows), width='stretch',
                   hide_index=True)
+
+
+# ── TAB 9: Parallelism rules ─────────────────────────────────────
+with tab_prules:
+    st.subheader("Parallelism rules — heuristics, constraints, and honest caveats")
+    st.info(
+        "**This is directional guidance, not a production procedure.** "
+        "Real hyperscalers combine the rules below with measured "
+        "collectives, workload traces, and per-architecture search "
+        "(GSPMD, Alpa, hand-tuned Megatron). Public disclosures "
+        "(DeepSeek-V3 training: PP16 + EP64 + no TP; vLLM's "
+        "PP-over-TP recommendation on non-NVLink GPUs) show that no "
+        "universal 'start TP at 8' rule survives contact with real "
+        "workloads. Use these tables as a starting point, then benchmark."
+    )
+
+    # ── Section 1: core principle ─────────────────────────────────
+    st.markdown("### 1. The core principle")
+    st.markdown(
+        "> **Rank the techniques by how much they communicate per unit "
+        "of compute, and assign the chattiest ones to the fastest links.**\n\n"
+        "Memory capacity is the feasibility filter (you fit or you don't). "
+        "Once you fit, every remaining choice is about minimizing "
+        "*exposed* communication (comm that couldn't be overlapped with "
+        "compute)."
+    )
+
+    st.divider()
+
+    # ── Section 2: Parallelism technique comparison ──────────────
+    st.markdown("### 2. Parallelism techniques — comm profile at a glance")
+    _p_compare = [
+        {"Technique": "DP (Data Parallelism)",
+         "Shards": "Nothing (or optimizer state w/ ZeRO)",
+         "Comm pattern": "AllReduce (gradients)",
+         "Frequency": "1× per training step",
+         "Per-link vol as degree ↑": "Flat, but overlappable",
+         "Hard ceiling": "Critical batch size"},
+        {"Technique": "TP (Tensor Parallelism)",
+         "Shards": "Weights + activations + KV by head",
+         "Comm pattern": "AllReduce",
+         "Frequency": "2× per layer",
+         "Per-link vol as degree ↑": "Flat (compute shrinks, comm doesn't)",
+         "Hard ceiling": "≤ NVLink domain size; efficient ≤ n_kv_heads"},
+        {"Technique": "PP (Pipeline Parallelism)",
+         "Shards": "Weights by layer + KV by layer",
+         "Comm pattern": "Point-to-point (activation hand-off)",
+         "Frequency": "Per stage boundary",
+         "Per-link vol as degree ↑": "Negligible per link",
+         "Hard ceiling": "n_layers"},
+        {"Technique": "CP (Context Parallelism)",
+         "Shards": "Activations + KV by position",
+         "Comm pattern": "Ring P2P (per hop)",
+         "Frequency": "Per layer, overlappable",
+         "Per-link vol as degree ↑": "Falls per hop (KV shard shrinks)",
+         "Hard ceiling": "seq_len / block_size"},
+        {"Technique": "EP (Expert Parallelism, MoE only)",
+         "Shards": "Expert weights",
+         "Comm pattern": "All-to-all (dispatch + combine)",
+         "Frequency": "2× per MoE layer",
+         "Per-link vol as degree ↑": "Falls per peer (total flat)",
+         "Hard ceiling": "n_experts"},
+    ]
+    st.dataframe(pd.DataFrame(_p_compare), width='stretch',
+                  hide_index=True)
+    st.caption(
+        "The TP row is the one that decides most configs. Its "
+        "AllReduce **volume doesn't shrink** as you add ranks — only "
+        "compute does. That's why TP has a practical ceiling (~8) "
+        "even without hitting a hard n_kv_heads limit."
+    )
+
+    st.divider()
+
+    # ── Section 3: dominant problem → first axis to try ──────────
+    st.markdown("### 3. What's hurting → which parallelism to reach for")
+    _p_problem = [
+        {"Dominant problem": "Weight memory doesn't fit",
+         "First to try": "TP (within a fast domain)",
+         "Second lever": "PP across nodes; FSDP/ZeRO in training",
+         "Why": "TP shards weights + activations; PP shards by depth"},
+        {"Dominant problem": "Optimizer / gradient memory (training)",
+         "First to try": "FSDP / ZeRO-2 or ZeRO-3",
+         "Second lever": "Increase DP degree",
+         "Why": "ZeRO shards optimizer state across DP ranks"},
+        {"Dominant problem": "Activation memory (long sequences)",
+         "First to try": "CP / sequence parallelism",
+         "Second lever": "Activation checkpointing",
+         "Why": "Activations scale with seq_len; CP shards positions"},
+        {"Dominant problem": "KV cache of ONE long sequence",
+         "First to try": "CP",
+         "Second lever": "TP over KV heads; PP by layer",
+         "Why": "DP can't split one sequence — CP shards positions"},
+        {"Dominant problem": "KV capacity for MANY sequences",
+         "First to try": "DP replicas + request scheduling",
+         "Second lever": "Attention-DP within EP groups",
+         "Why": "Independent sequences shard by DP for free"},
+        {"Dominant problem": "MoE expert weight memory",
+         "First to try": "EP",
+         "Second lever": "Redundant copies of hot experts",
+         "Why": "EP places different experts on different ranks"},
+        {"Dominant problem": "Single-request latency (TPOT)",
+         "First to try": "TP (up to ~8)",
+         "Second lever": "CP for very long prefill",
+         "Why": "TP divides weight-fetch time linearly"},
+        {"Dominant problem": "Aggregate throughput (many users)",
+         "First to try": "Outer DP",
+         "Second lever": "Higher B per replica",
+         "Why": "DP scales throughput linearly at zero comm cost"},
+        {"Dominant problem": "Model spans multiple nodes",
+         "First to try": "TP within node, PP across nodes",
+         "Second lever": "Consider CP if seq_len is also large",
+         "Why": "Slow scale-out links favor PP's point-to-point"},
+        {"Dominant problem": "Expert layer dominates (MoE)",
+         "First to try": "EP",
+         "Second lever": "Attention-DP + EP combined",
+         "Why": "vLLM supports DP-attn + EP-experts topology"},
+        {"Dominant problem": "Global batch already at critical size (training)",
+         "First to try": "TP × PP × CP × EP",
+         "Second lever": "Reduce DP further",
+         "Why": "Additional DP past critical batch wastes samples"},
+    ]
+    st.dataframe(pd.DataFrame(_p_problem), width='stretch',
+                  hide_index=True)
+
+    st.divider()
+
+    # ── Section 4: per-axis select/stop criteria ─────────────────
+    st.markdown("### 4. When to add / when to stop — per axis")
+    _p_criteria = [
+        {"Axis": "DP",
+         "Add when": "More independent batches/requests are available",
+         "Stop when": "Global batch exceeds critical batch, or replicas "
+                     "can't fit"},
+        {"Axis": "TP",
+         "Add when": "Layers/weights need sharding or TPOT needs improvement",
+         "Stop when": "Collectives dominate step time OR local GEMMs "
+                     "become too small OR TP > n_kv_heads (efficiency)"},
+        {"Axis": "PP",
+         "Add when": "Model depth must be distributed, or TP would cross "
+                     "slow links",
+         "Stop when": "Bubble fraction > ~10% OR stage imbalance grows"},
+        {"Axis": "CP",
+         "Add when": "A single sequence is long enough to need position "
+                     "sharding",
+         "Stop when": "Ring hops can't be hidden by compute OR local "
+                     "sequence blocks are too small for efficient kernels"},
+        {"Axis": "EP",
+         "Add when": "MoE model + expert weights should be distributed",
+         "Stop when": "Expert GEMMs become too small OR routing "
+                     "imbalance grows OR all-to-all latency dominates"},
+        {"Axis": "FSDP / ZeRO",
+         "Add when": "Training-state memory is the dominant problem",
+         "Stop when": "Parameter all-gather / reduce-scatter starts "
+                     "dominating step time"},
+        {"Axis": "Sequence parallelism",
+         "Add when": "TP is used during training (typically bundled)",
+         "Stop when": "Rarely disabled alone — coupled to TP support"},
+    ]
+    st.dataframe(pd.DataFrame(_p_criteria), width='stretch',
+                  hide_index=True)
+
+    st.divider()
+
+    # ── Section 5: naive rules vs nuanced reality ────────────────
+    st.markdown("### 5. Common misconceptions vs reality")
+    _p_miscon = [
+        {"Naive rule": "Memory is only a feasibility filter",
+         "Reality": "Memory headroom is also a continuous performance "
+                    "variable. More free HBM = larger B, more prefix "
+                    "cache, higher throughput — even after weights fit."},
+        {"Naive rule": "TP has a hard ceiling at n_kv_heads",
+         "Reality": "vLLM (and others) support KV-head replication when "
+                    "TP > n_kv_heads. It's an efficiency preference, "
+                    "not a correctness ceiling."},
+        {"Naive rule": "Always start TP at 8",
+         "Reality": "Start with the smallest TP that fits memory + "
+                    "meets latency. DeepSeek-V3 inference uses TP=4 "
+                    "(and TP=1 for some MLPs) explicitly to limit "
+                    "collective overhead."},
+        {"Naive rule": "EP comm falls as EP grows",
+         "Reality": "Per-peer traffic falls, total outgoing traffic is "
+                    "flat (~N·k·H). Large EP hurts via startup overhead, "
+                    "small local GEMMs, and imbalance."},
+        {"Naive rule": "Attention-DP shards KV 'for free'",
+         "Reality": "Only across DIFFERENT sequences. One 1M-token "
+                    "sequence still needs CP/TP/PP to shard — DP can't "
+                    "split it. Attention weights are still replicated "
+                    "across DP ranks."},
+        {"Naive rule": "DP > CP always if you have many requests",
+         "Reality": "For MANY short requests, yes. For ONE long "
+                    "sequence with strict prefill SLO, CP wins because "
+                    "DP can't help. Different dimensions solve "
+                    "different problems."},
+        {"Naive rule": "Avoid PP inside the domain — it's worthless",
+         "Reality": "vLLM specifically recommends PP over TP on GPUs "
+                    "without NVLink. DeepSeek-V3 training used PP=16 "
+                    "and no TP. PP shines when TP collectives cost more "
+                    "than PP bubbles."},
+        {"Naive rule": "FSDP/ZeRO-3 is THE training substrate",
+         "Reality": "DeepSeek-V3 used ZeRO-1 + PP16 + EP64 — no ZeRO-3. "
+                    "FSDP is one option; the right choice depends on "
+                    "the topology, batch size, and model shape."},
+        {"Naive rule": "Pipeline bubble = (P−1)/m",
+         "Reality": "Actual bubble fraction = (P−1)/(m + P−1). The "
+                    "(P−1)/m form is an approximation valid only when "
+                    "m ≫ P. Interleaved / 1F1B / zero-bubble schedules "
+                    "have their own formulas."},
+    ]
+    st.dataframe(pd.DataFrame(_p_miscon), width='stretch',
+                  hide_index=True)
+
+    st.divider()
+
+    # ── Section 6: inference decision procedure ──────────────────
+    st.markdown("### 6. Inference decision procedure (in order)")
+    _p_infer = [
+        {"Step": 1,
+         "Action": "Feasibility check",
+         "Criterion": "Compute weights + KV(B, S_kv). If it fits on "
+                      "one GPU, stop — use plain replication (DP only)."},
+        {"Step": 2,
+         "Action": "Set TP from TPOT budget",
+         "Criterion": "TP is the latency knob — divides weight-fetch "
+                      "time. Prefer TP=4–8 for latency-oriented; TP=1–2 "
+                      "for throughput-oriented. Try to keep TP ≤ n_kv_heads."},
+        {"Step": 3,
+         "Action": "Fill remaining domain with EP (if MoE)",
+         "Criterion": "EP is the cost knob for MoE. Cap at n_experts. "
+                      "Overlap dispatch/combine with a second micro-batch."},
+        {"Step": 4,
+         "Action": "Attention-DP over EP groups",
+         "Criterion": "Different attention replicas handle different "
+                      "sequence subsets → KV shards by sequence for "
+                      "free, no extra collective."},
+        {"Step": 5,
+         "Action": "CP only if attention-DP can't help",
+         "Criterion": "Trigger: one very long sequence or per-request "
+                      "KV exceeds one GPU. For many short concurrent "
+                      "requests, DP strictly wins."},
+        {"Step": 6,
+         "Action": "PP across racks if the model still won't fit",
+         "Criterion": "Only across slow scale-out boundaries. Inside a "
+                      "fast NVLink domain, PP is usually inferior to TP."},
+        {"Step": 7,
+         "Action": "Replicate the whole unit as outer DP",
+         "Criterion": "The throughput axis. Cap by traffic + critical "
+                      "batch (training) or by aggregate request "
+                      "arrival rate (inference)."},
+    ]
+    st.dataframe(pd.DataFrame(_p_infer), width='stretch',
+                  hide_index=True)
+
+    st.divider()
+
+    # ── Section 7: concrete inference cases ──────────────────────
+    st.markdown("### 7. Inference cases → recommended shape")
+    _p_cases = [
+        {"Case": "A. Fits on one GPU (with KV headroom)",
+         "Recommended": "TP=1, PP=1, CP=1; scale via outer DP",
+         "Notes": "Benchmark TP=2/4 if TPOT is tight"},
+        {"Case": "B. Fits within one NVLink node",
+         "Recommended": "TP = 2/4/8 (smallest that fits + meets SLO)",
+         "Notes": "Consider (TP=2, PP=2) for very deep models"},
+        {"Case": "C. Spans multiple nodes",
+         "Recommended": "TP = GPUs per fast domain, PP = "
+                        "number of nodes; outer DP for throughput",
+         "Notes": "vLLM's standard multi-node recipe"},
+        {"Case": "D. One request with extreme context (≫ L*)",
+         "Recommended": "TP × CP × PP (DP can't split one sequence)",
+         "Notes": "Long-context tier; heavy CP dominates"},
+        {"Case": "E. Many moderate-length requests",
+         "Recommended": "Minimum TP × PP per replica; maximize DP = "
+                        "total_GPUs / (TP × PP)",
+         "Notes": "Throughput-optimized; the mainline serving case"},
+    ]
+    st.dataframe(pd.DataFrame(_p_cases), width='stretch',
+                  hide_index=True)
+
+    st.divider()
+
+    # ── Section 8: training baseline recipe ──────────────────────
+    st.markdown("### 8. Training baseline recipe (in order)")
+    _p_train = [
+        {"Step": 1,
+         "Action": "Start with DP (or sharded DP / FSDP)",
+         "Rationale": "Simplest; shards optimizer state; scales at "
+                      "throughput cost only"},
+        {"Step": 2,
+         "Action": "Add TP if large layers don't fit or GEMMs stay "
+                   "efficient after sharding",
+         "Rationale": "TP divides weight memory + activation memory"},
+        {"Step": 3,
+         "Action": "Add CP when sequence activations dominate",
+         "Rationale": "Activation memory ~ seq_len × d; CP shards "
+                      "positions"},
+        {"Step": 4,
+         "Action": "Add PP when depth still doesn't fit or scaling "
+                   "across nodes",
+         "Rationale": "Depth sharding + slow-link tolerance"},
+        {"Step": 5,
+         "Action": "Add EP for MoE expert weights",
+         "Rationale": "Distributes expert weights; preserves tokens "
+                      "per expert"},
+        {"Step": 6,
+         "Action": "Use remaining devices as outer DP",
+         "Rationale": "Global batch = DP × micro-batch × grad-accum; "
+                      "cap by critical batch"},
+    ]
+    st.dataframe(pd.DataFrame(_p_train), width='stretch',
+                  hide_index=True)
+    st.caption(
+        "Total GPUs = **DP × TP × PP × CP × EP**. Global batch = "
+        "**DP × micro-batch × grad-accum**. The critical batch size "
+        "(NOT bandwidth) is what forces model parallelism at very "
+        "large scale — past that, additional DP wastes samples."
+    )
+
+    st.divider()
+
+    # ── Section 9: sanity checks for any config ──────────────────
+    st.markdown("### 9. Sanity checks for any candidate config")
+    _p_sanity = [
+        {"Check": "TP stays inside one NVLink / scale-up domain?",
+         "Pass criterion": "Yes",
+         "If not": "TP AllReduce crosses slow links → step time balloons"},
+        {"Check": "TP ≤ n_kv_heads?",
+         "Pass criterion": "Yes (or KV replication is acceptable)",
+         "If not": "Wastes KV BW to replication; usually not worth it"},
+        {"Check": "Any KV or weight state duplicated inside the "
+                  "scale-up domain?",
+         "Pass criterion": "No",
+         "If not": "You're wasting aggregate HBM bandwidth"},
+        {"Check": "Exposed (non-overlapped) comm fraction?",
+         "Pass criterion": "≤ ~10% of step time",
+         "If not": "Something is misplaced — swap axes across links"},
+        {"Check": "Pipeline bubble fraction (P−1)/(m+P−1)?",
+         "Pass criterion": "≤ ~10%",
+         "If not": "Need more microbatches or fewer stages"},
+        {"Check": "Predicted step time meets TPOT SLO?",
+         "Pass criterion": "Yes",
+         "If not": "Reduce B, add DP replicas, or accept slower SLO"},
+        {"Check": "B sits in the 2–3× B* sweet spot?",
+         "Pass criterion": "Yes",
+         "If not": "Below → underused compute; above → wasted latency"},
+    ]
+    st.dataframe(pd.DataFrame(_p_sanity), width='stretch',
+                  hide_index=True)
+
+    st.divider()
+
+    # ── Section 10: honest caveats ────────────────────────────────
+    st.markdown("### 10. One-paragraph rule + caveats")
+    st.markdown(
+        "> **First use the smallest parallelism group that makes the "
+        "model, training state, activations, and target KV capacity "
+        "fit. Use TP mainly inside the fastest scale-up domain, PP to "
+        "divide model depth — especially across slower boundaries — CP "
+        "to divide an individual long sequence, and EP to distribute "
+        "MoE experts while preserving enough tokens per expert. Use DP "
+        "across independent batches or requests for throughput. Then "
+        "benchmark several legal combinations using real kernel "
+        "shapes, collective measurements, workload distributions, and "
+        "end-to-end SLOs; choose based on exposed communication and "
+        "cost per token, not communication volume alone.**"
+    )
+    st.warning(
+        "**Honest caveat.** The specific configurations frontier labs "
+        "run are mostly unpublished, and they retune per architecture "
+        "+ hardware generation. What's stable is the *ordering "
+        "principle* and the *constraint list* — those follow from the "
+        "hardware, not from anyone's cleverness. Public counterexamples "
+        "(DeepSeek-V3, vLLM PP-on-non-NVLink) show that no single "
+        "recipe fits all workloads."
+    )
+
+    st.divider()
+
+    # ── Section 11: TP vs CP — the shared-budget tradeoff ─────────
+    st.markdown("### 11. TP vs CP: the shared-budget tradeoff")
+    st.markdown(
+        "**Fixing `TP × CP` doesn't equalize the two configurations.** "
+        "Larger TP additionally divides *model-weight* traffic; larger "
+        "CP only divides the *sequence* dimension. That's why (TP=8, "
+        "CP=2) often beats (TP=2, CP=8) during decode — not because "
+        "CP comm is heavier, but because CP doesn't touch weights."
+    )
+
+    st.markdown("#### 11a. Same TP·CP = 16 → different weight sharding")
+    _cptp_config = [
+        {"Config": "TP=8, CP=2",
+         "Weight shard / GPU": "W / 8",
+         "Context shard": "L / 2",
+         "KV shard / GPU": "KV / 16",
+         "Weight bytes / PE (Llama-like layer)": "~214 MB",
+         "KV read / PE @ 1M ctx": "~256 MB"},
+        {"Config": "TP=2, CP=8",
+         "Weight shard / GPU": "W / 2",
+         "Context shard": "L / 8",
+         "KV shard / GPU": "KV / 16",
+         "Weight bytes / PE (Llama-like layer)": "~856 MB (4× more)",
+         "KV read / PE @ 1M ctx": "~256 MB"},
+    ]
+    st.dataframe(pd.DataFrame(_cptp_config), width='stretch',
+                  hide_index=True)
+    st.caption(
+        "Concrete Llama-like layer: `d_model=8192`, `d_ff=28672`, "
+        "`H_kv=8`, `d_head=128`, BF16 → weights per layer ≈ 1.7 GB; "
+        "KV read at 1M context ≈ 4.1 GB per layer. **Both configs "
+        "give the same KV sharding**, but TP=8 reads 4× fewer weight "
+        "bytes per PE per token."
+    )
+
+    st.markdown("#### 11b. What each dimension actually shards")
+    _cptp_shard = [
+        {"Axis": "TP",
+         "Shards weights?": "✓ (by head)",
+         "Shards KV?": "✓ (by KV head)",
+         "Shards context (attention math)?": "✗"},
+        {"Axis": "CP",
+         "Shards weights?": "✗",
+         "Shards KV?": "✓ (by position)",
+         "Shards context (attention math)?": "✓"},
+        {"Axis": "PP",
+         "Shards weights?": "✓ (by layer)",
+         "Shards KV?": "✓ (by layer)",
+         "Shards context (attention math)?": "✗"},
+        {"Axis": "DP",
+         "Shards weights?": "✗ (or optimizer only in training w/ ZeRO)",
+         "Shards KV?": "✓ (across independent sequences)",
+         "Shards context (attention math)?": "✗"},
+    ]
+    st.dataframe(pd.DataFrame(_cptp_shard), width='stretch',
+                  hide_index=True)
+    st.caption(
+        "Key asymmetry: **CP doesn't touch the FFN weight footprint "
+        "at all**. FFN is a huge chunk of decode weight-fetch cost, "
+        "so if the FFN block is dominant, TP is the only axis that "
+        "helps."
+    )
+
+    st.markdown("#### 11c. When to prefer TP vs CP")
+    _cptp_prefer = [
+        {"Situation": "Weight fetch + FFN dominates step time",
+         "Prefer": "Higher TP",
+         "Why": "TP divides weight bandwidth linearly; CP doesn't touch it"},
+        {"Situation": "KV read + attention dominates",
+         "Prefer": "Higher CP",
+         "Why": "CP shards positions; attention scales O(S_kv)"},
+        {"Situation": "Moderate context, FFN-heavy model",
+         "Prefer": "Higher TP",
+         "Why": "Attention is a small fraction; FFN weight matters more"},
+        {"Situation": "Very long context (≫ L*), attention-heavy",
+         "Prefer": "Higher CP",
+         "Why": "KV read is the wall; CP lets multiple PEs share it"},
+        {"Situation": "TPOT-tight decode",
+         "Prefer": "Higher TP (up to ~n_kv_heads)",
+         "Why": "Divides the big cost (weights); collectives small at T_q=1"},
+        {"Situation": "Long-prompt prefill under TTFT SLO",
+         "Prefer": "Higher CP with `kv` ring variant",
+         "Why": "CP hops overlap with attention compute; bytes stay bounded"},
+    ]
+    st.dataframe(pd.DataFrame(_cptp_prefer), width='stretch',
+                  hide_index=True)
+
+    st.markdown("#### 11d. CP communication profile (prefill vs decode)")
+    _cptp_comm = [
+        {"Mode": "Prefill (T_q = prompt length, large)",
+         "CP comm bytes/hop":
+             "~O(S_local · H_kv · d_h) with kv-variant (bounded)",
+         "Overlap potential": "Good — hides in attention compute",
+         "Real cost": "Moderate",
+         "Notes": "kv-variant: bytes independent of T_q → good scaling; "
+                  "qoml-variant grows with T_q, worst case"},
+        {"Mode": "Decode (T_q = 1)",
+         "CP comm bytes/hop": "O(d_model) for merge (small)",
+         "Overlap potential": "Poor — only 1 query token",
+         "Real cost": "Sync-latency-dominated",
+         "Notes": "~2·(CP−1)·α per layer; α × hops beats bytes at T_q=1"},
+    ]
+    st.dataframe(pd.DataFrame(_cptp_comm), width='stretch',
+                  hide_index=True)
+    st.caption(
+        "So CP comm isn't necessarily *heavier* than TP comm — it's a "
+        "different profile. In decode, CP suffers from per-hop α "
+        "latency (small bytes, many syncs); in prefill, CP with the "
+        "right ring variant can be very cheap and overlappable."
+    )
+
+    st.divider()
+    st.markdown("#### 11e. Decode cost decomposition (mental model)")
+    st.markdown(
+        "> `T_decode ≈ T_weight_read(TP) + T_KV_read(TP, CP) + "
+        "T_TP_comm(TP) + T_CP_merge(CP)`\n\n"
+        "- **Larger TP**: cuts `T_weight_read`, cuts `T_KV_read` (by "
+        "KV head), grows `T_TP_comm`, may make GEMMs too small.\n"
+        "- **Larger CP**: cuts `T_KV_read` (by position) and speeds up "
+        "attention math, **doesn't touch weights or FFN**, grows "
+        "`T_CP_merge` sync overhead.\n\n"
+        "So the real question isn't 'which comm is heavier' — it's "
+        "**'which cost is dominating T_decode right now, and which "
+        "axis attacks it?'** For weight-fetch-bound decode with "
+        "moderate context, larger TP is the right lever. For "
+        "KV-bandwidth-bound decode with extreme context, larger CP is."
+    )

tests/analytical_visualization/test_stage_shapes.py

@@ -165,3 +165,23 @@ def test_ttft_tpot_section_wired_on_layout_tab():
     pos = src.index('"Full-model latency (TTFT & TPOT)"')
     tab_memory_pos = src.index("with tab_memory:")
     assert pos < tab_memory_pos, "TTFT section must live inside tab_layout"
+
+
+def test_parallelism_rules_tab_registered():
+    """The 'Parallelism rules' tab is wired into st.tabs and has a
+    corresponding `with tab_prules:` block. Guards the new tab existing."""
+    from pathlib import Path
+    src = (
+        Path(__file__).parent / "app.py"
+    ).resolve().read_text(encoding="utf-8")
+    assert src.count('"Parallelism rules"') == 1
+    assert "with tab_prules:" in src
+
+
+def test_tp_vs_cp_section_present():
+    """Section 11 'TP vs CP: the shared-budget tradeoff' must appear."""
+    from pathlib import Path
+    src = (
+        Path(__file__).parent / "app.py"
+    ).resolve().read_text(encoding="utf-8")
+    assert src.count("TP vs CP: the shared-budget tradeoff") == 1