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

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

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

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

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

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

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

docs/report/1H-codesign-paper/main.tex

@@ -15,6 +15,7 @@
 \usepackage{caption}
 \usepackage{subcaption}
 \captionsetup{font=small,labelfont=bf}
+\usepackage{lmodern}  % scalable Latin Modern fonts (required by microtype expansion)
 \usepackage{microtype}
 \usepackage{tikz}
 \usetikzlibrary{arrows.meta,positioning,calc,fit}
@@ -39,7 +40,12 @@ AGI Computing Lab, System Technology Group\\
 \input{sections/02-platform}
 \input{sections/03-gemm}
 \input{sections/04-allreduce}
-\input{sections/05-gqa}
+\input{sections/05-gqa}                  % section header + intro
+\input{sections/05a-roofline}            % 5.1 Roofline Analysis
+\input{sections/05b-capacity-planning}   % 5.2 Capacity Planning
+\input{sections/05c-parallelism-selection} % 5.3 Parallelism Selection
+\input{sections/05x-fused-kernel}        % 5.4-5.7 Placement + short/long ctx + composite
+\input{sections/05z-summary}             % 5.8 Summary
 \input{sections/06-agentic}
 \input{sections/07-hw-spec-search}
 \input{sections/08-discussion}

docs/report/1H-codesign-paper/sections/05-gqa.tex

@@ -24,533 +24,3 @@ kernel that uses the composite command and PE\_IPCQ at the same time.
 Multi-head attention (MHA) was studied in prior work and serves here as
 the established baseline rather than being re-derived.
 
-\subsection{Data Placement Policy}
-\label{sec:gqa-placement}
-
-Long-context decode is bound by the KV cache, so the first-order design
-question is how to place that cache---and the running softmax state it
-feeds---across the two hardware axes the machine exposes: the CUBEs and,
-within each CUBE, the PEs (here $C{=}8$ CUBEs $\times$ $P{=}8$ PEs, for
-$C\!\cdot\!P{=}64$ attention engines over one KV-head group on the
-LLaMA-3.1-70B target). Each axis can \emph{replicate} the KV cache or
-\emph{shard} it, and a shard can run along the sequence dimension
-$S_{kv}$ or the head dimension $d_{\text{head}}$. The cross product is a
-small, enumerable taxonomy; six placements span its meaningful corners
-(Figure~\ref{fig:gqa-kv-sharding}).
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_kv_sharding_diagram.png}
-\caption{The six KV-placement strategies, drawn on the
-$S_{kv}\!\times\!d_{\text{head}}$ KV tensor (rows = sequence, columns =
-head dimension). Cube colour bands and dashed PE dividers show which
-axis each level shards. Cases~1--3 either replicate the cache or shard
-it on a single axis (8-way at most); Cases~4--6 reach a full 64-way
-split, three different ways: Case~4 splits $S_{kv}$ across CUBEs and
-$d_{\text{head}}$ across PEs, Case~5 the mirror, and Case~6~$\star$
-splits $S_{kv}$ on \emph{both} axes.}
-\label{fig:gqa-kv-sharding}
-\end{figure}
-
-Two quantities decide which placement is viable, and they pull against
-each other. The first is \textbf{per-PE KV memory}. With a per-PE HBM
-budget of \SI{6.0}{\giga\byte} and \SI{1.76}{\giga\byte} of attention
-weights resident, the KV headroom is \SI{4.24}{\giga\byte} per PE. At a
-production context of $S_{kv}{=}1\,\text{M}$ tokens the unsharded cache
-is \SI{40}{\giga\byte}/PE (Case~1), an 8-way shard is \SI{5}{\giga\byte}
-(Cases~2--3)---both \emph{over} the headroom---while only the 64-way
-placements bring it to \SI{640}{\mega\byte}/PE (Cases~4--6), comfortably
-inside budget. Memory alone therefore eliminates Cases~1--3 at long
-context. The second quantity is \textbf{communication per token}, and it
-is what separates the three survivors. Sharding $d_{\text{head}}$
-(Cases~4--5) makes each PE hold only a slice of every head, so the
-$Q\!\cdot\!K^{\top}$ score is \emph{partial} and must be all-reduced
-across the slice owners on every token---a reduction whose volume scales
-with $S_{kv}$ ($\sim$\SI{166}{\mega\byte}/token analytically, intra-CUBE
-on the NoC for Case~4, inter-CUBE on UCIe for Case~5). Case~6~$\star$
-instead shards $S_{kv}$ on both axes, so every PE computes a
-\emph{complete} score over its own token range and only the small running
-softmax state $(m,\ell,O)$ is merged across PEs
-($\sim$\SI{6.2}{\mega\byte}/token)---a $\sim$27$\times$ lighter collective
-than the $d_{\text{head}}$-split designs.
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_summary.png}
-\caption{Long-context placement analysis at $S_{kv}{=}1\,\text{M}$ tokens.
-\emph{Left:} the per-PE HBM budget---\SI{1.76}{\giga\byte} of attention
-weights leave \SI{4.24}{\giga\byte} of KV headroom (red line).
-\emph{Middle:} per-PE KV memory per case (log scale); only the 64-way
-placements (Cases~4--6, \SI{640}{\mega\byte}) clear the headroom, while
-the unsharded (\SI{40}{\giga\byte}) and 8-way (\SI{5}{\giga\byte}) cases
-overflow. \emph{Right:} analytical communication per token (log scale);
-the $d_{\text{head}}$-split Cases~4--5 pay a partial-score all-reduce
-($\sim$\SI{166}{\mega\byte}/token) that the both-axes-$S_{kv}$ split of
-Case~6~$\star$ avoids ($\sim$\SI{6.2}{\mega\byte}/token, merging only the
-softmax state). On these two axes Case~6 (marked $\star$ in the figure)
-is the single placement that lands both inside the memory budget and at
-low per-token communication; whether that combination is the right one to
-pick is a regime-specific question taken up next.}
-\label{fig:gqa-budget}
-\end{figure}
-
-These two costs---per-PE KV memory and per-token communication---are
-intrinsic properties of each placement, fixed by how it shards the cache
-and independent of the workload regime. They do not by themselves name a
-winner: a short prompt where the whole cache fits on one PE values low
-communication and tolerates replication, whereas a million-token decode
-is bound by the memory wall and will pay communication to escape it.
-Which placement is appropriate is therefore a per-regime question, which
-the short- and long-context subsections that follow answer by running the
-options on the simulator and reading off latency, traffic, and the
-redundant compute each one induces.
-
-\subsection{Inference with Short-Context Length}
-\label{sec:gqa-short}
-
-The fused GQA kernel issues its matrix products as scheduler-managed
-composite commands and keeps the online-softmax merge and the cross-device
-KV reduction inside the kernel, on PE\_IPCQ. Two kernel families cover the
-two phases. The \emph{prefill} kernel splits the query tile across the PEs
-of a group and broadcasts each KV tile from the group's root PE over
-intra-CUBE IPCQ, so the HBM K/V read is paid once per group instead of
-once per PE. The \emph{decode} kernel sequence-shards the KV cache across
-the same group of PEs, runs per-PE local attention, then chain-reduces the
-partial $(m,\ell,O)$ triples back up to the group root. Two further
-primitives make long context practical: a \emph{lazy load} that issues the
-KV \textsf{DMA\_READ} and returns immediately, auto-waiting only at first
-use so KV load overlaps score computation; and per-tile \emph{scratch
-recycling} that keeps the running accumulators in a persistent arena while
-freeing per-tile temporaries, so the kernel fits the
-\SI{1}{\mebi\byte} scratch budget across many tiles. A further refinement
-restructures the decode step so the per-tile matrix products and the
-online-softmax merge are issued as \emph{composite} commands rather than
-hand-tiled primitives; \S\ref{sec:gqa-composite} measures it against the
-primitive baseline at long context.
-
-This is a different decomposition from the six-case placement taxonomy of
-\S\ref{sec:gqa-placement}: that taxonomy spreads a single KV-head group
-across all 64 PEs and is the right lens for long context
-(\S\ref{sec:gqa-long}), whereas short context---where the whole cache fits
-comfortably---turns instead on how the eight KV heads are distributed
-across the eight CUBEs. The design question for short context is therefore
-not whether to fuse softmax---that is settled---but how to distribute the
-eight KV heads across the eight CUBEs of a SIP. Four mappings cover the spectrum, named
-by KV-head count per CUBE: \textsf{1-kv-per-cube} dedicates one whole
-CUBE to each head and uses all eight PEs of that CUBE on the head's
-sequence shard; \textsf{2-kv-per-cube} packs two heads per CUBE with
-group size four; \textsf{4-kv-per-cube} packs four heads with group size
-two; \textsf{8-kv-per-cube} loads all heads onto a single CUBE with one
-PE per head. The mapping fixes a trade-off: as KV heads per CUBE grows, the per-CUBE KV
-footprint grows linearly and the mapping uses proportionally fewer CUBEs
-(eight down to one), but the intra-CUBE IPCQ broadcast/reduce cost shrinks
-to zero (8-kv-per-cube uses a single PE per head and has no group
-communication).
-
-A second axis is how each GEMM tile is issued. We isolate this with three
-\emph{composite tiers} applied to the same mapping: \textsf{without
-composite} uses primitive \textsf{tl.dot} and unfused softmax;
-\textsf{with composite (GEMM-only)} issues the GEMMs as
-\textsf{tl.composite(op="gemm")} commands so the scheduler can pipeline
-the tile stages but leaves softmax as primitives; \textsf{with composite
-+ softmax\_merge} additionally folds the softmax fold into the $P\!\cdot\!V$
-composite through a named \textsf{softmax\_merge} prologue recipe.
-
-We sweep all four mappings $\times$ three composite tiers $\times$
-$S_{kv}\in\{8\text{K},16\text{K},32\text{K},64\text{K}\}$ for both
-phases (prefill uses $T_q{=}8$ sliced tiles; decode uses $T_q{=}1$). The
-headline latency comparison is in Figure~\ref{fig:gqa-short-wall-baseline}.
-The mappings separate cleanly, and in proportion to how many PEs each
-dedicates to a head: \textsf{1-kv-per-cube} spreads one head across all
-eight PEs of its CUBE, \textsf{8-kv-per-cube} gives each head a single PE,
-and the four mappings form a $\{64,32,16,8\}$-active-PE ladder. Decode
-latency tracks that ladder inversely---at $S_{kv}{=}64$K the
-\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube} ratio is $7.4\times$
-($\SI{87.0}{}$ vs.\ $\SI{11.8}{\micro\second}$), and even at $8$K it is
-$4.8\times$ ($\SI{10.8}{}$ vs.\ $\SI{2.2}{\micro\second}$); prefill shows
-the same ordering more gently ($2.3\times$ at $64$K). The reason is that
-decode here is \emph{bandwidth-bound}: each active PE streams its KV shard
-at \SI{46}{}--\SI{76}{\percent} of the \SI{256}{\giga\byte\per\second}
-per-PE ceiling (rising with context), so a mapping's speed is set by how
-many PEs it puts to work---more PEs, more aggregate HBM bandwidth, lower
-latency. The $M{=}G{=}8$ score GEMM is skinny and leaves the MAC array
-lightly loaded throughout; the lever here is data movement, not compute.
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_short_context/wall_variant1_mode_compare.png}
-\caption{Wall-clock latency of the four short-context mappings, baseline
-kernel (no composite; log scale). The mappings separate along the
-$\{64,32,16,8\}$-active-PE ladder set by heads-per-CUBE:
-\textsf{1-kv-per-cube} (64 PEs) is fastest and \textsf{8-kv-per-cube}
-(8 PEs) slowest, by $4.8\times$ at $8$K growing to $7.4\times$ at $64$K in
-decode ($2.3\times$ at $64$K in prefill). Decode is bandwidth-bound, so
-latency scales inversely with the number of active PEs. Composite tiers
-track the baseline at this skinny shape
-(Figure~\ref{fig:gqa-short-a1-variant}).}
-\label{fig:gqa-short-wall-baseline}
-\end{figure}
-
-Figure~\ref{fig:gqa-short-tradeoff} shows what the mappings trade for that
-latency. The differentiator is \emph{density}: \textsf{8-kv-per-cube}
-concentrates all eight heads onto one CUBE, so its per-CUBE KV footprint is
-$8\times$ that of \textsf{1-kv-per-cube} (one CUBE per head), and it
-occupies a single one of the SIP's sixteen CUBEs against
-\textsf{1-kv-per-cube}'s eight. Per-PE HBM \emph{utilization}, by contrast,
-is similar across mappings---every active PE runs near the same
-\SI{46}{}--\SI{76}{\percent} of the per-PE ceiling---so the latency ladder
-comes from the \emph{count} of active PEs, not from any per-PE
-concentration. IPCQ traffic runs opposite to density:
-\textsf{1-kv-per-cube} pays the eight-PE chain-reduce while
-\textsf{8-kv-per-cube} (one PE per head) pays none, but absolute IPCQ volume
-stays under \SI{120}{\kibi\byte} per run, two orders of magnitude below the
-HBM traffic. The result is a genuine trade-off rather than a single winner:
-\textsf{1-kv-per-cube} minimizes per-request latency by spending the most
-hardware (eight CUBEs, 64 PEs) on one request, while denser mappings leave
-CUBEs free for other requests---the batched-serving axis taken up below.
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_short_context/per_cube_tradeoff_mode_compare.png}
-\caption{Per-CUBE trade-off across the four mappings (baseline kernel;
-log scale). Top row: per-CUBE KV footprint---\textsf{8-kv-per-cube} packs
-all eight heads onto one CUBE, an $8\times$ larger per-CUBE KV than
-\textsf{1-kv-per-cube} and one CUBE used against eight (the $8\times$ ratio
-is context-invariant; the absolute footprint grows with $S_{kv}$). Bottom
-row: IPCQ traffic per run---\textsf{1-kv-per-cube} pays the eight-PE
-chain-reduce while \textsf{8-kv-per-cube} (single PE per head) pays zero.
-The density (top) is what sets how many concurrent requests a SIP can hold.}
-\label{fig:gqa-short-tradeoff}
-\end{figure}
-
-The composite ablation in Figure~\ref{fig:gqa-short-a1-variant} shows the
-three tiers within \textsf{1-kv-per-cube}: wall-clock is nearly flat, with
-the GEMM-only composite tier a few percent \emph{slower}
-(\SI{14.0}{} vs.\ \SI{11.8}{\micro\second} at $64$K decode) and the
-\textsf{softmax\_merge} tier matching the baseline. This is the expected
-outcome for the decode-skinny shape---$M{=}G{=}8$ is well below the
-scheduler's $\textsf{TILE\_M}{=}32$ supertile (\S\ref{sec:gemm}), so the
-composite path pads $M$ by $4\times$ with zeros and the fusion has no slack
-to win back; the padding is pure overhead here. It shows up in
-Figure~\ref{fig:gqa-short-gemm-util} as inflated GEMM-engine busy time on
-the composite tiers---the padded-$M$ GEMM is charged in full against a
-sub-\SI{15}{\micro\second} decode wall, driving the accounted utilization
-well above the baseline's \SI{12}{}--\SI{18}{\percent}. Fusion benefit
-requires lifting $M$, either by the long-context Q-tile split
-(\S\ref{sec:gqa-composite}) or by batching multiple users---which is also
-the lever that converts a mapping's density into throughput.
-
-\paragraph{Latency versus batched throughput.} The latency ladder and the
-density trade-off pull in opposite directions once a SIP serves more than
-one request. Decode users are independent---each attends its own KV
-cache---and these mappings generate no cross-CUBE traffic, so a SIP runs
-several users at once on disjoint CUBE groups, up to $16/C$ users: two for
-\textsf{1-kv-per-cube}, sixteen for \textsf{8-kv-per-cube}. We measure this
-directly, launching $B$ concurrent users and reading aggregate latency off
-the shared clock (Figure~\ref{fig:gqa-batch}). The overlap is nearly
-perfect: aggregate latency stays within \SI{3}{}--\SI{5}{\percent} of the
-single-user latency all the way to a full SIP (\textsf{8-kv-per-cube} at
-$16$ users is \SI{11.4}{} vs.\ \SI{10.8}{\micro\second}), so that small
-residual is the only cross-user interference the shared fabric adds.
-Aggregate throughput at $8$K therefore rises almost linearly with $B$, from
-\SI{0.86}{}~requests\,/\,\si{\micro\second} (\textsf{1-kv-per-cube}, capped
-at two users) to \SI{1.41}{} (\textsf{8-kv-per-cube}, sixteen users): the
-\emph{dense} mapping wins on throughput even though it is the \emph{slowest}
-per request.
-
-The reason is worth spelling out, because both mappings fill the same
-hardware: at capacity \textsf{1-kv-per-cube} runs two users
-$\times$ 64 PEs and \textsf{8-kv-per-cube} sixteen users $\times$ 8 PEs, so
-each puts all 128 PEs of the SIP to work. What differs is per-PE
-\emph{efficiency}. Splitting one head across eight PEs
-(\textsf{1-kv-per-cube}) leaves each PE only $S_{kv}/8$ tokens---a single
-tile---so its fixed overheads (pipeline fill on the lazy KV load, and the
-eight-PE online-softmax chain-reduce) dominate the little streaming work it
-does, and it sustains just \SI{46}{\percent} of the per-PE HBM ceiling.
-Giving each head a whole PE (\textsf{8-kv-per-cube}) streams the full
-$S_{kv}$-token cache contiguously with no reduce, amortizing those overheads
-over $8\times$ more work and reaching \SI{76}{\percent}. The measured
-throughput ratio ($1.41/0.86 = 1.6\times$) matches the utilization ratio
-($0.76/0.46 = 1.7\times$) almost exactly: because decode is
-bandwidth-bound, aggregate throughput is set by total HBM efficiency, not
-by PE count. \textsf{1-kv-per-cube} is really spending hardware on
-\emph{latency}---eight PEs per head buy only a $4.8\times$ speedup, a $60\%$
-strong-scaling efficiency---and that same $40\%$ overhead is what caps its
-throughput once the SIP is full. The design conclusion is regime-dependent: \textsf{1-kv-per-cube} is the
-latency-optimal choice for single-stream or low-batch decode, while denser
-mappings are the throughput-optimal choice for high-batch short-context
-serving. At long context the per-request latencies already scale as $1/C$,
-so the same accounting predicts the throughputs converge; a measured
-long-context batch sweep is 2H work (\S\ref{sec:future}).
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_short_context/batch_scaling.png}
-\caption{Batch scaling at $S_{kv}{=}8$K: aggregate throughput (requests per
-\si{\micro\second}) versus the number of concurrent decode users on one SIP,
-each user on a disjoint CUBE group. Each mapping scales almost linearly with
-$B$ up to its SIP capacity $16/C$---\textsf{1-kv-per-cube} saturates at two
-users, \textsf{8-kv-per-cube} at sixteen. Solid segments are measured
-concurrent runs; the dashed extensions are the saturated-throughput ceiling
-beyond capacity, where further users run in waves at that sustained rate
-(so the SIP is full, not idle). The dense mapping reaches the highest
-ceiling; the spread mapping is latency-optimal but capacity-limited.
-Concurrent users overlap to within \SI{3}{}--\SI{5}{\percent} of the
-single-user latency.}
-\label{fig:gqa-batch}
-\end{figure}
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_short_context/wall_a1_variant_compare.png}
-\caption{Composite-tier ablation on \textsf{1-kv-per-cube} (log scale):
-wall-clock is nearly flat across tiers---the GEMM-only composite runs a few
-percent slower (padding overhead) and \textsf{softmax\_merge} matches the
-baseline. With $M{=}G{=}8 < \textsf{TILE\_M}{=}32$ the composite scheduler
-pads $M$ by $4\times$, leaving no fusion slack to recover at this shape.
-Fusion pays off only when $M$ fills the supertile---long-context Q-tile
-splitting or batched-$M$ inference.}
-\label{fig:gqa-short-a1-variant}
-\end{figure}
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_short_context/gemm_util_a1_variant_ablation.png}
-\caption{GEMM-engine busy fraction on \textsf{1-kv-per-cube} across the
-three composite tiers. The baseline runs at \SI{12}{}--\SI{18}{\percent};
-the composite tiers read much higher---up to and past
-\SI{100}{\percent} of the short decode wall---because the padded
-$8\!\to\!32$ $M$ dimension is charged as GEMM time, bookkeeping overhead
-rather than useful work. The engine is never the bottleneck at this shape;
-decode is KV-bandwidth-bound.}
-\label{fig:gqa-short-gemm-util}
-\end{figure}
-
-\subsection{Inference with Long-Context Length}
-\label{sec:gqa-long}
-
-% TODO: prefill long-context kernel implementation description
-% (Sequence-Parallel partition of S_kv, per-case mechanics).
-% Bench code: src/kernbench/benches/gqa_helpers/long_ctx/
-
-% TODO: prefill long-context performance figure.
-
-Long-context decode is the regime where the KV cache, not attention
-compute, sets serving cost, so the placement question of
-\S\ref{sec:gqa-placement} becomes decisive here. To pick the right
-placement for this regime we run each of the six options as a fused
-decode kernel on the simulator (one decode step on the LLaMA-3.1-70B
-single-KV-head-group target, $C{=}8$ CUBEs $\times$ $P{=}8$ PEs) at a
-tractable $S_{kv}{=}8192$, and read off end-to-end latency, on-device op
-traffic, and the redundant compute each one induces. The swept context is
-small enough that all six fit in memory at $S_{kv}{=}8192$; the
-placements the long-context memory budget rules out (Cases~1--3) are
-drawn in red, run here only to expose their issue and communication
-structure.
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_latency.png}
-\caption{Measured end-to-end decode latency per placement
-($S_{kv}{=}8192$; red = ruled out by the long-context memory budget,
-blue = the predicted Pareto choice). The fastest raw latency belongs to
-Case~3 (\SI{17.8}{\micro\second})---but Case~3 replicates the full KV
-cache into every CUBE, so it overflows the per-PE budget at production
-context and wastes 8$\times$ the compute
-(Figure~\ref{fig:gqa-6cases-par}). Among the placements that actually fit
-1\,M-token memory (Cases~4--6), Case~6~$\star$ is the fastest
-(\SI{30.6}{\micro\second}, versus \SI{31.4}{} and \SI{34.5}{\micro\second}
-for the $d_{\text{head}}$-split Cases~4 and~5)---making it the placement
-of choice for long-context decode.}
-\label{fig:gqa-6cases-lat}
-\end{figure}
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_parallelism.png}
-\caption{Redundant compute per placement, measured as active-PE
-$\times$ $S_{\text{local}}$ (PE-tokens; lower means less wasted work).
-The minimum is \num{8192} PE-tokens---one pass over the sequence.
-Case~3 inflates this 8$\times$ to \num{65536} by replicating the KV
-cache across all eight CUBEs so every CUBE redundantly re-attends the
-whole sequence; the $d_{\text{head}}$-split Cases~4--5 likewise carry
-\num{65536} because each token is processed across eight head slices.
-Case~6~$\star$ achieves the full 64-way split at the minimal
-\num{8192} PE-tokens---fully parallel, no replication.}
-\label{fig:gqa-6cases-par}
-\end{figure}
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_traffic.png}
-\caption{Measured on-device op traffic per placement. The unsharded
-Case~1 issues no IPCQ copies (each PE has the full cache, nothing to
-reduce); the single-axis Case~3 charges 168. Among the 64-way splits,
-the $d_{\text{head}}$-split Cases~4--5 charge the most---280 IPCQ copies
-and 8 DMA writes each, the partial-score all-reduce that head-slicing
-forces---while Case~6~$\star$ needs only 189 IPCQ copies and a single DMA
-write, because it merges just the running softmax state $(m,\ell,O)$
-rather than partial scores. This is the on-device collective traffic that
-PE\_IPCQ and the torus links of \S\ref{sec:allreduce} are provisioned to
-absorb at link speed.}
-\label{fig:gqa-6cases-traffic}
-\end{figure}
-
-Taken together the three panels select the placement. The only
-memory-feasible family at production context is the 64-way splits
-(Cases~4--6), and within it Case~6~$\star$ is both the fastest and the
-lightest-communicating, because it merges only the running softmax state
-rather than the partial scores that the $d_{\text{head}}$-split
-Cases~4--5 must all-reduce. Case~3's lower raw latency is beside the
-point---it replicates the full cache into every CUBE, overflowing the
-per-PE budget and wasting $8\times$ the compute. For long-context decode
-the placement of choice is therefore the both-axes sequence shard, and the
-cross-PE softmax reduction it does pay is precisely the traffic the
-communication-side codesign of this report is built to move quickly.
-
-\subsection{Use of Composite Commands}
-\label{sec:gqa-composite}
-
-The decode kernel of \S\ref{sec:gqa-long} issues its local attention as
-primitive operations: it walks each PE's $S_{\text{local}}$ token slice in
-\SI{1024}{}-token tiles, and for every tile issues a $Q\!\cdot\!K^{\top}$
-\textsf{dot}, the online-softmax primitives, and a $P\!\cdot\!V$
-\textsf{dot}, merging the running $(m,\ell,O)$ state by hand. The number
-of PE\_CPU commands this costs grows with the context. At a production
-context of $S_{kv}{=}1\,\text{M}$ tokens the Case-6 64-way split gives
-each PE $S_{\text{local}}{=}16384$ tokens, so its local attention is a
-$Q\!\cdot\!K^{\top}$ of $(8,128)\!\cdot\!(128,16384)$ and a
-$P\!\cdot\!V$ of $(8,16384)\!\cdot\!(16384,128)$---sixteen hand-issued
-tiles, each a fresh batch of CPU commands.
-
-The composite command lets the kernel hand that tiling to PE\_SCHEDULER.
-We compare three command forms of the \emph{same} Case-6 kernel---identical
-placement and identical $(m,\ell,O)$ reduce, differing only in how the
-local attention is issued:
-\begin{itemize}
-  \item \textbf{primitive}---the hand-tiled \textsf{dot}/softmax kernel
-    above (the baseline of \S\ref{sec:gqa-long}).
-  \item \textbf{composite}---each matrix product is one coarse
-    \textsf{composite} GEMM over the \emph{whole} $S_{\text{local}}$, with
-    $K$ and $V$ passed as HBM references so PE\_SCHEDULER streams and
-    tiles them on the fixed $32\!\times\!64\!\times\!32$ MAC tile; the
-    softmax stays primitive.
-  \item \textbf{composite\,+\,softmax\_merge}---additionally folds the
-    online-softmax merge and $P\!\cdot\!V$ into a single stateful
-    \textsf{softmax\_merge} recipe composite.
-\end{itemize}
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_decode_long_ctx_composite.png}
-\caption{Three command forms of the Case-6 decode kernel, swept over
-context length ($S_{\text{local}}{=}S_{kv}/64$ per PE). \emph{Right:}
-PE\_CPU commands issued. The hand-tiled primitive kernel rises
-$O(n_{\text{tiles}})$---from 96 commands at one tile to 426 at the
-1\,M-token, sixteen-tile production point---while both composite forms
-issue a context-\emph{independent} $O(1)$ count (94 and 98) that
-\emph{saturates}: one coarse descriptor offloads the entire per-tile
-fan-out. \emph{Left:} the consequence for wall-clock latency is none---all
-three land on the same curve (\SI{30.6}{}, \SI{231}{},
-\SI{461}{\micro\second} at 8\,K\,/\,64\,K\,/\,128\,K), because decode is
-bound by streaming the KV cache, not by issue. Command-count is measured
-at emit time (exact, to 1\,M); latency on the data-mode engine over the
-tractable range.}
-\label{fig:gqa-composite}
-\end{figure}
-
-Figure~\ref{fig:gqa-composite} reads off the two quantities that matter,
-and they point in opposite directions. The PE\_CPU command count (right)
-collapses from a context-growing $O(n_{\text{tiles}})$ to a flat $O(1)$:
-at 1\,M tokens the composite form issues \num{94} commands against the
-primitive kernel's \num{426}, a $4.5\times$ reduction
-($4.2\times$ in modeled dispatch cost), and---crucially---that number no
-longer grows with context. The wall-clock latency (left), by contrast,
-is unchanged across all three forms: decode is bound by streaming the KV
-cache out of HBM, so the command form does not move the critical path.
-
-That juxtaposition is the point. The composite command is not a latency
-optimization for this memory-bound decode; it is a \emph{CPU-issue}
-optimization. Its value is removing the per-tile dispatch work that would
-otherwise grow without bound as context grows, freeing PE\_CPU to run
-ahead and keep the engines fed---which is exactly what lets the
-data-movement cost analyzed next show through as the true bottleneck
-rather than being masked by issue overhead. The \textsf{softmax\_merge}
-recipe folds the online merge into the same descriptor; on this
-memory-bound path its marginal cost over the plain GEMM composite is small
-(98 vs.\ 94 commands), and like the plain composite it keeps the issued
-count flat as context scales.
-
-\paragraph{The compute-bound mirror: prefill.} Decode's verdict---command
-form is latency-neutral---is a property of its regime, not of the
-composite command. A decode step has $T_q{=}1$, so its score and context
-products are skinny ($M{=}G\,T_q{=}8$): the MAC array is barely fed and
-the kernel is bound by streaming the KV cache. Prefill is the opposite
-corner. It processes a block of query positions at once, so $M{=}G\,T_q$
-is large and tile-filling, the GEMMs carry real arithmetic intensity
-($\sim$$M$ flops/byte, well above the roofline ridge), and the kernel is
-\emph{compute-bound}. This is the regime the composite command was built
-for (\S\ref{sec:gemm}): it streams the per-HW-tile
-DMA$\rightleftarrows$compute pipeline so the MAC array stays fed, whereas
-the primitive kernel's blocking \textsf{tl.dot} serializes each tile's
-load and compute and starves the array between tiles. We run the same
-three command forms on a single-rank compute-bound prefill (FlashAttention
-$Q$-block $\times$ $S_{kv}$-tile, online softmax) and sweep the context
-length (Figure~\ref{fig:gqa-prefill-cb}).
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_prefill_compute_bound.png}
-\caption{Compute-bound prefill, three command forms, swept over context
-length ($M{=}8\,T_q$ tile-filling). \emph{Left:} end-to-end latency.
-\emph{Right:} MAC utilization (achieved $\div$ the
-\SI{8}{\tera\flop\per\second} per-PE peak). The hand-tiled primitive sits
-flat at $\sim$\SI{68}{\percent}---its serial load$\to$dot path leaves the
-MAC array idle between tiles regardless of context. The composite forms
-climb with context (\SI{67}{}$\to$\SI{80}{\percent} plain,
-\SI{67}{}$\to$\SI{83}{\percent} with the recipe) because a deeper $P\!\cdot
-\!V$ reduction gives more HW tiles to pipeline, and they convert that into
-wall-clock: at \num{1024} the recipe form is \SI{646.9}{} vs.\
-\SI{794.1}{\micro\second} (\SI{19}{\percent} faster). The margin
-\emph{grows} with context---the compute-bound mirror of the GEMM result of
-\S\ref{sec:gemm}.}
-\label{fig:gqa-prefill-cb}
-\end{figure}
-
-The two studies together state the composite command's value precisely. It
-has two distinct benefits, and which one matters is set by the workload's
-roofline position. The first is \emph{host-issue offload}: one macro
-command in place of $O(n_{\text{tiles}})$ fine ones, which removes
-PE\_CPU dispatch work and is regime-independent (it shows in the decode
-command count). The second is \emph{MAC-array feeding}: the
-scheduler-internal per-tile DMA$\rightleftarrows$compute pipeline, which
-only converts to latency when the workload is compute-bound enough to have
-a MAC array worth keeping busy (it shows in the prefill utilization). A
-memory-bound decode exercises only the first; a compute-bound prefill
-exercises both. The composite command is the single mechanism that
-delivers each where it applies.
-
-\subsection{Summary}
-\label{sec:gqa-analysis}
-
-The section's results line up into one picture of the fused kernel.
-Placement is decided by memory first and communication second: only the
-64-way splits fit production context, and among them the both-axes
-sequence shard (Case~6) minimizes the per-token collective by merging
-softmax state rather than partial scores. Command form is decided by
-roofline position: the composite command removes the per-tile issue work
-in both regimes, but converts to wall-clock only in compute-bound prefill,
-while memory-bound decode stays bound by KV streaming regardless of command
-form. The thread connecting the two is that, once composite issue makes
-GEMM cheap and leaves the MAC array idle, the fused kernel's latency is set
-almost entirely by data movement---streaming the KV cache and reducing
-partials across PEs---which is exactly the cost the communication-side
-work (on-device reduction, lazy load/compute overlap, fast TCM staging and
-torus links) is built to attack. What this implies for hardware investment
-across the report as a whole is taken up in the discussion.
-
-% TODO: cross-regime DP (data parallelism) applicability:
-%   - Does Case-4 long-context placement compose with batch-level DP
-%     without further changes?
-%   - Does the short-context placement compose the same way?
-%   - Implications for multi-user serving (single vs. mixed regimes).

