tests : add -INF blocks to the KQ mask in the FA tests

[ggerganov]

ref #16372 (comment)

The mask that we use in the FA tests is always filled with non-INF values, while in practice it often contains -INF blocks due to causal masking or due to padding. Extend the tests to add -INF values so we can exercise any logic that relies on detecting such blocks in the mask.

[Green-Sky]

I forgot to mention this, but it is possible that cuda fails this test.

[ggerganov]

@Green-Sky Do you mean you tried it locally and it failed?

[Green-Sky]

I just tried it locally and it works. I will have to go and look at the sd.cpp+chroma+fa bug again, so I actually know what fails there.

[Green-Sky]

Ok, did some more debugging. I am not sure if the issue really is with ggml, but flash attention just seems to die when we use the mask in the DiT for chroma.

The mask is:

• 0 in [0;19] (in my example)
• -inf in [20;511]
• and then padded with 0 again for [512;2815]

attention_ext L_q:2816 L_k:2816 n_head:24 C:3072 d_head:128 N:1

[ggerganov]

Does it show any error? What do you mean by "die"?

[Green-Sky]

Does it show any error? What do you mean by "die"?

Ah, sorry for that.
By "die", I mean silently fail. Pulling data from a backend is kind of a pain, but the resulting image is always black, so I assume it it pulled to (+/-) INF somewhere, or maybe zero.

[jeffbolznv]

I'm seeing some failures in the vulkan path with the current change. I'll look into it.

[jeffbolznv]

I'm seeing two classes of failures with vulkan. In the coopmat2 path, the backend will generate some negative infinity results, e.g.:

[FLASH_ATTN_EXT] inf mismatch: Vulkan0=-340282346638528859811704183484516925440.000000 CPU=0.019682   FLASH_ATTN_EXT(hsk=128,hsv=128,nh=4,nr23=[16,1],kv=512,nb=1,mask=1,sinks=0,max_bias=8.000000,logit_softcap=10.000000,prec=def,type_KV=q8_0,permute=[0,1,2,3]): FAIL

I can workaround this by clamping in the shader. But more concerning is cases where the CPU backend generates NaN:

[FLASH_ATTN_EXT] NaN at index 18432 (Vulkan0=nan CPU=-nan(ind))   FLASH_ATTN_EXT(hsk=80,hsv=80,nh=4,nr23=[4,1],kv=512,nb=35,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_KV=q8_0,permute=[0,2,1,3]): FAIL

It's easy to trigger this - just increase the frequency of the -inf blocks by 100x:

-    const int n_inf_blocks = 0.1*(ne0*ne1*ne2*ne3)/(blck0*blck1);
+    const int n_inf_blocks = 10*(ne0*ne1*ne2*ne3)/(blck0*blck1);

I imagine this may be happening when a whole row has -inf. I assume this will reprodue with other backends too, but I haven't verified it. It's not clear to me what the correct result is in this case.

[ggerganov]

I imagine this may be happening when a whole row has -inf. I assume this will reprodue with other backends too, but I haven't verified it. It's not clear to me what the correct result is in this case.

Yes, I have this fixed for CPU and Metal in: fdedda0

I think the correct thing to do is return 0 in this case.

[ggerganov]

With ad83f06, the CPU backend should no longer generate NaNs.

[jeffbolznv]

After #16447, the vulkan backend should pass the updated tests.