Question: why TMA is not used

[Idonthaveaname-wq]

I read the code of flash MLA, and I have some questions:

1. why not use tma to load QK, but use the SM80 copy_async.
2. To store data from register to global memory, it uses shared memory to change the layout and then read to register and then write to global.
   As far as I know, hopper supports data transfer from SHM to global directly. why is that used?

I don't have deep insight into cutlass and hopper, so curious about that. Thank you!

[lygztq]

For question 1, IMO, TMA can only load tiles of data with fixed block size, while cp.async can be used with a given cp_size. If you are using kvcache layout where the cache shape in seqlen (max_cache_len) is far beyond the actual seqlen in your request, the OOB behavior in TMA cannot be leveraged to mask out some NaN values in your cache tensor. And current implementation cannot mask out NaN values in tensor V by mask generated from softmax op (since 0 * NaN = NaN)