Apply offfset to KQ_max in CUDA flash attention

[ikawrakow]

This is to avoid (or at least reduce the chance for) overflow when accumulating the flash attention result as f16. Mainline have been doing this for a while. Basically the idea is that we add an offset $\Delta > 0$ to the running maximum of K*Q product. This offset maximum is then used when computing softmax. The factor of $e^\Delta$ that we have introduced that way cancels out at the end when we divide the final result by the sum of the row elements. As $\Delta > 0$, we reduce the magnitude of the K*Q product by $e^{-\Delta}$, so that the accumulated V*softmax(K*Q) is lower by that factor, thus reducing the risk for f16 range overflow.

I have been resisting it because the technique leads to reduced precision of the FA calculation. But I have now encountered a case where we get NaNs without a suitable offset. I was playing with MiniMax-M2 and by mistake started a perplexity calculation for a context of 64k tokens, which netted NaN values. Applying $\Delta = \ln 4$ removes the NaNs. Mainline uses $\Delta = \ln 8$, but I decided to go with $\ln 4 for now.

I know @Nexesenex will not be happy about this change as he is watching PPL values to the 4th significant digit, and PPL values do go up with this change. As an example, for Unsloth's MiniMax-M2.1-UD-Q2_K_XL model with context of 512, PPL goes up from 8.8240 on the main branch to 8.8450 with this PR. Mainline's PPL for this model is 8.8501.

[firecoperana]

I'm not concerned with PPL values, but for those who care, can we make this either opt-in, opt-out or as a parameter?

[Nexesenex]

@ikawrakow: I'm indeed looking raptly to unexpected PPL bumps beyond the realm of what's reasonable! Here, it's expected and amounts to +0.25% (which is the order of magnitude between IQ5_K to Q5_K and Q5_K to Q5_0. It's a bit high, but better than a NAN, that's for sure!

Now, if the NANs are observable in one case on one model at 64k context, maybe, and I say maybe, could it be useful to add the offset as an opt in like @firecoperana suggested, with an optional trigger at a given ctx length reached in the KV cache, if the amount of chores required to add that is reasonably low?

[ikawrakow]

@Nexesenex The model I'm running is Q2_K, so 0.25% change there does not make such a big difference. If instead I compare LlaMA-3-70B quantized with Q4_0, the difference is 0.02% (4.0567 vs 4.0574).