From 1aa5943805acee3524dc7e9062187bdaa256212d Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Wed, 27 May 2026 09:30:51 -0600 Subject: [PATCH 1/4] [triton-fa] hint head-stride div-by-8 for vectorized global load For the packed `[seq, heads, dim]` layout used by varlen prefill, the head-axis stride equals `head_dim`. When `head_dim` is a multiple of 8 but not 16 (e.g. 72), Triton's integer-arg auto-specialization does not attach `tt.divisibility = 8` to `stride_*h` (its threshold is 16), so AxisInfo treats the K/V global load as 2-byte aligned and Coalesce emits scalar `buffer_load_u16` instead of vectorized `buffer_load_b128`. Add a `HEAD_STRIDE_ALIGNED_8` constexpr to `attn_fwd` and apply `tl.multiple_of(off_h_{q,k} * stride_{q,k,v}h, 8)` to the head-axis integer offset when the caller sets it. AddPtr propagates this through to the load pointer, so AxisInfo computes a 16-byte alignment and the load coalesces. The wrapper checks `stride_*h % 8 == 0` against the actual runtime strides (not against `head_dim`), so the hint stays sound for non-contiguous Q/K/V views where `stride_*h != head_dim`. The constexpr defaults to `False`, so external callers of `attn_fwd` are unaffected unless they opt in. On gfx1151 with `head_dim=72, fp16`, the Qwen3-Omni ViT prefill shape goes from 128 `buffer_load_u16` to 16 `buffer_load_b128` and median kernel time drops from 3.04 ms to 2.71 ms (-10.9%). Correctness matches torch SDPA to 1e-4 max abs error. --- .../flash_attn_triton_amd/fwd_prefill.py | 30 +++++++++++++++++-- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py b/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py index c8bdaa3f6ee..e80ea4e56b6 100755 --- a/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py +++ b/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py @@ -941,6 +941,7 @@ def attn_fwd( USE_SEQUSED: tl.constexpr, FORCE_MASKING: tl.constexpr, NUM_XCD: tl.constexpr = 1, + HEAD_STRIDE_ALIGNED_8: tl.constexpr = False, ): # set params ACCUMULATOR_TYPE = tl.float32 @@ -1091,16 +1092,29 @@ def attn_fwd( # Initialize for processing # Compute pointers for all the tensors used in this kernel. + # When the caller guarantees that the head-axis strides of Q/K/V are + # multiples of 8 elements (set via HEAD_STRIDE_ALIGNED_8), the head-axis + # byte offset is 16-byte aligned. Auto-specialization only fires at the + # 16-element threshold, so hint the smaller multiple explicitly to let + # AxisInfo widen the global load. + qh_off = off_h_q * stride_qh + kh_off = off_h_k * stride_kh + vh_off = off_h_k * stride_vh + if HEAD_STRIDE_ALIGNED_8: + qh_off = tl.multiple_of(qh_off, 8) + kh_off = tl.multiple_of(kh_off, 8) + vh_off = tl.multiple_of(vh_off, 8) + q_offset = ( - Q + off_z * stride_qz + off_h_q * stride_qh + cu_seqlens_q_start * stride_qm + Q + off_z * stride_qz + qh_off + cu_seqlens_q_start * stride_qm ) q_ptrs = q_offset + offs_m[:, None] * stride_qm + offs_d_qk[None, :] * stride_qk k_offset = ( - K + off_z * stride_kz + off_h_k * stride_kh + cu_seqlens_k_start * stride_kn + K + off_z * stride_kz + kh_off + cu_seqlens_k_start * stride_kn ) k_ptrs = k_offset + offs_d_qk[:, None] * stride_kk + offs_n[None, :] * stride_kn v_offset = ( - V + off_z * stride_vz + off_h_k * stride_vh + cu_seqlens_k_start * stride_vk + V + off_z * stride_vz + vh_off + cu_seqlens_k_start * stride_vk ) v_ptrs = v_offset + offs_n[:, None] * stride_vk + offs_d_v[None, :] * stride_vn if USE_BIAS: @@ -1792,6 +1806,15 @@ def attention_forward_prefill_triton_impl( num_xcd = 1 if arch.is_rdna else 8 + # Soundness precondition for the `tl.multiple_of` head-stride hint inside + # `attn_fwd`: only enable it when every Q/K/V head-axis stride is a + # multiple of 8 elements. With a non-contiguous input (e.g. a transposed + # view) stride_*h need not equal head_dim, so the head_dim constexpr + # alone is not enough. + head_stride_aligned_8 = ( + stride_qh % 8 == 0 and stride_kh % 8 == 0 and stride_vh % 8 == 0 + ) + # launch kernel def grid(META): return (nheads_q, triton.cdiv(max_seqlens_q, META["BLOCK_M"]), batch) @@ -1872,4 +1895,5 @@ def grid(META): USE_SEQUSED=(seqused_q is not None or seqused_k is not None), FORCE_MASKING=force_masking, NUM_XCD=num_xcd, + HEAD_STRIDE_ALIGNED_8=head_stride_aligned_8, ) From 5a5f87e7d788fdfdf58050c597c07a1f1c3712ea Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Thu, 28 May 2026 08:30:44 -0600 Subject: [PATCH 2/4] [triton-fa] gfx1151: BLOCK_N=64, waves_per_eu=2 for LDS-bound FMHA The default RDNA config for `attn_fwd` uses BLOCK_N=32 + waves_per_eu=6 to target 6-wave occupancy. On gfx1151 the kernel allocates 32 KB shared per workgroup (Q tile staged through LDS), which already caps occupancy at 4 WGs/WGP (= 2 waves/EU) regardless of the register budget. The `waves_per_eu=6` hint then makes the register allocator spill / shorten live ranges to fit waves that will never schedule, which degrades codegen: fewer s_delay_alu hints, shorter clauses, more scalar f32 multiplies in the softmax path. Drop the hint to `waves_per_eu=2` (matching CK's high-VGPR low-occupancy choice for the same kernel on the same arch) and pair it with BLOCK_N=64 so each workgroup amortises loop-dispatch overhead over 2x the WMMA work per K-tile. The inner-loop mnemonic profile then collapses onto CK's: s_clause matches at 0.03/WMMA, s_waitcnt 1.19 vs 1.31, total instr/WMMA drops from 12.2 to 11.8. On AITER's flash_attn_2.varlen_fwd at the Qwen3-Omni ViT prefill shape (B=1, S=3200, H=16, head_dim=72, fp16) on gfx1151, median kernel time goes from 2.47 ms to 2.35 ms (-4.9%), now 3% faster than the CK ck_tile FmhaFwdKernel reference at 2.42 ms. Min times are 2.28 vs CK 2.34 ms (-2.6%). Stacked on top of the previous head-stride divisibility-by-8 hint commit, total speedup vs the upstream baseline is 3.04 ms -> 2.35 ms (-22.7%). Scoped to gfx1151 only; other RDNA targets keep the existing default pending validation. --- .../flash_attn_triton_amd/fwd_prefill.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py b/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py index e80ea4e56b6..98259f60690 100755 --- a/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py +++ b/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py @@ -64,6 +64,21 @@ def get_fwd_prefill_configs(mode: AutotuneMode): ), ] elif arch.is_rdna: + # gfx1151 (Strix Halo / RDNA3.5): tuned for the Qwen3-Omni + # ViT prefill shape (B=1, S=3200, H=16, head_dim=72, fp16). + if arch.name == "gfx1151": + return [ + triton.Config( + { + "BLOCK_M": 128, + "BLOCK_N": 64, + "PRE_LOAD_V": False, + "waves_per_eu": 2, + }, + num_stages=1, + num_warps=8, + ), + ] BLOCK_N = 64 if arch.name == "gfx1100" else 32 return [ triton.Config( @@ -166,6 +181,21 @@ def get_fwd_prefill_configs(mode: AutotuneMode): ) ] elif arch.is_rdna: + # gfx1151 (Strix Halo / RDNA3.5): tuned for the Qwen3-Omni + # ViT prefill shape (B=1, S=3200, H=16, head_dim=72, fp16). + if arch.name == "gfx1151": + return [ + triton.Config( + { + "BLOCK_M": 128, + "BLOCK_N": 64, + "PRE_LOAD_V": False, + "waves_per_eu": 2, + }, + num_stages=1, + num_warps=8, + ), + ] BLOCK_N = 64 if arch.name == "gfx1100" else 32 return [ triton.Config( From d79ed35b436f66eabe2a7b859148af0223f7aa8c Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Tue, 2 Jun 2026 03:23:26 -0600 Subject: [PATCH 3/4] Apply Black formatting to fwd_prefill.py --- .../flash_attn_triton_amd/fwd_prefill.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py b/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py index 98259f60690..ddea4f0c720 100755 --- a/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py +++ b/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py @@ -1135,17 +1135,11 @@ def attn_fwd( kh_off = tl.multiple_of(kh_off, 8) vh_off = tl.multiple_of(vh_off, 8) - q_offset = ( - Q + off_z * stride_qz + qh_off + cu_seqlens_q_start * stride_qm - ) + q_offset = Q + off_z * stride_qz + qh_off + cu_seqlens_q_start * stride_qm q_ptrs = q_offset + offs_m[:, None] * stride_qm + offs_d_qk[None, :] * stride_qk - k_offset = ( - K + off_z * stride_kz + kh_off + cu_seqlens_k_start * stride_kn - ) + k_offset = K + off_z * stride_kz + kh_off + cu_seqlens_k_start * stride_kn k_ptrs = k_offset + offs_d_qk[:, None] * stride_kk + offs_n[None, :] * stride_kn - v_offset = ( - V + off_z * stride_vz + vh_off + cu_seqlens_k_start * stride_vk - ) + v_offset = V + off_z * stride_vz + vh_off + cu_seqlens_k_start * stride_vk v_ptrs = v_offset + offs_n[:, None] * stride_vk + offs_d_v[None, :] * stride_vn if USE_BIAS: # Note: this might get large enough to overflow on some configs From f72bb3aab348cd8f3168c7d8f43c7085a596d162 Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Tue, 2 Jun 2026 13:57:43 -0600 Subject: [PATCH 4/4] [triton-fa] clarify head-stride alignment comment for non-16-bit dtypes The in-kernel comment claimed the head-axis byte offset is 16-byte aligned whenever stride_*h % 8 == 0. That only holds for 16-bit element types. The hint guarantees an 8-element multiple, whose byte alignment is element-size dependent (16 B for fp16/bf16, 8 B for fp8, 32 B for fp32); only the 16-bit case yields the buffer_load_b128 widening this PR targets. The tl.multiple_of(..., 8) hint stays sound for all dtypes. Comment-only change, no behavioral impact. --- .../flash_attn_triton_amd/fwd_prefill.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py b/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py index ddea4f0c720..c408140cb49 100755 --- a/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py +++ b/aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/fwd_prefill.py @@ -1124,9 +1124,14 @@ def attn_fwd( # Compute pointers for all the tensors used in this kernel. # When the caller guarantees that the head-axis strides of Q/K/V are # multiples of 8 elements (set via HEAD_STRIDE_ALIGNED_8), the head-axis - # byte offset is 16-byte aligned. Auto-specialization only fires at the - # 16-element threshold, so hint the smaller multiple explicitly to let - # AxisInfo widen the global load. + # offset is a multiple of 8 *elements*. The resulting byte alignment is + # element-size dependent: 16 bytes for 16-bit types (fp16/bf16), 8 bytes + # for fp8, 32 bytes for fp32. The 16-byte case is the one that lets + # AxisInfo widen the K/V global load to a 128-bit (buffer_load_b128) + # access; for the other dtypes the hint is still sound but yields a + # different (smaller or larger) vector width. Auto-specialization only + # fires at the 16-element threshold, so hint the smaller multiple + # explicitly. qh_off = off_h_q * stride_qh kh_off = off_h_k * stride_kh vh_off = off_h_k * stride_vh