diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index 4ba01862e..bc59daa8e 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -85,6 +85,7 @@ def build_flash_attn_dualwave_swp_module( has_bias=False, has_alibi=False, has_sink=False, + _xcd_swizzle=False, ): """Build an DUALWAVE_SWP flash_attn launcher for D=64/128 bf16/f16 on gfx950. @@ -146,6 +147,7 @@ def build_flash_attn_dualwave_swp_module( kv_cache_layout=kv_cache_layout, kv_vectorized=KV_VECTORIZED, return_lse=return_lse, + xcd_swizzle=_xcd_swizzle, ) traits.BLOCK_N_OUT // traits.BLOCK_N diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 442af8511..f716c5f29 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -28,7 +28,14 @@ import torch import torch.nn.functional as F # noqa: F401 (imported for callers' convenience) -from kernels.attention.flash_attn_utils import bias_addressing_error, dualwave_splitk_workspace_elems +# Re-export so callers only need to import from this module. +from kernels.attention.flash_attn_utils import ( + DUALWAVE_SWP_BLOCK_M, + MIN_Q_BLOCKS_XCD_SWIZZLE, + NUM_XCD_GFX950, + bias_addressing_error, + dualwave_splitk_workspace_elems, +) __all__ = ["flydsl_flash_attn_func", "dualwave_splitk_workspace_elems"] @@ -135,6 +142,7 @@ def _build_dense_dualwave( has_bias: bool = False, has_alibi: bool = False, has_sink: bool = False, + xcd_swizzle: bool = False, ): """Build (and cache) the dense gfx950 DUALWAVE_SWP launcher.""" from kernels.attention.flash_attn_gfx950 import build_flash_attn_dualwave_swp_module @@ -156,6 +164,7 @@ def _build_dense_dualwave( has_bias=has_bias, has_alibi=has_alibi, has_sink=has_sink, + _xcd_swizzle=xcd_swizzle, ) @@ -683,6 +692,10 @@ def flydsl_flash_attn_func( dualwave_swp_lazy_rescale: bool = True, dualwave_swp_setprio: bool = True, dualwave_swp_enable_stagger: bool = True, + # Re-derive (head, q_block) with head as the slow axis so one head's q-blocks + # stay on one XCD instead of every XCD re-streaming that head's K/V. None + # auto-selects on the shapes it helps; True/False force it. Dense non-fp8 only. + dualwave_swp_xcd_swizzle: Optional[bool] = None, # Debug: pass a pre-allocated float32[2] tensor to enable the lazy-rescale # branch counter (dualwave_swp_debug_lazy_counts=True). Only for dense mode. debug_counts: Optional[torch.Tensor] = None, @@ -1059,6 +1072,24 @@ def flydsl_flash_attn_func( or has_sink or (can_dualwave and _dense_routes_to_dualwave(B, Sq)) ): + # Workgroups map to XCDs as linear_id % 8, and linear_id is + # bx + by*H + bz*H*nqb, so with H % 8 == 0 a head-fast grid pins + # head h to XCD h % 8. The ~256 resident workgroups span all H + # heads within one batch, leaving each XCD to juggle H/8 K/V + # streams against its L2 slice. The head-slow remap in + # _init_dualwave_thread_mapping puts the resident window inside a + # single head instead: 1 stream. Measured penalty for leaving it + # off tracks H/8 (-6% at 8 streams, -3% at 4, nil at <=2), not the + # hit rate and not traffic volume. Bijective, so output is + # unchanged. NB: that function's own comment states the opposite + # rationale ("scatter across all XCDs") and is wrong. + num_q_blocks = -(-int(Sq) // DUALWAVE_SWP_BLOCK_M) + if dualwave_swp_xcd_swizzle is None: + xcd_swizzle = ( + not causal and H % NUM_XCD_GFX950 == 0 and num_q_blocks >= MIN_Q_BLOCKS_XCD_SWIZZLE + ) + else: + xcd_swizzle = dualwave_swp_xcd_swizzle exe = _build_dense_dualwave( num_heads=H, num_kv_heads=num_kv_heads, @@ -1076,6 +1107,7 @@ def flydsl_flash_attn_func( has_bias=has_bias, has_alibi=has_alibi, has_sink=has_sink, + xcd_swizzle=xcd_swizzle, ) else: block_m, flat_work_group_size, path_tag = _dense_generic_tile(B, Sq, H, D, dtype_str, q.device) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index ccf96a41e..2a269dbd8 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -32,6 +32,9 @@ # gfx950 (MI350/MI355X): 8 XCDs, each with a private ~4 MB L2. NUM_XCD_GFX950 = 8 MIN_Q_BLOCKS_XCD_SWIZZLE = 64 +# The dual-wave 8-wave CTA fixes the q-block height; callers need it to count +# q-blocks before any traits object exists. +DUALWAVE_SWP_BLOCK_M = 256 # s_waitcnt bitfield encoding _VMCNT_LO_MASK = 0xF _LGKMCNT_EXPCNT_BASE = 0x3F70 @@ -1607,7 +1610,7 @@ def _make_dualwave_swp_traits( ): """Build gfx950 DUALWAVE_SWP compile-time layout traits.""" # Tile shape and wave geometry follow the gfx950 dual-wave 8-wave CTA. - block_m = 256 + block_m = DUALWAVE_SWP_BLOCK_M block_n = 64 block_n_out = 64 k_sub_n = 32 diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index ec89e419b..a30a79709 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -4329,6 +4329,57 @@ def test_return_lse_rejects_fp8(): ) +@_requires_gfx950 +@pytest.mark.parametrize("H", [8, 16, 32, 64]) +def test_xcd_swizzle_is_bit_identical(H): + """The head-slow remap must not change a single bit of the output. + + It only re-derives (head, q_block) from the same linear workgroup id, so it + is bijective by construction -- but a mistake in the derivation would show + up as a permuted or partially-recomputed output rather than as an error, so + this pins it. S clears the auto-dispatch threshold (num_q_blocks >= 64 at + BLOCK_M=256) so both settings run on the shapes the remap targets. + """ + S = 64 * 256 + dtype = torch.bfloat16 + torch.manual_seed(H) + q = _rand_lse(1, S, H, 128, dtype=dtype) + k, v = torch.randn_like(q), torch.randn_like(q) + + def run(flag): + return flydsl_flash_attn_func(q, k, v, causal=False, dualwave_swp_xcd_swizzle=flag).clone() + + off, on = run(False), run(True) + torch.cuda.synchronize() + assert torch.equal(off, on) + + +@_requires_gfx950 +@pytest.mark.parametrize("xcd_swizzle", [None, True]) +def test_xcd_swizzle_heads_not_multiple_of_xcd(xcd_swizzle): + """H % 8 != 0 must fall back rather than mis-map, however the flag is set. + + The remap divides the linear workgroup id by the q-block count to recover + the head, which only lands each head on one XCD when the head count divides + evenly into the 8 XCDs. Two guards enforce that: the dispatch condition + below auto-selects against it, and _init_dualwave_thread_mapping re-checks + NUM_HEADS_Q % NUM_XCD_GFX950 independently -- so forcing the flag on is safe + and simply does not engage the remap. Both paths are checked here. + """ + S, H = 64 * 256, 12 + dtype = torch.bfloat16 + torch.manual_seed(H) + q = _rand_lse(1, S, H, 128, dtype=dtype) + k, v = torch.randn_like(q), torch.randn_like(q) + + out = flydsl_flash_attn_func(q, k, v, causal=False, dualwave_swp_xcd_swizzle=xcd_swizzle) + torch.cuda.synchronize() + ref = F.scaled_dot_product_attention( + q.transpose(1, 2).float(), k.transpose(1, 2).float(), v.transpose(1, 2).float() + ).transpose(1, 2) + torch.testing.assert_close(out.float(), ref, atol=_ATOL_BF16, rtol=0) + + if __name__ == "__main__": main()