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