# 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. # [2025-07-04] Version in Cute-DSL, for Hopper and Blackwell. You'll need install nvidia-cutlass-dsl==4.2.0. # [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 import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack from . import utils from .flash_fwd import FlashAttentionForwardSm80, 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, ) def maybe_contiguous(x): return x.contiguous() if x is not None and x.stride(-1) != 1 else x 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, 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: expected_out_shape = (*q_batch_seqlen_shape, num_head, head_dim_v) assert out.shape == expected_out_shape, ( f"out tensor shape {out.shape} does not match expected shape {expected_out_shape}" ) assert out.dtype == out_torch_dtype, ( f"out tensor dtype {out.dtype} does not match expected dtype {out_torch_dtype}" ) assert out.device == device, ( f"out tensor device {out.device} does not match input device {device}" ) assert out.is_cuda, "out tensor must be on CUDA 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: assert lse.shape == lse_shape, ( f"lse tensor shape {lse.shape} does not match expected shape {lse_shape}" ) assert lse.dtype == torch.float32, ( f"lse tensor dtype {lse.dtype} does not match expected dtype torch.float32" ) assert lse.device == device, ( f"lse tensor device {lse.device} does not match input device {device}" ) assert lse.is_cuda, "lse tensor must be on CUDA device" dtype = torch2cute_dtype_map[q.dtype] ( cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor, learnable_sink_tensor, ) = [ from_dlpack(t.detach(), assumed_align=4).mark_layout_dynamic(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 = ( from_dlpack(page_table.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=1) if page_table is not None else None ) compute_capability = ( torch.cuda.get_device_capability()[0] if _compute_capability is None else _compute_capability ) assert compute_capability in [9, 10], "Unsupported compute capability. Supported: 9.x, 10.x" 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).") m_block_size_block = m_block_size if compute_capability == 10: # TODO: This multiplier should really be q_stage, wire up in later PR # 1 cta handles 2*tile_m row m_block_size_block = 2 * m_block_size expected_m_blocks = (seqlen_q + m_block_size_block - 1) // m_block_size_block expected_n_blocks = (seqlen_k + n_block_size - 1) // n_block_size block_sparse_tensors = normalize_block_sparse_tensors( block_sparse_tensors, expected_count_shape=(batch_size, num_head, expected_m_blocks), expected_index_shape=(batch_size, num_head, expected_m_blocks, expected_n_blocks), ) sparse_tensors = to_cute_block_sparse_tensors(block_sparse_tensors) use_block_sparsity = 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 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 == 10: # TODO: fix the varlen case if ( pack_gqa and (128 % qhead_per_kvhead != 0) or (cu_seqlens_q is not None or seqused_q is not None) ): 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 num_splits < 1: max_seqlen_k = seqlen_k if cu_seqlens_k is None else (cu_seqlens_k[1:] - cu_seqlens_k[:-1]).max().item() max_seqlen_q = seqlen_q if cu_seqlens_q is None else (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item() seqlen_q_packgqa = max_seqlen_q * qhead_per_kvhead 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 - 1) // m_block_size 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) q_tensor, k_tensor, v_tensor, o_tensor = [ from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.ndim - 1) for t in (q, k, v, out if not is_split_kv else out_partial) ] if is_split_kv: lse_tensor = from_dlpack(lse_partial.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=lse_partial.ndim - 1) elif lse is not None: lse_tensor = from_dlpack(lse.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=lse.ndim - 1) else: lse_tensor = None # 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 score_mod is not None: if is_varlen: raise NotImplementedError( "score_mod with aux_tensors is not yet supported for varlen sequences. This will be fixed in a future PR." ) 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 pack_gqa: raise NotImplementedError( "mask_mod with aux_tensors is not yet supported with pack_gqa=True. 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." ) if pack_gqa: raise NotImplementedError( "Block sparsity is not yet supported with pack_gqa=True. This will be fixed in a future PR." ) if is_split_kv: raise NotImplementedError( "Block sparsity is not yet supported with SplitKV. TODO: partition sparse block lists per split." ) cute_aux_tensors = None if aux_tensors is not None: cute_aux_tensors = [from_dlpack(buf).mark_layout_dynamic() for buf in aux_tensors] compile_key = ( dtype, head_dim, head_dim_v, qhead_per_kvhead, causal, score_mod_hash, mask_mod_hash, use_block_sparsity, 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, 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: 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 == 10: 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, 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" ) # 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, ) _flash_attn_fwd.compile_cache[compile_key]( 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, ) 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, 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, deterministic: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: compute_capability = torch.cuda.get_device_capability()[0] assert compute_capability in [9, 10], "Unsupported compute capability. Supported: 9.x, 10.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 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 seqlen_q = None total_q = q.shape[0] 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 seqlen_k = None total_k = k.shape[0] num_head_kv = k.shape[-2] head_dim_v = v.shape[-1] 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 if compute_capability == 10: pack_gqa = False # override for now if compute_capability != 10: assert deterministic is False, "bwd deterministic only supported for sm100 for now" device = q.device # TODO: check if this is the right rounding dq = torch.empty_like(q) dk = torch.empty_like(k) dv = torch.empty_like(v) head_dim_rounded = (head_dim + 32 - 1) // 32 * 32 if cu_seqlens_q is None: seqlen_q_rounded = (seqlen_q + m_block_size - 1) // m_block_size * m_block_size 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) if qhead_per_kvhead > 1: head_dim_v_rounded = (head_dim_v + 32 - 1) // 32 * 32 if cu_seqlens_k is None: seqlen_k_rounded = (seqlen_k + n_block_size - 1) // n_block_size * n_block_size 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] q_tensor, k_tensor, v_tensor, o_tensor, do_tensor, dq_tensor, dk_tensor, dv_tensor = [ from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.ndim - 1) for t in (q, k, v, out, dout, dq, dk, dv) ] lse_tensor = from_dlpack(lse.detach(), assumed_align=4).mark_layout_dynamic( leading_dim=lse.ndim - 1 ) dq_accum_tensor, dpsum_tensor, lse_log2_tensor = [ from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.ndim - 1) for t in (dq_accum, dpsum, lse_log2) ] if qhead_per_kvhead > 1: dk_accum_tensor, dv_accum_tensor = [ from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.ndim - 1) for t in (dk_accum, dv_accum) ] cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor = [ from_dlpack(t.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=t.ndim - 1) if t is not None else None for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) ] 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 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) ] current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) # 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) if compile_key_pre not in _flash_attn_bwd.compile_cache_pre: fa_bwd_pre = FlashAttentionBackwardPreprocess( dtype, head_dim_v, 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, ) _flash_attn_bwd.compile_cache_pre[compile_key_pre]( o_tensor, do_tensor, dpsum_tensor, lse_tensor, lse_log2_tensor, dq_accum_tensor, cu_seqlens_q_tensor, seqused_q_tensor, current_stream, ) # Backward kernel: compute dk, dv, dq_accum. 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, ) else: 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, cluster_size, deterministic, ) num_threads = 384 if compile_key not in _flash_attn_bwd.compile_cache: 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, ) else: fa_bwd_obj = FlashAttentionBackwardSm100( head_dim, head_dim_v, is_causal=causal, 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, ) # 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 qhead_per_kvhead == 1 else dk_accum_tensor, dv_tensor if qhead_per_kvhead == 1 else dv_accum_tensor, softmax_scale, current_stream, cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor, mdQ_semaphore=dQ_semaphore_tensor, mdK_semaphore=dK_semaphore_tensor, mdV_semaphore=dV_semaphore_tensor, ) _flash_attn_bwd.compile_cache[compile_key]( q_tensor, k_tensor, v_tensor, do_tensor, lse_log2_tensor, dpsum_tensor, dq_accum_tensor, dk_tensor if qhead_per_kvhead == 1 else dk_accum_tensor, dv_tensor if qhead_per_kvhead == 1 else dv_accum_tensor, softmax_scale, current_stream, cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor, mdQ_semaphore=dQ_semaphore_tensor, mdK_semaphore=dK_semaphore_tensor, mdV_semaphore=dV_semaphore_tensor, ) num_threads = 256 if compute_capability == 9 else 128 # Postprocess kernel: convert dq_accum from float32 to dq in bf16/fp16 compile_key_post = (dtype, head_dim, m_block_size, num_threads, AtomLayoutMdQ, dQ_swapAB) if compile_key_post not in _flash_attn_bwd.compile_cache_post: arch = compute_capability * 10 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, ) _flash_attn_bwd.compile_cache_post[compile_key_post]( dq_accum_tensor, dq_tensor, softmax_scale, cu_seqlens_q_tensor, seqused_q_tensor, current_stream, ) if qhead_per_kvhead > 1: # Postprocess kernel: convert dk_accum & dv_accum from float32 to bf16/fp16 compile_key_post = (dtype, head_dim, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB) if compile_key_post not in _flash_attn_bwd.compile_cache_post: fa_bwd_post = FlashAttentionBackwardPostprocess( dtype, head_dim, 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, ) _flash_attn_bwd.compile_cache_post[compile_key_post]( dk_accum_tensor, dk_tensor, softmax_scale, cu_seqlens_k_tensor, seqused_k_tensor, current_stream, ) compile_key_post = ( dtype, head_dim_v, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB, ) if compile_key_post not in _flash_attn_bwd.compile_cache_post: fa_bwd_post = FlashAttentionBackwardPostprocess( dtype, head_dim_v, 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, ) _flash_attn_bwd.compile_cache_post[compile_key_post]( dv_accum_tensor, dv_tensor, cutlass.Float32(1.0), cu_seqlens_k_tensor, seqused_k_tensor, 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, 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, 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, ): out, lse = _flash_attn_fwd( q, k, v, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_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, ) 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 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 seqused_q == seqused_k == None assert ctx.softcap == 0.0 dq, dk, dv = _flash_attn_bwd( q, k, v, out, dout, lse, ctx.softmax_scale, ctx.causal, ctx.softcap, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, seqused_q=seqused_q, seqused_k=seqused_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, 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, ): return FlashAttnVarlenFunc.apply( q, k, v, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, page_table, softmax_scale, causal, window_size, learnable_sink, softcap, num_splits, pack_gqa, deterministic, ) 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) # Convert to cute tensors (using kernel-formatted tensors) out_partial_tensor = from_dlpack(out_partial.detach(), assumed_align=16).mark_layout_dynamic( leading_dim=4 if not is_varlen else 3 ) lse_partial_tensor = from_dlpack(lse_partial.detach(), assumed_align=4).mark_layout_dynamic( leading_dim=lse_partial.ndim - 2 ) out_tensor = from_dlpack(out.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=3 if not is_varlen else 2) lse_tensor = ( from_dlpack(lse.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=lse.ndim - 2) if lse is not None else None ) optional_tensors = [ from_dlpack(t.detach(), assumed_align=4).mark_layout_dynamic(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 ) 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: 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, ) _flash_attn_fwd_combine.compile_cache[compile_key]( out_partial_tensor, lse_partial_tensor, out_tensor, lse_tensor, cu_seqlens_tensor, seqused_tensor, num_splits_dynamic_tensor, semaphore_tensor, 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