native-sparse-attention-pytorch

The SparseAttention class implements the sparse attention pattern proposed by the Deepseek team. It allows for hardware-aligned and natively trainable sparse attention.

Key parameters for initialization:

• dim: The input dimension.
• dim_head: The dimension of each attention head.
• heads: The number of attention heads.
• sliding_window_size: Size of the sliding window.
• compress_block_size: Size of the compression blocks.
• compress_block_sliding_stride: Stride for the compression block sliding window.
• selection_block_size: Size of the selection blocks.
• num_selected_blocks: Number of blocks to be selected.

import torch
from native_sparse_attention_pytorch import SparseAttention

attn = SparseAttention(
    dim = 512,
    dim_head = 64,
    heads = 8,
    sliding_window_size = 2,
    compress_block_size = 4,
    compress_block_sliding_stride = 2,
    selection_block_size = 4,
    num_selected_blocks = 2
)

tokens = torch.randn(2, 31, 512)

attended = attn(tokens)

assert tokens.shape == attended.shape

When initializing SparseAttention, use the following parameters to control the sparsity and attention behavior:

The implementation provides utility functions to create block_mask objects compatible with PyTorch's flex_attention. These can be passed to SparseAttention.forward via sliding_window_flex_mask or fine_selection_flex_mask to accelerate computation.

• create_sliding_mask(seq_len, window_size, causal=True): Creates a mask for local sliding window attention.
• create_compress_mask(seq_len, kv_seq_len, compress_block_sliding_stride, mem_kv_len=0, causal=True): Creates a mask for the compressed attention pathway.
• create_fine_mask(seq_len, fine_block_size, causal=True): Returns a function that, when used with create_block_mask, defines the fine-grained selection mask based on provided selected_block_indices.

from torch.nn.attention.flex_attention import create_block_mask

# Example: Creating a sliding window mask for flex_attention
seq_len = 1024
window_size = 128

def sliding_mask(_, __, q_idx, kv_idx):
    distance = q_idx - kv_idx
    backward_sliding_mask = distance <= window_size
    forward_distance = 0 # assuming causal
    forward_sliding_mask = distance >= forward_distance
    return backward_sliding_mask & forward_sliding_mask

block_mask = create_block_mask(sliding_mask, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, _compile=True)

The Transformer class is the main entry point for building a model with sparse attention capabilities. You can choose between standard dense attention or various sparse attention mechanisms (sliding window, fine selection, etc.) using the use_sparse_attn flag and specific configuration kwargs.

Key configuration parameters:

• num_tokens: Vocabulary size.
• dim: Model dimension.
• depth: Number of layers.
• use_sparse_attn: Set to True to enable sparse attention modes.
• use_flex_sliding_window: Enables sliding window attention using PyTorch's flex_attention.
• use_flex_fine_selection: Enables fine-grained selection using flex_attention.
• use_triton_fine_selection: Enables fine-grained selection using a custom Triton kernel (cannot be used with use_flex_fine_selection).
• sparse_attn_kwargs: A dictionary passed to the underlying SparseAttention module to configure window sizes, block sizes, and selection parameters.

from native_sparse_attention_pytorch.transformer import Transformer

model = Transformer(
    num_tokens = 50257,
    dim = 512,
    depth = 12,
    use_sparse_attn = True,
    use_flex_sliding_window = True,
    sparse_attn_kwargs = dict(
        sliding_window_size = 64,
        compress_block_size = 4,
        compress_block_overlap_len = 0,
        selection_block_size = 4,
        num_selected_blocks = 4,
    )
)

The native_sparse_attn_forward function provides a Triton-accelerated implementation of sparse attention. It supports causal masking, sliding window attention, and grouped query attention (GQA).

Requirements:

• q, k, v, and kv_block_indices must be contiguous tensors.
• q, k, and v must be on a CUDA device.
• q, k, and v must have the same dtype (either torch.float16 or torch.bfloat16).
• The head dimension (dim) must be $\le 128$.
• kv_block_indices and kv_block_mask must have matching shapes.

Parameters:

• q: Query tensor of shape (batch, nheads, seqlen_q, dim).
• k: Key tensor of shape (batch, kv_heads, seqlen_k, dim).
• v: Value tensor of shape (batch, kv_heads, seqlen_k, dim).
• kv_block_indices: Indices of selected KV blocks.
• kv_block_mask: Mask for the selected KV blocks.
• block_size: Size of the attention block (must be a multiple of 16; default is 128).
• include_block_causal: Whether to include block-diagonal causal masking (default True).
• return_sliding_window_out: If True, returns both the standard output and a sliding window output.

Returns:

• o: The attention output tensor.
• slide_o: The sliding window attention output (only if return_sliding_window_out is True).
• lse: Log-sum-exp statistics for the attention.
• slide_lse: Log-sum-exp statistics for the sliding window (only if return_sliding_window_out is True).

from native_sparse_attention_pytorch.triton_native_sparse_attention import native_sparse_attn_forward

# Example usage (shapes are illustrative)
o, slide_o, lse, slide_lse = native_sparse_attn_forward(
    q=q, 
    k=k, 
    v=v, 
    kv_block_indices=kv_block_indices, 
    kv_block_mask=kv_block_mask,
    block_size=128,
    include_block_causal=True,
    return_sliding_window_out=True
)