cuda : fix KQ mask offset integer overflow in flash attention MMA kernel

[fairydreaming]

Overview

When using long context and large ubatch size KQ mask becomes very large during prompt processing with the number of elements exceeding INT32_MAX (for example in GLM-5.1 max context size is 202752, so with 16k ubatch we have 3321888768 KQ mask elements). This results in negative j_vram*stride_mask mask_h array offset in flash attention MMA kernel as confirmed by printfs:

j_sram*(nbatch_fa + 8) = 1600, j_vram*stride_mask = -1976369152
j_sram*(nbatch_fa + 8) = 1600, j_vram*stride_mask = -1976369152
j_sram*(nbatch_fa + 8) = 1600, j_vram*stride_mask = -1976369152
j_sram*(nbatch_fa + 8) = 1600, j_vram*stride_mask = -1976369152
j_sram*(nbatch_fa + 8) = 1640, j_vram*stride_mask = -1976221696
j_sram*(nbatch_fa + 8) = 1640, j_vram*stride_mask = -1976221696
j_sram*(nbatch_fa + 8) = 1640, j_vram*stride_mask = -1976221696
j_sram*(nbatch_fa + 8) = 1640, j_vram*stride_mask = -1976221696

Casting j_vram to int64_t so that the calculated offset is also int64_t resolves this problem.

Fixes #23574.

Requirements

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

[JohannesGaessler]

Suggested change

	const int64_t j_vram = fastmodulo(j0 + j_sram, ne01);

I would say it's better to do the upcast on this line instead because an overflow of a 32 bit signed integer here would be upcast to a negative 64 bit integer.