[KDA] Add FlashKDA CUDA training backend for chunk_kda

[xy200303]

Summary

Adds a CUDA training backend (flash_kda_train) for chunk_kda, dispatching the full fwd+bwd training step to the FlashKDA CUDA training kernels (WY representation, replicating the Triton pipeline stage by stage). The backend is opt-in (FLA_FLASH_KDA_TRAIN=1, default off) and conservatively gated: it only accepts grad-enabled bf16 calls with K=V=128, H=HV, chunk_size=64, no CP / no intermediate-states; everything else falls back to the Triton path.

Depends on #1111 — without the dispatch fix, top-level chunk_kda never reaches backend selection, so this backend would be dead code. The branch currently contains that commit; only the second commit ([KDA] Add FlashKDA CUDA training backend for chunk_kda) is new, and I will rebase once #1111 merges.

Kernel implementation: MoonshotAI/FlashKDA#28 (per-stage CUDA kernels, 105 stage/pipeline tests, fp64-gold precision sweep). Until that lands upstream, the package is installable from the fork branch:

pip install git+https://github.com/xy200303/FlashKDA.git@kda-train

The backend is discovered via find_spec("flash_kda_train_C"), so any distribution of the package works; nothing in fla pins the source repo.

Test plan

• tests/ops/test_kda.py — 88/88, including 5 new dispatch tests (dense/varlen × gate/safe_gate, fwd + all gradients) that spy-verify calls genuinely route to the CUDA backend rather than silently falling back
• scripts/find_dependent_tests.py dependents of fla/ops/kda/chunk.py were run; the 4 D64 modeling failures are pre-existing on a clean base (unrelated)
• FlashKDA side: tests/train/ 105/105 (per-stage vs Triton reference + end-to-end pipeline)

Benchmark

Hardware: RTX 5090 (sm_120a), CUDA 12.8, torch 2.8, bf16. End-to-end chunk_kda via benchmarks/ops/run.py (dispatch-verified routing):

mode	B	T	H	D	Triton (ms)	CUDA (ms)	speedup
fwd	1	8192	96	128	5.37	5.72	0.94×
fwd	2	16384	16	128	3.66	3.73	0.98×
fwd	4	2048	16	128	0.84	0.90	0.93×
fwd	4	4096	64	128	6.89	7.21	0.96×
fwdbwd	1	8192	96	128	19.71	19.12	1.03×
fwdbwd	2	16384	16	128	13.43	12.54	1.07×
fwdbwd	4	2048	16	128	6.58	3.04	2.17×
fwdbwd	4	4096	64	128	25.18	24.47	1.03×

fwd-only is slightly slower by design (the training path persists Aqk/Akk for backward); training steps are faster on every shape, up to 2.17× on small-grid shapes where the CUDA kernels' finer V-dim split wins occupancy. Per-stage breakdown and roofline analysis: MoonshotAI/FlashKDA#28.

Breaking changes

None. The backend is registered with default_enable=False and a strict verifier, so default behavior is bit-identical to today; systems without the flash_kda training package are unaffected.

Checklist

• I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
• I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
• This is not a minor/cosmetic-only PR (typo, formatting, style-only tweaks).
• Dependent tests pass locally or in CI; new behavior is covered by tests where applicable.
• Kernel changes include same-hardware before/after benchmark numbers (dense + varlen where applicable).

[imezx]

hey there, i wonder when will this be merged?

[xy200303]

hey there, i wonder when will this be merged?

I'm not too sure either, it would need the maintainer to test it. Do you need anything from this PR?

[imezx]

Do you need anything from this PR?

i just would like to take a try by install from source but it has 6 commits behind

[imezx]

i just tried yesterday and let it run for 24 hours, and it seems works fine. nothing breaks.

[xy200303]

i just tried yesterday and let it run for 24 hours, and it seems works fine. nothing breaks.

thanks you