Please use spas_sage2_attn_meansim_topk_cuda and block_sparse_sage2_attn_cuda APIs #100

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 ×ばつ64.

Feel free to try it out!