Please use spas_sage2_attn_meansim_topk_cuda and block_sparse_sage2_attn_cuda APIs

[jt-zhang]

We’ve updated the repository to support a convenient plug-and-play API, as well as custom block-sparse masks API.

Recommended API

We highly recommend using the spas_sage2_attn_meansim_topk_cuda and block_sparse_sage2_attn_cuda APIs. They are plug-and-play and customizable:

Plug-and-Play API

from spas_sage_attn import spas_sage2_attn_meansim_topk_cuda

attn_output = spas_sage2_attn_meansim_topk_cuda(q, k, v, topk=0.5, is_causal=False)

You can adjust topk to balance attention accuracy (higher topk is more accurate) and sparsity (lower topk is more sparse).

Customize your Block-Sparse Mask

from spas_sage_attn import block_sparse_sage2_attn_cuda

attn_output = block_sparse_sage2_attn_cuda(q, k, v, mask_id=None)

With this API, we support computing attention with any block-sparse mask per attention head. Specifically, the per-head attention mask mask_id has shape (batch_size, num_heads, ⌈seq_len / 128⌉, ⌈seq_len // 64⌉) and consists of 0 and 1. Currently, the block size is 128×64.

Feel free to try it out!

[dsouzaankit]

Hey @jt-zhang, is Sage attention 3 FP4 supported implicitly? Didn't catch an explicit connection to these: https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell/sageattn3/blackwell. Please let me know, thanks!