[icavan]

Summary

• keep the SM90 intracard-CP segment transition matrix canonically in FP32 registers
• retain a BF16 shared-memory mirror only for Tensor Core operands
• preserve the serial FlashKDA BF16 state recurrence and the existing shared-memory footprint
• add the public 64K split sweep plus a transition-sensitive FP64 regression

Root cause and implementation

The pre-scan formed both the local state and the CP-only affine transition through BF16 shared memory. Repeated segment composition can amplify transition rounding when the effective decay is close to non-contracting.

Four compute warps now jointly own the complete 128x128 transition as a distributed FP32 accumulator. Each tile publishes a BF16 mirror for the existing kd @ M Tensor Core product, while the canonical recurrence and final m_seg export stay FP32. This keeps the change local to CP transition composition; serial state semantics are unchanged.

Accuracy

Public issue reproducer, CP relative-L2 vs serial:

splits	output	final state	last 8K output
2	0	0	0
4	5.1775e-3	7.1517e-3	7.3234e-3
8	6.3413e-3	7.0972e-3	7.3068e-3
16	6.8596e-3	7.1433e-3	7.3176e-3

The divergence stays bounded as the split count increases.

On a transition-sensitive 16K FP64-oracle regression, split-16 CP error improves from approximately 1.45e-1 / 5.60e-1 to 5.48e-2 / 1.40e-2 for output/final-state relative L2.

Performance and resources

64K forced split-16 measurements on the SM90 validation environment:

heads	before	after
1	1.147 ms	1.120 ms
8	1.297 ms	1.270 ms

The pre-scan uses 255 registers, 0 bytes local memory, no spills, and the same shared-memory footprint.

Validation

• pre-commit run --all-files
• 64K issue regression at 2/4/8/16 splits
• transition-sensitive FP64 regression
• fixed and varlen CP-vs-serial cases
• initial/final state and transposed-state layouts
• non-aligned dense and varlen lengths
• fixed and varlen determinism checks
• final SM90 fence/memory-order audit

Related to #126

[KevinZeng08]

LGTM

[cherhh]

LGTM.