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

docs/environment-variables.md

@@ -492,7 +492,7 @@ declines to serve, see
 
 | Variable | Values | Default | Notes |
 |----------|--------|---------|-------|
-| `MLXCEL_PAGED_SLAB_BLOCKS` | positive integer blocks, used verbatim (neither floored nor capped); `0` pins the pool default | derived (`ceil(per_slot_ctx / block_size) * batch`, floored at the 32-block pool default and capped at the per-layer share of the KV block budget) | Paged pool slab size in blocks. The fused kernels read one contiguous buffer per side, so a layer whose rows outgrow its first slab silently returns to the gather path; the server derives the slab from `--ctx-size` and `--parallel` and logs the resolved value at startup (`Paged KV slab size: N blocks per layer`). Raise it when serving contexts longer than `--ctx-size` implies. `0` pins the historical 32-block pool default, which keeps pre-#899 allocation behaviour and, as a side effect, keeps the fused path unreachable for anything past one slab. A value that is not a non-negative integer warns and is ignored, so the derived size applies exactly as if the variable were unset (#1137). |
+| `MLXCEL_PAGED_SLAB_BLOCKS` | positive integer blocks, used verbatim (neither floored nor capped); `0` pins the pool default | derived (`ceil(per_slot_ctx / block_size) * batch`, floored at the 32-block pool default and capped at the per-layer share of the KV block budget) | Paged pool slab size in blocks. The fused kernels read one contiguous buffer per side, so a layer whose rows outgrow its first slab silently returns to the gather path; the server derives the slab from `--ctx-size` and `--parallel` and logs the resolved value at startup (`Paged KV slab size: N blocks per layer`). Raise it when serving contexts longer than `--ctx-size` implies. `0` pins the historical 32-block pool default, which keeps pre-#899 allocation behaviour and, as a side effect, keeps the fused path unreachable for anything past one slab. Slabs larger than 2^32 elements per side (for example more than 131,072 blocks of 32 tokens at 8 KV heads x 128 dims) are supported: the fused kernels address the pool with 64-bit offsets (#2153). A value that is not a non-negative integer warns and is ignored, so the derived size applies exactly as if the variable were unset (#1137). |
 | `MLXCEL_PAGED_V2_MIN_KV_TOKENS` | non-negative integer tokens | `4096` | Visible KV tokens a **lone-request** launch must reach before v2 is dispatched. 4096 is the weakest measured single-request win (1.08x); batch 1 at 1024 tokens is the only measured loss (0.91x), which is what the floor excludes. An unparseable value falls back to the default, so a typo cannot silently install an unmeasured policy. |
 | `MLXCEL_PAGED_V2_MIN_KV_TOKENS_PER_REQUEST` | non-negative integer tokens | `512` | Visible KV tokens **per request** a multi-request launch must reach. The smallest measured multi-request win is batch 4 at 1024 tokens per request (1.41x); the floor sits below that with margin so a prompt landing just under a nominal 1K is not declined for no measured reason. Same fallback-on-typo behaviour as the row above. To benchmark the declined regime, force the path on with `MLXCEL_PAGED_ATTENTION_NATIVE=1` rather than driving a floor to zero. |
 

src/lib/mlx-cpp/turbo/paged_attention.cpp

@@ -145,7 +145,8 @@ constexpr const char* PAGED_ATTENTION_DECODE_SOURCE = R"(
         uint block_idx = abs_pos / block_size;
         uint slot = abs_pos - block_idx * block_size;
         uint row = (uint)rows[row_off + block_idx];