docs/report/1H-codesign-paper/sections/05a-roofline.tex

@@ -0,0 +1,85 @@
+\subsection{Roofline Analysis: Batch and Context Regimes}
+\label{sec:roofline}
+
+Before deploying any sharding strategy, the workload's position on the
+decode roofline sets what \emph{can} be improved and what \emph{cannot}.
+Two quantities matter: the machine's \textbf{critical batch}
+$B^\ast = C \cdot b / (2 W)$ (where $C$ is per-PE FLOPs, $W$ is per-PE
+HBM bandwidth, $b$ is bytes per parameter), and its \textbf{balance
+context} $L^\ast = 2 N / (\mathrm{AI} \cdot \mathrm{kv\_bpt})$ (where
+$N$ is active parameters and $\mathrm{AI} = C/W$ is arithmetic
+intensity). Above $B^\ast$ the deployment leaves the memory-bound
+regime; above $L^\ast$ per-user KV streaming dominates and no amount
+of batching hides it. For Llama-3.1-70B on the default AHBM machine
+($C \!=\!  \SI{8}{TFLOPS}$/PE, $W\!=\!  \SI{256}{GB/s}$/PE), we get
+$B^\ast \!\approx\! 31$ and $L^\ast \!\approx\! \SI{13.4}{K}$ tokens.
+
+\paragraph{Short-context regime ($S_{kv} \!<\! L^\ast$).}
+Figure~\ref{fig:roofline-short} shows one decode step at
+$S_{kv}\!=\!  \SI{8}{K}$. The step-latency panel (left) makes it
+obvious that weight fetch dominates at low $B$: at $B\!=\!1$, one
+step spends $\sim\!\SI{535}{ms}$ streaming weights and only
+$\SI{28}{ms}$ on per-sequence compute and KV combined. The
+cost-per-token panel (right) is where the batch story lives: dividing
+step latency by $B$ shrinks the weight term as $1/B$, so per-token
+cost drops from $\SI{562}{ms}$ at $B\!=\!1$ to $\SI{29.7}{ms}$ at
+$B\!=\!256$ --- a $19\times$ reduction. This is the memory-bound-to-
+compute-bound crossover: past $B^\ast$, weight cost is fully
+amortized and per-token time asymptotes to the compute + KV floor
+($\sim\!\SI{27}{ms}$). \textbf{At short context, batching directly
+lowers cost per token.}
+
+\begin{figure*}[h]
+\centering
+\includegraphics[width=\textwidth]{roofline_short_context.png}
+\caption{Roofline decomposition for Llama-3.1-70B decode at
+$S_{kv}\!=\!  \SI{8}{K}$ (short, well below $L^\ast\!\approx\!  \SI{13.4}{K}$).
+\textbf{Left}: step latency vs.\ batch $B$; weight fetch is flat
+(one HBM sweep per step), compute and KV grow linearly with $B$.
+\textbf{Right}: per-token cost ($=$ step $\div B$); the $1/B$
+weight-fetch term amortizes rapidly, dropping total per-token time
+from \SI{562}{ms} at $B\!=\!1$ to \SI{29.7}{ms} at $B\!=\!256$
+--- a $19\times$ throughput gain from batching alone.}
+\label{fig:roofline-short}
+\end{figure*}
+
+\paragraph{Long-context regime ($S_{kv} \!\gg\! L^\ast$).}
+Figure~\ref{fig:roofline-long} shows the same decomposition at
+$S_{kv}\!=\!  \SI{1}{M}$ ($\sim\!78 \times L^\ast$). The step-latency
+panel shows KV fetch has swelled by two orders of magnitude:
+per-token KV cost is now $\sim\!\SI{1342}{ms}$, dwarfing both
+weight ($\SI{535}{ms}$ at $B\!=\!1$) and compute ($\SI{17}{ms}$).
+The cost-per-token panel is the punchline: increasing $B$ still
+shrinks the weight term but leaves the giant KV floor untouched,
+because \textbf{KV fetch is per-sequence} --- adding another user
+adds a full extra copy of their KV read. Total per-token cost falls
+only from \SI{1894}{ms} at $B\!=\!1$ to \SI{1361}{ms} at $B\!=\!256$
+--- a mere $28\%$ reduction despite $256\times$ the batch.
+\textbf{At long context, batching stops paying because per-user KV
+streaming dominates.}
+
+\begin{figure*}[h]
+\centering
+\includegraphics[width=\textwidth]{roofline_long_context.png}
+\caption{Same decomposition at $S_{kv}\!=\!  \SI{1}{M}$ (long,
+$\sim\!78\times L^\ast$). KV fetch has grown to \SI{1342}{ms}/token
+and is now the dominant term at every $B$. Batching only shrinks
+the weight component; the KV floor is unmovable because each new
+user brings a full per-sequence KV read. Total per-token cost falls
+just $28\%$ from $B\!=\!1$ to $B\!=\!256$, versus $19\times$ at
+short context.}
+\label{fig:roofline-long}
+\end{figure*}
+
+\paragraph{Implication for deployment.} The two regimes call for
+opposite strategies. In the short-context regime, the operator packs
+$B$ as high as HBM allows to sit on the compute floor --- this is
+where cost-per-token is minimized and hardware utilization is
+highest. In the long-context regime, per-user KV is the binding
+resource; batching offers little benefit, so the operator instead
+shrinks $\mathrm{kv\_bpt}$ (GQA / MQA / MLA, INT4 KV, sparse
+attention) and shards the sequence dimension itself (CP), routing
+long-context requests to a dedicated pool with more chips per user.
+The next two subsections (\S\ref{sec:capacity-planning},
+\S\ref{sec:parallelism-selection}) turn these regime observations
+into concrete sizing and sharding rules.

docs/report/1H-codesign-paper/sections/05b-capacity-planning.tex

@@ -0,0 +1,118 @@
+\subsection{Capacity Planning: LLM Serving on AHBM}
+\label{sec:capacity-planning}
+
+Once a fused kernel is chosen, the deployment question is orthogonal:
+\emph{how many cubes (or SIPs) does a given model, context length, and
+latency SLO require?} Hyperscalers decide this along three axes that
+must all be satisfied simultaneously. The total PE count is
+$\max(A, B, C) \times N_{\text{replicas}}$, where $A$ is the capacity
+floor, $B$ the KV-cache headroom, and $C$ the throughput SLO
+(Table~\ref{tab:capacity-axes}).
+
+\begin{table}[h]
+\centering
+\small
+\caption{Three-axis sizing. $N$ = active params, $b$ = bytes per param,
+$\mathrm{HBM}_{\mathrm{PE}}$ = per-PE HBM budget, $S_{kv}$ = context
+length, $\mathrm{kv\_bpt}$ = KV cache bytes per token per user
+($=2 \cdot h_{kv} \cdot d_{\text{head}} \cdot b \cdot L$).}
+\label{tab:capacity-axes}
+\begin{tabular}{@{}l l@{}}
+\toprule
+Axis & Formula \\
+\midrule
+A. Capacity floor & $\lceil N b / \mathrm{HBM}_{\mathrm{PE}} \rceil$ \\
+B. KV headroom & $\lceil (N b + u \cdot S_{kv} \cdot \mathrm{kv\_bpt}) / \mathrm{HBM}_{\mathrm{PE}} \rceil$ \\
+C. Throughput SLO & $N_{\text{replicas}} = \lceil n_{\text{users}} / B_{\mathrm{SLO}} \rceil$ \\
+\bottomrule
+\end{tabular}
+\end{table}
+
+\paragraph{SLO targets.} The per-token latency budget $B_{\mathrm{SLO}}$
+follows the use case. TTFT (Time To First Token) is dominated by
+prefill; TPOT (Time Per Output Token) by one decode step
+(Table~\ref{tab:slo-targets}).
+
+\begin{table}[h]
+\centering
+\small
+\caption{Typical SLO targets. Voice needs a tight per-token cadence;
+chat balances both; batch is priced on throughput not latency.}
+\label{tab:slo-targets}
+\begin{tabular}{@{}l l l@{}}
+\toprule
+Workload & TTFT & TPOT \\
+\midrule
+Voice / real-time & 200--300 ms & 10--25 ms \\
+Interactive chat & 300 ms -- 1 s & 20--50 ms \\
+Batch / offline & seconds -- minutes & N/A \\
+\bottomrule
+\end{tabular}
+\end{table}
+
+\paragraph{Regime-aware rules of thumb.} Context length relative to the
+machine's balance point $L^*=2N/(\mathrm{AI}\cdot\mathrm{kv\_bpt})$
+determines which strategy applies. Below $L^*$ the deployment is
+compute-bound and batch scales to $\sim\!2B^*$; above, KV streaming
+dominates and $B$ must shrink (Table~\ref{tab:regime-rules}).
+
+\begin{table*}[t]
+\centering
+\small
+\caption{Deployment strategy by context regime.}
+\label{tab:regime-rules}
+\begin{tabular}{@{}l l l@{}}
+\toprule
+Regime & Batch & Cost/token \\
+\midrule
+Short ($S_{kv} < L^*$) & $B \approx 2 B^*$ & low \\
+Long ($S_{kv} > L^*$) & smaller $B$, more CP & higher \\
+Extreme ($S_{kv} \gg L^*$) & dedicated pool, disagg.\ prefill & much higher \\
+\bottomrule
+\end{tabular}
+\end{table*}
+
+\paragraph{Sample deployment templates.} A single base model is
+typically routed to one of several sharding tiers depending on the
+request's expected context length; the API gateway dispatches to the
+appropriate replica pool (Table~\ref{tab:deploy-templates}).
+
+\begin{table*}[t]
+\centering
+\small
+\caption{Sample deployment tiers for one base model. TP is fixed to
+match $h_{kv}$; CP grows with context; $B$ shrinks to hold TPOT.}
+\label{tab:deploy-templates}
+\begin{tabular}{@{}l c c c c@{}}
+\toprule
+Tier & CP$\times$TP & Cubes/repl. & Max $S_{kv}$ & Typical $B$ \\
+\midrule
+Small & $1 \times 8$ & 1 & 32k & 64 \\
+Medium & $4 \times 8$ & 4 & 128k & 32 \\
+Large & $32 \times 8$ & 32 & 1M & 4 \\
+\bottomrule
+\end{tabular}
+\end{table*}
+
+\paragraph{What to do when each axis binds.} The dominant axis dictates
+the first lever to pull (Table~\ref{tab:axis-playbook}). Axes $A$ and
+$C$ scale nearly linearly with hardware; axis $B$ is the one where
+algorithmic techniques (GQA, MLA, INT4 KV, sparse attention) buy the
+most because they attack $\mathrm{kv\_bpt}$ directly rather than adding
+chips.
+
+\begin{table*}[t]
+\centering
+\small
+\caption{Playbook per binding axis.}
+\label{tab:axis-playbook}
+\begin{tabular}{@{}l l l@{}}
+\toprule
+Binding & First lever & Second lever \\
+\midrule
+A. Capacity & $\uparrow$ TP / PP & bigger-HBM chip; FP8 / INT4 weights \\
+B. KV & $\uparrow$ CP; $\downarrow B$ & GQA / MQA / MLA; INT4 KV; sparse attn \\
+C. Throughput & $\uparrow$ replicas (DP) & loosen SLO; smaller model; spec.\ decoding \\
+\bottomrule
+\end{tabular}
+\end{table*}

docs/report/1H-codesign-paper/sections/05c-parallelism-selection.tex

@@ -0,0 +1,91 @@
+\subsection{Parallelism Selection: TP $\times$ CP $\times$ PP $\times$ DP $\times$ EP}
+\label{sec:parallelism-selection}
+
+Given a capacity-feasible layout (\S\ref{sec:capacity-planning}), the
+remaining decision is \emph{how to shard}: which axes to enable and at
+what degrees. Memory is the feasibility filter; once you fit, the
+design problem is minimising \emph{exposed} communication (communication
+that could not be overlapped with compute). The core principle is to
+\textbf{rank techniques by how much they communicate per unit of
+compute, and assign the chattiest ones to the fastest links.}
+
+\begin{table*}[h]
+\centering
+\small
+\caption{Parallelism axes at a glance. TP AllReduce volume does not
+shrink with degree — only compute does — which sets its practical
+ceiling near a fast-link (NVLink / intra-SIP) domain.}
+\label{tab:parallelism-compare}
+\begin{tabular}{@{}l l l l l@{}}
+\toprule
+Axis & Shards & Comm pattern & Frequency & Hard ceiling \\
+\midrule
+DP & optimizer state (w/ ZeRO) & AllReduce (grads) & 1$\times$/step & critical batch \\
+TP & weights + acts + KV heads & AllReduce & 2$\times$/layer & fast-link domain; $\lesssim h_{kv}$ \\
+PP & weights + KV by layer & P2P activation hand-off & per stage boundary & $n_{\text{layers}}$ \\
+CP & activations + KV by position & Ring P2P per hop & per layer (overlappable) & $\mathrm{seq\_len} / \mathrm{block}$ \\
+EP & expert weights (MoE only) & All-to-all & 2$\times$/MoE layer & $n_{\text{experts}}$ \\
+\bottomrule
+\end{tabular}
+\end{table*}
+
+\paragraph{Symptom-driven axis selection.} The dominant problem dictates
+the first knob to try (Table~\ref{tab:parallelism-problem}).
+
+\begin{table}[h]
+\centering
+\small
+\caption{Which axis to reach for first, by dominant problem.}
+\label{tab:parallelism-problem}
+\begin{tabular}{@{}l l@{}}
+\toprule
+Dominant problem & First axis \\
+\midrule
+Weight memory doesn't fit & TP within a fast domain \\
+KV of ONE long sequence & CP \\
+KV for MANY sequences & DP replicas \\
+Single-request TPOT & TP ($\lesssim 8$) \\
+Aggregate throughput & Outer DP \\
+Model spans multiple nodes & TP intra-node, PP inter-node \\
+MoE expert weight memory & EP \\
+\bottomrule
+\end{tabular}
+\end{table}
+
+\paragraph{When to add, when to stop.} Every axis has a natural
+saturation point past which further degree hurts more than it helps
+(Table~\ref{tab:parallelism-criteria}). Sizing decisions are almost
+always constrained by two axes simultaneously (e.g.\ TP by NVLink
+domain and $h_{kv}$; CP by ring-hop hiding and per-rank block size).
+
+\begin{table*}[t]
+\centering
+\small
+\caption{Add / stop criteria per axis.}
+\label{tab:parallelism-criteria}
+\begin{tabular}{@{}l l l@{}}
+\toprule
+Axis & Add when & Stop when \\
+\midrule
+DP & more independent requests & global batch $>$ critical \\
+TP & weights need sharding or TPOT tight & collectives dominate; local GEMMs shrink \\
+PP & depth too large; TP would cross slow links & bubble $>$ $\sim$10\% \\
+CP & one sequence needs position sharding & ring hops can't overlap compute \\
+EP & MoE expert memory & expert GEMMs too small; routing imbalance \\
+\bottomrule
+\end{tabular}
+\end{table*}
+
+\paragraph{Common misconceptions.} Three widely repeated rules survive
+contact with real workloads only partially:
+\emph{(i)} memory is not just a feasibility filter but a continuous
+performance variable --- more free HBM enables larger $B$, more prefix
+cache, higher throughput even after weights fit;
+\emph{(ii)} the ``TP $\leq h_{kv}$'' ceiling is an efficiency
+preference, not a correctness constraint --- vLLM and others support
+KV-head replication and head-dim splitting for TP $>$ $h_{kv}$;
+\emph{(iii)} ``always start at TP=8'' is disproven by public
+disclosures (DeepSeek-V3 inference uses TP=4, some MLPs at TP=1;
+DeepSeek-V3 training uses PP=16 + EP=64 with no TP), which shows the
+right starting point is the smallest TP that fits memory and meets
+latency, followed by measurement.

docs/report/1H-codesign-paper/sections/05x-fused-kernel.tex

