Questions about KV loading

[carryyu]

Thanks for your work!

I noticed that in the Sparse implementation, TMA is not used for KV loading. Has there been any performance comparison between this approach and the use of TMA in FA3?

[yzh119]

You can change the codebase to use ampere style LDGSTS for contiguous KV-Cache.

As far as I remember, with CUDA 12.4, using LDGSTS can get around 570+ TFLOPs/s for long context causal attention while using TMA can get you ~600 TFLOPs/s. CUDA 12.3 might get you a higher TFLOPs/s because I haven't tried it by myself.

[carryyu]

谢谢，已收到！