diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index e72e7e163..586d4f530 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/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,6 +113,7 @@ 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, @@ -118,6 +121,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 @@ -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) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 5ebbd8207..ac9f62bc5 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/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) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index c8b5a7b47..b1014ddbd 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/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") 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, ) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 982e8876d..acb303a5b 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/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()