@@ -0,0 +1,515 @@
+\subsection{Data Placement Policy}
+\label{sec:gqa-placement}
+
+Long-context decode is bound by the KV cache, so the first-order design
+question is how to place that cache---and the running softmax state it
+feeds---across the two hardware axes the machine exposes: the CUBEs and,
+within each CUBE, the PEs (here $C{=}8$ CUBEs $\times$ $P{=}8$ PEs, for
+$C\!\cdot\!P{=}64$ attention engines over one KV-head group on the
+LLaMA-3.1-70B target). Each axis can \emph{replicate} the KV cache or
+\emph{shard} it, and a shard can run along the sequence dimension
+$S_{kv}$ or the head dimension $d_{\text{head}}$. The cross product is a
+small, enumerable taxonomy; six placements span its meaningful corners
+(Figure~\ref{fig:gqa-kv-sharding}).
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_kv_sharding_diagram.png}
+\caption{The six KV-placement strategies, drawn on the
+$S_{kv}\!\times\!d_{\text{head}}$ KV tensor (rows = sequence, columns =
+head dimension). Cube colour bands and dashed PE dividers show which
+axis each level shards. Cases~1--3 either replicate the cache or shard
+it on a single axis (8-way at most); Cases~4--6 reach a full 64-way
+split, three different ways: Case~4 splits $S_{kv}$ across CUBEs and
+$d_{\text{head}}$ across PEs, Case~5 the mirror, and Case~6~$\star$
+splits $S_{kv}$ on \emph{both} axes.}
+\label{fig:gqa-kv-sharding}
+\end{figure}
+
+Two quantities decide which placement is viable, and they pull against
+each other. The first is \textbf{per-PE KV memory}. With a per-PE HBM
+budget of \SI{6.0}{\giga\byte} and \SI{1.76}{\giga\byte} of attention
+weights resident, the KV headroom is \SI{4.24}{\giga\byte} per PE. At a
+production context of $S_{kv}{=}1\,\text{M}$ tokens the unsharded cache
+is \SI{40}{\giga\byte}/PE (Case~1), an 8-way shard is \SI{5}{\giga\byte}
+(Cases~2--3)---both \emph{over} the headroom---while only the 64-way
+placements bring it to \SI{640}{\mega\byte}/PE (Cases~4--6), comfortably
+inside budget. Memory alone therefore eliminates Cases~1--3 at long
+context. The second quantity is \textbf{communication per token}, and it
+is what separates the three survivors. Sharding $d_{\text{head}}$
+(Cases~4--5) makes each PE hold only a slice of every head, so the
+$Q\!\cdot\!K^{\top}$ score is \emph{partial} and must be all-reduced
+across the slice owners on every token---a reduction whose volume scales
+with $S_{kv}$ ($\sim$\SI{166}{\mega\byte}/token analytically, intra-CUBE
+on the NoC for Case~4, inter-CUBE on UCIe for Case~5). Case~6~$\star$
+instead shards $S_{kv}$ on both axes, so every PE computes a
+\emph{complete} score over its own token range and only the small running
+softmax state $(m,\ell,O)$ is merged across PEs
+($\sim$\SI{6.2}{\mega\byte}/token)---a $\sim$27$\times$ lighter collective
+than the $d_{\text{head}}$-split designs.
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_summary.png}
+\caption{Long-context placement analysis at $S_{kv}{=}1\,\text{M}$ tokens.
+\emph{Left:} the per-PE HBM budget---\SI{1.76}{\giga\byte} of attention
+weights leave \SI{4.24}{\giga\byte} of KV headroom (red line).
+\emph{Middle:} per-PE KV memory per case (log scale); only the 64-way
+placements (Cases~4--6, \SI{640}{\mega\byte}) clear the headroom, while
+the unsharded (\SI{40}{\giga\byte}) and 8-way (\SI{5}{\giga\byte}) cases
+overflow. \emph{Right:} analytical communication per token (log scale);
+the $d_{\text{head}}$-split Cases~4--5 pay a partial-score all-reduce
+($\sim$\SI{166}{\mega\byte}/token) that the both-axes-$S_{kv}$ split of
+Case~6~$\star$ avoids ($\sim$\SI{6.2}{\mega\byte}/token, merging only the
+softmax state). On these two axes Case~6 (marked $\star$ in the figure)
+is the single placement that lands both inside the memory budget and at
+low per-token communication; whether that combination is the right one to
+pick is a regime-specific question taken up next.}
+\label{fig:gqa-budget}
+\end{figure}
+
+These two costs---per-PE KV memory and per-token communication---are
+intrinsic properties of each placement, fixed by how it shards the cache
+and independent of the workload regime. They do not by themselves name a
+winner: a short prompt where the whole cache fits on one PE values low
+communication and tolerates replication, whereas a million-token decode
+is bound by the memory wall and will pay communication to escape it.
+Which placement is appropriate is therefore a per-regime question, which
+the short- and long-context subsections that follow answer by running the
+options on the simulator and reading off latency, traffic, and the
+redundant compute each one induces.
+
+\subsection{Inference with Short-Context Length}
+\label{sec:gqa-short}
+
+The fused GQA kernel issues its matrix products as scheduler-managed
+composite commands and keeps the online-softmax merge and the cross-device
+KV reduction inside the kernel, on PE\_IPCQ. Two kernel families cover the
+two phases. The \emph{prefill} kernel splits the query tile across the PEs
+of a group and broadcasts each KV tile from the group's root PE over
+intra-CUBE IPCQ, so the HBM K/V read is paid once per group instead of
+once per PE. The \emph{decode} kernel sequence-shards the KV cache across
+the same group of PEs, runs per-PE local attention, then chain-reduces the
+partial $(m,\ell,O)$ triples back up to the group root. Two further
+primitives make long context practical: a \emph{lazy load} that issues the
+KV \textsf{DMA\_READ} and returns immediately, auto-waiting only at first
+use so KV load overlaps score computation; and per-tile \emph{scratch
+recycling} that keeps the running accumulators in a persistent arena while
+freeing per-tile temporaries, so the kernel fits the
+\SI{1}{\mebi\byte} scratch budget across many tiles. A further refinement
+restructures the decode step so the per-tile matrix products and the
+online-softmax merge are issued as \emph{composite} commands rather than
+hand-tiled primitives; \S\ref{sec:gqa-composite} measures it against the
+primitive baseline at long context.
+
+This is a different decomposition from the six-case placement taxonomy of
+\S\ref{sec:gqa-placement}: that taxonomy spreads a single KV-head group
+across all 64 PEs and is the right lens for long context
+(\S\ref{sec:gqa-long}), whereas short context---where the whole cache fits
+comfortably---turns instead on how the eight KV heads are distributed
+across the eight CUBEs. The design question for short context is therefore
+not whether to fuse softmax---that is settled---but how to distribute the
+eight KV heads across the eight CUBEs of a SIP. Four mappings cover the spectrum, named
+by KV-head count per CUBE: \textsf{1-kv-per-cube} dedicates one whole
+CUBE to each head and uses all eight PEs of that CUBE on the head's
+sequence shard; \textsf{2-kv-per-cube} packs two heads per CUBE with
+group size four; \textsf{4-kv-per-cube} packs four heads with group size
+two; \textsf{8-kv-per-cube} loads all heads onto a single CUBE with one
+PE per head. The mapping fixes a trade-off: as KV heads per CUBE grows, the per-CUBE KV
+footprint grows linearly and the mapping uses proportionally fewer CUBEs
+(eight down to one), but the intra-CUBE IPCQ broadcast/reduce cost shrinks
+to zero (8-kv-per-cube uses a single PE per head and has no group
+communication).
+
+A second axis is how each GEMM tile is issued. We isolate this with three
+\emph{composite tiers} applied to the same mapping: \textsf{without
+composite} uses primitive \textsf{tl.dot} and unfused softmax;
+\textsf{with composite (GEMM-only)} issues the GEMMs as
+\textsf{tl.composite(op="gemm")} commands so the scheduler can pipeline
+the tile stages but leaves softmax as primitives; \textsf{with composite
++ softmax\_merge} additionally folds the softmax fold into the $P\!\cdot\!V$
+composite through a named \textsf{softmax\_merge} prologue recipe.
+
+We sweep all four mappings $\times$ three composite tiers $\times$
+$S_{kv}\in\{8\text{K},16\text{K},32\text{K},64\text{K}\}$ for both
+phases (prefill uses $T_q{=}8$ sliced tiles; decode uses $T_q{=}1$). The
+headline latency comparison is in Figure~\ref{fig:gqa-short-wall-baseline}.
+The mappings separate cleanly, and in proportion to how many PEs each
+dedicates to a head: \textsf{1-kv-per-cube} spreads one head across all
+eight PEs of its CUBE, \textsf{8-kv-per-cube} gives each head a single PE,
+and the four mappings form a $\{64,32,16,8\}$-active-PE ladder. Decode
+latency tracks that ladder inversely---at $S_{kv}{=}64$K the
+\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube} ratio is $7.4\times$
+($\SI{87.0}{}$ vs.\ $\SI{11.8}{\micro\second}$), and even at $8$K it is
+$4.8\times$ ($\SI{10.8}{}$ vs.\ $\SI{2.2}{\micro\second}$); prefill shows
+the same ordering more gently ($2.3\times$ at $64$K). The reason is that
+decode here is \emph{bandwidth-bound}: each active PE streams its KV shard
+at \SI{46}{}--\SI{76}{\percent} of the \SI{256}{\giga\byte\per\second}
+per-PE ceiling (rising with context), so a mapping's speed is set by how
+many PEs it puts to work---more PEs, more aggregate HBM bandwidth, lower
+latency. The $M{=}G{=}8$ score GEMM is skinny and leaves the MAC array
+lightly loaded throughout; the lever here is data movement, not compute.
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_short_context/wall_variant1_mode_compare.png}
+\caption{Wall-clock latency of the four short-context mappings, baseline
+kernel (no composite; log scale). The mappings separate along the
+$\{64,32,16,8\}$-active-PE ladder set by heads-per-CUBE:
+\textsf{1-kv-per-cube} (64 PEs) is fastest and \textsf{8-kv-per-cube}
+(8 PEs) slowest, by $4.8\times$ at $8$K growing to $7.4\times$ at $64$K in
+decode ($2.3\times$ at $64$K in prefill). Decode is bandwidth-bound, so
+latency scales inversely with the number of active PEs. Composite tiers
+track the baseline at this skinny shape
+(Figure~\ref{fig:gqa-short-a1-variant}).}
+\label{fig:gqa-short-wall-baseline}
+\end{figure}
+
+Figure~\ref{fig:gqa-short-tradeoff} shows what the mappings trade for that
+latency. The differentiator is \emph{density}: \textsf{8-kv-per-cube}
+concentrates all eight heads onto one CUBE, so its per-CUBE KV footprint is
+$8\times$ that of \textsf{1-kv-per-cube} (one CUBE per head), and it
+occupies a single one of the SIP's sixteen CUBEs against
+\textsf{1-kv-per-cube}'s eight. Per-PE HBM \emph{utilization}, by contrast,
+is similar across mappings---every active PE runs near the same
+\SI{46}{}--\SI{76}{\percent} of the per-PE ceiling---so the latency ladder
+comes from the \emph{count} of active PEs, not from any per-PE
+concentration. IPCQ traffic runs opposite to density:
+\textsf{1-kv-per-cube} pays the eight-PE chain-reduce while
+\textsf{8-kv-per-cube} (one PE per head) pays none, but absolute IPCQ volume
+stays under \SI{120}{\kibi\byte} per run, two orders of magnitude below the
+HBM traffic. The result is a genuine trade-off rather than a single winner:
+\textsf{1-kv-per-cube} minimizes per-request latency by spending the most
+hardware (eight CUBEs, 64 PEs) on one request, while denser mappings leave
+CUBEs free for other requests---the batched-serving axis taken up below.
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_short_context/per_cube_tradeoff_mode_compare.png}
+\caption{Per-CUBE trade-off across the four mappings (baseline kernel;
+log scale). Top row: per-CUBE KV footprint---\textsf{8-kv-per-cube} packs
+all eight heads onto one CUBE, an $8\times$ larger per-CUBE KV than
+\textsf{1-kv-per-cube} and one CUBE used against eight (the $8\times$ ratio
+is context-invariant; the absolute footprint grows with $S_{kv}$). Bottom
+row: IPCQ traffic per run---\textsf{1-kv-per-cube} pays the eight-PE
+chain-reduce while \textsf{8-kv-per-cube} (single PE per head) pays zero.
+The density (top) is what sets how many concurrent requests a SIP can hold.}
+\label{fig:gqa-short-tradeoff}
+\end{figure}
+
+The composite ablation in Figure~\ref{fig:gqa-short-a1-variant} shows the
+three tiers within \textsf{1-kv-per-cube}: wall-clock is nearly flat, with
+the GEMM-only composite tier a few percent \emph{slower}
+(\SI{14.0}{} vs.\ \SI{11.8}{\micro\second} at $64$K decode) and the
+\textsf{softmax\_merge} tier matching the baseline. This is the expected
+outcome for the decode-skinny shape---$M{=}G{=}8$ is well below the
+scheduler's $\textsf{TILE\_M}{=}32$ supertile (\S\ref{sec:gemm}), so the
+composite path pads $M$ by $4\times$ with zeros and the fusion has no slack
+to win back; the padding is pure overhead here. It shows up in
+Figure~\ref{fig:gqa-short-gemm-util} as inflated GEMM-engine busy time on
+the composite tiers---the padded-$M$ GEMM is charged in full against a
+sub-\SI{15}{\micro\second} decode wall, driving the accounted utilization
+well above the baseline's \SI{12}{}--\SI{18}{\percent}. Fusion benefit
+requires lifting $M$, either by the long-context Q-tile split
+(\S\ref{sec:gqa-composite}) or by batching multiple users---which is also
+the lever that converts a mapping's density into throughput.
+
+\paragraph{Latency versus batched throughput.} The latency ladder and the
+density trade-off pull in opposite directions once a SIP serves more than
+one request. Decode users are independent---each attends its own KV
+cache---and these mappings generate no cross-CUBE traffic, so a SIP runs
+several users at once on disjoint CUBE groups, up to $16/C$ users: two for
+\textsf{1-kv-per-cube}, sixteen for \textsf{8-kv-per-cube}. We measure this
+directly, launching $B$ concurrent users and reading aggregate latency off
+the shared clock (Figure~\ref{fig:gqa-batch}). The overlap is nearly
+perfect: aggregate latency stays within \SI{3}{}--\SI{5}{\percent} of the
+single-user latency all the way to a full SIP (\textsf{8-kv-per-cube} at
+$16$ users is \SI{11.4}{} vs.\ \SI{10.8}{\micro\second}), so that small
+residual is the only cross-user interference the shared fabric adds.
+Aggregate throughput at $8$K therefore rises almost linearly with $B$, from
+\SI{0.86}{}~requests\,/\,\si{\micro\second} (\textsf{1-kv-per-cube}, capped
+at two users) to \SI{1.41}{} (\textsf{8-kv-per-cube}, sixteen users): the
+\emph{dense} mapping wins on throughput even though it is the \emph{slowest}
+per request.
+
+The reason is worth spelling out, because both mappings fill the same
+hardware: at capacity \textsf{1-kv-per-cube} runs two users
+$\times$ 64 PEs and \textsf{8-kv-per-cube} sixteen users $\times$ 8 PEs, so
+each puts all 128 PEs of the SIP to work. What differs is per-PE
+\emph{efficiency}. Splitting one head across eight PEs
+(\textsf{1-kv-per-cube}) leaves each PE only $S_{kv}/8$ tokens---a single
+tile---so its fixed overheads (pipeline fill on the lazy KV load, and the
+eight-PE online-softmax chain-reduce) dominate the little streaming work it
+does, and it sustains just \SI{46}{\percent} of the per-PE HBM ceiling.
+Giving each head a whole PE (\textsf{8-kv-per-cube}) streams the full
+$S_{kv}$-token cache contiguously with no reduce, amortizing those overheads
+over $8\times$ more work and reaching \SI{76}{\percent}. The measured
+throughput ratio ($1.41/0.86 = 1.6\times$) matches the utilization ratio
+($0.76/0.46 = 1.7\times$) almost exactly: because decode is
+bandwidth-bound, aggregate throughput is set by total HBM efficiency, not
+by PE count. \textsf{1-kv-per-cube} is really spending hardware on
+\emph{latency}---eight PEs per head buy only a $4.8\times$ speedup, a $60\%$
+strong-scaling efficiency---and that same $40\%$ overhead is what caps its
+throughput once the SIP is full. The design conclusion is regime-dependent: \textsf{1-kv-per-cube} is the
+latency-optimal choice for single-stream or low-batch decode, while denser
+mappings are the throughput-optimal choice for high-batch short-context
+serving. At long context the per-request latencies already scale as $1/C$,
+so the same accounting predicts the throughputs converge; a measured
+long-context batch sweep is 2H work (\S\ref{sec:future}).
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_short_context/batch_scaling.png}
+\caption{Batch scaling at $S_{kv}{=}8$K: aggregate throughput (requests per
+\si{\micro\second}) versus the number of concurrent decode users on one SIP,
+each user on a disjoint CUBE group. Each mapping scales almost linearly with
+$B$ up to its SIP capacity $16/C$---\textsf{1-kv-per-cube} saturates at two
+users, \textsf{8-kv-per-cube} at sixteen. Solid segments are measured
+concurrent runs; the dashed extensions are the saturated-throughput ceiling
+beyond capacity, where further users run in waves at that sustained rate
+(so the SIP is full, not idle). The dense mapping reaches the highest
+ceiling; the spread mapping is latency-optimal but capacity-limited.
+Concurrent users overlap to within \SI{3}{}--\SI{5}{\percent} of the
+single-user latency.}
+\label{fig:gqa-batch}
+\end{figure}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_short_context/wall_a1_variant_compare.png}
+\caption{Composite-tier ablation on \textsf{1-kv-per-cube} (log scale):
+wall-clock is nearly flat across tiers---the GEMM-only composite runs a few
+percent slower (padding overhead) and \textsf{softmax\_merge} matches the
+baseline. With $M{=}G{=}8 < \textsf{TILE\_M}{=}32$ the composite scheduler
+pads $M$ by $4\times$, leaving no fusion slack to recover at this shape.
+Fusion pays off only when $M$ fills the supertile---long-context Q-tile
+splitting or batched-$M$ inference.}
+\label{fig:gqa-short-a1-variant}
+\end{figure}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_short_context/gemm_util_a1_variant_ablation.png}
+\caption{GEMM-engine busy fraction on \textsf{1-kv-per-cube} across the
+three composite tiers. The baseline runs at \SI{12}{}--\SI{18}{\percent};
+the composite tiers read much higher---up to and past
+\SI{100}{\percent} of the short decode wall---because the padded
+$8\!\to\!32$ $M$ dimension is charged as GEMM time, bookkeeping overhead
+rather than useful work. The engine is never the bottleneck at this shape;
+decode is KV-bandwidth-bound.}
+\label{fig:gqa-short-gemm-util}
+\end{figure}
+
+\subsection{Inference with Long-Context Length}
+\label{sec:gqa-long}
+
+% TODO: prefill long-context kernel implementation description
+% (Sequence-Parallel partition of S_kv, per-case mechanics).
+% Bench code: src/kernbench/benches/gqa_helpers/long_ctx/
+
+% TODO: prefill long-context performance figure.
+
+Long-context decode is the regime where the KV cache, not attention
+compute, sets serving cost, so the placement question of
+\S\ref{sec:gqa-placement} becomes decisive here. To pick the right
+placement for this regime we run each of the six options as a fused
+decode kernel on the simulator (one decode step on the LLaMA-3.1-70B
+single-KV-head-group target, $C{=}8$ CUBEs $\times$ $P{=}8$ PEs) at a
+tractable $S_{kv}{=}8192$, and read off end-to-end latency, on-device op
+traffic, and the redundant compute each one induces. The swept context is
+small enough that all six fit in memory at $S_{kv}{=}8192$; the
+placements the long-context memory budget rules out (Cases~1--3) are
+drawn in red, run here only to expose their issue and communication
+structure.
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_latency.png}
+\caption{Measured end-to-end decode latency per placement
+($S_{kv}{=}8192$; red = ruled out by the long-context memory budget,
+blue = the predicted Pareto choice). The fastest raw latency belongs to
+Case~3 (\SI{17.8}{\micro\second})---but Case~3 replicates the full KV
+cache into every CUBE, so it overflows the per-PE budget at production
+context and wastes 8$\times$ the compute
+(Figure~\ref{fig:gqa-6cases-par}). Among the placements that actually fit
+1\,M-token memory (Cases~4--6), Case~6~$\star$ is the fastest
+(\SI{30.6}{\micro\second}, versus \SI{31.4}{} and \SI{34.5}{\micro\second}
+for the $d_{\text{head}}$-split Cases~4 and~5)---making it the placement
+of choice for long-context decode.}
+\label{fig:gqa-6cases-lat}
+\end{figure}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_parallelism.png}
+\caption{Redundant compute per placement, measured as active-PE
+$\times$ $S_{\text{local}}$ (PE-tokens; lower means less wasted work).
+The minimum is \num{8192} PE-tokens---one pass over the sequence.
+Case~3 inflates this 8$\times$ to \num{65536} by replicating the KV
+cache across all eight CUBEs so every CUBE redundantly re-attends the
+whole sequence; the $d_{\text{head}}$-split Cases~4--5 likewise carry
+\num{65536} because each token is processed across eight head slices.
+Case~6~$\star$ achieves the full 64-way split at the minimal
+\num{8192} PE-tokens---fully parallel, no replication.}
+\label{fig:gqa-6cases-par}
+\end{figure}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_traffic.png}
+\caption{Measured on-device op traffic per placement. The unsharded
+Case~1 issues no IPCQ copies (each PE has the full cache, nothing to
+reduce); the single-axis Case~3 charges 168. Among the 64-way splits,
+the $d_{\text{head}}$-split Cases~4--5 charge the most---280 IPCQ copies
+and 8 DMA writes each, the partial-score all-reduce that head-slicing
+forces---while Case~6~$\star$ needs only 189 IPCQ copies and a single DMA
+write, because it merges just the running softmax state $(m,\ell,O)$
+rather than partial scores. This is the on-device collective traffic that
+PE\_IPCQ and the torus links of \S\ref{sec:allreduce} are provisioned to
+absorb at link speed.}
+\label{fig:gqa-6cases-traffic}
+\end{figure}
+
+Taken together the three panels select the placement. The only
+memory-feasible family at production context is the 64-way splits
+(Cases~4--6), and within it Case~6~$\star$ is both the fastest and the
+lightest-communicating, because it merges only the running softmax state
+rather than the partial scores that the $d_{\text{head}}$-split
+Cases~4--5 must all-reduce. Case~3's lower raw latency is beside the
+point---it replicates the full cache into every CUBE, overflowing the
+per-PE budget and wasting $8\times$ the compute. For long-context decode
+the placement of choice is therefore the both-axes sequence shard, and the
+cross-PE softmax reduction it does pay is precisely the traffic the
+communication-side codesign of this report is built to move quickly.
+
+\subsection{Use of Composite Commands}
+\label{sec:gqa-composite}
+
+The decode kernel of \S\ref{sec:gqa-long} issues its local attention as
+primitive operations: it walks each PE's $S_{\text{local}}$ token slice in
+\SI{1024}{}-token tiles, and for every tile issues a $Q\!\cdot\!K^{\top}$
+\textsf{dot}, the online-softmax primitives, and a $P\!\cdot\!V$
+\textsf{dot}, merging the running $(m,\ell,O)$ state by hand. The number
+of PE\_CPU commands this costs grows with the context. At a production
+context of $S_{kv}{=}1\,\text{M}$ tokens the Case-6 64-way split gives
+each PE $S_{\text{local}}{=}16384$ tokens, so its local attention is a
+$Q\!\cdot\!K^{\top}$ of $(8,128)\!\cdot\!(128,16384)$ and a
+$P\!\cdot\!V$ of $(8,16384)\!\cdot\!(16384,128)$---sixteen hand-issued
+tiles, each a fresh batch of CPU commands.
+
+The composite command lets the kernel hand that tiling to PE\_SCHEDULER.
+We compare three command forms of the \emph{same} Case-6 kernel---identical
+placement and identical $(m,\ell,O)$ reduce, differing only in how the
+local attention is issued:
+\begin{itemize}
+  \item \textbf{primitive}---the hand-tiled \textsf{dot}/softmax kernel
+    above (the baseline of \S\ref{sec:gqa-long}).
+  \item \textbf{composite}---each matrix product is one coarse
+    \textsf{composite} GEMM over the \emph{whole} $S_{\text{local}}$, with
+    $K$ and $V$ passed as HBM references so PE\_SCHEDULER streams and
+    tiles them on the fixed $32\!\times\!64\!\times\!32$ MAC tile; the
+    softmax stays primitive.
+  \item \textbf{composite\,+\,softmax\_merge}---additionally folds the
+    online-softmax merge and $P\!\cdot\!V$ into a single stateful
+    \textsf{softmax\_merge} recipe composite.
+\end{itemize}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_decode_long_ctx_composite.png}
+\caption{Three command forms of the Case-6 decode kernel, swept over
+context length ($S_{\text{local}}{=}S_{kv}/64$ per PE). \emph{Right:}
+PE\_CPU commands issued. The hand-tiled primitive kernel rises
+$O(n_{\text{tiles}})$---from 96 commands at one tile to 426 at the
+1\,M-token, sixteen-tile production point---while both composite forms
+issue a context-\emph{independent} $O(1)$ count (94 and 98) that
+\emph{saturates}: one coarse descriptor offloads the entire per-tile
+fan-out. \emph{Left:} the consequence for wall-clock latency is none---all
+three land on the same curve (\SI{30.6}{}, \SI{231}{},
+\SI{461}{\micro\second} at 8\,K\,/\,64\,K\,/\,128\,K), because decode is
+bound by streaming the KV cache, not by issue. Command-count is measured
+at emit time (exact, to 1\,M); latency on the data-mode engine over the
+tractable range.}
+\label{fig:gqa-composite}
+\end{figure}
+
+Figure~\ref{fig:gqa-composite} reads off the two quantities that matter,
+and they point in opposite directions. The PE\_CPU command count (right)
+collapses from a context-growing $O(n_{\text{tiles}})$ to a flat $O(1)$:
+at 1\,M tokens the composite form issues \num{94} commands against the
+primitive kernel's \num{426}, a $4.5\times$ reduction
+($4.2\times$ in modeled dispatch cost), and---crucially---that number no
+longer grows with context. The wall-clock latency (left), by contrast,
+is unchanged across all three forms: decode is bound by streaming the KV
+cache out of HBM, so the command form does not move the critical path.
+
+That juxtaposition is the point. The composite command is not a latency
+optimization for this memory-bound decode; it is a \emph{CPU-issue}
+optimization. Its value is removing the per-tile dispatch work that would
+otherwise grow without bound as context grows, freeing PE\_CPU to run
+ahead and keep the engines fed---which is exactly what lets the
+data-movement cost analyzed next show through as the true bottleneck
+rather than being masked by issue overhead. The \textsf{softmax\_merge}
+recipe folds the online merge into the same descriptor; on this
+memory-bound path its marginal cost over the plain GEMM composite is small
+(98 vs.\ 94 commands), and like the plain composite it keeps the issued
+count flat as context scales.
+
+\paragraph{The compute-bound mirror: prefill.} Decode's verdict---command
+form is latency-neutral---is a property of its regime, not of the
+composite command. A decode step has $T_q{=}1$, so its score and context
+products are skinny ($M{=}G\,T_q{=}8$): the MAC array is barely fed and
+the kernel is bound by streaming the KV cache. Prefill is the opposite
+corner. It processes a block of query positions at once, so $M{=}G\,T_q$
+is large and tile-filling, the GEMMs carry real arithmetic intensity
+($\sim$$M$ flops/byte, well above the roofline ridge), and the kernel is
+\emph{compute-bound}. This is the regime the composite command was built
+for (\S\ref{sec:gemm}): it streams the per-HW-tile
+DMA$\rightleftarrows$compute pipeline so the MAC array stays fed, whereas
+the primitive kernel's blocking \textsf{tl.dot} serializes each tile's
+load and compute and starves the array between tiles. We run the same
+three command forms on a single-rank compute-bound prefill (FlashAttention
+$Q$-block $\times$ $S_{kv}$-tile, online softmax) and sweep the context
+length (Figure~\ref{fig:gqa-prefill-cb}).
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_prefill_compute_bound.png}
+\caption{Compute-bound prefill, three command forms, swept over context
+length ($M{=}8\,T_q$ tile-filling). \emph{Left:} end-to-end latency.
+\emph{Right:} MAC utilization (achieved $\div$ the
+\SI{8}{\tera\flop\per\second} per-PE peak). The hand-tiled primitive sits
+flat at $\sim$\SI{68}{\percent}---its serial load$\to$dot path leaves the
+MAC array idle between tiles regardless of context. The composite forms
+climb with context (\SI{67}{}$\to$\SI{80}{\percent} plain,
+\SI{67}{}$\to$\SI{83}{\percent} with the recipe) because a deeper $P\!\cdot
+\!V$ reduction gives more HW tiles to pipeline, and they convert that into
+wall-clock: at \num{1024} the recipe form is \SI{646.9}{} vs.\
+\SI{794.1}{\micro\second} (\SI{19}{\percent} faster). The margin
+\emph{grows} with context---the compute-bound mirror of the GEMM result of
+\S\ref{sec:gemm}.}
+\label{fig:gqa-prefill-cb}
+\end{figure}
+
+The two studies together state the composite command's value precisely. It
+has two distinct benefits, and which one matters is set by the workload's
+roofline position. The first is \emph{host-issue offload}: one macro
+command in place of $O(n_{\text{tiles}})$ fine ones, which removes
+PE\_CPU dispatch work and is regime-independent (it shows in the decode
+command count). The second is \emph{MAC-array feeding}: the
+scheduler-internal per-tile DMA$\rightleftarrows$compute pipeline, which
+only converts to latency when the workload is compute-bound enough to have
+a MAC array worth keeping busy (it shows in the prefill utilization). A
+memory-bound decode exercises only the first; a compute-bound prefill
+exercises both. The composite command is the single mechanism that
+delivers each where it applies.
+
+% Summary moved to sections/05z-summary.tex so the three new
+% subsections (roofline, capacity planning, parallelism selection)
+% can appear between the technical results and the closing summary.
+
+% TODO: cross-regime DP (data parallelism) applicability:
+%   - Does Case-4 long-context placement compose with batch-level DP
+%     without further changes?
+%   - Does the short-context placement compose the same way?
+%   - Implications for multi-user serving (single vs. mixed regimes).

docs/report/1H-codesign-paper/sections/05z-summary.tex

@@ -0,0 +1,22 @@
+\subsection{Summary}
+\label{sec:gqa-analysis}
+
+The section's results line up into one picture of the fused kernel.
+Placement is decided by memory first and communication second: only the
+64-way splits fit production context, and among them the both-axes
+sequence shard (Case~6) minimizes the per-token collective by merging
+softmax state rather than partial scores. Command form is decided by
+roofline position: the composite command removes the per-tile issue work
+in both regimes, but converts to wall-clock only in compute-bound prefill,
+while memory-bound decode stays bound by KV streaming regardless of command
+form. The thread connecting the two is that, once composite issue makes
+GEMM cheap and leaves the MAC array idle, the fused kernel's latency is set
+almost entirely by data movement---streaming the KV cache and reducing
+partials across PEs---which is exactly the cost the communication-side
+work (on-device reduction, lazy load/compute overlap, fast TCM staging and
+torus links) is built to attack. Roofline framing
+(\S\ref{sec:roofline}) turns this observation into a sizing and sharding
+programme in the next two subsections: capacity planning
+(\S\ref{sec:capacity-planning}) and parallelism selection
+(\S\ref{sec:parallelism-selection}). What this implies for hardware
+investment across the report as a whole is taken up in the discussion.

tests/analytical_visualization/_gen_roofline_paper_figs.py

@@ -0,0 +1,104 @@
+"""Generate roofline figures for the paper's Roofline Analysis section.
+
+Produces two 2-panel PNGs into docs/report/1H-codesign-paper/figures/:
+  - roofline_short_context.png  (S_kv = 8K)   — batch drops per-token cost
+  - roofline_long_context.png   (S_kv = 1M)   — batch effect vanishes
+
+Each figure shows both the step-latency view (raw step time vs B, with
+weight/compute/KV components) and the cost-per-token view (÷B), which
+together make the batch-vs-context story visible at a glance.
+
+Run: python tests/analytical_visualization/_gen_roofline_paper_figs.py
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+# Ensure the repo root is on sys.path so `tests.analytical_visualization`
+# imports resolve when invoked as a script.
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from tests.analytical_visualization.chip_roofline import (
+    critical_batch,
+    per_token_latency_curve,
+    step_latency_curve,
+)
+from tests.analytical_visualization.model_config import MachineParams
+from tests.analytical_visualization.model_presets import PRESETS
+
+FIGURES_DIR = Path(__file__).resolve().parents[2] / (
+    "docs/report/1H-codesign-paper/figures"
+)
+FIGURES_DIR.mkdir(parents=True, exist_ok=True)
+
+MACHINE = MachineParams()
+MODEL = PRESETS["Llama 3 70B"].model
+B_RANGE = [1, 2, 4, 8, 16, 32, 64, 128, 256]
+
+
+def _plot_pair(s_kv: int, label: str, out_path: Path) -> None:
+    step = step_latency_curve(MACHINE, MODEL, B_RANGE, s_kv=s_kv)
+    tok = per_token_latency_curve(MACHINE, MODEL, B_RANGE, s_kv=s_kv)
+    xs = [p.batch for p in step]
+    b_star = critical_batch(MACHINE, MODEL)
+
+    fig, (axA, axB) = plt.subplots(1, 2, figsize=(11, 4.2))
+
+    # ── Left: step latency (undivided) ────────────────────────────────
+    axA.plot(xs, [p.weight_s * 1e3 for p in step], "^-",
+              color="#ffbe0b", label="Weight fetch (flat)")
+    axA.plot(xs, [p.compute_s * 1e3 for p in step], "o-",
+              color="#3a86ff", label="Compute (linear)")
+    axA.plot(xs, [p.kv_s * 1e3 for p in step], "s-",
+              color="#d90429", label="KV fetch (linear)")
+    axA.plot(xs, [p.total_s * 1e3 for p in step], "-",
+              color="#212529", linewidth=2.5, label="Total")
+    axA.set_xscale("log", base=2)
+    axA.set_yscale("log")
+    axA.set_xlabel("Batch size B")
+    axA.set_ylabel("Step time (ms)")
+    axA.set_title(f"Step latency ({label} context, $S_{{kv}}={s_kv:,}$)")
+    axA.grid(True, which="both", alpha=0.3)
+    axA.legend(fontsize=8, loc="upper left")
+
+    # ── Right: cost per token (÷ B) ───────────────────────────────────
+    axB.plot(xs, [p.weight_s * 1e3 for p in tok], "^-",
+              color="#ffbe0b", label="Weight fetch ($\\propto 1/B$)")
+    axB.plot(xs, [p.compute_s * 1e3 for p in tok], "o--",
+              color="#3a86ff", label="Compute (flat)")
+    axB.plot(xs, [p.kv_s * 1e3 for p in tok], "s--",
+              color="#d90429", label="KV fetch (flat)")
+    axB.plot(xs, [p.total_s * 1e3 for p in tok], "-",
+              color="#212529", linewidth=2.5, label="Total")
+    axB.axvline(b_star, linestyle=":", color="#2e7d32",
+                 label=f"$B^*={b_star:.0f}$")
+    axB.set_xscale("log", base=2)
+    axB.set_yscale("log")
+    axB.set_xlabel("Batch size B")
+    axB.set_ylabel("Per-token time (ms)")
+    axB.set_title(f"Cost per token ({label} context, $S_{{kv}}={s_kv:,}$)")
+    axB.grid(True, which="both", alpha=0.3)
+    axB.legend(fontsize=8, loc="upper right")
+
+    plt.tight_layout()
+    fig.savefig(out_path, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    print(f"  wrote {out_path.name}")
+
+
+def main() -> None:
+    print("Generating roofline figures for the paper ...")
+    _plot_pair(8_192, "short", FIGURES_DIR / "roofline_short_context.png")
+    _plot_pair(1_048_576, "long", FIGURES_DIR / "roofline_long_context.png")
+    print(f"Done. Output: {FIGURES_DIR}")
+
+
+if __name__ == "__main__":
+    main()

tests/analytical_visualization/auto_explore.py

@@ -419,6 +419,7 @@ def run_auto_explore(
     mode: str = "decode",
     include_attention: bool = True,
     include_ffn: bool = True,
+    b: int = 1,
 ) -> AutoExploreResult:
     """Enumerate all configs, score each, extract Pareto frontier.
 
