kernbench2/docs/report/1H-codesign-paper/sections/02-platform.tex

\section{The KernBench Platform}
\label{sec:platform}

All results in this report are produced on \emph{KernBench}, a
system-level discrete-event simulator for LLM kernels running on AHBM.
This section explains why the platform exists, the device and
execution model it presents to a kernel writer, how its latency model
turns that execution into a number, and the concrete hardware
configuration used for every experiment that follows.

\subsection{Why KernBench}
\label{sec:why}

In a production end-to-end (E2E) stack, kernel performance is entangled
with every layer above the hardware: the compiler's tiling and
scheduling choices, the framework's operator dispatch, the collective
library, and the runtime. Good E2E numbers require \emph{all} of those
layers to be co-optimized, which makes it hard to answer a narrower but
more fundamental question: \emph{given the hardware, how fast can a
well-written kernel be, and which hardware features actually make it
faster?}

KernBench is built to answer exactly that question. Kernels are written
and executed at the \emph{source level}---as algorithmic descriptions
in a small tile-oriented kernel API---with no dependency on a compiler
or any other software-stack layer. The simulator takes the kernel and a
hardware topology and reports the latency the modeled hardware would
deliver. This isolation is deliberate: it lets us study algorithm-level
optimizations (how to tile a GEMM, how to schedule a collective, how to
fuse an attention kernel) and the hardware features that support them,
without the confound of compiler maturity or framework overhead. The
cost is that KernBench numbers are \emph{not} E2E latencies; they are
the achievable-kernel latencies an ideal software stack would expose.

\subsection{Device and execution model}
\label{sec:exec}

KernBench models AHBM as a hierarchy of SIPs, CUBEs, and processing
elements (PEs). At the system level, multiple SIPs are connected
through inter-package links, while each SIP contains a collection of
CUBEs joined by an on-package interconnect
(Fig.~\ref{fig:sip-arch}).

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{sip_architecture.pdf}
\caption{Modeled hardware graph at the \emph{SIP} level (one
example configuration; specific parameters in \S\ref{sec:hw} /
Table~\ref{tab:hw}). The SIP holds a $4{\times}4$ mesh of CUBEs and
an IO chiplet, with each line a directed link labelled by its
physical distance and bandwidth.}
\label{fig:sip-arch}
\end{figure}

Each CUBE contains eight PEs, shared SRAM, HBM controllers, an
\textsf{M\_CPU} control processor, and an intra-CUBE router mesh
(Fig.~\ref{fig:cube-arch}). Together these components form the
execution substrate for all kernels evaluated in this report. This
organization reflects the memory-centric nature of AHBM: each PE is
paired with a dedicated slice of HBM bandwidth and local on-PE
storage (TCM), so compute lives next to the data it consumes rather
than fetching it through a far-away memory controller. The platform's
performance question is therefore not ``how many FLOPs can the chip
do'' but ``how well can a kernel keep each compute engine fed from
its locally-attached memory while moving the unavoidable traffic
between PEs efficiently.''

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{cube_architecture.pdf}
\caption{Modeled hardware graph at the \emph{CUBE} level---a
zoom-in on one CUBE node from Fig.~\ref{fig:sip-arch}. The CUBE
holds 8 PEs (each with its own HBM channels), the \textsf{M\_CPU}
control processor, a shared SRAM, an intra-CUBE NoC router mesh, and
UCIe links to neighbouring CUBEs.}
\label{fig:cube-arch}
\end{figure}

