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>
analytical-viz: parallelism rules — Section 11 TP vs CP shared-budget
2026-07-30 15:56:41 -07:00 · 2026-07-30 09:26:22 -07:00 · 2026-07-30 09:11:08 -07:00
13 changed files with 1469 additions and 532 deletions
@@ -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}
@@ -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).
@@ -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.
@@ -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*}
@@ -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.
@@ -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).
@@ -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.
@@ -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()
@@ -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."
    )
@@ -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