diff --git a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite.png b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite.png new file mode 100644 index 0000000..7078001 Binary files /dev/null and b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite.png differ diff --git a/docs/report/1H-codesign-paper/figures/gqa_prefill_compute_bound.png b/docs/report/1H-codesign-paper/figures/gqa_prefill_compute_bound.png new file mode 100644 index 0000000..83aa87a Binary files /dev/null and b/docs/report/1H-codesign-paper/figures/gqa_prefill_compute_bound.png differ diff --git a/docs/report/1H-codesign-paper/sections/05-gqa.tex b/docs/report/1H-codesign-paper/sections/05-gqa.tex index cb0edb4..db6731f 100644 --- a/docs/report/1H-codesign-paper/sections/05-gqa.tex +++ b/docs/report/1H-codesign-paper/sections/05-gqa.tex @@ -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} diff --git a/scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py b/scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py index 9c92733..cdf0684 100644 --- a/scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py +++ b/scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py @@ -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), diff --git a/scripts/paper/paper_plot_gqa_prefill_compute_bound.py b/scripts/paper/paper_plot_gqa_prefill_compute_bound.py new file mode 100644 index 0000000..f4023d7 --- /dev/null +++ b/scripts/paper/paper_plot_gqa_prefill_compute_bound.py @@ -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() diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite.png new file mode 100644 index 0000000..7078001 Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_compute_bound.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_compute_bound.png new file mode 100644 index 0000000..83aa87a Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_compute_bound.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json new file mode 100644 index 0000000..f4060d6 --- /dev/null +++ b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json @@ -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 + } + ] +} \ No newline at end of file diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_prefill_compute_bound.json b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_prefill_compute_bound.json new file mode 100644 index 0000000..90905aa --- /dev/null +++ b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_prefill_compute_bound.json @@ -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 + } + ] +} \ No newline at end of file diff --git a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_prefill_compute_bound.py b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_prefill_compute_bound.py new file mode 100644 index 0000000..4fb1248 --- /dev/null +++ b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_prefill_compute_bound.py @@ -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) diff --git a/src/kernbench/benches/gqa_helpers/long_ctx/gqa_prefill_compute_bound.py b/src/kernbench/benches/gqa_helpers/long_ctx/gqa_prefill_compute_bound.py new file mode 100644 index 0000000..1f816a1 --- /dev/null +++ b/src/kernbench/benches/gqa_helpers/long_ctx/gqa_prefill_compute_bound.py @@ -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) diff --git a/src/kernbench/benches/milestone_1h_gqa.py b/src/kernbench/benches/milestone_1h_gqa.py index 8a7cb87..a43d97c 100644 --- a/src/kernbench/benches/milestone_1h_gqa.py +++ b/src/kernbench/benches/milestone_1h_gqa.py @@ -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: diff --git a/tests/attention/test_gqa_prefill_compute_bound.py b/tests/attention/test_gqa_prefill_compute_bound.py new file mode 100644 index 0000000..366ba12 --- /dev/null +++ b/tests/attention/test_gqa_prefill_compute_bound.py @@ -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}"