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 tile pipeline
\[
\textsf{DMA\_READ}\rightarrow\textsf{FETCH}\rightarrow\textsf{GEMM}
\rightarrow\textsf{STORE}\rightarrow\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 DMA\,$\rightarrow$\,fetch\,$\rightarrow$\,GEMM\,$%
\rightarrow$\,store 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.