gqa: explain why Tile-0 / bootstrap can't be folded into the tile loop

Comment-only change. Adds a kernbench-only limitation note at the
Tile-0 / bootstrap block in all four GQA attention kernels:

  - persistent (m, l, O) must live outside tl.scratch_scope so
    subsequent tiles' merges can read them;
  - kernbench has no scratch-backed initializer (tl.zeros / tl.full
    return addr=0 handles), so we can't seed (-inf, 0, 0) and rely
    on tl.copy_to;
  - therefore Tile 0 must compute the initial running state directly.

prefill_long adds a third bullet noting the ring-step ordering reason
(k=0 must send W before any k>0 can recv E, so the (t=0, k=0) step has
to run outside the scoped loop where the send is conditional on C > 1).

Each block also notes that a real Triton port collapses Tile 0 into a
unified loop (SSA tensors stay live across iterations).

No behavior change; tests unchanged.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

src/kernbench/benches/_gqa_attention_decode_long.py

@@ -73,6 +73,17 @@ def gqa_attention_decode_long_kernel(
     KV_ROW_BYTES = d_head * 2  # f16
 
     # Tile 0: establishes persistent (m_local, l_local, O_local).
+    #
+    # Cannot be folded into the Tiles 1..N loop (kernbench-only limitation):
+    #   - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
+    #     otherwise scope teardown discards them before the next tile's
+    #     merge can read them;
+    #   - kernbench has no scratch-backed initializer — ``tl.zeros`` /
+    #     ``tl.full`` return addr=0 handles with no backing storage, so
+    #     they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0).
+    # So Tile 0 computes the initial running state directly; Tiles 1..N
+    # fold into it. Triton port: limitation does not apply (SSA tensors
+    # stay live across iterations) — a single unified loop suffices.
     tile_s0 = min(TILE_S_KV, S_local)
     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")

src/kernbench/benches/_gqa_attention_decode_short.py

@@ -72,6 +72,17 @@ def gqa_attention_decode_short_kernel(
     KV_ROW_BYTES = d_head * 2  # f16
 
     # Tile 0: establishes persistent (m_local, l_local, O_local).
+    #
+    # Cannot be folded into the Tiles 1..N loop (kernbench-only limitation):
+    #   - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
+    #     otherwise scope teardown discards them before the next tile's
+    #     merge can read them;
+    #   - kernbench has no scratch-backed initializer — ``tl.zeros`` /
+    #     ``tl.full`` return addr=0 handles with no backing storage, so
+    #     they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0).
+    # So Tile 0 computes the initial running state directly; Tiles 1..N
+    # fold into it. Triton port: limitation does not apply (SSA tensors
+    # stay live across iterations) — a single unified loop suffices.
     tile_s0 = min(TILE_S_KV, S_local)
     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")

src/kernbench/benches/_gqa_attention_prefill_long.py

@@ -87,8 +87,20 @@ def gqa_attention_prefill_long_kernel(
     Q = tl.load(q_ptr, shape=(T_q, d_head), dtype="f16")
 
     # ── Bootstrap: (t=0, k=0) — load my own tile 0, establish (m, ℓ, O) ──
-    # Outside ``scratch_scope`` so the persistent running state survives
-    # subsequent ring iterations.
+    #
+    # Cannot be folded into the (t, k) loop (kernbench-only limitation):
+    #   - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
+    #     otherwise scope teardown discards them before the next tile's
+    #     merge can read them;
+    #   - kernbench has no scratch-backed initializer — ``tl.zeros`` /
+    #     ``tl.full`` return addr=0 handles with no backing storage, so
+    #     they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0);
+    #   - the ring-step ordering also requires k=0 to send its loaded
+    #     tile W before any k>0 iteration can recv from E, so the very
+    #     first (t=0, k=0) step has to run outside the scoped loop body
+    #     where the send is conditional on ``C > 1``.
+    # Triton port: limitation does not apply (SSA tensors stay live
+    # across iterations); the bootstrap collapses into the loop.
     tile_s = min(TILE_S_KV, S_local)
     K_t = tl.load(k_ptr, shape=(d_head, tile_s), dtype="f16")
     V_t = tl.load(v_ptr, shape=(tile_s, d_head), dtype="f16")

src/kernbench/benches/_gqa_attention_prefill_short.py

@@ -54,6 +54,17 @@ def gqa_attention_prefill_short_kernel(
     KV_ROW_BYTES = d_head * 2  # f16
 
     # Tile 0: establishes persistent (m_local, l_local, O_local).
+    #
+    # Cannot be folded into the Tiles 1..N loop (kernbench-only limitation):
+    #   - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
+    #     otherwise scope teardown discards them before the next tile's
+    #     merge can read them;
+    #   - kernbench has no scratch-backed initializer — ``tl.zeros`` /
+    #     ``tl.full`` return addr=0 handles with no backing storage, so
+    #     they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0).
+    # So Tile 0 computes the initial running state directly; Tiles 1..N
+    # fold into it. Triton port: limitation does not apply (SSA tensors
+    # stay live across iterations) — a single unified loop suffices.
     tile_s0 = min(TILE_S_KV, S_local)
     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")