@@ -433,6 +434,8 @@ def run_auto_explore(
     all_scores: list[ConfigScore] = []
     total_enumerated = 0
     for topo in enumerate_configs(model, s_kv, mode):
+        # Inject the caller's batch B into every candidate.
+        topo.b = b
         total_enumerated += 1
         cfg = FullConfig(model=model, topo=topo, machine=machine)
         all_scores.append(score_config(

tests/analytical_visualization/auto_hardware.py

@@ -237,6 +237,7 @@ def _best_parallelism_for_hw(
     model: ModelConfig, machine: MachineParams,
     s_kv: int, mode: str,
     include_attention: bool = True, include_ffn: bool = True,
+    b: int = 1,
 ) -> ConfigScore | None:
     """Return the latency-minimum feasible parallelism for this HW.
 
@@ -245,14 +246,20 @@ def _best_parallelism_for_hw(
     matches the caller's attention/FFN/full choice.
     """
     best: ConfigScore | None = None
+    best_key: tuple | None = None
     for topo in _iter_reduced_parallelism(model, s_kv, mode):
+        topo.b = b
         cfg = FullConfig(model=model, topo=topo, machine=machine)
         s = score_config(cfg, include_attention=include_attention,
                           include_ffn=include_ffn)
         if not (s.fits_memory and s.placement_valid):
             continue
-        if best is None or s.total_latency_ns < best.total_latency_ns:
+        # Compound tiebreaker: prefer fewer PEs and lower HBM % when latency
+        # ties across variant knobs (cp_ring, placement, kv_shard_mode, etc.).
+        _k = (s.total_latency_ns, s.pes_used, s.hbm_utilization)
+        if best is None or _k < best_key:
             best = s
+            best_key = _k
     return best
 
 
@@ -260,13 +267,14 @@ def _best_parallelism_two_stage(
     model: ModelConfig, machine: MachineParams,
     s_kv: int, mode: str,
     include_attention: bool = True, include_ffn: bool = True,
+    b: int = 1,
 ) -> ConfigScore | None:
     """Fast fallback: use autosuggest's memory-min (CP,TP,PP) then score it."""
     sug = auto_suggest(model, machine, s_kv, mode)
     if not sug.fits:
         return None
     topo = TopologyConfig(
-        cp=sug.cp, tp=sug.tp, pp=sug.pp, dp=1,
+        cp=sug.cp, tp=sug.tp, pp=sug.pp, dp=1, b=b,
         s_kv=s_kv, mode=mode,
         kv_shard_mode="split",
         ffn_shard_scope="TP+CP",
@@ -361,6 +369,7 @@ def joint_explore(
     depth: str = "balanced",
     include_attention: bool = True,
     include_ffn: bool = True,
+    b: int = 1,
 ) -> JointExploreResult:
     """Sweep HW candidates × parallelism, return joint Pareto + sensitivity.
 
@@ -378,11 +387,13 @@ def joint_explore(
             par = _best_parallelism_two_stage(
                 model, machine, s_kv, mode,
                 include_attention=include_attention, include_ffn=include_ffn,
+                b=b,
             )
         else:
             par = _best_parallelism_for_hw(
                 model, machine, s_kv, mode,
                 include_attention=include_attention, include_ffn=include_ffn,
+                b=b,
             )
         if par is None:
             continue

tests/analytical_visualization/autosuggest.py

@@ -32,6 +32,8 @@ class Suggestion:
     fits: bool
     pes_used: int
     sips_used: int
+    cubes_used: int = 0            # picked to minimize this (see _score_candidate)
+    cp_placement: str = "cube"     # "pe" packs CP into intra-cube PEs when it fits
     reason: str = ""
 
 
@@ -58,10 +60,33 @@ def _iter_candidates(model: ModelConfig) -> Iterable[tuple[int, int, int]]:
 def _score_candidate(cp: int, tp: int, pp: int,
                      model: ModelConfig, machine: MachineParams,
                      s_kv: int, mode: str,
-                     slack_frac: float = 0.10) -> Suggestion:
+                     slack_frac: float = 0.10,
-    topo = TopologyConfig(cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode)
+                     b: int = 1,
+                     include_attention: bool = True,
+                     include_ffn: bool = True) -> Suggestion:
+    # For each triple, try both cp_placement options and keep the one
+    # with fewer cubes (breaks the historical "CP always across cubes"
+    # default when a smaller pack is possible). The pe placement is only
+    # valid when CP·TP fits within a single cube's PE count.
+    best_topo: TopologyConfig | None = None
+    best_placement = "cube"
+    _pe_per_cube_hw = TopologyConfig().pes_per_cube_hw  # instance-invariant HW const
+    _placements_to_try = ["cube"]
+    if cp * tp <= _pe_per_cube_hw:
+        _placements_to_try.append("pe")
+    for _place in _placements_to_try:
+        _t = TopologyConfig(
+            cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode, b=b,
+            cp_placement=_place,
+        )
+        if best_topo is None or _t.cubes_used < best_topo.cubes_used:
+            best_topo = _t
+            best_placement = _place
+    topo = best_topo
     cfg = FullConfig(model=model, topo=topo, machine=machine)
-    mem = compute_memory(cfg)
+    mem = compute_memory(
+        cfg, include_attention=include_attention, include_ffn=include_ffn,
+    )
 
     fits = (not mem.over_budget
             and mem.slack_bytes >= slack_frac * mem.budget_bytes)
@@ -80,27 +105,42 @@ def _score_candidate(cp: int, tp: int, pp: int,
         fits=fits,
         pes_used=topo.total_pes,
         sips_used=topo.sips_used,
+        cubes_used=topo.cubes_used,
+        cp_placement=best_placement,
         reason=reason,
     )
 
 
 def auto_suggest(model: ModelConfig, machine: MachineParams,
                  s_kv: int, mode: str = "decode",
-                 slack_frac: float = 0.10) -> Suggestion:
+                 slack_frac: float = 0.10,
-    """Return the smallest-PE-count (CP, TP, PP) that fits.
+                 b: int = 1,
+                 include_attention: bool = True,
+                 include_ffn: bool = True) -> Suggestion:
+    """Return the smallest deployment (fewer cubes, then fewer PEs)
+    that fits.
+
+    Sort key: (cubes_used ↑, pes_used ↑, pp ↑, tp ↑, cp ↑).
+    Fewer cubes wins first because a cube is the physical die-level
+    hardware unit; PE count is the tiebreaker. Preferring fewer PP then
+    TP then CP keeps the earlier historical bias (avoid pipeline
+    bubbles / head-dim splits) among equal-cube-and-PE ties.
 
     If no candidate fits, returns the best-effort one (highest slack,
     even if negative) with fits=False.
     """
-    best_fit: Suggestion | None = None
+    scored: list[Suggestion] = []
-    best_effort: Suggestion | None = None
-
     for cp, tp, pp in _iter_candidates(model):
-        s = _score_candidate(cp, tp, pp, model, machine, s_kv, mode, slack_frac)
+        scored.append(
-        if s.fits and best_fit is None:
+            _score_candidate(cp, tp, pp, model, machine, s_kv, mode,
-            best_fit = s
+                              slack_frac, b=b,
-            break  # candidates are pre-sorted by PE count
+                              include_attention=include_attention,
-        if best_effort is None or s.slack_gb > best_effort.slack_gb:
+                              include_ffn=include_ffn)
-            best_effort = s
+        )
+    scored.sort(key=lambda s: (s.cubes_used, s.pes_used, s.pp, s.tp, s.cp))
 
-    return best_fit or best_effort  # type: ignore
+    for s in scored:
+        if s.fits:
+            return s
+    # No fit — return the best-effort (largest slack, closest to fitting).
+    return max(scored, key=lambda s: s.slack_gb)

tests/analytical_visualization/chip_roofline.py

@@ -0,0 +1,518 @@
+"""Chip-level roofline math: AI, B*, L*, per-token latency curves.
+
+Surfaces the arithmetic-intensity story from LLM-serving practice:
+
+- **AI**  = C / W  (peak FLOPs per byte of HBM bandwidth).
+- **B\***  = C * b / (2 * W) * sparsity  — the critical batch size at
+  which weight-fetch time equals compute time for one decode step.
+  Sparsity = N_total / N_active (MoE factor; 1 for dense).
+- **L\***  = 2 * N_active / (AI * kv_bytes_per_token) — the balance
+  context length at which KV-read time equals compute time.
+- **B_knee(S_kv)** = B* / (1 - S_kv/L*) — the batch size where the
+  cost curve bends (weight-fetch drops below the compute+KV floor).
+  Diverges at S_kv = L* and no knee exists past it.
+
+Per-token decode-step latency, per PE (dense-approx, no comm):
+
+    t(B, S_kv)  =  N_active * b / (W * B)     <-- weight fetch, 1/B
+                 + 2 * N_active / C            <-- compute (peak), flat
+                 + S_kv * kv_bpt / W           <-- KV read, flat
+
+All numbers per PE / per one forward pass. **Peak roofline — no
+utilization factor.** Comm cost and TP/CP sharding are intentionally
+NOT in the roofline — this is the back-of-envelope chip-vs-model view
+the transcript talks about, not the full latency model that
+stage_latencies.py builds. With this convention weight_s == compute_s
+exactly at B*.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+
+from .model_config import FullConfig, MachineParams, ModelConfig
+
+
+# Per-token bytes of BF16 MAC arithmetic: one multiply + one add = 2 FLOPs.
+_FLOPS_PER_PARAM_PER_TOKEN = 2
+
+
+# ── Chip / model derived quantities ────────────────────────────────
+
+
+def arithmetic_intensity(machine: MachineParams) -> float:
+    """FLOPs per byte of HBM bandwidth. Peak roofline; utilization
+    is applied only in the compute-time formula, not here."""
+    return machine.peak_flops / machine.bw_hbm
+
+
+def total_active_params(model: ModelConfig) -> int:
+    """Full-model parameter count (attention + FFN, all layers).
+
+    Attention: 4 projections × hidden × H_q * d_head effective per layer.
+    (W_Q hidden×H_q*d_h, W_O H_q*d_h×hidden, W_K/W_V hidden×H_kv*d_h.)
+    FFN: 3 × hidden × ffn_dim per layer (gate, up, down).
+    """
+    m = model
+    attn = (
+        m.hidden * m.h_q * m.d_head            # W_Q
+        + m.hidden * m.h_kv * m.d_head * 2     # W_K + W_V
+        + m.h_q * m.d_head * m.hidden          # W_O
+    )
+    ffn = 3 * m.hidden * m.ffn_dim
+    return (attn + ffn) * m.layers
+
+
+def kv_bytes_per_token(model: ModelConfig) -> int:
+    """Bytes of KV cache one new token adds across ALL layers, per one
+    sequence, un-sharded (K + V, H_kv heads * d_h * bytes)."""
+    m = model
+    return 2 * m.h_kv * m.d_head * m.bytes_per_elem * m.layers
+
+
+def critical_batch(machine: MachineParams, model: ModelConfig,
+                    sparsity: float = 1.0) -> float:
+    """B* = C * b / (2 * W) * sparsity.
+
+    Sparsity = N_total / N_active (>= 1). Dense = 1. MoE 8-of-256 = 8.
+    """
+    b = model.bytes_per_elem
+    ai = arithmetic_intensity(machine)
+    return ai * b / _FLOPS_PER_PARAM_PER_TOKEN * sparsity
+
+
+def balance_context(machine: MachineParams, model: ModelConfig) -> float:
+    """L* = 2 * N_active / (AI * kv_bpt).
+
+    Context length (in tokens) at which per-step KV read matches the
+    per-step compute cost. Beyond L*, the KV term is dominant and no
+    batch size gets you compute-bound.
+    """
+    n_active = total_active_params(model)
+    ai = arithmetic_intensity(machine)
+    kv_bpt = kv_bytes_per_token(model)
+    return _FLOPS_PER_PARAM_PER_TOKEN * n_active / (ai * kv_bpt)
+
+
+def knee_batch(machine: MachineParams, model: ModelConfig,
+                s_kv: int) -> float | None:
+    """B_knee(S_kv) = B* / (1 - S_kv/L*).
+
+    Returns None when S_kv >= L* (no knee exists — the total-cost
+    curve never touches the compute floor).
+    """
+    b_star = critical_batch(machine, model)
+    l_star = balance_context(machine, model)
+    r = s_kv / l_star
+    if r >= 1.0:
+        return None
+    return b_star / (1.0 - r)
+
+
+# ── Per-token latency curves ───────────────────────────────────────
+
+
+@dataclass
+class RooflinePoint:
+    batch: int
+    weight_s: float   # weight fetch time, 1/B
+    compute_s: float  # compute time, flat
+    kv_s: float       # KV read time, flat
+    total_s: float
+
+
+def per_token_latency_curve(machine: MachineParams, model: ModelConfig,
+                             batch_range: list[int],
+                             s_kv: int) -> list[RooflinePoint]:
+    """Per-token decode-step latency curve across a range of batch sizes.
+
+    Returns one point per batch. All times per PE, dense-approx,
+    utilization from machine.compute_util. Comm and TP/CP sharding are
+    excluded — this is the roofline model.
+    """
+    n_active = total_active_params(model)
+    b = model.bytes_per_elem
+    weight_bytes = n_active * b
+    compute_flops = _FLOPS_PER_PARAM_PER_TOKEN * n_active
+    kv_read_bytes = s_kv * kv_bytes_per_token(model)
+
+    compute_s = compute_flops / machine.peak_flops
+    kv_s = kv_read_bytes / machine.bw_hbm
+
+    points: list[RooflinePoint] = []
+    for bs in batch_range:
+        weight_s = weight_bytes / machine.bw_hbm / max(1, bs)
+        points.append(RooflinePoint(
+            batch=bs,
+            weight_s=weight_s,
+            compute_s=compute_s,
+            kv_s=kv_s,
+            total_s=weight_s + compute_s + kv_s,
+        ))
+    return points
+
+
+def bound_regime(machine: MachineParams, model: ModelConfig,
+                  batch: int, s_kv: int) -> str:
+    """Which term dominates at the current (batch, S_kv) point.
+
+    Returns 'memory-bound' if weight_fetch is the largest term,
+    'kv-bound' if KV read is largest, 'compute-bound' if compute.
+    """
+    pts = per_token_latency_curve(machine, model, [batch], s_kv)
+    p = pts[0]
+    parts = {"memory-bound": p.weight_s,
+             "kv-bound": p.kv_s,
+             "compute-bound": p.compute_s}
+    return max(parts, key=parts.get)
+
+
+# ── Regime-dependent cost terms ────────────────────────────────────
+
+
+def t_mem_short(machine: MachineParams, model: ModelConfig,
+                 batch: int) -> float:
+    """Per-token weight-fetch time (short-context regime term).
+
+    N_active · b / (W · B). Shrinks as B grows — this is what
+    batching amortizes.
+    """
+    return (total_active_params(model) * model.bytes_per_elem
+            / (machine.bw_hbm * max(1, batch)))
+
+
+def t_mem_long(machine: MachineParams, model: ModelConfig,
+                s_kv: int) -> float:
+    """Per-token KV-read time (long-context regime term).
+
+    S_kv · kv_bpt / W. Independent of B — each sequence reads its
+    own KV cache; batching doesn't help.
+    """
+    return s_kv * kv_bytes_per_token(model) / machine.bw_hbm
+
+
+def t_com(machine: MachineParams, model: ModelConfig) -> float:
+    """Per-token compute time. Same in both regimes: 2·N/C, peak."""
+    return _FLOPS_PER_PARAM_PER_TOKEN * total_active_params(model) / machine.peak_flops
+
+
+# ── "Good" batch / context recommendations ─────────────────────────
+
+
+@dataclass
+class BatchRecommendation:
+    target: float          # Pope's rule: 2 × B*
+    b_star: float          # B* itself
+    effective: float       # what we recommend using
+    reason: str            # short explanation
+
+
+@dataclass
+class ContextRecommendation:
+    l_star: float          # balance context length
+    max_efficient: float   # same as l_star (compute-friendly ceiling)
+    utilization_at: float  # utilization at current s_kv
+    reason: str
+
+
+def good_batch(machine: MachineParams, model: ModelConfig,
+                sparsity: float = 1.0) -> BatchRecommendation:
+    """Recommended batch size: 2 × B* (Pope's rule of thumb).
+
+    Below B*: memory-bound, doubling B halves cost/token.
+    At 2×B*: 50% excess over compute floor — the sweet spot.
+    Beyond 3×B*: diminishing returns; latency keeps growing linearly.
+    """
+    b_star = critical_batch(machine, model, sparsity)
+    target = 2 * b_star
+    return BatchRecommendation(
+        target=target, b_star=b_star, effective=target,
+        reason=(f"2·B* = 2 · {b_star:.0f} = {target:.0f}. "
+                "Below B*: memory-bound (doubling B halves cost/token). "
+                "Beyond 3·B*: diminishing returns."),
+    )
+
+
+def good_context(machine: MachineParams, model: ModelConfig,
+                  s_kv: int) -> ContextRecommendation:
+    """Recommended max context: L* — the compute-friendly ceiling.
+
+    Below L*: KV read is cheap relative to compute → good utilization.
+    Above L*: KV bandwidth wall → utilization = 1/(1 + S_kv/L*).
+    """
+    l_star = balance_context(machine, model)
+    util = utilization_at(s_kv, l_star)
+    return ContextRecommendation(
+        l_star=l_star, max_efficient=l_star, utilization_at=util,
+        reason=(f"L* = {l_star:,.0f} tokens. Below L*: compute-bound "
+                f"(good util). At {s_kv:,} tokens: peak utilization ≈ "
+                f"{util*100:.1f}% (1 / (1 + S_kv/L*))."),
+    )
+
+
+def utilization_at(s_kv: int, l_star: float) -> float:
+    """Peak compute utilization at context length s_kv, given L*.
+
+    util = compute / (compute + KV_read) = 1 / (1 + S_kv/L*).
+    At S_kv=0: 100%. At S_kv=L*: 50%. At 2·L*: 33.3%. At 5·L*: 16.7%.
+    """
+    return 1.0 / (1.0 + s_kv / l_star)
+
+
+# ── Per-step latency (undivided by B) ─────────────────────────────
+
+
+@dataclass
+class StepLatencyPoint:
+    batch: int
+    weight_s: float   # N·b / W — flat in B (loaded once per step)
+    compute_s: float  # 2·N·B / C — linear in B
+    kv_s: float       # B · S_kv · kv_bpt / W — linear in B
+    total_s: float
+
+
+def step_latency_curve(machine: MachineParams, model: ModelConfig,
+                        batch_range: list[int],
+                        s_kv: int) -> list[StepLatencyPoint]:
+    """Total time of one decode step (one forward pass), across a
+    range of batch sizes. NOT divided by B — this is the SLO view.
+
+    step_weight = N·b / W                (batch-invariant)
+    step_compute = 2·N·B / C              (linear in B)
+    step_kv = B · S_kv · kv_bpt / W       (linear in B)
+
+    Per-token cost = step_total / B — the two views trade off:
+    bigger B lowers cost/token but raises step latency.
+    """
+    n_active = total_active_params(model)
+    b = model.bytes_per_elem
+    weight_bytes = n_active * b
+    kv_bpt = kv_bytes_per_token(model)
+
+    step_weight = weight_bytes / machine.bw_hbm
+
+    points: list[StepLatencyPoint] = []
+    for bs in batch_range:
+        B = max(1, bs)
+        step_compute = _FLOPS_PER_PARAM_PER_TOKEN * n_active * B / machine.peak_flops
+        step_kv = B * s_kv * kv_bpt / machine.bw_hbm
+        total = step_weight + step_compute + step_kv
+        points.append(StepLatencyPoint(
+            batch=bs,
+            weight_s=step_weight,
+            compute_s=step_compute,
+            kv_s=step_kv,
+            total_s=total,
+        ))
+    return points
+
+
+# ── PE memory budget curves ───────────────────────────────────────
+
+
+@dataclass
+class MemoryBudgetPoint:
+    axis_val: int          # S_kv or B being swept
+    weights_gb: float
+    kv_gb: float
+    transient_gb: float
+    used_gb: float
+    free_gb: float         # max(0, hbm_gb - used_gb)
+    over_budget: bool
+
+
+def _budget_point(weights_bytes: int, kv_bytes: int, transient_bytes: int,
+                   hbm_bytes: int, axis_val: int) -> MemoryBudgetPoint:
+    used = weights_bytes + kv_bytes + transient_bytes
+    return MemoryBudgetPoint(
+        axis_val=axis_val,
+        weights_gb=weights_bytes / 1e9,
+        kv_gb=kv_bytes / 1e9,
+        transient_gb=transient_bytes / 1e9,
+        used_gb=used / 1e9,
+        free_gb=max(0, hbm_bytes - used) / 1e9,
+        over_budget=(used > hbm_bytes),
+    )
+
+
+def memory_budget_curve_vs_skv(cfg: FullConfig,
+                                s_kv_range: list[int],
+                                batch: int) -> list[MemoryBudgetPoint]:
+    """Per-PE memory as S_kv sweeps. Uses cfg's current sharding
+    (CP, TP, PP). Sets topo.b = batch. Returns one point per S_kv."""
+    from .memory_layout import (
+        per_pe_kv_cache_bytes, per_pe_transient_bytes, per_pe_weight_bytes,
+    )
+    hbm_bytes = int(cfg.machine.pe_budget_bytes)
+    weights_bytes = per_pe_weight_bytes(cfg)
+    transient_bytes = per_pe_transient_bytes(cfg)
+    points: list[MemoryBudgetPoint] = []
+    for skv in s_kv_range:
+        swept_topo = replace(cfg.topo, s_kv=int(skv), b=max(1, int(batch)))
+        swept_cfg = FullConfig(model=cfg.model, topo=swept_topo,
+                                machine=cfg.machine)
+        kv_bytes = per_pe_kv_cache_bytes(swept_cfg)
+        points.append(_budget_point(weights_bytes, kv_bytes,
+                                     transient_bytes, hbm_bytes, int(skv)))
+    return points
+
+
+def memory_budget_curve_vs_batch(cfg: FullConfig,
+                                   b_range: list[int],
+                                   s_kv: int) -> list[MemoryBudgetPoint]:
+    """Per-PE memory as B sweeps. Sets topo.s_kv = s_kv."""
+    from .memory_layout import (
+        per_pe_kv_cache_bytes, per_pe_transient_bytes, per_pe_weight_bytes,
+    )
+    hbm_bytes = int(cfg.machine.pe_budget_bytes)
+    weights_bytes = per_pe_weight_bytes(cfg)
+    transient_bytes = per_pe_transient_bytes(cfg)
+    points: list[MemoryBudgetPoint] = []
+    for bs in b_range:
+        swept_topo = replace(cfg.topo, b=max(1, int(bs)), s_kv=int(s_kv))
+        swept_cfg = FullConfig(model=cfg.model, topo=swept_topo,
+                                machine=cfg.machine)
+        kv_bytes = per_pe_kv_cache_bytes(swept_cfg)
+        points.append(_budget_point(weights_bytes, kv_bytes,
+                                     transient_bytes, hbm_bytes, int(bs)))
+    return points
+
+
+# ── AI / B* sensitivity to hardware knobs ──────────────────────────
+
+
+@dataclass
+class AISensitivityPoint:
+    multiplier: float          # scale factor applied to the base machine
+    peak_tflops: float
+    bw_gbs: float
+    ai: float                  # C / W
+    b_star: float              # C·b/(2·W)
+
+
+# ── GPU count sizing (three-axis) ─────────────────────────────────
+
+
+@dataclass
+class GpuSizingResult:
+    # Inputs echoed back for display / debugging
+    n_users: int
+    avg_ctx_tokens: int
+    tpot_slo_s: float
+    # Per-replica values
+    pes_axis_a_capacity: int    # PEs to hold weights alone (bare floor)
+    pes_axis_b_kv: int          # PEs to hold weights + KV of this replica's users
+    pes_per_replica: int        # max(A, B)
+    # Throughput axis
+    b_at_slo: int               # largest per-replica B satisfying SLO
+    users_per_replica: int      # min(n_users, b_at_slo), at least 1
+    n_replicas: int             # ceil(n_users / users_per_replica)
+    # Grand total
+    total_pes: int
+    binding_axis: str           # 'capacity' | 'kv' | 'throughput'
+
+
+def max_batch_within_slo(machine: MachineParams, model: ModelConfig,
+                          s_kv: int, tpot_slo_s: float) -> int:
+    """Largest per-replica B such that decode step latency ≤ SLO.
+
+    step_latency(B) = N·b/W + 2·N·B/C + B·S_kv·kv_bpt/W
+                    = weight_fetch + B · (compute + kv_read)
+
+    Returns 0 if even B=1 exceeds SLO (weight fetch alone too large,
+    or per-sequence term already blows the budget).
+    """
+    n = total_active_params(model)
+    b_elem = model.bytes_per_elem
+    weight_time = n * b_elem / machine.bw_hbm
+    per_seq_time = (_FLOPS_PER_PARAM_PER_TOKEN * n / machine.peak_flops
+                    + s_kv * kv_bytes_per_token(model) / machine.bw_hbm)
+    if weight_time + per_seq_time > tpot_slo_s:
+        return 0
+    remaining = tpot_slo_s - weight_time
+    b_max = int(remaining / per_seq_time)
+    return max(1, b_max)
+
+
+def size_deployment(machine: MachineParams, model: ModelConfig,
+                     n_users: int, avg_ctx: int,
+                     tpot_slo_s: float) -> GpuSizingResult:
+    """Three-axis GPU count for a target workload.
+
+    Axis A — capacity floor: PEs to hold one replica's weights.
+    Axis B — KV headroom:    PEs to hold weights + all KV of the
+                             users assigned to this replica.
+    Axis C — throughput SLO: replicas needed so per-replica B ≤ b_at_slo.
+
+    Result: pes_per_replica = max(A, B); total = pes_per_replica × replicas.
+    Binding axis is whichever grew the count the most.
+    """
+    n = total_active_params(model)
+    b_elem = model.bytes_per_elem
+    weight_bytes = n * b_elem
+    hbm_pe = int(machine.pe_hbm_gb * 1e9)
+    kv_bpt = kv_bytes_per_token(model)
+
+    # Axis C: throughput
+    b_at_slo = max_batch_within_slo(machine, model, avg_ctx, tpot_slo_s)
+    users_per_replica = max(1, min(n_users, b_at_slo)) if b_at_slo > 0 else 1
+    n_replicas = (n_users + users_per_replica - 1) // users_per_replica
+
+    # Axis A: bare weights
+    pes_a = (weight_bytes + hbm_pe - 1) // hbm_pe
+
+    # Axis B: weights + KV load for this replica's users
+    kv_per_replica = users_per_replica * avg_ctx * kv_bpt
+    pes_b = (weight_bytes + kv_per_replica + hbm_pe - 1) // hbm_pe
+
+    pes_per_replica = max(pes_a, pes_b)
+    total = pes_per_replica * n_replicas
+
+    if n_replicas > 1:
+        binding = "throughput"
+    elif pes_b > pes_a:
+        binding = "kv"
+    else:
+        binding = "capacity"
+
+    return GpuSizingResult(
+        n_users=n_users,
+        avg_ctx_tokens=avg_ctx,
+        tpot_slo_s=tpot_slo_s,
+        pes_axis_a_capacity=int(pes_a),
+        pes_axis_b_kv=int(pes_b),
+        pes_per_replica=int(pes_per_replica),
+        b_at_slo=b_at_slo,
+        users_per_replica=int(users_per_replica),
+        n_replicas=int(n_replicas),
+        total_pes=int(total),
+        binding_axis=binding,
+    )
+
+
+def ai_sensitivity_curve(machine: MachineParams, model: ModelConfig,
+                          multipliers: list[float],
+                          axis: str = "flops") -> list[AISensitivityPoint]:
+    """Sweep FLOPs OR BW while holding the other fixed. Returns AI + B*
+    at each multiplier.
+
+    axis='flops' → scales peak_tflops_f16 (AI grows linearly).
+    axis='bw'    → scales bw_hbm_gbs (AI shrinks inversely).
+    """
+    if axis not in ("flops", "bw"):
+        raise ValueError(f"axis must be 'flops' or 'bw', got {axis!r}")
+    base_flops = machine.peak_tflops_f16
+    base_bw = machine.bw_hbm_gbs
+    points: list[AISensitivityPoint] = []
+    for k in multipliers:
+        if axis == "flops":
+            m = replace(machine, peak_tflops_f16=base_flops * k)
+        else:
+            m = replace(machine, bw_hbm_gbs=base_bw * k)
+        points.append(AISensitivityPoint(
+            multiplier=k,
+            peak_tflops=m.peak_tflops_f16,
+            bw_gbs=m.bw_hbm_gbs,
+            ai=arithmetic_intensity(m),
+            b_star=critical_batch(m, model),
+        ))
+    return points

tests/analytical_visualization/memory_layout.py

@@ -30,12 +30,19 @@ class MemoryBreakdown:
         return self.used_bytes > self.budget_bytes
 
 
-def per_pe_weight_bytes(cfg: FullConfig) -> int:
+def per_pe_weight_bytes(cfg: FullConfig,
+                        include_attention: bool = True,
+                        include_ffn: bool = True) -> int:
     """Attention + FFN weights per PE, bf16. Divided by TP, PP, and EP.
 
     EP divides the FFN (experts) across ranks; attention weights are
     unaffected. When kv_shard_mode='replicate' and TP > H_kv, each KV
     head is replicated (per-PE W_K/W_V size doesn't drop below one head).
+
+    Scope flags let callers count only attention or only FFN weights —
+    useful for "smallest deployment for just this block" sizing. Both
+    True (default) matches the traditional full-model per-PE weight
+    footprint.
     """
     m = cfg.model
     tp = cfg.topo.tp
@@ -50,18 +57,23 @@ def per_pe_weight_bytes(cfg: FullConfig) -> int:
         # Head-dim split: fractional head allowed (< 1 head bytes when TP > H_kv).
         hkv_per_pe_bytes = m.h_kv / tp
 
+    per_layer_attn = 0
+    if include_attention:
         per_layer_attn = (
             m.hidden * hq_per_pe * m.d_head +           # W_Q
             m.hidden * hkv_per_pe_bytes * m.d_head +    # W_K
             m.hidden * hkv_per_pe_bytes * m.d_head +    # W_V
             hq_per_pe * m.d_head * m.hidden             # W_O
         )
+
+    per_layer_ffn = 0
+    if include_ffn:
         # FFN divisor: TP (default) or TP*CP or TP*CP*DP if user opts in.
         # EP further divides (MoE experts).
         ffn_div = cfg.ffn_shard_divisor * ep
         per_layer_ffn = 3 * m.hidden * (m.ffn_dim // ffn_div)
-    per_layer = (per_layer_attn + per_layer_ffn) * m.bytes_per_elem
 
+    per_layer = (per_layer_attn + per_layer_ffn) * m.bytes_per_elem
     layers_per_stage = (m.layers + pp - 1) // pp
     return int(per_layer * layers_per_stage)
 
@@ -74,17 +86,19 @@ def per_pe_kv_cache_bytes(cfg: FullConfig) -> int:
       TP > H_kv, each KV head is duplicated across TP/H_kv ranks so
       per-PE storage doesn't shrink below 1 head.
     - PP: only layers_per_stage layers stored per PE.
+    - Batch (B): each concurrent request keeps its own KV cache slice.
     """
     m = cfg.model
     pp = cfg.topo.pp
     tp = cfg.topo.tp
+    B = max(1, cfg.topo.b)
     if cfg.topo.kv_shard_mode == "replicate":
         hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
     else:
         hkv_per_pe_bytes = m.h_kv / tp
     layers_per_stage = (m.layers + pp - 1) // pp
     per_layer = 2 * cfg.topo.s_local * hkv_per_pe_bytes * m.d_head * m.bytes_per_elem
-    return int(per_layer * layers_per_stage)
+    return int(per_layer * layers_per_stage * B)
 
 
 def per_pe_transient_bytes(cfg: FullConfig) -> int:
@@ -101,10 +115,19 @@ def per_pe_transient_bytes(cfg: FullConfig) -> int:
         return int(2 * tile_score + T_q * m.hidden * m.bytes_per_elem)
 
 
-def compute_memory(cfg: FullConfig) -> MemoryBreakdown:
+def compute_memory(cfg: FullConfig,
+                   include_attention: bool = True,
+                   include_ffn: bool = True) -> MemoryBreakdown:
+    """Per-PE memory breakdown. Scope flags let callers count only the
+    attention or only the FFN block (KV cache goes with the attention
+    block; transient activation buffer is small either way and left in).
+    """
+    kv_bytes = per_pe_kv_cache_bytes(cfg) if include_attention else 0
     return MemoryBreakdown(
-        weights_bytes=per_pe_weight_bytes(cfg),
+        weights_bytes=per_pe_weight_bytes(
-        kv_cache_bytes=per_pe_kv_cache_bytes(cfg),
+            cfg, include_attention=include_attention, include_ffn=include_ffn,
+        ),
+        kv_cache_bytes=kv_bytes,
         transient_bytes=per_pe_transient_bytes(cfg),
         budget_bytes=cfg.machine.pe_budget_bytes,
     )

tests/analytical_visualization/model_config.py

@@ -43,6 +43,7 @@ class TopologyConfig:
     pp: int = 1                 # pipeline stages (layer shard)
     dp: int = 1                 # data-parallel replicas
     ep: int = 1                 # expert-parallel degree (MoE only)
+    b: int = 1                  # batch size (concurrent requests per PP stage)
     s_kv: int = 4096
     mode: str = "decode"        # "decode" or "prefill"
     kv_shard_mode: str = "split"  # "split" or "replicate" (used when TP > H_kv)

tests/analytical_visualization/pe_weight_layout.py

@@ -114,6 +114,14 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
     cube_h = 6.0
     cube_gap = 0.6
 
+    # When TP fills the cube (TP>=8), cp_placement is "cube" and CP
+    # ranks span cubes — the intra-cube CP-color branch below is
+    # skipped. Surface which CP rank this figure is showing so users
+    # aren't guessing which of `CP` groups is drawn.
+    _cube_place_tag = ""
+    if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1:
+        _cube_place_tag = f"  [CP rank 0 of {cfg.topo.cp}]"
+
     for cube_idx in range(n_cubes):
         x0 = cube_idx * (cube_w + cube_gap)
         y0 = 0
@@ -125,7 +133,8 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
         )
         ax.add_patch(rect)
         ax.text(x0 + cube_w / 2, y0 + cube_h + 0.1,
-                f"Cube {cube_idx}  (PEs {cube_idx*8}-{cube_idx*8+7})",
+                f"Cube {cube_idx}  (PEs {cube_idx*8}-{cube_idx*8+7})"
+                + _cube_place_tag,
                 ha="center", va="bottom", fontsize=11, fontweight="bold")
 
         pe_pad_x = 0.15
@@ -136,6 +145,26 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
         pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
         pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
 
+        # When cp_placement=pe, CP ranks are packed intra-cube along with
+        # TP ranks. Color-code by CP rank so it's obvious which PEs belong
+        # to which sequence-shard group. Palette wraps beyond 8 CP ranks.
+        _cp_place = cfg.topo.cp_placement
+        _cp_val = cfg.topo.cp
+        _cp_palette = [
+            "#f0f7ff",  # baseline blue (also the fallback for cp_placement=cube)
+            "#e8f5e9",  # light green
+            "#fff3e0",  # light orange
+            "#f3e5f5",  # light purple
+            "#ffebee",  # light red
+            "#e0f7fa",  # light cyan
+            "#fce4ec",  # light pink
+            "#f9fbe7",  # light lime
+        ]
+        _cp_edge_palette = [
+            "#3a86ff", "#2e7d32", "#ef6c00", "#7b1fa2",
+            "#c62828", "#0097a7", "#c2185b", "#9e9d24",
+        ]
+
         for pr in range(PE_ROWS):
             for pc in range(PE_COLS):
                 pe_local = pr * PE_COLS + pc
@@ -145,17 +174,43 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
                 px = x0 + pe_pad_x + pc * (pe_w + pe_gap)
                 py = y0 + pe_pad_y + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
 
+                # Which (cp_rank, tp_rank) does this PE hold?
+                # When cp_placement=pe: cp_rank = pe_id // tp, tp_rank = pe_id % tp.
+                # When cp_placement=cube: every PE in this cube is one CP rank
+                # (drawn on its own tile), tp_rank = pe_id % tp.
+                if _cp_place == "pe" and _cp_val > 1:
+                    _cp_rank = pe_id_in_group // tp
+                    _tp_rank = pe_id_in_group % tp
+                    _fc = _cp_palette[_cp_rank % len(_cp_palette)]
+                    _ec = _cp_edge_palette[_cp_rank % len(_cp_edge_palette)]
+                    _cp_label = f"  |  CP={_cp_rank}"
+                else:
+                    _cp_rank = None
+                    _tp_rank = pe_id_in_group % tp
+                    _fc = "#f0f7ff"
+                    _ec = "#3a86ff"
+                    # cp_placement=="cube" with CP>1: the whole cube is
+                    # one CP rank (rank 0 by convention — see cube title
+                    # tag above). Add the locator so users know what
+                    # rank this per-PE data belongs to.
+                    if _cp_place == "cube" and _cp_val > 1:
+                        _cp_label = "  |  CP=0"
+                    else:
+                        _cp_label = ""
+
                 pe_rect = patches.Rectangle(
                     (px, py), pe_w, pe_h,
-                    facecolor="#f0f7ff", edgecolor="#3a86ff", linewidth=1.0,
+                    facecolor=_fc, edgecolor=_ec, linewidth=1.2,
                 )
                 ax.add_patch(pe_rect)
 
-                q_heads = _q_heads_for_pe(pe_id_in_group, tp, cfg.model.h_q)
+                # Head assignment uses TP rank (not raw pe_id) so that under
+                # cp_placement=pe, all CP ranks share the same head split.
+                q_heads = _q_heads_for_pe(_tp_rank, tp, cfg.model.h_q)
                 kv_heads, kv_note = _kv_heads_for_pe(
-                    pe_id_in_group, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
+                    _tp_rank, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
                 )
-                bytes_ = _per_pe_bytes(cfg, pe_id_in_group)
+                bytes_ = _per_pe_bytes(cfg, _tp_rank)
                 weights_gb = sum(v for k, v in bytes_.items()
                                  if k not in ("KV cache", "Transient")) / 1e9
                 kv_gb = bytes_["KV cache"] / 1e9
@@ -172,7 +227,7 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
                 ffn_mb = (bytes_["W_gate"] + bytes_["W_up"]
                           + bytes_["W_down"]) / 1e6
                 # Header (PE id and heads) bigger, then per-tensor breakdown
-                header = (f"PE {pe_id_in_group}\n"
+                header = (f"PE {pe_id_in_group}  TP={_tp_rank}{_cp_label}\n"
                           f"Q heads: {q_str}\n"
                           f"KV head: {kv_str}")
                 ax.text(px + 0.08, py + pe_h - 0.05, header,
@@ -211,9 +266,12 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
     ax.set_aspect("equal")
     ax.axis("off")
 
+    _title_cp = f", CP={cfg.topo.cp}"
+    if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1:
+        _title_cp += " (showing 1 of CP groups; others identical)"
     title = (
         f"Per-PE layout of one CP group  "
-        f"(TP={tp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})"
+        f"(TP={tp}{_title_cp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})"
         f" | KV mode: {cfg.topo.kv_shard_mode}"
     )
     ax.set_title(title, fontsize=11, fontweight="bold")

tests/analytical_visualization/stage_latencies.py

@@ -50,19 +50,22 @@ def stage_rmsnorm(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    bytes_ = T_q * d * b * 2  # load x + load weight
+    B = max(1, cfg.topo.b)
-    flops = 4 * T_q * d
+    # Activation memory (x) scales with B; weight (once) is fixed.
+    bytes_ = B * T_q * d * b + T_q * d * b
+    flops = 4 * B * T_q * d
     mem_s = bytes_ / cfg.machine.bw_hbm
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name="S1 RMSNorm",
-        formula=f"bytes = {T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM",
+        formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
         bound=bnd, visible_s=vis,
         flops=flops, mem_bytes=bytes_,
-        flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}",
+        flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}",
-        mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B",
+        mem_formula=f"B*T_q*d*b + T_q*d*b (weight) "
+                     f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B",
     )
 
 
@@ -76,21 +79,26 @@ def stage_wq(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     hq_per_pe = cfg.h_q_per_pe
     dh = cfg.model.d_head
-    flops = 2 * T_q * d * (hq_per_pe * dh)
+    # FLOPs scale with batch; weight bytes fixed (shared across batch).
+    flops = 2 * B * T_q * d * (hq_per_pe * dh)
     weight_B = d * (hq_per_pe * dh) * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name="S2 W_Q GEMM",
-        formula=f"FLOPs = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; "
+        formula=f"FLOPs = 2*B*T_q*d*(H_q/TP*d_h) "
-                f"weight = {weight_B/1e6:.1f} MB",
+                f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
         bound=bnd, visible_s=vis,
         flops=int(flops), mem_bytes=int(weight_B),
-        flops_formula=f"2*T_q*d*(H_q/TP*d_h) = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}",
+        flops_formula=f"2*B*T_q*d*(H_q/TP*d_h) "
-        mem_formula=f"d*(H_q/TP*d_h)*b = {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB",
+                       f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}",
+        mem_formula=f"d*(H_q/TP*d_h)*b (weight, B-invariant) "
+                     f"= {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB",
     )
 
 
@@ -98,9 +106,10 @@ def stage_wkv(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
     dh = cfg.model.d_head
-    flops_one = 2 * T_q * d * (hkv_per_pe * dh)
+    flops_one = 2 * B * T_q * d * (hkv_per_pe * dh)
     weight_B_one = d * (hkv_per_pe * dh) * b
     flops = 2 * flops_one
     weight_B = 2 * weight_B_one
@@ -108,30 +117,36 @@ def stage_wkv(cfg: FullConfig) -> StageCost:
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name="S3 W_K + W_V GEMM",
-        formula=f"FLOPs per proj = 2*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2",
+        formula=f"FLOPs per proj = 2*B*T_q*d*(H_kv/TP*d_h) "
+                f"= 2*{B}*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
         bound=bnd, visible_s=vis,
         flops=int(flops), mem_bytes=int(weight_B),
-        flops_formula=f"2*(2*T_q*d*(H_kv/TP*d_h)) = 2*(2*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}",
+        flops_formula=f"2*(2*B*T_q*d*(H_kv/TP*d_h)) "
-        mem_formula=f"2*(d*(H_kv/TP*d_h)*b) = 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB",
+                       f"= 2*(2*{B}*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}",
+        mem_formula=f"2*(d*(H_kv/TP*d_h)*b) (weights, B-invariant) "
+                     f"= 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB",
     )
 
 
 def stage_kv_append(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
     dh = cfg.model.d_head
-    bytes_ = 2 * T_q * hkv_per_pe * dh * b
+    bytes_ = 2 * B * T_q * hkv_per_pe * dh * b
     mem_s = bytes_ / cfg.machine.bw_hbm
     return StageCost(
         name="S4 KV cache append",
-        formula=f"bytes = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
+        formula=f"bytes = 2*B*T_q*(H_kv/TP)*d_h*b "
+                f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
         compute_s=0, memory_s=mem_s, comm_s=0,
         bound="memory", visible_s=mem_s,
         flops=0, mem_bytes=bytes_,
         flops_formula="0 (no matmul, just cache write)",
-        mem_formula=f"2*T_q*(H_kv/TP)*d_h*b = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
+        mem_formula=f"2*B*T_q*(H_kv/TP)*d_h*b "
+                     f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
     )
 
 
@@ -141,7 +156,8 @@ def _per_hop_qkT_pv(cfg: FullConfig) -> tuple[float, str]:
     S_local = cfg.topo.s_local
     dh = cfg.model.d_head
     hq_per_pe = cfg.h_q_per_pe
-    flops = 2 * T_q * S_local * dh * hq_per_pe
+    B = max(1, cfg.topo.b)
+    flops = 2 * B * T_q * S_local * dh * hq_per_pe
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
     formula = f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops:.2g} FLOPs/hop"
     return cmp_s, formula
@@ -165,19 +181,23 @@ def stage_qkT(cfg: FullConfig) -> StageCost:
     S_local = cfg.topo.s_local
     dh = cfg.model.d_head
     hq_per_pe = cfg.h_q_per_pe
-    flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe
+    B = max(1, cfg.topo.b)
+    flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe
     total_flops = flops_per_hop * passes
     _hop_word = "pass" if passes == 1 else "hops"
     return StageCost(
         name=f"S5 Q.K^T (x{passes} {_hop_word})",
-        formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop",
+        formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = "
+                f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = "
+                f"{flops_per_hop:.2g} FLOPs/hop",
         compute_s=total, memory_s=0, comm_s=0,
         bound="compute", visible_s=total,
         hop_multiplier=passes,
         flops=int(total_flops), mem_bytes=0,
         flops_formula=(
-            f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = "
+            f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = "
-            f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}"
+            f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = "
+            f"{total_flops:.3g}"
         ),
         mem_formula="0 (scores accumulated on-chip)",
     )
@@ -188,7 +208,8 @@ def stage_softmax(cfg: FullConfig) -> StageCost:
     S_local = cfg.topo.s_local
     hq_per_pe = cfg.h_q_per_pe
     b = cfg.model.bytes_per_elem
-    elems = hq_per_pe * T_q * S_local
+    B = max(1, cfg.topo.b)
+    elems = B * hq_per_pe * T_q * S_local
     bytes_ = elems * b * 2
     mem_s_per_hop = bytes_ / cfg.machine.bw_hbm
     passes = _cp_compute_passes(cfg)
@@ -197,15 +218,16 @@ def stage_softmax(cfg: FullConfig) -> StageCost:
     _hop_word = "pass" if passes == 1 else "hops"
     return StageCost(
         name=f"S6 softmax (x{passes} {_hop_word})",
-        formula=f"elts/hop = {hq_per_pe}*{T_q}*{S_local} = {elems:.2g}",
+        formula=f"elts/hop = B*(H_q/TP)*T_q*S_local "
+                f"= {B}*{hq_per_pe}*{T_q}*{S_local} = {elems:.2g}",
         compute_s=0, memory_s=total, comm_s=0,
         bound="memory", visible_s=total,
         hop_multiplier=passes,
         flops=0, mem_bytes=int(total_bytes),
         flops_formula="~O(elts) (negligible)",
         mem_formula=(
-            f"{passes}*2*b*(H_q/TP)*T_q*S_local = "
+            f"{passes}*2*b*B*(H_q/TP)*T_q*S_local = "
-            f"{passes}*2*{b}*{hq_per_pe}*{T_q}*{S_local} = "
+            f"{passes}*2*{b}*{B}*{hq_per_pe}*{T_q}*{S_local} = "
             f"{total_bytes/1e6:.2f} MB"
         ),
     )
@@ -219,19 +241,23 @@ def stage_pv(cfg: FullConfig) -> StageCost:
     S_local = cfg.topo.s_local
     dh = cfg.model.d_head
     hq_per_pe = cfg.h_q_per_pe
-    flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe
+    B = max(1, cfg.topo.b)
+    flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe
     total_flops = flops_per_hop * passes
     _hop_word = "pass" if passes == 1 else "hops"
     return StageCost(
         name=f"S7 P.V (x{passes} {_hop_word})",
-        formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop",
+        formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = "
+                f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = "
+                f"{flops_per_hop:.2g} FLOPs/hop",
         compute_s=total, memory_s=0, comm_s=0,
         bound="compute", visible_s=total,
         hop_multiplier=passes,
         flops=int(total_flops), mem_bytes=0,
         flops_formula=(
-            f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = "
+            f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = "
-            f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}"
+            f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = "
+            f"{total_flops:.3g}"
         ),
         mem_formula="0 (accumulated on-chip)",
     )
@@ -247,7 +273,8 @@ def stage_merge(cfg: FullConfig) -> StageCost:
     hq_per_pe = cfg.h_q_per_pe
     dh = cfg.model.d_head
     b = cfg.model.bytes_per_elem
-    flops = 6 * T_q * hq_per_pe * dh * max(0, cfg.topo.cp - 1)
+    B = max(1, cfg.topo.b)
+    flops = 6 * B * T_q * hq_per_pe * dh * max(0, cfg.topo.cp - 1)
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
 
     comm_s = 0.0
@@ -257,8 +284,8 @@ def stage_merge(cfg: FullConfig) -> StageCost:
 
     if cfg.topo.mode == "decode" and cfg.topo.cp > 1:
         cp = cfg.topo.cp
-        # (O + m + l) bytes per rank
+        # (O + m + l) bytes per rank — scales with B (per-request partials).
-        M = (T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+        M = B * ((T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
         if cfg.topo.cp_placement == "pe":
             bw, alpha, tier = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
         elif cfg.topo.sips_used > 1:
@@ -293,13 +320,15 @@ def stage_merge(cfg: FullConfig) -> StageCost:
 
     return StageCost(
         name=f"S8 online-softmax merge (x{max(0, cfg.topo.cp-1)}){name_suffix}",
-        formula=f"~6*{T_q}*{hq_per_pe}*{dh}*(C-1) = {flops:.2g} FLOPs",
+        formula=f"~6*B*T_q*(H_q/TP)*d_h*(C-1) "
+                f"= 6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} "
+                f"= {flops:.2g} FLOPs",
         compute_s=cmp_s, memory_s=0, comm_s=comm_s,
         bound=bound, visible_s=visible_s,
         flops=int(flops), mem_bytes=0, comm_bytes=comm_bytes,
         flops_formula=(
-            f"6*T_q*(H_q/TP)*d_h*(CP-1) = "
+            f"6*B*T_q*(H_q/TP)*d_h*(CP-1) = "
-            f"6*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}"
+            f"6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}"
         ),
         mem_formula="0 (in-register)",
         comm_formula=comm_formula or "0",
@@ -310,15 +339,18 @@ def stage_normalize(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     hq_per_pe = cfg.h_q_per_pe
     dh = cfg.model.d_head
-    flops = T_q * hq_per_pe * dh
+    B = max(1, cfg.topo.b)
+    flops = B * T_q * hq_per_pe * dh
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
     return StageCost(
         name="S9 normalize O/l",
-        formula=f"{T_q}*{hq_per_pe}*{dh} = {flops} divisions",
+        formula=f"B*T_q*(H_q/TP)*d_h = {B}*{T_q}*{hq_per_pe}*{dh} "
+                f"= {flops} divisions",
         compute_s=cmp_s, memory_s=0, comm_s=0,
         bound="trivial", visible_s=cmp_s,
         flops=int(flops), mem_bytes=0,
-        flops_formula=f"T_q*(H_q/TP)*d_h = {T_q}*{hq_per_pe}*{dh} = {flops}",
+        flops_formula=f"B*T_q*(H_q/TP)*d_h "
+                       f"= {B}*{T_q}*{hq_per_pe}*{dh} = {flops}",
         mem_formula="0",
     )
 
@@ -327,21 +359,25 @@ def stage_wo(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     hq_per_pe = cfg.h_q_per_pe
     dh = cfg.model.d_head
-    flops = 2 * T_q * (hq_per_pe * dh) * d
+    flops = 2 * B * T_q * (hq_per_pe * dh) * d
     weight_B = (hq_per_pe * dh) * d * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name="S10 W_O GEMM",
-        formula=f"FLOPs = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; "
+        formula=f"FLOPs = 2*B*T_q*(H_q/TP*d_h)*d "
-                f"weight = {weight_B/1e6:.1f} MB",
+                f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
         bound=bnd, visible_s=vis,
         flops=int(flops), mem_bytes=int(weight_B),
-        flops_formula=f"2*T_q*(H_q/TP*d_h)*d = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}",
+        flops_formula=f"2*B*T_q*(H_q/TP*d_h)*d "
-        mem_formula=f"(H_q/TP*d_h)*d*b = {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB",
+                       f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}",
+        mem_formula=f"(H_q/TP*d_h)*d*b (weight, B-invariant) "
+                     f"= {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB",
     )
 
 
@@ -369,12 +405,13 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     dh = cfg.model.d_head
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
 
     # ── Decode: single O/m/l all-reduce after local compute ──
     if cfg.topo.mode == "decode":
         cp = cfg.topo.cp
-        # per-rank O + m + l bytes
+        # per-rank O + m + l bytes — scales with B (per-request partials).
-        M = (T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+        M = B * ((T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
         if cfg.topo.cp_placement == "pe":
             bw = cfg.machine.bw_intra
             alpha = cfg.machine.alpha_intra
@@ -392,7 +429,7 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost:
         return StageCost(
             name=f"C1 O/m/l all-reduce ({tier}, x{cp} ranks)",
             formula=(f"decode: gather partial (O,m,l) once at end; "
-                     f"M = T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; "
+                     f"M = B*T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; "
                      f"AR: 2*(CP-1)/CP*M / BW + 2*(CP-1)*alpha"),
             compute_s=0, memory_s=0, comm_s=ar_time,
             bound="comm", visible_s=ar_time,
@@ -417,19 +454,19 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost:
 
     # ── Prefill: per-hop ring (K/V or Q+O/m/l) concurrent with S5-S7 ──
     if cfg.topo.cp_ring_variant == "qoml":
-        M_KV = (2 * T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+        M_KV = B * ((2 * T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
         variant_desc = "Q+O/m/l ring"
         formula_bytes = (
-            f"2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b "
+            f"B*(2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b) "
-            f"= 2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b} "
+            f"= {B}*(2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b}) "
             f"= {M_KV} B/hop"
         )
     else:
-        M_KV = 2 * S_local * hkv_per_pe * dh * b
+        M_KV = 2 * B * S_local * hkv_per_pe * dh * b
         variant_desc = "K/V ring"
         formula_bytes = (
-            f"2*S_local*(H_kv/TP)*d_h*b "
+            f"2*B*S_local*(H_kv/TP)*d_h*b "
-            f"= 2*{S_local}*{hkv_per_pe}*{dh}*{b} "
+            f"= 2*{B}*{S_local}*{hkv_per_pe}*{dh}*{b} "
             f"= {M_KV/1e6:.3f} MB/hop"
         )
 
@@ -501,7 +538,8 @@ def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    bytes_ = T_q * d * b
+    B = max(1, cfg.topo.b)
+    bytes_ = B * T_q * d * b
     tp = cfg.topo.tp
     tier = cfg.topo.tp_link_tier()  # "intra" | "inter" | "intersip"
     if tier == "intra":
@@ -514,7 +552,8 @@ def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
     total_comm_bytes = int(2 * (tp - 1) / tp * bytes_)
     return StageCost(
         name=f"C2 TP AllReduce W_O (TP={tp} ranks, {scope})",
-        formula=f"2*(TP-1)/TP * {T_q}*{d}*{b} B / BW + 2(TP-1)*alpha [{scope}]",
+        formula=f"2*(TP-1)/TP * B*T_q*d*b B / BW + 2(TP-1)*alpha "
+                f"[B={B}, {scope}]",
         compute_s=0, memory_s=0, comm_s=comm_time,
         bound="comm", visible_s=comm_time,
         flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
@@ -528,7 +567,8 @@ def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
             f"every rank ends up with the full hidden vector for the next\n"
             f"stage (RMSNorm -> FFN). Fires ONCE per layer over {scope} links.\n"
             f"---\n"
-            f"2*(TP-1)/TP * T_q*d*b = 2*({tp}-1)/{tp} * {T_q}*{d}*{b} "
+            f"2*(TP-1)/TP * B*T_q*d*b "
+            f"= 2*({tp}-1)/{tp} * {B}*{T_q}*{d}*{b} "
             f"= {total_comm_bytes/1e6:.2f} MB over {scope} at "
             f"{bw/1e9:.0f} GB/s"
         ),
@@ -552,8 +592,9 @@ def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost:
     S_local = cfg.topo.s_local
     hq_per_pe = cfg.h_q_per_pe
     b = cfg.model.bytes_per_elem
-    # score bytes per rank per hop
+    B = max(1, cfg.topo.b)
-    bytes_per_hop = hq_per_pe * T_q * S_local * b
+    # score bytes per rank per hop — scales with B (per-request scores).
+    bytes_per_hop = B * hq_per_pe * T_q * S_local * b
     # split-group AllReduce (intra-cube assumed if group fits)
     bw = cfg.machine.bw_intra
     alpha = cfg.machine.alpha_intra
@@ -562,7 +603,9 @@ def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost:
     total_comm_bytes = int(2 * (split - 1) / split * bytes_per_hop * cfg.topo.cp)
     return StageCost(
         name=f"C3 Score AllReduce ({split}-way, xCP hops)",
-        formula=f"2*({split}-1)/{split} * {hq_per_pe}*{T_q}*{S_local}*{b} / BW + latency",
+        formula=f"2*({split}-1)/{split} * B*(H_q/TP)*T_q*S_local*b / BW "
+                f"= 2*({split}-1)/{split} * {B}*{hq_per_pe}*{T_q}*{S_local}*{b} "
+                f"/ BW + latency",
         compute_s=0, memory_s=0, comm_s=total,
         bound="comm", visible_s=total,
         flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
@@ -626,17 +669,20 @@ def stage_ffn_rmsnorm(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    bytes_ = T_q * d * b * 2
+    B = max(1, cfg.topo.b)
-    flops = 4 * T_q * d
+    # Activation memory scales with B; weight (once) is fixed.
+    bytes_ = B * T_q * d * b + T_q * d * b
+    flops = 4 * B * T_q * d
     mem_s = bytes_ / cfg.machine.bw_hbm
     return StageCost(
         name="F1 RMSNorm (pre-FFN)",
-        formula=f"{T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM",
+        formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM",
         compute_s=0, memory_s=mem_s, comm_s=0,
         bound="memory", visible_s=mem_s,
         flops=flops, mem_bytes=bytes_,
-        flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}",
+        flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}",
-        mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B",
+        mem_formula=f"B*T_q*d*b + T_q*d*b (weight) "
+                     f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B",
     )
 
 
@@ -644,18 +690,24 @@ def _ffn_gemm(cfg: FullConfig, name: str, ffn_per_pe: int) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    flops = 2 * T_q * d * ffn_per_pe
+    B = max(1, cfg.topo.b)
+    # FLOPs scale with batch; weight (shared) is fixed.
+    flops = 2 * B * T_q * d * ffn_per_pe
     weight_B = d * ffn_per_pe * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name=name,
-        formula=f"FLOPs = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB",
+        formula=f"FLOPs = 2*B*T_q*d*(ffn/div) "
+                f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0,
         bound=bnd, visible_s=vis,
         flops=int(flops), mem_bytes=int(weight_B),
-        flops_formula=f"2*T_q*d*(ffn/div) = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}",
+        flops_formula=f"2*B*T_q*d*(ffn/div) "
-        mem_formula=f"d*(ffn/div)*b = {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB",
+                       f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}",
+        mem_formula=f"d*(ffn/div)*b (weight, B-invariant) "
+                     f"= {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB",
     )
 
 
@@ -674,17 +726,20 @@ def stage_ffn_swiglu(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
     b = cfg.model.bytes_per_elem
-    bytes_ = 3 * T_q * ffn_per_pe * b
+    B = max(1, cfg.topo.b)
-    flops = 3 * T_q * ffn_per_pe  # gate * silu(up) approximated as 3 flops/elt
+    bytes_ = 3 * B * T_q * ffn_per_pe * b
+    flops = 3 * B * T_q * ffn_per_pe  # gate * silu(up) approximated as 3 flops/elt
     mem_s = bytes_ / cfg.machine.bw_hbm
     return StageCost(
         name="F4 SwiGLU act",
-        formula=f"3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM",
+        formula=f"3*B*T_q*(ffn/div)*b "
+                f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM",
         compute_s=0, memory_s=mem_s, comm_s=0,
         bound="memory", visible_s=mem_s,
         flops=flops, mem_bytes=bytes_,
-        flops_formula=f"~3*T_q*(ffn/div) = 3*{T_q}*{ffn_per_pe} = {flops}",
+        flops_formula=f"~3*B*T_q*(ffn/div) = 3*{B}*{T_q}*{ffn_per_pe} = {flops}",
-        mem_formula=f"3*T_q*(ffn/div)*b = 3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B",
+        mem_formula=f"3*B*T_q*(ffn/div)*b "
+                     f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B",
     )
 
 
@@ -693,19 +748,24 @@ def stage_ffn_down(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
-    flops = 2 * T_q * ffn_per_pe * d
+    flops = 2 * B * T_q * ffn_per_pe * d
     weight_B = ffn_per_pe * d * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
     return StageCost(
         name="F5 W_down GEMM",
-        formula=f"FLOPs = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB",
+        formula=f"FLOPs = 2*B*T_q*(ffn/div)*d "
+                f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
         compute_s=cmp_s, memory_s=mem_s, comm_s=0,
         bound=bnd, visible_s=vis,
         flops=int(flops), mem_bytes=int(weight_B),
-        flops_formula=f"2*T_q*(ffn/div)*d = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}",
+        flops_formula=f"2*B*T_q*(ffn/div)*d "
-        mem_formula=f"(ffn/div)*d*b = {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB",
+                       f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}",
+        mem_formula=f"(ffn/div)*d*b (weight, B-invariant) "
+                     f"= {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB",
     )
 
 
@@ -723,7 +783,9 @@ def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    bytes_ = T_q * d * b
+    B = max(1, cfg.topo.b)
+    # AR reduces the batched FFN output — bytes scale with B.
+    bytes_ = B * T_q * d * b
     # Choose BW/alpha tier based on scope (rough): TP=intra-cube (or inter-cube),
     # +CP=inter-SIP possible, +DP=inter-SIP always.
     if "DP" in scope or "CP" in scope:
@@ -736,7 +798,8 @@ def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
     total_comm_bytes = int(2 * (divisor - 1) / divisor * bytes_)
     return StageCost(
         name=f"CF1 FFN AllReduce (scope={scope}, x{divisor})",
-        formula=f"2*({divisor}-1)/{divisor} * {T_q}*{d}*{b} B / BW + latency",
+        formula=f"2*({divisor}-1)/{divisor} * B*T_q*d*b B / BW + latency "
+                f"[B={B}]",
         compute_s=0, memory_s=0, comm_s=comm_time,
         bound="comm", visible_s=comm_time,
         flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
@@ -749,8 +812,8 @@ def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
             f"rank has the full FFN output for the next layer. Larger scope\n"
             f"= less FFN weight memory per PE but bigger AR bill.\n"
             f"---\n"
-            f"2*(div-1)/div * T_q*d*b = 2*({divisor}-1)/{divisor} * "
+            f"2*(div-1)/div * B*T_q*d*b = 2*({divisor}-1)/{divisor} * "
-            f"{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at "
+            f"{B}*{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at "
             f"{bw/1e9:.0f} GB/s (scope={scope})"
         ),
     )

tests/analytical_visualization/stage_shapes.py

@@ -0,0 +1,195 @@
+"""Per-stage tensor shape rows for the analytical visualization.
+
+Complements stage_latencies.py: while that module reports FLOPs/bytes/
+time per stage, this one reports the INPUT / WEIGHT / OUTPUT tensor
+shapes each stage operates on, per PE. Used for the 'per-stage shape'
+tables in the Streamlit app.
+
+Shape strings substitute the current cfg's numeric values so the table
+reads like an inspection of the actual deployment (no symbolic-only
+form).
+"""
+from __future__ import annotations
+
+from .model_config import FullConfig
+
+
+def _s(shape: tuple) -> str:
+    """Render a tuple of dim values as '(a, b, c)'."""
+    return "(" + ", ".join(str(x) for x in shape) + ")"
+
+
+def attn_stage_shape_rows(cfg: FullConfig) -> list[dict]:
+    """Per-PE attention shapes, one row per stage.
+
+    Rows follow the same S1..S10, C1..C3 order the latency table uses,
+    with the same conditional insertions (C1 only in prefill+CP>1, C3
+    only when kv_shard_mode='split' and TP > H_kv, C2 only when TP>1,
+    S8 merge only when CP>1).
+    """
+    m = cfg.model
+    t = cfg.topo
+    B = max(1, t.b)
+    T_q = t.T_q
+    S_local = t.s_local
+    d = m.hidden
+    dh = m.d_head
+    hq_pe = cfg.h_q_per_pe
+    hkv_pe = max(1, m.h_kv // t.tp)
+
+    rows: list[dict] = []
+
+    rows.append({
+        "Stage": "S1", "Op": "RMSNorm",
+        "Input (per PE)": _s((B, T_q, d)),
+        "Weight (per PE)": _s((d,)),
+        "Output (per PE)": _s((B, T_q, d)),
+    })
+    rows.append({
+        "Stage": "S2", "Op": "W_Q GEMM",
+        "Input (per PE)": _s((B, T_q, d)),
+        "Weight (per PE)": _s((d, hq_pe * dh)),
+        "Output (per PE)": _s((B, T_q, hq_pe * dh)),
+    })
+    rows.append({
+        "Stage": "S3", "Op": "W_K + W_V GEMM",
+        "Input (per PE)": _s((B, T_q, d)),
+        "Weight (per PE)": f"{_s((d, hkv_pe * dh))} x2 (K, V)",
+        "Output (per PE)": f"{_s((B, T_q, hkv_pe * dh))} x2 (K, V)",
+    })
+    rows.append({
+        "Stage": "S4", "Op": "KV cache append",
+        "Input (per PE)": f"{_s((B, T_q, hkv_pe * dh))} x2 (K, V)",
+        "Weight (per PE)": "-",
+        "Output (per PE)": (
+            f"KV cache: {_s((B, S_local, hkv_pe * dh))} x2 "
+            f"(extends by T_q={T_q} tokens)"
+        ),
+    })
+    rows.append({
+        "Stage": "S5", "Op": "Q · K^T",
+        "Input (per PE)": (
+            f"Q={_s((B, hq_pe, T_q, dh))}, "
+            f"K^T={_s((B, hkv_pe, dh, S_local))}"
+        ),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, S_local)),
+    })
+    rows.append({
+        "Stage": "S6", "Op": "softmax",
+        "Input (per PE)": _s((B, hq_pe, T_q, S_local)),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, S_local)),
+    })
+    rows.append({
+        "Stage": "S7", "Op": "P · V",
+        "Input (per PE)": (
+            f"P={_s((B, hq_pe, T_q, S_local))}, "
+            f"V={_s((B, hkv_pe, S_local, dh))}"
+        ),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, dh)),
+    })
+
+    if t.mode == "prefill" and t.cp > 1:
+        variant = t.cp_ring_variant
+        if variant == "kv":
+            payload = (
+                f"K,V shards rotate: {_s((B, S_local, hkv_pe, dh))} x2 per hop"
+            )
+        else:
+            payload = (
+                f"Q + running (O,m,l) rotate: {_s((B, hq_pe, T_q, dh))} per hop"
+            )
+        rows.append({
+            "Stage": "C1", "Op": f"CP ring ({variant})",
+            "Input (per PE)": payload,
+            "Weight (per PE)": "-",
+            "Output (per PE)": "(circulating; final owner reduces)",
+        })
+
+    if t.kv_shard_mode == "split" and t.tp > m.h_kv:
+        rows.append({
+            "Stage": "C3", "Op": "Score AllReduce (head-split)",
+            "Input (per PE)": _s((B, hq_pe, T_q, S_local)),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, hq_pe, T_q, S_local)),
+        })
+
+    if t.cp > 1:
+        rows.append({
+            "Stage": "S8", "Op": "online-softmax merge",
+            "Input (per PE)": (
+                f"O={_s((B, hq_pe, T_q, dh))}, "
+                f"m,l={_s((B, hq_pe, T_q))} each"
+            ),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, hq_pe, T_q, dh)),
+        })
+
+    rows.append({
+        "Stage": "S9", "Op": "normalize O / l",
+        "Input (per PE)": _s((B, hq_pe, T_q, dh)),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, dh)),
+    })
+    rows.append({
+        "Stage": "S10", "Op": "W_O GEMM",
+        "Input (per PE)": _s((B, T_q, hq_pe * dh)),
+        "Weight (per PE)": _s((hq_pe * dh, d)),
+        "Output (per PE)": _s((B, T_q, d)),
+    })
+
+    if t.tp > 1:
+        rows.append({
+            "Stage": "C2", "Op": f"TP AllReduce (W_O, TP={t.tp})",
+            "Input (per PE)": _s((B, T_q, d)),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, T_q, d)),
+        })
+
+    return rows
+
+
+def ffn_stage_shape_rows(cfg: FullConfig) -> list[dict]:
+    """Per-PE FFN shapes, one row per stage."""
+    m = cfg.model
+    t = cfg.topo
+    B = max(1, t.b)
+    T_q = t.T_q
+    d = m.hidden
+    divisor = cfg.ffn_shard_divisor * max(1, t.ep)
+    ffn_pe = max(1, m.ffn_dim // divisor)
+    scope = t.ffn_shard_scope
+
+    rows: list[dict] = [
+        {"Stage": "F1", "Op": "RMSNorm (pre-FFN)",
+         "Input (per PE)": _s((B, T_q, d)),
+         "Weight (per PE)": _s((d,)),
+         "Output (per PE)": _s((B, T_q, d))},
+        {"Stage": "F2", "Op": "W_gate GEMM",
+         "Input (per PE)": _s((B, T_q, d)),
+         "Weight (per PE)": _s((d, ffn_pe)),
+         "Output (per PE)": _s((B, T_q, ffn_pe))},
+        {"Stage": "F3", "Op": "W_up GEMM",
+         "Input (per PE)": _s((B, T_q, d)),
+         "Weight (per PE)": _s((d, ffn_pe)),
+         "Output (per PE)": _s((B, T_q, ffn_pe))},
+        {"Stage": "F4", "Op": "SwiGLU (gate * silu(up))",
+         "Input (per PE)": f"gate, up = {_s((B, T_q, ffn_pe))} each",
+         "Weight (per PE)": "-",
+         "Output (per PE)": _s((B, T_q, ffn_pe))},
+        {"Stage": "F5", "Op": "W_down GEMM",
+         "Input (per PE)": _s((B, T_q, ffn_pe)),
+         "Weight (per PE)": _s((ffn_pe, d)),
+         "Output (per PE)": _s((B, T_q, d))},
+    ]
+    if cfg.ffn_shard_divisor > 1:
+        rows.append({
+            "Stage": "CF1",
+            "Op": f"FFN AllReduce (scope={scope}, x{cfg.ffn_shard_divisor})",
+            "Input (per PE)": _s((B, T_q, d)),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, T_q, d)),
+        })
+    return rows