Within a PE, commands are dispatched by \textsf{PE\_CPU} to
\textsf{PE\_SCHED}, which routes work to specialized execution
engines---DMA, FETCH/STORE, GEMM, vector-math, and IPCQ
(Fig.~\ref{fig:pe-arch}). Commands come in two flavours. \emph{Atomic}
commands target a single engine---a plain DMA read, a GEMM tile, a
vector-math op, an IPCQ send/receive---and are the natural unit for
short or special-purpose work; the \textsf{PE\_CPU} itself also runs
control-plane work directly. \emph{Composite} commands, by contrast,
carry an ordered pipeline of operations across multiple engines
(\textsf{DMA\_READ} $\rightarrow$ \textsf{FETCH} $\rightarrow$
\textsf{GEMM}/\textsf{MATH} $\rightarrow$ \textsf{STORE} $\rightarrow$
\textsf{DMA\_WRITE}) that the \textsf{PE\_SCHED} tiles and streams
without per-tile redispatch. The composite form is the substrate for
the GEMM optimization (\S\ref{sec:gemm}) and, combined with on-PE
collectives, for fused attention (\S\ref{sec:gqa}).

\begin{figure*}[t]
\centering
\includegraphics[width=\linewidth]{pe_architecture.png}
\caption{Modeled PE architecture used in this report---one example
configuration whose specific parameters are listed in
\S\ref{sec:hw} (Table~\ref{tab:hw}); KernBench is not tied to this
particular decomposition and supports arbitrary PE block layouts
provided each component is given a port and bandwidth model.
\textsf{PE\_CPU} dispatches commands to the \textsf{PE\_SCHED}, which
routes tile-token streams through the \textsf{PE\_DMA},
\textsf{PE\_FETCH\_STORE}, and \textsf{GEMM}/\textsf{MATH} engines
along on-chip links; the \textsf{PE\_IPCQ} provides the control plane
for on-device collective communication.}
\label{fig:pe-arch}
\end{figure*}

KernBench is layered along the flow of a request:

\begin{itemize}
  \item The \textbf{runtime API} is host-facing and
  topology-agnostic---it deploys tensors and launches kernels but knows
  nothing about routing or interconnect.
  \item The \textbf{simulation engine} schedules discrete events and
  routes every request through the modeled graph.
  \item The \textbf{components} are device-side nodes that model
  hardware behaviour: the per-PE blocks (scheduler, DMA, GEMM and
  vector-math engines, TCM, IPCQ), the NoC routers, the HBM
  controllers, and the inter-chiplet links.
\end{itemize}

Data and timing are handled in two passes, so that a kernel's numeric
results and its latency are computed consistently but independently.

\subsection{Latency model: graph traversal and contention}
\label{sec:latency}

The modeled hardware hierarchy described above is represented
internally as a directed graph. Nodes correspond to hardware
components---PEs, routers, memory controllers, SRAM blocks, IO
chiplets---while edges represent communication links with associated
bandwidth (\si{\giga\byte\per\second}) and propagation
(\si{\nano\second}) attributes. Every operation in KernBench---DMA
transfers, remote memory accesses, collective communication, command
dispatch---is modelled as a traversal through this graph. End-to-end
latency is decomposed into four contributions accumulated along the
traversal path (Fig.~\ref{fig:latency-model}): \textbf{per-node
overhead} at each component, \textbf{per-edge transmission} on each
wire, \textbf{drain} (per-flit service occupancy) at the destination,
and \textbf{queuing delay} at the shared FIFOs that the wire and the
destination share between concurrent transactions.

\begin{figure*}[t]
\centering
\includegraphics[width=\textwidth]{latency_model.png}
\caption{Conceptual schematic of the latency model. Two source nodes
(Requester A/B) inject flits through a chain of routers into a
destination node; on the shared edge between routers, flits from the
two transactions are interleaved flit-by-flit by the wire's FIFO
arrival order. End-to-end latency is the sum of four contributions:
\textbf{per-node overhead} (the switch's fixed processing cost,
shown in light yellow---the same colour as the destination's
processing-logic block), \textbf{per-edge transmission}
($\textit{flit\_size}/\textit{BW}$ on each wire),
\textbf{drain} (per-flit service occupancy at the destination's
channel), and \textbf{queuing delay} (waiting in a FIFO when a
shared resource is busy). The places where queuing actually
accumulates are highlighted in green: the router output queue and
the destination's input queue. This is the model, not a
measurement; specific bandwidths and overheads are listed in
\S\ref{sec:hw} (Table~\ref{tab:hw}).}
\label{fig:latency-model}
\end{figure*}

