kernbench2/docs/report/1H-codesign-paper/sections/03-gemm.tex

\section{GEMM Acceleration via the Composite Command}
\label{sec:gemm}

GEMM is the compute core of every transformer block---the QKV
projections, the attention score and context products, and the
feed-forward matrices are all matrix multiplications. On a tiled
accelerator a single logical GEMM expands into many hardware tiles, and
each tile must be read from HBM, fetched into the register file, computed,
stored, and written back. If every one of those tile-stage steps were an
independently issued command, two costs would dominate. First, the host
and the PE control processor would pay a per-command issue overhead
$O(\text{tiles}\times\text{stages})$ times, which for a deep-$K$ reduction
is hundreds to thousands of issues. Second, with stages dispatched
one-at-a-time the scheduler cannot overlap the DMA of the next tile with
the compute of the current one---the pipeline never fills, and the MAC
array sits idle waiting for data. The hardware question is therefore: what
issue mechanism lets a single GEMM saturate the MAC array without drowning
in command overhead?

\subsection{Design}

The answer is the \emph{composite command}. A single command carries
the ordered five-stage tile pipeline (\textsf{DMA\_READ},
\textsf{FETCH}, \textsf{GEMM}, \textsf{STORE}, \textsf{DMA\_WRITE}),
and the PE scheduler splits the payload into hardware tiles
(here $32\times64\times32$), emitting one tile token per tile.
Subsequent stages are reached by \emph{token self-routing} between
the on-PE engines, so a tile flows through the chain without
returning to the scheduler between stages. Because the whole
pipeline is described by one command, the issue cost is paid once
per GEMM rather than once per tile-stage, and the scheduler is free
to keep every stage busy on different tiles simultaneously---tile
$i$'s GEMM overlaps tile $i{+}1$'s DMA read. A multi-operation composite
additionally lets an epilogue (for example a vector-math step) ride the
same tile loop, firing per $K$-tile, per output tile, or once per kernel
according to its declared scope; this is what later allows an attention
kernel to fuse its softmax work into the GEMM pipeline
(\S\ref{sec:gqa}).

\subsection{Results}

We sweep eight GEMM shapes spanning square, tall, wide, and deep-$K$
geometries under two operand-staging variants that bracket the
realistic LLM cases. In \textsf{load\_ref}, the activation $A$ is
pre-staged on chip and only the weight $W$ streams from HBM during
the composite (the ``activation in TCM, weights from HBM'' case
typical of decoding with a small batch). In \textsf{ref\_ref}, both
operands stream from HBM during the composite (the ``cold both
sides'' case, where the activation does not fit on-chip or is itself
the output of a producer composite that wrote back to HBM). Both
variants run the same composite command; only the up-front staging
of $A$ differs.

We evaluate the composite GEMM along two axes that together describe
how well it lands on the hardware: \emph{HBM bandwidth utilization}
(does the workload saturate the per-PE HBM bandwidth, i.e.\ is it
memory-bound on this configuration?) and \emph{achieved GEMM
throughput} (what fraction of the \SI{8}{\tera\flop\per\second}
per-PE peak does the kernel actually deliver). Both metrics use the
composite window as the denominator, so the up-front \textsf{tl.load}
of $A$ in \textsf{load\_ref} is excluded — only HBM traffic and
compute that happen \emph{inside} the composite count.

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gemm_hbm_bw_util.png}
\caption{Per-PE HBM bandwidth utilization during the composite window,
for the two staging variants. The ceiling is the per-PE HBM bandwidth
(\SI{256}{\giga\byte\per\second}; dashed line at \SI{100}{\percent}).
\textsf{ref\_ref} streams both $A$ and $W$ and so always sits at a
higher BW-utilization than \textsf{load\_ref} on the same shape;
both variants saturate ($>\SI{85}{\percent}$) once the kernel
has enough total bytes to keep the link busy (deep-$K$ $K{=}3072$,
or low-reuse output-dominated $M{=}128,K{=}8,N{=}128$). Small or
single-tile shapes do not have enough total traffic to saturate the
link and are bottlenecked by pipeline fill rather than HBM.}
\label{fig:gemm-bw}
\end{figure}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gemm_per_pe_tflops.png}
\caption{Per-PE GEMM throughput delivered during the composite window
(reference line at the \SI{8}{\tera\flop\per\second} per-PE engine
peak). At the deep-$K$ corner $M{=}32,K{=}3072,N{=}32$ the
activation-pre-staged kernel reaches
\SI{7.18}{\tera\flop\per\second} ($\sim\SI{90}{\percent}$ of peak);
the both-from-HBM variant drops to
\SI{3.83}{\tera\flop\per\second} ($\sim\SI{48}{\percent}$) at the
same shape because the second operand doubles HBM pressure and the
kernel hits the BW ceiling shown in Fig.~\ref{fig:gemm-bw}.
Under-tile shapes ($M{=}K{=}N{=}32$, $M{=}8$, $K{=}8$) are
hard-capped by tile-fill at \SI{50}{\percent}, \SI{25}{\percent},
\SI{12.5}{\percent} of peak regardless of staging.}
\label{fig:gemm-tflops}
\end{figure}

