large memory usage

[LzhinFdu]

Thanks for sharing this excellent implementation of ring attention.
Here are my test results on 2*A100 (with nvlink). Judging from the results, the memory usage of ring attention（ring_flash_attn_qkvpacked_func） seems to be very large. This is not as expected. Are there any possible problems?

[GeneZC]

The results of flash attention are somehow amazing... keep an eye on this.

[LzhinFdu]

The results of flash attention are somehow amazing... keep an eye on this.

thanks, I'll check it

[GeneZC]

And after reading the code, I have found that the ring attention should accept already-chunked qkv instead of the whole qkv. That is, qkv should be split into local qkv before fed into the ring attention. This might be a diff.

Not very certain that this is true, though.

[jzhang38]

@LzhinFdu @GeneZC Yeah you need to shard the sequence yourself before feeding them into ring-flash-attention.
I have an implementation here: https://github.com/jzhang38/EasyContext

[LzhinFdu]

And after reading the code, I have found that the ring attention should accept already-chunked qkv instead of the whole qkv. That is, qkv should be split into local qkv before fed into the ring attention. This might be a diff.

Not very certain that this is true, though.

That's right. Therefore, when comparing memory usage, the context used by Flash Attention should be doubled. Despite this, Flash Attention still maintains a significant lead.