\paragraph{The hardware as a graph.} The topology is compiled once at
configuration time into this graph and is never mutated during a run.
There are no hidden shortcuts, implicit bypasses, or magic paths: if a
request reaches its destination, the path it took is explicit in the
graph, and the latency it incurred is the sum of the per-node and
per-edge costs paid along that path. The same graph representation
applies recursively at every hierarchy level---system, SIP
(Fig.~\ref{fig:sip-arch}), CUBE (Fig.~\ref{fig:cube-arch}), and PE
(Fig.~\ref{fig:pe-arch}).

\paragraph{From graph to discrete-event simulation.} The graph is
driven by a discrete-event engine. Two kinds of events advance
simulation time: \emph{node events} (component switching overhead,
service completions such as an HBM channel commit or a GEMM tile
finish) and \emph{edge events} (the flit-by-flit serialization of a
payload across a bandwidth-limited link). The engine maintains a
priority queue of pending events ordered by their scheduled time and
fires them one at a time, with ties broken under a deterministic
policy so that the same kernel on the same topology always yields the
same trace. Per-request correlation IDs are stamped at injection and
carried through every hop, so the path from injection to completion
is fully traceable. Every nanosecond in a reported latency
corresponds to exactly one of these events on exactly one node or
edge---there is no slack in the budget.

\paragraph{Latency contributions.} Three kinds of latency accumulate
along a traversal: (i) \emph{per-node fixed overhead}---each component
carries a small switching cost (router decode, controller pickup,
scheduler handoff); (ii) \emph{per-edge transfer time}---each link's
payload is decomposed into fixed-size flits (default
\SI{256}{\byte}), and each flit arrives at
$\text{prop}+\text{flit\_bytes}/\text{bw}$ after the previous one,
giving wormhole semantics across multi-hop paths; and (iii)
\emph{per-service occupancy}---memory controllers, GEMM stages, and
collective engines hold the request for their service time before
releasing it downstream. Each of these is attached to a specific node
or edge in the graph; together they make up the entire latency budget.

\subsubsection{Congestion and contention modeling}
\label{sec:congestion}

The simulator's
sharpness comes from how it models contention for those nodes and
edges. \emph{Every directed edge has a FIFO}: an arriving flit takes
its bandwidth-limited transfer time on top of whatever earlier flits
are still being served, so a busy link queues later traffic behind
earlier traffic rather than transferring everything at peak BW.
\emph{HBM is modelled with per-pseudo-channel parallelism}: a
stateless array of channel-availability timestamps with address-based
channel selection captures the bank-level concurrency that real HBM
exposes, so 64 channels per CUBE deliver real parallelism on uniform
addresses but a hot channel surfaces as the bottleneck on skewed ones.
\emph{Every component has a serial worker}: a router carrying two
heavy streams interleaves them at flit granularity in arrival order
rather than fanning out for free, so two concurrent collectives sharing
a link share its bandwidth, not double it. Without these mechanisms
the simulator would simply re-confirm the peak-BW roofline; with them,
it reveals where the real bottlenecks form and which hardware levers
actually relieve them---which is exactly the question the codesign
work in this report turns on.

\paragraph{Command dispatch overhead model.} The cost incurred by
the \textsf{PE\_CPU} when it dispatches a command to one of the
accelerator engines is modelled structurally rather than with a
per-operation calibration table. The \textsf{PE\_CPU} charges, per
command,
\[
  d_{\text{cmd}} = \textsf{FIXED} + b_{\text{logical}} \cdot R,
\]
where $b_{\text{logical}}$ is the command's hardware-logical byte size,
\textsf{FIXED} captures the fixed per-command cost (queue-tail update,
completion registration) and $R$ captures the per-byte cost of
serializing the command descriptor into the scheduler queue. 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
not the bottleneck---data movement is---so this term stays small
relative to DMA and collective time.

