[RFC] Could you please provide the latest sample code for using the latest `parallel_nsa` function?

[Kyfafyd]

Proposal

R.T.

Rationale

No response

[Kyfafyd]

Also, what about the usage for naive_nsa_with_compression and parallel_nsa_with_compression?

[Hanyuezhuohua]

Thank you for your interest in our project. The latest sample code for using the latest parallel_nsa function has been provided, including both selected and sliding attention. You can directly try it following our instructions.

For the nsa_with_compression function, we will further integrate the compressed branch and online top-k selection into our kernel. However, this function is still under development for the parallel version.

[Kyfafyd]

Thanks for your response!
I would like to learn if nsa kernel can be applied for attention with 48 heads, which is not a power of 2.
Additionally, how to provide the input for g_slc, g_swa, block_indices, and block_counts in the attention computation of a vision transformer?

[Hanyuezhuohua]

The NSA kernel can be applied to attention with 48 query heads and 3 key-value heads. Since the query heads are grouped by key-value head, each group consists of 16 query heads.

In addition, our new nsa_with_compression function allows you to obtain block indices. Specifically, g_slc and g_swa are computed by passing the input through an MLP module with a sigmoid activation. The block_counts parameter is a user-defined constant that controls the sparsity ratio.