Can MHA be used in the DSA prefill stage?

[starwang1024]

The current sparse flash MLA implements the MQA version, where the last dimension of q and k increases from 192 to 576, adding computational overhead. Is it possible to replace MQA with MHA to reduce the computational load of sparse flash MLA during the prefill stage?

[llllvvuu]

I'm also curious about this. Could we do something like gather at topk indices, decompress only those KV, and run the standard varlen dense MHA?