feat(kda): add bitwise-aligned SM100 CuTeDSL forward

[icavan]

Description

This PR adds an SM100 CuTeDSL implementation of the KDA forward intra and recompute-W/U path.

The review boundary is a one-to-one port of the csrc API that consumes the same precomputed FP32 gk tensor. It preserves the csrc TF32/FP32 arithmetic order, including the single-accumulator K=32 Schur product, and is bitwise equal to csrc for the complete Aqk, Akk, KG, W, and U tensors.

Included components:

• csrc-boundary CuTeDSL intra for equal-length inputs;
• standalone FP32-workspace lower-triangular Akk inverse;
• csrc-aligned persistent recompute-W/U for BF16 and FP32 beta;
• equal-length, packed-uniform, and varlen experimental raw-gate routes;
• strict correctness tests, same-input benchmarks, and CUDA-Graph determinism harnesses.

The existing public csrc dispatch remains unchanged. The CuTeDSL specialization is currently limited to SM100, K=V=128, and chunk size 64.

Correctness

Tested on NVIDIA GB200 (SM100):

pytest -q tests/test_kda_sm100_intra_fused_cutedsl.py -x
pytest -q tests/test_kda_sm100_recompute_wu_cutedsl.py -x

Results:

• intra plus recompute-W/U: 5 passed;
• standalone recompute-W/U: 3 passed;
• BF16 and FP32 beta specializations are both covered;
• complete Aqk/Akk/KG/W/U tensors are compared with torch.equal, without numerical tolerance;
• cached output buffers are poisoned with NaNs and rerun to verify complete deterministic overwrite.

Bitwise results at the representative B=2,T=8192,H=64,K=V=128 shape:

Output	torch.equal	Mismatched elements	Max absolute difference
Aqk	true	0	0
Akk	true	0	0
KG	true	0	0
W	true	0	0
U	true	0	0

Performance

Complete csrc-boundary intra, Akk inverse, and recompute-W/U on the same precomputed FP32 gk input, with three warmup iterations and 20 CUDA-Event-timed iterations:

Shape	csrc	CuTeDSL	csrc / CuTeDSL
B=2, T=8192, H=64, K=V=128	0.9157 ms	0.8645 ms	1.059x

The standalone recompute-W/U path at the same representative T=8192 shape measures 0.3261 ms for csrc and 0.3129 ms for CuTeDSL, or 1.042x csrc throughput.

Appendix: 10,000,000-replay determinism and csrc alignment

The stress harness first requires complete bitwise equality with csrc, then captures CuTeDSL intra, Akk inverse, recompute-W/U, and a comparison of every Aqk/Akk/KG/W/U element against the csrc outputs in one CUDA Graph. Every replay therefore validates the full output, not a checksum or sample.

python \
  benchmarks/stress_kda_sm100_csrc_boundary_determinism.py \
  --iterations 10000000 --checkpoint 1000000 \
  --report-json /tmp/kda_sm100_csrc_boundary_10m.json

Stress configuration and result:

• device: NVIDIA GB200 (SM100);
• shape: B=1,T=256,H=4,K=V=128, BF16 beta;
• iterations: 10,000,000;
• exact element mismatches accumulated across all replays: 0;
• elapsed: 271.881 seconds;
• throughput: 36,780.9 iterations/second;
• status: passed.

Before replay, Aqk, Akk, KG, W, and U each had zero mismatched elements and zero maximum absolute difference from csrc. FP32 beta is covered by the strict pytest suite; this long replay uses the representative BF16 beta specialization.

Reviewer notes

• The exact boundary deliberately does not fuse gate activation or the chunk scan, because doing so changes the FP32 cumulative-gate reduction tree before TF32 MMA.
• fp32_akk_inv=True is required by the csrc-boundary API. The experimental in-CTA inverse is rejected rather than being included in the bitwise claim.
• The raw-gate fused varlen candidate and its separate determinism data are retained as experimental work and are clearly separated in the status document.
• Detailed source-alignment rationale, performance matrices, and rejected experiments are in docs/kda_sm100_cutedsl_fwd_status.md.

Checklist

• pre-commit run --all-files
• Signed-off commit
• SM100 correctness tests
• csrc bitwise comparison
• 10,000,000-replay determinism stress

[KevinZeng08]

typo: should be benchmarks/bench_kda_sm100_intra_fused.py

[KevinZeng08]

The standalone recompute-W/U path at the same representative T=8192 shape measures 0.3261 ms for csrc and 0.3129 ms for CuTeDSL, or 1.042x csrc throughput.

Does the acceleration come from fusing gate activation + cumsum into chunk_fwd_intra, or just the CuTeDSL speedup?

[KevinZeng08]

LGTM

[icavan]

Does the acceleration come from fusing gate activation + cumsum into chunk_fwd_intra, or just the CuTeDSL speedup?

Gate activation and chunk cumsum are deliberately outside this correctness boundary. The gain therefore comes from the CuTeDSL implementation and scheduling, especially the persistent recompute-W/U pipeline, vectorized KG stores, and removal of unnecessary launch/scheduling overhead. Standalone recompute-W/U alone is about 4.2% faster.