[Feature request] Sparse Attention

[Ageliss]

Recently, we see several awesome work focusing on kv cache compressing and they said can accelearte 1.7~2.3 times than FlashInfer, can you guys plz consider to surpport such features?

Same layer KV:
Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference

MoA: Mixture of Sparse Attention for Automatic Large Language Model Compression

Cross layer KV:
Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

[yzh119]

Sparse Attention is already implemented in the CUDA codebase, I should create pytorch api for that, please stay tuned.

[yzh119]

To be more clear, flashinfer implemented general (any block size: (B_r, B_c) ) block sparse flashattention on CUDA, and flashinfer's page attention implementation is based on these block sparse flashattention implementations (by representing page table as block sparse matrix). Quest's sparse attention implement is also based on flashinfer's CUDA APIs.

[Ageliss]

To be more clear, flashinfer implemented general (any block size: (B_r, B_c) ) block sparse flashattention on CUDA, and flashinfer's page attention implementation is based on these block sparse flashattention implementations (by representing page table as block sparse matrix). Quest's sparse attention implement is also based on flashinfer's CUDA APIs.

Good to know that, will Quest's sparse attention be merged into current flashinfer's CUDA APIs?

[yzh119]

Quest didn't implement anything new to FlashInfer on CUDA side (FlashInfer has already implemented all of the necessary CUDA sparse kernels), what Quest did is to call FlashInfer CUDA apis. So there is no need to "merge" Quest.

I'll create a PR for a general block sparse attention python API (which should be a superset of Quest's), please wait for a few hours.

[yzh119]

Please check https://docs.flashinfer.ai/api/python/sparse.html