[FEA]Request for Clarification and CPU Benchmark on HSTU Backward GQA Implementation

[ml1396588509-ops]

Regarding the GQA implementation of the HSTU backward operator, does it ensure consistency in the number of heads (nhead) between input and output through intra-group reduction? Additionally, as there is currently no corresponding CPU reference implementation for the backward GQA, could it be supplemented?

[shijieliu]

hi @ml1396588509-ops what do you mean by ensure consistency in the number of heads (nhead) between input and output through intra-group reduction? could you elaborate?

For the CPU refernece implementation, we dont have the plan to support it. Maybe you can refer the PyTorch HSTU impl and leverage pytorch autograd to achieve your goal?