Implementation:Sgl project Sglang Sparse Flash Attention

Knowledge Sources	Sgl_project_Sglang
Domains	Kernel, Attention, Sparse Computation
Last Updated	2026-02-10 00:00 GMT

Overview

Python interface for sparse Flash Attention with vertical-slash sparsity pattern support, enabling efficient long-context attention.

Description

The sparse_flash_attn.py module provides sparse attention kernels that exploit vertical (column) and slash (diagonal) sparsity patterns, as described in Appendix C.4.2 of the MInference paper. convert_vertical_slash_indexes converts raw sparsity index representations into block-level metadata: block_count and block_offset for slash (diagonal) patterns, and column_count and column_index for vertical (column) patterns. convert_vertical_slash_indexes_mergehead is a variant that supports per-head variable sparsity counts via vertical_indices_count and slash_indices_count. sparse_attn_func executes sparse attention for fixed-length batched inputs with shape (batch_size, seqlen, nheads, headdim), supporting causal masking, softcap, ALiBi slopes, and optional dropout. sparse_attn_varlen_func handles variable-length batched inputs using cumulative sequence lengths (cu_seqlens_q, cu_seqlens_k) with the same feature set. Both attention functions ensure input contiguity via the maybe_contiguous helper and delegate to torch.ops.sgl_kernel.fwd_sparse and torch.ops.sgl_kernel.varlen_fwd_sparse respectively.

Usage

Use these functions for efficient long-context attention in LLM serving where full attention is not necessary, particularly when vertical-slash sparsity patterns have been identified through attention profiling or predefined heuristics.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/python/sgl_kernel/sparse_flash_attn.py
• Lines: 1-293

Signature

def maybe_contiguous(x): ...

def convert_vertical_slash_indexes(
    q_seqlens: torch.Tensor,
    kv_seqlens: torch.Tensor,
    vertical_indexes: torch.Tensor,
    slash_indexes: torch.Tensor,
    context_size: int,
    block_size_M: int,
    block_size_N: int,
    causal: bool = True,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ...

def convert_vertical_slash_indexes_mergehead(
    q_seqlens: torch.Tensor,
    kv_seqlens: torch.Tensor,
    vertical_indexes: torch.Tensor,
    slash_indexes: torch.Tensor,
    vertical_indices_count: torch.Tensor,
    slash_indices_count: torch.Tensor,
    context_size: int,
    block_size_M: int,
    block_size_N: int,
    causal: bool = True,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ...

def sparse_attn_func(
    q, k, v,
    block_count, block_offset, column_count, column_index,
    dropout_p=0.0, softmax_scale=None, causal=False,
    softcap=0.0, alibi_slopes=None, deterministic=False,
    return_attn_probs=False, *, return_softmax_lse=False, out=None,
): ...

def sparse_attn_varlen_func(
    q, k, v,
    block_count, block_offset, column_count, column_index,
    cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
    dropout_p=0.0, softmax_scale=None, causal=False,
    softcap=0.0, alibi_slopes=None, deterministic=False,
    return_attn_probs=False, *, return_softmax_lse=False, out=None,
): ...

Import

from sgl_kernel.sparse_flash_attn import (
    convert_vertical_slash_indexes,
    convert_vertical_slash_indexes_mergehead,
    sparse_attn_func,
    sparse_attn_varlen_func,
)

I/O Contract

Inputs

Name	Type	Required	Description
q	torch.Tensor	Yes	Query tensor, shape (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)
k	torch.Tensor	Yes	Key tensor, shape (batch, seqlen, nheads_k, headdim) or (total_k, nheads_k, headdim)
v	torch.Tensor	Yes	Value tensor, same shape as k
block_count	torch.Tensor	Yes	Slash block counts, shape (batch, nheads, cdiv(seqlen, BLOCK_M))
block_offset	torch.Tensor	Yes	Slash block offsets, shape (batch, nheads, cdiv(seqlen, BLOCK_M), NNZ_S)
column_count	torch.Tensor	Yes	Vertical column counts, shape (batch, nheads, cdiv(seqlen, BLOCK_M))
column_index	torch.Tensor	Yes	Vertical column indices, shape (batch, nheads, cdiv(seqlen, BLOCK_M), NNZ_V)
q_seqlens	torch.Tensor	Yes (convert)	Query sequence lengths, shape (batch,)
kv_seqlens	torch.Tensor	Yes (convert)	KV sequence lengths, shape (batch,)
vertical_indexes	torch.Tensor	Yes (convert)	Raw vertical sparsity indices, shape (batch, nheads, NNZ_V)
slash_indexes	torch.Tensor	Yes (convert)	Raw slash sparsity indices, shape (batch, nheads, NNZ_S)
context_size	int	Yes (convert)	Maximum context size for block computation
block_size_M	int	Yes (convert)	Query block size
block_size_N	int	Yes (convert)	Key block size
softmax_scale	Optional[float]	No	QK^T scaling, defaults to 1/sqrt(headdim)
causal	bool	No	Whether to apply causal mask, default False
softcap	float	No	Softcap value, 0.0 to deactivate
alibi_slopes	Optional[torch.Tensor]	No	ALiBi positional bias slopes
cu_seqlens_q	torch.Tensor	Yes (varlen)	Cumulative query sequence lengths, shape (batch+1,)
cu_seqlens_k	torch.Tensor	Yes (varlen)	Cumulative key sequence lengths, shape (batch+1,)
max_seqlen_q	int	Yes (varlen)	Maximum query sequence length
max_seqlen_k	int	Yes (varlen)	Maximum key sequence length

Outputs

Name	Type	Description
block_count	torch.Tensor	Slash block counts (from convert functions)
block_offset	torch.Tensor	Slash block offsets (from convert functions)
column_count	torch.Tensor	Vertical column counts (from convert functions)
column_index	torch.Tensor	Vertical column indices (from convert functions)
out	torch.Tensor	Attention output, same shape as q
softmax_lse	torch.Tensor	Log-sum-exp values (if return_softmax_lse=True)

Usage Examples

from sgl_kernel.sparse_flash_attn import (
    convert_vertical_slash_indexes,
    sparse_attn_func,
)

# Step 1: Convert sparsity indices to block-level metadata
block_count, block_offset, column_count, column_index = \
    convert_vertical_slash_indexes(
        q_seqlens, kv_seqlens,
        vertical_indexes, slash_indexes,
        context_size=4096, block_size_M=64, block_size_N=64,
        causal=True
    )

# Step 2: Run sparse attention
out = sparse_attn_func(
    q, k, v,
    block_count, block_offset,
    column_count, column_index,
    causal=True
)

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime