CUDA: enable sparse fa for qwen4

[am17an]

Overview

Cont #27970. Enable sparse-fa for Qwen4. This model's attention is un-optimized at the moment, the entire kv-cache is re-scored everytime, we should fix that as well.

How this works ncols1=8 is take all the union of the tokens being used, so at max it will use see ncols1 * n_kv_max tokens, at twice this value we enable the sparse FA. For Qwen4 this value is 32768 ctx.

Additional information

Results on a DGX spark:

CPU	Model	Test	t/s qsa-blk-cache	t/s aman/qwen-sparse	Speedup
CPU	qwen4exp A3B IQ1_S - 1.5625 bpw	pp2048@d10000	615.70	663.17	1.08
CPU	qwen4exp A3B IQ1_S - 1.5625 bpw	pp2048@d20000	500.49	544.11	1.09
CPU	qwen4exp A3B IQ1_S - 1.5625 bpw	pp2048@d50000	405.18	469.03	1.16
CPU	qwen4exp A3B IQ1_S - 1.5625 bpw	pp2048@d100000	252.81	318.28	1.26
CPU	qwen4exp A3B IQ1_S - 1.5625 bpw	tg32@d10000	22.98	23.56	1.03
CPU	qwen4exp A3B IQ1_S - 1.5625 bpw	tg32@d20000	20.76	23.58	1.14
CPU	qwen4exp A3B IQ1_S - 1.5625 bpw	tg32@d50000	15.88	18.72	1.18
CPU	qwen4exp A3B IQ1_S - 1.5625 bpw	tg32@d100000	12.00	14.19	1.18

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES

[JohannesGaessler]

I don't really have anything to add. The implementation seems straightforwardly correct to me and is done well in terms of software architecture. One thing to keep in mind is that this is in essence the exact same infrastructure as would be needed for things like PagedAttention.

[JohannesGaessler]

It is good to enable sparse attention for ncols1 > 1 since that will drastically improve the prefill performance. However, if it is only enabled for 1 and 8 it will cause trouble in combination with speculative methods. My opinion is that we should compile the template specializations for batch sizes 2 and 4; if the compilation becomes too bloated we should shave off template specializations somewhere else.

It's also not clear to me why the new code for sparse attention with ncols1 > 1 would work for Qwen 4 but not the other models.

[am17an]

I only tested it with Qwen's headsizes, I guess it work anyway. I think I can expand the condition to be ncols==1, 2, 4, 8 and ncols2=8, 16.

[JohannesGaessler]

One thing I forgot: the optimal ncols1 value for prefill may not be the highest one at large context depths. As long as you're I/O bound then increasing ncols1 is always beneficial since you would need to load the KV data anyways and it doesn't matter if you waste a bit of compute because that's not the bottleneck anyways. But at some point you will become compute bound so reducing wasted compute can become worthwhile. And at large context depths K/V data may be less likely to be shared across one "group" as you called it here. So in that case running the template specialization with ncols1 == 4 may be faster than the one with ncols1 == 8. My intuition is though that on modern NVIDIA GPUs ncols1 == 8 is optimal.