The composite reaches the hardware roofline in two distinct regimes.
\emph{Memory-bound}: $K{=}3072$ deep-$K$ saturates HBM (\SI{88}{\percent}
load\_ref, \SI{94}{\percent} ref\_ref) and the low-reuse
$M{=}128,K{=}8,N{=}128$ corner saturates even harder
($>\SI{96}{\percent}$ in both), because back-to-back output
writes dominate. \emph{Compute-rich}: at the same deep-$K$ shape,
\textsf{load\_ref} delivers \SI{7.18}{\tera\flop\per\second}
($\sim\SI{90}{\percent}$ of the per-PE peak) — the configuration's
clean win. The \emph{distance between the two variants} at the same
shape is the cost of \emph{not} pre-staging $A$: at $K{=}3072$ the
throughput halves (\SI{7.18}{\tera\flop\per\second} $\to$
\SI{3.83}{\tera\flop\per\second}) because the second operand doubles
HBM pressure
and the kernel hits the BW ceiling. This is the operational
take-away — when the activation can live on-chip the composite
delivers near-peak GEMM; when it cannot the BW ceiling dominates and
throughput halves.

Figure~\ref{fig:gemm-util} validates that simulator-measured
efficiency tracks an analytic ideal-pipeline model (within
\SI{1.4}{ppt} across every swept shape), and
Figure~\ref{fig:gemm-stages} resolves the per-stage engine wall-clock
that backs the above interpretation.

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gemm_mac_utilization_theoretical_vs_measured.png}
\caption{GEMM MAC utilization and efficiency, theoretical vs.\ measured
(\textsf{load\_ref} staging). Tile-fill sets the ceiling: under-tile
shapes (marked $\ast$) such as $M{=}K{=}N{=}32$, $M{=}8$, and $K{=}8$
cannot fill the MAC tile and cap at \SI{50}{\percent},
\SI{25}{\percent}, \SI{12.5}{\percent}. For tile-filling shapes,
efficiency climbs with tile count---from
\textasciitilde\SI{15}{\percent} at one tile to \textasciitilde%
\SI{90}{\percent} measured at 48 tiles ($K{=}3072$)---and the
measured bars track the analytic ideal-pipeline prediction within
\SI{1.4}{ppt} across every shape.}
\label{fig:gemm-util}
\end{figure}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gemm_stage_breakdown.png}
\caption{Per-stage engine wall-clock (DMA in, Fetch, GEMM, DMA out)
under \textsf{load\_ref} staging. For the deep-$K$ shape ($K{=}3072$,
48 tiles) DMA in, Fetch, and GEMM are all close to
\textasciitilde\SI{785}{\nano\second}, running concurrently in a
tightly pipelined regime while DMA-out is negligible. For the
low-reuse shape ($M{=}128,K{=}8,N{=}128$, 16 output tiles) DMA-out
grows to \textasciitilde\SI{336}{\nano\second} while GEMM compute is
\textasciitilde\SI{262}{\nano\second}---a data-movement-bound regime
where the output write becomes the largest stage.}
\label{fig:gemm-stages}
\end{figure}

\subsection{Analysis and meaning}

The composite command does not manufacture bandwidth or MAC
throughput; it removes the two software-shaped obstacles between a
GEMM and its hardware roofline. By paying the issue cost once per
GEMM rather than once per tile-stage, and by chaining tile-stages
through token self-routing, it lets the scheduler keep every stage
busy on a different tile, so compute-rich shapes actually reach the
efficiency their arithmetic intensity allows
($\sim$\SI{90}{\percent} of GEMM peak at the deep-$K$
\textsf{load\_ref} corner) and data-bound shapes actually reach
their HBM-BW ceiling instead of stalling on command overhead.

The split between \textsf{load\_ref} and \textsf{ref\_ref} also makes
the hardware-software boundary explicit: when the activation fits in
on-chip storage the composite is BW-headroom-rich and saturates the
GEMM engine; when it does not, both operands compete for the same
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 kernel-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 kernel-level pattern --- async-load each operand,
overlap with compute, accumulate --- sufficient?

To answer this concretely we contrast composite against two
kernel-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 kernel-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
kernel-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$781 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{10}{\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 kernel-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 kernel-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
kernel 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.
Kernel-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.