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

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

\subsection{Why it is needed}

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 three operand-staging variants
(\textsf{ref\_ref}, both operands streamed from HBM; \textsf{load\_ref},
one operand resident in TCM; \textsf{load\_load}, both resident).
Figure~\ref{fig:gemm-util} reports MAC utilization and efficiency, and
Figure~\ref{fig:gemm-stages} breaks the kernel into per-stage engine
busy time.

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gemm_mac_utilization_theoretical_vs_measured.png}
\caption{GEMM MAC utilization and efficiency, theoretical vs.\ measured.
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{23}{\percent} at one tile to \textasciitilde%
\SI{78}{\percent} measured at 48 tiles ($K{=}3072$)---and the measured
bars track the analytic prediction within
\SIrange{10}{20}{\percent}.}
\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)
the Fetch and GEMM stages are large and comparable
(\textasciitilde\SI{770}{} and \SI{785}{\nano\second}) while DMA-out is
negligible---a compute-rich, well-pipelined regime. For the low-reuse
shape ($M{=}128,K{=}8,N{=}128$) DMA-out grows to
\textasciitilde\SI{350}{\nano\second} and compute is small---a
data-movement-bound regime.}
\label{fig:gemm-stages}
\end{figure}

Two regularities stand out. (i) Utilization is governed by how completely
the problem fills the MAC tile: the three under-tile shapes are hard-capped
well below \SI{100}{\percent}, independent of how the kernel is issued.
(ii) Among tile-filling shapes, efficiency is governed by tile count---more
tiles amortize the one-time pipeline fill and issue cost, so the deep-$K$
shape reaches \textasciitilde\SI{78}{\percent} of peak while a single-tile
shape reaches only \textasciitilde\SI{23}{\percent}. The stage breakdown
explains why: with 48 tiles the GEMM and Fetch stages overlap and stay
busy, whereas the low-reuse shape spends most of its time moving data in
and out.

\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 issue cost once per GEMM and chaining tile-stages
through token self-routing, it lets the scheduler keep the pipeline full,
so compute-rich shapes actually reach the efficiency their arithmetic
intensity allows ($\sim$\SI{78}{\percent} measured at 48 tiles), and
data-bound shapes actually reach their DMA bound instead of stalling on
command overhead. The close agreement between measured and theoretical
efficiency (Figure~\ref{fig:gemm-util}) is also the report's primary
validation that KernBench's latency model is faithful in the regime that
matters. The hardware implication is concrete: a single-command,
self-routing tile pipeline is the issue mechanism that makes the MAC array
usable, and it is a prerequisite---not a luxury---for the fused attention
kernel of \S\ref{sec:gqa}.