From 79ddb12b4292481cc4ca48d4a5dffada6a5c1b2c Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Wed, 10 Jun 2026 16:35:43 -0700 Subject: [PATCH] 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 --- .../benches/_gqa_attention_decode_long.py | 11 +++++++++++ .../benches/_gqa_attention_decode_short.py | 11 +++++++++++ .../benches/_gqa_attention_prefill_long.py | 16 ++++++++++++++-- .../benches/_gqa_attention_prefill_short.py | 11 +++++++++++ 4 files changed, 47 insertions(+), 2 deletions(-) diff --git a/src/kernbench/benches/_gqa_attention_decode_long.py b/src/kernbench/benches/_gqa_attention_decode_long.py index 5f3c9c3..576a4d1 100644 --- a/src/kernbench/benches/_gqa_attention_decode_long.py +++ b/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") diff --git a/src/kernbench/benches/_gqa_attention_decode_short.py b/src/kernbench/benches/_gqa_attention_decode_short.py index 3e244c1..48919f6 100644 --- a/src/kernbench/benches/_gqa_attention_decode_short.py +++ b/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") diff --git a/src/kernbench/benches/_gqa_attention_prefill_long.py b/src/kernbench/benches/_gqa_attention_prefill_long.py index b2cd21f..40a1ec9 100644 --- a/src/kernbench/benches/_gqa_attention_prefill_long.py +++ b/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") diff --git a/src/kernbench/benches/_gqa_attention_prefill_short.py b/src/kernbench/benches/_gqa_attention_prefill_short.py index 283e07e..c9a9c81 100644 --- a/src/kernbench/benches/_gqa_attention_prefill_short.py +++ b/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")