perf(core): adaptive selector for the native paged-attention decode kernel

docs/adr/0001-paged-attention-gather-vs-fused-kernel.md

@@ -121,11 +121,43 @@ Phase 6 lands the kernel gated off. The native path is opt-in through the `MLXCE
 
 The pool later moved from one growable tensor per layer to a list of fixed-size slab tensors (32 blocks per slab), so growth appends a slab instead of reallocating and copying the whole layer. The fused kernel reads one contiguous pool buffer per side, so `paged_decode_fused` declines (returns `None`, the caller falls back to gather) on any layer that has grown past a single slab. Layers within their first slab keep the kernel available. Teaching the kernel per-slab base pointers is possible if the trade-off above ever flips; given the kernel is gated off by default and loses in the long-context regime, the decline is the accepted state. The gather path splits the block-row list into per-slab runs and concatenates the per-run `take` results, which keeps it byte-identical and within run noise of the single-tensor layout.
 
+## Adaptive native-kernel selector (#331)
+
+The `use_native_paged_kernel` request from the scheduler previously meant "attempt the fused kernel, fall back only when it declines". That let `paged_decode_attention_pooled` dispatch native in regimes the Phase 6 table shows it losing: at `b=1` a single-slab layer ran the kernel at ~0.9x of gather. #331 replaces the request-means-attempt behaviour with an adaptive selector, `select_pooled_paged_dispatch(batch_size, visible_len, slab_count, backend)` in `src/lib/mlxcel-core/src/layers.rs`, which dispatches the kernel only inside the island the Phase 6 table measured it winning and uses gather everywhere else.
+
+The selector is native only when all four hold: the backend is Apple Silicon Metal (the kernel is Metal-only and the tables here are Metal); `batch_size >= 4` (Phase 6: `b=1` loses at 0.50x/0.39x, `b>=4` wins at 4096); `visible_len <= 4096` (Phase 6: 4096 wins, 16384 loses across every batch size, so the selector stops at the last measured winning context rather than extrapolating into 16384); and `slab_count <= 1` (the #235 decline above, hoisted into the selector so it never dispatches native for a layer the kernel would decline anyway). The decision is memoized in a last-key atomic cell: every layer in a decode step shares the same `(batch, visible_len, slab_count, backend)` key, so the pure selector runs at most once per distinct shape and later layers take a single relaxed atomic load. The pure function is a handful of integer comparisons, so the memo is a formality that pins "no per-token recompute" rather than a measured hot-path saving.
+
+`MLXCEL_PAGED_ATTENTION_NATIVE` now overrides the selector both ways: the original force-on set (`1`/`true`/`on`/`yes` and uppercase) still pins the kernel, and #331 adds a symmetric force-off set (`0`/`false`/`off`/`no`) that pins gather. An unset or unrecognised value defers to the selector.
+
+The chunked-slab reality (#235) narrows where native is reachable. Slab count is a per-layer property across all sequences sharing the pool, so a batched layer spans `B * ceil(visible_len / block_size)` rows; at `block_size` 32 any `B>=4` layer past ~256 tokens already exceeds one 32-row slab and the kernel declines. The batched moderate-context win the Phase 6 table recorded on the pre-#235 single-tensor pool is therefore unreachable today without the deferred multi-slab kernel. The reachable remnant is short-context batched decode (single-slab, `B>=4`), where the kernel still wins.
+
+Reachability caveat: the in-server decode path does not currently call `paged_decode_attention_pooled` at all. Pool-backed model layers gate out of the native-kernel arm (`is_paged_backed()`) and route through the per-sequence `update_and_fetch` pool intercept, which gathers the visible window and runs standard SDPA. The pooled entry point, and with it this selector, is exercised today by the kernel bench, the unit/FFI tests, and external mlxcel-core API consumers; inside `mlxcel-server` it is latent. Wiring the pooled path into the scheduler decode (or retiring it) is tracked in #710.
+
+**Hardware:** Apple M1 Ultra (Mac Studio), 128 GB, macOS 26.5. `--release --features metal,accelerate`, 50 timed iterations after 20 warmup, f16 pool, head_dim 128, 32 q-heads, 8 kv-heads, block 32. From `examples/paged_attention_kernel_bench.rs`. Fused = raw `paged_decode_fused`; speedup > 1 means the kernel beats gather:
+
+| batch | visible_len | slabs | selector | gather_us | fused_us | speedup |
+|------:|------------:|------:|:---------|----------:|:---------|--------:|
+| 1 | 512 | 1 | gather | 355 | 386 | 0.92x |
+| 1 | 1024 | 1 | gather | 352 | 391 | 0.90x |
+| 1 | 4096 | 4 | gather | 577 | declined | — |
+| 1 | 16384 | 16 | gather | 1201 | declined | — |
+| 4 | 4096 | 16 | gather | 1355 | declined | — |
+| 8 | 4096 | 32 | gather | 3442 | declined | — |
+| 4 | 128 | 1 | native | 490 | 393 | 1.25x |
+| 4 | 256 | 1 | native | 521 | 433 | 1.20x |
+| 8 | 128 | 1 | native | 747 | 427 | 1.75x |
+
+The table confirms the two design claims: single-slab `b=1` runs the kernel at a ~0.9x loss (which the selector now avoids by choosing gather, a small win over the old request-means-attempt path), and the only reachable batched island (single-slab `B>=4`) runs it at 1.20x–1.75x, where the selector chooses native. Everything at 1k/4k/16k with `B>=4` is multi-slab and declines to gather, so the selector and the kernel agree.
+
+The multi-slab native-kernel spike stays out of scope. It is worth building only if live traces show sustained requests in the `B>=4`, ~4k, multi-slab regime; that evidence does not exist yet, so this ADR keeps the decline (#235) as the accepted state and the selector confines native to the single-slab island.
+
 ## References
 
 - Epic #116, unified KV cache.
 - Issue #117, this Phase 0 spike.
 - `examples/page_gather_microbench.rs`, the microbench backing this ADR.
 - `src/lib/mlx-cpp/turbo/sparse_v_sdpa.metal`, the fused-kernel model for strategy (B).
 - `src/lib/mlxcel-core/src/layers.rs`, `paged_decode_attention_dense_compat`, the current dense decode path.
+- `src/lib/mlxcel-core/src/layers.rs`, `select_pooled_paged_dispatch` and `paged_decode_attention_pooled`, the adaptive selector (#331).
 - `src/lib/mlxcel-core/src/cache/paged.rs`, `PagedBlockPool` and `PagedKvLayout`.
+- `examples/paged_attention_kernel_bench.rs`, the fused-vs-gather bench and selector cross-check (#123, #331).

docs/environment-variables.md

@@ -211,6 +211,7 @@ recommended as normal deployment settings.
 | `MLXCEL_DISABLE_SOFTCAP_GQA_DECODE_GROUPED` | `1` disables, `0` enables | unset | Legacy rollback/override for grouped softcap-GQA decode. |
 | `MLXCEL_DISABLE_SINGLE_QUERY_MASKLESS` | truthy disables | maskless path on | Disables the single-query maskless attention path. |
 | `MLXCEL_EXPERIMENTAL_BOOL_CAUSAL_MASK` | truthy enables | off | Enables an experimental boolean causal-mask path. |
+| `MLXCEL_PAGED_ATTENTION_NATIVE` | `1`/`true`/`on`/`yes` force the native kernel; `0`/`false`/`off`/`no` force gather (case-insensitive); unset or any other value defers to the adaptive selector | selector-governed (native only for Metal + batch>=4 + ctx<=4096 + single-slab, gather otherwise) | Overrides the fused split-K Metal paged-attention decode kernel behind `paged_decode_attention_pooled` (epic #116 Phase 6, #123). Since #331 an unset value no longer means "always gather": `select_pooled_paged_dispatch` picks the kernel only inside the regime ADR 0001 measured it winning, and this variable still force-pins either arm for A/B testing. This pooled entry point is not reached by the `mlxcel serve` decode path today, which dispatches through the separate `DecodeBatchContext::use_native_paged_kernel` / block-table kernel instead; it is exercised by `examples/paged_attention_kernel_bench.rs` and the mlxcel-core unit/FFI test suites. See [ADR 0001](adr/0001-paged-attention-gather-vs-fused-kernel.md) and #710. |
 | `MLXCEL_SDPA_VECTOR_LARGE_D` | `0`/`false`/`off`/`no` disable; any other value or unset enables | on | **CUDA only.** Gates whether the CUDA `supports_sdpa_vector` check accepts head_dim 256/288 (gemma family, qwen3.5/3.6, baichuan-m1, paligemma2), routing their decode to the fused `sdpa_vector` kernels instead of the materializing SDPA fallback (issue #675). Disabling restores the prior fallback with no rebuild; used for the A/B in `benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv`. |
 | `MLXCEL_PIPELINE_GRANULARITY` | `off`, `layer`, `block:N` | `off` | Inserts layer-boundary async-eval hints for pipeline experiments. |
 | `MLXCEL_FUSED_MOE` | `0`/`false`/`off`/`no` disable; any other value or unset enables | on | Fused single-token decode-MoE kernel (#268), on by default since #282 (Metal) and #319 (CUDA, via `mx.fast.cuda_kernel`); validated on M1 Ultra, M5, and GB10. Set to `0` to force the proven `gather_qmm`/`SwitchGLU` path. Active for qwen3_moe, qwen3_next, dots.llm1, gemma4, qwen2_moe, mixtral, phimoe, lfm2, qwen3_vl_moe, and olmoe decode. |

docs/turbo-kv-cache.md

@@ -285,10 +285,16 @@ for a 4096-token prompt, versus 146 and 7.7 tok/s for the gather reference (1.9x
 and 10.9x). The gather reference degrades sharply with context because it
 re-materializes the visible window every step, which is why the live path uses
 the native kernel. The separate fused split-K Metal kernel
-(`MLXCEL_PAGED_ATTENTION_NATIVE`) is opt-in and stays off by default; with the
-chunked slab storage the pool is a list of fixed-size slab tensors, so the
-kernel also declines (falling back to gather) on any layer that has grown past
-one slab. See
+(`MLXCEL_PAGED_ATTENTION_NATIVE`, feeding `paged_decode_attention_pooled`) is a
+different code path from the block-table kernel above. Since #331 it is no
+longer a plain on/off switch: an adaptive selector dispatches it only inside
+the Metal / batch>=4 / ctx<=4096 / single-slab island ADR 0001 measured it
+winning, gather everywhere else, and the env var still force-pins either arm
+for A/B testing. The chunked slab storage narrows that island further, since
+the kernel declines (falling back to gather) once a layer has grown past one
+slab. Neither this kernel nor its selector is reached by the `mlxcel serve`
+decode path today, which stays on the block-table kernel described above; see
+ADR 0001's reachability caveat and #710. See
 [ADR 0001](adr/0001-paged-attention-gather-vs-fused-kernel.md).
 
 Pool growth appends fixed-size slabs instead of reallocating one big tensor

examples/paged_attention_kernel_bench.rs

@@ -13,19 +13,27 @@
 // limitations under the License.
 
 //! Fused paged-attention decode kernel throughput bench (epic #116 Phase 6,
-//! #123).
+//! #123; adaptive selector #331).
 //!
-//! Compares the fused Metal kernel (`paged_decode_attention_pooled` with the
-//! native path enabled, which reads scattered pool blocks directly) against the
-//! gather-then-SDPA reference (`paged_decode_attention_pooled_fallback`, which
-//! re-materialises a contiguous K/V every step) over a real `PagedBlockPool`.
+//! Compares the fused Metal kernel (raw `PagedBlockPool::paged_decode_fused`,
+//! which reads scattered pool blocks directly) against the gather-then-SDPA
+//! reference (`paged_decode_attention_pooled_fallback`, which re-materialises a
+//! contiguous K/V every step) over a real `PagedBlockPool`, and prints the
+//! decision the adaptive selector (`select_pooled_paged_dispatch`, #331) makes
+//! for each shape so both dispatch arms are visible.
 //!
 //! This is the #123 counterpart to `examples/page_gather_microbench.rs` (the
 //! ADR 0001 spike): the spike measured gather overhead against a contiguous
 //! SDPA lower bound; this measures the kernel that removes that gather. The
 //! sequences are grown with interleaved block writes so their physical pool
 //! rows are scattered (the gather pays its real scatter cost).
 //!
+//! The fused kernel reads one contiguous pool buffer per side, so with chunked
+//! slabs (#235) it can only run while a layer fits in a single 32-block slab
+//! (1024 tokens at block_size 32); past that it declines and the `fused` column
+//! reads `declined`. The selector encodes exactly that (plus the ADR 0001
+//! batch/context regime), so the two are cross-checked here.
+//!
 //! Run:
 //!   cargo run --release --features metal,accelerate \
 //!     --example paged_attention_kernel_bench
@@ -35,7 +43,10 @@
 use std::time::{Duration, Instant};
 
 use mlxcel_core::cache::{PagedBlockPool, PagedKvLayout, PagedSequenceState};
-use mlxcel_core::layers::{paged_decode_attention_pooled, paged_decode_attention_pooled_fallback};
+use mlxcel_core::layers::{
+    PagedDecodeBackend, PagedDecodeDispatch, paged_decode_attention_pooled_fallback,
+    select_pooled_paged_dispatch,
+};
 use mlxcel_core::{MlxArray, UniquePtr, astype, eval, from_slice_f32, synchronize_default};
 
 const HEAD_DIM: i32 = 128;
@@ -146,32 +157,83 @@ fn run_config(batch: usize, ctx: usize) {
         &[batch as i32, Q_HEADS, 1, HEAD_DIM],
     );
 
+    // Selector decision for this shape (Metal backend; this bench is
+    // Apple-Silicon only). Mirrors what production dispatch picks.
+    let visible_len = state_refs
+        .iter()
+        .map(|s| s.layer(LAYER).map_or(0, |l| l.visible_len()))
+        .max()
+        .unwrap_or(0);
+    let slabs = pool.slab_count(LAYER);
+    let decision =
+        select_pooled_paged_dispatch(batch, visible_len, slabs, PagedDecodeBackend::Metal);
+    let pick = match decision {
+        PagedDecodeDispatch::Native => "native",
+        PagedDecodeDispatch::Gather => "gather",
+    };
+
+    // Arm A: gather-then-SDPA reference (strategy A).
     let gather = time_body(|| {
         paged_decode_attention_pooled_fallback(&q, &pool, &state_refs, LAYER, scale).unwrap()
     });
-    let fused = time_body(|| {
-        paged_decode_attention_pooled(&q, &pool, &state_refs, LAYER, scale, true).unwrap()
-    });
-
     let g = per_call_us(gather);
-    let f = per_call_us(fused);
-    let speedup = g / f;
-    println!(
-        "  batch={batch:>2} ctx={ctx:>6}  gather={g:>9.1}us  fused={f:>9.1}us  speedup={speedup:>5.2}x"
-    );
+
+    // Arm B: raw fused kernel (strategy B), bypassing the selector. Declines
+    // (returns None) once the layer spans more than one slab, so probe first.
+    let fused_available = pool
+        .paged_decode_fused(&q, &state_refs, LAYER, scale)
+        .unwrap()
+        .is_some();
+
+    if fused_available {
+        let fused = time_body(|| {
+            pool.paged_decode_fused(&q, &state_refs, LAYER, scale)
+                .unwrap()
+                .unwrap()
+        });
+        let f = per_call_us(fused);
+        let speedup = g / f;
+        println!(
+            "  batch={batch:>2} ctx={ctx:>6} slabs={slabs:>3} select={pick}  gather={g:>9.1}us  fused={f:>9.1}us  speedup={speedup:>5.2}x"
+        );
+    } else {
+        println!(
+            "  batch={batch:>2} ctx={ctx:>6} slabs={slabs:>3} select={pick}  gather={g:>9.1}us  fused=declined(multi-slab)"
+        );
+    }
 }
 
 fn main() {
-    println!("=== mlxcel fused paged-attention decode kernel bench (#123) ===");
+    println!("=== mlxcel fused paged-attention decode kernel bench (#123, selector #331) ===");
     println!(
         "head_dim={HEAD_DIM} q_heads={Q_HEADS} kv_heads={KV_HEADS} block_size={BLOCK_SIZE} dtype=f16"
     );
     println!("warmup={WARMUP} iters={ITERS}");
     println!("Tip: run under `caffeinate -i` and let the machine cool between sweeps.");
+    println!(
+        "Note: the fused kernel runs only single-slab (<= {} tokens here); larger contexts decline and read `gather`.",
+        BLOCK_SIZE * 32
+    );
     println!();
-    for &ctx in &[4096usize, 16384] {
+    // Contexts 1k/4k/16k and batches 1/4/8 (issue #331 acceptance criterion 4);
+    // 512 and batch 16 are kept as extra single-slab / high-batch reference rows.
+    // The pool's slab count is per-layer across ALL sequences, so a batched
+    // layer spans B * ceil(ctx / block_size) rows: B>=4 at any of these contexts
+    // is already multi-slab and the kernel declines. select=gather throughout.
+    for &ctx in &[512usize, 1024, 4096, 16384] {
         for &batch in &[1usize, 4, 8, 16] {
             run_config(batch, ctx);
         }
     }
+
+    // Single-slab batched island: the only regime where a B>=4 layer still fits
+    // one 32-row slab (total rows = B * ceil(ctx / 32) <= 32), so the kernel is
+    // actually serviceable and the selector picks native. This is the reachable
+    // remnant of ADR 0001's "batched moderate-context" win after chunked slabs
+    // (#235) shrank the servable context; it exercises the native arm end to end.
+    println!();
+    println!("-- single-slab batched island (select=native, fused kernel live) --");
+    for &(batch, ctx) in &[(4usize, 128usize), (4, 256), (8, 128)] {
+        run_config(batch, ctx);
+    }
 }

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

@@ -1700,6 +1700,20 @@ impl PagedBlockPool {
         sum(&self.pool_k) + sum(&self.pool_v)
     }
 
+    /// Number of physical K slab tensors currently allocated for `layer_idx`
+    /// (0 before the layer's first write, then one per [`POOL_SLAB_BLOCKS`]-row
+    /// growth episode). The fused decode kernel reads one contiguous buffer per
+    /// side, so it can only run while this is `<= 1` (see
+    /// [`Self::paged_decode_fused`], which declines multi-slab layers). The
+    /// adaptive selector in `crate::layers` keys on this so it never dispatches
+    /// native for a layer the kernel would decline anyway.
+    ///
+    /// Used by: `crate::layers::paged_decode_attention_pooled` (adaptive
+    /// native-kernel selector, #331).
+    pub fn slab_count(&self, layer_idx: usize) -> usize {
+        self.pool_k.get(layer_idx).map(Vec::len).unwrap_or(0)
+    }
+
     /// Resolve the physical pool row for `block_id` on `layer_idx`, assigning
     /// one lazily on first write (reusing a freed row when available) and
     /// appending pool slabs if the new row exceeds capacity.

src/lib/mlxcel-core/src/ffi_tests.rs

@@ -800,9 +800,15 @@ fn test_fused_paged_decode_gqa_and_batched() {
     let q = concatenate(&q0, &q1, 0);
 
     let states: [&PagedSequenceState; 2] = [&s0, &s1];
-    let fused =
-        crate::layers::paged_decode_attention_pooled(&q, &pool, &states, layer_idx, scale, true)
-            .unwrap();
+    // Call the raw fused kernel directly, not through
+    // `paged_decode_attention_pooled`: the adaptive selector (#331) routes B=2
+    // to gather (native needs B>=4), which would turn this fused-vs-gather
+    // parity check into a gather-vs-gather tautology. The layer is single-slab
+    // here (8 rows <= POOL_SLAB_BLOCKS), so the kernel must serve it.
+    let fused = pool
+        .paged_decode_fused(&q, &states, layer_idx, scale)
+        .unwrap()
+        .expect("single-slab layer: fused kernel must serve it");
     let gather =
         crate::layers::paged_decode_attention_pooled_fallback(&q, &pool, &states, layer_idx, scale)
             .unwrap();