-        uint base = (row * block_size + slot) * stride_kv + kv_head * dim;
+        // 64-bit: one slab can hold more than 2^32 elements (#2153).
+        ulong base = ((ulong)row * block_size + slot) * (ulong)stride_kv + kv_head * dim;
 
         float partial = 0.0f;
         for (uint j = 0; j < dpt; j++) {
@@ -298,7 +299,8 @@ constexpr const char* PAGED_ATTENTION_DECODE_CUDA_SOURCE = R"(
         uint32_t block_idx = abs_pos / block_size;
         uint32_t slot = abs_pos - block_idx * block_size;
         uint32_t row = (uint32_t)rows[row_off + block_idx];
-        uint32_t base = (row * block_size + slot) * stride_kv + kv_head * dim;
+        // 64-bit: one slab can hold more than 2^32 elements (#2153).
+        uint64_t base = ((uint64_t)row * block_size + slot) * (uint64_t)stride_kv + kv_head * dim;
 
         float partial = 0.0f;
         for (uint32_t j = 0; j < dpt; j++) {

src/lib/mlx-cpp/turbo/paged_attention_hip.h

@@ -136,7 +136,8 @@ inline constexpr const char* PAGED_ATTENTION_DECODE_HIP_SOURCE = R"(
         uint32_t block_idx = abs_pos / block_size;
         uint32_t slot = abs_pos - block_idx * block_size;
         uint32_t row = (uint32_t)rows[row_off + block_idx];
-        uint32_t base = (row * block_size + slot) * stride_kv + kv_head * dim;
+        // 64-bit: one slab can hold more than 2^32 elements (#2153).
+        uint64_t base = ((uint64_t)row * block_size + slot) * (uint64_t)stride_kv + kv_head * dim;
 
         float partial = 0.0f;
         for (uint32_t j = 0; j < dpt; j++) {
@@ -297,7 +298,8 @@ inline constexpr const char* PAGED_ATTENTION_V2_PARTIAL_HIP_SOURCE = R"(
         uint32_t page_off = abs_pos / page_size;
         uint32_t entry = abs_pos - page_off * page_size;
         uint32_t row = (uint32_t)indices[page_begin + page_off];
-        uint32_t base = (row * page_size + entry) * stride_kv + kv_head * dim;
+        // 64-bit: one slab can hold more than 2^32 elements (#2153).
+        uint64_t base = ((uint64_t)row * page_size + entry) * (uint64_t)stride_kv + kv_head * dim;
 
         float k_reg[DimsPerThread];
         float v_reg[DimsPerThread];

src/lib/mlx-cpp/turbo/paged_attention_v2.cpp

@@ -172,7 +172,8 @@ constexpr const char* PAGED_ATTENTION_V2_PARTIAL_SOURCE = R"(
         uint page_off = abs_pos / page_size;
         uint entry = abs_pos - page_off * page_size;
         uint row = (uint)indices[page_begin + page_off];
-        uint base = (row * page_size + entry) * stride_kv + kv_head * dim;
+        // 64-bit: one slab can hold more than 2^32 elements (#2153).
+        ulong base = ((ulong)row * page_size + entry) * (ulong)stride_kv + kv_head * dim;
 
         float k_reg[DimsPerThread];
         float v_reg[DimsPerThread];
@@ -367,7 +368,8 @@ constexpr const char* PAGED_ATTENTION_V2_PARTIAL_CUDA_SOURCE = R"(
         uint32_t page_off = abs_pos / page_size;
         uint32_t entry = abs_pos - page_off * page_size;
         uint32_t row = (uint32_t)indices[page_begin + page_off];
-        uint32_t base = (row * page_size + entry) * stride_kv + kv_head * dim;
+        // 64-bit: one slab can hold more than 2^32 elements (#2153).
+        uint64_t base = ((uint64_t)row * page_size + entry) * (uint64_t)stride_kv + kv_head * dim;
 
         float k_reg[DimsPerThread];
         float v_reg[DimsPerThread];

src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp

@@ -32,6 +32,7 @@
 #include <mlx/backend/cuda/cuda.h>
 #include <mlx/backend/metal/metal.h>
 
+#include <cstdint>
 #include <mutex>
 #include <optional>
 #include <string>
@@ -273,6 +274,19 @@ std::vector<mlx::core::array> paged_attention_merge_states(
         num_outputs = 0;
     }
 
+    // Every body indexes `v_in` and `out_v` in 32-bit unsigned arithmetic,
+    // which wraps past UINT32_MAX elements and would read or write a different
+    // row of the same buffer (issue #2153). The partials are `chunks * H * D`
+    // f32, far below that in practice, so this refuses instead of widening the
+    // bodies; callers treat a refused launch as "fall back to gather".
+    const uint64_t out_elems = static_cast<uint64_t>(num_outputs) *
+        static_cast<uint64_t>(heads) * static_cast<uint64_t>(dim);
+    if (static_cast<uint64_t>(v_in.size()) > UINT32_MAX ||
+        out_elems > UINT32_MAX) {
+        throw std::invalid_argument(
+            "[paged_attention_merge_states] v_in or the merged output exceeds "
+            "UINT32_MAX elements, which the merge kernel cannot index");
+    }
 
     // Refuses when this backend has no port, naming the entry point and the
     // fallback from the table (issues #1885, #1801).

src/lib/mlxcel-core/src/cache/paged_batch_decode.rs

@@ -448,3 +448,7 @@ fn slice_row(arr: &MlxArray, b: usize) -> UniquePtr<MlxArray> {
 #[cfg(test)]
 #[path = "paged_batch_decode_tests.rs"]
 mod paged_batch_decode_tests;
+
+#[cfg(test)]
+#[path = "paged_pool_offset_tests.rs"]
+mod paged_pool_offset_tests;