[NVIDIA] upstream FA4 (#15182)

Co-authored-by: Qiaolin-Yu <liin1211@outlook.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>

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)

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)
+        _validate_tensor(
-        assert (
+            out,
-            out.shape == expected_out_shape
+            "out",
-        ), f"out tensor shape {out.shape} does not match expected shape {expected_out_shape}"
+            (*q_batch_seqlen_shape, num_head, head_dim_v),
-        assert (
+            out_torch_dtype,
-            out.dtype == out_torch_dtype
+            device,
-        ), 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 = (
@@ -194,103 +256,134 @@ def _flash_attn_fwd(
             else None
         )
     elif lse is not None:
-        assert (
+        _validate_tensor(lse, "lse", lse_shape, torch.float32, device)
-            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]
-    q_tensor, k_tensor, v_tensor, o_tensor = [
+    compute_capability = (
-        from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(
+        _get_device_capability() if _compute_capability is None else _compute_capability
-            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
     )
+
+    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
-    compute_capability = (
+    else:
-        torch.cuda.get_device_capability()[0]
+        causal, local = False, False
-        if _compute_capability is None
+
-        else _compute_capability
-    )
-    assert compute_capability in [
-        9,
-        10,
-    ], "Unsupported compute capability. Supported: 9.x, 10.x"
     current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
 
-    if compute_capability == 9:  # TODO: tune block size according to hdim
+    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 (
-            pack_gqa
+            head_dim == head_dim_v == 128
-            and (128 % qhead_per_kvhead != 0)
+            and not causal
-            or (cu_seqlens_q is not None or seqused_q is not None)
+            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
     )
+
+    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 use_block_sparsity:
-    if buffers is not None:
+        if is_varlen:
-        cute_buffers = [from_dlpack(buf) for buf in buffers]
+            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,
+        score_mod_hash,
-        buffers is not None,
+        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:
+        elif compute_capability in [10, 11]:
-            assert page_size in [
-                None,
-                128,
-            ], "Only page_size=128 is supported for paged KV on SM 10.0"
             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,25 +540,55 @@ 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,
+        q,
-        k_tensor,
+        k,
-        v_tensor,
+        v,
-        o_tensor,
+        out if not is_split_kv else out_partial,
-        lse_tensor,
+        lse_partial if is_split_kv else lse,
         softmax_scale,
         current_stream,
-        cu_seqlens_q_tensor,
+        cu_seqlens_q,
-        cu_seqlens_k_tensor,
+        cu_seqlens_k,
-        seqused_q_tensor,
+        seqused_q,
-        seqused_k_tensor,
+        seqused_k,
-        page_table_tensor,
+        page_table,
         window_size_left,
         window_size_right,
-        learnable_sink_tensor,
+        learnable_sink,
-        cute_buffers,
+        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
 
@@ -404,6 +596,186 @@ def _flash_attn_fwd(
 _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

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():

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",