\paragraph{Accuracy.} The model is precise about the effects that
dominate kernel latency on this class of hardware: per-edge bandwidth
occupancy and flit-level serialization, HBM pseudo-channel parallelism,
and per-component switching overhead. Two independent cross-checks
drawn from the experiments in this report confirm that this precision
translates into physically reasonable kernel latencies. First, in the
GEMM study (\S\ref{sec:gemm}), simulator-measured MAC efficiency
tracks an analytic ideal-pipeline model within roughly
\SIrange{10}{20}{\percent} across a wide range of tile counts; the
residual gap is attributable to pipeline-fill and DMA effects the
analytic model omits. Second, in the all-reduce study
(\S\ref{sec:allreduce}, Fig.~\ref{fig:allreduce-cmp}), simulator
latency for a 2D-torus over six devices follows the expected
startup-plus-per-packet shape across the entire payload sweep---tight
at small payloads where startup dominates, and within a single-digit
multiplicative factor at the largest payloads, where the residual gap
is explained by per-router switching the analytic shape elides. A
single-device point from an external full-system simulator (FSIM) at
the largest payload sits an order of magnitude above the KernBench
multi-device torus, illustrating the well-known gap between an
achievable-kernel number and a full end-to-end-stack number rather
than a model error. The known simplifications---FIFO router
arbitration (instead of round-robin), HBM scheduler without
write-buffer reordering, no bank conflict, no refresh or thermal
effects, and no upstream backpressure---are the price of a
deterministic, inspectable model. These simplifications bound the
absolute accuracy, but the agreement with both analytic models and the
external full-system simulator data above indicates that KernBench is
sufficiently accurate for evaluating the \emph{relative}
hardware--software design trade-offs (tiling A vs.\ B, topology X
vs.\ Y, with vs.\ without composite command, mesh vs.\ torus) that
are the primary objective of this work.

A third source of confidence comes from directly probing the
simulator's per-traversal behaviour. Running \texttt{kernbench probe}
on the modelled topology issues a sequence of PE-to-HBM DMA reads at
progressively greater hop distances and reports the per-component
overhead, per-edge serialization, and per-PC drain that the model
charges (Table~\ref{tab:probe-pe-dma}). Three properties stand out.
First, the reported latency increases \emph{monotonically} with hop
count---from \SI{141}{\nano\second} at the local HBM slice to
\SI{677}{\nano\second} at the worst-case remote-CUBE slice---with the
increment per added hop matching the per-router overhead and the
UCIe-link cost in Table~\ref{tab:hw}. Second, the bandwidth-saturation
curves track the wire's serialization model: at \SI{1}{\mebi\byte}
transfers the local PE DMA reaches \SI{100}{\percent} of the per-edge
bandwidth limit, while the longest cross-CUBE path saturates at
\SI{97.8}{\percent}---exactly the gap that the per-edge propagation
and per-router overhead model would predict. Third, the simulator
passes the internal-consistency invariants the probe builds in
(monotonic hop progression; D2H reads at least as long as the
equivalent H2D writes because of the reverse-path acknowledgement;
cross-CUBE best less than cross-CUBE worst). These are properties
that hold \emph{only} when the underlying graph traversal, FIFO
contention, and propagation models are internally
self-consistent---a model error in any of them would surface as a
non-monotonic or under-utilising curve here long before it polluted a
kernel-level measurement.

