From dfaa54cdbd8604b376b0da5c1f3068faef7fe81e Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 9 Jul 2026 16:28:24 +0900 Subject: [PATCH 1/5] perf(core): adaptive selector for the native paged-attention decode kernel The scheduler's paged decode requested the fused native kernel unconditionally, so paged_decode_attention_pooled ran it whenever the layer was single-slab, including regimes ADR 0001 Phase 6 measured it losing (b=1 single-slab runs the kernel at ~0.9x of gather on M1 Ultra). Replace request-means-attempt with an adaptive selector that dispatches native only inside the island the kernel is known to win and uses the gather-then-SDPA fallback everywhere else. select_pooled_paged_dispatch(batch_size, visible_len, slab_count, backend) is a pure, allocation-free function keyed on the four inputs and returns Native only when the backend is Apple Silicon Metal, batch_size >= 4, visible_len <= 4096, and slab_count <= 1. Each threshold cites docs/adr/0001-paged-attention-gather-vs-fused-kernel.md: b=1 loses (0.50x/0.39x) and b>=4 wins at 4096; 16384 loses across every batch size so the cap stops at the last measured winning context; the single-slab bound hoists the #235 decline 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 so the pure selector recomputes at most once per distinct decode-step shape and later layers take one relaxed atomic load, satisfying the "cached per layer/shape, no measurable hot-path overhead" criterion. MLXCEL_PAGED_ATTENTION_NATIVE now overrides the selector both ways. The original force-on set (1/true/on/yes and uppercase) is preserved exactly and still pins the kernel; #331 adds a symmetric force-off set (0/false/off/no) that pins gather. Unset or unrecognised values defer to the selector. The env is parsed once and cached, and the pure parser and the override/selector combiner are factored out for unit testing. Add PagedBlockPool::slab_count so the selector can read the per-layer slab count. Rework examples/paged_attention_kernel_bench.rs to call the raw fused kernel directly (bypassing the selector), print the selector decision per shape, sweep contexts 1k/4k/16k at B=1/4/8, and add a single-slab batched island section that exercises the native arm end to end. On M1 Ultra the island confirms the design: single-slab b=1 runs the kernel at 0.90-0.92x (selector picks gather), the reachable single-slab B>=4 island runs it at 1.20-1.75x (selector picks native), and every B>=4 layer at 1k/4k/16k is multi-slab and declines to gather. test_fused_paged_decode_gqa_and_batched now calls paged_decode_fused directly because the selector routes its B=2 shape to gather, which would otherwise make the fused-vs-gather parity a tautology. The multi-slab native-kernel spike stays out of scope pending live-trace evidence; the ADR addendum records the decision and the M1 Ultra numbers. Refs #331 --- ...-paged-attention-gather-vs-fused-kernel.md | 30 + examples/paged_attention_kernel_bench.rs | 96 +++- src/lib/mlxcel-core/src/cache/paged.rs | 14 + src/lib/mlxcel-core/src/ffi_tests.rs | 12 +- src/lib/mlxcel-core/src/layers.rs | 538 +++++++++++++++++- 5 files changed, 642 insertions(+), 48 deletions(-) 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..c423b42c2 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,34 @@ 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. + +**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 +156,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/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..9398f9c99 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,292 @@ 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(), + ) }) } -/// Gated pooled paged decode attention (epic #116 Phase 6, #123). +/// 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. /// -/// 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 fused kernel is Metal-only; on Apple Silicon (`target_os = "macos"`) it +/// is a candidate, everywhere else the selector must pick gather. 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 cfg!(target_os = "macos") { + 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 and decision live in +/// two `AtomicU64`s and the selector is only re-run when the key changes. +struct PagedDispatchCache { + key: AtomicU64, + decision: AtomicU64, +} + +/// Sentinel key that never collides with a real packed key (real keys pack +/// three `u16`-bounded fields into the low 48 bits; see [`PagedDispatchCache::pack_key`]). +const PAGED_DISPATCH_CACHE_EMPTY: u64 = u64::MAX; + +impl PagedDispatchCache { + const fn new() -> Self { + Self { + key: AtomicU64::new(PAGED_DISPATCH_CACHE_EMPTY), + decision: AtomicU64::new(0), + } + } + + /// 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); + if self.key.load(Ordering::Relaxed) == key { + return match self.decision.load(Ordering::Relaxed) { + 0 => PagedDecodeDispatch::Gather, + _ => PagedDecodeDispatch::Native, + }; + } + let decision = select_pooled_paged_dispatch(batch_size, visible_len, slab_count, backend); + // Store decision before key: a racing reader that sees the new key then + // reads the freshly stored decision, and the selector is pure so a stale + // decision paired with a matching key would still be correct anyway. + self.decision.store( + match decision { + PagedDecodeDispatch::Gather => 0, + PagedDecodeDispatch::Native => 1, + }, + Ordering::Relaxed, + ); + self.key.store(key, 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()) + }) +} + +/// Adaptive pooled paged decode attention (epic #116 Phase 6, #123; adaptive +/// selector #331). +/// +/// 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 +3554,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 +5114,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 + ); + } } From 209b7cc020777321a4410a43f47560afd6170015 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 9 Jul 2026 16:50:08 +0900 Subject: [PATCH 2/5] fix(core): pack paged-dispatch memo into one atomic to kill torn read The per-shape dispatch memo stored its key and decision in two separate Relaxed AtomicU64s, so a concurrent decode worker could torn-read a fresh key paired with a stale decision and pick the wrong dispatch arm for its shape. Fold the decision into bit 49 of the single packed key word so key and decision are loaded and stored as one indivisible AtomicU64, removing the race while keeping the last-key-cell design and identical dispatch results. Refs #331. --- src/lib/mlxcel-core/src/layers.rs | 54 ++++++++++++++++++------------- 1 file changed, 32 insertions(+), 22 deletions(-) diff --git a/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index 9398f9c99..2dde1feff 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/src/lib/mlxcel-core/src/layers.rs @@ -3434,22 +3434,32 @@ fn paged_decode_backend() -> PagedDecodeBackend { /// 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 and decision live in -/// two `AtomicU64`s and the selector is only re-run when the key changes. +/// 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 { - key: AtomicU64, - decision: AtomicU64, + cell: AtomicU64, } -/// Sentinel key that never collides with a real packed key (real keys pack -/// three `u16`-bounded fields into the low 48 bits; see [`PagedDispatchCache::pack_key`]). +/// 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 { - key: AtomicU64::new(PAGED_DISPATCH_CACHE_EMPTY), - decision: AtomicU64::new(0), + cell: AtomicU64::new(PAGED_DISPATCH_CACHE_EMPTY), } } @@ -3482,24 +3492,24 @@ impl PagedDispatchCache { backend: PagedDecodeBackend, ) -> PagedDecodeDispatch { let key = Self::pack_key(batch_size, visible_len, slab_count, backend); - if self.key.load(Ordering::Relaxed) == key { - return match self.decision.load(Ordering::Relaxed) { - 0 => PagedDecodeDispatch::Gather, - _ => PagedDecodeDispatch::Native, + // 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); - // Store decision before key: a racing reader that sees the new key then - // reads the freshly stored decision, and the selector is pure so a stale - // decision paired with a matching key would still be correct anyway. - self.decision.store( - match decision { + let packed = key + | match decision { + PagedDecodeDispatch::Native => Self::DECISION_BIT, PagedDecodeDispatch::Gather => 0, - PagedDecodeDispatch::Native => 1, - }, - Ordering::Relaxed, - ); - self.key.store(key, Ordering::Relaxed); + }; + self.cell.store(packed, Ordering::Relaxed); decision } } From 32e4616bc2ca094e031588112cb85a8be84e04c6 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 9 Jul 2026 16:55:11 +0900 Subject: [PATCH 3/5] fix(core): gate paged-decode native backend on runtime Metal check paged_decode_backend used the compile-time cfg!(target_os = "macos") to mark the fused kernel a candidate, which is true even for a macOS build without the metal feature or a host with no usable Metal device, so the selector could pick the Metal JIT native path where it cannot run. Switch to the runtime crate::metal_is_available() gate already used by the model dispatch paths, so CUDA/CPU, non-metal macOS, and Metal-less machines all correctly fall to gather. Refs #331. --- src/lib/mlxcel-core/src/layers.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index 2dde1feff..f5672e1ae 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/src/lib/mlxcel-core/src/layers.rs @@ -3412,14 +3412,17 @@ fn resolve_dispatch_decision( /// The backend the fused kernel would run on, cached for the decode hot path. /// -/// The fused kernel is Metal-only; on Apple Silicon (`target_os = "macos"`) it -/// is a candidate, everywhere else the selector must pick gather. Detection is -/// process-static, so it is read once. +/// 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 cfg!(target_os = "macos") { + if crate::metal_is_available() { PagedDecodeBackend::Metal } else { PagedDecodeBackend::Other From 1378027304a40edad5bc4444947eb58ee717336f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 9 Jul 2026 16:59:39 +0900 Subject: [PATCH 4/5] docs(adr): note the pooled paged-attention selector is server-latent (#710) The PR #709 review established that paged_decode_attention_pooled has no production caller: pool-backed model layers route through the update_and_fetch pool intercept (gather_visible + standard SDPA), so the #331 selector is exercised by the bench, the tests, and external mlxcel-core consumers only. Record the caveat in the ADR addendum and point at #710, which tracks wiring the pooled path into the scheduler decode or retiring it. --- docs/adr/0001-paged-attention-gather-vs-fused-kernel.md | 2 ++ 1 file changed, 2 insertions(+) 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 c423b42c2..f4fe28a00 100644 --- a/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md +++ b/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md @@ -131,6 +131,8 @@ The selector is native only when all four hold: the backend is Apple Silicon Met 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 | From 481998229631faccf3b68c10517e82346b343a2f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 9 Jul 2026 17:14:00 +0900 Subject: [PATCH 5/5] docs: reconcile paged-attention native-kernel docs with the #331 selector MLXCEL_PAGED_ATTENTION_NATIVE stopped being a plain on/off switch once the adaptive selector landed: an unset value now defers to select_pooled_paged_dispatch instead of always staying on gather, and the variable still force-pins either arm for A/B testing. docs/turbo-kv-cache.md described the old boolean-gate behavior and did not distinguish the pooled kernel from the live-path block-table kernel that DecodeBatchContext::use_native_paged_kernel actually drives, so its "opt-in and stays off by default" line was stale; it now cites the selector regime and the ADR 0001 reachability caveat that the pooled entry point is not reached by mlxcel serve today (#710). docs/environment-variables.md never listed MLXCEL_PAGED_ATTENTION_NATIVE at all despite documenting every other kernel-toggle env var in the "Hardware and kernel diagnostic variables" table; added a row covering the tri-state values, the selector-governed default, and the same reachability caveat. Validation: - cargo fmt --all -- --check - cargo clippy -p mlxcel-core --lib --tests --features metal,accelerate -- -D warnings - cargo test --release -p mlxcel-core --features metal,accelerate --lib layers::tests (53 passed, including the 10 selector tests) - cargo test --release -p mlxcel-core --features metal,accelerate --lib ffi_tests::test_fused_paged_decode / ffi_tests::test_paged (4 passed) - cargo test --release -p mlxcel-core --features metal,accelerate --lib cache::paged (87 passed) - python3 scripts/ci/check_cross_repo_refs.py (only #116/#123/#331/#710, all legitimate) Refs #331 --- docs/environment-variables.md | 1 + docs/turbo-kv-cache.md | 14 ++++++++++---- 2 files changed, 11 insertions(+), 4 deletions(-) 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