paper(gqa): single-rank decode-streaming study — composite's masked bandwidth win

Adds the memory-bound mirror of the compute-bound prefill figure. New single-rank decode sweep (T_q=1, M=8) reuses the prefill_compute_bound kernels, sweeping per-rank S_kv with NO cross-CUBE reduce — isolating the local attention that the 64-way Case-6 decode masks under its reduce tail.

Finding: even memory-bound decode benefits from the composite command. Its scheduler-streamed concurrent per-tile DMAs reach ~233 GB/s (91% of the 256 GB/s per-rank roofline) while the primitive's blocking tl.dot serializes one tile DMA and plateaus at ~166 GB/s — a 25-28% latency win that widens with context. This refines the paper's 'decode is latency-neutral' claim: neutral at 64-way production scale (reduce-dominated), but the composite extracts bandwidth at the local level. The bandwidth roofline here mirrors the MAC roofline in prefill.

Adds: gqa_decode_streaming.py sweep (wired into the milestone umbrella as GQA_1H_SWEEPS=decode_streaming), paper_plot_gqa_decode_streaming.py, the figure, and a new fig:gqa-decode-stream + paragraph in 05-gqa.tex; rebuilt main.pdf. Prefill (Experiment B) re-verified byte-identical (794.1/668.1/646.9us).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-23 19:39:45 -07:00
parent faf011dae5
commit a4a2683aad
8 changed files with 514 additions and 3 deletions
Binary file not shown.
Binary file not shown.

After

Width:  |  Height:  |  Size: 136 KiB

