milnor_gpu: fix >2^32-element master via cubecl 64-bit addressing

[JoeyBF]

Problem

The batched Milnor multiply silently corrupted results once the shared masks admissible-master crossed 2³² u16 elements (~8.6 GB, around S_2 stem ~160–180). The buffer length and per-R offsets were u32, so cubecl's default U32 address_type truncated them (len as u32) → out-of-range reads → dx != 0 (d²≠0) panics, e.g. deterministically at (180,92). (160,82) stayed just under the ceiling and was clean.

Diagnosis (instrumented): masks reached 4,452,195,742 u16 (8.29 GiB) > RESIDENT_MAX_CAP = 2³²−1; the ArrayArg length metadata also truncated (u32). Confirmed not the row-reduce, not the resident basis, not a cross-stream race — pure u32 addressing overflow.

Fix — idiomatic cubecl 0.10 64-bit addressing (no hand-rolled buffer splitting)

cubecl has AddressType {U32, U64}; a plain #[cube(launch)] defaults to U32. Two CUDA-backend quirks shaped the fix:

1. address_type = "dynamic" miscompiles for us — it picks U32 for small blocks and then narrows the Array<u64> offset arrays on read (usize::cast_from(u64) under u32 usize), corrupting results. So the multiply is static "u64".
2. checked-mode + U64 won't compile — the bounds clamp emits min(u64, u64), ambiguous for NVRTC. So it uses launch_unchecked (safe: every access is in-bounds by construction — uploaded need_* prefix + per-column j guards).

Changes (all in the batched multiply path):

• multiply_batch_kernel → #[cube(launch_unchecked, address_type = "u64")].
• RInfo.cs_off/mk_off u32→u64; r_cs_offset/r_mk_offset bound as Array<u64>.
• resident grow copies (copy_into_u16/_u32) → #[cube(launch_unchecked, address_type = "dynamic")].
• drop RESIDENT_MAX_CAP + clamps; assert out_len ≤ u32::MAX loudly.

Validation (H200)

• (160,82): rc=0 clean, 517 s vs 486 s baseline (+6.4% for static-u64 multiply, partly offset by unchecked dropping the clamp — no regression).
• (180,92): ran 1500 s with zero dx panics (was a deterministic panic ~130–180 s), i.e. correct well past the 2³² masks boundary. A to-completion rc=0 run is in progress.

🤖 Generated with Claude Code