Can MHA be used in the DSA prefill stage?

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?

Start by tracing the sparse flash MLA implementation for the DSA prefill stage and compare the current MQA formulation with the proposed MHA approach. Determine whether MHA can reduce computational load without changing the required attention results; done means the feasibility and measured trade-offs are documented or implemented with validation.

Written by the indexing model from the issue text.