Add CUDA training kernels (fwd+bwd) for KDA

[xy200303]

Summary

Adds a complete CUDA training path (forward and backward) for KDA. The kernels replicate the FLA Triton chunk_kda training pipeline stage by stage (WY representation, CuTe SM80 MMA atoms, cp.async pipelining), so the numerics track the Triton reference closely:

• csrc/train/ — 9 CUDA kernels merged into a single flash_kda_train_C extension: gate cumsum (kda_gate fwd+bwd), intra attention (Aqk/Akk fwd, block-triangular solve, bwd), WY recompute (w/u), chunk o, chunk h fwd + bwd dh/dh0, bwd dAv, fused bwd dq/dk/dg/db.
• flash_kda/train/pipeline.py — end-to-end chunk_kda_train_fwd / chunk_kda_train_bwd orchestration mirroring the FLA recompute path, dense and varlen (cu_seqlens) alike.
• Same single-source/multi-arch build model as csrc/smxx (SM80-era instructions, arch handled via -gencode).

Performance (RTX 5090, bf16, fwd+bwd training step)

Full details and repro command in BENCHMARK_TRAIN_RTX5090.md (this branch).

T=8192, H=96, D=128

Case	flash_kda_train fwd (ms)	fla_chunk_kda fwd (ms)	fwd speedup	flash_kda_train fwd+bwd (ms)	fla_chunk_kda fwd+bwd (ms)	fwd+bwd speedup
Fixed	6.1385	5.4185	0.88×	21.2520	23.4677	1.10×
Varlen, seq_lens=[1300, 547, 2048, 963, 271, 3063]	6.1990	5.4671	0.88×	21.3378	23.0464	1.08×
Varlen, seq_lens=1024 x 8	6.1546	5.4465	0.88×	21.1676	23.0427	1.09×

T=8192, H=64, D=128

Case	flash_kda_train fwd (ms)	fla_chunk_kda fwd (ms)	fwd speedup	flash_kda_train fwd+bwd (ms)	fla_chunk_kda fwd+bwd (ms)	fwd+bwd speedup
Fixed	4.0432	3.5908	0.89×	13.9893	15.3251	1.10×
Varlen, seq_lens=[1300, 547, 2048, 963, 271, 3063]	4.1186	3.5776	0.87×	14.1548	15.0392	1.06×
Varlen, seq_lens=1024 x 8	4.0600	3.5983	0.89×	13.9681	15.1306	1.08×

fwd-only is slightly slower (0.87×–0.89×) by design: the training path persists Aqk/Akk for backward, which pure-inference use does not amortize — inference stays on the existing kernels. Training (fwd+bwd) is faster on every measured case; an FLA-side end-to-end run shows up to 2.17× on small-grid shapes where the CUDA kernels' finer V-dim split wins occupancy.

Stage-level breakdown (CUDA vs Triton, B=2 T=16384 H=16 D=128, 20 reps)

Reproduce with python benchmarks/bench_train_stages.py:

stage	Triton (ms)	CUDA (ms)	ratio
gate_cumsum	0.305	0.297	1.03×
fwd_intra	1.549	1.668	0.93×
recompute_w_u	0.910	0.950	0.96×
fwd_h	0.739	0.700	1.06×
fwd_o	0.712	0.694	1.03×
bwd_dAv	0.428	0.429	1.00×
bwd_dhu	1.127	0.899	1.25×
bwd_wy_dqkg	2.707	1.925	1.41×
bwd_intra	1.604	1.638	0.98×
reverse_cumsum	0.369	0.383	0.96×
gate_bwd	0.759	0.770	0.99×
TOTAL	11.207	10.353	1.08×

The gain comes mainly from the two heavy backward kernels (bwd_wy_dqkg 1.41×, bandwidth-saturated at ~1.47 TB/s; bwd_dhu 1.25×). Stages still below 1.0× (fwd_intra, recompute_w_u) measure at the bandwidth floor (0.94–0.97 TB/s) where the Triton kernels sit too — memory-pattern bound, not scheduling slack.

Precision (vs fp64 gold, fwd + all gradients)

Sweep of the official g/bias configurations (tests/test_fwd.py style), comparing o, ht and all gradients (dq/dk/dv/dg/db/dh0) of both flash_kda_train and fla chunk_kda against a differentiable fp64 naive recurrence. Generated by tests/train/test_train_vs_fla.py (this branch):

Per-position max/mean error, RMSE ratio and error distributions are on par with or below the Triton reference in every configuration (see the summary row at the bottom, e.g. dq: chunk 4.33e-03 vs flash 3.88e-03; dh0: chunk 2.98e-03 vs flash 2.04e-03).

