Skip to content

Add differentiable CPU backend for KDA research - #27

Open
xy200303 wants to merge 1 commit into
MoonshotAI:masterfrom
xy200303:feat/cpu-training-backend
Open

Add differentiable CPU backend for KDA research#27
xy200303 wants to merge 1 commit into
MoonshotAI:masterfrom
xy200303:feat/cpu-training-backend

Conversation

@xy200303

Copy link
Copy Markdown

Summary

  • add a differentiable recurrent KDA backend with a portable PyTorch reference path and an optional C++/ATen CPU operator
  • provide analytical first-order CPU gradients, variable-length sequence handling, grouped value heads, FLA-style gates, and initial/final states
  • support CPU-only and compiler-free installations while preserving the existing CUTLASS CUDA inference backend
  • document the CPU backend as a research and prototyping path for integrating KDA into conventional PyTorch models and exploring model innovations

Verification

  • FLASH_KDA_BUILD_CUDA=0 FLASH_KDA_BUILD_CPU=1 python setup.py build_ext --inplace
  • python -m pytest tests/test_cpu.py -q with the compiled extension: 7 passed
  • python -m pytest tests/test_cpu.py -q without the extension: 5 passed, 2 skipped
  • git diff --check

Notes

  • the intended use is research, model prototyping, architecture experimentation, correctness validation, and small workloads
  • the C++ implementation currently composes ATen tensor operations; it is not a fused AVX/NEON kernel and does not claim production-grade long-sequence CPU throughput
  • the analytical C++ backward supports first-order gradients; higher-order gradients can use use_cpp_backend=False
  • Linux and the minimum supported PyTorch 2.4 still need CI coverage; current compiled validation used Windows and PyTorch 2.12.0+cpu
  • native CUDA backward remains future work; the existing CUTLASS path remains the high-performance CUDA inference backend

Closes #26

@xy200303
xy200303 marked this pull request as ready for review July 29, 2026 09:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add a differentiable CPU backend for KDA research and prototyping

1 participant