diff --git a/flash_attn/cute/flash_bwd_sm100.py b/flash_attn/cute/flash_bwd_sm100.py index 11db2dab563..273cd44f6f2 100644 --- a/flash_attn/cute/flash_bwd_sm100.py +++ b/flash_attn/cute/flash_bwd_sm100.py @@ -82,13 +82,7 @@ def __init__( assert self.tile_hdim <= 128 or (self.tile_hdim == 192 and self.tile_hdimv == 128) assert self.tile_hdimv <= 128 - self.use_2cta_instrs = bool( - use_2cta_instrs - and cluster_size == 2 - and score_mod is None - and score_mod_bwd is None - and mask_mod is None - ) + self.use_2cta_instrs = bool(use_2cta_instrs and cluster_size == 2) self.cta_group_size = 2 if self.use_2cta_instrs else 1 assert self.tile_hdim != 192 or self.use_2cta_instrs, "Must use 2CTA for hdim 192" @@ -2758,9 +2752,15 @@ def apply_score_mod( fastdiv_mods=(None, None), ): """Apply forward score modification for SM100 backward pass.""" - # In bwd, S is computed as K @ Q.T so dimensions are (tile_n, tile_m) - cS = cute.make_identity_tensor((self.tile_n, self.tile_m)) - cS = cute.domain_offset((n_block * self.tile_n, m_block * self.tile_m), cS) + # In bwd, S is computed as K @ Q.T so dimensions are (tile_n, tile_m). + # With 2CTA, partition_C must see the full cluster tile so each CTA + # gets its own half of the tile. + cluster_tile_n = self.tile_n * self.cta_group_size + cluster_n_block = n_block // self.cta_group_size + cS = cute.make_identity_tensor((cluster_tile_n, self.tile_m)) + cS = cute.domain_offset( + (cluster_n_block * cluster_tile_n, m_block * self.tile_m), cS + ) tScS = thr_mma_S.partition_C(cS) tScS_idx = thr_copy_t2r.partition_D(tScS) @@ -2976,13 +2976,13 @@ def compute_loop( seqlen, n_block // self.cluster_shape_mnk[0] ) mask = AttentionMaskCls(seqlen) - n_block_for_cluster = n_block // self.cta_group_size + cluster_n_block = n_block // self.cta_group_size # TODO: condition mask_seqlen mask_fn = partial( mask.apply_mask_sm100_transposed, tScS_t2r=tScS_t2r, t0ScS_t2r=t0ScS_t2r, - n_block=n_block_for_cluster, + n_block=cluster_n_block, mask_seqlen=True, mask_causal=self.is_causal, mask_local=self.is_local, @@ -3194,9 +3194,12 @@ def compute_loop( if const_expr(self.score_mod_bwd is not None): tSrS_pre_cur = tSrS_pre[None, stage, 0, 0] - cS_bwd = cute.make_identity_tensor((self.tile_n, self.tile_m)) + cluster_tile_n = self.tile_n * self.cta_group_size + cluster_n_block = n_block // self.cta_group_size + cS_bwd = cute.make_identity_tensor((cluster_tile_n, self.tile_m)) cS_bwd = cute.domain_offset( - (n_block * self.tile_n, m_block * self.tile_m), cS_bwd + (cluster_n_block * cluster_tile_n, m_block * self.tile_m), + cS_bwd, ) tScS_bwd = thr_mma_S.partition_C(cS_bwd) tScS_idx_bwd = thr_copy_t2r.partition_D(tScS_bwd) diff --git a/flash_attn/cute/interface.py b/flash_attn/cute/interface.py index 0bd4651190c..9e423c16ed9 100644 --- a/flash_attn/cute/interface.py +++ b/flash_attn/cute/interface.py @@ -1323,9 +1323,6 @@ def _flash_attn_bwd( requested_disable_2cta = utils._get_disable_2cta_default() disable_2cta = ( requested_disable_2cta - or score_mod is not None - or score_mod_bwd is not None - or mask_mod is not None or block_sparse_tensors is not None ) cluster_size = 2 if head_dim >= 128 and not disable_2cta else 1 diff --git a/tests/cute/test_flash_attn.py b/tests/cute/test_flash_attn.py index 21ed3a48d57..4b5fe669004 100644 --- a/tests/cute/test_flash_attn.py +++ b/tests/cute/test_flash_attn.py @@ -349,11 +349,10 @@ def test_flash_attn_output( and not has_qv and not dv > 256 and not attention_chunk != 0 - and softcap == 0.0 and ( (dv == d and d <= 128) or (d == 192 and dv == 128) - or (IS_SM100 and d == 256 and dv == 256) + or (IS_SM100 and d == 256 and dv == 256 and softcap == 0.0) ) and learnable_sink is None # and False