[Question] Roadmap for supporting additional K/V head dimensions and backward kernels

[dmm19941210]

Hi, thanks for open-sourcing FlashKDA!

I am currently evaluating it as a backend for fla.ops.kda.chunk_kda.
From the README and API, my understanding is that the current implementation is mainly targeting inference-mode forward execution, and seems to assume K/V head dimension = 128.

I have a couple of questions regarding the roadmap:

1. Are there any plans to support K/V head dimensions other than 128?
   For example, we have seen some KDA configurations where the per-head K/V dimension is 64 rather than 128. We are not specifically asking only about D=64, but more generally whether FlashKDA is expected to support additional common head dimensions in the future.
   If support for other dimensions is planned, would it likely be implemented as separate CUDA/CUTLASS specializations, or as a more general implementation?
2. Are there any plans to add backward/autograd support, or is the project intended to remain an inference-only forward backend?

Thanks again for your work, and have a great weekend ahead!

[esp0r]

1. The current kernel is tuned for D=128. If we are supporting D=64 in the future, it will be a separate CUDA/CUTLASS specialization
2. There are plans for backward support. The backward kernel is still under development.

[dmm19941210]

Thanks for the clarification! Looking forward to hearing any good news from your side.

[hjlee1371]

Hi, is there any update on backward pass support? When it comes to implementing the backward pass, I would also appreciate some insights into the numerical design decisions—especially regarding FlashKDA's choices like the lower_bound and the bf16 state matrix. Thank you!