# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. # [2025-07-04] Version in Cute-DSL, for Hopper and Blackwell. You'll need install nvidia-cutlass-dsl==4.2.0. # Supported features: # - BF16 & FP16 dtype # - noncausal & causal attention # - MHA, GQA, MQA # - hdim 64, 96, 128. # - (hdim_qk, hdim_v) = (192, 128) for Blackwell (i.e. DeepSeek shape) # - varlen # - sliding window # - bwd pass for Ampere (will also run on Hopper/Blackwell, but will be slow) # Features not supported yet: # - split (i.e. FlashDecoding) # - tuned block sizes # - paged KV # - append KV to existing KV cache # - FP8 # - bwd pass optimized for Hopper/Blackwell import math from typing import Optional, Tuple, Callable import torch from sglang.jit_kernel.utils import cache_once import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import sglang.jit_kernel.flash_attention.cute.utils as utils from .cute_dsl_utils import to_cute_tensor from .flash_fwd import FlashAttentionForwardSm90 from .flash_fwd_sm100 import FlashAttentionForwardSm100 from .flash_bwd_preprocess import FlashAttentionBackwardPreprocess from .flash_bwd import FlashAttentionBackwardSm80 from .flash_bwd_sm90 import FlashAttentionBackwardSm90 from .flash_bwd_sm100 import FlashAttentionBackwardSm100 from .flash_bwd_postprocess import FlashAttentionBackwardPostprocess from .flash_fwd_combine import FlashAttentionForwardCombine from .block_sparsity import ( BlockSparseTensorsTorch, to_cute_block_sparse_tensors, normalize_block_sparse_tensors, get_block_sparse_expected_shapes, get_block_sparse_expected_shapes_bwd, get_block_sparse_broadcast_pattern, ) @cache_once def _get_device_capability(): """Cached device capability check.""" return torch.cuda.get_device_capability()[0] def maybe_contiguous(x): return x.contiguous() if x is not None and x.stride(-1) != 1 else x def _validate_tensor(t, name, expected_shape, expected_dtype, expected_device): assert t.shape == expected_shape, f"{name} shape {t.shape} != expected {expected_shape}" assert t.dtype == expected_dtype, f"{name} dtype {t.dtype} != expected {expected_dtype}" assert t.device == expected_device, f"{name} device {t.device} != expected {expected_device}" assert t.is_cuda, f"{name} must be on CUDA" torch2cute_dtype_map = { torch.float16: cutlass.Float16, torch.bfloat16: cutlass.BFloat16, torch.float32: cutlass.Float32, } def num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, max_splits): # If num_n_blocks is too small, use 1 split. For example, we never split for hdim = 128 and seqlen_k = 512. if num_n_blocks <= 4: return 1 # NOTE: We should revisit this heuristic after persistence is supported for split KV. # Sometimes, it's ideal to over-schedule splits for better efficiency. return min(num_SMs // total_mblocks, max_splits, num_n_blocks) def _flash_attn_fwd( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_k: Optional[torch.Tensor] = None, seqused_q: Optional[torch.Tensor] = None, seqused_k: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_k: Optional[int] = None, page_table: Optional[torch.Tensor] = None, softmax_scale: Optional[float] = None, causal: bool = False, softcap: Optional[float] = None, window_size_left: Optional[int] = None, window_size_right: Optional[int] = None, learnable_sink: Optional[torch.Tensor] = None, # m_block_size: int = 128, # n_block_size: int = 64, # num_threads: int = 128, m_block_size: int = 128, n_block_size: int = 128, num_threads: int = 384, num_splits: int = 1, pack_gqa: Optional[bool] = None, _compute_capability: Optional[int] = None, score_mod: Optional[Callable] = None, mask_mod: Optional[Callable] = None, block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None, return_lse: bool = False, out: Optional[torch.Tensor] = None, lse: Optional[torch.Tensor] = None, aux_tensors: Optional[list[torch.Tensor]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Forward pass for FlashAttention. Args: ... score_mod: A callable that takes the attention scores and applies a modification. mask_mod: A callable that takes token position information and selectively masks block_sparse_tensors: A tuple of tensors used for block sparsity. return_lse: Whether to return the log softmax of the attention scores. If set to True will always calculate out: Optional pre-allocated output tensor. If None, will be allocated internally. lse: Optional pre-allocated log-sum-exp tensor. If None, will be allocated when needed. aux_tensors: Some score_mods will want to read from global aux_tensors. This is how we thread them through to the inner kernel. """ q, k, v = [maybe_contiguous(t) for t in (q, k, v)] num_head, head_dim = q.shape[-2:] if cu_seqlens_q is None: batch_size, seqlen_q = q.shape[:2] total_q = batch_size * seqlen_q else: batch_size = cu_seqlens_q.shape[0] - 1 seqlen_q = None total_q = q.shape[0] if page_table is not None: assert cu_seqlens_k is None, "page_table is not supported with cu_seqlens_k" assert page_table.dtype == torch.int32, "page_table must be int32" assert page_table.stride(-1) == 1, "page_table must be contiguous in the last dimension" max_num_pages_per_seq = page_table.shape[1] assert page_table.shape == (batch_size, max_num_pages_per_seq) num_pages, page_size = k.shape[:2] seqlen_k = num_pages * page_size else: num_pages, page_size = None, None seqlen_k = k.shape[-3] num_head_kv = k.shape[-2] head_dim_v = v.shape[-1] if cu_seqlens_k is None: if page_table is None: assert k.shape == (batch_size, seqlen_k, num_head_kv, head_dim) assert v.shape == (batch_size, seqlen_k, num_head_kv, head_dim_v) else: assert k.shape == (num_pages, page_size, num_head_kv, head_dim) assert v.shape == (num_pages, page_size, num_head_kv, head_dim_v) else: assert k.shape == (seqlen_k, num_head_kv, head_dim) assert v.shape == (seqlen_k, num_head_kv, head_dim_v) assert cu_seqlens_k.shape == (batch_size + 1,), ( "cu_seqlens_k must have shape (batch_size + 1,)" ) if cu_seqlens_q is not None: assert cu_seqlens_q.shape == (batch_size + 1,), ( "cu_seqlens_q must have shape (batch_size + 1,)" ) assert seqused_q is None or seqused_q.shape == (batch_size,), ( "seqused_q must have shape (batch_size,)" ) assert seqused_k is None or seqused_k.shape == (batch_size,), ( "seqused_k must have shape (batch_size,)" ) assert q.dtype in [torch.float16, torch.bfloat16], "inputs must be float16 or bfloat16" assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype" for t in [cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k]: if t is not None: assert t.dtype == torch.int32, ( "cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k must be int32" ) assert t.stride(0) == 1, ( "cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k must be contiguous" ) if learnable_sink is not None: assert learnable_sink.shape == (num_head,) assert learnable_sink.dtype == torch.bfloat16, "learnable_sink must be bfloat16" assert all( t is None or t.is_cuda for t in ( q, k, v, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, page_table, learnable_sink, ) ), "inputs must be on CUDA device" assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv" assert head_dim <= 256, "head_dim must be less than or equal to 256" alignment = 16 // q.element_size() assert head_dim % alignment == 0, f"head_dim must be divisible by {alignment}" assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}" if softmax_scale is None: softmax_scale = 1.0 / math.sqrt(head_dim) if softcap == 0.0: softcap = None qhead_per_kvhead = num_head // num_head_kv if pack_gqa is None: pack_gqa = qhead_per_kvhead > 1 out_torch_dtype = q.dtype device = q.device q_batch_seqlen_shape = (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,) lse_shape = (batch_size, num_head, seqlen_q) if cu_seqlens_q is None else (num_head, total_q) requires_grad = q.requires_grad or k.requires_grad or v.requires_grad if out is None: out = torch.empty( *q_batch_seqlen_shape, num_head, head_dim_v, dtype=out_torch_dtype, device=device ) else: _validate_tensor(out, "out", (*q_batch_seqlen_shape, num_head, head_dim_v), out_torch_dtype, device) if lse is None: lse = ( torch.empty(lse_shape, dtype=torch.float32, device=device) if requires_grad or return_lse else None ) elif lse is not None: _validate_tensor(lse, "lse", lse_shape, torch.float32, device) dtype = torch2cute_dtype_map[q.dtype] compute_capability = ( _get_device_capability() if _compute_capability is None else _compute_capability ) assert compute_capability in [9, 10, 11], "Unsupported compute capability. Supported: 9.x, 10.x, 11.x" use_block_sparsity = block_sparse_tensors is not None if mask_mod is None: if causal: window_size_right = 0 local = window_size_left is not None or window_size_right is not None if window_size_left is not None or window_size_right is not None: if window_size_left is None and window_size_right == 0: causal, local = True, False window_size_right = None else: causal, local = False, True else: causal, local = False, False current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) if compute_capability == 9: # TODO: tune block size according to hdim. if head_dim == head_dim_v == 128 and not causal and not local and not use_block_sparsity: n_block_size = 192 if compute_capability in [10, 11]: if ( pack_gqa and (128 % qhead_per_kvhead != 0) ): pack_gqa = False # TODO: fix GQA + SplitKV + non-varlen if pack_gqa and num_splits != 1 and cu_seqlens_q is None: pack_gqa = False if max_seqlen_q is None: max_seqlen_q = seqlen_q if cu_seqlens_q is None else total_q if max_seqlen_k is None: max_seqlen_k = seqlen_k seqlen_q_packgqa = max_seqlen_q * qhead_per_kvhead if compute_capability == 10: q_stage = 2 if seqlen_q_packgqa > m_block_size else 1 else: q_stage = 1 if num_splits < 1: m_block_size_effective = q_stage * m_block_size seqlen_k_loaded = max_seqlen_k if not local else max(0, min(max_seqlen_k, window_size_right + window_size_left + 1 + m_block_size)) num_n_blocks = (seqlen_k_loaded + n_block_size - 1) // n_block_size num_m_blocks = (seqlen_q_packgqa + m_block_size_effective - 1) // m_block_size_effective total_mblocks = batch_size * num_head_kv * num_m_blocks num_splits = num_splits_heuristic( total_mblocks, torch.cuda.get_device_properties(device).multi_processor_count, num_n_blocks, 128, ) is_split_kv = num_splits > 1 if is_split_kv: out_partial = torch.empty(num_splits, *q_batch_seqlen_shape, num_head, head_dim_v, dtype=torch.float32, device=device) lse_partial = torch.empty(num_splits, *lse_shape, dtype=torch.float32, device=device) # hash score and mask mods for compile cache score_mod_hash = utils.hash_callable(score_mod) if score_mod is not None else False mask_mod_hash = utils.hash_callable(mask_mod) if mask_mod is not None else False if softcap is not None: assert score_mod is None, "softcap and score_mod cannot be used together" score_mod = utils.create_softcap_scoremod(softcap) is_varlen = ( cu_seqlens_q is not None or cu_seqlens_k is not None or seqused_q is not None or seqused_k is not None ) if mask_mod is not None: if is_varlen: raise NotImplementedError( "mask_mod with aux_tensors is not yet supported for varlen sequences. This will be fixed in a future PR." ) if use_block_sparsity: if is_varlen: raise NotImplementedError( "Block sparsity is not yet supported for varlen sequences. This will be fixed in a future PR." ) # NB: pack_gqa requires block sparse head dim == 1 (broadcasted) if pack_gqa and block_sparse_tensors.mask_block_cnt.shape[1] != 1: pack_gqa = False if is_split_kv: raise NotImplementedError( "Block sparsity is not yet supported with SplitKV. TODO: partition sparse block lists per split." ) # See get_broadcast_dims for why this is needed in compile key block_sparse_broadcast_pattern = None normalized_block_sparse_tensors = None if block_sparse_tensors is not None: if seqlen_q is None: raise ValueError("Block sparsity requires fixed-length sequences (seqlen_q must be known).") expected_count_shape, expected_index_shape = get_block_sparse_expected_shapes( batch_size, num_head, seqlen_q, seqlen_k, m_block_size, n_block_size, q_stage, ) normalized_block_sparse_tensors = normalize_block_sparse_tensors( block_sparse_tensors, expected_count_shape=expected_count_shape, expected_index_shape=expected_index_shape, ) block_sparse_broadcast_pattern = get_block_sparse_broadcast_pattern( normalized_block_sparse_tensors ) compile_key = ( dtype, head_dim, head_dim_v, qhead_per_kvhead, causal, score_mod_hash, mask_mod_hash, use_block_sparsity, block_sparse_broadcast_pattern, len(aux_tensors) if aux_tensors is not None else 0, lse is None, cu_seqlens_q is None, cu_seqlens_k is None, seqused_q is None, seqused_k is None, page_table is not None, window_size_left is not None, window_size_right is not None, learnable_sink is not None, m_block_size, n_block_size, q_stage, num_threads, is_split_kv, pack_gqa, compute_capability, page_size not in [None, 128], # paged KV non-TMA ) if compile_key not in _flash_attn_fwd.compile_cache: ( cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor, learnable_sink_tensor, ) = [ to_cute_tensor(t, assumed_align=4, leading_dim=0) if t is not None else None for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink) ] page_table_tensor = ( to_cute_tensor(page_table, assumed_align=4, leading_dim=1) if page_table is not None else None ) q_tensor, k_tensor, v_tensor, o_tensor = [ to_cute_tensor(t) for t in (q, k, v, out if not is_split_kv else out_partial) ] if is_split_kv: lse_tensor = to_cute_tensor(lse_partial, assumed_align=4) elif lse is not None: lse_tensor = to_cute_tensor(lse, assumed_align=4) else: lse_tensor = None sparse_tensors = None if normalized_block_sparse_tensors is not None: sparse_tensors = to_cute_block_sparse_tensors(normalized_block_sparse_tensors) 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] if compute_capability == 9: assert page_table is None, "paged KV not supported on SM 9.0" assert not is_split_kv, "SplitKV not supported on SM 9.0" # fa_fwd = FlashAttentionForwardSm80( fa_fwd = FlashAttentionForwardSm90( dtype, head_dim, head_dim_v, qhead_per_kvhead, is_causal=causal, is_local=local, pack_gqa=pack_gqa, tile_m=m_block_size, tile_n=n_block_size, # num_stages=1, num_stages=2, num_threads=num_threads, Q_in_regs=False, intra_wg_overlap=True, mma_pv_is_rs=True, mask_mod=mask_mod, score_mod=score_mod, has_aux_tensors=aux_tensors is not None, ) elif compute_capability in [10, 11]: fa_fwd = FlashAttentionForwardSm100( head_dim, head_dim_v, qhead_per_kvhead=qhead_per_kvhead, is_causal=causal, is_local=local, is_split_kv=is_split_kv, pack_gqa=pack_gqa, m_block_size=m_block_size, n_block_size=n_block_size, q_stage=q_stage, is_persistent=not causal and not local and cu_seqlens_q is None and seqused_q is None and not is_split_kv, score_mod=score_mod, mask_mod=mask_mod, has_aux_tensors=aux_tensors is not None, paged_kv_non_tma=page_size not in [None, 128], is_varlen_q=cu_seqlens_q is not None or seqused_q is not None, ) else: raise ValueError( f"Unsupported compute capability: {compute_capability}. Supported: 9.x, 10.x, 11.x" ) # TODO: check @can_implement _flash_attn_fwd.compile_cache[compile_key] = cute.compile( fa_fwd, q_tensor, k_tensor, v_tensor, o_tensor, lse_tensor, softmax_scale, current_stream, cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor, page_table_tensor, window_size_left, window_size_right, learnable_sink_tensor, sparse_tensors, cute_aux_tensors, options="--enable-tvm-ffi", ) _flash_attn_fwd.compile_cache[compile_key]( q, k, v, out if not is_split_kv else out_partial, lse_partial if is_split_kv else lse, softmax_scale, current_stream, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, page_table, window_size_left, window_size_right, learnable_sink, normalized_block_sparse_tensors, aux_tensors, ) if is_split_kv: _flash_attn_fwd_combine( out_partial, lse_partial.transpose(-1, -2), out, lse.transpose(-1, -2) if lse is not None else None, cu_seqlens_q, seqused_q, ) return out, lse _flash_attn_fwd.compile_cache = {} def _flash_attn_bwd( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor, dout: torch.Tensor, lse: torch.Tensor, softmax_scale: Optional[float] = None, causal: bool = False, softcap: float = 0.0, window_size_left: Optional[int] = None, window_size_right: Optional[int] = None, m_block_size: int = 64, n_block_size: int = 128, num_threads: int = 256, pack_gqa: bool = False, num_stages_Q: int = 2, num_stages_dO: int = 2, SdP_swapAB: bool = False, dKV_swapAB: bool = False, dQ_swapAB: bool = False, AtomLayoutMSdP: int = 2, AtomLayoutNdKV: int = 2, AtomLayoutMdQ: int = 2, V_in_regs: bool = False, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_k: Optional[torch.Tensor] = None, seqused_q: Optional[torch.Tensor] = None, seqused_k: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_k: Optional[int] = None, deterministic: bool = False, dq: Optional[torch.Tensor] = None, dk: Optional[torch.Tensor] = None, dv: Optional[torch.Tensor] = None, score_mod: Optional[Callable] = None, score_mod_bwd: Optional[Callable] = None, mask_mod: Optional[Callable] = None, aux_tensors: Optional[list[torch.Tensor]] = None, block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: compute_capability = _get_device_capability() assert compute_capability in [9, 10, 11], "Unsupported compute capability. Supported: 9.x, 10.x, 11.x" if compute_capability == 9: m_block_size = 80 if not causal else 64 n_block_size = 128 num_stages_Q = 2 num_stages_dO = 2 num_stages_PdS = 2 SdP_swapAB = True dKV_swapAB = False dQ_swapAB = not causal AtomLayoutMSdP = 1 AtomLayoutNdKV = 2 AtomLayoutMdQ = 1 cluster_size = 1 assert window_size_left is None and window_size_right is None, "local not supported yet on 9.x" else: m_block_size = 128 n_block_size = 128 dQ_swapAB = False dKV_swapAB = False AtomLayoutMdQ = 1 AtomLayoutNdKV = 1 # TODO: support cluster size 2 cluster_size = 1 q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = [ maybe_contiguous(t) for t in (q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) ] num_head, head_dim = q.shape[-2:] if cu_seqlens_q is None: batch_size, seqlen_q = q.shape[:2] total_q = batch_size * seqlen_q else: batch_size = cu_seqlens_q.shape[0] - 1 total_q = q.shape[0] seqlen_q = max_seqlen_q if max_seqlen_q is not None else total_q if cu_seqlens_k is None: batch_size, seqlen_k = k.shape[:2] total_k = batch_size * seqlen_k else: batch_size = cu_seqlens_k.shape[0] - 1 total_k = k.shape[0] seqlen_k = max_seqlen_k if max_seqlen_k is not None else total_k num_head_kv = k.shape[-2] head_dim_v = v.shape[-1] if causal: window_size_right = 0 local = window_size_left is not None or window_size_right is not None if local: if window_size_left is None and window_size_right == 0: causal, local = True, False window_size_right = None else: causal, local = False, True use_block_sparsity = block_sparse_tensors is not None # SM90 block-sparse backward: tile_m=64 is the GCD between a m_block_size that fits, # the base block_m of 128 from forward, and block-sparse size for subtiling. if compute_capability == 9 and use_block_sparsity: m_block_size = 64 # dQ_swapAB tuning: use False when m_block_size=64 (same as causal case) dQ_swapAB = False # NB: this could be derived from the block_sparse_tensors but for now we hardcode it to 2 subtile_factor = 2 sparse_block_size_q = subtile_factor * m_block_size seqlen_q_rounded = (seqlen_q + m_block_size - 1) // m_block_size * m_block_size seqlen_k_rounded = (seqlen_k + n_block_size - 1) // n_block_size * n_block_size if cu_seqlens_k is None: assert k.shape == (batch_size, seqlen_k, num_head_kv, head_dim) assert v.shape == (batch_size, seqlen_k, num_head_kv, head_dim_v) else: assert k.shape == (total_k, num_head_kv, head_dim) assert v.shape == (total_k, num_head_kv, head_dim_v) assert cu_seqlens_k.shape == (batch_size + 1,), ( "cu_seqlens_k must have shape (batch_size + 1,)" ) if cu_seqlens_q is not None: assert cu_seqlens_q.shape == (batch_size + 1,), ( "cu_seqlens_q must have shape (batch_size + 1,)" ) assert out.shape == (total_q, num_head, head_dim_v) assert dout.shape == (total_q, num_head, head_dim_v) assert lse.shape == (num_head, total_q), "lse must have shape (num_head, total_q)" else: assert out.shape == (batch_size, seqlen_q, num_head, head_dim_v) assert dout.shape == (batch_size, seqlen_q, num_head, head_dim_v) assert lse.shape == (batch_size, num_head, seqlen_q), ( "lse must have shape (batch_size, num_head, seqlen_q)" ) assert q.dtype in [torch.float16, torch.bfloat16], "inputs must be float16 or bfloat16" assert q.dtype == k.dtype == v.dtype == out.dtype == dout.dtype, ( "inputs must have the same dtype" ) for t in [cu_seqlens_q, cu_seqlens_k]: if t is not None: assert t.dtype == torch.int32, "cu_seqlens_q, cu_seqlens_k must be int32" assert lse.dtype == torch.float32, "lse must be float32" assert all( t is None or t.is_cuda for t in (q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k) ), "inputs must be on CUDA device" assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv" assert head_dim <= 256, "head_dim must be less than or equal to 256" alignment = 16 // q.element_size() assert head_dim % alignment == 0, f"head_dim must be divisible by {alignment}" assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}" if softmax_scale is None: softmax_scale = 1.0 / math.sqrt(head_dim) qhead_per_kvhead = num_head // num_head_kv if pack_gqa is None: pack_gqa = qhead_per_kvhead > 1 # pack_gqa backward not yet supported in bwd pack_gqa = False if compute_capability not in [10, 11]: assert deterministic is False, "bwd deterministic only supported for sm100/sm110 for now" if score_mod is not None: assert score_mod_bwd is not None, "score_mod_bwd is required when score_mod is provided" assert softcap == 0.0, "softcap and score_mod are mutually exclusive (different log2 scaling)" assert cu_seqlens_q is None and cu_seqlens_k is None, ( "varlen + score_mod not supported in bwd yet" ) device = q.device out_torch_dtype = q.dtype if dq is None: dq = torch.empty_like(q) else: _validate_tensor(dq, "dq", q.shape, out_torch_dtype, device) if dk is None: dk = torch.empty_like(k) else: _validate_tensor(dk, "dk", k.shape, out_torch_dtype, device) if dv is None: dv = torch.empty_like(v) else: _validate_tensor(dv, "dv", v.shape, out_torch_dtype, device) head_dim_rounded = (head_dim + 32 - 1) // 32 * 32 if cu_seqlens_q is None: dq_accum = torch.empty( batch_size, num_head, seqlen_q_rounded * head_dim_rounded, dtype=torch.float32, device=device, ) dpsum = torch.empty( batch_size, num_head, seqlen_q_rounded, dtype=torch.float32, device=device ) lse_log2 = torch.empty( batch_size, num_head, seqlen_q_rounded, dtype=torch.float32, device=device ) else: total_q_rounded_padded = ( (total_q + cu_seqlens_q.shape[0] * m_block_size - 1) // m_block_size * m_block_size ) dq_accum = torch.empty( num_head, total_q_rounded_padded * head_dim_rounded, dtype=torch.float32, device=device ) dpsum = torch.empty(num_head, total_q_rounded_padded, dtype=torch.float32, device=device) lse_log2 = torch.empty(num_head, total_q_rounded_padded, dtype=torch.float32, device=device) dKV_postprocess = qhead_per_kvhead > 1 if dKV_postprocess: head_dim_v_rounded = (head_dim_v + 32 - 1) // 32 * 32 if cu_seqlens_k is None: num_n_blocks = seqlen_k_rounded // n_block_size if cluster_size == 2 and num_n_blocks % cluster_size != 0: seqlen_k_rounded = seqlen_k_rounded + n_block_size dk_accum = torch.zeros( batch_size, num_head_kv, seqlen_k_rounded * head_dim_rounded, dtype=torch.float32, device=device, ) dv_accum = torch.zeros( batch_size, num_head_kv, seqlen_k_rounded * head_dim_v_rounded, dtype=torch.float32, device=device, ) else: total_k_rounded_padded = ( (total_k + cu_seqlens_k.shape[0] * n_block_size - 1) // n_block_size * n_block_size ) num_n_blocks = total_k_rounded_padded // n_block_size if cluster_size == 2 and num_n_blocks % cluster_size != 0: total_k_rounded_padded = total_k_rounded_padded + n_block_size dk_accum = torch.zeros( num_head_kv, total_k_rounded_padded * head_dim_rounded, dtype=torch.float32, device=device, ) dv_accum = torch.zeros( num_head_kv, total_k_rounded_padded * head_dim_v_rounded, dtype=torch.float32, device=device, ) dtype = torch2cute_dtype_map[q.dtype] current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) if deterministic: dQ_semaphore = torch.zeros(batch_size, num_head, seqlen_q_rounded // m_block_size, 1, dtype=torch.int32, device="cuda") else: dQ_semaphore = None if deterministic and qhead_per_kvhead > 1: dK_semaphore = torch.zeros(batch_size, num_head_kv, seqlen_k_rounded // n_block_size, 2, dtype=torch.int32, device="cuda") dV_semaphore = torch.zeros(batch_size, num_head_kv, seqlen_k_rounded // n_block_size, 2, dtype=torch.int32, device="cuda") else: dK_semaphore = None dV_semaphore = None # Preprocess kernel: compute (o * dout).sum(dim=-1), lse * log2_e, and zero out dq_accum. compile_key_pre = ( compute_capability, dtype, head_dim_v, m_block_size, num_threads, cu_seqlens_q is None, seqused_q is None, ) if compile_key_pre not in _flash_attn_bwd.compile_cache_pre: o_tensor, do_tensor = [to_cute_tensor(t) for t in (out, dout)] dq_accum_tensor, dpsum_tensor, lse_log2_tensor = [ to_cute_tensor(t) for t in (dq_accum, dpsum, lse_log2) ] lse_tensor = to_cute_tensor(lse, assumed_align=4) cu_seqlens_q_tensor, seqused_q_tensor = [ to_cute_tensor(t, assumed_align=4) if t is not None else None for t in (cu_seqlens_q, seqused_q) ] arch = compute_capability * 10 fa_bwd_pre = FlashAttentionBackwardPreprocess( dtype, head_dim_v, arch, m_block_size, num_threads=num_threads, ) # TODO: check @can_implement _flash_attn_bwd.compile_cache_pre[compile_key_pre] = cute.compile( fa_bwd_pre, o_tensor, do_tensor, dpsum_tensor, lse_tensor, lse_log2_tensor, dq_accum_tensor, cu_seqlens_q_tensor, seqused_q_tensor, current_stream, options="--enable-tvm-ffi", ) _flash_attn_bwd.compile_cache_pre[compile_key_pre]( out, dout, dpsum, lse, lse_log2, dq_accum, cu_seqlens_q, seqused_q, current_stream, ) # NB num_threads application for 3 kernels # There are pre, main, post processing kernels, currenlty num_threads is only actually # used for the pre proc, and then we hard code to 384 for the main and post proc, and we do # before cache key gen num_threads = 384 # Backward kernel: compute dk, dv, dq_accum. score_mod_hash = utils.hash_callable(score_mod) if score_mod else False 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 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] block_sparse_broadcast_pattern = None normalized_block_sparse_tensors = None if block_sparse_tensors is not None: expected_count_shape, expected_index_shape = get_block_sparse_expected_shapes_bwd( batch_size, num_head, seqlen_q, seqlen_k, m_block_size, n_block_size, subtile_factor, ) normalized_block_sparse_tensors = normalize_block_sparse_tensors( block_sparse_tensors, expected_count_shape=expected_count_shape, expected_index_shape=expected_index_shape, context="_flash_attn_bwd", hint=lambda: ( f"Backward expects Q-direction block-sparse tensors (q_mask_cnt/q_mask_idx, and optionally full_q_cnt/full_q_idx). " f"Regenerate the backward BlockMask with BLOCK_SIZE=({sparse_block_size_q}, {n_block_size}) " f"(sparse_block_size_q={sparse_block_size_q})." ), ) block_sparse_broadcast_pattern = get_block_sparse_broadcast_pattern( normalized_block_sparse_tensors ) if compute_capability == 9: compile_key = ( compute_capability, dtype, head_dim, head_dim_v, qhead_per_kvhead, causal, softcap != 0.0, m_block_size, n_block_size, num_threads, pack_gqa, num_stages_Q, num_stages_dO, SdP_swapAB, dKV_swapAB, dQ_swapAB, AtomLayoutMSdP, AtomLayoutNdKV, AtomLayoutMdQ, V_in_regs, cu_seqlens_q is None, cu_seqlens_k is None, seqused_q is None, seqused_k is None, score_mod_hash, score_mod_bwd_hash, mask_mod_hash, num_aux_tensors, use_block_sparsity, block_sparse_broadcast_pattern, ) else: compile_key = ( compute_capability, dtype, head_dim, head_dim_v, qhead_per_kvhead, causal, window_size_left is not None, window_size_right is not None, softcap != 0.0, m_block_size, n_block_size, num_threads, pack_gqa, cluster_size, deterministic, score_mod_hash, score_mod_bwd_hash, mask_mod_hash, num_aux_tensors, use_block_sparsity, block_sparse_broadcast_pattern, cu_seqlens_q is None, cu_seqlens_k is None, seqused_q is None, seqused_k is None, ) if compile_key not in _flash_attn_bwd.compile_cache: q_tensor, k_tensor, v_tensor, do_tensor, dq_tensor, dk_tensor, dv_tensor = [ to_cute_tensor(t) for t in (q, k, v, dout, dq, dk, dv) ] dq_accum_tensor, dpsum_tensor, lse_log2_tensor = [ to_cute_tensor(t) for t in (dq_accum, dpsum, lse_log2) ] if dKV_postprocess: dk_accum_tensor, dv_accum_tensor = [ to_cute_tensor(t) for t in (dk_accum, dv_accum) ] cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor = [ to_cute_tensor(t, assumed_align=4) if t is not None else None for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) ] dQ_semaphore_tensor, dK_semaphore_tensor, dV_semaphore_tensor = [ utils.convert_from_dlpack_leading_static(t.detach(), leading_dim=3, alignment=4, stride_order=t.dim_order()) if t is not None else None for t in (dQ_semaphore, dK_semaphore, dV_semaphore) ] fa_bwd_sm80 = FlashAttentionBackwardSm80( dtype, head_dim, head_dim_v, qhead_per_kvhead, m_block_size, n_block_size, num_stages_Q, num_stages_dO, num_threads, pack_gqa, causal, SdP_swapAB, dKV_swapAB, dQ_swapAB, AtomLayoutMSdP, AtomLayoutNdKV, AtomLayoutMdQ, V_in_regs=V_in_regs, ) if compute_capability == 9: fa_bwd_obj = FlashAttentionBackwardSm90( dtype, head_dim, head_dim_v, qhead_per_kvhead, causal, m_block_size, n_block_size, num_stages_Q, num_stages_dO, num_stages_PdS, SdP_swapAB, dKV_swapAB, dQ_swapAB, AtomLayoutMSdP, AtomLayoutNdKV, AtomLayoutMdQ, num_threads, V_in_regs=V_in_regs, score_mod=score_mod, score_mod_bwd=score_mod_bwd, mask_mod=mask_mod, has_aux_tensors=aux_tensors is not None, subtile_factor=subtile_factor, ) else: fa_bwd_obj = FlashAttentionBackwardSm100( head_dim, head_dim_v, is_causal=causal, is_local=local, qhead_per_kvhead=qhead_per_kvhead, # tile_m=m_block_size, # tile_n=n_block_size, cluster_size=cluster_size, # cluster_size=1, deterministic=deterministic, score_mod=score_mod, score_mod_bwd=score_mod_bwd, mask_mod=mask_mod, has_aux_tensors=aux_tensors is not None, subtile_factor=subtile_factor, ) # Block sparse tensors for backward use Q-direction indexing (transposed from forward). # sparse_block_size_q = subtile_factor * tile_m matches BlockMask granularity. sparse_tensors_compile = None if normalized_block_sparse_tensors is not None: sparse_tensors_compile = to_cute_block_sparse_tensors(normalized_block_sparse_tensors) # TODO: check @can_implement _flash_attn_bwd.compile_cache[compile_key] = cute.compile( fa_bwd_obj, q_tensor, k_tensor, v_tensor, do_tensor, lse_log2_tensor, dpsum_tensor, dq_accum_tensor, dk_tensor if not dKV_postprocess else dk_accum_tensor, dv_tensor if not dKV_postprocess else dv_accum_tensor, softmax_scale, current_stream, cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor, None, # softcap - not yet supported in backward window_size_left, window_size_right, dQ_semaphore_tensor, dK_semaphore_tensor, dV_semaphore_tensor, cute_aux_tensors, sparse_tensors_compile, options="--enable-tvm-ffi", ) _flash_attn_bwd.compile_cache[compile_key]( q, k, v, dout, lse_log2, dpsum, dq_accum, dk if not dKV_postprocess else dk_accum, dv if not dKV_postprocess else dv_accum, softmax_scale, current_stream, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, None, # softcap - not yet supported in backward window_size_left, window_size_right, dQ_semaphore, dK_semaphore, dV_semaphore, aux_tensors, normalized_block_sparse_tensors, ) num_threads = 256 if compute_capability == 9 else 128 arch = compute_capability * 10 # Postprocess kernel: convert dq_accum from float32 to dq in bf16/fp16 compile_key_post = ( compute_capability, dtype, head_dim, m_block_size, num_threads, AtomLayoutMdQ, dQ_swapAB, cu_seqlens_q is None, seqused_q is None, ) if compile_key_post not in _flash_attn_bwd.compile_cache_post: dq_accum_tensor = to_cute_tensor(dq_accum) dq_tensor = to_cute_tensor(dq) cu_seqlens_q_tensor, seqused_q_tensor = [ to_cute_tensor(t, assumed_align=4) if t is not None else None for t in (cu_seqlens_q, seqused_q) ] fa_bwd_post = FlashAttentionBackwardPostprocess( dtype, head_dim, arch, m_block_size, num_threads, AtomLayoutMdQ, dQ_swapAB ) # TODO: check @can_implement _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile( fa_bwd_post, dq_accum_tensor, dq_tensor, softmax_scale, cu_seqlens_q_tensor, seqused_q_tensor, current_stream, options="--enable-tvm-ffi", ) _flash_attn_bwd.compile_cache_post[compile_key_post]( dq_accum, dq, softmax_scale, cu_seqlens_q, seqused_q, current_stream, ) if dKV_postprocess: # Postprocess kernel: convert dk_accum & dv_accum from float32 to bf16/fp16 compile_key_post = ( compute_capability, dtype, head_dim, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB, cu_seqlens_k is None, seqused_k is None, ) if compile_key_post not in _flash_attn_bwd.compile_cache_post: dk_accum_tensor = to_cute_tensor(dk_accum) dk_tensor = to_cute_tensor(dk) cu_seqlens_k_tensor, seqused_k_tensor = [ to_cute_tensor(t, assumed_align=4) if t is not None else None for t in (cu_seqlens_k, seqused_k) ] arch = compute_capability * 10 fa_bwd_post = FlashAttentionBackwardPostprocess( dtype, head_dim, arch, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB ) # TODO: check @can_implement _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile( fa_bwd_post, dk_accum_tensor, dk_tensor, softmax_scale, cu_seqlens_k_tensor, seqused_k_tensor, current_stream, options="--enable-tvm-ffi", ) _flash_attn_bwd.compile_cache_post[compile_key_post]( dk_accum, dk, softmax_scale, cu_seqlens_k, seqused_k, current_stream, ) compile_key_post = ( compute_capability, dtype, head_dim_v, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB, cu_seqlens_k is None, seqused_k is None, ) if compile_key_post not in _flash_attn_bwd.compile_cache_post: dv_accum_tensor = to_cute_tensor(dv_accum) dv_tensor = to_cute_tensor(dv) cu_seqlens_k_tensor, seqused_k_tensor = [ to_cute_tensor(t, assumed_align=4) if t is not None else None for t in (cu_seqlens_k, seqused_k) ] arch = compute_capability * 10 fa_bwd_post = FlashAttentionBackwardPostprocess( dtype, head_dim_v, arch, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB ) # TODO: check @can_implement _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile( fa_bwd_post, dv_accum_tensor, dv_tensor, cutlass.Float32(1.0), cu_seqlens_k_tensor, seqused_k_tensor, current_stream, options="--enable-tvm-ffi", ) _flash_attn_bwd.compile_cache_post[compile_key_post]( dv_accum, dv, 1.0, cu_seqlens_k, seqused_k, current_stream, ) return dq, dk, dv _flash_attn_bwd.compile_cache_pre = {} _flash_attn_bwd.compile_cache = {} _flash_attn_bwd.compile_cache_post = {} class FlashAttnFunc(torch.autograd.Function): @staticmethod def forward( ctx, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, softmax_scale: Optional[float] = None, causal: bool = False, window_size: Tuple[Optional[int], Optional[int]] = (None, None), learnable_sink: Optional[torch.Tensor] = None, softcap: float = 0.0, num_splits: int = 1, pack_gqa: Optional[bool] = None, deterministic: bool = False, mask_mod: Optional[Callable] = None, full_block_cnt: Optional[torch.Tensor] = None, full_block_idx: Optional[torch.Tensor] = None, mask_block_cnt: Optional[torch.Tensor] = None, mask_block_idx: Optional[torch.Tensor] = None, ): # Only create block sparse tensors if at least one block sparse parameter is provided block_sparse_tensors = None if any(t is not None for t in [full_block_cnt, full_block_idx, mask_block_cnt, mask_block_idx]): block_sparse_tensors = BlockSparseTensorsTorch( full_block_cnt=full_block_cnt, full_block_idx=full_block_idx, mask_block_cnt=mask_block_cnt, mask_block_idx=mask_block_idx, ) out, lse = _flash_attn_fwd( q, k, v, softmax_scale=softmax_scale, causal=causal, window_size_left=window_size[0], window_size_right=window_size[1], learnable_sink=learnable_sink, softcap=softcap, num_splits=num_splits, pack_gqa=pack_gqa, mask_mod=mask_mod, block_sparse_tensors=block_sparse_tensors ) ctx.save_for_backward(q, k, v, out, lse) ctx.softmax_scale = softmax_scale ctx.causal = causal ctx.window_size = window_size ctx.softcap = softcap ctx.deterministic = deterministic return out, lse @staticmethod def backward(ctx, dout, *args): q, k, v, out, lse = ctx.saved_tensors dq, dk, dv = _flash_attn_bwd( q, k, v, out, dout, lse, ctx.softmax_scale, ctx.causal, ctx.softcap, window_size_left=ctx.window_size[0], window_size_right=ctx.window_size[1], deterministic=ctx.deterministic, ) return dq, dk, dv, *((None,) * 20) # Extra Nones is fine class FlashAttnVarlenFunc(torch.autograd.Function): @staticmethod def forward( ctx, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens_q: Optional[torch.Tensor], cu_seqlens_k: Optional[torch.Tensor], seqused_q: Optional[torch.Tensor] = None, seqused_k: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_k: Optional[int] = None, page_table: Optional[torch.Tensor] = None, softmax_scale: Optional[float] = None, causal: bool = False, window_size: Tuple[Optional[int], Optional[int]] = (None, None), learnable_sink: Optional[torch.Tensor] = None, softcap: float = 0.0, num_splits: int = 1, pack_gqa: Optional[bool] = None, deterministic: bool = False, score_mod: Optional[Callable] = None, aux_tensors: Optional[list] = None, ): out, lse = _flash_attn_fwd( q, k, v, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, page_table=page_table, softmax_scale=softmax_scale, causal=causal, window_size_left=window_size[0], window_size_right=window_size[1], learnable_sink=learnable_sink, softcap=softcap, num_splits=num_splits, pack_gqa=pack_gqa, score_mod=score_mod, aux_tensors=aux_tensors, return_lse=True, ) ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) ctx.softmax_scale = softmax_scale ctx.causal = causal ctx.window_size = window_size ctx.softcap = softcap ctx.deterministic = deterministic ctx.max_seqlen_q = max_seqlen_q ctx.max_seqlen_k = max_seqlen_k return out, lse @staticmethod def backward(ctx, dout, *args): q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors assert ctx.softcap == 0.0 dq, dk, dv = _flash_attn_bwd( q, k, v, out, dout, lse, ctx.softmax_scale, ctx.causal, ctx.softcap, window_size_left=ctx.window_size[0], window_size_right=ctx.window_size[1], cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, seqused_q=seqused_q, seqused_k=seqused_k, max_seqlen_q=ctx.max_seqlen_q, max_seqlen_k=ctx.max_seqlen_k, deterministic=ctx.deterministic, ) return dq, dk, dv, *((None,) * 20) def flash_attn_func( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, softmax_scale: Optional[float] = None, causal: bool = False, window_size: Tuple[Optional[int], Optional[int]] = (None, None), learnable_sink: Optional[torch.Tensor] = None, softcap: float = 0.0, num_splits: int = 1, pack_gqa: Optional[bool] = None, deterministic: bool = False, mask_mod: Optional[Callable] = None, full_block_cnt: Optional[torch.Tensor] = None, full_block_idx: Optional[torch.Tensor] = None, mask_block_cnt: Optional[torch.Tensor] = None, mask_block_idx: Optional[torch.Tensor] = None, ): return FlashAttnFunc.apply( q, k, v, softmax_scale, causal, window_size, learnable_sink, softcap, num_splits, pack_gqa, deterministic, mask_mod, full_block_cnt, full_block_idx, mask_block_cnt, mask_block_idx, ) def flash_attn_varlen_func( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_k: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_k: Optional[int] = None, seqused_q: Optional[torch.Tensor] = None, seqused_k: Optional[torch.Tensor] = None, page_table: Optional[torch.Tensor] = None, softmax_scale: Optional[float] = None, causal: bool = False, window_size: Tuple[Optional[int], Optional[int]] = (None, None), learnable_sink: Optional[torch.Tensor] = None, softcap: float = 0.0, num_splits: int = 1, pack_gqa: Optional[bool] = None, deterministic: bool = False, score_mod: Optional[Callable] = None, aux_tensors: Optional[list] = None, ): return FlashAttnVarlenFunc.apply( q, k, v, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, max_seqlen_q, max_seqlen_k, page_table, softmax_scale, causal, window_size, learnable_sink, softcap, num_splits, pack_gqa, deterministic, score_mod, aux_tensors, ) def _flash_attn_fwd_combine( out_partial: torch.Tensor, lse_partial: torch.Tensor, out: torch.Tensor, lse: Optional[torch.Tensor] = None, cu_seqlens: Optional[torch.Tensor] = None, seqused: Optional[torch.Tensor] = None, num_splits_dynamic_ptr: Optional[torch.Tensor] = None, semaphore_to_reset: Optional[torch.Tensor] = None, ) -> None: """Forward combine kernel for split attention computation. Combines partial outputs and log-sum-exp values from multiple splits of attention computation into final outputs. Args: out_partial: Partial outputs tensor (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim) if there's cu_seqlens lse_partial: Partial LSE tensor (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads) if there's cu_seqlens out: Output tensor (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim) if there's cu_seqlens lse: Output LSE tensor (batch, seqlen, nheads) or (total_q, nheads) if there's cu_seqlens. cu_seqlens: Cumulative sequence lengths for variable length sequences seqused: Used sequence lengths for each batch num_splits_dynamic_ptr: Dynamic number of splits per batch semaphore_to_reset: Semaphore for synchronization k_block_size: Block size for head dimension Returns: None """ # Input validation assert out_partial.dim() in [4, 5], "out_partial must have 4 or 5 dimensions" assert lse_partial.dim() in [3, 4], "lse_partial must have 3 or 4 dimensions" assert out_partial.dtype in [torch.float16, torch.bfloat16, torch.float32], ( "out_partial must be fp16, bf16, or fp32" ) assert lse_partial.dtype == torch.float32, "lse_partial must be fp32" assert out_partial.is_cuda and lse_partial.is_cuda, "tensors must be on CUDA device" assert out_partial.stride(-1) == 1, "out_partial must be contiguous in the last dimension" assert lse_partial.stride(-2) == 1, "lse_partial must be contiguous in the seqlen dimension" assert lse_partial.shape == out_partial.shape[:-1] # Determine if this is variable length based on dimensions is_varlen = out_partial.dim() == 4 # Validate output tensor shapes and types assert out.shape == out_partial.shape[1:], "out shape mismatch" if lse is not None: assert lse.shape == lse_partial.shape[1:], "lse shape mismatch" assert lse.dtype == torch.float32, "lse must be fp32" # Validate optional tensors for t, name in [ (cu_seqlens, "cu_seqlens"), (seqused, "seqused"), (num_splits_dynamic_ptr, "num_splits_dynamic_ptr"), ]: if t is not None: assert t.dtype == torch.int32, f"{name} must be int32" assert t.is_cuda, f"{name} must be on CUDA device" assert t.is_contiguous(), f"{name} must be contiguous" head_dim = out_partial.shape[-1] num_splits = out_partial.shape[0] assert num_splits <= 256 # If hdim is 96 or 192, it's faster to round them to 128 or 256 respectively # so that kBlockM is smaller and we have more parallelism. k_block_size = 64 if head_dim <= 64 else 128 # We want kBlockM to be as small as possible to maximize parallelism. # E.g., if hdim is 64, we want kBlockM to be 16 so that we can use 256 threads, each reading 4 elements (floats). m_block_size = 8 if k_block_size % 128 == 0 else (16 if k_block_size % 64 == 0 else 32) log_max_splits = max(math.ceil(math.log2(num_splits)), 4) if m_block_size == 8: # If kBlockM == 8 then the minimum number of splits is 32. # TODO: we can deal w this by using 128 threads instead log_max_splits = max(log_max_splits, 5) current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) # Create combine kernel configuration dtype = torch2cute_dtype_map[out.dtype] dtype_partial = torch2cute_dtype_map[out_partial.dtype] compile_key = ( dtype, dtype_partial, head_dim, m_block_size, k_block_size, log_max_splits, cu_seqlens is not None, seqused is not None, lse is not None, ) if compile_key not in _flash_attn_fwd_combine.compile_cache: out_partial_tensor = to_cute_tensor( out_partial, leading_dim=4 if not is_varlen else 3 ) lse_partial_tensor = to_cute_tensor( lse_partial, assumed_align=4, leading_dim=lse_partial.ndim - 2 ) out_tensor = to_cute_tensor(out, leading_dim=3 if not is_varlen else 2) lse_tensor = ( to_cute_tensor(lse, assumed_align=4, leading_dim=lse.ndim - 2) if lse is not None else None ) optional_tensors = [ to_cute_tensor(t, assumed_align=4, leading_dim=0) if t is not None else None for t in (cu_seqlens, seqused, num_splits_dynamic_ptr, semaphore_to_reset) ] cu_seqlens_tensor, seqused_tensor, num_splits_dynamic_tensor, semaphore_tensor = ( optional_tensors ) fa_combine = FlashAttentionForwardCombine( dtype=dtype, dtype_partial=dtype_partial, head_dim=head_dim, m_block_size=m_block_size, k_block_size=k_block_size, log_max_splits=log_max_splits, ) # Check if implementation is supported if not fa_combine.can_implement( dtype, dtype_partial, head_dim, m_block_size, k_block_size, log_max_splits, num_threads=256, ): raise RuntimeError( "FlashAttention combine kernel cannot be implemented with given parameters" ) _flash_attn_fwd_combine.compile_cache[compile_key] = cute.compile( fa_combine, out_partial_tensor, lse_partial_tensor, out_tensor, lse_tensor, cu_seqlens_tensor, seqused_tensor, num_splits_dynamic_tensor, semaphore_tensor, current_stream, options="--enable-tvm-ffi", ) _flash_attn_fwd_combine.compile_cache[compile_key]( out_partial, lse_partial, out, lse, cu_seqlens, seqused, num_splits_dynamic_ptr, semaphore_to_reset, current_stream, ) _flash_attn_fwd_combine.compile_cache = {} def flash_attn_combine( out_partial: torch.Tensor, lse_partial: torch.Tensor, out: Optional[torch.Tensor] = None, out_dtype: Optional[torch.dtype] = None, cu_seqlens: Optional[torch.Tensor] = None, seqused: Optional[torch.Tensor] = None, return_lse: bool = True, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Flash Attention combine function for split attention computation. Combines partial outputs and log-sum-exp values from multiple splits of attention computation into final outputs. This is the main user-facing interface for the combine kernel. Args: out_partial: Partial outputs tensor with shape: - (num_splits, batch_size, seqlen, num_heads, head_size) for regular batched input - (num_splits, total_q, num_heads, head_size) for variable length input lse_partial: Partial LSE tensor with shape: - (num_splits, batch_size, seqlen, num_heads) for regular batched input - (num_splits, total_q, num_heads) for variable length input out: Optional output tensor. If None, will be created automatically. out_dtype: Optional output dtype. If None, will use fp16/bf16 based on input. cu_seqlens: Cumulative sequence lengths for variable length sequences seqused: Used sequence lengths for each batch return_lse: Whether to return the combined LSE tensor. Default is True. Returns: Tuple of (out, lse) where: - out: Combined output tensor with shape (batch_size, seqlen, num_heads, head_size) or (total_q, num_heads, head_size) for varlen - lse: Combined log-sum-exp tensor with shape (batch_size, seqlen, num_heads) or (total_q, num_heads) for varlen. None if return_lse=False Note: This function expects the input tensors to be in the format produced by split attention computation, where the first dimension is num_splits. The permuting from user format to kernel format is now done inside the kernel. """ # Input validation assert out_partial.dim() in [4, 5], "out_partial must have 4 or 5 dimensions" assert lse_partial.dim() in [3, 4], "lse_partial must have 3 or 4 dimensions" assert out_partial.dtype == torch.float32, "out_partial must be fp32 (from accumulation)" assert lse_partial.dtype == torch.float32, "lse_partial must be fp32" # Determine if this is variable length based on dimensions is_varlen = out_partial.dim() == 4 if is_varlen: # Variable length: (num_splits, total_q, num_heads, head_size) num_splits, total_q, num_heads, head_size = out_partial.shape assert lse_partial.shape == (num_splits, total_q, num_heads), ( "lse_partial shape mismatch for varlen" ) batch_size = 1 # Treat as single batch for varlen seqlen = total_q else: # Regular batched: (num_splits, batch_size, seqlen, num_heads, head_size) num_splits, batch_size, seqlen, num_heads, head_size = out_partial.shape assert lse_partial.shape == (num_splits, batch_size, seqlen, num_heads), ( "lse_partial shape mismatch" ) # Determine output dtype if out_dtype is None: out_dtype = out_partial.dtype # Create output if not provided device = out_partial.device if out is None: if is_varlen: out = torch.empty(total_q, num_heads, head_size, dtype=out_dtype, device=device) else: out = torch.empty( batch_size, seqlen, num_heads, head_size, dtype=out_dtype, device=device ) # Create lse output only if requested if return_lse: if is_varlen: lse = torch.empty(num_heads, total_q, dtype=torch.float32, device=device).transpose( 0, 1 ) else: lse = torch.empty( batch_size, num_heads, seqlen, dtype=torch.float32, device=device ).transpose(1, 2) else: lse = None _flash_attn_fwd_combine( out_partial, lse_partial, out, lse, cu_seqlens, seqused, ) return out, lse