Implementation:Sgl project Sglang FA4 Interface

Knowledge Sources	Sgl_project_Sglang
Domains	Flash Attention, GPU Kernels, CUTLASS DSL
Last Updated	2026-02-10 00:00 GMT

Overview

Flash Attention 4 interface using the CUTLASS Cute DSL for JIT-compiled attention kernels targeting Hopper (SM90) and Blackwell (SM100/SM110) architectures.

Description

_fa4_interface.py implements the cutting-edge Flash Attention 4 (FA4) backend that leverages the CUTLASS Cute DSL to JIT-compile attention kernels via TVM FFI. This provides an alternative to the C++-based FA3 path, achieving maximum performance on the latest NVIDIA architectures.

The module contains two primary functions:

_flash_attn_fwd is the core forward attention function that handles the complete FA4 pipeline:

• Architecture Detection: Caches device capability checks to select between FlashAttentionForwardSm90 (Hopper) and FlashAttentionForwardSm100 (Blackwell/SM110) implementations
• Kernel Compilation and Caching: Uses a compile_cache dictionary keyed by a comprehensive tuple of configuration parameters (dtype, head_dim, causal, GQA ratio, block sizes, etc.) to avoid recompilation
• Split-KV Support: Implements a num_splits_heuristic based on total M-blocks and SM count, with a combine kernel (FlashAttentionForwardCombine) that reduces partial outputs
• Block Sparsity: Supports BlockSparseTensorsTorch for sparse attention patterns, with shape normalization and broadcasting
• GQA Packing: Automatically enables GQA packing when qhead_per_kvhead > 1, with architecture-specific constraints (SM100 requires 128-divisible heads)
• Tensor Conversion: Converts between PyTorch and Cute tensors via to_cute_tensor with configurable alignment and dynamic layout marking

flash_attn_varlen_func is the public API wrapped with the @warmup_flash_attn decorator. The warmup decorator:

• Runs multiple warmup passes on first call covering global/causal/local attention, LSE on/off, and optional GQA packing
• Clones arguments to avoid modifying user tensors
• Releases GPU memory between warmup passes via torch.cuda.empty_cache()
• Can be disabled via the SGLANG_DISABLE_FA4_WARMUP environment variable

The function supports:

• Variable-length sequences via cu_seqlens_q and cu_seqlens_k
• Paged KV cache via page_table (int32, with page_size constraints)
• Causal and sliding window masking via window_size tuple
• Softcapping via softcap parameter (converted to a score_mod internally)
• Custom score/mask modifications via score_mod and aux_tensors
• Learnable sink tokens via learnable_sink (bfloat16)
• FP16 and BF16 input dtypes

Usage

Use this module when serving models on Hopper or Blackwell GPUs where FA4's JIT-compiled kernels can provide better performance than the pre-compiled FA3 path. It is typically accessed indirectly through the flash_attn.py module which dispatches between FA3 and FA4 based on the ver parameter.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/python/sgl_kernel/_fa4_interface.py
• Lines: 1-940

Signature

def _flash_attn_fwd(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    cu_seqlens_q: Optional[torch.Tensor] = None,
    cu_seqlens_k: Optional[torch.Tensor] = None,
    seqused_q: Optional[torch.Tensor] = None,
    seqused_k: Optional[torch.Tensor] = None,
    max_seqlen_q: Optional[int] = None,
    max_seqlen_k: Optional[int] = None,
    page_table: Optional[torch.Tensor] = None,
    softmax_scale: Optional[float] = None,
    causal: bool = False,
    softcap: Optional[float] = None,
    window_size_left: Optional[int] = None,
    window_size_right: Optional[int] = None,
    learnable_sink: Optional[torch.Tensor] = None,
    m_block_size: int = 128,
    n_block_size: int = 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]:

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,
    return_softmax_lse: Optional[bool] = False,
    score_mod: Optional[Callable] = None,
    aux_tensors: Optional[list] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:

def num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, max_splits) -> int:

def maybe_contiguous(x):

def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False):

Import

from sgl_kernel._fa4_interface import flash_attn_varlen_func

I/O Contract

Inputs

Name	Type	Required	Description
q	torch.Tensor (fp16/bf16)	Yes	Query tensor: (batch, seqlen_q, num_heads, head_dim) or (total_q, num_heads, head_dim)
k	torch.Tensor (fp16/bf16)	Yes	Key tensor: (batch, seqlen_k, num_heads_kv, head_dim) or paged layout
v	torch.Tensor (fp16/bf16)	Yes	Value tensor: same layout as k, head_dim_v may differ
cu_seqlens_q	torch.Tensor (int32)	No	Cumulative query sequence lengths (batch_size+1,)
cu_seqlens_k	torch.Tensor (int32)	No	Cumulative key sequence lengths (batch_size+1,)
page_table	torch.Tensor (int32)	No	Page table for paged KV cache (batch_size, max_pages)
softmax_scale	float	No	Attention scale factor, defaults to 1/sqrt(head_dim)
causal	bool	No	Whether to apply causal attention mask
window_size	Tuple[Optional[int], Optional[int]]	No	Sliding window (left, right) sizes
softcap	float	No	Softcapping value (0.0 to disable)
num_splits	int	No	Number of KV splits (0 for heuristic, 1 for no split)

Outputs

Name	Type	Description
out	torch.Tensor	Attention output, same shape and dtype as q (with head_dim_v)
lse	torch.Tensor (float32)	Log-sum-exp values (only if return_softmax_lse=True)

Usage Examples

from sgl_kernel._fa4_interface import flash_attn_varlen_func

# Variable-length attention with causal masking
out = flash_attn_varlen_func(
    q=query_tensor,       # (total_q, num_heads, head_dim)
    k=key_tensor,         # (total_k, num_heads_kv, head_dim)
    v=value_tensor,       # (total_k, num_heads_kv, head_dim_v)
    cu_seqlens_q=cu_q,    # (batch_size + 1,), int32
    cu_seqlens_k=cu_k,    # (batch_size + 1,), int32
    causal=True,
    softmax_scale=1.0 / math.sqrt(128),
)

# With paged KV cache and sliding window
out, lse = flash_attn_varlen_func(
    q=query_tensor,
    k=paged_key_cache,     # (num_pages, page_size, num_heads_kv, head_dim)
    v=paged_value_cache,
    cu_seqlens_q=cu_q,
    seqused_k=cache_seqlens,
    page_table=page_table,  # (batch_size, max_num_pages)
    window_size=(256, 0),   # left=256, right=0 (causal with window)
    return_softmax_lse=True,
)

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime