Enable the XCD head-slow remap on the bf16 dense attention path

The DUALWAVE_SWP dense kernel launches
grid=(NUM_HEADS_Q, num_q_blocks, grid_z), which is head-fast: one head's
q-blocks scatter across all 8 XCDs and every XCD re-streams that head's
K/V from HBM. With 256 CUs at one workgroup/CU the number of concurrent
distinct K/V streams is the head count, so at H=32/64 the working set
outgrows L2 and throughput decays as S grows.

Measured on MI355X, bf16, D=128, non-causal, B=2 H=32: the L2 hit rate
falls from 86% at S=27k to 37% at S=180k, where the kernel reads 1182 GB
from HBM against a compulsory 8.9 GB -- 133x -- while aiter holds 87% and
28x over the same range.

The fix already existed. _init_dualwave_thread_mapping re-derives
(head, q_block) with head as the slow axis, and _make_dualwave_swp_traits
already accepted xcd_swizzle; only the fp8 builder ever passed it, gated
on dtype_str == "fp8". This plumbs the flag through the bf16 builder and
reuses the fp8 dispatch condition (non-causal, H % NUM_XCD_GFX950 == 0,
num_q_blocks >= MIN_Q_BLOCKS_XCD_SWIZZLE, no split-K).

With the remap on, FlyDSL's memory profile lands on aiter's: 89.2% hit
and 28.1x re-stream at S=180k, against aiter's 87.1% and 28.2x.
Throughput goes flat at 1264-1272 TFLOP/s (84.5% of this box's 1497
TFLOP/s bf16 GEMM ceiling) across S=27k..240k, 2.7-3.1% above aiter at
every point. Best case is +8.5% at S=239,580 H=64.

The remap is bijective, so the output is bit-identical to the unpatched
kernel -- pinned by a new test. H=8 and H=16 already fit L2 and show no
decay; the remap measures as a no-op there (0.999-1.000x), so the
dispatch does not gate on head count.

dualwave_swp_xcd_swizzle joins the existing dualwave_swp_* build options:
None auto-selects, True/False force it, which is what lets the test drive
both paths through the public API.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: John Qin <yanyuan.qin@amd.com>

Author: JohnQinAMD

kernels/attention/flash_attn_gfx950.py

@@ -67,6 +67,7 @@ def build_flash_attn_dualwave_swp_module(
     paged=False,
     kv_cache_layout="linear",
     return_lse=False,
+    _xcd_swizzle=False,
 ):
     """Build an DUALWAVE_SWP flash_attn launcher for D=64/128 bf16/f16 on gfx950.
 
@@ -118,6 +119,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
     _dualwave_swp_cache_tag = traits.cache_tag

kernels/attention/flash_attn_interface.py

@@ -29,7 +29,11 @@
 import torch.nn.functional as F  # noqa: F401  (imported for callers' convenience)
 
 # Re-export so callers only need to import from this module.
-from kernels.attention.flash_attn_utils import dualwave_splitk_workspace_elems
+from kernels.attention.flash_attn_utils import (
+    MIN_Q_BLOCKS_XCD_SWIZZLE,
+    NUM_XCD_GFX950,
+    dualwave_splitk_workspace_elems,
+)
 
 __all__ = ["flydsl_flash_attn_func", "dualwave_splitk_workspace_elems"]
 
@@ -133,6 +137,7 @@ def _build_dense_dualwave(
     debug_lazy_counts: bool,
     enable_stagger: bool,
     return_lse: 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
@@ -151,6 +156,7 @@ def _build_dense_dualwave(
         dualwave_swp_debug_lazy_counts=debug_lazy_counts,
         dualwave_swp_enable_stagger=enable_stagger,
         return_lse=return_lse,
+        _xcd_swizzle=xcd_swizzle,
     )
 
 
@@ -625,6 +631,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,
@@ -873,6 +883,20 @@ def flydsl_flash_attn_func(
                         "flydsl_flash_attn_func: debug_counts requires the gfx950 DUALWAVE_SWP path"
                     )
                 if debug_lazy or (can_dualwave and _dense_routes_to_dualwave(B, Sq)):
+                    # Head-fast grid scatters one head's q-blocks across all XCDs, so
+                    # every XCD re-streams that head's K/V and L2 residency collapses
+                    # as S grows (measured: 86% hit at S=27k down to 37% at S=180k,
+                    # 133x the compulsory HBM read). The head-slow remap in
+                    # _init_dualwave_thread_mapping fixes it and is bijective, so the
+                    # output is bit-identical. The fp8 path already dispatches this
+                    # way; bf16 never plumbed the flag through.
+                    num_q_blocks = (int(Sq) + 255) // 256
+                    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,
@@ -887,6 +911,7 @@ def flydsl_flash_attn_func(
                         debug_lazy_counts=debug_lazy,
                         enable_stagger=dualwave_swp_enable_stagger,
                         return_lse=return_lse,
+                        xcd_swizzle=xcd_swizzle,
                     )
                 else:
                     block_m, flat_work_group_size, path_tag = _dense_generic_tile(B, Sq, H, D, dtype_str, q.device)

tests/kernels/test_flash_attn_fwd.py

@@ -3471,5 +3471,52 @@ 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)
+
+    outs = [
+        flydsl_flash_attn_func(q, k, v, causal=False, dualwave_swp_xcd_swizzle=flag).clone() for flag in (False, True)
+    ]
+    torch.cuda.synchronize()
+    assert torch.equal(outs[0], outs[1])
+
+
+@_requires_gfx950
+def test_xcd_swizzle_skipped_when_heads_not_multiple_of_xcd():
+    """H % 8 != 0 must fall back rather than mis-map.
+
+    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. The guard is in the dispatch condition, so this
+    checks the kernel still produces a correct answer for such a shape.
+    """
+    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)
+    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()