trainer: parallel build_cdf scan + tiled (bitwise-identical) mingru backward

[Infatoshi]

What this is

Two env-agnostic optimizations to the compiled trainer, split out of #614 so each reviews on its own merits. Both come with parity proofs against the existing kernels; no training semantics change.

1. Parallel build_cdf (src/pufferlib.cu)

The prioritized-replay CDF was a <<<1,1>>> kernel — one thread serially prefix-summing B floats, once per minibatch update. At B=8192 that's 90us x 32 updates = ~2.9ms per train call of pure serialization; at B=262144 it's 2.85ms per update. Replaced with a two-level block scan using a preallocated block-sums buffer (registered in PrioBuffers): 7.5us at B=8192, 7.2us at B=262144, matching the serial result to 2e-7 rel, monotonic, deterministic, still allocation- and sync-free. A host assert documents the two-level limit (B <= 1M). This is the same fix the 5.0 branch makes via cub::DeviceScan in curriculum.cu — backported to 4.0's pipeline.

2. Tiled mingru_scan_backward (src/models.cu)

The BPTT backward launched B*H threads each walking T serially — at minibatch 8192 (B=64, H=256) that's 16k threads on a GPU that wants ~20x more, and the kernel sits on the critical path 3 layers x N updates per train call. The tiled variant assigns one block per (batch row, 16-wide hidden tile), recomputes the T/4 chunks in parallel from the checkpoints the forward already writes, and replays the cheap carries in reference order in-block — so gradients are bitwise identical to the original kernel (the affine chunk composition is mathematically clean but fp32 reassociation moves near-cancelled gate gradients; keeping reference order costs little and removes all numerical risk). 4x faster at B=64, 1.3x (bf16) at B=256. Originals kept as _ref for the parity harness. Plus a small fusion of scan-adjacent elementwise ops (193 -> 186 nodes per captured update).

Measured end-to-end

On the craftax env at h256x3, 8192 envs, horizon 128 (measured on the #614 GPU env; the patches themselves are env-independent):

minibatch	before	after
32768 (32 updates)	8.09M	8.64M train SPS
8192 (128 updates)	3.92M	4.81M train SPS

Gains scale with updates-per-batch since both patches attack per-update serial cost.

Notes for 5.0

• build_cdf is already fixed there (CUB); this brings 4.0 to parity.
• The tiled backward ports directly — 5.0's PrefixScan/kernels are structurally identical; only the scratch-tensor registration needs reapplying.
• Two measured dead ends worth recording so they don't get re-attempted: a warp-specialized persistent fused update kernel is 3.6-4.8x slower than the current chain (cudagraphs already remove the launch overhead it targets; cuBLAS beats in-kernel WMMA at 256-wide shapes), and further launch-fusion beyond this PR measured neutral-to-negative (cublasLt epilogues drift fp32 by 3e-3; H-specialized launch bounds were slower).

🤖 Generated with Claude Code