kernbench2/docs/report/1H-codesign-paper/sections/04-allreduce.tex

\section{PE\_IPCQ and Collective Communication}
\label{sec:allreduce}

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

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}

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
\caption{All-reduce latency (ns) across six devices, by topology and
per-PE payload. Lower is better; the torus wins at every size.}
\label{tab:allreduce}
\small
\begin{tabular}{@{}rrrr@{}}
\toprule
\textbf{Bytes/PE} & \textbf{2D mesh} & \textbf{Ring 1D} & \textbf{2D torus} \\
\midrule
256     & 4189  & 3883  & 2957 \\
4{,}096  & 5566  & 5240  & 4031 \\
16{,}384 & 10016 & 9376  & 7396 \\
65{,}536 & 27821 & 25925 & 20858 \\
98{,}304 & 39690 & 36957 & 29833 \\
\bottomrule
\end{tabular}
\end{table}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{allreduce_comparison.png}
\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}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{allreduce_buffer_kind.png}
\caption{Effect of IPCQ staging-buffer placement (2D torus). At
\SI{64}{\kibi\byte}/PE, TCM staging (\SI{20858}{\nano\second}) beats HBM
(\SI{24074}{\nano\second}) by \textasciitilde\SI{13}{\percent} and SRAM
(\SI{33194}{\nano\second}) by \textasciitilde\SI{37}{\percent}; at small
payloads the three are indistinguishable.}
\label{fig:allreduce-buf}
\end{figure}

Three findings follow. First, topology matters and the effect grows with
size: at \SI{96}{\kilo\byte}/PE the 2D torus is \SI{25}{\percent} faster
than the mesh and \SI{19}{\percent} faster than the ring, because its
wrap-around links shorten the worst-case reduction path. Second, the
measured curves grow smoothly with payload and stay within a small
constant factor of the analytic torus model, with the gap reflecting real
link serialization that the analytic model idealizes away. Third, where
the IPCQ staging buffer lives is a first-order knob: keeping it in on-PE
TCM is materially faster than HBM or shared SRAM at large payloads
(Figure~\ref{fig:allreduce-buf}).

\subsection{Analysis and meaning}

PE\_IPCQ turns the collective from an off-device, contention-prone
operation into an on-device primitive whose latency tracks the
interconnect's physical limits. The control/data split keeps the engine
small while reusing PE\_DMA for movement, and the virtual-channel split is
what lets a reduction proceed without stalling the compute DMA that feeds
the very GEMMs producing the partials---the same property the fused
attention kernel relies on when it interleaves KV reduction with score
computation (\S\ref{sec:gqa}). For the hardware roadmap the results argue
two things: provisioning wrap-around (torus) inter-device links is worth
roughly a \SI{20}{}--\SI{25}{\percent} collective-latency reduction at
scale, and giving the IPCQ a fast on-PE staging buffer (TCM) rather than
forcing it through HBM or SRAM is worth another \SI{14}{}--%
\SI{38}{\percent} at large payloads.