Why block size have to be 128?

[suhmily10]

Is that possible parsing block size as parameter?

[GeneZC7]

same question

[KevinZeng08]

Same question. Could you explain that? @JerryGJX

[JerryGJX]

Oh, that's because FlashAttention tiles tokens in chunks of 64 or 128. By setting the block size to 128, we can skip the main loop at the block level. Otherwise, additional operations would be required at each step. This is not an ideal design, but modifying it would require changes to how data copying and MAA is handled with CuTe. This could be a todo.

[suhmily10]

Oh, that's because FlashAttention tiles tokens in chunks of 64 or 128. By setting the block size to 128, we can skip the main loop at the block level. Otherwise, additional operations would be required at each step. This is not an ideal design, but modifying it would require changes to how data copying and MAA is handled with CuTe. This could be a todo.

is it possible to block in head dim instead of seq dim?

[wrencanfly]

thanks! Can I change it to 64 instead?