cuda : support for sparse KV indices in MMA flash attention kernel

[fairydreaming]

Overview

This PR adds support for sparse KV indices (called top-k) in MMA flash attention kernel. Originally I wrote this for my DeepSeek V3.2 implementation (#21149). I think it may be still useful for someone or at least serve as an inspiration for implementing sparse attention in llama.cpp.

Additional information

Currently the sparse path is only enabled for:

• DeepSeek V3.2 and similar models (GLM 5/5.1/5.2) since they calculate top-k indices tensor that can be directly used,
• DeepSeek V4 CSA layers as they also use lightning indexer.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: I used an AI assistant to help me debug issues with the kernel

[jeffbolznv]

Please add comments documenting how this input affects the calculations.

[fairydreaming]

Edit: DeepSeek V4 shall work correctly now with sparse attention (from my experience it shows improved performance starting from context length around 64k).

[ggerganov]

The main thing that I am wondering is if we need to explicitly pass the top_k indices to the flash attention.

The alternative approach that I would like to consider is to just pass a boolean hint to the flash attention op: "is the attention sparse?". Then in the backend, I would check this hint, and if it set, I will first run a quick kernel on top of the KQ mask to collect the non-masked indices and pass those to the FA kernels. I think this could be more generic than the proposed "top-k hint". The only question is if it will be efficient enough to collect the indices this way.

@tarruda Initially, I thought that your sparse approach from earlier (#26512) was doing the described idea above. LMK if you give this a shot - it should be quite light on the ggml changes (just a hint for the existing ggml_flash_attn_ext) and for the model graph - we can initially just always set the hint to true for the DSv4 model. The rest would be only changes in the Metal backend around the FA operator.

[fairydreaming]

@ggerganov I don't like it because of the memory usage. If the model does not use compressed attention like DSv4 the memory needed for 1M context length and for example 8k ubatch (for fast PP) is 1M * 8k * 2 (f16) = 16GB - filled mostly with -INF in DSA based models.

I thought about this a bit and I wonder if perhaps we could simply support KQ mask tensor in two formats - dense and sparse. FA implementation would check the mask tensor type - if F16/F32 it would be dense (value-based), if I32 it would be sparse (indices-based).

So DSA-based models instead of modifying dense KQ mask based on top-k tensor like it's currently done would:

1. prepare top-k tensor,
2. apply causality on top-k indices with some new OP (masked indices would be -1 or something)
3. pass masked top-k as sparse KQ mask to FA

This is simple and keeps the memory usage low. We could convert the mask from sparse format to dense only for backends that do not support sparse KQ mask. It doesn't require any FA API changes (except for new allowed KQ mask tensor type/shape).

Thoughts?

[tarruda]

I thought about this a bit and I wonder if perhaps we could simply support KQ mask tensor in two formats - dense and sparse.

Not sure if it is the same thing you are talking about, but in an earlier branch had implemented a env var gated selection for dense attention:
73e8829

[ggerganov]

apply causality on top-k indices with some new OP (masked indices would be -1 or something)

It would work, but that specific part might be complicated because the causality-related information is on the host side in the llama_kv_cells structure, which is not easy to pass to the backends. There could be some simplified causality logic where it just masks based solely on the indices (i.e. assuming that they are correlated with the sequence position of the data), but this IMO is quite limiting - cannot work with unified KV cache and not compatible with multi-modal cases where sometimes the causality is not simpler triangular matrix. Apart from this, the approach could work.

Regarding the memory concern - if the attention is not compressed, then the KV cache would already be an enormous amount of memory, multiple times bigger than the F16 KQ mask. Plus I think that beyond ubatch of 2048 no setup actually improves (at least I am not aware of one).

[fairydreaming]

apply causality on top-k indices with some new OP (masked indices would be -1 or something)

It would work, but that specific part might be complicated because the causality-related information is on the host side in the llama_kv_cells structure, which is not easy to pass to the backends. There could be some simplified causality logic where it just masks based solely on the indices (i.e. assuming that they are correlated with the sequence position of the data), but this IMO is quite limiting - cannot work with unified KV cache and not compatible with multi-modal cases where sometimes the causality is not simpler triangular matrix. Apart from this, the approach could work.

@ggerganov I see, thanks for sharing your opinion.

Regarding the memory concern - if the attention is not compressed, then the KV cache would already be an enormous amount of memory, multiple times bigger than the F16 KQ mask. Plus I think that beyond ubatch of 2048 no setup actually improves (at least I am not aware of one).

OK, maybe my setup is a bit special since I use -cmoe all the time. For example with Kimi K3 that I'm testing now:

$ ./bin/llama-batched-bench -m /mnt/md0/models/Kimi-K3-Q2_K.gguf -b 8192 -ub 2048 -npl 1 -npp 8192 -ntg 128 -fa 1 -cmoe --no-repack -c 65536
0.00.507.719 W load: special_eos_id is not in special_eog_ids - the tokenizer config may be incorrect
0.00.556.050 W llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --no-mmap for better performance

llama_batched_bench: n_kv_max = 65536, n_batch = 8192, n_ubatch = 2048, flash_attn = 1, is_pp_shared = 0, is_tg_separate = 0, n_gpu_layers = -1, n_threads = 32, n_threads_batch = 32

|    PP |     TG |    B |   N_KV |   T_PP s | S_PP t/s |   T_TG s | S_TG t/s |      T s |    S t/s |
|-------|--------|------|--------|----------|----------|----------|----------|----------|----------|
|  8192 |    128 |    1 |   8320 |  162.743 |    50.34 |   17.982 |     7.12 |  180.725 |    46.04 |

$ ./bin/llama-batched-bench -m /mnt/md0/models/Kimi-K3-Q2_K.gguf -b 8192 -ub 4096 -npl 1 -npp 8192 -ntg 128 -fa 1 -cmoe --no-repack -c 65536
0.00.501.776 W load: special_eos_id is not in special_eog_ids - the tokenizer config may be incorrect
0.00.549.747 W llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --no-mmap for better performance

llama_batched_bench: n_kv_max = 65536, n_batch = 8192, n_ubatch = 4096, flash_attn = 1, is_pp_shared = 0, is_tg_separate = 0, n_gpu_layers = -1, n_threads = 32, n_threads_batch = 32

|    PP |     TG |    B |   N_KV |   T_PP s | S_PP t/s |   T_TG s | S_TG t/s |      T s |    S t/s |
|-------|--------|------|--------|----------|----------|----------|----------|----------|----------|
|  8192 |    128 |    1 |   8320 |   93.468 |    87.64 |   17.977 |     7.12 |  111.445 |    74.66 |

$ ./bin/llama-batched-bench -m /mnt/md0/models/Kimi-K3-Q2_K.gguf -b 8192 -ub 8192 -npl 1 -npp 8192 -ntg 128 -fa 1 -cmoe --no-repack -c 65536
0.00.521.082 W load: special_eos_id is not in special_eog_ids - the tokenizer config may be incorrect
0.00.568.962 W llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --no-mmap for better performance

llama_batched_bench: n_kv_max = 65536, n_batch = 8192, n_ubatch = 8192, flash_attn = 1, is_pp_shared = 0, is_tg_separate = 0, n_gpu_layers = -1, n_threads = 32, n_threads_batch = 32

|    PP |     TG |    B |   N_KV |   T_PP s | S_PP t/s |   T_TG s | S_TG t/s |      T s |    S t/s |
|-------|--------|------|--------|----------|----------|----------|----------|----------|----------|
|  8192 |    128 |    1 |   8320 |   57.457 |   142.58 |   18.062 |     7.09 |   75.519 |   110.17 |

[ggerganov]

Yes, the -cmoe case actually would benefit from larger ubatches - that's true.

Anyway, yes it's a drawback to have a large mask. But the architecture to support that is already in place - it won't hurt to try to take the most of it. The sparse integer mask would likely require quite a few extra logic in the libllama - unless I am not seeing a simple approach to accommodate the existing llama_kv_cache to it.