Add intermediate states and exact CP transition prototype

[0z5a]

Summary

• add an optional BF16 intermediate-state buffer that snapshots the recurrent state after every 16-token FlashKDA chunk, including fixed-length and ragged layout helpers
• add exact high-precision affine KDA transition algebra (S' = S @ A + B) with composition, exclusive prefix scans, and a CP=2 NCCL correctness path
• add a fused Triton D=V=128 segment-summary path that keeps each transition row register-resident across chunks, with a PyTorch reference fallback
• cover state dtype dispatch, tail chunks, varlen snapshots, non-power-of-two prefix scans, and distributed start-state reconstruction

H200 results

Measured on 2x NVIDIA H200 NVL with H=8, D=V=128, and 512 tokens per CP rank:

Path	Before	After	Speedup
low-rank transition core	22.146 ms	0.656 ms	33.74x
transition summary including preprocessing	23.056 ms	1.776 ms	12.98x

For CP=2, the current NCCL all-gather plus local prefix composition takes 0.372 ms for a 1 MiB transition per rank. The two-stage FlashKDA comparison against a serial 256-token run has rank-1 output max error 5.72e-4 and mean error 3.27e-6.

Validation

python -m pytest -q tests/test_cp_transition.py
7 passed

python -m pytest -q tests/test_fwd.py -k intermediate_state
2 passed, 4 deselected

torchrun --standalone --nproc-per-node=2 tests/distributed_cp_smoke.py
distributed CP + FlashKDA end-to-end smoke test passed on 2 GPUs

The SM90A extension builds successfully on H200, and git diff --check passes.

Scope and limitations

• the distributed correctness path intentionally uses an O(P) all-gather, which is practical for small CP degrees; a recursive-doubling transport remains a follow-up optimization
• "exact" refers to the high-precision affine recurrence; forcing a BF16 materialization after every chunk is nonlinear and therefore cannot be reproduced bitwise by a fixed (A, B) transition
• this PR provides the forward transition and scan prototype, not a complete training/backward CP runtime