diff --git a/docs/report/1H-codesign-paper/build/main.pdf b/docs/report/1H-codesign-paper/build/main.pdf index d6f1319..803e246 100644 Binary files a/docs/report/1H-codesign-paper/build/main.pdf and b/docs/report/1H-codesign-paper/build/main.pdf differ diff --git a/docs/report/1H-codesign-paper/sections/05-gqa.tex b/docs/report/1H-codesign-paper/sections/05-gqa.tex index 1a71ab5..0eda884 100644 --- a/docs/report/1H-codesign-paper/sections/05-gqa.tex +++ b/docs/report/1H-codesign-paper/sections/05-gqa.tex @@ -128,9 +128,14 @@ 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. -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 +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 @@ -155,10 +160,12 @@ 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}. -At every context length the \textsf{1-kv-per-cube} mapping is the -fastest, and the $\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube}$ ratio -grows monotonically with $S_{kv}$ (decode: $0.98\times$ at $8$K, -$1.06\times$ at $64$K; prefill follows the same trend). The wall gap is +At short context the two mappings are effectively tied, and they separate +as context grows: the $\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube}$ +latency ratio rises monotonically with $S_{kv}$ (decode: $0.98\times$ at +$8$K---within measurement noise---to $1.06\times$ at $64$K, where +\textsf{1-kv-per-cube} is the faster; prefill follows the same trend). The +wall gap is modest because short-context decode is dominated by setup overhead and its $M{=}G{=}8$ skinny shape leaves the MAC array idle most of the time; where the mappings really separate is per-CUBE pressure. @@ -167,10 +174,11 @@ where the mappings really separate is per-CUBE pressure. \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). Prefill and decode follow the same ordering at -every context length and the \textsf{1-kv-per-cube} mapping is the -fastest throughout; the gap to \textsf{8-kv-per-cube} grows from -$0.98\times$ at $8$K to $1.06\times$ at $64$K. The composite and +kernel (no composite). Prefill and decode follow the same ordering; the +two mappings are effectively tied at $8$K and separate as context grows, +the \textsf{8-kv-per-cube}/\textsf{1-kv-per-cube} ratio widening from +$0.98\times$ at $8$K to $1.06\times$ at $64$K (\textsf{1-kv-per-cube} +faster at the large end). The composite and softmax\_merge tiers track the baseline within a few percent at this shape (see Figure~\ref{fig:gqa-short-a1-variant}).} \label{fig:gqa-short-wall-baseline} @@ -320,22 +328,17 @@ absorb at link speed.} \label{fig:gqa-6cases-traffic} \end{figure} -The measurements select the right placement for this regime. The fastest -raw latency (Case~3, \SI{17.8}{\micro\second}) comes from replicating the -full KV cache into every CUBE---which is exactly the placement the -long-context memory budget forbids, and which the parallelism panel shows -wastes 8$\times$ the compute. Restricting attention to the placements that -fit production-context memory (the 64-way splits, Cases~4--6), the choice -is Case~6~$\star$: it is the fastest of the three -(\SI{30.6}{\micro\second}), and the op-count panel shows why---its -softmax-state-only reduction charges 189 IPCQ copies and one DMA write -against the 280 copies and 8 DMA writes the $d_{\text{head}}$-split -Cases~4--5 pay for their partial-score all-reduce. For long-context -decode, then, the appropriate data placement is the both-axes sequence -shard (Case~6): it is the cheapest-communicating member of the only -memory-feasible family, 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. +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} @@ -459,24 +462,24 @@ 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{Comprehensive Analysis} +\subsection{Summary} \label{sec:gqa-analysis} -These panels are the clearest statement of the codesign thesis in the -report. Because the composite command keeps GEMM issue cheap and the MAC -array barely occupied, the fused attention kernel's latency is set almost -entirely by data movement: streaming the KV cache and reducing partials -across devices. That is precisely the cost that the communication-side -work targets---PE\_IPCQ for the on-device reduction, the lazy load for -load/compute overlap, fast TCM staging and torus links for the reduction -itself. In other words, the two enablers are not independent features that -happen to appear in the same kernel; the GEMM optimization is what -\emph{exposes} the data-movement bottleneck (by removing the compute and -issue overhead that would otherwise hide it), and the communication -optimization is what \emph{attacks} it. For an attention-dominated decoder -the meaningful hardware investments are therefore the ones that move data -faster and reduce it on-device---not additional MAC throughput, which this -workload cannot use. +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