CUDA + ggml: add sparse-fa for DSV4/GLM - #27970
Conversation
ggerganov
left a comment
There was a problem hiding this comment.
Should we also try instead of extracting indices from the mask, to directly construct a dense KV cache (i.e. get_rows-style) and run the existing FA kernels without modifications?
| const bool need_f16_V = type_V == GGML_TYPE_F16; | ||
| constexpr size_t nbytes_shared = 0; | ||
| launch_fattn<D, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false); | ||
| launch_fattn<D, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false, false); |
There was a problem hiding this comment.
Where are the "vec" FA kernels located? Just from the filenames, it looks as if the sparse indices are not used in the "vec" case.
There was a problem hiding this comment.
Yeah we use the mma kernel for the vec (bs=1) case when f16 kv cache is used. I think the vec kernel is only used for quantized kv cache - I may be wrong
| static constexpr bool ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse( | ||
| const int DKQ, const int DV, const int ncols1, const int ncols2) { | ||
| return (DKQ == 512 && DV == 512 && ncols1 == 1 && ncols2 == 8) || | ||
| (DKQ == 576 && DV == 512 && ncols1 == 1 && ncols2 == 16); | ||
| } |
There was a problem hiding this comment.
Do we need to restrict to only these head sizes?
Could you remind me what ncols1 and ncols2 refer to in the CUDA backend?
There was a problem hiding this comment.
ncols1 is the number of query tokens and ncols2 is the k/v tokens per block. It's restricted here because I only ran test this for dsv4, for qwen4 it was slower for prefill and faster for decode but qwen4 has other issues w.r.t to the mask so I left for a future PR
At least for the CUDA backend this was slower than not doing anything at all for dsv4. |
|
I have done a lot of attempts at optimizing sparse attention in DSv4 in a Vulkan focused llamacpp fork, and the biggest improvements was using dense attention for the continuous block of keys, sparse attention for the selected keys and then ultimately combining the result. The commit for this exact change can be found here: Nathanw1014@4bbe53e Then I did some further memory usage improvements by changing the implementation to a tiled version in a follow-up commit here: Nathanw1014@80728f3 These changes were vibe coded with oversight from me, but the idea itself might be useful for optimizing sparse attention. My apologies if this PR already does something similar. I am on my phone so I haven't really look at the code yet, but I thought I would share my experiences with optimizing sparse attention when I saw this PR. |
|
I have a Metal implementation that works with DSv4. But I also want to test it with Qwen4. Do you have a suggestion for a patch that enables sparse attention with the Qwen4 graph? |
|
@ggerganov you can just pass n_kv_max (which should be |
|
I run some tests on my container |
|
Rebased onto current master and benchmarked against it on an RTX PRO 6000 Blackwell, DeepSeek-V4-Flash IQ2_XXS, f16 KV cache, single sequence. Clear win at long context, small cost below 32k, which matches the shape of your own table.
Two passes per binary with the run order alternated, so the long context numbers are not a thermal artifact. One correctness note: in Patch : |
|
@ServeurpersoCom thanks, fixed in be029fd. |
|
Optional nits : Small cleanup patch on top of the PDL fix, nothing behavioural except one extra test case. It names the block size and the two activation thresholds, replaces the serial prefix sum with a warp scan over the eight warp counts, adds a static_assert for the ncols1 == 1 assumption the index addressing relies on, and drops a dead VARS_TO_STR18 plus a test case that was listed twice. The one real change is in the mask generator: it now reduces one row in every thirty two to a single finite entry, which is the degenerate case of the first token of a sequence and was not covered before. Built and ran the full FLASH_ATTN_EXT suite on a Blackwell card, 2948 pass, one fewer than before because of the removed duplicate. Take or leave any of it, none of it is load bearing. Running a greedy check next and I will approve after that. |
JohannesGaessler
left a comment
There was a problem hiding this comment.
Any comments regarding performance are only suggestions.
|
|
||
| // Use finite mask entries as a sparse K/V set. Set 0 to disable. | ||
| // n_kv_max must bound the number of finite entries in every mask row. | ||
| GGML_API void ggml_flash_attn_ext_set_sparse( |
There was a problem hiding this comment.
I don't feel strongly about this but wouldn't ggml_flash_attn_ext_set_n_kv_max be the more appropriate name?
| const int32_t index = i < i_sup ? indices[i] : -1; | ||
| src = index >= 0 ? KV + int64_t(index)*stride_KV + k*h2_per_chunk : zero; |
There was a problem hiding this comment.
It's not great to have a condition here. It would likely be preferable to pad indices with safe values that result in some redundant work.
| // Skip unused kernel variants for faster compilation: | ||
| if (use_logit_softcap && !(DKQ == 128 || DKQ == 256 || DKQ == 512)) { | ||
| NO_DEVICE_CODE; | ||
| return; | ||
| } | ||
| if (DKQ == 192 && ncols2 != 8 && ncols2 != 16) { | ||
| NO_DEVICE_CODE; | ||
| return; | ||
| } | ||
| #ifdef VOLTA_MMA_AVAILABLE | ||
| if (ncols1*ncols2 < 32) { | ||
| NO_DEVICE_CODE; | ||
| return; | ||
| } | ||
| #endif // VOLTA_MMA_AVAILABLE | ||
|
|
||
| #if __CUDA_ARCH__ == GGML_CUDA_CC_TURING | ||
| if (ncols1*ncols2 > 32) { | ||
| NO_DEVICE_CODE; | ||
| return; | ||
| } | ||
| #endif // __CUDA_ARCH__ == GGML_CUDA_CC_TURING | ||
|
|
||
| #if defined(AMD_WMMA_AVAILABLE) | ||
| if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 128) { | ||
| NO_DEVICE_CODE; | ||
| return; | ||
| } | ||
| #endif // defined(AMD_WMMA_AVAILABLE) | ||
|
|
||
| #if defined(AMD_MFMA_AVAILABLE) | ||
| if (ncols1*ncols2 < 16 || DKQ > 256) { | ||
| NO_DEVICE_CODE; | ||
| return; | ||
| } | ||
| #endif // defined(AMD_MFMA_AVAILABLE) |
There was a problem hiding this comment.
Please skip unused kernel templates here, I would suggest you re-use ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse.
There was a problem hiding this comment.
does shall_use_sparse etc already do this?
There was a problem hiding this comment.
No, because shall_use_sparse is a host function function intended for kernel selection logic.
There was a problem hiding this comment.
In terms of program logic shall_use_sparse must be strictly a subset of may_use_sparse.
There was a problem hiding this comment.
Okay it should be there
|
|
||
| const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); | ||
| const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; | ||
| const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt)*ne11 : nullptr; |
There was a problem hiding this comment.
| const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt)*ne11 : nullptr; | |
| const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt*ncols1)*ne11 : nullptr; |
I think the way you're calculating the indices may be incorrect here though this defect would as of right now not manifest as a bug due to ncols1 == 1.
There was a problem hiding this comment.
Sorry, after thinking about it some more I'm not sure what I suggested here is correct. For > 1 tokens it's not clear to me what the best way to handle indices is.
|
Greedy output diverges from master, as expected since compacting the selected entries changes the accumulation order, so I checked perplexity instead. Same 153806 token corpus, run at c=32768, well past the 4096 KV length needed to arm the sparse path:
0.0006 apart against a 0.0045 confidence interval, same direction on all four chunks. No sign of anything dropped. Good for me, I will leave the rest to @JohannesGaessler. |
Overview
Based on @fairydreaming's PR #25917, it is the same idea but instead just relies on a hint from the API about the max number of live kv-entries per token. Further sparse attention methods like QSA(qwen4) can be enrolled at a later stage as they use similar mechanisms. Performance of 2-bit DSV4 quant on DGX spark. More testing will be appreciated
Additional information
Requirements