\begin{table}[t]
\centering
\caption{PE\,$\to$\,HBM DMA latency probe at varying hop distances
(\SI{32}{\kibi\byte} transfer; output captured from \texttt{kernbench
probe}). \emph{Util\%} is the achieved bandwidth as a fraction of the
path's bottleneck-edge bandwidth.}
\label{tab:probe-pe-dma}
\small
\begin{tabular}{@{}lrrr@{}}
\toprule
Case & Latency~(\si{\nano\second}) & Util\% (\,32\,KiB) & Util\% (\,1\,MiB) \\
\midrule
PE\,$\to$\,local HBM             & 141.0 & 90.8 & 100.0 \\
PE\,$\to$\,same-half HBM         & 147.9 & 86.6 & ~99.9 \\
PE\,$\to$\,cross-half HBM        & 161.2 & 79.4 & ~99.8 \\
PE\,$\to$\,cross-CUBE (best)     & 330.5 & 77.5 & ~99.6 \\
PE\,$\to$\,cross-CUBE (worst)    & 677.1 & 37.8 & ~97.8 \\
\bottomrule
\end{tabular}
\end{table}

\subsection{Modeled hardware configuration}
\label{sec:hw}

Table~\ref{tab:hw} summarizes the hardware configuration
used throughout this report. The intent of this
configuration is not to model a specific product, but to
represent a realistic memory-centric accelerator and to
provide a consistent baseline for evaluating hardware--
software co-design mechanisms.

The compute capability within each CUBE is provisioned
such that inference workloads can effectively saturate
the available HBM bandwidth. The aggregate compute
throughput is therefore balanced against memory-system
bandwidth rather than being intentionally over- or
under-provisioned. HBM bandwidth is distributed evenly
across the PEs within a CUBE, with each PE responsible
for servicing approximately one-eighth of the CUBE's
memory bandwidth through its dedicated HBM channels.

The on-chip interconnect is configured using bandwidth
and latency parameters representative of commercially
available mesh-network IPs. The goal is not to study a
particular NoC implementation, but rather to evaluate
kernel behavior under a realistic baseline communication
fabric.

For die-to-die communication, UCIe-A bandwidth
characteristics are used as the reference point for
inter-CUBE links. Inter-SIP communication is modeled
using PCIe-class links. Together, these assumptions
provide a representative communication hierarchy for
evaluating collective communication and distributed
kernel execution.

Unless otherwise stated, all experiments use this
configuration. Workload-specific parameters such as
matrix dimensions, sequence lengths, and collective
payload sizes are introduced in their respective sections.

\begin{table}[t]
\centering
\caption{Modeled hardware configuration (shared by all experiments).}
\label{tab:hw}
\small
\begin{tabular}{@{}ll@{}}
\toprule
\textbf{Parameter} & \textbf{Value} \\
\midrule
\multicolumn{2}{@{}l}{\emph{Hierarchy}} \\
SIPs & 2 (1D ring) \\
CUBEs per SIP & 16 ($4\times4$ mesh) \\
PEs per CUBE & 8 (4 corners $\times$ 2) \\
PEs total & 256 \\
\midrule
\multicolumn{2}{@{}l}{\emph{Processing element (PE)}} \\
GEMM engine peak & \SI{8}{\tera\flop\per\second} (f16) \\
TCM (on-PE) & \SI{16}{\mega\byte}, \SI{512}{\giga\byte\per\second} R/W \\
\quad kernel scratch & \SI{1}{\mega\byte} \\
DMA engines & 1 read + 1 write \\
CPU / scheduler overhead & \SI{2}{\nano\second} / \SI{1}{\nano\second} \\
\midrule
\multicolumn{2}{@{}l}{\emph{Memory (per CUBE)}} \\
HBM capacity & \SI{48}{\giga\byte} (8 slices) \\
HBM aggregate BW & \SI{1024}{\giga\byte\per\second} \\
HBM pseudo-channels & 64 (8 per PE), \SI{32}{\giga\byte\per\second} each \\
SRAM (shared) & \SI{32}{\mega\byte}, \SI{128}{\giga\byte\per\second} link \\
HBM burst & \SI{256}{\byte} \\
\midrule
\multicolumn{2}{@{}l}{\emph{Interconnect}} \\
Intra-CUBE NoC link & \SI{256}{\giga\byte\per\second}, \SI{2}{\nano\second}/router \\
Inter-CUBE (UCIe PHY) & \SI{512}{\giga\byte\per\second}, \SI{8}{\nano\second}, XY routing \\
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 \\
per-byte rate $R$ & 0.0625 cycles/byte (\SI{16}{\byte\per\cycle}) \\
composite size cap & \SI{1024}{\byte} \\
\bottomrule
\end{tabular}
\end{table}