From f0c7e20591763938501fc97ae04e77d68e20dd8a Mon Sep 17 00:00:00 2001 From: drisspg Date: Tue, 16 Jun 2026 17:39:47 +0000 Subject: [PATCH 1/2] Enable 2CTA for SM100 block-sparse backward stack-info: PR: https://github.com/Dao-AILab/flash-attention/pull/2661, branch: drisspg/stack/45 --- flash_attn/cute/block_sparse_utils.py | 707 ++++++++++++++---- flash_attn/cute/block_sparsity.py | 59 +- flash_attn/cute/flash_bwd_sm100.py | 265 +++++-- flash_attn/cute/flash_fwd_sm100.py | 7 + flash_attn/cute/interface.py | 58 +- flash_attn/cute/mask.py | 2 + flash_attn/cute/seqlen_info.py | 2 + .../cute/sm100_hd256_2cta_fmha_forward.py | 4 + tests/cute/test_mask_mod.py | 680 +++++++++++++++-- tests/cute/test_mask_mod_varlen.py | 72 ++ 10 files changed, 1564 insertions(+), 292 deletions(-) diff --git a/flash_attn/cute/block_sparse_utils.py b/flash_attn/cute/block_sparse_utils.py index d00ee34e27b..7ae8b959232 100644 --- a/flash_attn/cute/block_sparse_utils.py +++ b/flash_attn/cute/block_sparse_utils.py @@ -27,11 +27,22 @@ def _get_curr_blocksparse_tensors_varlen( m_block: cutlass.Int32, blocksparse_tensors: BlockSparseTensors, seqlen_info: SeqlenInfoQK, + kv_subtile_factor: cutlass.Constexpr[int] = 1, ) -> Tuple[cutlass.Int32, cute.Tensor, cutlass.Int32, Optional[cute.Tensor]]: """Varlen path: tensors are 2D [nheads, total_m_blocks] / [nheads, total_n_blocks].""" mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx, *_ = blocksparse_tensors curr_m_block = seqlen_info.m_block_offset + m_block - curr_block_idx_offset = seqlen_info.block_idx_offset + m_block * seqlen_info.num_n_blocks + sparse_num_n_blocks = ( + seqlen_info.num_n_blocks + if const_expr(kv_subtile_factor == 1) + else (seqlen_info.num_n_blocks + kv_subtile_factor - 1) // kv_subtile_factor + ) + batch_block_idx_offset = ( + seqlen_info.block_idx_offset + if const_expr(seqlen_info.has_cu_block_idx_offsets) + else seqlen_info.m_block_offset * sparse_num_n_blocks + ) + curr_block_idx_offset = batch_block_idx_offset + m_block * sparse_num_n_blocks curr_mask_block_cnt = mask_block_cnt[head_idx, curr_m_block] curr_mask_block_idx = cute.domain_offset(curr_block_idx_offset, mask_block_idx[head_idx, None]) if const_expr(full_block_cnt is not None): @@ -72,11 +83,12 @@ def get_curr_blocksparse_tensors( m_block: cutlass.Int32, blocksparse_tensors: BlockSparseTensors, seqlen_info: SeqlenInfoQK, + kv_subtile_factor: cutlass.Constexpr[int] = 1, ) -> Tuple[cutlass.Int32, cute.Tensor, cutlass.Int32, Optional[cute.Tensor]]: """Extract head, m_block, and batch-local blocksparsity data from blocksparse_tensors""" if const_expr(len(blocksparse_tensors.mask_block_cnt.shape) == 2): return _get_curr_blocksparse_tensors_varlen( - head_idx, m_block, blocksparse_tensors, seqlen_info + head_idx, m_block, blocksparse_tensors, seqlen_info, kv_subtile_factor ) return _get_curr_blocksparse_tensors(batch_idx, head_idx, m_block, blocksparse_tensors) @@ -123,13 +135,30 @@ def get_curr_blocksparse_tensors( # to ack/advance, and arrives `mbar_P_full_O_rescaled` when MMA can proceed. # # Backward (SM100): -# - Empty KV tile: for a given `n_block`, `total_m_block_cnt == 0` means no Q tiles contribute. -# - Both the load and compute loops guard all pipeline work on `process_tile`, so empty tiles -# skip producer/consumer operations entirely (no per-tile mbarrier phase handshake like forward). +# - Empty KV tile: for a given sparse `n_block`, `loop_count == 0` means no Q tiles contribute. +# - Load helpers guard their prologue/mainloop/tail on `loop_count > 0`; MMA/softmax/relay use +# the same block-sparse count as `process_tile`. Empty tiles therefore skip both producer and +# consumer operations entirely, with no synthetic per-tile mbarrier handshake like forward. +# - For non-empty tiles, producer tails are part of the non-empty load contract: the load warp emits +# exactly the Q/dO/LSE/dPsum/Qt/Kt items that the MMA and softmax consumers wait on, then tails the +# advanced producer states. Empty tiles leave producer and consumer states unchanged for the next tile. # - In the `not dKV_postprocess` path, dK/dV for empty KV tiles are explicitly written as zeros # even when `process_tile == False` (see `flash_bwd_sm100.py` `should_zero_dKV`). +@cute.jit +def sparse_physical_n_block_forward( + block_indices: cute.Tensor, + offset, + kv_subtile_factor: cutlass.Constexpr[int], +): + """Map forward physical N offsets to physical N tile indices.""" + sparse_offset = offset // kv_subtile_factor + subtile_offset = offset - sparse_offset * kv_subtile_factor + coarse_block = block_indices[sparse_offset] + return coarse_block * kv_subtile_factor + subtile_offset + + @cute.jit def load_block_list( block_indices: cute.Tensor, @@ -141,6 +170,7 @@ def load_block_list( pipeline_k, pipeline_v, intra_wg_overlap: cutlass.Constexpr, + kv_subtile_factor: cutlass.Constexpr[int] = 1, ): """Iterate over the sparse blocks and load K, V into the pipeline. For the intra_wg_overlap case, we overlap the loads of K and V. And this @@ -158,23 +188,40 @@ def load_block_list( """ if block_count > 0: + total_blocks = block_count * kv_subtile_factor if const_expr(not intra_wg_overlap): - for offset in cutlass.range(block_count): - n_block = block_indices[block_count - 1 - offset] + for offset in cutlass.range(total_blocks): + n_block = sparse_physical_n_block_forward( + block_indices, + total_blocks - 1 - offset, + kv_subtile_factor, + ) pipeline_k.producer_acquire(kv_producer_state) load_K(src_idx=n_block, producer_state=kv_producer_state) pipeline_v.producer_acquire(kv_producer_state) load_V(src_idx=n_block, producer_state=kv_producer_state) kv_producer_state.advance() else: - n_block_first = block_indices[block_count - 1] + n_block_first = sparse_physical_n_block_forward( + block_indices, + total_blocks - 1, + kv_subtile_factor, + ) if const_expr(not first_block_preloaded): pipeline_k.producer_acquire(kv_producer_state) load_K(src_idx=n_block_first, producer_state=kv_producer_state) - for idx in cutlass.range(block_count - 1, unroll=1): - n_block_prev = block_indices[block_count - 1 - idx] - n_block = block_indices[block_count - 2 - idx] + for idx in cutlass.range(total_blocks - 1, unroll=1): + n_block_prev = sparse_physical_n_block_forward( + block_indices, + total_blocks - 1 - idx, + kv_subtile_factor, + ) + n_block = sparse_physical_n_block_forward( + block_indices, + total_blocks - 1 - (idx + 1), + kv_subtile_factor, + ) kv_producer_state_prev = kv_producer_state.clone() kv_producer_state.advance() pipeline_k.producer_acquire(kv_producer_state) @@ -192,10 +239,11 @@ def finish_overlap_v_load( load_V, pipeline_v, kv_producer_state, + kv_subtile_factor: cutlass.Constexpr[int] = 1, ): """Load the final V block after overlapped K/V loads.""" if block_count > 0: - n_block_last = block_indices[0] + n_block_last = block_indices[0] * kv_subtile_factor pipeline_v.producer_acquire(kv_producer_state) load_V(src_idx=n_block_last, producer_state=kv_producer_state) kv_producer_state.advance() @@ -233,6 +281,7 @@ def produce_block_sparse_loads( intra_wg_overlap: cutlass.Constexpr, qhead_per_kvhead: cutlass.Constexpr[int] = 1, q_subtile_factor: cutlass.Constexpr[int] = 1, + kv_subtile_factor: cutlass.Constexpr[int] = 1, ): """Iterate over the mask and full block lists for a single tile. @@ -263,6 +312,7 @@ def produce_block_sparse_loads( m_block_sparse, blocksparse_tensors, seqlen_info, + kv_subtile_factor, ) mask_empty = curr_mask_block_cnt == 0 @@ -280,6 +330,7 @@ def produce_block_sparse_loads( pipeline_k=pipeline_k, pipeline_v=pipeline_v, intra_wg_overlap=intra_wg_overlap, + kv_subtile_factor=kv_subtile_factor, ) if const_expr(intra_wg_overlap) and curr_full_block_cnt > 0: @@ -289,6 +340,7 @@ def produce_block_sparse_loads( load_V, pipeline_v, kv_producer_state, + kv_subtile_factor=kv_subtile_factor, ) else: # Masked blocks present. When overlap is disabled this fully drains the list. @@ -302,6 +354,7 @@ def produce_block_sparse_loads( pipeline_k=pipeline_k, pipeline_v=pipeline_v, intra_wg_overlap=intra_wg_overlap, + kv_subtile_factor=kv_subtile_factor, ) if full_empty: @@ -312,13 +365,18 @@ def produce_block_sparse_loads( load_V, pipeline_v, kv_producer_state, + kv_subtile_factor=kv_subtile_factor, ) else: if const_expr(intra_wg_overlap): # Bridge the masked list to the full list by overlapping the pending masked V # with the first full K load. - n_block_mask_last = curr_mask_block_idx[0] - n_block_full_first = curr_full_block_idx[curr_full_block_cnt - 1] + n_block_mask_last = curr_mask_block_idx[0] * kv_subtile_factor + n_block_full_first = sparse_physical_n_block_forward( + curr_full_block_idx, + curr_full_block_cnt * kv_subtile_factor - 1, + kv_subtile_factor, + ) kv_producer_state_prev = kv_producer_state.clone() kv_producer_state.advance() pipeline_k.producer_acquire(kv_producer_state) @@ -336,6 +394,7 @@ def produce_block_sparse_loads( pipeline_k=pipeline_k, pipeline_v=pipeline_v, intra_wg_overlap=intra_wg_overlap, + kv_subtile_factor=kv_subtile_factor, ) kv_producer_state = finish_overlap_v_load( @@ -344,6 +403,7 @@ def produce_block_sparse_loads( load_V, pipeline_v, kv_producer_state, + kv_subtile_factor=kv_subtile_factor, ) else: # Non-overlap path with both lists: run the full list normally. @@ -357,6 +417,7 @@ def produce_block_sparse_loads( pipeline_k=pipeline_k, pipeline_v=pipeline_v, intra_wg_overlap=intra_wg_overlap, + kv_subtile_factor=kv_subtile_factor, ) return kv_producer_state @@ -572,12 +633,17 @@ def load_block_list_sm100( load_K, load_V, pipeline_kv, + kv_subtile_factor: cutlass.Constexpr[int] = 1, ): - """SM100 version of load_block_list (no intra_wg_overlap, no extra_tx_count).""" - block_count = block_end - block_begin - if block_count > 0: + """SM100 sparse load loop over physical N tiles from coarse metadata.""" + physical_block_count = block_end - block_begin + if physical_block_count > 0: # First iteration: load Q alongside K if requested - n_block_first = block_indices[block_end - 1] + n_block_first = sparse_physical_n_block_forward( + block_indices, + block_end - 1, + kv_subtile_factor, + ) if const_expr(load_q_with_first): # SM100 loads Q0 and optionally Q1 @@ -593,8 +659,12 @@ def load_block_list_sm100( kv_producer_state.advance() # Remaining blocks - for offset in cutlass.range(1, block_count): - n_block = block_indices[block_end - 1 - offset] + for offset in cutlass.range(1, physical_block_count): + n_block = sparse_physical_n_block_forward( + block_indices, + block_end - 1 - offset, + kv_subtile_factor, + ) load_K(block=n_block, producer_state=kv_producer_state, page_idx=None) kv_producer_state.advance() load_V(block=n_block, producer_state=kv_producer_state, page_idx=None) @@ -622,6 +692,7 @@ def produce_block_sparse_loads_sm100( q_producer_phase: Int32, qhead_per_kvhead: cutlass.Constexpr, q_subtile_factor: cutlass.Constexpr, + kv_subtile_factor: cutlass.Constexpr[int] = 1, ): """SM100 entry point for sparse block iteration. @@ -645,10 +716,19 @@ def produce_block_sparse_loads_sm100( m_block_sparse, blocksparse_tensors, seqlen_info, + kv_subtile_factor, ) - mask_begin, mask_end = split_block_range(curr_mask_block_cnt, split_idx, num_splits) - full_begin, full_end = split_block_range(curr_full_block_cnt, split_idx, num_splits) + mask_begin, mask_end = split_block_range( + curr_mask_block_cnt * kv_subtile_factor, + split_idx, + num_splits, + ) + full_begin, full_end = split_block_range( + curr_full_block_cnt * kv_subtile_factor, + split_idx, + num_splits, + ) mask_empty = mask_begin == mask_end full_empty = full_begin == full_end @@ -667,6 +747,7 @@ def produce_block_sparse_loads_sm100( load_K=load_K, load_V=load_V, pipeline_kv=pipeline_kv, + kv_subtile_factor=kv_subtile_factor, ) q_phase_flipped = not full_empty else: @@ -682,6 +763,7 @@ def produce_block_sparse_loads_sm100( load_K=load_K, load_V=load_V, pipeline_kv=pipeline_kv, + kv_subtile_factor=kv_subtile_factor, ) q_phase_flipped = True @@ -698,6 +780,7 @@ def produce_block_sparse_loads_sm100( load_K=load_K, load_V=load_V, pipeline_kv=pipeline_kv, + kv_subtile_factor=kv_subtile_factor, ) if q_phase_flipped: @@ -717,6 +800,7 @@ def get_total_block_count( qhead_per_kvhead: cutlass.Constexpr, q_subtile_factor: cutlass.Constexpr, seqlen_info: SeqlenInfoQK, + kv_subtile_factor: cutlass.Constexpr[int] = 1, ): m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor) ( @@ -730,10 +814,19 @@ def get_total_block_count( m_block_sparse, blocksparse_tensors, seqlen_info, + kv_subtile_factor, ) - mask_begin, mask_end = split_block_range(curr_mask_block_cnt, split_idx, num_splits) - full_begin, full_end = split_block_range(curr_full_block_cnt, split_idx, num_splits) + mask_begin, mask_end = split_block_range( + curr_mask_block_cnt * kv_subtile_factor, + split_idx, + num_splits, + ) + full_begin, full_end = split_block_range( + curr_full_block_cnt * kv_subtile_factor, + split_idx, + num_splits, + ) return mask_end - mask_begin + full_end - full_begin @@ -855,6 +948,112 @@ def handle_block_sparse_empty_tile_correction_sm100( ) +@cute.jit +def softmax_block_sparse_sm100_list( + block_indices: cute.Tensor, + block_end, + split_block_cnt, + kv_subtile_factor: cutlass.Constexpr[int], + softmax_step: Callable, + mask_fn_base: Callable, + mma_si_consumer_phase: Int32, + si_corr_producer_phase: Int32, + s0_s1_sequence_phase: Int32, + check_m_boundary: bool, + is_first_block: bool, +): + """Run one reverse sparse list while masking only the first coarse KV fragment.""" + n_block = sparse_physical_n_block_forward( + block_indices, + block_end - 1, + kv_subtile_factor, + ) + if is_first_block: + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + n_block, + is_first=True, + mask_fn=partial( + mask_fn_base, + mask_seqlen=True, + check_q_boundary=check_m_boundary, + ), + ) + else: + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + n_block, + is_first=False, + mask_fn=partial( + mask_fn_base, + mask_seqlen=True, + check_q_boundary=check_m_boundary, + ), + ) + + first_fragment_count = cutlass.min( + split_block_cnt, + (block_end - 1) % kv_subtile_factor + 1, + ) + for j in cutlass.range_constexpr(1, kv_subtile_factor): + if j < first_fragment_count: + n_block = sparse_physical_n_block_forward( + block_indices, + block_end - 1 - j, + kv_subtile_factor, + ) + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + n_block, + mask_fn=partial( + mask_fn_base, + mask_seqlen=True, + check_q_boundary=check_m_boundary, + ), + ) + for i in cutlass.range(first_fragment_count, split_block_cnt): + n_block = sparse_physical_n_block_forward( + block_indices, + block_end - 1 - i, + kv_subtile_factor, + ) + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + n_block, + mask_fn=partial( + mask_fn_base, + mask_seqlen=False, + check_q_boundary=check_m_boundary, + ), + ) + + return mma_si_consumer_phase, si_corr_producer_phase, s0_s1_sequence_phase + + @cute.jit def softmax_block_sparse_sm100( blocksparse_tensors: BlockSparseTensors, @@ -877,6 +1076,7 @@ def softmax_block_sparse_sm100( check_m_boundary: bool, qhead_per_kvhead: cutlass.Constexpr, q_subtile_factor: cutlass.Constexpr[int] = 1, + kv_subtile_factor: cutlass.Constexpr[int] = 1, ): warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor) @@ -892,10 +1092,19 @@ def softmax_block_sparse_sm100( m_block_sparse, blocksparse_tensors, seqlen_info, + kv_subtile_factor, ) - mask_begin, mask_end = split_block_range(curr_mask_block_cnt, split_idx, num_splits) - full_begin, full_end = split_block_range(curr_full_block_cnt, split_idx, num_splits) + mask_begin, mask_end = split_block_range( + curr_mask_block_cnt * kv_subtile_factor, + split_idx, + num_splits, + ) + full_begin, full_end = split_block_range( + curr_full_block_cnt * kv_subtile_factor, + split_idx, + num_splits, + ) split_mask_block_cnt = mask_end - mask_begin split_full_block_cnt = full_end - full_begin total_block_cnt = split_mask_block_cnt + split_full_block_cnt @@ -904,80 +1113,42 @@ def softmax_block_sparse_sm100( sm_stats_barrier.arrive_w_index(index=stage_idx * 4 + warp_idx) else: if split_mask_block_cnt > 0: - mask_n_block = curr_mask_block_idx[mask_end - 1] ( mma_si_consumer_phase, si_corr_producer_phase, s0_s1_sequence_phase, - ) = softmax_step( + ) = softmax_block_sparse_sm100_list( + curr_mask_block_idx, + mask_end, + split_mask_block_cnt, + kv_subtile_factor, + softmax_step, + mask_fn, mma_si_consumer_phase, si_corr_producer_phase, s0_s1_sequence_phase, - mask_n_block, - is_first=True, - mask_fn=partial(mask_fn, mask_seqlen=True, check_q_boundary=check_m_boundary), + check_m_boundary, + True, ) - for i in cutlass.range(1, split_mask_block_cnt): - mask_n_block = curr_mask_block_idx[mask_end - 1 - i] - ( - mma_si_consumer_phase, - si_corr_producer_phase, - s0_s1_sequence_phase, - ) = softmax_step( - mma_si_consumer_phase, - si_corr_producer_phase, - s0_s1_sequence_phase, - mask_n_block, - mask_fn=partial(mask_fn, mask_seqlen=False, check_q_boundary=check_m_boundary), - ) if split_full_block_cnt > 0: - full_n_block = curr_full_block_idx[full_end - 1] - if split_mask_block_cnt == 0: - ( - mma_si_consumer_phase, - si_corr_producer_phase, - s0_s1_sequence_phase, - ) = softmax_step( - mma_si_consumer_phase, - si_corr_producer_phase, - s0_s1_sequence_phase, - full_n_block, - is_first=True, - mask_fn=partial( - mask_fn_none, mask_seqlen=True, check_q_boundary=check_m_boundary - ), - ) - else: - ( - mma_si_consumer_phase, - si_corr_producer_phase, - s0_s1_sequence_phase, - ) = softmax_step( - mma_si_consumer_phase, - si_corr_producer_phase, - s0_s1_sequence_phase, - full_n_block, - is_first=False, - mask_fn=partial( - mask_fn_none, mask_seqlen=True, check_q_boundary=check_m_boundary - ), - ) - for i in cutlass.range(1, split_full_block_cnt): - full_n_block = curr_full_block_idx[full_end - 1 - i] - ( - mma_si_consumer_phase, - si_corr_producer_phase, - s0_s1_sequence_phase, - ) = softmax_step( - mma_si_consumer_phase, - si_corr_producer_phase, - s0_s1_sequence_phase, - full_n_block, - mask_fn=partial( - mask_fn_none, mask_seqlen=False, check_q_boundary=check_m_boundary - ), - ) + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_block_sparse_sm100_list( + curr_full_block_idx, + full_end, + split_full_block_cnt, + kv_subtile_factor, + softmax_step, + mask_fn_none, + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + check_m_boundary, + split_mask_block_cnt == 0, + ) return ( mma_si_consumer_phase, @@ -1019,12 +1190,12 @@ def get_total_q_block_count_bwd( @cute.jit -def produce_block_sparse_q_loads_bwd_sm100( +def produce_block_sparse_q_loads_bwd_sm100_default( blocksparse_tensors: BlockSparseTensors, batch_idx, head_idx, n_block, - # Pipeline states (will be returned after advancing) + # Pipeline states returned after advancing producer_state_Q_LSE, producer_state_dO_dPsum, # Pipelines @@ -1052,12 +1223,18 @@ def produce_block_sparse_q_loads_bwd_sm100( # Subtiling factor and bounds q_subtile_factor: cutlass.Constexpr = 1, m_block_max: int = 0, + # Optional 2CTA state for hdim <= 128 + use_2cta_instrs: cutlass.Constexpr = False, + producer_state_Qt=None, + producer_state_Kt=None, + pipeline_Qt=None, + pipeline_Kt=None, + load_Qt=None, + load_Kt=None, + load_dOt=None, + tma_copy_bytes_dO=0, ): - """SM100 backward block sparse loading with subtiling. - - Returns updated (producer_state_Q_LSE, producer_state_dO_dPsum). - First iteration loads K/V alongside Q/dO; subsequent iterations load only Q/dO. - """ + """Produce SM100 backward block-sparse Q/dO loads for 1CTA and non-hdim192 2CTA.""" ( curr_q_cnt, curr_q_idx, @@ -1067,79 +1244,323 @@ def produce_block_sparse_q_loads_bwd_sm100( ) = get_block_sparse_iteration_info_bwd( blocksparse_tensors, batch_idx, head_idx, n_block, q_subtile_factor, m_block_max ) - - for iter_idx in cutlass.range(loop_count, unroll=1): - m_block, _ = get_m_block_from_iter_bwd( - iter_idx, + # 2 cta peels the first loop for Qt path; so we guard the whole block if loopcount == 0 + if loop_count > Int32(0): + first_m_block, _ = get_m_block_from_iter_bwd( + Int32(0), curr_q_cnt, curr_q_idx, curr_full_cnt, curr_full_idx, - q_subtile_factor, - m_block_max, + q_subtile_factor=q_subtile_factor, + m_block_max=m_block_max, ) - m_block_safe = m_block + # with q_subtile > 1 we need to guard against fully OOB regions if m_block_max > 0: - m_block_safe = cutlass.min(m_block, m_block_max - 1) - - if iter_idx == 0: - # First block: load K/V alongside Q/dO - if const_expr(should_load_Q): - pipeline_Q.producer_acquire(producer_state_Q_LSE, extra_tx_count=tma_copy_bytes_K) - load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q_LSE)) - load_Q(m_block_safe, producer_state=producer_state_Q_LSE) - pipeline_Q.producer_commit(producer_state_Q_LSE) - pipeline_LSE.producer_acquire(producer_state_Q_LSE) - with cute.arch.elect_one(): - copy_stats( - gLSE[None, m_block_safe], - sLSE[None, producer_state_Q_LSE.index], - mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_Q_LSE), - ) - producer_state_Q_LSE.advance() - if const_expr(should_load_dO): - pipeline_dO.producer_acquire( - producer_state_dO_dPsum, extra_tx_count=tma_copy_bytes_V + first_m_block = cutlass.min(first_m_block, m_block_max - 1) + + if const_expr(should_load_Q): + pipeline_Q.producer_acquire(producer_state_Q_LSE, extra_tx_count=tma_copy_bytes_K) + load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q_LSE)) + load_Q(first_m_block, producer_state=producer_state_Q_LSE) + pipeline_Q.producer_commit(producer_state_Q_LSE) + + pipeline_LSE.producer_acquire(producer_state_Q_LSE) + with cute.arch.elect_one(): + copy_stats( + gLSE[None, first_m_block], + sLSE[None, producer_state_Q_LSE.index], + mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_Q_LSE), ) - load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_dO_dPsum)) - load_dO(m_block_safe, producer_state=producer_state_dO_dPsum) - pipeline_dO.producer_commit(producer_state_dO_dPsum) - pipeline_dPsum.producer_acquire(producer_state_dO_dPsum) - with cute.arch.elect_one(): - copy_stats( - gdPsum[None, m_block_safe], - sdPsum[None, producer_state_dO_dPsum.index], - mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dO_dPsum), - ) - producer_state_dO_dPsum.advance() - else: - # Subsequent blocks: just load Q/dO (K/V already loaded) + producer_state_Q_LSE.advance() + + if const_expr(should_load_dO): + pipeline_dO.producer_acquire( + producer_state_dO_dPsum, + extra_tx_count=tma_copy_bytes_V + tma_copy_bytes_dO + if const_expr(load_dOt is not None) + else tma_copy_bytes_V, + ) + load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_dO_dPsum)) + load_dO(first_m_block, producer_state=producer_state_dO_dPsum) + if const_expr(load_dOt is not None): + load_dOt(first_m_block, producer_state=producer_state_dO_dPsum) + pipeline_dO.producer_commit(producer_state_dO_dPsum) + + pipeline_dPsum.producer_acquire(producer_state_dO_dPsum) + with cute.arch.elect_one(): + copy_stats( + gdPsum[None, first_m_block], + sdPsum[None, producer_state_dO_dPsum.index], + mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dO_dPsum), + ) + producer_state_dO_dPsum.advance() + + if const_expr(use_2cta_instrs): + assert load_Kt is not None and pipeline_Kt is not None + assert producer_state_Kt is not None + pipeline_Kt.producer_acquire(producer_state_Kt) + load_Kt(tma_bar_ptr=pipeline_Kt.producer_get_barrier(producer_state_Kt)) + pipeline_Kt.producer_commit(producer_state_Kt) + producer_state_Kt.advance() + + prev_m_block = first_m_block + for iter_idx in cutlass.range(Int32(1), loop_count, unroll=1): + m_block, _ = get_m_block_from_iter_bwd( + iter_idx, + curr_q_cnt, + curr_q_idx, + curr_full_cnt, + curr_full_idx, + q_subtile_factor=q_subtile_factor, + m_block_max=m_block_max, + ) + if m_block_max > 0: + m_block = cutlass.min(m_block, m_block_max - 1) if const_expr(should_load_Q): + if const_expr(load_Qt is not None): + assert pipeline_Qt is not None and producer_state_Qt is not None + pipeline_Qt.producer_acquire(producer_state_Qt) + load_Qt(prev_m_block, producer_state=producer_state_Qt) + pipeline_Qt.producer_commit(producer_state_Qt) + producer_state_Qt.advance() + pipeline_Q.producer_acquire(producer_state_Q_LSE) - load_Q(m_block_safe, producer_state=producer_state_Q_LSE) + load_Q(m_block, producer_state=producer_state_Q_LSE) pipeline_Q.producer_commit(producer_state_Q_LSE) + pipeline_LSE.producer_acquire(producer_state_Q_LSE) with cute.arch.elect_one(): copy_stats( - gLSE[None, m_block_safe], + gLSE[None, m_block], sLSE[None, producer_state_Q_LSE.index], mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_Q_LSE), ) producer_state_Q_LSE.advance() + if const_expr(should_load_dO): - pipeline_dO.producer_acquire(producer_state_dO_dPsum) - load_dO(m_block_safe, producer_state=producer_state_dO_dPsum) + pipeline_dO.producer_acquire( + producer_state_dO_dPsum, + extra_tx_count=tma_copy_bytes_dO if const_expr(load_dOt is not None) else 0, + ) + load_dO(m_block, producer_state=producer_state_dO_dPsum) + if const_expr(load_dOt is not None): + load_dOt(m_block, producer_state=producer_state_dO_dPsum) pipeline_dO.producer_commit(producer_state_dO_dPsum) + pipeline_dPsum.producer_acquire(producer_state_dO_dPsum) with cute.arch.elect_one(): copy_stats( - gdPsum[None, m_block_safe], + gdPsum[None, m_block], sdPsum[None, producer_state_dO_dPsum.index], mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dO_dPsum), ) producer_state_dO_dPsum.advance() + prev_m_block = m_block + + if const_expr(should_load_Q): + if const_expr(load_Qt is not None): + assert pipeline_Qt is not None and producer_state_Qt is not None + pipeline_Qt.producer_acquire(producer_state_Qt) + load_Qt(prev_m_block, producer_state=producer_state_Qt) + pipeline_Qt.producer_commit(producer_state_Qt) + producer_state_Qt.advance() + + pipeline_Q.producer_tail(producer_state_Q_LSE.clone()) + pipeline_LSE.producer_tail(producer_state_Q_LSE) + if const_expr(load_Qt is not None): + pipeline_Qt.producer_tail(producer_state_Qt) + if const_expr(should_load_dO): + pipeline_dO.producer_tail(producer_state_dO_dPsum.clone()) + pipeline_dPsum.producer_tail(producer_state_dO_dPsum) + + return producer_state_Q_LSE, producer_state_dO_dPsum, producer_state_Qt, producer_state_Kt + + +@cute.jit +def produce_block_sparse_q_loads_bwd_sm100_2cta_hdim192( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + n_block, + # Pipeline states returned after advancing + producer_state_Q_Qt, + producer_state_O_Ot, + producer_state_LSE, + producer_state_dPsum, + # Pipelines + pipeline_Q, + pipeline_LSE, + pipeline_dO, + pipeline_dPsum, + pipeline_Qt, + # Load functions + load_K, + load_V, + load_Q, + load_dO, + load_Qt, + load_Kt, + load_dOt, + copy_stats, + # Global tensors for LSE/dPsum + gLSE, + sLSE, + gdPsum, + sdPsum, + # TMA copy bytes for extra_tx_count + tma_copy_bytes_K, + tma_copy_bytes_V, + # Subtiling factor and bounds + q_subtile_factor: cutlass.Constexpr = 1, + m_block_max: int = 0, +): + """Produce SM100 backward block-sparse Q/dO loads for the hdim192 2CTA schedule.""" + ( + curr_q_cnt, + curr_q_idx, + curr_full_cnt, + curr_full_idx, + loop_count, + ) = get_block_sparse_iteration_info_bwd( + blocksparse_tensors, batch_idx, head_idx, n_block, q_subtile_factor, m_block_max + ) + + if loop_count > Int32(0): + first_m_block, _ = get_m_block_from_iter_bwd( + Int32(0), + curr_q_cnt, + curr_q_idx, + curr_full_cnt, + curr_full_idx, + q_subtile_factor=q_subtile_factor, + m_block_max=m_block_max, + ) - return producer_state_Q_LSE, producer_state_dO_dPsum + # with q_subtile > 1 we need to guard against fully OOB regions + if m_block_max > 0: + first_m_block = cutlass.min(first_m_block, m_block_max - 1) + + # K & Q (for S) + pipeline_Q.producer_acquire( + producer_state_Q_Qt, + extra_tx_count=tma_copy_bytes_K, + ) + load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q_Qt)) + load_Q(first_m_block, producer_state=producer_state_Q_Qt) + pipeline_Q.producer_commit(producer_state_Q_Qt) + producer_state_Q_Qt.advance() + + # LSE + pipeline_LSE.producer_acquire(producer_state_LSE) + with cute.arch.elect_one(): + copy_stats( + gLSE[None, first_m_block], + sLSE[None, producer_state_LSE.index], + mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_LSE), + ) + producer_state_LSE.advance() + + # dOt + V, for dP.T = V @ dO.T + pipeline_dO.producer_acquire( + producer_state_O_Ot, + extra_tx_count=tma_copy_bytes_V, + ) + load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_O_Ot)) + load_dOt(first_m_block, producer_state=producer_state_O_Ot) + pipeline_dO.producer_commit(producer_state_O_Ot) + producer_state_O_Ot.advance() + + # dPsum + pipeline_dPsum.producer_acquire(producer_state_dPsum) + with cute.arch.elect_one(): + copy_stats( + gdPsum[None, first_m_block], + sdPsum[None, producer_state_dPsum.index], + mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dPsum), + ) + producer_state_dPsum.advance() + + # Qt, for dK = dS.T @ Q + pipeline_Qt.producer_acquire( + producer_state_Q_Qt, + extra_tx_count=tma_copy_bytes_K, + ) + load_Qt(first_m_block, producer_state=producer_state_Q_Qt) + load_Kt(tma_bar_ptr=pipeline_Qt.producer_get_barrier(producer_state_Q_Qt)) + pipeline_Qt.producer_commit(producer_state_Q_Qt) + producer_state_Q_Qt.advance() + + # dO, for dV = P.T @ dO + pipeline_dO.producer_acquire(producer_state_O_Ot) + load_dO(first_m_block, producer_state=producer_state_O_Ot) + pipeline_dO.producer_commit(producer_state_O_Ot) + producer_state_O_Ot.advance() + + # 2CTA: [lse | Q | dOt | dPsum | Qt | dO] + for iter_idx in cutlass.range(Int32(1), loop_count, unroll=1): + m_block, _ = get_m_block_from_iter_bwd( + iter_idx, + curr_q_cnt, + curr_q_idx, + curr_full_cnt, + curr_full_idx, + q_subtile_factor=q_subtile_factor, + m_block_max=m_block_max, + ) + if m_block_max > 0: + m_block = cutlass.min(m_block, m_block_max - 1) + + # LSE + pipeline_LSE.producer_acquire(producer_state_LSE) + with cute.arch.elect_one(): + copy_stats( + gLSE[None, m_block], + sLSE[None, producer_state_LSE.index], + mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_LSE), + ) + producer_state_LSE.advance() + + # Q + pipeline_Q.producer_acquire(producer_state_Q_Qt) + load_Q(m_block, producer_state=producer_state_Q_Qt) + pipeline_Q.producer_commit(producer_state_Q_Qt) + producer_state_Q_Qt.advance() + + # dPsum + pipeline_dPsum.producer_acquire(producer_state_dPsum) + with cute.arch.elect_one(): + copy_stats( + gdPsum[None, m_block], + sdPsum[None, producer_state_dPsum.index], + mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dPsum), + ) + producer_state_dPsum.advance() + + # dOt, for dP.T = V @ dO.T + pipeline_dO.producer_acquire(producer_state_O_Ot) + load_dOt(m_block, producer_state=producer_state_O_Ot) + pipeline_dO.producer_commit(producer_state_O_Ot) + producer_state_O_Ot.advance() + + # Qt, for dK = dS.T @ Q + pipeline_Qt.producer_acquire(producer_state_Q_Qt) + load_Qt(m_block, producer_state=producer_state_Q_Qt) + pipeline_Qt.producer_commit(producer_state_Q_Qt) + producer_state_Q_Qt.advance() + + # dO, for dV = P.T @ dO + pipeline_dO.producer_acquire(producer_state_O_Ot) + load_dO(m_block, producer_state=producer_state_O_Ot) + pipeline_dO.producer_commit(producer_state_O_Ot) + producer_state_O_Ot.advance() + + pipeline_Q.producer_tail(producer_state_Q_Qt) + pipeline_LSE.producer_tail(producer_state_LSE) + pipeline_dO.producer_tail(producer_state_O_Ot) + pipeline_dPsum.producer_tail(producer_state_dPsum) + + return producer_state_Q_Qt, producer_state_O_Ot, producer_state_LSE, producer_state_dPsum @cute.jit diff --git a/flash_attn/cute/block_sparsity.py b/flash_attn/cute/block_sparsity.py index 009886e835a..b0f12dc8581 100644 --- a/flash_attn/cute/block_sparsity.py +++ b/flash_attn/cute/block_sparsity.py @@ -195,6 +195,32 @@ def get_sparse_q_block_size( return min_block_size +def get_kv_subtile_factor( + block_sparse_tensors: BlockSparseTensorsTorch | None, + n_block_size: int, +) -> int: + """Return the number of physical KV tiles covered by one sparse KV block.""" + if block_sparse_tensors is None or block_sparse_tensors.block_size is None: + return 1 + sparse_block_size_kv = block_sparse_tensors.block_size[1] + if sparse_block_size_kv % n_block_size != 0: + raise ValueError( + "Block sparsity expects sparse_block_size[1] " + f"to be a multiple of tile_n={n_block_size}; got {sparse_block_size_kv}." + ) + return sparse_block_size_kv // n_block_size + + +def block_sparse_bwd_supports_2cta( + block_sparse_tensors: BlockSparseTensorsTorch | None, + n_block_size: int, +) -> bool: + """Return whether sparse KV metadata constrains backward away from 2CTA.""" + if block_sparse_tensors is None: + return True + return get_kv_subtile_factor(block_sparse_tensors, n_block_size) % 2 == 0 + + def _expand_sparsity_tensor( tensor: torch.Tensor, expected_shape: Tuple[int, ...], @@ -310,7 +336,7 @@ def infer_block_sparse_expected_shapes( Expectations: - mask_block_cnt is (B, H, M) and mask_block_idx is (B, H, M, N). - Batch/head dims may be 1 for broadcast, or match the requested sizes. - - sparse_block_size_kv must match tile_n. + - sparse_block_size_kv must be a multiple of tile_n. - sparse_block_size_q must be a multiple of q_stage * tile_m. - If sparse_block_size_q is omitted and seqlen_q/num_m_blocks is ambiguous, the caller must provide block_size to disambiguate. TODO will make this required in a future PR. @@ -319,8 +345,10 @@ def infer_block_sparse_expected_shapes( base_n_block = n_block_size if sparse_block_size_kv is None: sparse_block_size_kv = base_n_block - if sparse_block_size_kv != base_n_block: - raise ValueError(f"Block sparse tensors{context} require BLOCK_SIZE_KV={base_n_block}.") + if sparse_block_size_kv % base_n_block != 0: + raise ValueError( + f"Block sparse tensors{context} require BLOCK_SIZE_KV to be a multiple of {base_n_block}." + ) if tensors.mask_block_idx is None: raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.") num_m_blocks = tensors.mask_block_idx.shape[2] @@ -392,6 +420,7 @@ def get_block_sparse_expected_shapes_bwd( m_block_size: int, n_block_size: int, q_subtile_factor: int, + kv_subtile_factor: int = 1, ) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]: """Return (expected_count_shape, expected_index_shape) for backward block sparse normalization. @@ -400,8 +429,9 @@ def get_block_sparse_expected_shapes_bwd( by q_subtile_factor * m_block_size. """ sparse_block_size_q = q_subtile_factor * m_block_size + sparse_block_size_kv = kv_subtile_factor * n_block_size expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q) - expected_n_blocks = ceildiv(seqlen_k, n_block_size) + expected_n_blocks = ceildiv(seqlen_k, sparse_block_size_kv) expected_count_shape = (batch_size, num_head, expected_n_blocks) expected_index_shape = (batch_size, num_head, expected_n_blocks, expected_m_blocks) return expected_count_shape, expected_index_shape @@ -525,7 +555,8 @@ def normalize_block_sparse_config( seqlen_k: int, block_size: tuple[int, int], q_stage: int, -) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None, int]: + allow_kv_subtile: bool = False, +) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None, int, int]: """Validate the block-sparse config, infer expected shapes, and normalize. Handles both fixed-length (3D `[B, H, M]` / 4D `[B, H, M, N]`) and varlen @@ -538,7 +569,12 @@ def normalize_block_sparse_config( sparse_block_size_q, sparse_block_size_kv = None, n_block_size else: sparse_block_size_q, sparse_block_size_kv = tensors.block_size - if sparse_block_size_kv != n_block_size: + if sparse_block_size_kv % n_block_size != 0: + raise ValueError( + f"Block sparsity requires sparse_block_size[1] to be a multiple of tile_n={n_block_size}." + ) + kv_subtile_factor = sparse_block_size_kv // n_block_size + if kv_subtile_factor != 1 and not allow_kv_subtile: raise ValueError( f"Block sparsity requires sparse_block_size[1]={n_block_size} to match tile_n." ) @@ -579,6 +615,7 @@ def normalize_block_sparse_config( normalized_tensors, get_block_sparse_broadcast_pattern(normalized_tensors), q_subtile_factor, + kv_subtile_factor, ) @@ -591,6 +628,7 @@ def normalize_block_sparse_config_bwd( seqlen_k: int, block_size: tuple[int, int], q_subtile_factor: int, + kv_subtile_factor: int = 1, ) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None]: m_block_size, n_block_size = block_size if tensors.block_size is None: @@ -602,9 +640,11 @@ def normalize_block_sparse_config_bwd( f"Block sparsity expects sparse_block_size_q={q_subtile_factor * m_block_size} " f"for q_subtile_factor={q_subtile_factor}." ) - if sparse_block_size_kv != n_block_size: + expected_sparse_block_size_kv = kv_subtile_factor * n_block_size + if sparse_block_size_kv != expected_sparse_block_size_kv: raise ValueError( - f"Block sparsity expects sparse_block_size[1]={n_block_size} to match tile_n." + f"Block sparsity expects sparse_block_size[1]={expected_sparse_block_size_kv} " + f"for kv_subtile_factor={kv_subtile_factor}." ) expected_count_shape, expected_index_shape = get_block_sparse_expected_shapes_bwd( batch_size, @@ -614,6 +654,7 @@ def normalize_block_sparse_config_bwd( m_block_size, n_block_size, q_subtile_factor, + kv_subtile_factor, ) normalized_tensors = normalize_block_sparse_tensors( tensors, @@ -623,7 +664,7 @@ def normalize_block_sparse_config_bwd( hint=lambda: ( f"Backward expects Q-direction block-sparse tensors (q_mask_cnt/q_mask_idx, " f"and optionally full_q_cnt/full_q_idx). Regenerate the backward BlockMask with " - f"BLOCK_SIZE=({q_subtile_factor * m_block_size}, {n_block_size})." + f"BLOCK_SIZE=({q_subtile_factor * m_block_size}, {expected_sparse_block_size_kv})." ), ) return normalized_tensors, get_block_sparse_broadcast_pattern(normalized_tensors) diff --git a/flash_attn/cute/flash_bwd_sm100.py b/flash_attn/cute/flash_bwd_sm100.py index f0d39f0c6b8..799b3a1b414 100644 --- a/flash_attn/cute/flash_bwd_sm100.py +++ b/flash_attn/cute/flash_bwd_sm100.py @@ -41,7 +41,8 @@ get_total_q_block_count_bwd, get_block_sparse_iteration_info_bwd, get_m_block_from_iter_bwd, - produce_block_sparse_q_loads_bwd_sm100, + produce_block_sparse_q_loads_bwd_sm100_2cta_hdim192, + produce_block_sparse_q_loads_bwd_sm100_default, ) @@ -67,6 +68,7 @@ def __init__( mask_mod: cutlass.Constexpr | None = None, has_aux_tensors: cutlass.Constexpr = False, q_subtile_factor: cutlass.Constexpr[int] = 1, + kv_subtile_factor: cutlass.Constexpr[int] = 1, ): # padding head_dim to a multiple of 16 as k_block_size hdim_multiple_of = 16 @@ -120,6 +122,8 @@ def __init__( self.mask_mod = mask_mod self.has_aux_tensors = has_aux_tensors self.q_subtile_factor = q_subtile_factor + self.kv_subtile_factor = kv_subtile_factor + assert self.kv_subtile_factor == 1 or self.kv_subtile_factor % self.cta_group_size == 0 # For score_mod, use vec_size=1 (like forward) to handle per-element indices if cutlass.const_expr(has_aux_tensors): self.vec_size: cutlass.Constexpr = 1 @@ -921,12 +925,12 @@ class SharedStorage: seqlen_k_divmod = FastDivmodDivisor(seqlen_k) fastdiv_mods = (seqlen_q_divmod, seqlen_k_divmod) self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None) - - if const_expr(self.use_2cta_instrs): - assert blocksparse_tensors is None, ( - "2-CTA mode does not support block sparsity. " - "Please create kernel with use_2cta_instrs=False for block sparse attention." + if const_expr(self.use_block_sparsity and self.use_2cta_instrs): + # Both CTAs of a cluster must map to the same sparse KV column or they deadlock. + assert self.kv_subtile_factor % self.cta_group_size == 0, ( + "2-CTA block-sparse backward requires kv_subtile_factor % cta_group_size == 0" ) + # 2-CTA: 231424 and 1-CTA: 232448 # print("SMEM: ", self.shared_storage.size_in_bytes()) if const_expr(self.use_block_sparsity or aux_data.tensors is not None): @@ -1439,6 +1443,7 @@ def kernel( block_info, SeqlenInfoCls, TileSchedulerCls, + blocksparse_tensors, ) # LOAD @@ -1630,6 +1635,7 @@ def relay( block_info: BlockInfo, SeqlenInfoCls: Callable, TileSchedulerCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors] = None, ): cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) dS_cluster_phase = Int32(0) @@ -1647,9 +1653,20 @@ def relay( process_tile = ( const_expr(not self.is_local and not self.is_varlen_q) or m_block_min < m_block_max ) + num_iters = m_block_max - m_block_min + if const_expr(self.use_block_sparsity): + assert blocksparse_tensors is not None + num_iters = get_total_q_block_count_bwd( + blocksparse_tensors, + batch_idx, + head_idx, + n_block // self.kv_subtile_factor, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + process_tile = num_iters > Int32(0) if process_tile: - num_iters = m_block_max - m_block_min for _ in cutlass.range(num_iters, unroll=1): # Wait for dS_xchg from peer CTA cute.arch.mbarrier_wait(dS_cluster_full_mbar_ptr, phase=dS_cluster_phase) @@ -1755,6 +1772,7 @@ def load( ) head_idx_kv = head_idx // self.qhead_per_kvhead n_block_cta_group = n_block // self.cta_group_size + n_block_sparse = n_block // self.kv_subtile_factor # GMEM tensors (varlen-aware) mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] @@ -1834,6 +1852,7 @@ def load( single_stage=True, ) + load_dOt = None if const_expr(tma_atom_dOt is not None): gdOt = cute.local_tile( mdOt_cur, cute.select(self.mma_tiler_vdo, mode=[1, 2]), (None, 0) @@ -1863,6 +1882,7 @@ def load( load_dO = copy_utils.tma_producer_copy_fn(load_dO, pipeline_dO) # (4) dK += dS.T @ Q (2-CTA: needs separate Qt load) + load_Qt = None if const_expr(tma_atom_Qt is not None): gQt = cute.local_tile( mQt_cur, cute.select(self.mma_tiler_dsq, mode=[1, 2]), (0, None) @@ -1879,6 +1899,7 @@ def load( load_Qt = copy_utils.tma_producer_copy_fn(load_Qt, pipeline_Qt) # (5) dQ = dS @ K + load_Kt = None if const_expr(self.use_2cta_instrs): gKt = cute.local_tile( mKt_cur, cute.select(self.mma_tiler_dsk, mode=[1, 2]), (0, n_block_cta_group) @@ -1903,59 +1924,18 @@ def load( # gdPsum = cute.logical_divide(gdPsum, (64,))[(None, block_in_cluster_coord_vmnk[1]), None] # copy_stats = partial(cute.copy, copy_atom_stats, mcast_mask=q_do_mcast_mask) - # some tiles might be empty due to block sparsity - if const_expr(self.use_block_sparsity): - total_m_block_cnt = get_total_q_block_count_bwd( - blocksparse_tensors, - batch_idx, - head_idx, - n_block, - q_subtile_factor=self.q_subtile_factor, - m_block_max=m_block_max, - ) - process_tile = total_m_block_cnt > Int32(0) - else: + if const_expr(not self.use_block_sparsity): process_tile = ( const_expr(not self.is_local and not self.is_varlen_q) or m_block_min < m_block_max ) - if process_tile: - if const_expr(self.use_block_sparsity): - producer_state_Q_LSE, producer_state_dO_dPsum = ( - produce_block_sparse_q_loads_bwd_sm100( - blocksparse_tensors, - batch_idx, - head_idx, - n_block, - producer_state_Q_LSE, - producer_state_dO_dPsum, - pipeline_Q, - pipeline_LSE, - pipeline_dO, - pipeline_dPsum, - load_K, - load_V, - load_Q, - load_dO, - copy_stats, - gLSE, - sLSE, - gdPsum, - sdPsum, - self.tma_copy_bytes["K"], - self.tma_copy_bytes["V"], - should_load_Q=should_load_Q, - should_load_dO=should_load_dO, - q_subtile_factor=self.q_subtile_factor, - m_block_max=m_block_max, - ) - ) - else: + if process_tile: first_m_block = m_block_min if const_expr(self.use_2cta_instrs and self.tile_hdim == 192): #### Prologue #### assert should_load_Q and should_load_dO + assert load_dOt is not None and load_Qt is not None # K & Q (for S) pipeline_Q.producer_acquire( producer_state_Q_Qt, @@ -2059,6 +2039,10 @@ def load( pipeline_dO.producer_commit(producer_state_O_Ot) producer_state_O_Ot.advance() + pipeline_Q.producer_tail(producer_state_Q_Qt) + pipeline_LSE.producer_tail(producer_state_LSE) + pipeline_dO.producer_tail(producer_state_O_Ot) + pipeline_dPsum.producer_tail(producer_state_dPsum) else: #### Prologue #### if const_expr(should_load_Q): @@ -2088,7 +2072,7 @@ def load( pipeline_dO.producer_acquire( producer_state_dO_dPsum, extra_tx_count=self.tma_copy_bytes["V"] + self.tma_copy_bytes["dO"] - if const_expr(tma_atom_dOt is not None) + if const_expr(load_dOt is not None) else self.tma_copy_bytes["V"], ) load_V( @@ -2097,7 +2081,7 @@ def load( ) ) load_dO(first_m_block, producer_state=producer_state_dO_dPsum) - if const_expr(tma_atom_dOt is not None): + if const_expr(load_dOt is not None): load_dOt(first_m_block, producer_state=producer_state_dO_dPsum) pipeline_dO.producer_commit(producer_state_dO_dPsum) @@ -2121,7 +2105,7 @@ def load( #### Main Loop #### for m_block in cutlass.range(m_block_min + 1, m_block_max, unroll=1): if const_expr(should_load_Q): - if const_expr(tma_atom_Qt is not None): + if const_expr(load_Qt is not None): pipeline_Qt.producer_acquire(producer_state_Qt) load_Qt(m_block - 1, producer_state=producer_state_Qt) pipeline_Qt.producer_commit(producer_state_Qt) @@ -2148,11 +2132,11 @@ def load( pipeline_dO.producer_acquire( producer_state_dO_dPsum, extra_tx_count=self.tma_copy_bytes["dO"] - if const_expr(tma_atom_dOt is not None) + if const_expr(load_dOt is not None) else 0, ) load_dO(m_block, producer_state=producer_state_dO_dPsum) - if const_expr(tma_atom_dOt is not None): + if const_expr(load_dOt is not None): load_dOt(m_block, producer_state=producer_state_dO_dPsum) pipeline_dO.producer_commit(producer_state_dO_dPsum) @@ -2170,27 +2154,104 @@ def load( #### Tail #### if const_expr(should_load_Q): - if const_expr(tma_atom_Qt is not None): + if const_expr(load_Qt is not None): pipeline_Qt.producer_acquire(producer_state_Qt) load_Qt(m_block_max - 1, producer_state=producer_state_Qt) pipeline_Qt.producer_commit(producer_state_Qt) producer_state_Qt.advance() + pipeline_Q.producer_tail(producer_state_Q_LSE.clone()) + pipeline_LSE.producer_tail(producer_state_Q_LSE) + if const_expr(load_Qt is not None): + pipeline_Qt.producer_tail(producer_state_Qt) + if const_expr(should_load_dO): + pipeline_dO.producer_tail(producer_state_dO_dPsum.clone()) + pipeline_dPsum.producer_tail(producer_state_dO_dPsum) + + else: + assert blocksparse_tensors is not None if const_expr(self.use_2cta_instrs and self.tile_hdim == 192): - pipeline_Q.producer_tail(producer_state_Q_Qt) - pipeline_LSE.producer_tail(producer_state_LSE) - pipeline_dO.producer_tail(producer_state_O_Ot) - pipeline_dPsum.producer_tail(producer_state_dPsum) + assert should_load_Q and should_load_dO + assert load_dOt is not None and load_Qt is not None + assert load_Kt is not None and pipeline_Qt is not None + ( + producer_state_Q_Qt, + producer_state_O_Ot, + producer_state_LSE, + producer_state_dPsum, + ) = produce_block_sparse_q_loads_bwd_sm100_2cta_hdim192( + blocksparse_tensors, + batch_idx, + head_idx, + n_block_sparse, + producer_state_Q_Qt, + producer_state_O_Ot, + producer_state_LSE, + producer_state_dPsum, + pipeline_Q, + pipeline_LSE, + pipeline_dO, + pipeline_dPsum, + pipeline_Qt, + load_K, + load_V, + load_Q, + load_dO, + load_Qt, + load_Kt, + load_dOt, + copy_stats, + gLSE, + sLSE, + gdPsum, + sdPsum, + self.tma_copy_bytes["K"], + self.tma_copy_bytes["V"], + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) else: - if const_expr(should_load_Q): - pipeline_Q.producer_tail(producer_state_Q_LSE.clone()) - pipeline_LSE.producer_tail(producer_state_Q_LSE) - if const_expr(tma_atom_Qt is not None): - pipeline_Qt.producer_tail(producer_state_Qt) - if const_expr(should_load_dO): - pipeline_dO.producer_tail(producer_state_dO_dPsum.clone()) - pipeline_dPsum.producer_tail(producer_state_dO_dPsum) - + ( + producer_state_Q_LSE, + producer_state_dO_dPsum, + producer_state_Qt, + producer_state_Kt, + ) = produce_block_sparse_q_loads_bwd_sm100_default( + blocksparse_tensors, + batch_idx, + head_idx, + n_block_sparse, + producer_state_Q_LSE, + producer_state_dO_dPsum, + pipeline_Q, + pipeline_LSE, + pipeline_dO, + pipeline_dPsum, + load_K, + load_V, + load_Q, + load_dO, + copy_stats, + gLSE, + sLSE, + gdPsum, + sdPsum, + self.tma_copy_bytes["K"], + self.tma_copy_bytes["V"], + should_load_Q, + should_load_dO, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + use_2cta_instrs=self.use_2cta_instrs, + producer_state_Qt=producer_state_Qt, + producer_state_Kt=producer_state_Kt, + pipeline_Qt=pipeline_Qt, + pipeline_Kt=pipeline_Kt, + load_Qt=load_Qt, + load_Kt=load_Kt, + load_dOt=load_dOt, + tma_copy_bytes_dO=self.tma_copy_bytes["dO"], + ) tile_scheduler.prefetch_next_work() tile_scheduler.advance_to_next_work() work_tile = tile_scheduler.get_current_work() @@ -2365,7 +2426,7 @@ def mma( blocksparse_tensors, batch_idx, head_idx, - n_block, + n_block // self.kv_subtile_factor, q_subtile_factor=self.q_subtile_factor, m_block_max=m_block_max, ) @@ -2391,7 +2452,7 @@ def mma( # 4. dV = P.T @ dO # 5. dQ = dS @ K - main_loop_iters = m_block_max - m_block_min + main_loop_iters = block_iter_count # empty waits # pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc) @@ -3018,7 +3079,7 @@ def compute_loop( blocksparse_tensors, batch_idx, head_idx, - n_block, + n_block // self.kv_subtile_factor, q_subtile_factor=self.q_subtile_factor, m_block_max=m_block_max, ) @@ -3452,6 +3513,32 @@ def _dq_semaphore_lock_value( else: assert curr_dq_write_order_full is not None lock_value = curr_dq_write_order_full[sparse_iter - curr_q_cnt] + if const_expr(self.kv_subtile_factor > self.cta_group_size): + groups_per_sparse_block = self.kv_subtile_factor // self.cta_group_size + local_group = n_block % groups_per_sparse_block + if const_expr(self.spt): + # [NOTE] KV_subtile determ + spt + # dq_write_order stores one rank per sparse block; each physical tile + # derives its slot as rank * groups_per_sparse_block + (n_block % groups_per_sparse_block). + # W/ kv_subtile the tail sparse column can have a physical tile that is + # never scheduled; e.g. kv_tile = 128, seqlen = 1023, KV_Block = 384 -> 3 + # sparse blocks. The last block is covered in [768, 896), [896, 1024) and + # then [1024, 1152) which no CTA ever runs. Since the highest tile gets + # the lowest lock value under spt, we would hang! + # In this case we locally reverse, [N+2, N+1, N*] where N* is not scheduled + # -> [N+1, N, N+2*]. Ahh but won't the vacant N+2 slot stall the next sparse + # block's CTAs? It will, so the writer holding N+1 bumps the increment by 2 + # instead of 1 (i.e. 1 + #unscheduled) :) + total_groups = cute.ceil_div( + seqlen.seqlen_k, self.tile_n * self.cta_group_size + ) + groups_in_own_block = cutlass.min( + groups_per_sparse_block, + total_groups + - (n_block // groups_per_sparse_block) * groups_per_sparse_block, + ) + local_group = groups_in_own_block - 1 - local_group + lock_value = lock_value * groups_per_sparse_block + local_group return lock_value @cute.jit @@ -3511,6 +3598,7 @@ def dQacc_reduce( while work_tile.is_valid_tile: n_block, head_idx, batch_idx, _ = work_tile.tile_idx n_block_cta_group = n_block // self.cta_group_size # for 2cta + n_block_sparse = n_block // self.kv_subtile_factor seqlen = SeqlenInfoCls(batch_idx) m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block_cta_group) if const_expr(not seqlen.has_cu_seqlens_q): @@ -3531,6 +3619,26 @@ def dQacc_reduce( delay_semaphore_release = not self.tile_hdim == 192 and not self.use_block_sparsity + dq_sem_release_inc = Int32(1) + if const_expr( + self.deterministic + and self.use_block_sparsity + and self.spt + and self.kv_subtile_factor > self.cta_group_size + ): + # A truncated tail block's last writer releases the missing increments, + # see: [NOTE] KV_subtile determ + spt + groups_per_sparse_block = self.kv_subtile_factor // self.cta_group_size + total_groups = cute.ceil_div(seqlen.seqlen_k, self.tile_n * self.cta_group_size) + tail_sparse_block_idx = (total_groups - 1) // groups_per_sparse_block + groups_in_tail = total_groups - tail_sparse_block_idx * groups_per_sparse_block + is_tail_bridge_group = ( + n_block_cta_group // groups_per_sparse_block == tail_sparse_block_idx + and n_block_cta_group % groups_per_sparse_block == 0 + ) + if is_tail_bridge_group: + dq_sem_release_inc = Int32(1) + groups_per_sparse_block - groups_in_tail + curr_q_cnt = Int32(0) curr_q_idx = None curr_full_cnt = Int32(0) @@ -3553,7 +3661,7 @@ def dQacc_reduce( blocksparse_tensors, batch_idx, head_idx, - n_block, + n_block_sparse, q_subtile_factor=self.q_subtile_factor, m_block_max=m_block_max, ) @@ -3563,12 +3671,12 @@ def dQacc_reduce( if const_expr(blocksparse_tensors.dq_write_order is not None): assert blocksparse_tensors.dq_write_order is not None curr_dq_write_order = blocksparse_tensors.dq_write_order[ - batch_idx, head_idx, n_block, None + batch_idx, head_idx, n_block_sparse, None ] if const_expr(blocksparse_tensors.dq_write_order_full is not None): assert blocksparse_tensors.dq_write_order_full is not None curr_dq_write_order_full = blocksparse_tensors.dq_write_order_full[ - batch_idx, head_idx, n_block, None + batch_idx, head_idx, n_block_sparse, None ] # dQacc_reduce mainloop @@ -3678,7 +3786,10 @@ def dQacc_reduce( self.reduce_sync_barrier.arrive_and_wait() if not m_block_oob_upper: barrier.arrive_inc( - mdQ_semaphore_cur[m_block, None].iterator, tidx, cta_rank_in_cluster, 1 + mdQ_semaphore_cur[m_block, None].iterator, + tidx, + cta_rank_in_cluster, + dq_sem_release_inc, ) if process_tile: diff --git a/flash_attn/cute/flash_fwd_sm100.py b/flash_attn/cute/flash_fwd_sm100.py index 2e0a91b39aa..24a38ade384 100644 --- a/flash_attn/cute/flash_fwd_sm100.py +++ b/flash_attn/cute/flash_fwd_sm100.py @@ -129,6 +129,7 @@ def __init__( is_split_kv: bool = False, pack_gqa: bool = False, q_subtile_factor: int = 1, + kv_subtile_factor: int = 1, m_block_size: int = 128, n_block_size: int = 128, q_stage: cutlass.Constexpr[int] = 2, @@ -192,6 +193,7 @@ def __init__( ) self.use_correction_warps_for_epi = not self.use_tma_O self.q_subtile_factor = q_subtile_factor + self.kv_subtile_factor = kv_subtile_factor assert not (self.is_split_kv and self.head_dim_v_padded >= 192), ( "SplitKV is not supported for hdim >= 192" ) @@ -1529,6 +1531,7 @@ def load( q_producer_phase, self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, self.q_subtile_factor, + self.kv_subtile_factor, ) @@ -1672,6 +1675,7 @@ def mma( self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, self.q_subtile_factor, seqlen_info=seqlen, + kv_subtile_factor=self.kv_subtile_factor, ) process_tile = block_iter_count > Int32(0) else: @@ -2043,6 +2047,7 @@ def softmax_loop( self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, self.q_subtile_factor, seqlen_info=seqlen, + kv_subtile_factor=self.kv_subtile_factor, ) has_work = tile_block_count > Int32(0) else: @@ -2114,6 +2119,7 @@ def softmax_loop( check_m_boundary, self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, self.q_subtile_factor, + self.kv_subtile_factor, ) if not empty_tile: sScale[tidx + stage * self.m_block_size] = softmax.row_sum[0] @@ -2461,6 +2467,7 @@ def correction_loop( self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, self.q_subtile_factor, seqlen_info=seqlen, + kv_subtile_factor=self.kv_subtile_factor, ) has_work = total_block_count > Int32(0) else: diff --git a/flash_attn/cute/interface.py b/flash_attn/cute/interface.py index 1098911ec7d..379e3c24f69 100644 --- a/flash_attn/cute/interface.py +++ b/flash_attn/cute/interface.py @@ -54,6 +54,8 @@ from flash_attn.cute.utils import AuxData from flash_attn.cute.block_sparsity import ( BlockSparseTensorsTorch, + block_sparse_bwd_supports_2cta, + get_kv_subtile_factor, get_sparse_q_block_size, to_cute_block_sparse_tensors, normalize_block_sparse_config, @@ -638,11 +640,13 @@ def _flash_attn_fwd( block_sparse_broadcast_pattern = None normalized_block_sparse_tensors = None q_subtile_factor = 1 + kv_subtile_factor = 1 if block_sparse_tensors is not None: ( normalized_block_sparse_tensors, block_sparse_broadcast_pattern, q_subtile_factor, + kv_subtile_factor, ) = normalize_block_sparse_config( block_sparse_tensors, batch_size=batch_size, @@ -651,6 +655,7 @@ def _flash_attn_fwd( seqlen_k=seqlen_k, block_size=(tile_m, tile_n), q_stage=q_stage, + allow_kv_subtile=arch // 10 in [10, 11], ) if aux_tensors is not None: aux_tensor_metadata = get_aux_tensor_metadata(aux_tensors) @@ -739,6 +744,7 @@ def _flash_attn_fwd( page_size not in [None, tile_n], # paged KV non-TMA use_2cta_instrs, q_subtile_factor, + kv_subtile_factor, mma_pv_is_rs, intra_wg_overlap, use_clc_scheduler, @@ -939,6 +945,7 @@ def _flash_attn_fwd( paged_kv_non_tma=page_size not in [None, tile_n], is_varlen_q=cu_seqlens_q is not None or seqused_q is not None, q_subtile_factor=q_subtile_factor, + kv_subtile_factor=kv_subtile_factor, use_2cta_instrs=use_2cta_instrs, use_clc_scheduler=use_clc_scheduler, ) @@ -1294,8 +1301,12 @@ def _flash_attn_bwd( arch = _get_device_arch() assert arch // 10 in [9, 10, 11, 12], "Unsupported compute capability. Supported: 9.x, 10.x, 11.x, 12.x" sparse_q = None - if block_sparse_tensors is not None and arch // 10 == 9: - sparse_q = block_sparse_tensors.block_size[0] if block_sparse_tensors.block_size is not None else 128 + kv_subtile_factor = 1 + if block_sparse_tensors is not None: + if block_sparse_tensors.block_size is not None: + sparse_q = block_sparse_tensors.block_size[0] + elif arch // 10 == 9: + sparse_q = 128 num_head, head_dim = q.shape[-2:] head_dim_v = v.shape[-1] @@ -1366,12 +1377,31 @@ def _flash_attn_bwd( AtomLayoutMdQ = 1 AtomLayoutNdKV = 1 requested_disable_2cta = utils._get_disable_2cta_default() - disable_2cta = ( - requested_disable_2cta - or block_sparse_tensors is not None - ) - cluster_size = 2 if head_dim >= 128 and not disable_2cta else 1 - use_2cta_instrs = cluster_size==2 + block_sparse_supports_2cta = True + if block_sparse_tensors is not None: + kv_subtile_factor = get_kv_subtile_factor( + block_sparse_tensors, + n_block_size, + ) + block_sparse_supports_2cta = block_sparse_bwd_supports_2cta( + block_sparse_tensors, + n_block_size, + ) + disable_2cta = requested_disable_2cta or not block_sparse_supports_2cta + if block_sparse_tensors is not None and head_dim == 192 and disable_2cta: + reason = ( + "2CTA was disabled by request" + if requested_disable_2cta + else ( + f"sparse_block_size[1] must cover an even number of tile_n={n_block_size} " + f"tiles; got factor {kv_subtile_factor}" + ) + ) + raise ValueError( + f"SM100 block-sparse backward with head_dim=192 requires 2CTA; {reason}." + ) + use_2cta_instrs = head_dim >= 128 and not disable_2cta + cluster_size = 2 if use_2cta_instrs else 1 use_dedicated_hd256_kernel = arch // 10 in [10, 11] and head_dim == 256 and head_dim_v == 256 use_2cta_instrs = use_2cta_instrs or use_dedicated_hd256_kernel @@ -1596,14 +1626,15 @@ def _flash_attn_bwd( score_mod_bwd_hash = utils.hash_callable(score_mod_bwd) if score_mod_bwd else False mask_mod_hash = utils.hash_callable(mask_mod) if mask_mod else False num_aux_tensors = len(aux_tensors) if aux_tensors else 0 + aux_tensor_metadata = get_aux_tensor_metadata(aux_tensors) if aux_tensors is not None else None aux_scalar_metadata = tuple(type(s) for s in aux_scalars) if aux_scalars is not None else None cute_aux_tensors = None if aux_tensors is not None: - cute_aux_tensors = [to_cute_tensor(buf, assumed_align=None, fully_dynamic=True) for buf in aux_tensors] + cute_aux_tensors = [to_cute_aux_tensor(buf) for buf in aux_tensors] block_sparse_broadcast_pattern = None normalized_block_sparse_tensors = None - if block_sparse_tensors is not None: + if use_block_sparsity: ( normalized_block_sparse_tensors, block_sparse_broadcast_pattern, @@ -1615,6 +1646,7 @@ def _flash_attn_bwd( seqlen_k=seqlen_k, block_size=(m_block_size, n_block_size), q_subtile_factor=q_subtile_factor, + kv_subtile_factor=kv_subtile_factor, ) if deterministic: if normalized_block_sparse_tensors.dq_write_order is None: @@ -1674,8 +1706,10 @@ def _flash_attn_bwd( score_mod_bwd_hash, mask_mod_hash, num_aux_tensors, + aux_tensor_metadata, aux_scalar_metadata, use_block_sparsity, + q_subtile_factor, block_sparse_broadcast_pattern, get_broadcast_dims(q), get_broadcast_dims(k), @@ -1701,12 +1735,15 @@ def _flash_attn_bwd( pack_gqa, cluster_size, use_2cta_instrs, + q_subtile_factor, + kv_subtile_factor, deterministic, spt, score_mod_hash, score_mod_bwd_hash, mask_mod_hash, num_aux_tensors, + aux_tensor_metadata, aux_scalar_metadata, use_block_sparsity, block_sparse_broadcast_pattern, @@ -1845,6 +1882,7 @@ def _flash_attn_bwd( mask_mod=mask_mod, has_aux_tensors=aux_tensors is not None, q_subtile_factor=q_subtile_factor, + kv_subtile_factor=kv_subtile_factor, ) # Block sparse tensors for backward use Q-direction indexing (transposed from forward). diff --git a/flash_attn/cute/mask.py b/flash_attn/cute/mask.py index 312cb06150e..94a8031d14f 100644 --- a/flash_attn/cute/mask.py +++ b/flash_attn/cute/mask.py @@ -1111,6 +1111,7 @@ def get_trip_start_count_via_block_info( has_cu_seqlens_k=False, has_seqused_q=False, has_seqused_k=False, + has_cu_block_idx_offsets=False, ) n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen_info, blk_coord[0]) return n_block_min, n_block_max - n_block_min @@ -1154,6 +1155,7 @@ def get_trip_mask_bounds_via_block_info( has_cu_seqlens_k=False, has_seqused_q=False, has_seqused_k=False, + has_cu_block_idx_offsets=False, ) n_block_min, _ = block_info.get_n_block_min_max(seqlen_info, blk_coord[0]) n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask( diff --git a/flash_attn/cute/seqlen_info.py b/flash_attn/cute/seqlen_info.py index c8ba5672664..7110c8f2b78 100644 --- a/flash_attn/cute/seqlen_info.py +++ b/flash_attn/cute/seqlen_info.py @@ -78,6 +78,7 @@ class SeqlenInfoQK: has_cu_seqlens_k: cutlass.Constexpr[bool] has_seqused_q: cutlass.Constexpr[bool] has_seqused_k: cutlass.Constexpr[bool] + has_cu_block_idx_offsets: cutlass.Constexpr[bool] = False @staticmethod def create( @@ -142,6 +143,7 @@ def create( has_cu_seqlens_k=mCuSeqlensK is not None, has_seqused_q=mSeqUsedQ is not None, has_seqused_k=mSeqUsedK is not None, + has_cu_block_idx_offsets=mCuBlockIdxOffsets is not None, ) def offset_batch_Q( diff --git a/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py b/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py index b21fc16c70c..877646ee48e 100644 --- a/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py +++ b/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py @@ -44,6 +44,7 @@ def __init__( is_split_kv: bool = False, pack_gqa: bool = False, q_subtile_factor: int = 1, + kv_subtile_factor: int = 1, m_block_size: int = 128, n_block_size: int = 128, q_stage: int = 2, @@ -71,6 +72,9 @@ def __init__( assert q_subtile_factor == 1, ( "SM100 forward with head_dim=256 does not support q_subtile_factor" ) + assert kv_subtile_factor == 1, ( + "SM100 forward with head_dim=256 does not support kv_subtile_factor" + ) assert m_block_size == 128 and n_block_size == 128, ( "SM100 dedicated kernel only supports tile_m=128 and tile_n=128" ) diff --git a/tests/cute/test_mask_mod.py b/tests/cute/test_mask_mod.py index 710b7c3f202..3a6566b7365 100644 --- a/tests/cute/test_mask_mod.py +++ b/tests/cute/test_mask_mod.py @@ -22,11 +22,17 @@ from torch.nn.attention.flex_attention import create_block_mask, flex_attention import torch.nn.functional as F -from flash_attn.cute.interface import _flash_attn_fwd, _flash_attn_bwd, flash_attn_func +from flash_attn.cute.interface import ( + _flash_attn_fwd, + _flash_attn_bwd, + flash_attn_func, +) from flash_attn.cute.block_sparsity import ( BlockSparseTensorsTorch, + block_sparse_bwd_supports_2cta, fast_sampling, - normalize_block_sparse_config, + get_kv_subtile_factor, + normalize_block_sparse_config_bwd, compute_dq_write_order, compute_dq_write_order_from_block_mask, ) @@ -1009,6 +1015,216 @@ def test_sm100_block_sparse_sink_all_masked(): assert torch.allclose(lse, expected, atol=0.0, rtol=0.0) +def make_empty_block_sparse_tensors(sparse_block_size_kv: int) -> BlockSparseTensorsTorch: + """Build shape-only metadata for block-sparse dispatch helper tests.""" + return BlockSparseTensorsTorch( + mask_block_cnt=torch.empty(0, dtype=torch.int32), + mask_block_idx=torch.empty(0, dtype=torch.int32), + block_size=(256, sparse_block_size_kv), + ) + + +def test_sm100_block_sparse_bwd_kv_subtile_selects_cta_policy(): + n_block_size = 128 + + assert get_kv_subtile_factor(None, n_block_size) == 1 + assert block_sparse_bwd_supports_2cta(None, n_block_size) + assert get_kv_subtile_factor( + make_empty_block_sparse_tensors(n_block_size), n_block_size + ) == 1 + assert not block_sparse_bwd_supports_2cta( + make_empty_block_sparse_tensors(n_block_size), n_block_size + ) + assert get_kv_subtile_factor( + make_empty_block_sparse_tensors(2 * n_block_size), n_block_size + ) == 2 + assert block_sparse_bwd_supports_2cta( + make_empty_block_sparse_tensors(2 * n_block_size), n_block_size + ) + assert get_kv_subtile_factor( + make_empty_block_sparse_tensors(3 * n_block_size), n_block_size + ) == 3 + assert not block_sparse_bwd_supports_2cta( + make_empty_block_sparse_tensors(3 * n_block_size), n_block_size + ) + assert get_kv_subtile_factor( + make_empty_block_sparse_tensors(4 * n_block_size), n_block_size + ) == 4 + assert block_sparse_bwd_supports_2cta( + make_empty_block_sparse_tensors(4 * n_block_size), n_block_size + ) + with pytest.raises(ValueError, match=r"multiple of tile_n=128; got 192"): + get_kv_subtile_factor( + make_empty_block_sparse_tensors(192), n_block_size + ) + + +def test_block_sparse_bwd_normalize_accepts_odd_kv_subtile_for_1cta(): + tensors = BlockSparseTensorsTorch( + mask_block_cnt=torch.zeros((1, 1, 1), device="cuda", dtype=torch.int32), + mask_block_idx=torch.zeros((1, 1, 1, 1), device="cuda", dtype=torch.int32), + block_size=(256, 384), + ) + normalized, _ = normalize_block_sparse_config_bwd( + tensors, + batch_size=1, + num_head=1, + seqlen_q=256, + seqlen_k=384, + block_size=(128, 128), + q_subtile_factor=2, + kv_subtile_factor=3, + ) + + assert normalized.block_size == (256, 384) + + +@pytest.mark.skipif(COMPUTE_CAPABILITY != 10, reason="SM100-only test") +@pytest.mark.parametrize( + "headdim,headdim_v,seqlen_q,seqlen_k,sparse_tile_m,sparse_tile_n,expected_use_2cta", + [ + (128, 128, 384, 768, 256, 384, False), + (128, 128, 384, 768, 256, 512, True), + (192, 128, 384, 384, 256, 256, True), + (192, 128, 1024, 1024, 512, 512, True), + ], +) +def test_sm100_block_sparse_bwd_kv_subtile_actual_kernel( + headdim, + headdim_v, + seqlen_q, + seqlen_k, + sparse_tile_m, + sparse_tile_n, + expected_use_2cta, +): + from flash_attn.cute import flash_bwd_sm100 + + torch.manual_seed(124) + batch_size = 1 + nheads = 1 + dtype = torch.bfloat16 + tile_m = 128 + tile_n = 128 + + mask_mod_cute, mask_mod_flex = get_mask_pair( + "causal", seqlen_q=seqlen_q, seqlen_k=seqlen_k, window_size=None + ) + tensors = create_tensors( + batch_size, seqlen_q, seqlen_k, nheads, nheads, headdim, headdim_v, dtype + ) + block_sparse_mask_fwd, block_sparse_mask_bwd, block_mask = _build_block_sparse_masks_for_bwd( + mask_mod_flex=mask_mod_flex, + batch_size=batch_size, + nheads=nheads, + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + tile_m=tile_m, + tile_n=tile_n, + spt=False, + sparse_tile_m=sparse_tile_m, + sparse_tile_n=sparse_tile_n, + return_block_mask=True, + ) + + out_cute, lse_cute = _flash_attn_fwd( + q=tensors["q"], + k=tensors["k"], + v=tensors["v"], + out=tensors["out"], + lse=tensors["lse"], + cu_seqlens_q=None, + cu_seqlens_k=None, + seqused_q=None, + seqused_k=None, + page_table=None, + softmax_scale=1.0 / math.sqrt(headdim), + causal=False, + softcap=None, + window_size_left=None, + window_size_right=None, + learnable_sink=None, + tile_mn=(tile_m, tile_n), + pack_gqa=False, + _arch=None, + score_mod=None, + mask_mod=mask_mod_cute, + block_sparse_tensors=block_sparse_mask_fwd, + return_lse=True, + ) + grad_out = torch.randn_like(out_cute) + observed = {} + original_init = flash_bwd_sm100.FlashAttentionBackwardSm100.__init__ + + def wrapped_init(self, *args, **kwargs): + observed["use_2cta_instrs"] = kwargs.get("use_2cta_instrs") + observed["kernel_q_subtile_factor"] = kwargs.get("q_subtile_factor") + observed["kernel_kv_subtile_factor"] = kwargs.get("kv_subtile_factor") + return original_init(self, *args, **kwargs) + + def wrapped_normalize(*args, **kwargs): + observed["q_subtile_factor"] = kwargs.get("q_subtile_factor") + observed["kv_subtile_factor"] = kwargs.get("kv_subtile_factor") + return normalize_block_sparse_config_bwd(*args, **kwargs) + + compile_cache = _flash_attn_bwd.compile_cache + _flash_attn_bwd.compile_cache = get_jit_cache("test_mask_mod.kv_subtile_bwd") + try: + with ( + mock.patch.object(flash_bwd_sm100.FlashAttentionBackwardSm100, "__init__", wrapped_init), + mock.patch( + "flash_attn.cute.interface.normalize_block_sparse_config_bwd", + side_effect=wrapped_normalize, + ), + ): + dq_cute, dk_cute, dv_cute = run_cute_mask_bwd( + tensors["q"], + tensors["k"], + tensors["v"], + out_cute, + lse_cute, + grad_out, + mask_mod_cute, + block_sparse_mask_bwd=block_sparse_mask_bwd, + tile_m=tile_m, + tile_n=tile_n, + ) + finally: + _flash_attn_bwd.compile_cache.clear() + _flash_attn_bwd.compile_cache = compile_cache + + expected_q_subtile_factor = sparse_tile_m // tile_m + expected_kv_subtile_factor = sparse_tile_n // tile_n + assert observed == { + "kernel_q_subtile_factor": expected_q_subtile_factor, + "kernel_kv_subtile_factor": expected_kv_subtile_factor, + "q_subtile_factor": expected_q_subtile_factor, + "kv_subtile_factor": expected_kv_subtile_factor, + "use_2cta_instrs": expected_use_2cta, + } + out_ref_fp32, dq_ref_fp32, dk_ref_fp32, dv_ref_fp32 = run_flex_reference_bwd( + tensors["q"], tensors["k"], tensors["v"], block_mask, grad_out, dtype=torch.float32 + ) + out_pt, dq_pt, dk_pt, dv_pt = run_flex_reference_bwd( + tensors["q"], tensors["k"], tensors["v"], block_mask, grad_out + ) + + assert_fwd_matches_reference(out_cute, out_ref_fp32, out_pt) + assert_bwd_matches_reference( + dq_cute, + dk_cute, + dv_cute, + dq_ref_fp32, + dk_ref_fp32, + dv_ref_fp32, + dq_pt, + dk_pt, + dv_pt, + dtype, + min(seqlen_q, seqlen_k), + ) + + @pytest.mark.skipif(COMPUTE_CAPABILITY != 10, reason="SM100-only test") def test_sm100_block_sparse_q_stage1(): from flash_attn.cute import flash_fwd_sm100 @@ -1148,6 +1364,81 @@ def test_sm100_block_sparse_coarse_blocks(): ) +@pytest.mark.skipif(COMPUTE_CAPABILITY != 10, reason="SM100-only test") +@pytest.mark.parametrize("headdim", [128, 192]) +def test_sm100_block_sparse_coarse_kv_masks_tail_subtiles(headdim): + """Exercise ragged coarse-KV blocks whose expanded physical subtiles include K padding.""" + torch.manual_seed(13005) + seqlen_q = 257 + seqlen_k = 513 + nheads = 1 + headdim_v = 128 + dtype = torch.bfloat16 + tile_m = 128 + tile_n = 128 + sparse_tile_m = 256 + sparse_tile_n = 256 + batch_size = 1 + + mask_mod_cute, mask_mod_flex = get_mask_pair( + "mini_causal", seqlen_q=seqlen_q, seqlen_k=seqlen_k, window_size=None + ) + tensors = create_tensors( + batch_size, seqlen_q, seqlen_k, nheads, nheads, headdim, headdim_v, dtype + ) + + bm = create_block_mask( + mask_mod_flex, + batch_size, + nheads, + seqlen_q, + seqlen_k, + device="cuda", + BLOCK_SIZE=(sparse_tile_m, sparse_tile_n), + ) + ( + _seq_q, + _seq_k, + kv_mask_cnt, + kv_mask_idx, + full_kv_cnt, + full_kv_idx, + *_, + ) = bm.as_tuple() + + block_sparse_mask_fwd = BlockSparseTensorsTorch( + mask_block_cnt=kv_mask_cnt, + mask_block_idx=kv_mask_idx, + full_block_cnt=full_kv_cnt, + full_block_idx=full_kv_idx, + block_size=(sparse_tile_m, sparse_tile_n), + ) + + out_cute, _ = _flash_attn_fwd( + q=tensors["q"], + k=tensors["k"], + v=tensors["v"], + out=tensors["out"], + lse=tensors["lse"], + softmax_scale=1.0 / math.sqrt(headdim), + causal=False, + tile_mn=(tile_m, tile_n), + pack_gqa=False, + mask_mod=mask_mod_cute, + block_sparse_tensors=block_sparse_mask_fwd, + return_lse=True, + ) + out_ref_fp32 = compute_reference_flex_attn( + {name: tensor.float() for name, tensor in tensors.items()}, + mask_mod_flex, + (sparse_tile_m, sparse_tile_n), + ) + out_ref = compute_reference_flex_attn( + tensors, mask_mod_flex, (sparse_tile_m, sparse_tile_n) + ) + assert_fwd_matches_reference(out_cute, out_ref_fp32, out_ref) + + @pytest.mark.skipif(COMPUTE_CAPABILITY != 10, reason="SM100-only test") def test_sm100_block_sparse_coarse_blocks_mismatch(): torch.manual_seed(0) @@ -1195,41 +1486,31 @@ def test_sm100_block_sparse_coarse_blocks_mismatch(): block_size=(sparse_tile_m, tile_n), ) - observed = {} - original_normalize = normalize_block_sparse_config - - def wrapped_normalize(*args, **kwargs): - normalized, pattern, q_subtile_factor = original_normalize(*args, **kwargs) - observed["q_subtile_factor"] = q_subtile_factor - return normalized, pattern, q_subtile_factor - - with mock.patch("flash_attn.cute.interface.normalize_block_sparse_config", wrapped_normalize): - out_cute, _ = _flash_attn_fwd( - q=tensors["q"], - k=tensors["k"], - v=tensors["v"], - out=tensors["out"], - lse=tensors["lse"], - cu_seqlens_q=None, - cu_seqlens_k=None, - seqused_q=None, - seqused_k=None, - page_table=None, - softmax_scale=1.0 / math.sqrt(headdim), - causal=False, - softcap=None, - window_size_left=None, - window_size_right=None, - learnable_sink=None, - tile_mn=(tile_m, tile_n), - pack_gqa=False, - _arch=None, - score_mod=None, - mask_mod=mask_mod_cute, - block_sparse_tensors=block_sparse_mask_fwd, - return_lse=True, - ) - assert observed.get("q_subtile_factor") == 2 + out_cute, _ = _flash_attn_fwd( + q=tensors["q"], + k=tensors["k"], + v=tensors["v"], + out=tensors["out"], + lse=tensors["lse"], + cu_seqlens_q=None, + cu_seqlens_k=None, + seqused_q=None, + seqused_k=None, + page_table=None, + softmax_scale=1.0 / math.sqrt(headdim), + causal=False, + softcap=None, + window_size_left=None, + window_size_right=None, + learnable_sink=None, + tile_mn=(tile_m, tile_n), + pack_gqa=False, + _arch=None, + score_mod=None, + mask_mod=mask_mod_cute, + block_sparse_tensors=block_sparse_mask_fwd, + return_lse=True, + ) tensors_fp32 = { k: v.float() if v.dtype in [torch.float16, torch.bfloat16] else v @@ -1781,7 +2062,7 @@ def causal_mask(b, h, q_idx, kv_idx): pt_error = (out_ref - out_ref_fp32).abs().max().item() cute_error = (out_fwd - out_ref_fp32).abs().max().item() - print(f"\nGQA expand stride=0 test:") + print("\nGQA expand stride=0 test:") print(f" Forward: kernel err={cute_error:.2e}, ref err={pt_error:.2e}, atol={fwd_atol:.2e}") assert cute_error <= rtol * pt_error + fwd_atol, ( f"Forward error {cute_error:.2e} exceeds {rtol}x ref error {pt_error:.2e} + {fwd_atol:.2e}" @@ -2036,8 +2317,12 @@ def _build_block_sparse_masks_for_bwd( tile_m, tile_n, spt, + sparse_tile_m=None, + sparse_tile_n=None, + return_block_mask=False, ): - sparse_tile_m = 2 * tile_m if COMPUTE_CAPABILITY == 10 else tile_m + sparse_tile_m = sparse_tile_m or (2 * tile_m if COMPUTE_CAPABILITY == 10 else tile_m) + sparse_tile_n = sparse_tile_n or tile_n bm = create_block_mask( mask_mod_flex, batch_size, @@ -2045,7 +2330,7 @@ def _build_block_sparse_masks_for_bwd( seqlen_q, seqlen_k, device="cuda", - BLOCK_SIZE=(sparse_tile_m, tile_n), + BLOCK_SIZE=(sparse_tile_m, sparse_tile_n), ) ( _seq_q, @@ -2066,21 +2351,24 @@ def _build_block_sparse_masks_for_bwd( mask_block_idx=kv_mask_idx, full_block_cnt=full_kv_cnt, full_block_idx=full_kv_idx, - block_size=(sparse_tile_m, tile_n), + block_size=(sparse_tile_m, sparse_tile_n), ) block_sparse_mask_bwd = BlockSparseTensorsTorch( mask_block_cnt=q_mask_cnt, mask_block_idx=q_mask_idx, full_block_cnt=full_q_cnt, full_block_idx=full_q_idx, - block_size=(sparse_tile_m, tile_n), + block_size=(sparse_tile_m, sparse_tile_n), ) dq_write_order = compute_dq_write_order_from_block_mask(bm, spt=spt) - return block_sparse_mask_fwd, block_sparse_mask_bwd._replace( + block_sparse_mask_bwd = block_sparse_mask_bwd._replace( dq_write_order=dq_write_order[0], dq_write_order_full=dq_write_order[1], spt=spt, ) + if return_block_mask: + return block_sparse_mask_fwd, block_sparse_mask_bwd, bm + return block_sparse_mask_fwd, block_sparse_mask_bwd @pytest.mark.skipif(COMPUTE_CAPABILITY not in (10, 11), reason="deterministic bwd only supported on sm100/sm110") @@ -2262,6 +2550,285 @@ def _setup_block_sparse_deterministic_validation_case(): return q, k, v, out_cute, lse_cute, torch.randn_like(out_cute), block_sparse_mask_bwd, tile_m, tile_n +@pytest.mark.skipif(COMPUTE_CAPABILITY != 10, reason="SM100-only deterministic coarse-KV repro") +def test_block_sparse_bwd_deterministic_kv_subtile_repro(): + torch.manual_seed(42) + batch_size = 1 + nheads = 1 + seqlen_q = 384 + seqlen_k = 1024 + headdim = 128 + tile_m = 128 + tile_n = 128 + sparse_tile_m = 256 + sparse_tile_n = 512 + dtype = torch.bfloat16 + + def mask_mod_flex(b, h, q_idx, kv_idx): + return q_idx >= 0 + + tensors = create_tensors( + batch_size, seqlen_q, seqlen_k, nheads, nheads, headdim, headdim, dtype + ) + block_sparse_mask_fwd, block_sparse_mask_bwd, block_mask = _build_block_sparse_masks_for_bwd( + mask_mod_flex=mask_mod_flex, + batch_size=batch_size, + nheads=nheads, + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + tile_m=tile_m, + tile_n=tile_n, + spt=False, + sparse_tile_m=sparse_tile_m, + sparse_tile_n=sparse_tile_n, + return_block_mask=True, + ) + out_cute, lse_cute = _flash_attn_fwd( + q=tensors["q"], + k=tensors["k"], + v=tensors["v"], + out=tensors["out"], + lse=tensors["lse"], + softmax_scale=1.0 / math.sqrt(headdim), + tile_mn=(tile_m, tile_n), + mask_mod=None, + block_sparse_tensors=block_sparse_mask_fwd, + return_lse=True, + ) + grad_out = torch.randn_like(out_cute) + + dq0, dk0, dv0 = run_cute_mask_bwd( + tensors["q"], + tensors["k"], + tensors["v"], + out_cute, + lse_cute, + grad_out, + None, + block_sparse_mask_bwd=block_sparse_mask_bwd, + tile_m=tile_m, + tile_n=tile_n, + deterministic=True, + ) + dq1, dk1, dv1 = run_cute_mask_bwd( + tensors["q"], + tensors["k"], + tensors["v"], + out_cute, + lse_cute, + grad_out, + None, + block_sparse_mask_bwd=block_sparse_mask_bwd, + tile_m=tile_m, + tile_n=tile_n, + deterministic=True, + ) + dq_ref_kernel, dk_ref_kernel, dv_ref_kernel = run_cute_mask_bwd( + tensors["q"], + tensors["k"], + tensors["v"], + out_cute, + lse_cute, + grad_out, + None, + block_sparse_mask_bwd=block_sparse_mask_bwd, + tile_m=tile_m, + tile_n=tile_n, + deterministic=False, + ) + out_ref_fp32, dq_ref_fp32, dk_ref_fp32, dv_ref_fp32 = run_flex_reference_bwd( + tensors["q"], tensors["k"], tensors["v"], block_mask, grad_out, dtype=torch.float32 + ) + out_pt, dq_pt, dk_pt, dv_pt = run_flex_reference_bwd( + tensors["q"], tensors["k"], tensors["v"], block_mask, grad_out + ) + + assert_fwd_matches_reference(out_cute, out_ref_fp32, out_pt) + dq_ref = dq_ref_fp32.to(dtype) + pt_dq_err = (dq_pt - dq_ref).abs().max().item() + cute_dq_err = (dq0 - dq_ref).abs().max().item() + assert cute_dq_err <= 2 * pt_dq_err + 1e-5 + torch.testing.assert_close(dq0, dq_ref_kernel, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(dk0, dk_ref_kernel, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(dv0, dv_ref_kernel, rtol=2e-2, atol=2e-2) + assert torch.equal(dq1, dq0) + assert torch.equal(dk1, dk0) + assert torch.equal(dv1, dv0) + + +@pytest.mark.skipif(COMPUTE_CAPABILITY != 10, reason="SM100-only deterministic coarse-KV SPT test") +@pytest.mark.parametrize( + "seqlen,sparse_tile_n", + [ + (1152, 384), # f=3 odd -> 1CTA, coarse grid tiles the schedule exactly + (1024, 512), # f=4 -> 2CTA, exact cover + (1023, 384), # f=3 -> 1CTA, truncated tail sparse block (9 implied vs 8 scheduled) + (641, 512), # f=4 -> 2CTA, truncated tail sparse block (4 implied vs 3 cta-groups) + (1793, 768), # f=6 -> 2CTA, 3 groups per sparse block, truncated tail (2 of 3) + (1089, 512), # f=4 -> 2CTA, odd scheduled tile count (cluster pad) + truncated tail + ], +) +def test_block_sparse_bwd_deterministic_spt_kv_subtile(seqlen, sparse_tile_n): + """Deterministic SPT with coarse KV blocks: covers the in-kernel lock expansion + (local_group reversal within the scheduled group count) and the semaphore bridge + over unscheduled tail groups, for both 1CTA and 2CTA paths. Asserts run-to-run + bitwise determinism and closeness to the non-deterministic reference.""" + torch.manual_seed(42) + batch_size = 1 + nheads = 2 + headdim = 128 + tile_m = tile_n = 128 + dtype = torch.bfloat16 + + mask_mod_cute, mask_mod_flex = get_mask_pair("causal", seqlen_q=seqlen, seqlen_k=seqlen) + tensors = create_tensors(batch_size, seqlen, seqlen, nheads, nheads, headdim, headdim, dtype) + block_sparse_mask_fwd, block_sparse_mask_bwd, block_mask = _build_block_sparse_masks_for_bwd( + mask_mod_flex=mask_mod_flex, + batch_size=batch_size, + nheads=nheads, + seqlen_q=seqlen, + seqlen_k=seqlen, + tile_m=tile_m, + tile_n=tile_n, + spt=True, + sparse_tile_m=256, + sparse_tile_n=sparse_tile_n, + return_block_mask=True, + ) + out_cute, lse_cute = _flash_attn_fwd( + q=tensors["q"], + k=tensors["k"], + v=tensors["v"], + out=tensors["out"], + lse=tensors["lse"], + softmax_scale=1.0 / math.sqrt(headdim), + tile_mn=(tile_m, tile_n), + mask_mod=mask_mod_cute, + block_sparse_tensors=block_sparse_mask_fwd, + return_lse=True, + ) + grad_out = torch.randn_like(out_cute) + + def bwd(deterministic): + return run_cute_mask_bwd( + tensors["q"], + tensors["k"], + tensors["v"], + out_cute, + lse_cute, + grad_out, + mask_mod_cute, + block_sparse_mask_bwd=block_sparse_mask_bwd, + tile_m=tile_m, + tile_n=tile_n, + deterministic=deterministic, + ) + + dq0, dk0, dv0 = bwd(True) + dq1, dk1, dv1 = bwd(True) + dq_nd, dk_nd, dv_nd = bwd(False) + + assert torch.equal(dq1, dq0) + assert torch.equal(dk1, dk0) + assert torch.equal(dv1, dv0) + torch.testing.assert_close(dq0, dq_nd, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(dk0, dk_nd, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(dv0, dv_nd, rtol=2e-2, atol=2e-2) + + _, dq_ref_fp32, dk_ref_fp32, dv_ref_fp32 = run_flex_reference_bwd( + tensors["q"], tensors["k"], tensors["v"], block_mask, grad_out, dtype=torch.float32 + ) + _, dq_pt, dk_pt, dv_pt = run_flex_reference_bwd( + tensors["q"], tensors["k"], tensors["v"], block_mask, grad_out + ) + assert_bwd_matches_reference( + dq0, dk0, dv0, + dq_ref_fp32, dk_ref_fp32, dv_ref_fp32, + dq_pt, dk_pt, dv_pt, + dtype, min_seqlen=seqlen, + ) + + +@pytest.mark.skipif(COMPUTE_CAPABILITY != 10, reason="SM100-only compile-key regression") +def test_sm100_block_sparse_bwd_q_subtile_compile_key(): + """Two bwd calls differing only in sparse_block_size_q must not share a kernel. + + q_subtile_factor is a kernel constexpr; if it is missing from the compile key the + second call silently reuses the first kernel and produces wrong gradients. + """ + batch_size = 1 + nheads = 2 + seqlen = 1024 + headdim = 128 + tile_m = tile_n = 128 + dtype = torch.bfloat16 + + compile_cache = _flash_attn_bwd.compile_cache + _flash_attn_bwd.compile_cache = get_jit_cache("test_mask_mod.q_subtile_compile_key") + try: + for sparse_tile_m in (256, 512): + torch.manual_seed(7) + mask_mod_cute, mask_mod_flex = get_mask_pair( + "causal", seqlen_q=seqlen, seqlen_k=seqlen + ) + tensors = create_tensors( + batch_size, seqlen, seqlen, nheads, nheads, headdim, headdim, dtype + ) + sparse_fwd, sparse_bwd, block_mask = _build_block_sparse_masks_for_bwd( + mask_mod_flex=mask_mod_flex, + batch_size=batch_size, + nheads=nheads, + seqlen_q=seqlen, + seqlen_k=seqlen, + tile_m=tile_m, + tile_n=tile_n, + spt=False, + sparse_tile_m=sparse_tile_m, + sparse_tile_n=tile_n, + return_block_mask=True, + ) + out_cute, lse_cute = _flash_attn_fwd( + q=tensors["q"], + k=tensors["k"], + v=tensors["v"], + out=tensors["out"], + lse=tensors["lse"], + softmax_scale=1.0 / math.sqrt(headdim), + tile_mn=(tile_m, tile_n), + mask_mod=mask_mod_cute, + block_sparse_tensors=sparse_fwd, + return_lse=True, + ) + grad_out = torch.randn_like(out_cute) + dq, dk, dv = run_cute_mask_bwd( + tensors["q"], + tensors["k"], + tensors["v"], + out_cute, + lse_cute, + grad_out, + mask_mod_cute, + block_sparse_mask_bwd=sparse_bwd, + tile_m=tile_m, + tile_n=tile_n, + ) + _, dq_ref_fp32, dk_ref_fp32, dv_ref_fp32 = run_flex_reference_bwd( + tensors["q"], tensors["k"], tensors["v"], block_mask, grad_out, + dtype=torch.float32, + ) + _, dq_pt, dk_pt, dv_pt = run_flex_reference_bwd( + tensors["q"], tensors["k"], tensors["v"], block_mask, grad_out + ) + assert_bwd_matches_reference( + dq, dk, dv, + dq_ref_fp32, dk_ref_fp32, dv_ref_fp32, + dq_pt, dk_pt, dv_pt, + dtype, min_seqlen=seqlen, + ) + finally: + _flash_attn_bwd.compile_cache = compile_cache + + @pytest.mark.skipif(COMPUTE_CAPABILITY not in (10, 11), reason="deterministic bwd only supported on sm100/sm110") def test_block_sparse_bwd_deterministic_missing_dq_write_order_raises(): q, k, v, out_cute, lse_cute, grad_out, block_sparse_mask_bwd, tile_m, tile_n = ( @@ -2399,26 +2966,30 @@ def test_block_sparse_splitkv_matches_unsplit(): torch.manual_seed(123) batch_size = 1 nheads = 4 - seqlen = 2048 + seqlen_q = 513 + seqlen_k = 769 headdim = 64 tile_m = 128 tile_n = 128 dtype = torch.bfloat16 sparse_tile_m = 2 * tile_m + sparse_tile_n = 2 * tile_n - mask_mod_cute, mask_mod_flex = get_mask_pair("causal", seqlen_q=seqlen, seqlen_k=seqlen) + mask_mod_cute, mask_mod_flex = get_mask_pair( + "causal", seqlen_q=seqlen_q, seqlen_k=seqlen_k + ) tensors = create_tensors( - batch_size, seqlen, seqlen, nheads, nheads, headdim, headdim, dtype + batch_size, seqlen_q, seqlen_k, nheads, nheads, headdim, headdim, dtype ) bm = create_block_mask( mask_mod_flex, batch_size, nheads, - seqlen, - seqlen, + seqlen_q, + seqlen_k, device="cuda", - BLOCK_SIZE=(sparse_tile_m, tile_n), + BLOCK_SIZE=(sparse_tile_m, sparse_tile_n), ) (_, _, kv_mask_cnt, kv_mask_idx, full_kv_cnt, full_kv_idx, *_) = bm.as_tuple() block_sparse_fwd = BlockSparseTensorsTorch( @@ -2426,7 +2997,7 @@ def test_block_sparse_splitkv_matches_unsplit(): mask_block_idx=kv_mask_idx, full_block_cnt=full_kv_cnt, full_block_idx=full_kv_idx, - block_size=(sparse_tile_m, tile_n), + block_size=(sparse_tile_m, sparse_tile_n), ) out_unsplit, lse_unsplit = _flash_attn_fwd( @@ -2452,18 +3023,21 @@ def test_block_sparse_splitkv_matches_unsplit(): causal=False, mask_mod=mask_mod_cute, block_sparse_tensors=block_sparse_fwd, - num_splits=3, + num_splits=5, return_lse=True, ) - out_ref = compute_reference_flex_attn(tensors, mask_mod_flex, block_size=(sparse_tile_m, tile_n)) + out_ref = compute_reference_flex_attn( + tensors, mask_mod_flex, block_size=(sparse_tile_m, sparse_tile_n) + ) out_ref_fp32 = compute_reference_flex_attn( {name: tensor.float() for name, tensor in tensors.items()}, mask_mod_flex, - block_size=(sparse_tile_m, tile_n), + block_size=(sparse_tile_m, sparse_tile_n), ) assert_fwd_matches_reference(out_split, out_ref_fp32, out_ref) + assert torch.allclose(out_split, out_unsplit, atol=4e-3, rtol=4e-3) assert torch.allclose(lse_split, lse_unsplit, atol=2e-3, rtol=2e-3) diff --git a/tests/cute/test_mask_mod_varlen.py b/tests/cute/test_mask_mod_varlen.py index 6e37e9ed4b8..9ffc6b23cf0 100644 --- a/tests/cute/test_mask_mod_varlen.py +++ b/tests/cute/test_mask_mod_varlen.py @@ -1044,6 +1044,78 @@ def make_cu_seqlens(seqlens): ) +@pytest.mark.skipif(COMPUTE_CAPABILITY not in (10, 11), reason="SM100/SM110 coarse KV forward only") +@pytest.mark.parametrize("seqlens_k", [[512, 512], [384, 384], [128, 128]]) +@pytest.mark.parametrize("varlen_k", [False, True]) +def test_varlen_block_sparse_coarse_kv_metadata_stride_repro(seqlens_k, varlen_k): + torch.manual_seed(42) + device = "cuda" + seqlens_q = [512, 512] + num_heads = 1 + head_dim = 128 + dtype = torch.bfloat16 + physical_tile_n = 128 + sparse_tile_m = 256 + sparse_tile_n = 256 + + q = torch.randn(sum(seqlens_q), num_heads, head_dim, device=device, dtype=dtype) + cu_seqlens_q = torch.tensor( + [0] + list(torch.tensor(seqlens_q).cumsum(0).tolist()), + device=device, + dtype=torch.int32, + ) + if varlen_k: + k = torch.randn(sum(seqlens_k), num_heads, head_dim, device=device, dtype=dtype) + v = torch.randn_like(k) + cu_seqlens_k = torch.tensor( + [0] + list(torch.tensor(seqlens_k).cumsum(0).tolist()), + device=device, + dtype=torch.int32, + ) + else: + k = torch.randn( + len(seqlens_k), max(seqlens_k), num_heads, head_dim, device=device, dtype=dtype + ) + v = torch.randn_like(k) + cu_seqlens_k = None + mask_mod = get_mask_pair("block_diagonal")[0] + block_sparse_tensors = _make_block_sparse_tensors( + mask_mod=mask_mod, + seqlens_q=seqlens_q, + seqlens_k=seqlens_k, + num_heads=num_heads, + tile_m=sparse_tile_m, + tile_n=sparse_tile_n, + device=device, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + ) + + out_with_block_sparsity = _run_fwd( + q, + k, + v, + mask_mod, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + block_sparse_tensors=block_sparse_tensors, + ) + out_no_block_sparsity = _run_fwd( + q, + k, + v, + mask_mod, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + ) + + max_err = (out_with_block_sparsity - out_no_block_sparsity).abs().max().item() + assert max_err <= 0.01, ( + f"varlen coarse-KV block-sparse output differs from mask-mod-only by {max_err} " + f"with physical tile_n={physical_tile_n} and sparse tile_n={sparse_tile_n}" + ) + + VARLEN_BLOCK_SPARSE_SPLITKV_SEQLENS = [ ([128], [2048]), ([96], [1536]), From 538aab77209c44124b9ddc7966af3835adfce5eb Mon Sep 17 00:00:00 2001 From: drisspg Date: Wed, 17 Jun 2026 20:54:46 -0700 Subject: [PATCH 2/2] Stop the log spew stack-info: PR: https://github.com/Dao-AILab/flash-attention/pull/2664, branch: drisspg/stack/46 --- flash_attn/cute/flash_bwd_sm100.py | 12 ++++++------ flash_attn/cute/flash_fwd_mla_sm100.py | 4 ++-- flash_attn/cute/flash_fwd_sm100.py | 4 ++-- .../sm100_hd256_2cta_fmha_backward_dkdvkernel.py | 4 ++-- .../cute/sm100_hd256_2cta_fmha_backward_dqkernel.py | 4 ++-- flash_attn/cute/sm100_hd256_2cta_fmha_forward.py | 4 ++-- 6 files changed, 16 insertions(+), 16 deletions(-) diff --git a/flash_attn/cute/flash_bwd_sm100.py b/flash_attn/cute/flash_bwd_sm100.py index 799b3a1b414..ec61b4c2111 100644 --- a/flash_attn/cute/flash_bwd_sm100.py +++ b/flash_attn/cute/flash_bwd_sm100.py @@ -1121,10 +1121,10 @@ def kernel( dQ_cluster_empty_mbar_ptr = storage.dQ_cluster_empty_mbar_ptr.data_ptr() if const_expr(self.use_2cta_instrs): - dS_cluster_full_mbar_ptr = storage.dS_cluster_full_mbar_ptr - dS_cluster_empty_mbar_ptr = storage.dS_cluster_empty_mbar_ptr - dS_cluster_leader_mbar_ptr = storage.dS_cluster_leader_mbar_ptr - dQaccum_empty_mbar_ptr = storage.dQaccum_empty_mbar_ptr + dS_cluster_full_mbar_ptr = storage.dS_cluster_full_mbar_ptr.ptr + dS_cluster_empty_mbar_ptr = storage.dS_cluster_empty_mbar_ptr.ptr + dS_cluster_leader_mbar_ptr = storage.dS_cluster_leader_mbar_ptr.ptr + dQaccum_empty_mbar_ptr = storage.dQaccum_empty_mbar_ptr.ptr else: dS_cluster_full_mbar_ptr = None dS_cluster_empty_mbar_ptr = None @@ -1156,11 +1156,11 @@ def kernel( * len((self.mma_warp_id, *self.compute_warp_ids, *self.reduce_warp_ids)), ) tmem = cutlass.utils.TmemAllocator( - storage.tmem_holding_buf, + storage.tmem_holding_buf.ptr, barrier_for_retrieve=tmem_alloc_barrier, allocator_warp_id=self.mma_warp_id, is_two_cta=self.use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, ) # UMMA producers and AsyncThread consumers diff --git a/flash_attn/cute/flash_fwd_mla_sm100.py b/flash_attn/cute/flash_fwd_mla_sm100.py index 84c349c5e3a..643885e19a5 100644 --- a/flash_attn/cute/flash_fwd_mla_sm100.py +++ b/flash_attn/cute/flash_fwd_mla_sm100.py @@ -821,11 +821,11 @@ def kernel( num_threads=self.num_mma_threads + self.num_softmax_threads + self.num_epilogue_threads, ) tmem = cutlass.utils.TmemAllocator( - storage.tmem_holding_buf, + storage.tmem_holding_buf.ptr, barrier_for_retrieve=tmem_alloc_barrier, allocator_warp_id=self.mma_warp_id, is_two_cta=self.use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.mbar_ptr_tmem_dealloc, + two_cta_tmem_dealloc_mbar_ptr=storage.mbar_ptr_tmem_dealloc.ptr, ) # ==== Prefetch TMA descriptors ==== diff --git a/flash_attn/cute/flash_fwd_sm100.py b/flash_attn/cute/flash_fwd_sm100.py index 24a38ade384..4ed7155fc6e 100644 --- a/flash_attn/cute/flash_fwd_sm100.py +++ b/flash_attn/cute/flash_fwd_sm100.py @@ -885,11 +885,11 @@ def kernel( ) # Tensor memory dealloc barrier init tmem = cutlass.utils.TmemAllocator( - storage.tmem_holding_buf, + storage.tmem_holding_buf.ptr, barrier_for_retrieve=tmem_alloc_barrier, allocator_warp_id=self.mma_warp_id, is_two_cta=self.use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, ) ThreadCooperativeGroup = partial(pipeline.CooperativeGroup, pipeline.Agent.Thread) diff --git a/flash_attn/cute/sm100_hd256_2cta_fmha_backward_dkdvkernel.py b/flash_attn/cute/sm100_hd256_2cta_fmha_backward_dkdvkernel.py index 885ae336f5f..5abeb4a0663 100644 --- a/flash_attn/cute/sm100_hd256_2cta_fmha_backward_dkdvkernel.py +++ b/flash_attn/cute/sm100_hd256_2cta_fmha_backward_dkdvkernel.py @@ -1029,11 +1029,11 @@ def dkdv_bwd( ) tmem = utils.TmemAllocator( - storage.tmem_holding_buf, + storage.tmem_holding_buf.ptr, barrier_for_retrieve=tmem_alloc_barrier, allocator_warp_id=self.load_warp_id, is_two_cta=True, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, ) tmem.allocate(self.tmem_alloc_cols) diff --git a/flash_attn/cute/sm100_hd256_2cta_fmha_backward_dqkernel.py b/flash_attn/cute/sm100_hd256_2cta_fmha_backward_dqkernel.py index 25d6a91de70..e60d1b5ffff 100644 --- a/flash_attn/cute/sm100_hd256_2cta_fmha_backward_dqkernel.py +++ b/flash_attn/cute/sm100_hd256_2cta_fmha_backward_dqkernel.py @@ -751,11 +751,11 @@ def kernel( # Tensor memory dealloc barrier init tmem = utils.TmemAllocator( - storage.tmem_holding_buf, + storage.tmem_holding_buf.ptr, barrier_for_retrieve=self.tmem_alloc_barrier, allocator_warp_id=self.epilogue_warp_ids[0], is_two_cta=True, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, ) tmem.allocate(self.tmem_alloc_cols) tmem.wait_for_alloc() diff --git a/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py b/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py index 877646ee48e..7e78f3c2eac 100644 --- a/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py +++ b/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py @@ -680,11 +680,11 @@ def kernel( ).make_participants() # Tensor memory dealloc barrier init tmem = utils.TmemAllocator( - storage.tmem_holding_buf, + storage.tmem_holding_buf.ptr, barrier_for_retrieve=self.tmem_alloc_barrier, allocator_warp_id=self.correction_warp_ids[0], is_two_cta=True, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, ) tmem.allocate(self.tmem_alloc_cols) tmem.wait_for_alloc()