Flash Attention without Triangular Mask

[pearlli98]

I am interested in implementing a version of flash attention in triton without masking.

I infer from the pytest in the tutorial that the flash attention example generate outputs that are masked with a triangular matrix. I also infer from the code that it would be this specific line inside the kernel that apply the triangular mask.

However, when i remove both of these operations, the outputs from triton and pytorch are not the same. Does anyone know how I can change the code to make this happen? Thank you!

[ptillet]

You also need to change the bounds of the loop to make sure each query iterates through all the keys

[pearlli98]

@ptillet Thank you for your response! I changed the for loop bound to for start_n in range(0, N_CTX, BLOCK_N): and right now the results seem to assign for the most part. I have some follow-up questions on this:

1. Does the change correspond to what you meant by "iterating through all the keys"?
2. I do see some numerical differences on the forward pass output for up to 0.0015. Do you think this range is expected?

Thank you so much :)

[ptillet]

Yes, and yes :)