[Fix] Restore the gfx950 XCD mapping for bf16 dense attention

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
 

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)

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

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