Closed yihaocs closed 4 months ago
same problem on 8*A100 flash-attn 2.5.6 triton 2.2.0 torch 2.2.2
Interesting, we have tested this in A100.
But this is because the block size is too large so it does not fit in SRAM. You can reduce the block size to make it fit: https://github.com/RulinShao/LightSeq/blob/main/lightseq/lightseq_async_attn.py#L252 e.g. reduce this to 64, 64. If you encounter the same problem in backward, you can reduce the backward block size as well: https://github.com/RulinShao/LightSeq/blob/main/lightseq/lightseq_async_attn.py#L342
Thanks! But looks weird if we can get it to work for A100 not for H100.
Hi,
Thanks for the great work and the code. I tried to run the test on H100, but it gave me the following error:
The Triton and pytorch packages I am using the latest nightly version installed from Pytorch website, and flash-attn is installed from source:
Do you have any idea what the problem could be?
Thanks!