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 f4fe28a0..8569de99 100644 --- a/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md +++ b/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md @@ -151,6 +151,29 @@ The table confirms the two design claims: single-slab `b=1` runs the kernel at a 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. +## Pooled entry point retired to a library-only API (#710) + +#710 resolved the reachability caveat above. The pooled decode entry point (`paged_decode_attention_pooled` and its `select_pooled_paged_dispatch` selector) is retired to a library-only API rather than wired into the scheduler decode. The occupancy derivation below shows the selector's winning island is both unreachable under the shipped serving defaults and transient even when it is entered, so a wire-in would thread a new batched-decode arm through every transformer family (high blast radius, jitter-class parity risk) for a sliver of accelerated steps that real chat serving leaves within its first few generated tokens. + +### Occupancy derivation (Apple M1 Ultra) + +The kernel wins only on a single-slab layer (`slab_count <= 1`, the #235 decline hoisted into the selector). A slab is `POOL_SLAB_BLOCKS = 32` block rows (`src/lib/mlxcel-core/src/cache/paged.rs`) and a block holds `DEFAULT_PAGED_BLOCK_SIZE = 32` tokens (`src/server/batch/scheduler.rs`), so one slab is 32 x 32 = 1024 token rows per layer across every sequence resident in the shared pool: `PagedBlockPool::slab_count` returns the layer's whole slab-list length, not one sequence's rows. The pool stays single-slab only while the total allocated block rows satisfy `sum_i ceil(len_i / 32) <= 32`. For a batch of `B` equal-length sequences that is `B * ceil(len / 32) <= 32`. + +Combined with the selector's `batch_size >= 4` floor and the shipped `--parallel 4` default (#714, `n_parallel = 4`), the reachable island at `B = 4` is `ceil(len / 32) <= 8`, i.e. `len <= 256` tokens per sequence counting the prompt, the chat-template framing, and generated tokens together. At `B = 8` it tightens to `len <= 128`, and at `B = 16` to `len <= 64`. + +Two facts make that island negligible in production. First, it is unreachable under the shipped defaults: the #714 serving-throughput bench drives 512-token prompts, and at `B = 4` a 512-token sequence needs `ceil(512 / 32) = 16` blocks, so the pool holds `4 x 16 = 64` rows = 2 slabs and every layer is multi-slab from the first decode step. The #331 bench table above shows exactly this, with every `B >= 4` row at `ctx >= 512` (`ctx 4096`, 16-32 slabs) reading `declined` and only the `ctx 128 / 256` rows at `B = 4 / 8` staying single-slab and native. Real chat serving passes 256 total tokens per sequence almost immediately, since the system prompt and chat-template framing alone are tens of tokens before the first user turn. Second, the island is transient even when entered: the pool only appends slabs (#235: growth appends, existing slabs are never freed), so once total rows cross 32 the layer stays multi-slab for the rest of the request. A request that starts inside the island leaves it permanently within its first few dozen generated tokens and never returns, so the best case for a wire-in is a short burst of accelerated steps at the very start of short-prompt batched requests, after which every step is gather anyway. + +An instrumented occupancy run adds nothing here. The pooled path has no server caller, so its measured production occupancy is definitionally zero today, and the geometry above is a closed-form bound on the hypothetical wired-in occupancy rather than a noisy sample that Apple Silicon thermal drift would blur. + +### Decision + +Retire `paged_decode_attention_pooled` and `select_pooled_paged_dispatch` (with their `MLXCEL_PAGED_ATTENTION_NATIVE` override and per-shape memo) to a library-only API: kept and tested for external mlxcel-core consumers and `examples/paged_attention_kernel_bench.rs`, and documented as not on the `mlxcel-server` decode path. The fused kernel (`PagedBlockPool::paged_decode_fused`), the selector, and the bench all stay, since they are tested, benchmarked library surface and the selector remains the correct dispatch gate for any consumer that does call the pooled entry point. `MLXCEL_PAGED_ATTENTION_NATIVE` remains a library-consumer control and an A/B pin for the bench, not a server knob. Nothing is deleted: the `use_native_paged_kernel` scheduler request is left in place because it also gates the live `paged_decode_attention_dense_compat` block-table decode path (`src/models/llama3.rs` and the other families), so it is shared plumbing, not dead code. + +### What reopens this + +- A serving workload that sustains `B >= 4` with total contexts at or under ~256 tokens per sequence (short-prompt, high-concurrency batched decode), where the single-slab island is entered often enough that a wire-in's transient burst pays for the cross-family blast radius. +- A multi-slab native kernel (the deferred #235 per-slab-base-pointer spike), which would lift the `slab_count <= 1` constraint and make the batched moderate-context win reachable at real context lengths. Live traces showing sustained `B >= 4`, ~4k, multi-slab decode would justify building it, per the #235 note above. + ## References - Epic #116, unified KV cache. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 5ccdc898..5583bbdb 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -217,7 +217,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_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 the library-only `paged_decode_attention_pooled` entry point (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. #710 retired this entry point to a library-only API: it is not on the `mlxcel serve` decode path (which dispatches through the separate `DecodeBatchContext::use_native_paged_kernel` block-table kernel), so this variable is a control for external mlxcel-core consumers and `examples/paged_attention_kernel_bench.rs`, not a server knob. 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 2e4b06f9..0be625fa 100644 --- a/docs/turbo-kv-cache.md +++ b/docs/turbo-kv-cache.md @@ -295,9 +295,11 @@ 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 +slab. #710 retired this pooled entry point to a library-only API: neither this +kernel nor its selector is on the `mlxcel serve` decode path (which stays on the +block-table kernel described above), and `MLXCEL_PAGED_ATTENTION_NATIVE` is a +control for external mlxcel-core consumers and the kernel bench, not a server +knob. See ADR 0001's #710 decision record, [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/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index f5672e1a..229b8053 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/src/lib/mlxcel-core/src/layers.rs @@ -3313,6 +3313,11 @@ const NATIVE_MAX_SLABS: usize = 1; /// Pure regime selector for the pooled paged-attention decode (issue #331). /// +/// Library-only: this selector governs [`paged_decode_attention_pooled`], which +/// #710 retired to a library-only API. It is not on the `mlxcel-server` decode +/// path; it is kept for external mlxcel-core consumers and +/// `examples/paged_attention_kernel_bench.rs`. See ADR 0001 (#710). +/// /// 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 @@ -3343,6 +3348,11 @@ pub fn select_pooled_paged_dispatch( /// Process-wide `MLXCEL_PAGED_ATTENTION_NATIVE` override for the fused /// paged-attention kernel (#123, extended to tri-state in #331). +/// +/// Library-only control: this variable steers only the library-only +/// [`paged_decode_attention_pooled`] entry point (ADR 0001 #710), not the live +/// `mlxcel-server` block-table decode path (which uses the separate +/// `DecodeBatchContext::use_native_paged_kernel` request). #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum NativePagedOverride { /// Force the fused kernel, bypassing the adaptive selector (the original @@ -3433,6 +3443,9 @@ fn paged_decode_backend() -> PagedDecodeBackend { /// Cheap per-shape memoization of the pooled-decode dispatch decision (#331, /// acceptance criterion 3). /// +/// Serves only the library-only [`paged_decode_attention_pooled`] path (ADR 0001 +/// #710); it is not on the `mlxcel-server` decode path. +/// /// 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 @@ -3543,6 +3556,18 @@ fn resolve_pooled_paged_dispatch( /// Adaptive pooled paged decode attention (epic #116 Phase 6, #123; adaptive /// selector #331). /// +/// **Library-only (not on the `mlxcel-server` decode path).** #710 retired this +/// entry point to a library API. The in-server decode routes pool-backed layers +/// through the per-sequence `update_and_fetch` pool intercept (gather + standard +/// SDPA, `src/models/llama3.rs` `is_paged_backed()`), and production models call +/// [`paged_decode_attention_dense_compat`] / [`paged_decode_attention_dense_fallback`], +/// never this function. It is kept and tested for external mlxcel-core consumers +/// and `examples/paged_attention_kernel_bench.rs`. ADR 0001 (#710) records the +/// occupancy derivation: the selector's single-slab batched island (`B>=4`, at +/// most ~256 tokens per sequence at `block_size` 32) is unreachable under the +/// shipped serving defaults and transient even when entered, so a scheduler +/// wire-in is not worth its cross-family blast radius. +/// /// 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).