Add a differentiable CPU backend for KDA research and prototyping

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.

Start by reviewing the existing CUTLASS CUDA path and the proposed native PyTorch and C++/ATen CPU entry points described in the issue. The work is complete when the CPU backend supports autograd, variable-length sequences, grouped value heads, gates, initial/final states, and CPU-only or compiler-free installation modes while leaving the CUDA path unchanged.