paper(1H): reflect D8 single-op cost model in §2/§3.4 + figure/diagram regen
Two strands bundled as the 1H-codesign-paper refresh unit: (A) This session — single-op cost-model reflection (depends on2d8271c): - §2 Table 2 (tab:hw): split "FIXED per command" into "FIXED per single-op command" (8 cycles) and "FIXED per composite command" (40 cycles); §2 dispatch-overhead prose updated to the two-class split. - §3.4 (sec:gemm-vs-async): rename paragraph headers + prose to async-full / async-tiled; "atomic" -> "single-op" throughout; reframe mechanism #3 from the old DMA-only fast-path to the single-op fast-path. Headline narrative now: even with EVERY single-op cmd (96 DMA + 48 dot + 47 add) charged the light 8-cycle FIXED, composite still wins ~2.8x at K=3072 purely on command-count structure (1 vs 192 commands) -- down from the pre-D8 ~6.3x, and explicitly NOT a modelling artifact. Numbers refreshed from the regenerated sweep: async-full 3.83->3.91, async-tiled 1.14->~2.53, under-tile corner 1.06->1.21, depth-2 vs depth-inf spread <1%. New figure wired in. - build/main.pdf rebuilt (tectonic); pdftotext-verified (no broken refs; Table 2 split, single-op terms, 2.8x/2.53/192-host-commands all present). (B) Prior-session paper work riding along uncommitted: §4 all-reduce deep-edit, §5 GQA, §6 discussion trims; milestone_1h_ccl.py plot label "FSIM" -> "H2 2025 SW queue baseline"; regenerated diagrams under docs/diagrams/** and gemm output PNGs under 1H_milestone_output/gemm/. (Composite-window gemm plots are unaffected by D8 — D8 only changes single-op dispatch FIXED, which the composite window excludes.) All TODO items for the D8 single-op extension are now complete and pushed across 3 commits (2d8271ccost-model+ADR+tests,821bbf2bench harness, this paper refresh). Full regression green (826 passed, 1 skipped). No remaining work. NOTE for review (carried from2d8271c): ADR-0065's "2x CPU-offload win" headline for GQA decode opt2 may want a refresh to the post-D8 ~1.87x. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Binary file not shown.
Binary file not shown.
|
Before Width: | Height: | Size: 86 KiB After Width: | Height: | Size: 92 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 75 KiB |
Binary file not shown.
@@ -243,14 +243,18 @@ serializing the command descriptor into the scheduler queue. This
|
||||
\SI{2}{\nano\second} / \SI{1}{\nano\second} \textsf{PE\_CPU} /
|
||||
\textsf{PE\_SCHED} fixed costs: the latter are per-component
|
||||
traversal overheads each command pays when it transits those nodes,
|
||||
whereas the 40-cycle \textsf{FIXED} term is the command-descriptor
|
||||
issue cost. The
|
||||
default anchoring (\textsf{FIXED} $= 40$ cycles, $R = 0.0625$
|
||||
cycles/byte, i.e.\ \SI{16}{\byte\per\cycle}, at \SI{1}{\giga\hertz})
|
||||
places a typical composite at roughly \SI{43}{\nano\second}, and a
|
||||
hard cap on a composite's descriptor size prevents the model from
|
||||
rewarding arbitrarily large fused commands beyond what real descriptor
|
||||
queues accept. In the configurations measured here, command issue is
|
||||
whereas the \textsf{FIXED} term here is the command-descriptor
|
||||
issue cost. \textsf{FIXED} depends on the command class: a single-op
|
||||
command --- one engine operation, i.e.\ a single DMA descriptor, GEMM,
|
||||
or elementwise issue --- carries a lighter 8-cycle \textsf{FIXED},
|
||||
while the composite command, which the scheduler expands into a
|
||||
multi-stage tile-feeder plan, carries 40 cycles
|
||||
(\S\ref{sec:gemm-vs-async}). With the composite \textsf{FIXED} $= 40$
|
||||
cycles and $R = 0.0625$ cycles/byte (i.e.\ \SI{16}{\byte\per\cycle},
|
||||
at \SI{1}{\giga\hertz}) a typical composite lands at roughly
|
||||
\SI{43}{\nano\second}, and a hard cap on a composite's descriptor size
|
||||
prevents the model from rewarding arbitrarily large fused commands
|
||||
beyond what real descriptor queues accept. In the configurations measured here, command issue is
|
||||
not the bottleneck---data movement is---so this term stays small
|
||||
relative to DMA and collective time.
|
||||
|
||||
@@ -452,7 +456,8 @@ Inter-CUBE (UCIe PHY) & \SI{512}{\giga\byte\per\second}, \SI{8}{\nano\second}, X
|
||||
Inter-SIP (PCIe) & \SI{768}{\giga\byte\per\second} per endpoint \\
|
||||
\midrule
|
||||
\multicolumn{2}{@{}l}{\emph{Command-issue cost model (defaults)}} \\
|
||||
FIXED per command & 40 cycles \\
|
||||
FIXED per single-op command & 8 cycles \\
|
||||
FIXED per composite command & 40 cycles \\
|
||||
per-byte rate $R$ & 0.0625 cycles/byte (\SI{16}{\byte\per\cycle}) \\
|
||||
composite size cap & \SI{1024}{\byte} \\
|
||||
\bottomrule
|
||||
|
||||
@@ -171,3 +171,231 @@ HBM port and the kernel is BW-bound. This is the lever the GQA
|
||||
kernel of \S\ref{sec:gqa} reaches for next — keeping the right
|
||||
working set on-chip so the composite pipeline lands in the
|
||||
compute-rich regime rather than the BW-bound one.
|
||||
|
||||
\subsection{Why composite, and not user-orchestrated async loading?}
|
||||
\label{sec:gemm-vs-async}
|
||||
|
||||
A reader familiar with double-buffered GEMM kernels on conventional
|
||||
hardware may ask: why a hardware-side composite command at all? Why
|
||||
isn't the obvious user-level pattern --- async-load each operand,
|
||||
overlap with compute, accumulate --- sufficient?
|
||||
|
||||
To answer this concretely we contrast composite against two
|
||||
user-orchestrated baselines that have access to the same single-op
|
||||
primitives the platform exposes (\textsf{tl.load} for async DMA into
|
||||
TCM, \textsf{tl.dot} for a single-op GEMM command on TCM-resident
|
||||
operands, \textsf{tl.store} for a DMA write-back). Both baselines
|
||||
pre-stage activation $A$ identically to \textsf{load\_ref}, so the
|
||||
window measured is the engine-pipeline window with $A$'s up-front DMA
|
||||
excluded for all three kernels.
|
||||
|
||||
\paragraph{Async-full (naive).} A single \textsf{tl.load(A)} followed by
|
||||
a single \textsf{tl.load(B)} (async, queued behind $A$),
|
||||
\textsf{tl.dot(A, B)}, and \textsf{tl.store(out)}. The kernel issues
|
||||
four commands total. The decisive constraint is that
|
||||
\textsf{tl.dot}'s \textsf{\_await\_pending(b)} blocks the GEMM
|
||||
command until the \emph{entire} $B$ has landed in TCM --- there is no
|
||||
way at the runtime API surface to express ``start computing on
|
||||
tile 0 of $B$ while tile 1 is still in flight.'' Load-of-$B$ and
|
||||
GEMM therefore serialize.
|
||||
|
||||
\paragraph{Async-tiled (chunked prefetch).} The user-level workaround is to
|
||||
split $B$ along $K$ into \textsf{TILE\_K}-sized chunks, issue async
|
||||
\textsf{tl.load}s for those chunks, issue one \textsf{tl.dot} per
|
||||
chunk (each blocking only on its own $b_i$), and accumulate via
|
||||
\textsf{out = out + tl.dot(...)}. This is the standard
|
||||
double-buffered-GEMM pattern transcribed to the single-op primitives.
|
||||
We measure two versions of this kernel that differ only in
|
||||
\emph{prefetch depth}:
|
||||
\textbf{depth-2 (TCM-bounded)} keeps at most two B-chunks in flight
|
||||
at any time by issuing the next \textsf{tl.load} just before each
|
||||
\textsf{tl.dot}; \textbf{depth-$\infty$ (queue-all)} issues all $N_K$
|
||||
B-chunk loads up front so the DMA engine has the deepest possible
|
||||
request queue. For a $K{=}3072$ shape both versions emit
|
||||
48 $A$-chunk loads + 48 $B$-chunk loads + 48 \textsf{tl.dot}s + 47
|
||||
elementwise adds + 1 store = 192 host-side commands; the only
|
||||
difference is the temporal interleaving of B-load and dot dispatches.
|
||||
|
||||
\paragraph{Why two depths.} The depth distinction matters because the
|
||||
async-full kernel and the depth-$\infty$ async-tiled kernel both pin the
|
||||
\emph{entire} $B$ in TCM simultaneously --- $K \cdot N \cdot 2$
|
||||
bytes. The on-PE scratch is capped at \SI{1}{\mebi\byte}
|
||||
(\texttt{topology.yaml: pe\_tcm.kernel\_scratch\_mb=1}), so an LLM-scale
|
||||
attention $B = K_{\text{KV}} \times d_{\text{head}}$ at
|
||||
$K_{\text{KV}}=4096, d_{\text{head}}=128$ already needs
|
||||
\SI{1}{\mebi\byte}, and at $K_{\text{KV}}=8192$ it needs \SI{2}{\mebi\byte}
|
||||
--- past the cap. async-full and queue-all async-tiled are therefore
|
||||
not just slower than composite but \emph{architecturally infeasible}
|
||||
at LLM context length. The depth-2 async-tiled kernel is the only
|
||||
user-level
|
||||
variant whose peak TCM footprint stays
|
||||
$O(2 \cdot \textsf{TILE\_K} \cdot N)$ regardless of $K$, the same
|
||||
order as composite's per-tile streaming buffer. It is the apples-to-apples
|
||||
comparison.
|
||||
|
||||
\paragraph{Per-PE throughput.} Figure~\ref{fig:gemm-async} reports
|
||||
per-PE TFLOP/s for all four kernels side-by-side. The single-op
|
||||
fast-path in the dispatch cost model (\S\ref{sec:congestion}) is
|
||||
enabled, so every single-op command the async kernels emit --- DMA
|
||||
descriptors and \textsf{tl.dot}/\textsf{tl.add} alike --- is charged the
|
||||
light 8-cycle \textsf{FIXED}, not the 40-cycle composite control-path
|
||||
cost; neither async-tiled variant carries an inflated per-command cost.
|
||||
|
||||
\paragraph{The $K{=}3072$ corner, concretely.} We take this shape
|
||||
($M{=}32, K{=}3072, N{=}32$) as the running example throughout the
|
||||
mechanism discussion because it is the regime where the four
|
||||
kernels spread the widest --- the deepest $K$ in the sweep maps
|
||||
onto $K / \textsf{TILE\_K} = 48$ hardware tiles, so per-tile costs
|
||||
amplify into the largest measurable gap. The work content is
|
||||
identical for all four kernels: $\sim$6.3 M f16 MACs and
|
||||
$\sim$386 KiB of $B$ traffic from HBM, which together require
|
||||
$\sim$786 ns of GEMM-engine compute and $\sim$750 ns of DMA on a
|
||||
saturated per-PE link. What differs is the number of host commands
|
||||
the same work is decomposed into --- 2 for composite (one
|
||||
\textsf{tl.load(A)} plus one composite), 4 for async-full, and 192
|
||||
for either async-tiled variant (48 $A$-loads + 48 $B$-loads + 48
|
||||
\textsf{tl.dot}s + 47 elementwise adds + 1 store). The
|
||||
engine-pipeline-window throughput tracks that decomposition closely:
|
||||
composite reaches \SI{7.18}{\tera\flop\per\second} (post-overlap
|
||||
limit, only \SI{12}{\percent} below the \SI{8}{\tera\flop\per\second}
|
||||
per-PE GEMM peak), async-full \SI{3.91}{\tera\flop\per\second}
|
||||
(DMA and compute serialize on a single big dot), and both async-tiled
|
||||
variants $\sim$\SI{2.53}{\tera\flop\per\second} (192 commands' worth of
|
||||
structural dispatch cost --- even at the light per-command rate ---
|
||||
accumulates on the wall). The next paragraph attributes those gaps to
|
||||
specific simulator mechanisms.
|
||||
|
||||
\begin{figure*}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gemm_composite_vs_async_tflops.png}
|
||||
\caption{Per-PE achieved TFLOP/s for the same shape sweep run under
|
||||
four issuance patterns: composite (one command, scheduler streams
|
||||
per-tile internally), async-full (one \textsf{tl.dot} on
|
||||
fully-loaded $B$), async-tiled with depth-2 double-buffer
|
||||
(TCM-bounded; the only user-level variant that scales to LLM
|
||||
context length), and async-tiled with depth-$\infty$
|
||||
(all B-tiles queued up front; included as a sanity check that the
|
||||
prefetch depth is \emph{not} what separates the async-tiled kernel from
|
||||
composite). All curves exclude $A$'s up-front DMA from the
|
||||
measurement window. Composite wins at every full-tile shape; the
|
||||
depth-2 and depth-$\infty$ async-tiled kernels deliver
|
||||
\emph{indistinguishable} throughput (e.g. \SI{2.522}{} vs.\
|
||||
\SI{2.544}{\tera\flop\per\second} at $K{=}3072$), confirming that
|
||||
prefetch depth is not the lever --- the structural per-command
|
||||
dispatch cost is. The gap between composite and the async-tiled
|
||||
kernels grows with $K_{\text{useful}}$
|
||||
(\textbf{$\sim$$2.8\times$} at $K{=}3072$, where the async-tiled
|
||||
kernels emit 192 host commands while composite emits one).
|
||||
The under-tile corner $M{=}128,K{=}8,N{=}128$ inverts: composite's
|
||||
per-tile orchestration overhead exceeds the per-tile useful work,
|
||||
and all three async kernels beat it.}
|
||||
\label{fig:gemm-async}
|
||||
\end{figure*}
|
||||
|
||||
\paragraph{Decomposing the gap.} Three structural mechanisms separate
|
||||
composite from the user-level baselines, and they layer.
|
||||
|
||||
\emph{1. Inter-engine token routing happens below the host-side
|
||||
dispatch path.} The composite encodes the full
|
||||
\textsf{DMA\_READ}$\to$\textsf{FETCH}$\to$\textsf{GEMM}$\to$\textsf{STORE}$\to$\textsf{DMA\_WRITE}
|
||||
pipeline once. The scheduler's tile-feeder loop then emits one
|
||||
\emph{tile token} per HW tile inside that one composite, and each
|
||||
token self-routes between engines after each stage finishes. The
|
||||
per-tile token routing is a scheduler-internal event, not a fresh
|
||||
host command, so it does not pay the structural CPU dispatch cost.
|
||||
At $K{=}3072$ the composite emits 48 tile tokens that flow
|
||||
fully-pipelined through five stages each --- 240 inter-engine
|
||||
hand-offs total --- behind a single command from the host's point of
|
||||
view.
|
||||
|
||||
\emph{2. \textsf{tl.dot} cannot replicate that per-tile pipeline at the
|
||||
user level.} A single-op GEMM command is handled on the GEMM engine as
|
||||
a single monolithic compute timeout for the supplied $M{\times}K{\times}N$;
|
||||
there is no internal token loop that would let a streaming DMA of
|
||||
$B[i{+}1]$ overlap with the GEMM of $B[i]$ inside one
|
||||
\textsf{tl.dot}. The user can only recover inter-tile overlap by
|
||||
emitting one \textsf{tl.dot} per chunk --- which is exactly what the
|
||||
async-tiled baseline does, at the price of $N$ host-side commands.
|
||||
|
||||
\emph{3. The host-side dispatch cost the async-tiled baseline pays is
|
||||
structural, not modelling slack.} KernBench charges every host-emitted
|
||||
command a structural CPU dispatch cost $d_{\text{cmd}} = \textsf{FIXED} +
|
||||
b_{\text{logical}} \cdot R$ (\S\ref{sec:congestion}). The cost model is
|
||||
deliberately charitable to the async kernels here: only a
|
||||
\textsf{CompositeCmd} pays the 40-cycle control-path \textsf{FIXED} (it
|
||||
alone drives a scheduler-built tile-feeder plan), while \emph{every}
|
||||
single-op command --- the 96 DMA descriptors, the 48 \textsf{tl.dot}s,
|
||||
and the 47 elementwise adds alike --- pays only the light 8-cycle
|
||||
\textsf{FIXED}, calibrated to descriptor-ring-push / single-instruction
|
||||
issue patterns in modern accelerators (NVIDIA Hopper TMA $\sim$1 ISA
|
||||
cycle, a single \textsf{mma.sync} one instruction, AMD AQL packet writes
|
||||
$\sim$5--15 cycles). So the async-tiled baseline is \emph{not} penalized
|
||||
by an inflated per-command cost on \emph{any} of its operations. Yet it
|
||||
still emits 192 host commands against composite's one, so $\sim$192
|
||||
light dispatches accumulate to $\sim$\SI{1.5}{\micro\second} of
|
||||
structural PE\_CPU time that composite never pays --- composite hides
|
||||
its 48 tiles' 240 inter-engine hand-offs as scheduler-internal events
|
||||
(mechanism 1). Layered on top, the dependency chain on the running
|
||||
accumulator serializes the 47 adds on the math engine. The net result
|
||||
is that the async-tiled kernel runs $\sim$2.8$\times$ slower than
|
||||
composite at $K{=}3072$ despite getting the inter-chunk overlap right
|
||||
--- down from $\sim$6.3$\times$ under the earlier uniform-40 cost model,
|
||||
because D8 removed the per-command overcharge, but \emph{not} closed:
|
||||
the residual gap is the command-count structure (one composite vs.\ 192
|
||||
single-ops), not a modelling artifact.
|
||||
|
||||
\emph{4. Prefetch depth is not the lever, command count is.} The
|
||||
depth-2 and depth-$\infty$ async-tiled kernels land within \SI{1}{\percent}
|
||||
of each other at every measured shape (Figure~\ref{fig:gemm-async}).
|
||||
This is the diagnostic against a natural objection: ``surely the
|
||||
async-tiled kernel was just under-prefetching; deepen the queue and the
|
||||
DMA$\to$GEMM overlap recovers.'' Deepening the prefetch queue
|
||||
changes \emph{when} DMA descriptors hit the engine but not their
|
||||
total number or the structural dispatch cost they each pay. The
|
||||
$\sim$2.8$\times$ gap to composite is not a prefetch-depth gap; it
|
||||
is the gap between
|
||||
``one composite command with internal per-tile token routing'' and
|
||||
``$N_K$ host-side dot/add commands, each charged separately.''
|
||||
The depth-2 kernel additionally constrains peak TCM occupancy to
|
||||
$O(2 \cdot \textsf{TILE\_K} \cdot N)$, matching composite's per-tile
|
||||
streaming buffer; the depth-$\infty$ kernel needs the full $B$ in
|
||||
TCM, which makes it infeasible at LLM context length even if the
|
||||
throughput were competitive.
|
||||
|
||||
\paragraph{Where the composite advantage doesn't apply.} The shape
|
||||
$M{=}128,K{=}8,N{=}128$ inverts the picture: composite delivers
|
||||
\SI{0.66}{\tera\flop\per\second} and both async kernels reach
|
||||
\SI{1.21}{\tera\flop\per\second}. The reason is consistent with the
|
||||
analysis above and explicit in the simulator state. Composite emits
|
||||
16 tile tokens (one per output tile) for this shape, each carrying
|
||||
$K_{\text{useful}}{=}8$ MACs across a TILE\_K$=64$ pipeline ---
|
||||
$12.5\%$ of the hardware tile's MAC slots are useful, the rest is
|
||||
K-padding. The per-tile inter-engine hand-off cost stays the same
|
||||
regardless. When the per-tile useful work is small enough that
|
||||
hand-off overhead exceeds the GEMM work itself, a single-op
|
||||
\textsf{tl.dot} --- which submits one monolithic GEMM command with no
|
||||
per-tile orchestration --- wins. The take-away is the bound on
|
||||
composite's value: it amortizes useful per-tile compute, not
|
||||
padding-dominated under-tile shapes. Real kernels at this corner are
|
||||
better served by reshape-into-batched-GEMM transforms that move
|
||||
under-tile $K$ into a tile-filling dimension before reaching the GEMM
|
||||
engine, which is exactly what the GQA decode kernel of
|
||||
\S\ref{sec:gqa} does for its $K_{\text{useful}}=\text{head\_dim}=128$
|
||||
inner reduction.
|
||||
|
||||
\paragraph{Summary of the comparison.} Composite is the only one of
|
||||
the four kernels to combine (a) macro-command dispatch at the host
|
||||
boundary (amortizing the structural CPU cost across all the work a
|
||||
single GEMM does), (b) scheduler-internal per-HW-tile streaming of
|
||||
DMA$\rightleftarrows$compute, and (c) TCM-bounded streaming buffer.
|
||||
User-orchestrated async kernels can have any two of those, not all
|
||||
three: async-full pays one host dispatch (a) but forfeits per-tile
|
||||
overlap (b) and pins all of $B$ in TCM (c); depth-$\infty$ async-tiled
|
||||
achieves inter-chunk overlap but at $N_K$ host dispatches and
|
||||
full-$B$ TCM occupancy; depth-2 async-tiled fixes the TCM
|
||||
footprint (c) but still pays $N_K$ host dispatches. The two corners
|
||||
where async catches up
|
||||
(small-$K$ where composite has nothing useful to amortize;
|
||||
under-tile shapes where per-tile useful work is sub-token) are
|
||||
diagnostic of \emph{where composite is the wrong tool}, not of a
|
||||
slack the user kernel could close.
|
||||
|
||||
@@ -1,46 +1,145 @@
|
||||
\section{PE\_IPCQ and Collective Communication}
|
||||
\label{sec:allreduce}
|
||||
|
||||
\subsection{Why it is needed}
|
||||
|
||||
Distributing a transformer across devices turns every tensor-parallel
|
||||
layer into a collective: partial results computed on different PEs, CUBEs,
|
||||
and SIPs must be summed and redistributed with an all-reduce. If that
|
||||
collective is handled by the host or by a generic DMA path, three problems
|
||||
appear. The reduction traffic competes with the kernel's own compute DMA
|
||||
on the same links, causing head-of-line blocking; there is no efficient
|
||||
peer-to-peer ring primitive, so data takes extra hops; and the ordering is
|
||||
hard to make deterministic. The hardware question is how to perform
|
||||
collectives \emph{on the device}, overlapped with compute and reproducible
|
||||
run-to-run.
|
||||
and SIPs must be summed and redistributed with an all-reduce. Underneath
|
||||
the algorithm this is fundamentally a PE-to-PE problem---many short
|
||||
messages flowing between neighbors as the reduction proceeds. The natural
|
||||
software realization is a per-direction ring buffer whose head and tail
|
||||
pointers the producer and consumer update atomically and poll, but our
|
||||
H2 2025 report measured this scheme end-to-end and found that the
|
||||
atomic-pointer traffic together with the consumer's polling loop dominate
|
||||
the per-message cost: the queue itself becomes the bottleneck well before
|
||||
the link runs out of bandwidth, and a pure-SW collective spends most of
|
||||
its time on metadata rather than on actually moving partials. A second,
|
||||
orthogonal problem is link sharing---if the collective rides the kernel's
|
||||
generic DMA path it competes with the GEMM's compute traffic on the same
|
||||
wires, so a large tile transfer head-of-line-blocks a pending reduction.
|
||||
This section asks how a dedicated hardware primitive can lift queue
|
||||
management off the software path entirely and, in the same design, stop
|
||||
collective traffic from stalling behind compute traffic, so the all-reduce
|
||||
runs overlapped with compute and at the interconnect's physical limit.
|
||||
|
||||
\subsection{Design}
|
||||
|
||||
KernBench models a dedicated per-PE collective engine, \textbf{PE\_IPCQ}
|
||||
(inter-PE communication queue). It is a control-plane block: it owns the
|
||||
ring-buffer address arithmetic, head/tail pointers, peer-pointer caches,
|
||||
backpressure, and the four-direction (N/S/E/W) neighbor map, with eight
|
||||
ring buffers per PE (four directions $\times$ \{tx, rx\}). Crucially,
|
||||
PE\_IPCQ does \emph{not} move data itself---it delegates the actual
|
||||
transfer to PE\_DMA, keeping a clean control/data split. To stop
|
||||
collective traffic from blocking compute, PE\_DMA is extended into a
|
||||
two-channel virtual-channel model: \texttt{vc\_compute} carries tile
|
||||
load/store for GEMM and vector math, \texttt{vc\_comm} carries IPCQ sends,
|
||||
each with an independent state machine. The same physical link is shared
|
||||
but progresses in chunks (\SI{256}{\byte}), so a large GEMM DMA does not
|
||||
lock the link end-to-end against a pending reduction. On top of this
|
||||
substrate the collective runs a hierarchical local-reduce / global
|
||||
all-reduce-broadcast schedule across whatever inter-device topology the
|
||||
configuration specifies.
|
||||
The proposed block is \textbf{PE\_IPCQ} (inter-PE communication queue),
|
||||
a small controller dropped into every PE next to PE\_DMA, PE\_GEMM,
|
||||
PE\_MATH, and PE\_TCM. It is a control-plane block---it holds no
|
||||
payload data---and consists of three pieces: a per-direction
|
||||
\emph{QPair register file} (\textasciitilde\SI{576}{\byte} of flip-flops
|
||||
covering up to eight directions), a combinational slot-address
|
||||
generator and backpressure comparator, and a credit injector/receiver
|
||||
wired to the NoC. Each QPair holds the local pointers
|
||||
(\texttt{my\_head}, \texttt{my\_tail}), shadowed views of the peer's
|
||||
(\texttt{peer\_head\_cache}, \texttt{peer\_tail\_cache}), the local and
|
||||
peer rx-buffer physical bases, ring depth and slot size (both
|
||||
power-of-two), and the peer's credit-target address. The ring data
|
||||
itself lives in a reserved slot region of TCM (or PE-local HBM, or
|
||||
cube-shared SRAM), addressed by the QPair registers; PE\_IPCQ never
|
||||
touches the bytes, only the pointers. The PE\_CPU sees the controller
|
||||
as an MMIO peripheral and the N/S/E/W direction labels as logical
|
||||
ports---one kernel image runs across a 1D ring, 2D mesh, or 2D torus
|
||||
because the topology only changes which peers the QPair registers
|
||||
point at.
|
||||
|
||||
\emph{Initialization.} The host CCL backend brings the whole machine
|
||||
to a usable state before any kernel runs.
|
||||
\texttt{init\_process\_group(backend="ahbm")} loads \texttt{ccl.yaml},
|
||||
resolves the algorithm + topology + buffer\_kind + slot configuration,
|
||||
allocates an rx ring-buffer region on every participating PE so every
|
||||
rank now knows every other rank's \texttt{rx\_base\_pa}, and fans out
|
||||
an \texttt{IpcqInitMsg} that writes the QPair register file on each
|
||||
PE\_IPCQ over MMIO. The same fan-out wires the per-direction
|
||||
credit-return channel: each PE\_IPCQ records its peer's credit-target
|
||||
address and a back-pointer that the credit receiver will use to update
|
||||
the right QPair. After init, every register the runtime needs is
|
||||
preloaded; nothing in the kernel path allocates memory, walks a table,
|
||||
or talks to the host.
|
||||
|
||||
\emph{Send.} When the kernel executes \texttt{tl.send(dir="E",
|
||||
src\_addr, nbytes)}, the PE\_CPU performs a single MMIO write into
|
||||
PE\_IPCQ describing the request (direction, source address/space,
|
||||
length, sender handle). The controller evaluates backpressure in one
|
||||
combinational compare---\texttt{(my\_head $-$ peer\_tail\_cache) $<$
|
||||
n\_slots}---and, on a hit, the slot-address generator returns
|
||||
\texttt{dst = peer\_rx\_base\_pa + (my\_head \% n\_slots) $\times$
|
||||
slot\_size} in one to two cycles. PE\_IPCQ then emits one
|
||||
\texttt{IpcqDmaToken} on its dedicated port to PE\_DMA's
|
||||
\texttt{vc\_comm} channel, carrying the data descriptor (\texttt{src,
|
||||
dst, nbytes}) plus a small piggyback header (\texttt{sender\_seq =
|
||||
my\_head, src\_coord, direction}). \texttt{my\_head} is incremented in
|
||||
the same cycle---a local flip-flop bump, not a cross-PE atomic---and
|
||||
\texttt{tl.send} returns fire-and-forget. If the backpressure compare
|
||||
fails, the controller stalls the CPU in either of two modes selected
|
||||
at init: \texttt{poll} (CPU re-reads a status CSR) or \texttt{sleep}
|
||||
(controller asserts a wake event when a credit arrives). Both modes
|
||||
are benchmarked in the results.
|
||||
|
||||
\emph{Transport.} PE\_DMA is the only block that touches the fabric,
|
||||
and it has been extended for IPCQ in two ways. First, it now exposes
|
||||
two virtual channels---\texttt{vc\_compute} for GEMM/Math
|
||||
\texttt{TileToken}s and \texttt{vc\_comm} for \texttt{IpcqDmaToken}s
|
||||
---with independent state machines and a chunk-level
|
||||
(\SI{256}{\byte}) weighted round-robin arbiter on the shared physical
|
||||
link. A large GEMM tile DMA can no longer monopolize the link against
|
||||
a pending IPCQ send, enabling compute and communication to make
|
||||
progress independently as an architectural property of the design. Second, on the sender side PE\_DMA packs the piggyback
|
||||
header into the same flit train as the data, and on the receiver side
|
||||
it runs the I6 \emph{atomic terminal handler}: pay the per-flit
|
||||
bottleneck-BW drain, then, in one indivisible block, (i) write the
|
||||
payload into \texttt{MemoryStore} at \texttt{dst\_addr} (the
|
||||
receiver's ring slot) and (ii) forward an \texttt{IpcqMetaArrival}
|
||||
$\{\texttt{sender\_seq, dst\_addr}\}$ to the local PE\_IPCQ over a
|
||||
PE-internal wire. No yield, no PE\_CPU interaction, no cache-coherence
|
||||
round trip---the bytes land in TCM and the metadata reaches the
|
||||
controller in the same cycle.
|
||||
|
||||
\emph{Receive and credit return.} PE\_IPCQ's Meta Extractor
|
||||
range-matches the incoming \texttt{dst\_addr} against each direction's
|
||||
$[\texttt{rx\_base}, \texttt{rx\_base} + \texttt{n\_slots} \times
|
||||
\texttt{slot\_size})$ window---unambiguous even when two directions
|
||||
share a peer, as in a 2-rank bidirectional ring---and updates
|
||||
\texttt{peer\_head\_cache[d] := max(prev, sender\_seq + 1)}, releasing
|
||||
any \texttt{tl.recv} blocked on direction $d$. When the kernel
|
||||
eventually consumes the slot, the controller increments
|
||||
\texttt{my\_tail} and emits a \SI{16}{\byte} credit packet on the
|
||||
dedicated fast-path channel preloaded at init; the latency is the
|
||||
real fabric path (per-node overhead + edge propagation + 16 B /
|
||||
bottleneck BW), so an in-cube credit returns faster than a cross-SIP
|
||||
credit and the model carries no magic constants. The sender's PE\_IPCQ
|
||||
absorbs the credit, advances \texttt{peer\_tail\_cache}, and
|
||||
de-asserts backpressure if it was stalled. The pointer-synchronization
|
||||
problem the software baseline solved with explicit atomic RMWs and
|
||||
polling loops has been split in two and dissolved into existing
|
||||
traffic: \emph{head} updates ride on the DMA payload itself, with the
|
||||
receiver's PE\_DMA performing the data + metadata write in a single
|
||||
atomic step, so the sender never blocks waiting for the receiver to
|
||||
see it; \emph{tail} updates ride on a dedicated 16 B credit on a
|
||||
side-channel, with no software in the loop. Send becomes a single
|
||||
MMIO write, receive becomes a flip-flop read, backpressure becomes
|
||||
one 64-bit subtract, and the per-message cost in IPCQ is dominated by
|
||||
the link traversal rather than by metadata bookkeeping---which is
|
||||
what the H2 2025 measurement showed the pure-software queue could
|
||||
not achieve. On top of this substrate the collective runs a
|
||||
hierarchical local-reduce / global all-reduce-broadcast schedule
|
||||
across whatever inter-device topology the configuration specifies,
|
||||
with the IPCQ ring buffer placed in on-PE TCM, PE-local HBM, or
|
||||
cube-shared SRAM---the third knob the results section sweeps.
|
||||
|
||||
\subsection{Results}
|
||||
|
||||
Following the milestone-evaluation convention, the collective sweep builds
|
||||
its own six-device (six-SIP, $2\times3$) configurations---distinct from
|
||||
the two-SIP default of Table~\ref{tab:hw}---and measures all-reduce
|
||||
latency as a function of payload size for three inter-device topologies:
|
||||
a 1D ring, a 2D mesh (no wrap), and a 2D torus. Table~\ref{tab:allreduce}
|
||||
and Figure~\ref{fig:allreduce-cmp} report the result.
|
||||
All measurements in this section run on the PE\_IPCQ substrate
|
||||
described above; the topology sweep is intended to characterize how
|
||||
effectively the proposed mechanism exposes the underlying interconnect's
|
||||
properties, not to compare PE\_IPCQ against an alternative
|
||||
communication primitive. Following the milestone-evaluation convention,
|
||||
the collective sweep builds its own six-device (six-SIP, $2\times3$)
|
||||
configurations---distinct from the two-SIP default of
|
||||
Table~\ref{tab:hw}---and measures all-reduce latency as a function of
|
||||
payload size for three inter-device topologies: a 1D ring, a 2D mesh
|
||||
(no wrap), and a 2D torus. Table~\ref{tab:allreduce} and
|
||||
Figure~\ref{fig:allreduce-cmp} report the result.
|
||||
|
||||
\begin{table}[t]
|
||||
\centering
|
||||
@@ -64,13 +163,19 @@ per-PE payload. Lower is better; the torus wins at every size.}
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{allreduce_comparison.png}
|
||||
\caption{All-reduce latency vs.\ per-PE payload for the three topologies,
|
||||
against the analytic torus model and an external full-system simulator
|
||||
(FSIM) reference. The measured torus tracks the analytic curve within a
|
||||
small constant factor; the FSIM single-device point
|
||||
(\SI{366}{\micro\second}) sits an order of magnitude above the
|
||||
KernBench algorithmic latency, illustrating the difference between an
|
||||
achievable-kernel number and a full end-to-end-stack number.}
|
||||
\caption{All-reduce latency vs.\ per-PE payload for the three PE\_IPCQ
|
||||
topologies, against the analytic torus model. The isolated point in
|
||||
the top panel (\SI{366}{\micro\second}) is the only data point
|
||||
available from the H2 2025 software-queue measurement campaign---an
|
||||
intra-device all-reduce across 16 CUBEs on a single device. A true
|
||||
6-device inter-SIP measurement under the same software queue was not
|
||||
collected, so this point in fact \emph{understates} the SW-queue cost
|
||||
for the multidevice configuration KernBench measures here; even so the
|
||||
proposed PE\_IPCQ inter-device curves sit roughly an order of
|
||||
magnitude below it, which is the headline SW-vs-HW comparison this
|
||||
section makes. The measured torus tracks the analytic curve within a
|
||||
small constant factor, the gap reflecting real link serialization that
|
||||
the analytic model idealizes away.}
|
||||
\label{fig:allreduce-cmp}
|
||||
\end{figure}
|
||||
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
\section{Fused Grouped-Query Attention}
|
||||
\label{sec:gqa}
|
||||
|
||||
\subsection{Why it is needed}
|
||||
|
||||
Attention is the 1H focus, and it is where the two preceding optimizations
|
||||
have to come together. Grouped-Query Attention (GQA) shrinks the KV cache
|
||||
by sharing each KV head across a group of query heads (here $h_q=8$ query
|
||||
@@ -16,12 +14,31 @@ realizing it as a fast \emph{fused} kernel needs both building blocks from
|
||||
this report: efficient GEMM issue (\S\ref{sec:gemm}) for the
|
||||
$Q\!\cdot\!K^{\top}$ and $P\!\cdot\!V$ products, and an efficient on-device
|
||||
reduction (\S\ref{sec:allreduce}) for the multi-user and
|
||||
sequence-parallel KV reductions. This section is the capstone: the fused
|
||||
sequence-parallel KV reductions. \emph{Fused} here is meant in the
|
||||
FlashAttention sense---$Q\!\cdot\!K^{\top}$, the online softmax, and
|
||||
$P\!\cdot\!V$ collapse into a single kernel that never materializes the
|
||||
score matrix---and, beyond that, the cross-device KV reduction is absorbed
|
||||
into the same kernel (on PE\_IPCQ) rather than issued as a separate
|
||||
all-reduce. This section is the capstone: the fused
|
||||
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{Design}
|
||||
\subsection{Data Placement Policy}
|
||||
\label{sec:gqa-placement}
|
||||
|
||||
% TODO: compare data-placement options that apply across both the
|
||||
% short- and long-context regimes. Candidate axes:
|
||||
% - KV cache: per-CUBE shard vs. replicate; per-PE shard vs. replicate
|
||||
% - Q / W_qkv / W_o weights: static partition across CUBEs and PEs
|
||||
% - Workspace (m, l, O softmax state): scratch arena placement
|
||||
% The 4-case taxonomy used for long-context decode in
|
||||
% \S\ref{sec:gqa-long} (Cube-{SP,Repl} x PE-{TP,SP}) is one instantiation
|
||||
% of this framework; the short-context mapping in \S\ref{sec:gqa-short}
|
||||
% is another.
|
||||
|
||||
\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
|
||||
@@ -41,7 +58,22 @@ that restructures the decode step into two stateful composites (a named
|
||||
\textsf{softmax\_merge} recipe) is designed but not yet wired into the
|
||||
measured path; results below reflect the implemented kernel only.
|
||||
|
||||
\subsection{Results: long-context decode and parallelism strategies}
|
||||
% TODO: CUBE <-> KV-head mapping diagram for the short-context regime
|
||||
% (h_kv=8 KV heads -> 8 CUBEs, 1:1; intra-CUBE PE usage).
|
||||
% Bench code: src/kernbench/benches/gqa_helpers/short_ctx/
|
||||
|
||||
% TODO: prefill performance figure (latency, stage breakdown).
|
||||
% TODO: decode performance figure (latency, stage breakdown).
|
||||
% Bench output for short_ctx to be generated.
|
||||
|
||||
\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.
|
||||
|
||||
The four headline panels above stress the kernel at moderate context
|
||||
lengths. Long-context decode---the regime where KV cache size, not
|
||||
@@ -116,7 +148,8 @@ on-device collective traffic that PE\_IPCQ and the torus links of
|
||||
``slower'' strategy is in fact the one that fully cashes in the
|
||||
communication-side codesign work of this report.
|
||||
|
||||
\subsection{Analysis and meaning}
|
||||
\subsection{Comprehensive Analysis}
|
||||
\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
|
||||
@@ -133,3 +166,9 @@ 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.
|
||||
|
||||
% 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).
|
||||
|
||||
@@ -40,7 +40,10 @@ hardware levers for this workload:
|
||||
GQA panels leave the GEMM and vector-math engines two to three orders of
|
||||
magnitude below the DMA engine in busy time. Adding MAC area would not move
|
||||
decode latency; the workload cannot use it. This is the single most
|
||||
actionable finding for an attention-dominated roadmap.
|
||||
actionable finding for an attention-dominated roadmap. The implication is
|
||||
that future hardware investment should prioritize communication and
|
||||
memory-system efficiency over additional compute throughput for
|
||||
attention-dominated inference workloads.
|
||||
|
||||
\paragraph{Caveats.} These conclusions are achievable-kernel results from a
|
||||
deterministic model, not E2E measurements; absolute numbers carry the
|
||||
|
||||
Reference in New Issue
Block a user