Missing sparse backward kernel.

[EricLina]

Thank you for the excellent work on the highly optimized attention kernels in FlashMLA. The performance of the sparse forward kernel flash_mla_sparse_fwd for the prefill stage is particularly impressive.

I've noticed that while the library provides a complete forward and backward pass for dense attention, but the sparse attention implementation appears to be forward-only. Will the sparse attention backward kernel be released?

Thanks.

[sunying2018]

same question, thanks!

[IISuperluminaLII]

This would take MUCH MUCH more amount of time to implement than using dense backwards. Why? the main point is to use CUTLASS gemm optimizations. Additionally sparse backwards needs atomic accumulation which is quite difficult to impediment (Thus CUTLASS issue)

If the Qs and Ks are aligned then sure, not the case for SPARSE backwards right... so since we are dealing with "misaligned" Ks and Vs the wall-time is exponentially if not MUCH larger(STRICTLY SPEAKING ON SPARSE COMPUTATIONS)

[IISuperluminaLII]

I mean CUTLASS is the reason they probably don't ship windows whl or source with windows links...

If I had a PhD and free time (2-6 months + 2 years of cuda classes) then maybe I would consider implementing it, but I have neither. I was just curious if anyone had tried it lol