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>

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}"