diff --git a/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md b/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md index 9f2437958..f4fe28a00 100644 --- a/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md +++ b/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md @@ -121,6 +121,36 @@ 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. @@ -128,4 +158,6 @@ The pool later moved from one growable tensor per layer to a list of fixed-size - `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). diff --git a/docs/environment-variables.md b/docs/environment-variables.md index d6e84da9c..c0489734c 100644 --- a/docs/environment-variables.md +++ b/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. | diff --git a/docs/turbo-kv-cache.md b/docs/turbo-kv-cache.md index 4932633ad..e2ba6cf43 100644 --- a/docs/turbo-kv-cache.md +++ b/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 diff --git a/examples/paged_attention_kernel_bench.rs b/examples/paged_attention_kernel_bench.rs index c93c9d43e..5cdd215a5 100644 --- a/examples/paged_attention_kernel_bench.rs +++ b/examples/paged_attention_kernel_bench.rs @@ -13,12 +13,14 @@ // 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 @@ -26,6 +28,12 @@ //! 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); + } } diff --git a/src/lib/mlxcel-core/src/cache/paged.rs b/src/lib/mlxcel-core/src/cache/paged.rs index c26ecf57b..701857fe4 100644 --- a/src/lib/mlxcel-core/src/cache/paged.rs +++ b/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. diff --git a/src/lib/mlxcel-core/src/ffi_tests.rs b/src/lib/mlxcel-core/src/ffi_tests.rs index 819bbba8e..e7cfd3398 100644 --- a/src/lib/mlxcel-core/src/ffi_tests.rs +++ b/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(); diff --git a/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index 48d3cea39..f5672e1ae 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/src/lib/mlxcel-core/src/layers.rs @@ -25,7 +25,7 @@ use crate::ffi; use crate::ffi::MlxArray; use cxx::UniquePtr; use std::sync::OnceLock; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; pub use crate::cache::{ChunkedKVCache, KVCache, KVCacheMode, RotatingKVCache}; @@ -3260,38 +3260,305 @@ pub fn paged_decode_attention_pooled_fallback( Ok(result) } -/// Process-wide env override for the fused paged-attention kernel (#123). +/// Where the pooled paged-attention decode should run for a given shape. /// -/// `MLXCEL_PAGED_ATTENTION_NATIVE=1` (or `true` / `on` / `yes`) force-enables -/// the native kernel regardless of the per-config `use_native_paged_kernel` -/// flag, so operators can A/B the kernel without rebuilding. Read once and -/// cached so the decode hot path never touches the environment. -fn native_paged_kernel_env() -> bool { +/// `Native` dispatches the fused Metal kernel +/// ([`crate::cache::PagedBlockPool::paged_decode_fused`], ADR 0001 strategy B); +/// `Gather` uses the gather-then-SDPA reference +/// ([`paged_decode_attention_pooled_fallback`], strategy A). The two paths agree +/// within RMS < 5e-3, so the choice is a pure performance switch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PagedDecodeDispatch { + /// Fused native kernel (strategy B). + Native, + /// Gather-then-SDPA reference (strategy A). + Gather, +} + +/// Compute backend the pooled decode runs on. The fused kernel is a Metal JIT +/// kernel and ADR 0001's regime table was measured on Apple Silicon Metal, so +/// only [`PagedDecodeBackend::Metal`] is ever a native candidate. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PagedDecodeBackend { + /// Apple Silicon Metal (the fused kernel's home). + Metal, + /// Anything else (CUDA/CPU): no fused-kernel evidence, always gather. + Other, +} + +// ── Adaptive native-kernel selector thresholds (issue #331) ────────────────── +// +// All three cite docs/adr/0001-paged-attention-gather-vs-fused-kernel.md, +// "Phase 6 outcome (#123)" (the fused-vs-gather speedup table) and the +// "Chunked slab storage interaction (#235)" addendum. + +/// Minimum batch size for a native dispatch. ADR 0001 Phase 6: the fused kernel +/// wins only once decode is batched (`b>=4` is `1.30x` at ctx 4096), and loses +/// at `b=1` (`0.50x`). +const NATIVE_MIN_BATCH: usize = 4; + +/// Maximum visible context for a native dispatch. ADR 0001 Phase 6: the fused +/// kernel wins at ctx 4096 but loses across every batch size at ctx 16384 +/// (`0.83x`–`0.91x`). 4096 is the largest context with measured evidence of a +/// win, so the selector stops there rather than extrapolating into the losing +/// 16384 regime. +const NATIVE_MAX_VISIBLE_LEN: usize = 4096; + +/// Maximum slab count for a native dispatch. ADR 0001 "#235": the fused kernel +/// reads one contiguous pool buffer per side, so +/// [`crate::cache::PagedBlockPool::paged_decode_fused`] declines any layer that +/// has grown past a single slab. Selecting native above this would always fall +/// back to gather; the selector short-circuits it instead. +const NATIVE_MAX_SLABS: usize = 1; + +/// Pure regime selector for the pooled paged-attention decode (issue #331). +/// +/// Returns [`PagedDecodeDispatch::Native`] only inside the island where ADR 0001 +/// Phase 6 measured the fused kernel winning: Apple Silicon Metal, batched +/// decode (`batch_size >= `[`NATIVE_MIN_BATCH`]`), moderate context +/// (`visible_len <= `[`NATIVE_MAX_VISIBLE_LEN`]`), and a single-slab layer the +/// kernel will not decline (`slab_count <= `[`NATIVE_MAX_SLABS`]`). Everything +/// else, including the `b=1` long-context regime the ADR named as a loss, routes +/// to [`PagedDecodeDispatch::Gather`]. +/// +/// Pure and allocation-free (a handful of integer comparisons) so it is cheap +/// to memoize on the decode hot path and trivially unit-testable without MLX. +#[must_use] +pub fn select_pooled_paged_dispatch( + batch_size: usize, + visible_len: usize, + slab_count: usize, + backend: PagedDecodeBackend, +) -> PagedDecodeDispatch { + let native = backend == PagedDecodeBackend::Metal + && slab_count <= NATIVE_MAX_SLABS + && batch_size >= NATIVE_MIN_BATCH + && visible_len <= NATIVE_MAX_VISIBLE_LEN; + if native { + PagedDecodeDispatch::Native + } else { + PagedDecodeDispatch::Gather + } +} + +/// Process-wide `MLXCEL_PAGED_ATTENTION_NATIVE` override for the fused +/// paged-attention kernel (#123, extended to tri-state in #331). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum NativePagedOverride { + /// Force the fused kernel, bypassing the adaptive selector (the original + /// #123 force-on semantics, preserved exactly). + ForceNative, + /// Force the gather fallback, bypassing the adaptive selector. The #331 + /// escape hatch for operators who want to pin the pre-kernel behaviour. + ForceGather, + /// No override: the adaptive selector decides. + Auto, +} + +/// Pure parse of a `MLXCEL_PAGED_ATTENTION_NATIVE` value into an override. +/// +/// The force-on value set (`1` / `true` / `on` / `yes`, and their uppercase +/// forms) is unchanged from #123, so an existing `MLXCEL_PAGED_ATTENTION_NATIVE=1` +/// still force-enables the kernel regardless of the selector. #331 adds a +/// symmetric force-off set (`0` / `false` / `off` / `no`) that pins the gather +/// fallback. Any other value (or `None`) yields [`NativePagedOverride::Auto`], +/// letting the selector govern dispatch. +fn parse_native_paged_override(value: Option<&str>) -> NativePagedOverride { + match value.map(str::trim) { + Some("1" | "true" | "on" | "yes" | "TRUE" | "ON" | "YES") => { + NativePagedOverride::ForceNative + } + Some("0" | "false" | "off" | "no" | "FALSE" | "OFF" | "NO") => { + NativePagedOverride::ForceGather + } + _ => NativePagedOverride::Auto, + } +} + +/// Read `MLXCEL_PAGED_ATTENTION_NATIVE` once and cache it so the decode hot path +/// never touches the environment. +fn native_paged_override() -> NativePagedOverride { use std::sync::OnceLock; - static ENABLED: OnceLock = OnceLock::new(); - *ENABLED.get_or_init(|| { - std::env::var("MLXCEL_PAGED_ATTENTION_NATIVE") - .map(|v| { - matches!( - v.trim(), - "1" | "true" | "on" | "yes" | "TRUE" | "ON" | "YES" - ) - }) - .unwrap_or(false) + static OVERRIDE: OnceLock = OnceLock::new(); + *OVERRIDE.get_or_init(|| { + parse_native_paged_override( + std::env::var("MLXCEL_PAGED_ATTENTION_NATIVE") + .ok() + .as_deref(), + ) + }) +} + +/// Pure combiner: apply the env override, otherwise the caller's willingness to +/// run native, otherwise the (lazily computed) adaptive selector. Force cases +/// never invoke `selector`, so the caller can skip deriving its inputs. +fn resolve_dispatch_decision( + over: NativePagedOverride, + use_native_requested: bool, + selector: impl FnOnce() -> PagedDecodeDispatch, +) -> PagedDecodeDispatch { + match over { + NativePagedOverride::ForceNative => PagedDecodeDispatch::Native, + NativePagedOverride::ForceGather => PagedDecodeDispatch::Gather, + NativePagedOverride::Auto => { + if use_native_requested { + selector() + } else { + PagedDecodeDispatch::Gather + } + } + } +} + +/// The backend the fused kernel would run on, cached for the decode hot path. +/// +/// The fused kernel is a Metal JIT kernel, so it is a candidate only when Metal +/// is available at runtime ([`crate::metal_is_available`], the same gate the +/// model dispatch paths use). This correctly falls to gather for a non-Metal +/// build (CUDA/CPU), a macOS build without the `metal` feature, and a machine +/// with no usable Metal device, none of which a compile-time `target_os` check +/// would catch. Detection is process-static, so it is read once. +fn paged_decode_backend() -> PagedDecodeBackend { + use std::sync::OnceLock; + static BACKEND: OnceLock = OnceLock::new(); + *BACKEND.get_or_init(|| { + if crate::metal_is_available() { + PagedDecodeBackend::Metal + } else { + PagedDecodeBackend::Other + } + }) +} + +/// Cheap per-shape memoization of the pooled-decode dispatch decision (#331, +/// acceptance criterion 3). +/// +/// Within a decode step every layer shares the same `(batch_size, visible_len, +/// slab_count, backend)` key, so the pure selector is recomputed at most once +/// per distinct shape and every subsequent layer takes the cached decision via +/// a single relaxed atomic load + compare. This is the "last-key cell" the issue +/// asks for, not a per-token locking map: the packed key (bits `0..=48`) and the +/// decision (bit `49`) live in one `AtomicU64`, so a reader observes them as one +/// indivisible word and can never pair a fresh key with a stale decision. The +/// selector is only re-run when the key changes. +struct PagedDispatchCache { + cell: AtomicU64, +} + +/// Sentinel value that never collides with a real packed cell: real cells pack +/// three `u16`-bounded fields plus the backend bit into bits `0..=48` (see +/// [`PagedDispatchCache::pack_key`]) and the decision into bit `49`, so bits +/// `50..=63` stay zero and the all-ones sentinel is unambiguous. +const PAGED_DISPATCH_CACHE_EMPTY: u64 = u64::MAX; + +impl PagedDispatchCache { + /// Bit carrying the decision alongside the packed key: set means + /// [`PagedDecodeDispatch::Native`], clear means [`PagedDecodeDispatch::Gather`]. + /// The key uses bits `0..=48` ([`Self::pack_key`]), so bit `49` is free. + const DECISION_BIT: u64 = 1u64 << 49; + /// Mask covering the packed-key bits (`0..=48`), used to compare a cell's + /// key half against a freshly packed key while ignoring the decision bit. + const KEY_MASK: u64 = Self::DECISION_BIT - 1; + + const fn new() -> Self { + Self { + cell: AtomicU64::new(PAGED_DISPATCH_CACHE_EMPTY), + } + } + + /// Pack the selector inputs into a single `u64` key. Each field is saturated + /// to `u16::MAX` first: the exact large value does not change the decision + /// (both `visible_len` and `slab_count` are far past every native threshold + /// once they exceed `u16::MAX`), and saturating keeps the packing lossless + /// for every value that could flip the outcome. + fn pack_key( + batch_size: usize, + visible_len: usize, + slab_count: usize, + backend: PagedDecodeBackend, + ) -> u64 { + let b = batch_size.min(u16::MAX as usize) as u64; + let v = visible_len.min(u16::MAX as usize) as u64; + let s = slab_count.min(u16::MAX as usize) as u64; + let k = match backend { + PagedDecodeBackend::Metal => 0u64, + PagedDecodeBackend::Other => 1u64, + }; + b | (v << 16) | (s << 32) | (k << 48) + } + + fn select( + &self, + batch_size: usize, + visible_len: usize, + slab_count: usize, + backend: PagedDecodeBackend, + ) -> PagedDecodeDispatch { + let key = Self::pack_key(batch_size, visible_len, slab_count, backend); + // Key and decision share one word, so a single relaxed load is + // torn-free: a hit returns the decision packed with this exact key, + // never a stale decision paired with a fresh key. + let cell = self.cell.load(Ordering::Relaxed); + if cell != PAGED_DISPATCH_CACHE_EMPTY && (cell & Self::KEY_MASK) == key { + return if cell & Self::DECISION_BIT != 0 { + PagedDecodeDispatch::Native + } else { + PagedDecodeDispatch::Gather + }; + } + let decision = select_pooled_paged_dispatch(batch_size, visible_len, slab_count, backend); + let packed = key + | match decision { + PagedDecodeDispatch::Native => Self::DECISION_BIT, + PagedDecodeDispatch::Gather => 0, + }; + self.cell.store(packed, Ordering::Relaxed); + decision + } +} + +static PAGED_DISPATCH_CACHE: PagedDispatchCache = PagedDispatchCache::new(); + +/// Resolve the dispatch for `paged_decode_attention_pooled`: honour the env +/// override first, otherwise run the memoized adaptive selector over the shape +/// derived from `pool`/`states`. +fn resolve_pooled_paged_dispatch( + pool: &crate::cache::PagedBlockPool, + states: &[&crate::cache::PagedSequenceState], + layer_idx: usize, + use_native_requested: bool, +) -> PagedDecodeDispatch { + resolve_dispatch_decision(native_paged_override(), use_native_requested, || { + let batch_size = states.len(); + let visible_len = states + .iter() + .map(|s| s.layer(layer_idx).map_or(0, |l| l.visible_len())) + .max() + .unwrap_or(0); + let slab_count = pool.slab_count(layer_idx); + PAGED_DISPATCH_CACHE.select(batch_size, visible_len, slab_count, paged_decode_backend()) }) } -/// Gated pooled paged decode attention (epic #116 Phase 6, #123). +/// Adaptive pooled paged decode attention (epic #116 Phase 6, #123; adaptive +/// selector #331). /// -/// When the native kernel is enabled (the per-config `use_native_paged_kernel` -/// flag or the `MLXCEL_PAGED_ATTENTION_NATIVE` env override), dispatches to the -/// fused Metal kernel ([`crate::cache::PagedBlockPool::paged_decode_fused`]), -/// which reads scattered pool blocks directly with no gather copy (ADR 0001 -/// strategy B). Otherwise, and whenever the kernel declines (the layer's pool -/// tensors are not yet allocated, or no sequence has visible tokens), it falls -/// back to [`paged_decode_attention_pooled_fallback`], the gather-then-SDPA -/// reference. The two paths agree within RMS < 5e-3, so the gate is a pure -/// performance switch with no behavioural change. +/// The per-config `use_native_paged_kernel` flag marks the caller as *willing* +/// to run the fused Metal kernel +/// ([`crate::cache::PagedBlockPool::paged_decode_fused`], ADR 0001 strategy B). +/// The actual dispatch is then decided by [`select_pooled_paged_dispatch`], +/// which only picks the kernel inside the regime ADR 0001 Phase 6 measured it +/// winning (Metal, `batch >= 4`, `visible_len <= 4096`, single-slab layer) and +/// otherwise uses [`paged_decode_attention_pooled_fallback`], the +/// gather-then-SDPA reference (strategy A). This keeps the long-context and +/// `b=1` regimes on gather, where the ADR shows the kernel loses. +/// +/// `MLXCEL_PAGED_ATTENTION_NATIVE` overrides the selector both ways: the +/// original force-on values (`1`/`true`/`on`/`yes`) still pin the kernel, and +/// #331 adds force-off values (`0`/`false`/`off`/`no`) that pin gather. Whenever +/// native is chosen but the kernel declines (pool tensors not yet allocated, no +/// visible tokens, or a multi-slab layer), it falls back to gather, so the two +/// paths stay interchangeable (RMS < 5e-3). pub fn paged_decode_attention_pooled( q: &MlxArray, pool: &crate::cache::PagedBlockPool, @@ -3300,7 +3567,8 @@ pub fn paged_decode_attention_pooled( scale: f32, use_native_paged_kernel: bool, ) -> Result, String> { - if (use_native_paged_kernel || native_paged_kernel_env()) + if resolve_pooled_paged_dispatch(pool, states, layer_idx, use_native_paged_kernel) + == PagedDecodeDispatch::Native && let Some(out) = pool.paged_decode_fused(q, states, layer_idx, scale)? { return Ok(out); @@ -4859,4 +5127,231 @@ mod tests { ); } } + + // ── Adaptive pooled paged-attention selector (issue #331) ──────────────── + // + // Thresholds cite docs/adr/0001-paged-attention-gather-vs-fused-kernel.md + // Phase 6 (#123): the fused kernel wins at Metal / batch >= 4 / ctx 4096 / + // single-slab, and loses at b=1, ctx 16384, or a multi-slab layer. + + use super::{ + NativePagedOverride, PagedDecodeBackend, PagedDecodeDispatch, PagedDispatchCache, + parse_native_paged_override, resolve_dispatch_decision, select_pooled_paged_dispatch, + }; + + #[test] + fn selector_b1_long_context_is_gather() { + // ADR 0001 Phase 6: b=1 loses (0.50x at 4096, 0.39x at 16384). The + // single-sequence regime must never dispatch native (acceptance + // criterion 2: no B=1 long-context regression). + for &ctx in &[1usize, 1024, 4096, 16384, 32768] { + assert_eq!( + select_pooled_paged_dispatch(1, ctx, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Gather, + "b=1 ctx={ctx} must stay on gather" + ); + } + } + + #[test] + fn selector_batched_moderate_context_single_slab_is_native() { + // ADR 0001 Phase 6: the fused kernel wins at ctx 4096 for b>=4 + // (1.30x/1.01x/1.36x at b=4/8/16), single-slab. + for &b in &[4usize, 8, 16] { + for &ctx in &[1usize, 512, 1024, 4096] { + assert_eq!( + select_pooled_paged_dispatch(b, ctx, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Native, + "b={b} ctx={ctx} single-slab Metal must dispatch native" + ); + } + } + } + + #[test] + fn selector_batched_long_context_is_gather() { + // ADR 0001 Phase 6: at ctx 16384 the kernel loses for every batch size + // (0.83x/0.82x/0.91x), so context past 4096 stays on gather. + for &b in &[4usize, 8, 16] { + for &ctx in &[4097usize, 8192, 16384, 32768] { + assert_eq!( + select_pooled_paged_dispatch(b, ctx, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Gather, + "b={b} ctx={ctx} must stay on gather past the 4096 cap" + ); + } + } + } + + #[test] + fn selector_multi_slab_is_gather() { + // ADR 0001 "#235": the fused kernel reads one contiguous buffer per side + // and declines multi-slab layers; the selector short-circuits them even + // in the otherwise-winning batched/4096 regime. + for &slabs in &[2usize, 3, 8] { + assert_eq!( + select_pooled_paged_dispatch(8, 4096, slabs, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Gather, + "slab_count={slabs} must stay on gather" + ); + } + } + + #[test] + fn selector_non_metal_backend_is_gather() { + // The fused kernel is Metal-only; off Apple Silicon there is no evidence + // and no kernel, so every shape stays on gather. + for &b in &[1usize, 4, 8, 16] { + assert_eq!( + select_pooled_paged_dispatch(b, 4096, 1, PagedDecodeBackend::Other), + PagedDecodeDispatch::Gather, + "b={b} on a non-Metal backend must stay on gather" + ); + } + } + + #[test] + fn selector_batch_and_context_boundaries() { + // batch boundary: 3 -> gather, 4 -> native (NATIVE_MIN_BATCH). + assert_eq!( + select_pooled_paged_dispatch(3, 4096, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Gather + ); + assert_eq!( + select_pooled_paged_dispatch(4, 4096, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Native + ); + // context boundary: 4096 -> native, 4097 -> gather (NATIVE_MAX_VISIBLE_LEN). + assert_eq!( + select_pooled_paged_dispatch(4, 4096, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Native + ); + assert_eq!( + select_pooled_paged_dispatch(4, 4097, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Gather + ); + // slab boundary: 1 -> native, 2 -> gather (NATIVE_MAX_SLABS). + assert_eq!( + select_pooled_paged_dispatch(4, 4096, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Native + ); + assert_eq!( + select_pooled_paged_dispatch(4, 4096, 2, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Gather + ); + } + + #[test] + fn native_override_parses_force_on_off_and_auto() { + // Force-on set is preserved exactly from #123. + for v in ["1", "true", "on", "yes", "TRUE", "ON", "YES", " 1 "] { + assert_eq!( + parse_native_paged_override(Some(v)), + NativePagedOverride::ForceNative, + "{v:?} must force native" + ); + } + // #331 force-off set. + for v in ["0", "false", "off", "no", "FALSE", "OFF", "NO", " 0 "] { + assert_eq!( + parse_native_paged_override(Some(v)), + NativePagedOverride::ForceGather, + "{v:?} must force gather" + ); + } + // Unset or unrecognised -> selector decides. + for v in [None, Some(""), Some("maybe"), Some("2")] { + assert_eq!( + parse_native_paged_override(v), + NativePagedOverride::Auto, + "{v:?} must defer to the selector" + ); + } + } + + #[test] + fn override_pins_dispatch_both_ways_bypassing_selector() { + use std::cell::Cell; + + // Force-native pins Native even when the selector would pick Gather, and + // the selector closure is never invoked. + let called = Cell::new(false); + let out = resolve_dispatch_decision(NativePagedOverride::ForceNative, true, || { + called.set(true); + PagedDecodeDispatch::Gather + }); + assert_eq!(out, PagedDecodeDispatch::Native); + assert!(!called.get(), "force-native must not consult the selector"); + + // Force-gather pins Gather even when the selector would pick Native. + let called = Cell::new(false); + let out = resolve_dispatch_decision(NativePagedOverride::ForceGather, true, || { + called.set(true); + PagedDecodeDispatch::Native + }); + assert_eq!(out, PagedDecodeDispatch::Gather); + assert!(!called.get(), "force-gather must not consult the selector"); + } + + #[test] + fn auto_defers_to_selector_and_respects_native_request() { + // Auto + caller willing -> whatever the selector returns. + assert_eq!( + resolve_dispatch_decision(NativePagedOverride::Auto, true, || { + PagedDecodeDispatch::Native + }), + PagedDecodeDispatch::Native + ); + assert_eq!( + resolve_dispatch_decision(NativePagedOverride::Auto, true, || { + PagedDecodeDispatch::Gather + }), + PagedDecodeDispatch::Gather + ); + // Auto + caller opted out of native -> gather, selector not consulted. + use std::cell::Cell; + let called = Cell::new(false); + let out = resolve_dispatch_decision(NativePagedOverride::Auto, false, || { + called.set(true); + PagedDecodeDispatch::Native + }); + assert_eq!(out, PagedDecodeDispatch::Gather); + assert!(!called.get(), "opt-out must not consult the selector"); + } + + #[test] + fn dispatch_cache_returns_selector_decision_across_keys() { + // The last-key memo must agree with the pure selector for both a + // repeated key (cache hit) and a changed key (recompute), and must not + // let stale keys leak the wrong decision. + let cache = PagedDispatchCache::new(); + + // Miss then hit on a Gather shape (b=1). + assert_eq!( + cache.select(1, 4096, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Gather + ); + assert_eq!( + cache.select(1, 4096, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Gather + ); + + // Key change flips to a Native shape (b=8). + assert_eq!( + cache.select(8, 4096, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Native + ); + // Back to the Gather key recomputes correctly (no stale Native). + assert_eq!( + cache.select(1, 4096, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Gather + ); + + // Saturated fields (values past u16::MAX) never spuriously match a small + // key and stay on gather (well past every native threshold). + assert_eq!( + cache.select(1, 1 << 20, 9, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Gather + ); + } }