TMA loads turned off for paged KV cache kernel

[christindbose]

I noticed that the TMA loading feature is turned off for batched inference during cascade attention. Turning it on results in jit compilation errors. Any idea why this optimization is not enabled for the Hopper?

flashinfer/include/flashinfer/attention/hopper/prefill_sm90.cuh

Line 578 in 7e98d8b

AttentionKernelTraits</*USE_TMA_LOAD_KV=*/false, HEAD_DIM_QK, HEAD_DIM_VO,

[yzh119]

Because we are using LDGSTS (ampere-style cp.async) for page attention, this is designed for small page size (such as 1) because TMA failed to achieve maximum bandwidth for small block of memory access.

For larger page size, yes using TMA is better.