fix(paged-attention): use 64-bit pool offsets in the fused decode kernels #2153

Problem / Background

The fused paged-attention kernels compute the K/V pool element offset in 32-bit unsigned arithmetic on every backend. Once one pool slab holds more than 2^32 elements, the offset wraps and the kernel reads a different row of the same buffer: no fault, silently wrong attention output. Found during the review of PR #2103 (HIP port, issue #2068), which copied the CUDA arithmetic unchanged.

The wrap point is slab_tokens * Hkv * D > 2^32. For Llama-3.1-8B (Hkv 8, D 128) that is 4,194,304 token slots in one layer's slab (131,072 blocks at block size 32), i.e. 8 GiB per side per layer at f16. The server sizes the slab as ceil(per_slot_ctx / block_size) * batch (resolve_paged_slab_blocks, src/execution/memory_estimate.rs:1385-1430), capped only by the per-layer share of the block budget when one is set, and MLXCEL_PAGED_SLAB_BLOCKS is used verbatim. A large-memory host serving long contexts at high --parallel reaches it, and nothing refuses the launch.

Current Behavior

The pool base offset, (row * block_size + slot) * stride_kv + kv_head * dim, is 32-bit in all six fused bodies:

• v1 decode: Metal uint at src/lib/mlx-cpp/turbo/paged_attention.cpp:142,148 (PAGED_ATTENTION_DECODE_SOURCE), CUDA uint32_t at paged_attention.cpp:295,301, HIP uint32_t at src/lib/mlx-cpp/turbo/paged_attention_hip.h:131,137.
• v2 partial: Metal at src/lib/mlx-cpp/turbo/paged_attention_v2.cpp:168,175, CUDA at paged_attention_v2.cpp:363,370, HIP at paged_attention_hip.h:291,298.
• Sparse decode (src/lib/mlxcel-core/src/paged_v2/sparse.rs:586-589) reshapes the dense [B, H, Cap, D] allocation to [B*H*Cap, 1, 1, D] and launches v2 partial with page size 1, so it wraps at B * H * Cap * D > 2^32 (MiniMax-M3: H = 5 including the index head, D = 128, so B * Cap > 6,710,886).
• Merge (paged_attention_v2_merge.cpp:93 Metal, :138 CUDA, paged_attention_hip.h:420 HIP) indexes v_in[(i * heads + h) * dim + d] in 32 bits. Its input is the partial buffer (chunks * Hq * D), which is far smaller, but it is unguarded.

The launchers' host checks (paged_attention.cpp:513-528, paged_attention_v2.cpp:643-656, paged_attention_v2_merge.cpp:253-263, added by #2103) validate rank and axes 1 to 3 but not the element count.

Proposed Solution

Widen the pool address, not every index. In each of the six fused bodies, compute the per-token base in 64 bits and keep the loop counters 32-bit:

• Metal: ulong base = ((ulong)row * block_size + slot) * (ulong)stride_kv + kv_head * dim; (v2: page_size and entry in place of block_size and slot).
• CUDA and HIP: the same with uint64_t.
• Every read through base (k_pool[base + d], v_pool[base + d]) uses the 64-bit value.

The rows/indices inputs stay int32 (row counts stay far below 2^31). For merge, add a host-side refusal in paged_attention_v2_merge.cpp next to the existing check: throw std::invalid_argument naming the entry point when v_in.size() or the output size exceeds UINT32_MAX; it reaches Rust as Err through the existing Result bridge, and callers already treat a refused launch as fall back to gather.

Rejected: a host-side refusal for the pool kernels. It would send exactly the long-context, large-batch shapes the fused path exists for back to gather. A 64-bit multiply per token per lane is negligible in a memory-bound kernel; the benchmark below confirms it.

Keep the dtype keys and the select_kernel_port launch pattern untouched, so make verify-kernel-dtype-keys verify-kernel-port-dispatch stays green with the same EXPECTED_IN_SCOPE count.

Scope

• In scope: the six fused bodies and the merge host check listed above, tests, and a short note in docs/environment-variables.md on MLXCEL_PAGED_SLAB_BLOCKS that slabs past 2^32 elements are supported.
• Out of scope: the gather fallback, which uses MLX ops with 64-bit indexing, and the dense SDPA path.

Implementation Notes

• Reuse: the existing refusal blocks in the three launchers for the merge size check; src/lib/mlxcel-core/src/test_support/kernel_ports.rs skip helper for backends without a port.
• Constraints: Metal, CUDA and HIP bodies must stay arithmetically identical to each other apart from the type spelling. No change to template arguments, grid shape or JIT names beyond what the source change implies.
• Edge cases: row * block_size must be widened before the multiply, not after; a pool exactly at 2^32 elements must address its last element correctly.
• Error handling: merge oversize returns Err from the bridge with the entry-point name; no new panics.

Acceptance Criteria

• All six bodies compute the pool base in 64 bits; merge refuses v_in or output sizes above UINT32_MAX with an Err.
• A fast test pins the widening: it reads the six source strings with include_str! (paths relative to CARGO_MANIFEST_DIR: ../mlx-cpp/turbo/paged_attention.cpp, paged_attention_v2.cpp, paged_attention_hip.h) and asserts each base = assignment is declared ulong or uint64_t; reverting any one body makes it fail, and the PR reports that revert-to-fail run.
• An #[ignore]d GPU test paged_pool_past_u32_elements_matches_gather in src/lib/mlxcel-core/src/cache/paged_batch_decode_tests.rs (needs about 17 GiB) builds an f16 pool with 131_073 blocks of [32, 8, 128], writes distinct K/V only into the last block, and checks that v1 and v2 decode over that block match the gather path; it fails on the 32-bit code and passes after the change. The PR records the run on at least one backend.
• examples/paged_attention_kernel_bench at batch 4 and 16K shows no decode regression beyond 3% on the backends available to the implementer, medians of 3.
• The production paged decode path (src/lib/mlxcel-core/src/cache/paged.rs, single_slab_tensors callers) uses the changed kernels with no caller change.

Verification

make verify-kernel-dtype-keys verify-kernel-port-dispatch
cargo test -p mlxcel-core --profile test-fast --features metal,accelerate -- paged --test-threads=1
cargo test -p mlxcel-core --profile test-fast --features metal,accelerate -- --ignored paged_pool_past_u32 --test-threads=1
make verify   # Metal host; on gfx1151 run make verify-rocm instead

Related: #2103, #2068, #899 (slab sizing), #904 (sparse decode).

Progress (PR #2186): the GPU test lives in cache/paged_pool_offset_tests.rs because paged_batch_decode_tests.rs is already past the file-size limit. The kernel bench keeps a 32-block slab and measures gather at batch 4 / 16K, so the no-regression check used the server paged decode on gfx1151 instead (14.0 vs 14.0 tok/s, medians of 3, guarded). Merge oversize is also declined in PagedDecodePlan::validate, because the batched path panics on a launch Err rather than falling back.