SpargeAttention Documentation

The following models are supported via the plug-and-play inference scripts:

• WAN: Wan2.1-T2V-1.3B, Wan2.1-T2V-14B, and Wan2.2-T2V-A14B (Diffusers versions).
• LTX-Video: 0.9.7-dev and spatial upscaler.
• HunyuanVideo
• CogVideoX-2B

To install SpargeAttention, ensure your environment meets the base requirements and then use pip or setup.py to install the package. Parallel compilation via ninja is recommended.

# Base environment requirements:
# python>=3.9, torch>=2.3.0
# CUDA: >=12.8 (Blackwell), >=12.4 (Ada fp8), >=12.3 (Hopper fp8), >=12.0 (Ampere)

# Installation steps:
pip install ninja
python setup.py install
# or
pip install -e .

You can run pre-configured inference scripts for various video models (Wan, LTX-Video, HunyuanVideo, CogVideoX-2B) that have the SpargeAttn API integrated. These scripts should be executed as modules from the repository root.

Execution Command:

python -m inference_examples.<script_name> [--flags]

Output Directory Structure: Videos are saved to inference_examples/videos/<model>/<run_dir>/. The <run_dir> name depends on the selected mode:

• --mode full: original/
• --mode cdfthreshd: cdfthreshd-<value>/
• --mode topk: topk-<value>/

python -m inference_examples.wan_infer \
  --model wan2_2-14b \
  --mode topk --value 0.4 \
  --start 0 --end 1

The recommended way to use SpargeAttention is via the spas_sage2_attn_meansim_topk_cuda API. This is a drop-in replacement for torch.nn.functional.scaled_dot_product_attention.

You can adjust the topk parameter to balance attention accuracy (higher topk) and sparsity (lower topk).

from spas_sage_attn import spas_sage2_attn_meansim_topk_cuda

# Replace torch.nn.functional.scaled_dot_product_attention with this:
attn_output = spas_sage2_attn_meansim_topk_cuda(q, k, v, topk=0.5, is_causal=False)

If you need to compute attention using a specific block-sparse mask per attention head, use the block_sparse_sage2_attn_cuda API.

Mask Requirements:

• The mask_id parameter must have the shape (batch_size, num_heads, ⌈seq_len / 128⌉, ⌈seq_len // 64⌉).
• The mask must consist of 0 and 1 values.
• The current fixed block size is 128×64.

from spas_sage_attn import block_sparse_sage2_attn_cuda

# mask_id shape: (batch_size, num_heads, ceil(seq_len / 128), ceil(seq_len / 64))
attn_output = block_sparse_sage2_attn_cuda(q, k, v, mask_id=None)

When running the inference scripts, use the following flags to control model selection and SpargeAttn behavior:

# Example for Wan
python -m inference_examples.wan_infer \
  --model wan2_2-14b \
  --mode topk --value 0.4 \
  --start 0 --end 1

# Example for LTX-Video
python -m inference_examples.ltx_infer --mode topk --value 0.5

SpargeAttention provides several APIs based on different versions of SageAttention. The sage2 versions are highly recommended for better performance.

- spas_sage2_attn_meansim_topk_cuda: SpargeAttn based on SageAttention2 (RECOMMENDED)
- spas_sage2_attn_meansim_cuda: SpargeAttn based on SageAttention2 (NOT RECOMMENDED)
- spas_sage_attn_meansim_topk_cuda: SpargeAttn based on SageAttention (RECOMMENDED)
- spas_sage_attn_meansim_cuda: SpargeAttn based on SageAttention (NOT RECOMMENDED)