[Perf] Enable the XCD swizzle and pin the softmax basis for gfx950 bf16 attention

kernels/attention/flash_attn_gfx950.py

@@ -61,12 +61,14 @@ def build_flash_attn_dualwave_swp_module(
     dualwave_swp_setprio=True,
     dualwave_swp_debug_lazy_counts=False,
     dualwave_swp_enable_stagger=True,
+    dualwave_swp_fixed_basis=False,
     num_kv_splits=1,
     varlen=False,
     cross_seqlen=False,
     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.
 
@@ -111,13 +113,15 @@ def build_flash_attn_dualwave_swp_module(
         dualwave_swp_setprio=dualwave_swp_setprio,
         dualwave_swp_debug_lazy_counts=dualwave_swp_debug_lazy_counts,
         dualwave_swp_enable_stagger=dualwave_swp_enable_stagger,
+        dualwave_swp_fixed_basis=dualwave_swp_fixed_basis,
         num_kv_splits=num_kv_splits,
         varlen=varlen,
         cross_seqlen=cross_seqlen,
         paged=paged,
         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
@@ -263,7 +267,7 @@ def _main_body():
                     v_s_0 = softmax_helper.seq_pad_mask_if_needed(v_s_0, softmax_helper.split_tile(0))
                 else:
                     v_s_0 = softmax_helper.seq_pad_mask_if_needed(v_s_0, fx.Index(0))
-            m_row_pro = softmax_helper.reduce_max(v_s_0)
+            m_row_pro = softmax_helper.reduce_max_basis(v_s_0)
             if const_expr(traits.CAUSAL):
                 # Floor fully-masked rows (-inf) to finite so exp2 yields 0, not NaN.
                 m_row_pro = softmax_helper.floor_masked_max(m_row_pro)

kernels/attention/flash_attn_interface.py

@@ -29,7 +29,12 @@
 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 (
+    DUALWAVE_SWP_BLOCK_M,
+    MIN_Q_BLOCKS_XCD_SWIZZLE,
+    NUM_XCD_GFX950,
+    dualwave_splitk_workspace_elems,
+)
 
 __all__ = ["flydsl_flash_attn_func", "dualwave_splitk_workspace_elems"]
 
@@ -133,6 +138,8 @@ def _build_dense_dualwave(
     debug_lazy_counts: bool,
     enable_stagger: bool,
     return_lse: bool = False,
+    xcd_swizzle: bool = False,
+    fixed_basis: 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 +158,8 @@ 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,
+        dualwave_swp_fixed_basis=fixed_basis,
     )
 
 
@@ -625,6 +634,11 @@ def flydsl_flash_attn_func(
     dualwave_swp_lazy_rescale: bool = True,
     dualwave_swp_setprio: bool = True,
     dualwave_swp_enable_stagger: bool = True,
+    # Head-slow workgroup remap; None auto-selects. Bit-identical.
+    dualwave_swp_xcd_swizzle: Optional[bool] = None,
+    # Pin the softmax basis to the prologue tile max and fold the scale into
+    # sub_m. More accurate but not bit-identical, so opt-in.
+    dualwave_swp_fixed_basis: bool = False,
     # 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 +887,16 @@ 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)):
+                    # Workgroups map to XCDs as linear_id % 8, so a head-fast grid
+                    # with H % 8 == 0 pins each head to one XCD and leaves it H/8
+                    # concurrent K/V streams. The head-slow remap leaves one.
+                    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,
@@ -887,6 +911,8 @@ def flydsl_flash_attn_func(
                         debug_lazy_counts=debug_lazy,
                         enable_stagger=dualwave_swp_enable_stagger,
                         return_lse=return_lse,
+                        xcd_swizzle=xcd_swizzle,
+                        fixed_basis=dualwave_swp_fixed_basis,
                     )
                 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

@@ -31,6 +31,7 @@
 from kernels.common.kernels_common import dtype_to_elem_type
 
 _LOG2E = host_math.log2(host_math.e)
+DUALWAVE_SWP_BLOCK_M = 256
 # gfx950 (MI350/MI355X): 8 XCDs, each with a private ~4 MB L2.
 NUM_XCD_GFX950 = 8
 MIN_Q_BLOCKS_XCD_SWIZZLE = 64
@@ -1429,6 +1430,7 @@ class DualwaveSwpTraits:
     DUALWAVE_SWP_SETPRIO: bool
     DUALWAVE_SWP_DEBUG_LAZY_COUNTS: bool
     DUALWAVE_SWP_ENABLE_STAGGER: bool
+    DUALWAVE_SWP_FIXED_BASIS: bool
     NUM_KV_SPLITS: int
     SPLITK: bool
     PAGED: bool
@@ -1495,6 +1497,7 @@ def cache_tag(self):
             self.DUALWAVE_SWP_SETPRIO,
             self.DUALWAVE_SWP_DEBUG_LAZY_COUNTS,
             self.DUALWAVE_SWP_ENABLE_STAGGER,
+            self.DUALWAVE_SWP_FIXED_BASIS,
             self.NUM_KV_SPLITS,
             self.SPLITK,
             self.PAGED,
@@ -1519,6 +1522,7 @@ def _make_dualwave_swp_traits(
     dualwave_swp_setprio=True,
     dualwave_swp_debug_lazy_counts=False,
     dualwave_swp_enable_stagger=True,
+    dualwave_swp_fixed_basis=False,
     num_kv_splits=1,
     varlen=False,
     cross_seqlen=False,
@@ -1530,7 +1534,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
@@ -1627,6 +1631,7 @@ def _make_dualwave_swp_traits(
         DUALWAVE_SWP_SETPRIO=bool(dualwave_swp_setprio),
         DUALWAVE_SWP_DEBUG_LAZY_COUNTS=bool(dualwave_swp_debug_lazy_counts),
         DUALWAVE_SWP_ENABLE_STAGGER=bool(dualwave_swp_enable_stagger),
+        DUALWAVE_SWP_FIXED_BASIS=bool(dualwave_swp_fixed_basis),
         NUM_KV_SPLITS=num_kv_splits,
         SPLITK=splitk,
         PAGED=paged,
@@ -3651,6 +3656,9 @@ def load_all(self):
         return Vec(q_all, (traits.K_STEPS_QK * traits.MFMA_LANE_K,), self.elem_dtype)
 
     def scale_all(self, q_all_bf16):
+        if const_expr(self.traits.DUALWAVE_SWP_FIXED_BASIS):
+            # sub_m carries the scale, so Q is never narrowed a second time.
+            return q_all_bf16
         traits = self.traits
         fm_fast_attr = ir.Attribute.parse("#llvm.fastmath<fast>")
         v64bf16_type = Vec.make_type(traits.K_STEPS_QK * traits.MFMA_LANE_K, self.elem_dtype)
@@ -3707,9 +3715,15 @@ class DualwaveSoftmaxHelper(DualwaveKernelContext):
     def __init__(self, ctx):
         super().__init__(ctx)
 
-    def reduce_max(self, v_s):
+    def reduce_max_basis(self, v_s):
         return _score_pair_max(v_s, self.c_neg_inf, self.fm_fast)
 
+    def reduce_max(self, v_s):
+        if const_expr(self.traits.DUALWAVE_SWP_FIXED_BASIS):
+            # -inf leaves fmax(m_row, m_tile_max) == m_row, pinning the basis.
+            return self.c_neg_inf
+        return self.reduce_max_basis(v_s)
+
     def floor_masked_max(self, row_max):
         return _fmax(row_max, self.c_neg_floor, self.fm_fast)
 
@@ -3726,6 +3740,10 @@ def reduce_sum(self, l_row, v_p):
         return _fadd(l_row, _score_pair_sum(v_p, self.c_zero_f, self.fm_fast), self.fm_fast)
 
     def sub_m(self, v_s, row_max):
+        if const_expr(self.traits.DUALWAVE_SWP_FIXED_BASIS):
+            return _scale_sub_score_pair(
+                _score_lists_to_vecs(v_s), row_max, self.c_sm_scale_log2e, self.c_zero_f, self.fm_fast
+            )
         return _sub_score_pair(v_s, row_max, self.fm_fast)
 
     def cast_p(self, v_p):
@@ -3742,6 +3760,8 @@ def scale_o(self, v_o, scale_scalar):
         _scale_o_accs(v_o, scale_scalar, self.traits, self.fm_fast)
 
     def rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p):
+        if const_expr(self.traits.DUALWAVE_SWP_FIXED_BASIS):
+            return v_o, m_row, l_row, v_p
         m_new = _fmax(m_row, m_tile_max, self.fm_fast)
         corr = rocdl.exp2(T.f32, as_mlir_value(_fsub(m_row, m_new, self.fm_fast)))
         self.scale_o(v_o, corr)
@@ -3774,6 +3794,8 @@ def _lazy_rescale_o_rescale(self, _n, *_st, v_o, m_row, l_row, m_tile_max, v_p):
         return out
 
     def lazy_rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p):
+        if const_expr(self.traits.DUALWAVE_SWP_FIXED_BASIS):
+            return v_o, m_row, l_row, v_p
         traits = self.traits
         lane = self.lane
         debug_counts_rsrc = self.debug_counts_rsrc
@@ -4238,15 +4260,19 @@ def _store_empty_split():
         _store_empty_split()
 
     def _store_lse_row(self, m_row, l_row, q_row):
-        # LSE = m_row * ln2 + ln(l_row); natural log, scale folded (m_row is
-        # sm_scale*log2e-scaled). Fully-masked row has l_row == 0 -> -inf.
+        # LSE = m * ln2 + ln(l_row), natural log. m_row is log2-scaled unless the
+        # basis is pinned, in which case it is a raw logit and scales here.
+        # Fully-masked row has l_row == 0 -> -inf.
         traits = self.traits
         lse_base_i64 = fx.Int64(fx.ptrtoint(fx.get_iter(self.LSE)))
         lse_per_batch_elems = fx.Index(traits.NUM_HEADS_Q) * self.seq_len_v
         lse_per_batch_bytes = lse_per_batch_elems * fx.Index(4)
         lse_rsrc = _make_ws_rsrc(lse_base_i64, self.batch_idx * lse_per_batch_bytes, lse_per_batch_bytes)
+        m_log2 = m_row
+        if const_expr(self.traits.DUALWAVE_SWP_FIXED_BASIS):
+            m_log2 = _fmul(m_row, self.c_sm_scale_log2e, self.fm_fast)
         lse_val = _fadd(
-            _fmul(m_row, self.c_ln2_f, self.fm_fast),
+            _fmul(m_log2, self.c_ln2_f, self.fm_fast),
             fmath.log(as_mlir_value(l_row), fastmath=self.fm_fast),
             self.fm_fast,
         )

tests/kernels/test_flash_attn_fwd.py

@@ -3471,5 +3471,94 @@ 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 remap only relabels workgroups, so the output must not move a bit."""
+    S = 64 * 256  # clears the auto-dispatch threshold; below it this is vacuous
+    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 to the head-fast path, forced or auto."""
+    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)
+
+
+@_requires_gfx950
+@pytest.mark.parametrize("causal", [False, True])
+@pytest.mark.parametrize("H", [8, 32])
+def test_fixed_basis_improves_accuracy(H, causal):
+    """Pinning the basis drops a bf16 narrowing of Q, so the error must fall."""
+    S = 8192
+    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)
+    ref = F.scaled_dot_product_attention(
+        q.transpose(1, 2).float(), k.transpose(1, 2).float(), v.transpose(1, 2).float(), is_causal=causal
+    ).transpose(1, 2)
+
+    def err(fixed):
+        out = flydsl_flash_attn_func(q, k, v, causal=causal, dualwave_swp_fixed_basis=fixed)
+        torch.cuda.synchronize()
+        return ((out.float() - ref).norm() / ref.norm()).item(), out
+
+    e_stock, _ = err(False)
+    e_fixed, out_fixed = err(True)
+
+    assert torch.isfinite(out_fixed).all()
+    # Dropping the scale out of Q removes a whole bf16 narrowing, worth ~17-18%
+    # measured. Assert a much weaker 5% so this tracks the direction, not the
+    # exact figure, but still fails if the scale creeps back into Q.
+    assert e_fixed < e_stock * 0.95, f"stock {e_stock:.4e} -> fixed {e_fixed:.4e}"
+
+
+@_requires_gfx950
+@pytest.mark.parametrize("causal", [False, True])
+def test_fixed_basis_improves_lse(causal):
+    """LSE reads m directly, so it gains most from scaling m exactly."""
+    B, S, H, D = 1, 4096, 8, 128
+    torch.manual_seed(7)
+    q = _rand_lse(B, S, H, D, dtype=torch.bfloat16)
+    k, v = torch.randn_like(q), torch.randn_like(q)
+
+    scores = (q.transpose(1, 2).float() @ k.transpose(1, 2).float().transpose(-1, -2)) / D**0.5
+    if causal:
+        scores = scores.masked_fill(torch.triu(torch.ones(S, S, device=q.device, dtype=torch.bool), 1), float("-inf"))
+    ref = torch.logsumexp(scores, dim=-1)
+
+    def err(fixed):
+        _, lse = flydsl_flash_attn_func(q, k, v, causal=causal, return_lse=True, dualwave_swp_fixed_basis=fixed)
+        torch.cuda.synchronize()
+        return (lse.float() - ref).abs().max().item()
+
+    e_stock, e_fixed = err(False), err(True)
+    assert e_fixed < e_stock, f"stock {e_stock:.3e} -> fixed {e_fixed:.3e}"
+    assert e_fixed < 1e-4, f"LSE error {e_fixed:.3e} is far from the float32 floor"
+
+
 if __name__ == "__main__":
     main()