ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

paper(1H): reflect D8 single-op cost model in §2/§3.4 + figure/diagram regen

Browse Source 
Two strands bundled as the 1H-codesign-paper refresh unit:

(A) This session — single-op cost-model reflection (depends on 2d8271c):
  - §2 Table 2 (tab:hw): split "FIXED per command" into "FIXED per
    single-op command" (8 cycles) and "FIXED per composite command"
    (40 cycles); §2 dispatch-overhead prose updated to the two-class split.
  - §3.4 (sec:gemm-vs-async): rename paragraph headers + prose to
    async-full / async-tiled; "atomic" -> "single-op" throughout; reframe
    mechanism #3 from the old DMA-only fast-path to the single-op
    fast-path. Headline narrative now: even with EVERY single-op cmd
    (96 DMA + 48 dot + 47 add) charged the light 8-cycle FIXED, composite
    still wins ~2.8x at K=3072 purely on command-count structure (1 vs
    192 commands) -- down from the pre-D8 ~6.3x, and explicitly NOT a
    modelling artifact. Numbers refreshed from the regenerated sweep:
    async-full 3.83->3.91, async-tiled 1.14->~2.53, under-tile corner
    1.06->1.21, depth-2 vs depth-inf spread <1%. New figure wired in.
  - build/main.pdf rebuilt (tectonic); pdftotext-verified (no broken
    refs; Table 2 split, single-op terms, 2.8x/2.53/192-host-commands
    all present).

(B) Prior-session paper work riding along uncommitted: §4 all-reduce
    deep-edit, §5 GQA, §6 discussion trims; milestone_1h_ccl.py plot
    label "FSIM" -> "H2 2025 SW queue baseline"; regenerated diagrams
    under docs/diagrams/** and gemm output PNGs under
    1H_milestone_output/gemm/. (Composite-window gemm plots are
    unaffected by D8 — D8 only changes single-op dispatch FIXED, which
    the composite window excludes.)

All TODO items for the D8 single-op extension are now complete and
pushed across 3 commits (2d8271c cost-model+ADR+tests, 821bbf2 bench
harness, this paper refresh). Full regression green (826 passed, 1
skipped). No remaining work.

NOTE for review (carried from 2d8271c): ADR-0065's "2x CPU-offload win"
headline for GQA decode opt2 may want a refresh to the post-D8 ~1.87x.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook Kang
2026-06-17 15:12:31 -07:00
parent 821bbf26a2
commit b6315c3c90
31 changed files with 568 additions and 185 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/report/1H-codesign-paper/sections/03-gemm.tex
+228
View File
@@ -171,3 +171,231 @@ 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 user-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 user-level pattern --- async-load each operand,
overlap with compute, accumulate --- sufficient?
To answer this concretely we contrast composite against two
user-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 user-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
user-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$750 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{12}{\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 user-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 user-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
user 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.
User-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.