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

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!

Start with the README and the API for fla.ops.kda.chunk_kda, then inspect how the current inference-mode forward path handles K/V head dimension 128. The issue does not name implementation files, tests, or a defined acceptance condition; completion would require an agreed roadmap or concrete scope for additional dimensions and backward/autograd support.