tests/analytical_visualization/tensor_sharding.py

@@ -279,12 +279,38 @@ def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0,
                   row_label=f"rows split by {ffn_txt}",
                   note="row-parallel")
 
-    # ── KV cache: (S_kv, H_kv*d_head), 2D shard ─
+    # ── KV cache: (B, S_kv, H_kv*d_head), 2D shard on S_kv x heads ─
+    # Batch (B) is a stacking dimension — each concurrent request holds
+    # its own KV slice per PE, so total KV per PE = B * (per-slice size).
     kv_row_splits = cp
     kv_col_splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
+    _B = max(1, cfg.topo.b)
+    _batch_suffix = f", B={_B}" if _B > 1 else ""
+    _batch_note = (f"; batch B={_B} => {_B}x KV bytes per PE"
+                    if _B > 1 else "")
+    # Draw B-1 offset "shadow" rectangles behind the KV cache to give
+    # a visual hint of the batch stacking dimension. Cap at 5 shadows
+    # so a large B doesn't overwhelm the diagram.
+    if _B > 1:
+        _kv_shadow_x = px(3)
+        _kv_shadow_y = py(0)
+        _n_shadows = min(_B - 1, 5)
+        _offset_step = 0.15
+        for _si in range(_n_shadows, 0, -1):
+            _dx = _si * _offset_step
+            _dy = _si * _offset_step
+            _shadow = patches.Rectangle(
+                (_kv_shadow_x + _dx, _kv_shadow_y + _dy),
+                cell_w, cell_h,
+                facecolor="#f5f5f5",
+                edgecolor="#adb5bd",
+                linewidth=0.8, linestyle="--", alpha=0.55,
+            )
+            ax.add_patch(_shadow)
     _draw_tensor(ax, px(3), py(0), cell_w, cell_h,
                   "KV cache (K and V)",
-                  f"({m.h_kv * m.d_head}, S_kv={cfg.topo.s_kv:,})",
+                  f"({m.h_kv * m.d_head}, S_kv={cfg.topo.s_kv:,}"
+                  f"{_batch_suffix})",
                   row_splits=kv_row_splits,
                   col_splits=kv_col_splits if show_physical else 1,
                   my_row_shard=kv_row_shard,
@@ -292,7 +318,8 @@ def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0,
                   row_label=f"rows: S split by {cp_txt}",
                   col_label=f"cols: heads split by {tp_txt}",
                   col_line_color=_TP_LINE_COLOR,
-                  note="2D shard: seq axis (CP) x head axis (TP)",
+                  note=f"2D shard: seq axis (CP) x head axis (TP)"
+                        f"{_batch_note}",
                   cell_annot=_kv_ann)
     # In non-physical mode, KV was drawn with col_splits=1; overlay TP col
     # splits as dotted lines to hint at the 2D nature. In physical mode
