Motivation
FlashKDA currently focuses on high-performance CUTLASS CUDA inference. Researchers who want to introduce KDA into conventional PyTorch models also need an accessible CPU path for architecture exploration, correctness validation, unit tests, and small-scale experiments.
A differentiable CPU implementation would make KDA easier to use as a reusable model component during research, including environments where Hopper-class GPUs or a CUDA toolchain are unavailable.
Proposed scope
- Add a native PyTorch recurrent KDA reference implementation with autograd support.
- Add an optional C++/ATen CPU recurrent operator with an analytical first-order backward.
- Support arbitrary key/value dimensions, grouped value heads, variable-length sequences, FLA-style gates, and initial/final states.
- Keep the existing CUTLASS path unchanged as the high-performance CUDA inference backend.
- Allow CPU-only and compiler-free installation modes.
Intended use
This backend is primarily for research, model prototyping, model-architecture innovation, correctness checks, and small workloads. It is not intended to claim production-grade CPU throughput: the proposed C++ implementation composes ATen operations and is not yet a fused AVX/NEON kernel.
Higher-order gradients can use the pure PyTorch reference implementation. A native CUDA backward remains separate future work.
I have an implementation ready and will link the pull request to this issue.
Motivation
FlashKDA currently focuses on high-performance CUTLASS CUDA inference. Researchers who want to introduce KDA into conventional PyTorch models also need an accessible CPU path for architecture exploration, correctness validation, unit tests, and small-scale experiments.
A differentiable CPU implementation would make KDA easier to use as a reusable model component during research, including environments where Hopper-class GPUs or a CUDA toolchain are unavailable.
Proposed scope
Intended use
This backend is primarily for research, model prototyping, model-architecture innovation, correctness checks, and small workloads. It is not intended to claim production-grade CPU throughput: the proposed C++ implementation composes ATen operations and is not yet a fused AVX/NEON kernel.
Higher-order gradients can use the pure PyTorch reference implementation. A native CUDA backward remains separate future work.
I have an implementation ready and will link the pull request to this issue.