[inureyes]

Summary

The fused paged-attention kernels computed the per-token pool base in 32 bits, so a layer slab past 2^32 elements per side (131,072 blocks of [32, 8, 128] for Llama-3.1-8B) silently read another row of the same buffer. The base is now 64-bit in all six bodies, and the 32-bit merge kernel is guarded by size checks instead.

What changed

Kernel bodies: only the pool-base declaration changes, plus a one-line comment above it. The variable type and the casts change; template args, grid, JIT names and dtype keys are unchanged (make verify-kernel-dtype-keys: 9 in scope).

Backend	v1 decode	v2 partial
Metal (ulong)	paged_attention.cpp:149	paged_attention_v2.cpp:176
CUDA (uint64_t)	paged_attention.cpp:303	paged_attention_v2.cpp:372
HIP (uint64_t)	paged_attention_hip.h:140	paged_attention_hip.h:302

Each line becomes T base = ((T)row * page + slot) * (T)stride_kv + kv_head * dim;, with row widened before the multiply. The k_pool[base + d] / v_pool[base + d] reads use the 64-bit value unchanged. The same arithmetic in fused_norm.cpp and fused_rope_append.cpp already uses ulong (Metal) and uint64_t (CUDA). Sparse decode runs on v2 partial, so it gets the same fix.

• Merge: paged_attention_merge_states throws std::invalid_argument (an Err in Rust) when v_in or the merged output exceeds UINT32_MAX elements. PagedDecodePlan::validate rejects such plans before launch. The batched decode panics on a launch Err, so with this check those batches go to gather instead.
• Docs: the MLXCEL_PAGED_SLAB_BLOCKS row in docs/environment-variables.md.
• Tests are in a new cache/paged_pool_offset_tests.rs, not paged_batch_decode_tests.rs, which is already 599 lines.

Verification

On gfx1151, rebased on 9c0ae2e9 (after #2181). Review follow-up c3fef346 changes comments and docs only:

• paged_pool_past_u32_elements_matches_gather (ignored, about 17 GiB): f16 pool of 131,073 blocks, K/V only in the last block, v1 and v2 checked against a host reference. With main's bodies it fails with max error 0.456 for both kernels. With this PR it passes (guarded runs). It also confirms the last block reads back and row 0 (where a 32-bit base lands) is zero.
• Source pin test: with only the HIP v2 body reverted it fails, naming that line.
• Merge refusal test (lazy broadcasts) and plan-bound test pass.
• cargo test -p mlxcel-core --features rocm --lib -- paged: 264 passed. Clippy on lib and tests is clean, and the script gates plus verify-rocm-overlay pass.
• Decode: Llama-3.1-8B-Instruct-4bit, server --parallel 4 --ctx-size 131072, 4 clients at ~16K prompt, 128 tokens, fused v2 on both arms (the log shows fused v2 launch). A build with main's kernel bodies (this branch with the three kernel files reverted) and this PR were run as a pair inside one rocm_gpu_guard.sh window, three times, and all attempts were clean. Per-request decode tok/s was 14.0/13.9/14.0 before and 14.1/14.0/14.0 after. Both medians are 14.0, so there is no regression. examples/paged_attention_kernel_bench was not used: it keeps the 32-block default slab, so at batch 4 / 16K it measures gather.
• make verify-rocm on c3fef346 (with MLXCEL_ROCM_SMOKE_MODEL=models/mlx/Qwen3-0.6B-4bit): script gates, workspace clippy and smoke pass. Tests: 152 suites, 11984 passed, 5 failed, 383 ignored. The 5 failures are the Gemma 3 lookahead tests in server::batch::scheduler (from perf(scheduler): restore Gemma 3 decode lookahead via ring rewind #2182). They fail the same way with every file of this PR reverted to origin/main 9c0ae2e9, pass when run alone, and are filed as test(scheduler): Gemma 3 lookahead tests fail in the full mlxcel lib run on ROCm #2187.

Found while writing the test: on ROCm, strided MLX copies of 2^32 or more elements fail to launch. That is filed as #2184, and the test avoids it.

Not verified: Metal and CUDA are not available on this host. Their v1 and v2 partial pool-base lines in the table above changed and were only reviewed, not compiled or run.

Closes #2153

[inureyes]

Implementation Review Summary

Intent

Address fused paged-attention pool reads with 64-bit offsets in all six bodies, and refuse merge sizes the 32-bit merge bodies cannot index.

Findings Addressed

None needed. No CRITICAL or HIGH findings.

Remaining Items

• The merge launcher check leaves a double blank line before the port selection (paged_attention_v2_merge.cpp:290) (LOW), cosmetic
• The doc example "131,072 blocks at 8 KV heads x 128 dims" (and the matching module doc in paged_pool_offset_tests.rs) is exactly 2^32 elements, which the old 32-bit base still addressed; the first failing size is 131,073 blocks, as the GPU test uses (LOW), wording only
• The PagedDecodePlan::validate comment names only the merge kernel, but the check also protects the v2 partial kernel's 32-bit out_v index on merge-free plans (LOW), comment only

Verification

• All stated requirements implemented
• No placeholder/mock code remaining
• Integrated into project code flow
• Project conventions followed
• Existing modules reused where applicable
• No unintended structural changes
• Tests pass (not run in this review; the host GPU was busy with make verify-rocm)

Notes: all six base lines widen row before the multiply and match across Metal (ulong) and CUDA/HIP (uint64_t), with precedent in fused_norm.cpp and fused_rope_append.cpp. No other pool index in those bodies is 32-bit. The plan check sends batched, sparse and cascade paths to gather through their existing PlanRejected/CascadeFailed outcomes, and MLA split-KV drops to absorbed_decode via .ok(). The source-pin, merge-refusal and plan tests each fail when the fix is reverted.

[inureyes]

Security and performance review

No CRITICAL or HIGH findings; no fix commits.

Checked:

• All six pool-base expressions widen row before the first multiply, and stride_kv before the second. slot/entry and kv_head * dim promote safely. Every K/V load uses base + d, so no 32-bit pool index remains in the six bodies.
• The other 32-bit indices are bounded. v1 q/out stay within B * Hq * D. v2 out_v/out_lse stay within num_chunks * Hq * D, and PagedDecodePlan::validate now caps that at u32. validate runs in the batched path (cache/paged.rs) and in launch_with_lse, so both cascade sub-plans are covered too. The merge launcher refuses anything still too large before launch, and MLX is lazy, so unevaluated partials are never run.
• Overflow in the host checks: the merge check multiplies in uint64_t after num_outputs is clamped to 0 or more. In validate, num_chunks <= 65535 is checked before workspace_partial_v_elems(), so the usize product cannot wrap.

LOW / informational (no action requested):

• Metal has no native 64-bit integer multiply, so each token costs a few extra ALU ops per lane in the loop. The loop is bandwidth-bound and ROCm measured no change (14.0 vs 14.0 tok/s), but Metal has not been measured. If it ever shows up, the row base can be recomputed only when block_idx changes.
• Pre-existing, not introduced here: rows/indices values are cast from i32 without a bound check against the pool row count, and the v1 launcher allows v_pool to have fewer rows than k_pool. The values come from the internal allocator, so they are not attacker-controlled. Note that a corrupt row now reads at its true 64-bit offset where it used to wrap.