@@ -311,8 +338,11 @@ def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0,
     ax.set_aspect("auto")
     ax.axis("off")
 
+    _B_title = max(1, cfg.topo.b)
+    _b_title_suffix = f", B={_B_title}" if _B_title > 1 else ""
     ax.set_title(
-        f"Tensor sharding view (PE {my_pe})  |  TP={tp}, CP={cp}, "
+        f"Tensor sharding view (PE {my_pe})  |  TP={tp}, CP={cp}"
+        f"{_b_title_suffix}, "
         f"FFN scope={cfg.topo.ffn_shard_scope}  |  "
         f"blue = this PE's shard",
         fontsize=11, fontweight="bold",

tests/analytical_visualization/test_auto_explore.py

@@ -19,8 +19,9 @@ from tests.analytical_visualization.auto_explore import (
     run_auto_explore,
     score_config,
 )
+from tests.analytical_visualization.autosuggest import auto_suggest
 from tests.analytical_visualization.model_config import (
-    FullConfig, MachineParams,
+    FullConfig, MachineParams, TopologyConfig,
 )
 from tests.analytical_visualization.model_presets import PRESETS
 
@@ -208,6 +209,39 @@ def test_ffn_only_latency_less_than_full_and_different_from_attn():
     )
 
 
+def test_auto_suggest_cubes_match_default_topology():
+    """A TopologyConfig built from auto_suggest's returned (cp, tp, pp,
+    cp_placement) with all other fields left at their dataclass defaults
+    must produce the same cubes_used the Suggestion reports. Guards the
+    sidebar apply: any session_state key auto_suggest doesn't explore
+    (tp_placement, kv_shard_mode, cp_ring_variant, ep, ffn_shard_scope)
+    must be reset to the default before rerun, else the applied topology
+    won't match what the memory-min search actually scored.
+    """
+    machine = MachineParams()
+    matrix = [
+        ("Llama 3 70B", 8192),
+        ("Llama 3 70B", 32768),
+        ("Llama 3 8B",  131072),
+        ("Qwen 3 32B",  32768),
+    ]
+    for name, skv in matrix:
+        m = PRESETS[name].model
+        for ia, ff in [(True, False), (False, True), (True, True)]:
+            sug = auto_suggest(m, machine, s_kv=skv, mode="decode",
+                               include_attention=ia, include_ffn=ff)
+            topo = TopologyConfig(
+                cp=sug.cp, tp=sug.tp, pp=sug.pp,
+                s_kv=skv, mode="decode", b=1,
+                cp_placement=sug.cp_placement,
+            )
+            assert topo.cubes_used == sug.cubes_used, (
+                f"{name} skv={skv} scope=(ia={ia},ff={ff}): "
+                f"topo.cubes_used={topo.cubes_used} != "
+                f"sug.cubes_used={sug.cubes_used}"
+            )
+
+
 def test_parallelism_sensitivity_includes_baseline_value():
     """Each row's values contain the baseline_value."""
     from tests.analytical_visualization.auto_explore import (

tests/analytical_visualization/test_chip_roofline.py

@@ -0,0 +1,519 @@
+"""Tests for chip_roofline: AI, B*, L*, per-token latency curves."""
+from __future__ import annotations
+
+import math
+
+import pytest
+
+from tests.analytical_visualization.chip_roofline import (
+    ai_sensitivity_curve,
+    arithmetic_intensity,
+    balance_context,
+    bound_regime,
+    critical_batch,
+    good_batch,
+    good_context,
+    knee_batch,
+    kv_bytes_per_token,
+    max_batch_within_slo,
+    memory_budget_curve_vs_batch,
+    memory_budget_curve_vs_skv,
+    per_token_latency_curve,
+    size_deployment,
+    step_latency_curve,
+    t_com,
+    t_mem_long,
+    t_mem_short,
+    total_active_params,
+    utilization_at,
+)
+from tests.analytical_visualization.model_config import (
+    FullConfig, MachineParams, ModelConfig, TopologyConfig,
+)
+from tests.analytical_visualization.model_presets import PRESETS
+
+
+# ── Arithmetic intensity ────────────────────────────────────────────
+
+
+def test_arithmetic_intensity_matches_ratio():
+    """AI = peak_flops / bw_hbm — pure ratio, no util factor."""
+    m = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
+    assert arithmetic_intensity(m) == pytest.approx(8e12 / 256e9)
+
+
+def test_ai_scales_with_flops_and_bandwidth():
+    m1 = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
+    m2 = MachineParams(peak_tflops_f16=16.0, bw_hbm_gbs=256.0)
+    m3 = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=128.0)
+    assert arithmetic_intensity(m2) == pytest.approx(2 * arithmetic_intensity(m1))
+    assert arithmetic_intensity(m3) == pytest.approx(2 * arithmetic_intensity(m1))
+
+
+# ── Critical batch B* ───────────────────────────────────────────────
+
+
+def test_critical_batch_h100_reference():
+    """H100-class: 989 TFLOPs BF16, 3.35 TB/s HBM3 → B* ≈ 295."""
+    m = MachineParams(peak_tflops_f16=989.0, bw_hbm_gbs=3350.0)
+    model = PRESETS["Llama 3 8B"].model  # bf16 (b=2)
+    b_star = critical_batch(m, model)
+    assert b_star == pytest.approx(295, rel=0.05), b_star
+
+
+def test_critical_batch_scales_with_sparsity():
+    """MoE 8× sparsity gives 8× B*."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    b_dense = critical_batch(m, model, sparsity=1)
+    b_moe8 = critical_batch(m, model, sparsity=8)
+    assert b_moe8 == pytest.approx(8 * b_dense)
+
+
+def test_critical_batch_bf16_formula():
+    """For BF16 (b=2), B* = AI in flops-per-byte units → numerically."""
+    m = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
+    model = ModelConfig(bytes_per_elem=2)
+    ai = arithmetic_intensity(m)
+    assert critical_batch(m, model) == pytest.approx(ai), (
+        critical_batch(m, model), ai,
+    )
+
+
+# ── Balance context L* ──────────────────────────────────────────────
+
+
+def test_balance_context_positive_finite():
+    """L* for a real model on real machine is a positive finite number."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    l_star = balance_context(m, model)
+    assert 0 < l_star < 1e9, l_star
+
+
+def test_balance_context_scales_with_bandwidth():
+    """L* = 2N*W/(C*kv_bpt): doubling HBM BW doubles L*.
+    Slower memory hits the KV wall at a shorter context. Machine-only
+    change so N_active and kv_bpt stay fixed."""
+    fast = MachineParams(bw_hbm_gbs=512.0)
+    slow = MachineParams(bw_hbm_gbs=256.0)
+    model = PRESETS["Llama 3 8B"].model
+    assert (balance_context(fast, model)
+            == pytest.approx(2 * balance_context(slow, model)))
+
+
+# ── Knee ────────────────────────────────────────────────────────────
+
+
+def test_knee_equals_b_star_at_short_context():
+    """S_kv → 0: B_knee → B*."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    b_star = critical_batch(m, model)
+    assert knee_batch(m, model, s_kv=1) == pytest.approx(b_star, rel=1e-3)
+
+
+def test_knee_diverges_at_balance_context():
+    """S_kv = L*: knee is infinite; S_kv > L*: no knee (None)."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    l_star = balance_context(m, model)
+    assert knee_batch(m, model, s_kv=int(2 * l_star)) is None
+    # At 0.9 * L*, knee ~= 10 * B*; assert at least 5x to survive rounding.
+    below = knee_batch(m, model, s_kv=int(0.9 * l_star))
+    assert below is not None and below > 5 * critical_batch(m, model)
+
+
+def test_knee_slides_right_with_context():
+    """S_kv = L*/2 → B_knee = 2 * B*."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    b_star = critical_batch(m, model)
+    l_star = balance_context(m, model)
+    got = knee_batch(m, model, s_kv=int(l_star / 2))
+    assert got == pytest.approx(2 * b_star, rel=0.01), (got, 2 * b_star)
+
+
+# ── Per-token latency curve ─────────────────────────────────────────
+
+
+def test_latency_curve_monotonically_decreasing():
+    """total_s must strictly decrease as batch increases (weight/B shrinks)."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    pts = per_token_latency_curve(m, model, [1, 2, 4, 8, 16, 32, 64],
+                                    s_kv=1024)
+    totals = [p.total_s for p in pts]
+    for a, b in zip(totals, totals[1:]):
+        assert b < a, (a, b)
+
+
+def test_latency_curve_asymptotes_to_compute_plus_kv():
+    """At very large B, total → compute + KV (weight/B → 0)."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    pts = per_token_latency_curve(m, model, [10_000_000], s_kv=1024)
+    p = pts[0]
+    assert p.total_s == pytest.approx(p.compute_s + p.kv_s, rel=1e-3)
+
+
+def test_latency_at_b_star_is_roughly_2x_floor():
+    """At B*, weight = compute (dense, S_kv small), so total ≈ 2 * compute."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    b_star = int(round(critical_batch(m, model)))
+    pts = per_token_latency_curve(m, model, [b_star], s_kv=1)
+    p = pts[0]
+    # weight_s ≈ compute_s at B*
+    assert p.weight_s == pytest.approx(p.compute_s, rel=0.05), (
+        p.weight_s, p.compute_s,
+    )
+
+
+# ── Regime classification ───────────────────────────────────────────
+
+
+def test_regime_memory_bound_at_b1():
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    assert bound_regime(m, model, batch=1, s_kv=1024) == "memory-bound"
+
+
+def test_regime_kv_bound_past_balance_context():
+    """S_kv > L* with reasonable batch → KV dominates."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    l_star = balance_context(m, model)
+    b_star = int(round(critical_batch(m, model)))
+    assert bound_regime(m, model, batch=b_star * 4,
+                         s_kv=int(3 * l_star)) == "kv-bound"
+
+
+# ── Component sanity ────────────────────────────────────────────────
+
+
+def test_total_active_params_llama3_8b_ballpark():
+    """Llama 3 8B — attn+FFN alone (no embeddings/LM head) is ~6.9B,
+    HF-reported total 8.03B. Guard the ballpark."""
+    n = total_active_params(PRESETS["Llama 3 8B"].model)
+    assert 6.5e9 < n < 8e9, n
+
+
+def test_kv_bytes_per_token_llama3_8b():
+    """Llama 3 8B: 2*8*128*2 bytes/layer/token × 32 layers = 131072 bytes."""
+    kv_bpt = kv_bytes_per_token(PRESETS["Llama 3 8B"].model)
+    assert kv_bpt == 2 * 8 * 128 * 2 * 32
+
+
+# ── Regime terms + good-batch / good-context recommendations ──────
+
+
+def test_t_mem_short_shrinks_with_batch():
+    """t_mem_short = N·b/(W·B) — doubling B halves it."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    assert t_mem_short(m, model, 2) == pytest.approx(
+        t_mem_short(m, model, 1) / 2
+    )
+
+
+def test_t_mem_long_equals_t_com_at_l_star():
+    """L* is defined by t_mem_long(L*) == t_com. Cross-check."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    l_star = balance_context(m, model)
+    assert t_mem_long(m, model, int(round(l_star))) == pytest.approx(
+        t_com(m, model), rel=0.01,
+    )
+
+
+def test_t_com_is_batch_independent():
+    """t_com = 2·N/C. Constant — doesn't depend on B or S_kv."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    tc = t_com(m, model)
+    assert tc == pytest.approx(2 * total_active_params(model) / m.peak_flops)
+
+
+def test_good_batch_is_two_b_star():
+    """Recommendation is 2 × B* (Pope's rule)."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    rec = good_batch(m, model)
+    assert rec.effective == pytest.approx(2 * critical_batch(m, model))
+    assert rec.b_star == pytest.approx(critical_batch(m, model))
+
+
+def test_good_batch_moe_scales_with_sparsity():
+    """MoE sparsity=8 → recommended B is 8× dense recommendation."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    r_dense = good_batch(m, model, sparsity=1.0)
+    r_moe = good_batch(m, model, sparsity=8.0)
+    assert r_moe.effective == pytest.approx(8 * r_dense.effective)
+
+
+def test_good_context_returns_l_star_as_ceiling():
+    """Recommendation is L*; utilization drops past it."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    rec = good_context(m, model, s_kv=100)
+    assert rec.l_star == pytest.approx(balance_context(m, model))
+    assert rec.max_efficient == pytest.approx(rec.l_star)
+
+
+def test_utilization_at_l_star_is_half():
+    """At S_kv = L*, compute and KV read take equal time → 50% util."""
+    l_star = 3000.0
+    assert utilization_at(int(l_star), l_star) == pytest.approx(0.5)
+
+
+def test_utilization_at_2_l_star_is_one_third():
+    """At S_kv = 2·L*, util = 1/(1+2) = 33.3%."""
+    l_star = 3000.0
+    assert utilization_at(int(2 * l_star), l_star) == pytest.approx(1 / 3, rel=1e-3)
+
+
+def test_utilization_at_zero_context_is_100_percent():
+    """Zero context, no KV to read → all time is compute."""
+    assert utilization_at(0, 3000.0) == pytest.approx(1.0)
+
+
+def test_utilization_drops_monotonically_with_context():
+    """util(S_kv) is strictly decreasing."""
+    l_star = 3000.0
+    vals = [utilization_at(s, l_star) for s in [100, 1000, 3000, 6000, 30000]]
+    for a, b in zip(vals, vals[1:]):
+        assert b < a
+
+
+# ── Step latency (undivided by B) ──────────────────────────────────
+
+
+def test_step_weight_is_batch_invariant():
+    """step_weight = N·b/W — same for every B (loaded once per step)."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    pts = step_latency_curve(m, model, [1, 8, 64, 512], s_kv=1024)
+    ws = [p.weight_s for p in pts]
+    assert all(w == pytest.approx(ws[0]) for w in ws), ws
+
+
+def test_step_compute_linear_in_batch():
+    """step_compute = 2·N·B/C — doubling B doubles compute."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    pts = step_latency_curve(m, model, [1, 2, 8], s_kv=1024)
+    assert pts[1].compute_s == pytest.approx(2 * pts[0].compute_s)
+    assert pts[2].compute_s == pytest.approx(8 * pts[0].compute_s)
+
+
+def test_step_kv_linear_in_batch():
+    """step_kv = B · S_kv · kv_bpt / W — linear in B."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    pts = step_latency_curve(m, model, [1, 4], s_kv=1024)
+    assert pts[1].kv_s == pytest.approx(4 * pts[0].kv_s)
+
+
+def test_step_total_equals_per_token_times_batch():
+    """The step and per-token views are just ÷ B / × B of each other."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    B_range = [1, 4, 16, 64]
+    s_kv = 1024
+    step_pts = step_latency_curve(m, model, B_range, s_kv=s_kv)
+    tok_pts = per_token_latency_curve(m, model, B_range, s_kv=s_kv)
+    for sp, tp in zip(step_pts, tok_pts):
+        assert sp.total_s == pytest.approx(tp.total_s * sp.batch)
+
+
+def test_step_and_per_token_identical_at_b1():
+    """At B=1, step latency == per-token cost (dividing by 1)."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    sp = step_latency_curve(m, model, [1], s_kv=1024)[0]
+    tp = per_token_latency_curve(m, model, [1], s_kv=1024)[0]
+    assert sp.total_s == pytest.approx(tp.total_s)
+    assert sp.weight_s == pytest.approx(tp.weight_s)
+    assert sp.compute_s == pytest.approx(tp.compute_s)
+    assert sp.kv_s == pytest.approx(tp.kv_s)
+
+
+# ── PE memory budget sweeps ────────────────────────────────────────
+
+
+def _budget_cfg():
+    """Helper: real cfg with enough sharding to leave headroom."""
+    return FullConfig(
+        model=PRESETS["Llama 3 8B"].model,
+        topo=TopologyConfig(cp=2, tp=8, pp=1, b=1, s_kv=8192),
+        machine=MachineParams(),
+    )
+
+
+def test_kv_grows_linear_with_skv():
+    cfg = _budget_cfg()
+    pts = memory_budget_curve_vs_skv(cfg, s_kv_range=[1024, 2048, 8192],
+                                       batch=1)
+    assert pts[1].kv_gb == pytest.approx(2 * pts[0].kv_gb, rel=1e-3)
+    assert pts[2].kv_gb == pytest.approx(8 * pts[0].kv_gb, rel=1e-3)
+
+
+def test_kv_grows_linear_with_batch():
+    cfg = _budget_cfg()
+    pts = memory_budget_curve_vs_batch(cfg, b_range=[1, 2, 8],
+                                          s_kv=8192)
+    assert pts[1].kv_gb == pytest.approx(2 * pts[0].kv_gb, rel=1e-3)
+    assert pts[2].kv_gb == pytest.approx(8 * pts[0].kv_gb, rel=1e-3)
+
+
+def test_weights_flat_across_skv_sweep():
+    """Weight bytes are batch/S_kv-invariant — sharding fixed."""
+    cfg = _budget_cfg()
+    pts = memory_budget_curve_vs_skv(cfg, s_kv_range=[128, 8192, 1_000_000],
+                                       batch=1)
+    ws = [p.weights_gb for p in pts]
+    assert all(w == pytest.approx(ws[0]) for w in ws), ws
+
+
+def test_over_budget_flag_trips_past_hbm():
+    """Push S_kv until KV blows the HBM budget; over_budget must flip."""
+    cfg = _budget_cfg()
+    pts = memory_budget_curve_vs_skv(
+        cfg, s_kv_range=[1024, 10_000_000], batch=64,
+    )
+    # 10M tokens × B=64 KV should overflow the 6 GB budget.
+    assert not pts[0].over_budget, pts[0]
+    assert pts[-1].over_budget, pts[-1]
+
+
+def test_free_gb_never_negative():
+    cfg = _budget_cfg()
+    pts = memory_budget_curve_vs_skv(
+        cfg, s_kv_range=[128, 10_000_000], batch=32,
+    )
+    for p in pts:
+        assert p.free_gb >= 0, p
+
+
+# ── AI sensitivity to FLOPs / BW ──────────────────────────────────
+
+
+def test_ai_scales_linearly_with_flops():
+    """Doubling FLOPs doubles AI."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    pts = ai_sensitivity_curve(m, model, [1.0, 2.0, 4.0], axis="flops")
+    assert pts[1].ai == pytest.approx(2 * pts[0].ai)
+    assert pts[2].ai == pytest.approx(4 * pts[0].ai)
+
+
+def test_ai_scales_inversely_with_bw():
+    """Doubling BW halves AI."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    pts = ai_sensitivity_curve(m, model, [1.0, 2.0], axis="bw")
+    assert pts[1].ai == pytest.approx(pts[0].ai / 2)
+
+
+def test_ai_sensitivity_b_star_tracks_ai_for_bf16():
+    """For BF16, B* == AI. Curve should mirror."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model  # BF16
+    pts = ai_sensitivity_curve(m, model, [0.5, 1.0, 4.0], axis="flops")
+    for p in pts:
+        assert p.b_star == pytest.approx(p.ai)
+
+
+def test_ai_sensitivity_bad_axis_raises():
+    with pytest.raises(ValueError):
+        ai_sensitivity_curve(MachineParams(), PRESETS["Llama 3 8B"].model,
+                              [1.0], axis="bogus")
+
+
+# ── GPU sizing (three-axis) ────────────────────────────────────────
+
+
+def test_capacity_floor_scales_with_model_size():
+    """A model with 2× params needs ~2× PEs for weights alone (axis A)."""
+    m = MachineParams()
+    small = PRESETS["Llama 3 8B"].model
+    big = PRESETS["Llama 3 70B"].model
+    r_small = size_deployment(m, small, n_users=1, avg_ctx=1024,
+                                tpot_slo_s=1.0)
+    r_big = size_deployment(m, big, n_users=1, avg_ctx=1024,
+                              tpot_slo_s=1.0)
+    ratio = r_big.pes_axis_a_capacity / r_small.pes_axis_a_capacity
+    # 70B / 8B ≈ 8.75× params — allow generous tolerance
+    assert 5 < ratio < 12, (ratio, r_small, r_big)
+
+
+def test_kv_headroom_scales_with_users_and_context():
+    """Doubling n_users OR avg_ctx grows KV load roughly linearly."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    base = size_deployment(m, model, n_users=64, avg_ctx=4096,
+                             tpot_slo_s=1.0)
+    twice_users = size_deployment(m, model, n_users=128, avg_ctx=4096,
+                                    tpot_slo_s=1.0)
+    twice_ctx = size_deployment(m, model, n_users=64, avg_ctx=8192,
+                                  tpot_slo_s=1.0)
+    # KV bytes doubled either way → axis-B PE count must strictly grow.
+    assert twice_users.pes_axis_b_kv > base.pes_axis_b_kv
+    assert twice_ctx.pes_axis_b_kv > base.pes_axis_b_kv
+
+
+def test_binding_axis_flips_with_workload():
+    """Huge N_users + short ctx → throughput binds.
+    Huge ctx + few users → kv (or capacity) binds."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    big_users = size_deployment(m, model, n_users=10_000, avg_ctx=1024,
+                                  tpot_slo_s=0.05)
+    long_ctx = size_deployment(m, model, n_users=1, avg_ctx=500_000,
+                                 tpot_slo_s=1.0)
+    assert big_users.binding_axis == "throughput", big_users
+    assert long_ctx.binding_axis in ("kv", "capacity"), long_ctx
+
+
+def test_max_batch_shrinks_with_shorter_slo():
+    """Tighter SLO → smaller B fits."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    loose = max_batch_within_slo(m, model, s_kv=1024, tpot_slo_s=1.0)
+    tight = max_batch_within_slo(m, model, s_kv=1024, tpot_slo_s=0.05)
+    assert tight < loose, (tight, loose)
+
+
+def test_max_batch_returns_zero_when_weight_time_exceeds_slo():
+    """If weight-fetch alone > SLO, no batch fits."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 70B"].model
+    # 70B in BF16 = 140 GB. Reading at 256 GB/s takes ~547 ms.
+    # SLO of 100 ms cannot be met.
+    assert max_batch_within_slo(m, model, s_kv=100,
+                                  tpot_slo_s=0.1) == 0
+
+
+def test_total_pes_equals_per_replica_times_replicas():
+    """Sanity: total_pes = pes_per_replica × n_replicas."""
+    m = MachineParams()
+    model = PRESETS["Llama 3 8B"].model
+    r = size_deployment(m, model, n_users=500, avg_ctx=4096,
+                          tpot_slo_s=0.1)
+    assert r.total_pes == r.pes_per_replica * r.n_replicas
+
+
+# ── App wiring: tab exists on the Streamlit app ────────────────────
+
+
+def test_roofline_tab_registered_in_app():
+    """The new 'Chip roofline & B*' tab is listed in st.tabs and gets a
+    corresponding `with tab_roofline:` block."""
+    from pathlib import Path
+    src = (Path(__file__).parent / "app.py").resolve().read_text(
+        encoding="utf-8"
+    )
+    assert src.count('"Chip roofline & B*"') == 1
+    assert "with tab_roofline:" in src

tests/analytical_visualization/test_pe_weight_layout.py

@@ -0,0 +1,101 @@
+"""Tests for draw_pe_layout — CP rank locator visibility across TP tiers.
+
+The PE-level view has always color-labeled CP ranks when they're packed
+intra-cube (`cp_placement="pe"`). When TP >= 8, TP fills the cube and CP
+gets pushed to cube-level, so the intra-cube color branch is skipped —
+which used to silently drop the CP rank locator entirely. These tests
+guard the fix: a `[CP rank 0 of N]` cube tag + `CP=0` per-PE label are
+always present when CP > 1, regardless of placement.
+"""
+from __future__ import annotations
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from tests.analytical_visualization.model_config import (
+    FullConfig, MachineParams, TopologyConfig,
+)
+from tests.analytical_visualization.model_presets import PRESETS
+from tests.analytical_visualization.pe_weight_layout import draw_pe_layout
+
+
+def _cfg(cp: int, tp: int, cp_placement: str | None = None):
+    model = PRESETS["Llama 3 8B"].model
+    kwargs = dict(cp=cp, tp=tp, pp=1, s_kv=8192, mode="decode")
+    if cp_placement is not None:
+        kwargs["cp_placement"] = cp_placement
+    return FullConfig(
+        model=model, topo=TopologyConfig(**kwargs), machine=MachineParams(),
+    )
+
+
+def _all_text(fig) -> str:
+    """Concatenate all text artists (labels + title) into one blob."""
+    ax = fig.gca()
+    parts = [t.get_text() for t in ax.texts]
+    parts.append(ax.get_title())
+    return "\n".join(parts)
+
+
+def test_high_tp_cp_gt_1_shows_rank_locator():
+    """TP=8, CP=4 -> cp_placement forced to 'cube'; CP rank locator
+    must still appear in cube tag + per-PE header + figure title."""
+    cfg = _cfg(cp=4, tp=8)
+    assert cfg.topo.cp_placement == "cube"  # sanity
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        assert "CP rank 0 of 4" in text, text
+        assert "CP=0" in text, text
+        assert "CP=4" in text, text  # figure title dim summary
+    finally:
+        plt.close(fig)
+
+
+def test_intra_cube_cp_placement_still_shows_per_rank_labels():
+    """cp_placement='pe' (small TP, CP>1) — the old branch already
+    colors + labels each CP rank. Regression check: the new tag does
+    NOT replace those, and each CP=0..CP-1 label is still emitted."""
+    cfg = _cfg(cp=4, tp=2)  # 4*2=8 fits intra-cube
+    assert cfg.topo.cp_placement == "cube"  # default, but auto/manual can differ
+    # Force pe placement to exercise the intra-cube branch.
+    cfg.topo.cp_placement = "pe"
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        for r in range(4):
+            assert f"CP={r}" in text, (r, text)
+        # The 'CP rank 0 of N' tag is only for the cube-placement fallback.
+        assert "CP rank 0 of" not in text, text
+    finally:
+        plt.close(fig)
+
+
+def test_cp_1_no_locator_added():
+    """CP=1: nothing to locate; no CP tag or label should appear."""
+    cfg = _cfg(cp=1, tp=8)
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        assert "CP rank" not in text, text
+        # The header 'CP=0' should not be added when CP=1.
+        # (Figure title mentions 'CP=1' — that's an OK dim summary,
+        # not a locator per-PE.)
+        assert "  |  CP=0" not in text, text
+    finally:
+        plt.close(fig)
+
+
+def test_tp_16_still_shows_cp_rank_across_spilled_cubes():
+    """TP=16 (spills to 2 cubes per CP group), CP=2 -> both cubes
+    for CP rank 0 must carry the '[CP rank 0 of 2]' tag."""
+    cfg = _cfg(cp=2, tp=16)
+    assert cfg.topo.cp_placement == "cube"
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        assert text.count("CP rank 0 of 2") >= 2, text
+        assert "CP=0" in text, text
+    finally:
+        plt.close(fig)

tests/analytical_visualization/test_stage_shapes.py

@@ -0,0 +1,187 @@
+"""Tests for stage_shapes: per-PE shape rows for attention + FFN stages."""
+from __future__ import annotations
+
+from tests.analytical_visualization.model_config import (
+    FullConfig, MachineParams, TopologyConfig,
+)
+from tests.analytical_visualization.model_presets import PRESETS
+from tests.analytical_visualization.stage_shapes import (
+    attn_stage_shape_rows,
+    ffn_stage_shape_rows,
+)
+
+
+def _cfg(cp=2, tp=8, pp=1, b=1, mode="decode", s_kv=8192):
+    model = PRESETS["Llama 3 8B"].model
+    topo = TopologyConfig(cp=cp, tp=tp, pp=pp, b=b, s_kv=s_kv, mode=mode)
+    return FullConfig(model=model, topo=topo, machine=MachineParams())
+
+
+# ── Attention shapes ────────────────────────────────────────────────
+
+
+def test_attn_rows_have_required_columns():
+    """Every attention row must define Stage, Op, Input, Weight, Output."""
+    rows = attn_stage_shape_rows(_cfg())
+    required = {"Stage", "Op", "Input (per PE)",
+                "Weight (per PE)", "Output (per PE)"}
+    assert rows, "empty attention shape rows"
+    for r in rows:
+        missing = required - r.keys()
+        assert not missing, f"row {r.get('Stage')} missing {missing}"
+        for k in required:
+            assert r[k], f"row {r['Stage']} col {k} is empty"
+
+
+def test_attn_row_prefixes_present_in_order():
+    """S1..S10 always appear; C1 only when prefill+CP>1; C2 when TP>1;
+    C3 when kv_shard_mode=split and TP>H_kv; S8 only when CP>1."""
+    cfg = _cfg(cp=2, tp=8, mode="decode")
+    prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
+    # decode with CP=2, TP=8 (h_kv=8 for Llama 3 8B) → no C1, no C3; C2 yes; S8 yes.
+    assert "C1" not in prefixes, prefixes
+    assert "C3" not in prefixes, prefixes
+    assert "S8" in prefixes, prefixes
+    assert "C2" in prefixes, prefixes
+    # Ensure S1..S10 sequence appears in relative order.
+    s_order = [p for p in prefixes if p.startswith("S")]
+    assert s_order == sorted(s_order, key=lambda s: int(s[1:])), s_order
+
+
+def test_attn_cp1_omits_s8_merge():
+    """CP=1 → no online-softmax merge stage."""
+    cfg = _cfg(cp=1, tp=8, mode="decode")
+    prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
+    assert "S8" not in prefixes
+
+
+def test_attn_prefill_cp_gt_1_inserts_c1():
+    """Prefill mode with CP>1 shows the CP-ring stage."""
+    cfg = _cfg(cp=4, tp=2, mode="prefill", s_kv=8192)
+    prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
+    assert "C1" in prefixes
+
+
+def test_attn_shape_reflects_batch():
+    """B appears as the leading dim in each activation shape."""
+    for b in (1, 4, 32):
+        rows = attn_stage_shape_rows(_cfg(b=b))
+        s1 = next(r for r in rows if r["Stage"] == "S1")
+        assert s1["Input (per PE)"].startswith(f"({b},"), s1
+        assert s1["Output (per PE)"].startswith(f"({b},"), s1
+
+
+def test_attn_s5_shape_uses_tq_and_slocal():
+    """S5 (Q·K^T) shape must reference both T_q and S_local values."""
+    cfg = _cfg(cp=2, tp=8, mode="decode", s_kv=8192)
+    T_q = cfg.topo.T_q
+    S_local = cfg.topo.s_local
+    rows = attn_stage_shape_rows(cfg)
+    s5 = next(r for r in rows if r["Stage"] == "S5")
+    # Output shape (B, H_q/TP, T_q, S_local)
+    assert str(T_q) in s5["Output (per PE)"], s5
+    assert str(S_local) in s5["Output (per PE)"], s5
+
+
+# ── FFN shapes ──────────────────────────────────────────────────────
+
+
+def test_ffn_rows_have_required_columns():
+    """Every FFN row must define Stage, Op, Input, Weight, Output."""
+    rows = ffn_stage_shape_rows(_cfg())
+    required = {"Stage", "Op", "Input (per PE)",
+                "Weight (per PE)", "Output (per PE)"}
+    assert rows, "empty ffn shape rows"
+    for r in rows:
+        missing = required - r.keys()
+        assert not missing, f"row {r.get('Stage')} missing {missing}"
+
+
+def test_ffn_row_prefixes_f1_to_f5_always_present():
+    """F1..F5 always emitted; CF1 only when FFN divisor > 1."""
+    rows = ffn_stage_shape_rows(_cfg(cp=2, tp=8))
+    prefixes = [r["Stage"] for r in rows]
+    for p in ("F1", "F2", "F3", "F4", "F5"):
+        assert p in prefixes, prefixes
+
+
+def test_ffn_gate_up_output_matches_ffn_per_pe():
+    """F2/F3 output must have ffn_per_pe as the last dim."""
+    cfg = _cfg(cp=2, tp=8)
+    ffn_pe = max(1, cfg.model.ffn_dim
+                 // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+    rows = ffn_stage_shape_rows(cfg)
+    f2 = next(r for r in rows if r["Stage"] == "F2")
+    assert str(ffn_pe) in f2["Output (per PE)"], (ffn_pe, f2)
+
+
+# ── App wiring: consolidated shapes section lives on layout tab ───
+
+
+def test_shapes_section_is_on_physical_layout_tab():
+    """The 'All tensor shapes (per PE)' expander must live inside
+    tab_layout (Physical layout), not tab_stages. Guards the move so
+    the section doesn't accidentally end up back on the latency tab.
+    """
+    from pathlib import Path
+    app_path = (
+        Path(__file__).parent / "app.py"
+    ).resolve()
+    src = app_path.read_text(encoding="utf-8")
+    # Section appears exactly once.
+    assert src.count('"All tensor shapes (per PE)"') == 1, src.count(
+        '"All tensor shapes (per PE)"'
+    )
+    # And it sits inside tab_layout, before tab_memory begins.
+    section_pos = src.index('"All tensor shapes (per PE)"')
+    tab_memory_pos = src.index("with tab_memory:")
+    tab_stages_pos = src.index("with tab_stages:")
+    assert section_pos < tab_memory_pos, (
+        "shapes section must be inside tab_layout (before tab_memory)"
+    )
+    assert section_pos < tab_stages_pos, (
+        "shapes section must be inside tab_layout (before tab_stages)"
+    )
+    # The old per-stage shape rendering under tab_stages must be gone —
+    # no attn_stage_shape_rows / ffn_stage_shape_rows call may sit
+    # between 'with tab_stages:' and the next 'with tab_' block.
+    stages_block = src[tab_stages_pos:src.index("# ── TAB", tab_stages_pos + 1)]
+    assert "attn_stage_shape_rows(cfg)" not in stages_block, (
+        "per-stage shape table must not remain on Per-stage latency tab"
+    )
+    assert "ffn_stage_shape_rows(cfg)" not in stages_block, (
+        "per-stage FFN shape table must not remain on Per-stage latency tab"
+    )
+
+
+def test_ttft_tpot_section_wired_on_layout_tab():
+    """The 'Full-model latency (TTFT & TPOT)' subheader must live on
+    tab_layout, before tab_memory. Guard the section presence + placement."""
+    from pathlib import Path
+    src = (
+        Path(__file__).parent / "app.py"
+    ).resolve().read_text(encoding="utf-8")
+    assert src.count('"Full-model latency (TTFT & TPOT)"') == 1
+    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

tests/analytical_visualization/topology_map.py

@@ -258,12 +258,27 @@ def _draw_one_sip(ax, cfg: FullConfig, sip_idx: int,
         pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
         pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
         pes_used = cfg.topo.pes_per_cube_used if is_used else 0
+        # When cp_placement=pe, multiple CP ranks live inside this cube's PEs
+        # (cp_rank = pe_id // tp). Color each PE by its cp_rank so all four
+        # groups are visible; otherwise fall back to the whole-cube pe_fill.
+        _cp_packed = is_used and cfg.topo.cp_placement == "pe" and cfg.topo.cp > 1
         for pr in range(PE_ROWS):
             for pc in range(PE_COLS):
                 pe_id = pr * PE_COLS + pc
                 px = x + pe_pad + pc * (pe_w + pe_gap)
                 py = y + pe_pad + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
-                fill = pe_fill if pe_id < pes_used else _INACTIVE_PE
+                if pe_id >= pes_used:
+                    fill = _INACTIVE_PE
+                elif _cp_packed:
+                    _pe_cp_rank = pe_id // cfg.topo.tp
+                    _pe_color, _pe_pe_fill = _cp_color(
+                        info[0] if is_used else 0,
+                        _pe_cp_rank,
+                        cfg.topo.cp,
+                    )
+                    fill = _pe_pe_fill
+                else:
+                    fill = pe_fill
                 pe_rect = patches.Rectangle(
                     (px, py), pe_w, pe_h,
                     facecolor=fill, edgecolor="#666", linewidth=0.3,