Test plan

• tests/train/ — 105/105 passing (RTX 5090): every stage compared point-by-point against the FLA Triton reference, plus end-to-end pipeline tests (dense + varlen, gate/safe_gate configs, full gradient checks).
• tests/train/test_train_vs_fla.py — fp64-gold precision sweep above (asserts included).

Benchmark

• Hardware: NVIDIA RTX 5090 (sm_120a), CUDA 12.8, torch 2.8, bf16
• Command: python benchmarks/generate_train_benchmark_md.py
• Workload: use_gate_in_kernel=True, lower_bound=-5, post-sigmoid beta, fp32 initial_state, chunk_size=64

Breaking changes

None — purely additive (csrc/train/, flash_kda/train/, tests/train/, benchmark scripts/docs).

[hanxdmech-ship-it]

您好，之前的 提交 显示前向性能优于triton 版本，请问为什么没有基于上述前向的路径做反向设计？因为每个 chunk 中 16 token 不利于反向并行计算，或者上述前向路径不利于反向并行？

[xy200303]

感谢提问。简短回答：与 16-token 分块不利于并行无关——训练路径在 chunk=64 内部同样以 16 为 sub-chunk 做 WY 分块求逆（intra.cu 的 kBC=16，复刻 FLA 的 sub_chunk solve），16 本来就是 WY 求逆内部对 tensor core 友好的粒度。没有基于推理前向路径做反向，是因为两者的设计目标根本不同：

1. 推理前向不产出反向所需的任何中间量。 推理 kernel 是为"零中间量落盘"深度定制的两阶段融合实现：K1 把 decay/cumsum/16×16 求逆结果写入易失 workspace；K2 在 smem 持有 128×128 state 沿 chunk 串行扫描、边走边覆盖，最终只输出 o 和 final_state。而反向（bwd_dhu、bwd_wy_dqkg 等）需要逐 chunk 的 h、v_new 以及 chunk=64 语义的 Aqk/Akk，推理路径一概没有保存；要用它做反向就得按训练需要的粒度重算这些量，等于重写。

2. 数值前提不兼容。 推理选 chunk=16 的根本原因是让 exp(cumsum(g)) 落在 bf16 可表示范围内（配合 lower_bound=-5），省去 FLA chunk=64 需要的 intra-chunk rescaling，state 也是 bf16 累加——这些只对推理成立。训练反向需要 fp32 gate 累积和 fp32 的 per-chunk state，并且本 PR 的验证策略是逐 stage 与 FLA Triton 参考对拍，必须保持 chunk=64 的 recompute 语义才能逐阶段对齐。

3. 并行结构上反向其实延续了推理的思路。 训练路径的 state 扫描阶段（chunk_h fwd / bwd_dhu）与推理 K2 形态相同——(seq, head) 网格 + V 维 split + CTA 内沿 chunk 串行扫描，区别只是每个 chunk 的 h/dh 写回 gmem 供后续阶段使用。反向的主要收益恰恰来自这些重写的 kernel（bwd_wy_dqkg 1.41×、bwd_dhu 1.25×）；若沿用 T=16，state 扫描长度会 ×4，反而不利。

[xy200303]

补充说明：我们也在另一条分支上尝试基于 WY-16 前向思路的训练实现（专用 FP32 训练前向 + chunk 边界检查点 + 窗口化 chunk 并行反向）。初步结果是在短序列上相比 Triton 有明显优势，但长序列上反向随 T 增长退化明显、目前效果不好，还在定位并行度瓶颈。所以本 PR 先走语义与 FLA 完全对齐、各序列长度下都稳定的 chunk=64 路线；WY-16 训练路线待长序列问题解决后再考虑合入。

[hanxdmech-ship-it]

补充说明：我们也在另一条分支上尝试基于 WY-16 前向思路的训练实现（专用 FP32 训练前向 + chunk 边界检查点 + 窗口化 chunk 并行反向）。初步结果是在短序列上相比 Triton 有明显优势，但长序列上反向随 T 增长退化明显、目前效果不好，还在定位并行度瓶颈。所以本 PR 先走语义与 FLA 完全对齐、各序列长度下都稳定的 chunk=64 路线；WY-16 训练路线待长序列问题解决后再考虑合入。

我也尝试基于WY-16 前向思路的训练实现，发现了相同的类似规律： 短序列上相比 Triton 有明显优势，但长序列上反向随 T 增长退化明显。 但是不知道 kimi flash kda 在内部训练是如何实现的？
另外 请问在sm90系列设备上测试验证过吗？

[xy200303]

SM90 还没有真机验证。本 PR 的全部测试与 benchmark 都在 RTX 5090（sm_120a）上完成