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:
2026-06-23 17:04:26 -07:00
parent cd6f0ed91d
commit 24c705419e
13 changed files with 970 additions and 5 deletions
@@ -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),