@@ -284,7 +284,9 @@ 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 for this memory-bound decode \emph{at the 64-way production scale} (a
single-rank caveat follows, Figure~\ref{fig:gqa-decode-stream}); 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
@@ -295,8 +297,49 @@ 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
\paragraph{Isolating the rank: the masked streaming win.} That neutrality
is a property of the \emph{full 64-way} critical path, not of the local
attention: at production scale the inter-CUBE $(m,\ell,O)$ reduce tail and
shared-HBM contention set the wall clock, so a faster local attention does
not surface. Stripping those away---a single rank, no cross-CUBE reduce,
swept over the per-rank context $S_{kv}$ (so $S_{kv}{=}16$\,K here is the
per-PE load of a 1M-token, 64-way-sharded decode)---exposes the local
attention directly (Figure~\ref{fig:gqa-decode-stream}), and the composite
\emph{does} win, by \SI{25}{}--\SI{28}{\percent}. The reason is the
memory-bound mirror of prefill: its scheduler-streamed concurrent per-tile
DMAs keep the HBM pipeline full and reach \SI{233}{\giga\byte\per\second}
---\SI{91}{\percent} of the per-rank \SI{256}{\giga\byte\per\second}
roofline---whereas the primitive kernel's blocking \textsf{tl.dot}
serializes one tile DMA at a time and plateaus at
\SI{166}{\giga\byte\per\second}. So even for memory-bound decode the
composite is not \emph{only} a CPU-issue optimization---it also extracts
bandwidth---but that latency benefit materializes only when the local
attention is on the critical path, which at 64-way production scale it is
not.
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_decode_streaming.png}
\caption{Single-rank memory-bound decode ($T_q{=}1$, $M{=}8$), three
command forms, swept over per-rank context. \emph{Left:} end-to-end
latency---the composite forms run \SI{25}{}--\SI{28}{\percent} below the
primitive, a gap that widens with context. \emph{Right:} achieved HBM
bandwidth against the per-rank \SI{256}{\giga\byte\per\second} roofline.
The primitive's blocking load$\rightarrow$dot serializes the KV stream and
plateaus at \SI{166}{\giga\byte\per\second}; the composite forms pipeline
concurrent per-tile DMAs through the scheduler and reach
\SI{233}{\giga\byte\per\second}. This is the memory-bound mirror of the
prefill result (Figure~\ref{fig:gqa-prefill-cb}): there the composite
approaches the MAC roofline, here the bandwidth roofline. Capped at
$16$\,K---the plain composite materializes the full $(M,S_{kv})$ scores in
TCM, so beyond that only the \textsf{softmax\_merge} recipe, which tiles
the softmax, stays within scratch.}
\label{fig:gqa-decode-stream}
\end{figure}
\paragraph{The compute-bound mirror: prefill.} Decode's \emph{production}
verdict---command form is latency-neutral at 64-way scale---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
@@ -0,0 +1,124 @@
"""Comparative figure for the memory-bound decode-streaming composite study.
Reads sweep_decode_streaming.json (emitted by milestone-1h-gqa, sweep
``decode_streaming``) and writes one two-panel PNG:
gqa_decode_streaming.png
Left — end-to-end single-rank decode latency (µs) vs per-rank context.
Right — achieved HBM bandwidth (GB/s) vs context, against the
256 GB/s per-rank roofline.
The memory-bound mirror of the compute-bound prefill figure. With T_q=1
the GEMMs are skinny (M=8) and the kernel is bound by streaming the KV
cache. Isolating a single rank (no inter-CUBE reduce) reveals what the
64-way Case-6 decode masks: the composite command still wins, not by
feeding the MAC array but by keeping the DMA pipeline full — its
scheduler-streamed concurrent tile DMAs extract ~230 GB/s (near the
256 GB/s roofline) while the primitive kernel's blocking tl.dot serializes
one tile DMA at a time and plateaus at ~166 GB/s. That bandwidth gap is a
~25-28 % latency win that grows nowhere near prefill's compute-bound
margin but is decidedly not zero.
Run (after the bench):
GQA_1H_RUN=1 GQA_1H_SWEEPS=decode_streaming python -m kernbench.cli.main \\
run --bench milestone-1h-gqa --topology topology.yaml
python scripts/paper/paper_plot_gqa_decode_streaming.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_decode_streaming.json"
_PAPER_FIG_DIR = (
_REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
)
# Per-rank HBM roofline: 8 pseudo-channels × 32 GB/s (topology.yaml
# hbm_ctrl.num_pcs / pc_bw_gbs; = pe_dma_to_noc_bw_gbs).
_PEAK_HBM_GBS = 256.0
_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["s_kv"], 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["s_kv_points"]
fig, (ax_lat, ax_bw) = 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, bw = _series(rows, v, "achieved_bw_gbs")
ax_bw.plot(xs, bw, marker=marker, color=color, label=label, lw=2)
for ax in (ax_lat, ax_bw):
ax.set_xscale("log", base=2)
ax.set_xticks(ctxs)
ax.set_xticklabels([_ctx_label(c) for c in ctxs])
ax.set_xlabel(r"per-rank context length $S_{kv}$ ($T_q{=}1$)")
ax.grid(True, ls=":", alpha=0.5)
ax.legend(fontsize=9)
ax_lat.set_ylabel("end-to-end decode latency (µs)")
ax_lat.set_title("Single-rank memory-bound decode latency per command form")
ax_bw.set_ylabel("achieved HBM bandwidth (GB/s)")
ax_bw.set_title(
"HBM bandwidth — composite keeps the DMA pipe full; primitive plateaus"
)
ax_bw.axhline(_PEAK_HBM_GBS, color="#888", ls="--", lw=1, alpha=0.7)
ax_bw.text(ctxs[0], _PEAK_HBM_GBS - 8, "256 GB/s roofline",
fontsize=8, color="#555", va="top")
ax_bw.set_ylim(0, _PEAK_HBM_GBS * 1.08)
fig.suptitle(
"Memory-bound decode streaming — use of composite commands\n"
"single-rank, GQA single-KV-head group ($h_q{=}8$, $d_{\\text{head}}"
"{=}128$); $M{=}8$ skinny, KV-streaming-bound",
fontsize=11,
)
fig.tight_layout(rect=(0, 0, 1, 0.92))
out = _FIG_DIR / "gqa_decode_streaming.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()
Binary file not shown.

After

Width:  |  Height:  |  Size: 136 KiB

@@ -0,0 +1,172 @@
{
"version": 1,
"variants": [
"primitive",
"composite",
"composite_extended"
],
"s_kv_points": [
2048,
4096,
8192,
16384
],
"rows": [
{
"variant": "primitive",
"s_kv": 2048,
"M": 8,
"latency_ns": 6188.437999999816,
"gemm_busy_ns": 1048.576000000001,
"dma_busy_ns": 6322.0,
"kv_bytes": 1048576.0,
"achieved_bw_gbs": 169.44114169036374,
"achieved_tflops": 1.3555291335229098,
"mac_util": 0.16944114169036373,
"dma_occupancy": 1.021582505957107
},
{
"variant": "composite",
"s_kv": 2048,
"M": 8,
"latency_ns": 4836.115999999989,
"gemm_busy_ns": 6629.632000000123,
"dma_busy_ns": 267480.720000013,
"kv_bytes": 1048576.0,
"achieved_bw_gbs": 216.82192900253062,
"achieved_tflops": 1.734575432020245,
"mac_util": 0.2168219290025306,
"dma_occupancy": 55.30899589671001
},
{
"variant": "composite_extended",
"s_kv": 2048,
"M": 8,
"latency_ns": 4737.751999999986,
"gemm_busy_ns": 8481.408000000116,
"dma_busy_ns": 251398.7400000122,
"kv_bytes": 1048576.0,
"achieved_bw_gbs": 221.32353065335693,
"achieved_tflops": 1.7705882452268553,
"mac_util": 0.22132353065335691,
"dma_occupancy": 53.06287454472352
},
{
"variant": "primitive",
"s_kv": 4096,
"M": 8,
"latency_ns": 12510.006000000496,
"gemm_busy_ns": 2097.152000000002,
"dma_busy_ns": 12602.0,
"kv_bytes": 2097152.0,
"achieved_bw_gbs": 167.6379691584414,
"achieved_tflops": 1.3411037532675312,
"mac_util": 0.1676379691584414,
"dma_occupancy": 1.0073536335633653
},
{
"variant": "composite",
"s_kv": 4096,
"M": 8,
"latency_ns": 9360.883999999554,
"gemm_busy_ns": 19520.799999939867,
"dma_busy_ns": 1059043.5999999335,
"kv_bytes": 2097152.0,
"achieved_bw_gbs": 224.03354213128802,
"achieved_tflops": 1.792268337050304,
"mac_util": 0.224033542131288,
"dma_occupancy": 113.13499878857424
},
{
"variant": "composite_extended",
"s_kv": 4096,
"M": 8,
"latency_ns": 9198.135999999562,
"gemm_busy_ns": 39531.583999941715,
"dma_busy_ns": 1026582.7399999356,
"kv_bytes": 2097152.0,
"achieved_bw_gbs": 227.99749862364504,
"achieved_tflops": 1.8239799889891604,
"mac_util": 0.22799749862364505,
"dma_occupancy": 111.60769312390951
},
{
"variant": "primitive",
"s_kv": 8192,
"M": 8,
"latency_ns": 25153.141999997388,
"gemm_busy_ns": 4194.304000000004,
"dma_busy_ns": 25162.0,
"kv_bytes": 4194304.0,
"achieved_bw_gbs": 166.7506985807354,
"achieved_tflops": 1.334005588645883,
"mac_util": 0.16675069858073538,
"dma_occupancy": 1.0003521627637062
},
{
"variant": "composite",
"s_kv": 8192,
"M": 8,
"latency_ns": 18419.187999999034,
"gemm_busy_ns": 64177.11999976149,
"dma_busy_ns": 4214541.840000685,
"kv_bytes": 4194304.0,
"achieved_bw_gbs": 227.713838416776,
"achieved_tflops": 1.821710707334208,
"mac_util": 0.227713838416776,
"dma_occupancy": 228.81257523409317
},
{
"variant": "composite_extended",
"s_kv": 8192,
"M": 8,
"latency_ns": 18128.439999999027,
"gemm_busy_ns": 169650.68799976518,
"dma_busy_ns": 4149323.2200006745,
"kv_bytes": 4194304.0,
"achieved_bw_gbs": 231.365964197704,
"achieved_tflops": 1.850927713581632,
"mac_util": 0.231365964197704,
"dma_occupancy": 228.88473691067168
},
{
"variant": "primitive",
"s_kv": 16384,
"M": 8,
"latency_ns": 50439.414000008954,
"gemm_busy_ns": 8388.607999999076,
"dma_busy_ns": 50282.0,
"kv_bytes": 8388608.0,
"achieved_bw_gbs": 166.31057609032712,
"achieved_tflops": 1.330484608722617,
"mac_util": 0.16631057609032712,
"dma_occupancy": 0.9968791469304357
},
{
"variant": "composite",
"s_kv": 16384,
"M": 8,
"latency_ns": 36535.79600000568,
"gemm_busy_ns": 228978.46400288073,
"dma_busy_ns": 16815028.239995122,
"kv_bytes": 8388608.0,
"achieved_bw_gbs": 229.59970545047648,
"achieved_tflops": 1.8367976436038118,
"mac_util": 0.22959970545047648,
"dma_occupancy": 460.2343477063565
},
{
"variant": "composite_extended",
"s_kv": 16384,
"M": 8,
"latency_ns": 35989.048000005685,
"gemm_busy_ns": 701990.3680028584,
"dma_busy_ns": 16684294.09999516,
"kv_bytes": 8388608.0,
"achieved_bw_gbs": 233.0877993771515,
"achieved_tflops": 1.864702395017212,
"mac_util": 0.2330877993771515,
"dma_occupancy": 463.5936493789034
}
]
}
@@ -0,0 +1,168 @@
"""milestone-1h-gqa: memory-bound decode streaming composite-command study.
Single-rank companion to the compute-bound prefill study
(``gqa_prefill_compute_bound``). Same three command-form kernels
(``_gqa_prefill_compute_bound``), same single-rank (C=P=1) harness, but
driven with **T_q=1** (decode) and swept over a large S_kv. With T_q=1 the
score / context GEMMs are skinny (M = G·T_q = 8), so the MAC array is
barely fed and the kernel is **memory-bound** — streaming the KV cache out
of HBM dominates. This is the regime mirror of prefill: command form is
*latency-neutral* here, because the bottleneck is data movement, not
issue.
Running single-rank (no cross-CUBE (m,,O) reduce) isolates the
local-attention streaming vs. issue trade-off cleanly — unlike the 64-way
Case-6 decode sweep, whose small-S_kv end-to-end latency is masked by the
inter-CUBE reduce tail. S_kv here is the *per-rank* context, so S_kv=64K
single-rank is the per-PE load of a 4M-token, 64-way-sharded decode.
Records per (variant, S_kv) the end-to-end latency, GEMM/DMA engine busy
time, MAC utilization (achieved ÷ peak), and DMA occupancy (dma_busy ÷
e2e), so the comparative plot can show all three command forms landing on
the same latency curve while the MAC array stays floored and the DMA
channel stays saturated.
Runs in data mode (engine latency). Gated via the umbrella
``GQA_1H_SWEEPS=decode_streaming``.
"""
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_decode_streaming.json"
# Same kernels as the prefill study — they are general attention kernels;
# the regime is set by T_q (=1 here → memory-bound).
_VARIANT_KERNELS = {
"primitive": gqa_prefill_primitive_kernel,
"composite": gqa_prefill_composite_kernel,
"composite_extended": gqa_prefill_composite_ext_kernel,
}
_VARIANTS = ("primitive", "composite", "composite_extended")
# Per-rank context length. T_q=1, so M = G·T_q = 8 (skinny, memory-bound)
# at every point. Capped at 16K: the *plain* composite materializes the
# full (M, S_kv) scores in TCM scratch (~48·S_kv B incl. the exp transients),
# which overflows the 1 MB kernel scratch beyond ~21K — itself a sign that
# the recipe (composite_extended), which tiles the softmax, is what long
# context actually needs. S_kv is per-rank, so 16K single-rank already
# equals the per-PE load of a 1M-token, 64-way-sharded decode.
_S_KV_POINTS = (2048, 4096, 8192, 16384)
_T_Q = 1
_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, s_kv: int):
kernel = _VARIANT_KERNELS[variant]
panel = f"decode_stream_{variant}_s{s_kv}"
def _bench_fn(ctx):
dp = DPPolicy(cube="replicate", pe="replicate",
num_cubes=1, num_pes=1)
q = ctx.zeros((_T_Q, _H_Q * _D_HEAD),
dtype="f16", dp=dp, name=f"{panel}_q")
k = ctx.zeros((s_kv, _H_KV * _D_HEAD),
dtype="f16", dp=dp, name=f"{panel}_k")
v = ctx.zeros((s_kv, _H_KV * _D_HEAD),
dtype="f16", dp=dp, name=f"{panel}_v")
o = ctx.empty((_T_Q, _H_Q * _D_HEAD),
dtype="f16", dp=dp, name=f"{panel}_o")
ctx.launch(panel, kernel, q, k, v, o,
_T_Q, s_kv, _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, s_kv: 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, s_kv),
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-decode-streaming {variant}@{s_kv} 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")
G = _H_Q // _H_KV
M = G * _T_Q
# Useful attention flops (Q·Kᵀ + P·V), single rank.
useful_flops = 4.0 * M * _D_HEAD * s_kv
# KV-cache bytes streamed from HBM (K + V, f16) — the memory-bound
# denominator. achieved_bw = bytes / e2e (GB/s) measures how close the
# command form gets to the HBM roofline (the memory-bound mirror of
# MAC utilization in the compute-bound prefill study).
kv_bytes = 2.0 * s_kv * _D_HEAD * 2.0
return {
"variant": variant,
"s_kv": s_kv,
"M": M,
"latency_ns": e2e,
"gemm_busy_ns": gemm,
"dma_busy_ns": dma,
"kv_bytes": kv_bytes,
"achieved_bw_gbs": (kv_bytes / e2e) if e2e > 0 else 0.0,
"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,
"dma_occupancy": (dma / e2e) if e2e > 0 else 0.0,
}
def run_sweep(topology: str = "topology.yaml") -> int:
"""Drive all (variant, S_kv) decode-streaming panels; write sweep.json."""
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
rows = [
_run_panel(variant, s_kv, topology)
for s_kv in _S_KV_POINTS
for variant in _VARIANTS
]
sweep = {
"version": 1,
"variants": list(_VARIANTS),
"s_kv_points": list(_S_KV_POINTS),
"rows": rows,
}
_SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
print(f" gqa-decode-streaming: {len(rows)} rows -> {_SWEEP_JSON}")
return len(rows)
if __name__ == "__main__":
run_sweep()
@@ -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_decode_streaming import (
run_sweep as _run_decode_streaming_sweep,
)
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
run_sweep as _run_prefill_cb_sweep,
)
@@ -66,6 +69,7 @@ def run(torch) -> None:
"decode": _run_decode_sweep,
"composite": _run_composite_sweep,
"prefill_cb": _run_prefill_cb_sweep,
"decode_streaming": _run_decode_streaming_sweep,
}
unknown = [s for s in sweeps if s not in runners]
if unknown: