diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 1862e8707..81a020d43 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -92,7 +92,7 @@ FetchContent_Populate(repo-flashinfer) FetchContent_Declare( repo-flash-attention GIT_REPOSITORY https://github.com/sgl-project/sgl-attn - GIT_TAG f20a52329482ddca4a627b2f028f88c2959ee299 + GIT_TAG f866ec34002250e74c8bbcbcffa0e1ae71300b2d GIT_SHALLOW OFF ) FetchContent_Populate(repo-flash-attention) diff --git a/sgl-kernel/python/sgl_kernel/_fa4_interface.py b/sgl-kernel/python/sgl_kernel/_fa4_interface.py index ffd0563f5..1b6ab5305 100644 --- a/sgl-kernel/python/sgl_kernel/_fa4_interface.py +++ b/sgl-kernel/python/sgl_kernel/_fa4_interface.py @@ -1,4 +1,4 @@ -# Adapted from https://github.com/Dao-AILab/flash-attention/blob/54d8aa6751fc9d5f0357854079261913d5df1f9d/flash_attn/cute/interface.py +# Adapted from https://github.com/Dao-AILab/flash-attention/blob/5d4c9537a1e0f1adcc3e4c3e11ae46fe94a18b11/flash_attn/cute/interface.py # Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. # [2025-10-14] Version in Cute-DSL, for Hopper and Blackwell. You'd need to install nvidia-cutlass-dsl==4.2.1. @@ -9,6 +9,7 @@ import gc import logging import math import os +from functools import lru_cache from typing import Callable, Optional, Tuple logger = logging.getLogger(__name__) @@ -20,14 +21,50 @@ import cutlass.cute as cute import torch from cutlass.cute.runtime import from_dlpack from flash_attn_origin.cute import utils +from flash_attn_origin.cute.block_sparsity import ( + BlockSparseTensorsTorch, + get_block_sparse_expected_shapes, + normalize_block_sparse_tensors, + to_cute_block_sparse_tensors, +) from flash_attn_origin.cute.flash_fwd import FlashAttentionForwardSm90 +from flash_attn_origin.cute.flash_fwd_combine import FlashAttentionForwardCombine from flash_attn_origin.cute.flash_fwd_sm100 import FlashAttentionForwardSm100 +@lru_cache(maxsize=None) +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" + + +def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False): + """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1.""" + tensor = from_dlpack(t.detach(), assumed_align=assumed_align, enable_tvm_ffi=True) + if fully_dynamic: + return tensor.mark_layout_dynamic() + if leading_dim == -1: + leading_dim = t.ndim - 1 + return tensor.mark_layout_dynamic(leading_dim=leading_dim) + + torch2cute_dtype_map = { torch.float16: cutlass.Float16, torch.bfloat16: cutlass.BFloat16, @@ -35,6 +72,16 @@ torch2cute_dtype_map = { } +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, @@ -43,6 +90,8 @@ def _flash_attn_fwd( 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, @@ -56,14 +105,29 @@ def _flash_attn_fwd( 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: Callable | None = 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, - buffers: Optional[list[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: @@ -101,6 +165,7 @@ def _flash_attn_fwd( 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, @@ -127,6 +192,7 @@ def _flash_attn_fwd( 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 ( @@ -175,17 +241,13 @@ def _flash_attn_fwd( 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" + _validate_tensor( + out, + "out", + (*q_batch_seqlen_shape, num_head, head_dim_v), + out_torch_dtype, + device, + ) if lse is None: lse = ( @@ -194,103 +256,134 @@ def _flash_attn_fwd( 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" + _validate_tensor(lse, "lse", lse_shape, torch.float32, device) dtype = torch2cute_dtype_map[q.dtype] - 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) - ] - lse_tensor = ( - from_dlpack(lse.detach(), assumed_align=4).mark_layout_dynamic( - leading_dim=lse.ndim - 1 - ) - if lse is not None - else None - ) - ( - 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 - ) - 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 compute_capability = ( - torch.cuda.get_device_capability()[0] - if _compute_capability is None - else _compute_capability + _get_device_capability() if _compute_capability is None else _compute_capability ) + assert compute_capability in [ 9, 10, - ], "Unsupported compute capability. Supported: 9.x, 10.x" + 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 - # Perf heuristic from upstream: hdim=128, noncausal, non-local benefits from larger n_block - if head_dim == head_dim_v == 128 and not causal and not local: - n_block_size = 192 - if compute_capability == 10: - # TODO: fix the varlen case + if compute_capability == 9: # TODO: tune block size according to hdim. if ( - pack_gqa - and (128 % qhead_per_kvhead != 0) - or (cu_seqlens_q is not None or seqused_q is not None) + 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) - if score_mod is not None: - 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 - ) + 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( - "score_mod with buffers is not yet supported for varlen sequences. This will be fixed in a future PR." + "mask_mod with aux_tensors is not yet supported for varlen sequences. This will be fixed in a future PR." ) - cute_buffers = None - if buffers is not None: - cute_buffers = [from_dlpack(buf) for buf in buffers] + 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." + ) compile_key = ( dtype, @@ -298,8 +391,10 @@ def _flash_attn_fwd( head_dim_v, qhead_per_kvhead, causal, - utils.hash_callable(score_mod) if score_mod is not None else None, - buffers is not None, + 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, @@ -311,13 +406,74 @@ def _flash_attn_fwd( 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 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, + ) + ) + compile_time_normalized = normalize_block_sparse_tensors( + block_sparse_tensors, + expected_count_shape=expected_count_shape, + expected_index_shape=expected_index_shape, + ) + sparse_tensors = to_cute_block_sparse_tensors(compile_time_normalized) + + 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, @@ -333,34 +489,40 @@ def _flash_attn_fwd( 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_buffers=buffers is not None, + has_aux_tensors=aux_tensors is not None, ) - elif compute_capability == 10: - assert page_size in [ - None, - 128, - ], "Only page_size=128 is supported for paged KV on SM 10.0" + 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 seqused_q is None + and not is_split_kv, score_mod=score_mod, - has_buffers=buffers is not None, + 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" + f"Unsupported compute capability: {compute_capability}. Supported: 9.x, 10.x, 11.x" ) # TODO: check @can_implement - # TODO caching for buffers; cute_buffers _flash_attn_fwd.compile_cache[compile_key] = cute.compile( fa_fwd, q_tensor, @@ -378,32 +540,242 @@ def _flash_attn_fwd( window_size_left, window_size_right, learnable_sink_tensor, - cute_buffers, + sparse_tensors, + cute_aux_tensors, + options="--enable-tvm-ffi", + ) + + # Expand block sparse tensors to match actual head count (may be broadcast from 1) + normalized_block_sparse_tensors = None + if block_sparse_tensors is not None: + 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, ) _flash_attn_fwd.compile_cache[compile_key]( - q_tensor, - k_tensor, - v_tensor, - o_tensor, - lse_tensor, + 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_tensor, - cu_seqlens_k_tensor, - seqused_q_tensor, - seqused_k_tensor, - page_table_tensor, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + page_table, window_size_left, window_size_right, - learnable_sink_tensor, - cute_buffers, + 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_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 warmup_flash_attn(f): """ Decorator for flash_attn_varlen_func: @@ -537,8 +909,11 @@ def flash_attn_varlen_func( 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, return_softmax_lse: Optional[bool] = False, + score_mod: Optional[Callable] = None, + aux_tensors: Optional[list] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: out, lse = _flash_attn_fwd( q, @@ -555,8 +930,11 @@ def flash_attn_varlen_func( window_size_right=window_size[1], learnable_sink=learnable_sink, softcap=softcap, + num_splits=num_splits, pack_gqa=pack_gqa, return_lse=return_softmax_lse, + score_mod=score_mod, + aux_tensors=aux_tensors, ) return (out, lse) if return_softmax_lse else out diff --git a/sgl-kernel/python/sgl_kernel/flash_attn.py b/sgl-kernel/python/sgl_kernel/flash_attn.py index c3ffbc540..12afdb790 100644 --- a/sgl-kernel/python/sgl_kernel/flash_attn.py +++ b/sgl-kernel/python/sgl_kernel/flash_attn.py @@ -68,6 +68,8 @@ def flash_attn_with_kvcache( sm_margin=0, # Can be tuned if some SMs are used for communication return_softmax_lse=False, sinks=None, + score_mod=None, + aux_tensors=None, ver=3, ): """ @@ -149,6 +151,8 @@ def flash_attn_with_kvcache( to automatically determine the number of splits. Don't change this unless you know what you are doing. return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + score_mod [optional]: A callable that takes the attention scores and applies a modification. + aux_tensors [optional]: Some score_mods will want to read from global aux_tensors. This is how we thread them through to the inner kernel. Return: out: (batch_size, seqlen, nheads, headdim). @@ -176,6 +180,7 @@ def flash_attn_with_kvcache( if window_size == (-1, -1): window_size = (None, None) + return flash_attn_varlen_func_v4( q=q, k=k_cache, @@ -186,10 +191,13 @@ def flash_attn_with_kvcache( causal=causal, window_size=window_size, softcap=softcap, + num_splits=num_splits, pack_gqa=pack_gqa, return_softmax_lse=return_softmax_lse, learnable_sink=sinks, page_table=page_table, + score_mod=score_mod, + aux_tensors=aux_tensors, ) assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" @@ -286,6 +294,8 @@ def flash_attn_varlen_func( sm_margin=0, return_softmax_lse=False, sinks=None, + score_mod=None, + aux_tensors=None, ver=3, ): if ver == 4: @@ -311,6 +321,8 @@ def flash_attn_varlen_func( pack_gqa=pack_gqa, learnable_sink=sinks, return_softmax_lse=return_softmax_lse, + score_mod=score_mod, + aux_tensors=aux_tensors, ) if not is_fa3_supported(): diff --git a/sgl-kernel/tests/test_flash_attention_4.py b/sgl-kernel/tests/test_flash_attention_4.py index 2296d71aa..45a5c9cf9 100644 --- a/sgl-kernel/tests/test_flash_attention_4.py +++ b/sgl-kernel/tests/test_flash_attention_4.py @@ -526,7 +526,7 @@ def attention_ref( # @pytest.mark.parametrize('d', [56, 80]) # @pytest.mark.parametrize('d', [32, 40, 64, 80, 96, 128]) # @pytest.mark.parametrize("d", [64, 96, 128]) -@pytest.mark.parametrize("d", [128, 192]) +@pytest.mark.parametrize("d", [64, 128]) # @pytest.mark.parametrize("d", [192]) @pytest.mark.parametrize( "seqlen_q,seqlen_k",