paper(gqa): two-regime "use of composite commands" — decode + compute-bound prefill
Complete the composite-command study across both attention regimes and write up the result, resolving when the composite command helps latency. Decode (memory-bound, T_q=1, M=8): the command form is latency-neutral — the kernel is bound by KV streaming and the MAC array is near-idle, so the composite's only benefit here is host-issue offload (O(n_tiles) -> O(1) PE_CPU commands). Figure: gqa_decode_long_ctx_composite (latency flat, command count saturates). Prefill (compute-bound, large M=G*T_q tile-filling): add three command-form variants of a single-rank FlashAttention prefill kernel (_gqa_prefill_compute_bound: primitive / composite / composite_extended). Here the composite's scheduler-internal per-tile DMA<->compute pipeline keeps the MAC array fed while the primitive's blocking tl.dot starves it, so composite wins on MAC utilization (68% flat -> 83%) and wall-clock (19% faster at ctx=1024), with the margin growing with context (deeper P.V reduction = more tiles to pipeline) — the compute-bound mirror of the GEMM result in section 3. Figure: gqa_prefill_compute_bound (latency + MAC util). Sweep wired as GQA_1H_SWEEPS=prefill_cb; tests cover structure, e2e, and composite<primitive in the compute-bound regime. The synthesis: composite has two benefits set by roofline position — host-issue offload (always) and MAC-array feeding (compute-bound only). Decode exercises the first, prefill both. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Binary file not shown.
|
After Width: | Height: | Size: 136 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 145 KiB |
@@ -122,9 +122,10 @@ overlaps score computation; and per-tile \emph{scratch recycling} that
|
||||
keeps the running softmax accumulators ($m,\ell,O$) in a persistent arena
|
||||
while freeing per-tile temporaries, so the kernel fits the
|
||||
\SI{1}{\mebi\byte} scratch budget across many tiles. A further refinement
|
||||
that restructures the decode step into two stateful composites (a named
|
||||
\textsf{softmax\_merge} recipe) is designed but not yet wired into the
|
||||
measured path; results below reflect the implemented kernel only.
|
||||
restructures the decode step so the per-tile matrix products and the
|
||||
online-softmax merge are issued as \emph{composite} commands rather than
|
||||
hand-tiled primitives; \S\ref{sec:gqa-composite} measures it against the
|
||||
primitive baseline at long context.
|
||||
|
||||
% TODO: CUBE <-> KV-head mapping diagram for the short-context regime
|
||||
% (h_kv=8 KV heads -> 8 CUBEs, 1:1; intra-CUBE PE usage).
|
||||
@@ -221,6 +222,128 @@ memory-feasible family, and the cross-PE softmax reduction it does pay is
|
||||
precisely the traffic the communication-side codesign of this report is
|
||||
built to move quickly.
|
||||
|
||||
\subsection{Use of Composite Commands}
|
||||
\label{sec:gqa-composite}
|
||||
|
||||
The decode kernel of \S\ref{sec:gqa-long} issues its local attention as
|
||||
primitive operations: it walks each PE's $S_{\text{local}}$ token slice in
|
||||
\SI{1024}{}-token tiles, and for every tile issues a $Q\!\cdot\!K^{\top}$
|
||||
\textsf{dot}, the online-softmax primitives, and a $P\!\cdot\!V$
|
||||
\textsf{dot}, merging the running $(m,\ell,O)$ state by hand. The number
|
||||
of PE\_CPU commands this costs grows with the context. At a production
|
||||
context of $S_{kv}{=}1\,\text{M}$ tokens the Case-6 64-way split gives
|
||||
each PE $S_{\text{local}}{=}16384$ tokens, so its local attention is a
|
||||
$Q\!\cdot\!K^{\top}$ of $(8,128)\!\cdot\!(128,16384)$ and a
|
||||
$P\!\cdot\!V$ of $(8,16384)\!\cdot\!(16384,128)$---sixteen hand-issued
|
||||
tiles, each a fresh batch of CPU commands.
|
||||
|
||||
The composite command lets the kernel hand that tiling to PE\_SCHEDULER.
|
||||
We compare three command forms of the \emph{same} Case-6 kernel---identical
|
||||
placement and identical $(m,\ell,O)$ reduce, differing only in how the
|
||||
local attention is issued:
|
||||
\begin{itemize}
|
||||
\item \textbf{primitive}---the hand-tiled \textsf{dot}/softmax kernel
|
||||
above (the baseline of \S\ref{sec:gqa-long}).
|
||||
\item \textbf{composite}---each matrix product is one coarse
|
||||
\textsf{composite} GEMM over the \emph{whole} $S_{\text{local}}$, with
|
||||
$K$ and $V$ passed as HBM references so PE\_SCHEDULER streams and
|
||||
tiles them on the fixed $32\!\times\!64\!\times\!32$ MAC tile; the
|
||||
softmax stays primitive.
|
||||
\item \textbf{composite\,+\,softmax\_merge}---additionally folds the
|
||||
online-softmax merge and $P\!\cdot\!V$ into a single stateful
|
||||
\textsf{softmax\_merge} recipe composite.
|
||||
\end{itemize}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_composite.png}
|
||||
\caption{Three command forms of the Case-6 decode kernel, swept over
|
||||
context length ($S_{\text{local}}{=}S_{kv}/64$ per PE). \emph{Right:}
|
||||
PE\_CPU commands issued. The hand-tiled primitive kernel rises
|
||||
$O(n_{\text{tiles}})$---from 96 commands at one tile to 426 at the
|
||||
1\,M-token, sixteen-tile production point---while both composite forms
|
||||
issue a context-\emph{independent} $O(1)$ count (94 and 98) that
|
||||
\emph{saturates}: one coarse descriptor offloads the entire per-tile
|
||||
fan-out. \emph{Left:} the consequence for wall-clock latency is none---all
|
||||
three land on the same curve (\SI{30.6}{}, \SI{231}{},
|
||||
\SI{461}{\micro\second} at 8\,K\,/\,64\,K\,/\,128\,K), because decode is
|
||||
bound by streaming the KV cache, not by issue. Command-count is measured
|
||||
at emit time (exact, to 1\,M); latency on the data-mode engine over the
|
||||
tractable range.}
|
||||
\label{fig:gqa-composite}
|
||||
\end{figure}
|
||||
|
||||
Figure~\ref{fig:gqa-composite} reads off the two quantities that matter,
|
||||
and they point in opposite directions. The PE\_CPU command count (right)
|
||||
collapses from a context-growing $O(n_{\text{tiles}})$ to a flat $O(1)$:
|
||||
at 1\,M tokens the composite form issues \num{94} commands against the
|
||||
primitive kernel's \num{426}, a $4.5\times$ reduction
|
||||
($4.2\times$ in modeled dispatch cost), and---crucially---that number no
|
||||
longer grows with context. The wall-clock latency (left), by contrast,
|
||||
is unchanged across all three forms: decode is bound by streaming the KV
|
||||
cache out of HBM, so the command form does not move the critical path.
|
||||
|
||||
That juxtaposition is the point. The composite command is not a latency
|
||||
optimization for this memory-bound decode; it is a \emph{CPU-issue}
|
||||
optimization. Its value is removing the per-tile dispatch work that would
|
||||
otherwise grow without bound as context grows, freeing PE\_CPU to run
|
||||
ahead and keep the engines fed---which is exactly what lets the
|
||||
data-movement cost analyzed next show through as the true bottleneck
|
||||
rather than being masked by issue overhead. The \textsf{softmax\_merge}
|
||||
recipe folds the online merge into the same descriptor; on this
|
||||
memory-bound path its marginal cost over the plain GEMM composite is small
|
||||
(98 vs.\ 94 commands), and like the plain composite it keeps the issued
|
||||
count flat as context scales.
|
||||
|
||||
\paragraph{The compute-bound mirror: prefill.} Decode's verdict---command
|
||||
form is latency-neutral---is a property of its regime, not of the
|
||||
composite command. A decode step has $T_q{=}1$, so its score and context
|
||||
products are skinny ($M{=}G\,T_q{=}8$): the MAC array is barely fed and
|
||||
the kernel is bound by streaming the KV cache. Prefill is the opposite
|
||||
corner. It processes a block of query positions at once, so $M{=}G\,T_q$
|
||||
is large and tile-filling, the GEMMs carry real arithmetic intensity
|
||||
($\sim$$M$ flops/byte, well above the roofline ridge), and the kernel is
|
||||
\emph{compute-bound}. This is the regime the composite command was built
|
||||
for (\S\ref{sec:gemm}): it streams the per-HW-tile
|
||||
DMA$\rightleftarrows$compute pipeline so the MAC array stays fed, whereas
|
||||
the primitive kernel's blocking \textsf{tl.dot} serializes each tile's
|
||||
load and compute and starves the array between tiles. We run the same
|
||||
three command forms on a single-rank compute-bound prefill (FlashAttention
|
||||
$Q$-block $\times$ $S_{kv}$-tile, online softmax) and sweep the context
|
||||
length (Figure~\ref{fig:gqa-prefill-cb}).
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_prefill_compute_bound.png}
|
||||
\caption{Compute-bound prefill, three command forms, swept over context
|
||||
length ($M{=}8\,T_q$ tile-filling). \emph{Left:} end-to-end latency.
|
||||
\emph{Right:} MAC utilization (achieved $\div$ the
|
||||
\SI{8}{\tera\flop\per\second} per-PE peak). The hand-tiled primitive sits
|
||||
flat at $\sim$\SI{68}{\percent}---its serial load$\to$dot path leaves the
|
||||
MAC array idle between tiles regardless of context. The composite forms
|
||||
climb with context (\SI{67}{}$\to$\SI{80}{\percent} plain,
|
||||
\SI{67}{}$\to$\SI{83}{\percent} with the recipe) because a deeper $P\!\cdot
|
||||
\!V$ reduction gives more HW tiles to pipeline, and they convert that into
|
||||
wall-clock: at \num{1024} the recipe form is \SI{646.9}{} vs.\
|
||||
\SI{794.1}{\micro\second} (\SI{19}{\percent} faster). The margin
|
||||
\emph{grows} with context---the compute-bound mirror of the GEMM result of
|
||||
\S\ref{sec:gemm}.}
|
||||
\label{fig:gqa-prefill-cb}
|
||||
\end{figure}
|
||||
|
||||
The two studies together state the composite command's value precisely. It
|
||||
has two distinct benefits, and which one matters is set by the workload's
|
||||
roofline position. The first is \emph{host-issue offload}: one macro
|
||||
command in place of $O(n_{\text{tiles}})$ fine ones, which removes
|
||||
PE\_CPU dispatch work and is regime-independent (it shows in the decode
|
||||
command count). The second is \emph{MAC-array feeding}: the
|
||||
scheduler-internal per-tile DMA$\rightleftarrows$compute pipeline, which
|
||||
only converts to latency when the workload is compute-bound enough to have
|
||||
a MAC array worth keeping busy (it shows in the prefill utilization). A
|
||||
memory-bound decode exercises only the first; a compute-bound prefill
|
||||
exercises both. The composite command is the single mechanism that
|
||||
delivers each where it applies.
|
||||
|
||||
\subsection{Comprehensive Analysis}
|
||||
\label{sec:gqa-analysis}
|
||||
|
||||
|
||||
@@ -113,8 +113,8 @@ def main() -> None:
|
||||
)
|
||||
ax_cmd.set_ylabel("PE_CPU commands issued")
|
||||
ax_cmd.set_title(
|
||||
"PE_CPU command count — primitive rises O(n$_\\mathrm{tiles}$); "
|
||||
"composite saturates O(1)"
|
||||
"PE_CPU command count per command form\n"
|
||||
"(primitive O(n$_\\mathrm{tiles}$) rises; composite O(1) saturates)"
|
||||
)
|
||||
ax_cmd.axvline(1 << 20, color="#888", ls="--", lw=1, alpha=0.7)
|
||||
ax_cmd.annotate("1M production\ncontext", xy=(1 << 20, 0),
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Comparative figure for the compute-bound prefill composite study.
|
||||
|
||||
Reads sweep_prefill_compute_bound.json (emitted by milestone-1h-gqa,
|
||||
sweep ``prefill_cb``) and writes one two-panel PNG:
|
||||
|
||||
gqa_prefill_compute_bound.png
|
||||
Left — end-to-end prefill latency (µs) vs context length.
|
||||
Right — MAC utilization (achieved / 8 TFLOP·s⁻¹ per-PE peak) vs context.
|
||||
|
||||
Unlike memory-bound decode (where command form is latency-neutral), in
|
||||
compute-bound prefill the composite command keeps the MAC array fed by
|
||||
streaming DMA↔compute per HW tile, so it wins on both latency and
|
||||
utilization — and the margin grows with context (deeper P·V reduction =
|
||||
more tiles to pipeline).
|
||||
|
||||
Run (after the bench):
|
||||
GQA_1H_RUN=1 GQA_1H_SWEEPS=prefill_cb python -m kernbench.cli.main run \\
|
||||
--bench milestone-1h-gqa --topology topology.yaml
|
||||
python scripts/paper/paper_plot_gqa_prefill_compute_bound.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt # noqa: E402
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
_FIG_DIR = (
|
||||
_REPO_ROOT / "src" / "kernbench" / "benches"
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
)
|
||||
_SWEEP_JSON = _FIG_DIR / "sweep_prefill_compute_bound.json"
|
||||
_PAPER_FIG_DIR = (
|
||||
_REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
|
||||
)
|
||||
|
||||
_VARIANT_STYLE = {
|
||||
"primitive": ("primitive (tl.dot, hand-tiled)", "#c0504d", "o"),
|
||||
"composite": ("composite GEMM", "#3b6ea5", "s"),
|
||||
"composite_extended": ("composite + softmax_merge", "#4f8a4f", "^"),
|
||||
}
|
||||
_ORDER = ("primitive", "composite", "composite_extended")
|
||||
|
||||
|
||||
def _ctx_label(c: int) -> str:
|
||||
return f"{c // 1024}K" if c >= 1024 else str(c)
|
||||
|
||||
|
||||
def _series(rows, variant, key):
|
||||
pts = sorted(((r["ctx_len"], r[key]) for r in rows
|
||||
if r["variant"] == variant), key=lambda t: t[0])
|
||||
return [p[0] for p in pts], [p[1] for p in pts]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
sweep = json.loads(_SWEEP_JSON.read_text())
|
||||
rows = sweep["rows"]
|
||||
ctxs = sweep["ctx_points"]
|
||||
|
||||
fig, (ax_lat, ax_util) = plt.subplots(1, 2, figsize=(13.0, 4.8))
|
||||
|
||||
for v in _ORDER:
|
||||
label, color, marker = _VARIANT_STYLE[v]
|
||||
xs, lat = _series(rows, v, "latency_ns")
|
||||
ax_lat.plot(xs, [y / 1e3 for y in lat], marker=marker,
|
||||
color=color, label=label, lw=2)
|
||||
xs, util = _series(rows, v, "mac_util")
|
||||
ax_util.plot(xs, [u * 100 for u in util], marker=marker,
|
||||
color=color, label=label, lw=2)
|
||||
|
||||
for ax in (ax_lat, ax_util):
|
||||
ax.set_xscale("log", base=2)
|
||||
ax.set_xticks(ctxs)
|
||||
ax.set_xticklabels([_ctx_label(c) for c in ctxs])
|
||||
ax.set_xlabel(r"context length (= $T_q$ = $S_{kv}$)")
|
||||
ax.grid(True, ls=":", alpha=0.5)
|
||||
ax.legend(fontsize=9)
|
||||
|
||||
ax_lat.set_ylabel("end-to-end prefill latency (µs)")
|
||||
ax_lat.set_title("Compute-bound prefill latency per command form")
|
||||
ax_util.set_ylabel("MAC utilization (% of 8 TFLOP·s⁻¹ peak)")
|
||||
ax_util.set_title(
|
||||
"MAC utilization — composite keeps the array fed; primitive starves"
|
||||
)
|
||||
ax_util.axhline(100, color="#888", ls="--", lw=1, alpha=0.6)
|
||||
|
||||
fig.suptitle(
|
||||
"Compute-bound prefill attention — use of composite commands\n"
|
||||
"single-rank, GQA single-KV-head group ($h_q{=}8$, $d_{\\text{head}}"
|
||||
"{=}128$); $M{=}8T_q$ tile-filling",
|
||||
fontsize=11,
|
||||
)
|
||||
fig.tight_layout(rect=(0, 0, 1, 0.92))
|
||||
|
||||
out = _FIG_DIR / "gqa_prefill_compute_bound.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
print(f"wrote {out}")
|
||||
|
||||
if _PAPER_FIG_DIR.is_dir():
|
||||
dst = _PAPER_FIG_DIR / out.name
|
||||
dst.write_bytes(out.read_bytes())
|
||||
print(f"copied {dst}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
BIN
Binary file not shown.
|
After Width: | Height: | Size: 136 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 145 KiB |
@@ -0,0 +1,258 @@
|
||||
{
|
||||
"version": 2,
|
||||
"variants": [
|
||||
"primitive",
|
||||
"composite",
|
||||
"composite_extended"
|
||||
],
|
||||
"s_kv_opcount": [
|
||||
8192,
|
||||
65536,
|
||||
131072,
|
||||
262144,
|
||||
524288,
|
||||
1048576
|
||||
],
|
||||
"s_kv_latency": [
|
||||
8192,
|
||||
32768,
|
||||
65536,
|
||||
131072
|
||||
],
|
||||
"rows": [
|
||||
{
|
||||
"variant": "primitive",
|
||||
"S_kv": 8192,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 96,
|
||||
"pe_cpu_dispatch_cycles": 930,
|
||||
"latency_ns": 30646.00300000079
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"S_kv": 8192,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 94,
|
||||
"pe_cpu_dispatch_cycles": 978,
|
||||
"latency_ns": 30566.1840000007
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"S_kv": 8192,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 98,
|
||||
"pe_cpu_dispatch_cycles": 1032,
|
||||
"latency_ns": 30483.359500000362
|
||||
},
|
||||
{
|
||||
"variant": "primitive",
|
||||
"S_kv": 65536,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 96,
|
||||
"pe_cpu_dispatch_cycles": 930,
|
||||
"latency_ns": 231579.37900000165
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"S_kv": 65536,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 94,
|
||||
"pe_cpu_dispatch_cycles": 978,
|
||||
"latency_ns": 231270.18400000152
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"S_kv": 65536,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 98,
|
||||
"pe_cpu_dispatch_cycles": 1032,
|
||||
"latency_ns": 231211.1740000015
|
||||
},
|
||||
{
|
||||
"variant": "primitive",
|
||||
"S_kv": 131072,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 118,
|
||||
"pe_cpu_dispatch_cycles": 1145,
|
||||
"latency_ns": 461119.51900000183
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"S_kv": 131072,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 94,
|
||||
"pe_cpu_dispatch_cycles": 978,
|
||||
"latency_ns": 460651.3170000027
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"S_kv": 131072,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 98,
|
||||
"pe_cpu_dispatch_cycles": 1032,
|
||||
"latency_ns": 460552.9805000025
|
||||
},
|
||||
{
|
||||
"variant": "primitive",
|
||||
"S_kv": 262144,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 162,
|
||||
"pe_cpu_dispatch_cycles": 1575,
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"S_kv": 262144,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 94,
|
||||
"pe_cpu_dispatch_cycles": 978,
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"S_kv": 262144,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 98,
|
||||
"pe_cpu_dispatch_cycles": 1032,
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
"variant": "primitive",
|
||||
"S_kv": 524288,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 250,
|
||||
"pe_cpu_dispatch_cycles": 2435,
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"S_kv": 524288,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 94,
|
||||
"pe_cpu_dispatch_cycles": 978,
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"S_kv": 524288,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 98,
|
||||
"pe_cpu_dispatch_cycles": 1032,
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
"variant": "primitive",
|
||||
"S_kv": 1048576,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 426,
|
||||
"pe_cpu_dispatch_cycles": 4155,
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"S_kv": 1048576,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 94,
|
||||
"pe_cpu_dispatch_cycles": 978,
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"S_kv": 1048576,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1,
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 98,
|
||||
"pe_cpu_dispatch_cycles": 1032,
|
||||
"latency_ns": null
|
||||
}
|
||||
]
|
||||
}
|
||||
+105
@@ -0,0 +1,105 @@
|
||||
{
|
||||
"version": 1,
|
||||
"variants": [
|
||||
"primitive",
|
||||
"composite",
|
||||
"composite_extended"
|
||||
],
|
||||
"ctx_points": [
|
||||
256,
|
||||
512,
|
||||
1024
|
||||
],
|
||||
"rows": [
|
||||
{
|
||||
"variant": "primitive",
|
||||
"ctx_len": 256,
|
||||
"M": 2048,
|
||||
"latency_ns": 48857.361999999725,
|
||||
"gemm_busy_ns": 33554.43200000003,
|
||||
"dma_busy_ns": 20880.000000000004,
|
||||
"achieved_tflops": 5.4942683151825005,
|
||||
"mac_util": 0.6867835393978126
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"ctx_len": 256,
|
||||
"M": 2048,
|
||||
"latency_ns": 49887.505999999594,
|
||||
"gemm_busy_ns": 34531.32799999314,
|
||||
"dma_busy_ns": 1095773.440000072,
|
||||
"achieved_tflops": 5.380815308746887,
|
||||
"mac_util": 0.6726019135933609
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"ctx_len": 256,
|
||||
"M": 2048,
|
||||
"latency_ns": 50280.91399999841,
|
||||
"gemm_busy_ns": 129096.19199996963,
|
||||
"dma_busy_ns": 633563.5200000411,
|
||||
"achieved_tflops": 5.338714725830331,
|
||||
"mac_util": 0.6673393407287914
|
||||
},
|
||||
{
|
||||
"variant": "primitive",
|
||||
"ctx_len": 512,
|
||||
"M": 4096,
|
||||
"latency_ns": 197496.81800000605,
|
||||
"gemm_busy_ns": 134217.72800000012,
|
||||
"dma_busy_ns": 66880.0,
|
||||
"achieved_tflops": 5.436755056985105,
|
||||
"mac_util": 0.6795943821231382
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"ctx_len": 512,
|
||||
"M": 4096,
|
||||
"latency_ns": 177341.23400000148,
|
||||
"gemm_busy_ns": 141168.63999995415,
|
||||
"dma_busy_ns": 8567063.039998509,
|
||||
"achieved_tflops": 6.054665346469796,
|
||||
"mac_util": 0.7568331683087245
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"ctx_len": 512,
|
||||
"M": 4096,
|
||||
"latency_ns": 174702.7699999963,
|
||||
"gemm_busy_ns": 1076198.3999996716,
|
||||
"dma_busy_ns": 6594394.87999886,
|
||||
"achieved_tflops": 6.146106464139193,
|
||||
"mac_util": 0.7682633080173992
|
||||
},
|
||||
{
|
||||
"variant": "primitive",
|
||||
"ctx_len": 1024,
|
||||
"M": 8192,
|
||||
"latency_ns": 794090.4820000405,
|
||||
"gemm_busy_ns": 536870.9120000004,
|
||||
"dma_busy_ns": 234240.0,
|
||||
"achieved_tflops": 5.4086623544239645,
|
||||
"mac_util": 0.6760827943029956
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"ctx_len": 1024,
|
||||
"M": 8192,
|
||||
"latency_ns": 668110.7059999453,
|
||||
"gemm_busy_ns": 873011.1999923651,
|
||||
"dma_busy_ns": 67794150.3999821,
|
||||
"achieved_tflops": 6.428526376585786,
|
||||
"mac_util": 0.8035657970732233
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"ctx_len": 1024,
|
||||
"M": 8192,
|
||||
"latency_ns": 646937.2019999702,
|
||||
"gemm_busy_ns": 8870907.903995967,
|
||||
"dma_busy_ns": 59655820.79998447,
|
||||
"achieved_tflops": 6.638924586068552,
|
||||
"mac_util": 0.829865573258569
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
"""Compute-bound prefill attention — 3 command-form variants (single-rank).
|
||||
|
||||
Companion to the memory-bound decode study
|
||||
(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp*``). Prefill processes a
|
||||
block of T_q query positions at once, so the score / context GEMMs have a
|
||||
large M = G·T_q and high arithmetic intensity (~M flops/byte) — the
|
||||
workload is **compute-bound** (above the roofline ridge), unlike T_q=1
|
||||
decode (M=8, memory-bound). This is the regime where the composite
|
||||
command's value shows: it streams DMA↔compute per HW tile to keep the MAC
|
||||
array fed, while the primitive kernel serializes load→dot and starves it.
|
||||
|
||||
Single-rank (C=P=1): no cross-device reduce — the focus is the per-PE
|
||||
GEMM-issue mechanism. FlashAttention 2-D tiling (Q-block × S_kv-tile,
|
||||
online softmax) bounds the TCM scratch.
|
||||
|
||||
Three forms, differing only in how each Q-block's local attention is
|
||||
issued (placement/softmax identical):
|
||||
primitive tl.dot per S_kv-tile + primitive online merge.
|
||||
composite one coarse composite GEMM per Q-block over the whole
|
||||
S_kv (K, V as HBM refs → PE_SCHEDULER tiles/streams).
|
||||
composite_extended Q·Kᵀ composite + softmax_merge recipe composite.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import _merge_running
|
||||
|
||||
M_BLOCK = 128 # query rows per FlashAttention Q-block (tile-filling).
|
||||
TILE_S_KV = 256 # primitive per-tile S_kv width.
|
||||
_SEED_S_KV = 64 # composite_extended recipe seed slice.
|
||||
|
||||
|
||||
def _qblock_bounds(T_q: int, h_q: int, h_kv: int):
|
||||
G = h_q // h_kv
|
||||
M_total = G * T_q
|
||||
n_qblocks = (M_total + M_BLOCK - 1) // M_BLOCK
|
||||
return G, M_total, n_qblocks
|
||||
|
||||
|
||||
# ── primitive: hand-tiled monolithic dots + online merge ─────────────
|
||||
|
||||
|
||||
def gqa_prefill_primitive_kernel(
|
||||
q_ptr, k_ptr, v_ptr, o_ptr,
|
||||
T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
|
||||
) -> None:
|
||||
G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
|
||||
ROW = d_head * 2 # f16
|
||||
n_tiles = (S_kv + TILE_S_KV - 1) // TILE_S_KV
|
||||
|
||||
for qb in range(n_qblocks):
|
||||
# Per-Q-block scratch_scope: this block's running (m,ℓ,O) plus its
|
||||
# tile-0 transients are freed before the next block, so scratch
|
||||
# stays O(one block) rather than accumulating across all blocks.
|
||||
with tl.scratch_scope():
|
||||
m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
|
||||
q_off = qb * M_BLOCK * ROW
|
||||
Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
|
||||
|
||||
tile_s0 = min(TILE_S_KV, S_kv)
|
||||
K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
|
||||
V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
|
||||
scores = tl.dot(Q, K_T)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
exp_s = tl.exp(scores - m_local)
|
||||
l_local = tl.sum(exp_s, axis=-1)
|
||||
O_local = tl.dot(exp_s, V)
|
||||
|
||||
for ti in range(1, n_tiles):
|
||||
start = ti * TILE_S_KV
|
||||
tile_s = min(TILE_S_KV, S_kv - start)
|
||||
with tl.scratch_scope():
|
||||
K_T_t = tl.load(k_ptr + start * ROW,
|
||||
shape=(d_head, tile_s), dtype="f16")
|
||||
V_t = tl.load(v_ptr + start * ROW,
|
||||
shape=(tile_s, d_head), dtype="f16")
|
||||
scores_t = tl.dot(Q, K_T_t)
|
||||
m_tile = tl.max(scores_t, axis=-1)
|
||||
exp_t = tl.exp(scores_t - m_tile)
|
||||
l_tile = tl.sum(exp_t, axis=-1)
|
||||
O_tile = tl.dot(exp_t, V_t)
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr + q_off, O_final)
|
||||
|
||||
|
||||
# ── composite: one coarse composite GEMM per Q-block ─────────────────
|
||||
|
||||
|
||||
def gqa_prefill_composite_kernel(
|
||||
q_ptr, k_ptr, v_ptr, o_ptr,
|
||||
T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
|
||||
) -> None:
|
||||
G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
|
||||
ROW = d_head * 2
|
||||
|
||||
for qb in range(n_qblocks):
|
||||
with tl.scratch_scope():
|
||||
m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
|
||||
q_off = qb * M_BLOCK * ROW
|
||||
Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
|
||||
K_T = tl.ref(k_ptr, shape=(d_head, S_kv), dtype="f16")
|
||||
V = tl.ref(v_ptr, shape=(S_kv, d_head), dtype="f16")
|
||||
|
||||
scores = tl.composite(op="gemm", a=Q, b=K_T) # Q·Kᵀ, one command
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
exp_s = tl.exp(scores - m_local)
|
||||
l_local = tl.sum(exp_s, axis=-1)
|
||||
O_local = tl.zeros((m_blk, d_head), dtype="f16")
|
||||
tl.composite(op="gemm", a=exp_s, b=V, out=O_local) # P·V, one command
|
||||
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr + q_off, O_final)
|
||||
|
||||
|
||||
# ── composite_extended: Q·Kᵀ composite + softmax_merge recipe ────────
|
||||
|
||||
|
||||
def gqa_prefill_composite_ext_kernel(
|
||||
q_ptr, k_ptr, v_ptr, o_ptr,
|
||||
T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
|
||||
) -> None:
|
||||
G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
|
||||
ROW = d_head * 2
|
||||
|
||||
for qb in range(n_qblocks):
|
||||
with tl.scratch_scope():
|
||||
m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
|
||||
q_off = qb * M_BLOCK * ROW
|
||||
Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
|
||||
|
||||
seed = min(_SEED_S_KV, S_kv)
|
||||
K_T0 = tl.load(k_ptr, shape=(d_head, seed), dtype="f16")
|
||||
V0 = tl.load(v_ptr, shape=(seed, d_head), dtype="f16")
|
||||
scores0 = tl.dot(Q, K_T0)
|
||||
m_local = tl.max(scores0, axis=-1)
|
||||
exp0 = tl.exp(scores0 - m_local)
|
||||
l_local = tl.sum(exp0, axis=-1)
|
||||
O_local = tl.dot(exp0, V0)
|
||||
|
||||
rest = S_kv - seed
|
||||
if rest > 0:
|
||||
K_T1 = tl.ref(k_ptr + seed * ROW,
|
||||
shape=(d_head, rest), dtype="f16")
|
||||
V1 = tl.ref(v_ptr + seed * ROW,
|
||||
shape=(rest, d_head), dtype="f16")
|
||||
scores1 = tl.composite(op="gemm", a=Q, b=K_T1)
|
||||
tl.composite(
|
||||
prologue=[{"op": "softmax_merge", "s": scores1,
|
||||
"m": m_local, "l": l_local, "O": O_local}],
|
||||
op="gemm", b=V1, out=O_local,
|
||||
epilogue=[{"op": "add", "other": O_local}],
|
||||
)
|
||||
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr + q_off, O_final)
|
||||
@@ -0,0 +1,142 @@
|
||||
"""milestone-1h-gqa: compute-bound prefill composite-command study.
|
||||
|
||||
Three command-form variants of a single-rank compute-bound prefill
|
||||
attention kernel (``_gqa_prefill_compute_bound``), swept over context
|
||||
length S_kv = T_q. Unlike the memory-bound decode study, prefill has a
|
||||
large M = G·T_q, so the score / context GEMMs are compute-bound — the
|
||||
regime where the composite command keeps the MAC array fed (DMA↔compute
|
||||
pipelining) and the hand-tiled primitive starves it on the serial
|
||||
load→dot path.
|
||||
|
||||
Records per (variant, context) the end-to-end latency, the GEMM-engine
|
||||
busy time, and the MAC occupancy (gemm_busy / e2e) so the comparative
|
||||
plot can show composite winning on both latency and utilization.
|
||||
|
||||
Runs in data mode (engine latency). Gated via the umbrella
|
||||
``GQA_1H_SWEEPS=prefill_cb``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
|
||||
gqa_prefill_composite_ext_kernel,
|
||||
gqa_prefill_composite_kernel,
|
||||
gqa_prefill_primitive_kernel,
|
||||
)
|
||||
from kernbench.policy.placement.dp import DPPolicy
|
||||
|
||||
_OUTPUT_DIR = (
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
)
|
||||
_SWEEP_JSON = _OUTPUT_DIR / "sweep_prefill_compute_bound.json"
|
||||
|
||||
_VARIANT_KERNELS = {
|
||||
"primitive": gqa_prefill_primitive_kernel,
|
||||
"composite": gqa_prefill_composite_kernel,
|
||||
"composite_extended": gqa_prefill_composite_ext_kernel,
|
||||
}
|
||||
_VARIANTS = ("primitive", "composite", "composite_extended")
|
||||
|
||||
# Context length S_kv = T_q (prefill processes T_q tokens against S_kv=T_q
|
||||
# keys). M = G·T_q = 8·T_q is tile-filling/compute-bound at every point.
|
||||
_CTX_POINTS = (256, 512, 1024)
|
||||
|
||||
_H_Q, _H_KV, _D_HEAD = 8, 1, 128
|
||||
_PEAK_TFLOPS = 8.0 # per-PE f16 GEMM peak (topology.yaml pe_gemm.peak_tflops_f16)
|
||||
|
||||
|
||||
def _run_panel_fn(variant: str, ctx_len: int):
|
||||
kernel = _VARIANT_KERNELS[variant]
|
||||
panel = f"prefill_cb_{variant}_c{ctx_len}"
|
||||
|
||||
def _bench_fn(ctx):
|
||||
dp = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=1, num_pes=1)
|
||||
q = ctx.zeros((ctx_len, _H_Q * _D_HEAD),
|
||||
dtype="f16", dp=dp, name=f"{panel}_q")
|
||||
k = ctx.zeros((ctx_len, _H_KV * _D_HEAD),
|
||||
dtype="f16", dp=dp, name=f"{panel}_k")
|
||||
v = ctx.zeros((ctx_len, _H_KV * _D_HEAD),
|
||||
dtype="f16", dp=dp, name=f"{panel}_v")
|
||||
o = ctx.empty((ctx_len, _H_Q * _D_HEAD),
|
||||
dtype="f16", dp=dp, name=f"{panel}_o")
|
||||
ctx.launch(panel, kernel, q, k, v, o,
|
||||
ctx_len, ctx_len, _H_Q, _H_KV, _D_HEAD, 1, 1,
|
||||
_auto_dim_remap=False)
|
||||
|
||||
return _bench_fn
|
||||
|
||||
|
||||
def _end_to_end_ns(op_log) -> float:
|
||||
if not op_log:
|
||||
return 0.0
|
||||
return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
|
||||
|
||||
|
||||
def _engine_busy_ns(op_log, suffix: str) -> float:
|
||||
return sum(r.t_end - r.t_start
|
||||
for r in op_log if r.component_id.endswith("." + suffix))
|
||||
|
||||
|
||||
def _run_panel(variant: str, ctx_len: int, topology: str) -> dict:
|
||||
from kernbench.runtime_api.bench_runner import run_bench
|
||||
from kernbench.runtime_api.types import resolve_device
|
||||
from kernbench.sim_engine.engine import GraphEngine
|
||||
from kernbench.topology.builder import resolve_topology
|
||||
|
||||
topo = resolve_topology(topology)
|
||||
result = run_bench(
|
||||
topology=topo, bench_fn=_run_panel_fn(variant, ctx_len),
|
||||
device=resolve_device(None),
|
||||
engine_factory=lambda t, d: GraphEngine(
|
||||
getattr(t, "topology_obj", t), enable_data=True,
|
||||
),
|
||||
)
|
||||
if not result.completion.ok:
|
||||
raise RuntimeError(
|
||||
f"gqa-prefill-cb {variant}@{ctx_len} failed: {result.completion}"
|
||||
)
|
||||
op_log = result.engine.op_log
|
||||
e2e = _end_to_end_ns(op_log)
|
||||
gemm = _engine_busy_ns(op_log, "pe_gemm")
|
||||
dma = _engine_busy_ns(op_log, "pe_dma")
|
||||
# Useful attention flops (Q·Kᵀ + P·V), single rank.
|
||||
G = _H_Q // _H_KV
|
||||
M = G * ctx_len
|
||||
useful_flops = 4.0 * M * _D_HEAD * ctx_len
|
||||
# MAC utilization = achieved / peak. ``achieved_tflops`` uses wall-clock
|
||||
# (useful_flops / e2e) so it is bounded by peak even when the composite
|
||||
# path overlaps many tile GEMMs (gemm_busy is a sum over overlapping ops
|
||||
# and is kept only for diagnostics).
|
||||
return {
|
||||
"variant": variant,
|
||||
"ctx_len": ctx_len,
|
||||
"M": M,
|
||||
"latency_ns": e2e,
|
||||
"gemm_busy_ns": gemm,
|
||||
"dma_busy_ns": dma,
|
||||
"achieved_tflops": (useful_flops / e2e / 1e3) if e2e > 0 else 0.0,
|
||||
"mac_util": (useful_flops / e2e / 1e3 / _PEAK_TFLOPS) if e2e > 0 else 0.0,
|
||||
}
|
||||
|
||||
|
||||
def run_sweep(topology: str = "topology.yaml") -> int:
|
||||
"""Drive all (variant, context) prefill panels; write sweep.json."""
|
||||
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
rows = [
|
||||
_run_panel(variant, ctx_len, topology)
|
||||
for ctx_len in _CTX_POINTS
|
||||
for variant in _VARIANTS
|
||||
]
|
||||
sweep = {
|
||||
"version": 1,
|
||||
"variants": list(_VARIANTS),
|
||||
"ctx_points": list(_CTX_POINTS),
|
||||
"rows": rows,
|
||||
}
|
||||
_SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
|
||||
print(f" gqa-prefill-cb: {len(rows)} rows -> {_SWEEP_JSON}")
|
||||
return len(rows)
|
||||
@@ -28,6 +28,9 @@ from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import (
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (
|
||||
run_sweep as _run_composite_sweep,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
|
||||
run_sweep as _run_prefill_cb_sweep,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases import (
|
||||
run_sweep as _run_prefill_sweep,
|
||||
)
|
||||
@@ -62,6 +65,7 @@ def run(torch) -> None:
|
||||
"prefill": _run_prefill_sweep,
|
||||
"decode": _run_decode_sweep,
|
||||
"composite": _run_composite_sweep,
|
||||
"prefill_cb": _run_prefill_cb_sweep,
|
||||
}
|
||||
unknown = [s for s in sweeps if s not in runners]
|
||||
if unknown:
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Tests for the compute-bound prefill composite-command study.
|
||||
|
||||
Unlike memory-bound decode (where command form is latency-neutral), in
|
||||
compute-bound prefill the composite command keeps the MAC array fed by
|
||||
pipelining DMA↔compute per HW tile, so it wins on wall-clock latency. The
|
||||
rigorous claim here is therefore the *opposite* of the decode study:
|
||||
``composite_latency < primitive_latency`` at a compute-bound context.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.common.pe_commands import CompositeCmd
|
||||
from kernbench.triton_emu.tl_context import TLContext, run_kernel
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
|
||||
gqa_prefill_composite_ext_kernel,
|
||||
gqa_prefill_composite_kernel,
|
||||
gqa_prefill_primitive_kernel,
|
||||
)
|
||||
|
||||
_TOPOLOGY = Path(__file__).resolve().parents[2] / "topology.yaml"
|
||||
_H_Q, _H_KV, _D_HEAD = 8, 1, 128
|
||||
|
||||
|
||||
def _n_composites(kernel, T_q: int, S_kv: int) -> int:
|
||||
tl = TLContext(pe_id=0, num_programs=1, scratch_base=1 << 61,
|
||||
scratch_size=1 << 20)
|
||||
run_kernel(kernel, tl, 0x1000, 0x2000, 0x3000, 0x4000,
|
||||
T_q, S_kv, _H_Q, _H_KV, _D_HEAD, 1, 1)
|
||||
return sum(1 for c in tl.commands if isinstance(c, CompositeCmd))
|
||||
|
||||
|
||||
def test_command_form_composite_counts():
|
||||
"""One Q-block (T_q=16 → M=128): the primitive issues no composites;
|
||||
each composite form issues exactly two (Q·Kᵀ and P·V / recipe)."""
|
||||
assert _n_composites(gqa_prefill_primitive_kernel, 16, 256) == 0
|
||||
assert _n_composites(gqa_prefill_composite_kernel, 16, 256) == 2
|
||||
assert _n_composites(gqa_prefill_composite_ext_kernel, 16, 256) == 2
|
||||
|
||||
|
||||
def _latency_ns(variant: str, ctx_len: int) -> float:
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
|
||||
_run_panel,
|
||||
)
|
||||
return _run_panel(variant, ctx_len, str(_TOPOLOGY))["latency_ns"]
|
||||
|
||||
|
||||
def test_three_variants_complete_in_data_mode():
|
||||
"""All three prefill forms run end-to-end through the engine."""
|
||||
for variant in ("primitive", "composite", "composite_extended"):
|
||||
assert _latency_ns(variant, 256) > 0
|
||||
|
||||
|
||||
def test_composite_faster_in_compute_bound_prefill():
|
||||
"""At a compute-bound context (512), the composite command's DMA↔compute
|
||||
pipelining beats the primitive's serial load→dot — the opposite of the
|
||||
memory-bound decode study, where the two tie."""
|
||||
prim = _latency_ns("primitive", 512)
|
||||
comp = _latency_ns("composite", 512)
|
||||
assert comp < prim, f"composite {comp} not < primitive {prim}"
|
||||
Reference in New Issue
Block a user