docs(core): retire pooled paged-attention decode to library-only API#720
Merged
Conversation
PR #709's review established that paged_decode_attention_pooled has no mlxcel-server caller: the scheduler routes pool-backed layers through the per-sequence update_and_fetch gather intercept (src/models/llama3.rs is_paged_backed()), and production models call paged_decode_attention_dense_compat / _fallback, never the pooled entry point. #710 records the wire-in-vs-retire decision as retire to a library-only API. Evidence (ADR 0001, #710 section, Apple M1 Ultra): the fused kernel's winning island is single-slab only. A slab is POOL_SLAB_BLOCKS = 32 block rows and a block holds DEFAULT_PAGED_BLOCK_SIZE = 32 tokens, so one slab is 1024 token rows per layer across every sequence resident in the shared pool. The selector requires slab_count <= 1 and batch >= 4, so at the shipped --parallel 4 default (#714) the reachable island is at most ~256 total tokens per sequence. The #714 serving bench drives 512-token prompts, which are already 2 slabs at B=4 from the first decode step, and the pool only appends slabs (#235), so a request that starts inside the island leaves it permanently within a few dozen generated tokens. Wiring the pooled path in would thread a new batched-decode arm through every transformer family for that transient sliver. The change is documentation-only. paged_decode_attention_pooled, select_pooled_paged_dispatch, the MLXCEL_PAGED_ATTENTION_NATIVE override, and the per-shape dispatch memo now carry library-only doc comments stating they are not on the mlxcel-server decode path. The fused kernel, the selector, and examples/paged_attention_kernel_bench.rs stay as tested, benchmarked library surface. Nothing is deleted: use_native_paged_kernel is left in place because it also gates the live paged_decode_attention_dense_compat block-table decode path, so it is shared plumbing, not dead code. docs/environment-variables.md and docs/turbo-kv-cache.md now state MLXCEL_PAGED_ATTENTION_NATIVE is a library-consumer control and an A/B pin for the bench, not a server knob. Validated on Apple M1 Ultra with cargo check and cargo clippy --lib --tests (metal,accelerate, -D warnings), the layers::tests selector/override suite (55 passed), and ffi_tests::test_fused_paged_decode (2 passed). Refs #710
5 tasks
inureyes
added a commit
that referenced
this pull request
Jul 10, 2026
…ment ref Address two LOW review findings on the #634 CUDA paged-attention port. Add host-side unit tests for the GPU-independent selector logic that previously only the GPU-gated runtime test exercised: `selector_cuda_backend_native_on_single_slab_only` (CUDA + single-slab selects Native across the b=1 and long-context regimes Metal routes to gather; multi-slab declines to Gather), `dispatch_cache_cuda_tag_does_not_alias_metal_or_other` (a same-shape roundtrip proving the widened 2-bit backend tag keys three distinct cells), and `dispatch_cache_cuda_pack_key_roundtrips_without_collision` (the packed keys differ per backend and stay clear of the decision bit and empty sentinel). Fix the doc reference that attributed the pooled-decode retirement to #710: #720 (commit 8b2329c) retired the entry point; #710 is the issue it resolved under ADR 0001. The `select_pooled_paged_dispatch` and `paged_decode_attention_pooled` doc comments now credit #720 with #710 as the resolved issue. Verified with `cargo check --features cuda --lib --tests` and the narrow selectors: `layers::tests` 56 passed (3 new), `ffi_tests::test_fused_paged_decode` 3 passed. Refs #634
inureyes
added a commit
that referenced
this pull request
Jul 10, 2026
…DR docs MLXCEL_PAGED_ATTENTION_NATIVE's default-column and description in docs/environment-variables.md, and the matching paragraph in docs/turbo-kv-cache.md, described the fused kernel and its selector as Metal-only, which #634 makes stale: the kernel now also runs on CUDA and the selector's CUDA arm dispatches native for any single-slab layer instead of the Metal batch>=4/ctx<=4096 gate. Both docs now describe the per-backend selector behavior. ADR 0001 gets a short addendum, following the addendum convention already used in ADR 0002, recording the CUDA port and pointing at the new benchmark note; the ADR's original decision text is unchanged. Add docs/benchmark_results/paged-attention-cuda-port-gb10-2026-07-10.md with the GB10 kernel A/B (1.31x-6.20x fused-vs-gather on single-slab shapes, multi-slab declines), the parity result (worst RMS 5.45e-5 over 60 configs), and the sanity A/B confirming mlxcel-bench-decode is unaffected, plus the #720/#235 reachability caveat. Validation: - cargo fmt --check on the PR's touched Rust files (already clean, no changes needed) - manual scan of added doc lines for em dashes and trailing whitespace (none found) Refs #634
5 tasks
inureyes
added a commit
that referenced
this pull request
Jul 10, 2026
) ## Summary Port the fused paged-attention decode kernel from Metal-only (`mlx::core::fast::metal_kernel`) to `mx.fast.cuda_kernel` so CUDA no longer falls back to gather-then-SDPA for the pooled paged decode path, wire the runtime gate to allow native on CUDA behind the existing `MLXCEL_PAGED_ATTENTION_NATIVE` control, keep the gather fallback selectable, and add a native-vs-fallback numerical parity test across the issue's case matrix. ## What changed `src/lib/mlx-cpp/turbo/paged_attention.cpp`: add `PAGED_ATTENTION_DECODE_CUDA_SOURCE`, a `cuda_kernel` port of the split-K flash-decoding scheme. One CUDA block per (batch, query head), `(32, NumSplits)` threads where the 32 warp lanes partition the head dimension and NumSplits warps sweep strided token stripes; online-softmax accumulation in registers; a butterfly `__shfl_xor_sync` all-reduce replacing Metal `simd_sum` (every lane needs the score, so the reduction is an all-reduce, not a reduce-to-lane-0); `__syncthreads()` combine over `__shared__` scratch. Grid `(32, NumSplits, B*Hq)` over threadgroup `(32, NumSplits, 1)` ceil-divides to exact blocks, so no padded threads and no grid-padding guards. The launcher selects the CUDA body via `!metal::is_available()`, mirroring the #631/#727 SSM and MoE ports. The port matches the Metal kernel's contract exactly (f32 Q/out, f16 K/V pool via `const __half*`, `*_shape` metadata, template args, grid). `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`, `mlx_cxx_bridge.h`, `src/lib/mlxcel-core/src/lib.rs`: add a `cuda_is_available()` FFI probe over `mlx::core::cu::is_available()` (backend-agnostic, `no_cuda` stub off CUDA). `src/lib/mlxcel-core/src/layers.rs`: add `PagedDecodeBackend::Cuda`; `paged_decode_backend()` returns it when CUDA is available; `select_pooled_paged_dispatch` makes native the default on CUDA for any single-slab layer (the Metal-measured batch/context ceilings do not apply on CUDA); widen the dispatch cache's backend tag to two bits (decision bit moves to bit 50). The cached `MLXCEL_PAGED_ATTENTION_NATIVE` override is unchanged and still forces native (=1) or gather (=0) both ways; Metal dispatch is untouched. `src/lib/mlxcel-core/src/ffi_tests.rs`: add `test_fused_paged_decode_native_vs_fallback_matrix` (head dims 64/80/96/128, GQA 1:1/4:1/8:1, page-boundary cases, batch > 1); the older `test_fused_paged_decode_gqa_and_batched` no longer skips on non-Metal. `examples/paged_attention_kernel_bench.rs`: pick the backend at runtime so the `select=` column reflects CUDA; document the CUDA run command. ## Dtype and scope The Metal source is fp16-only (f16 K/V pool, f32 Q/out), so the port is fp16-only as well and says so. Quantized KV pages are out of scope here and interact with #635. ## Reachability (server vs library) The fused kernel is reached only via the library-only `paged_decode_attention_pooled` -> `PagedBlockPool::paged_decode_fused` path, exercised by `examples/paged_attention_kernel_bench.rs` and external mlxcel-core consumers. #720 retired the pooled decode from the server: `mlxcel-server --decode-storage paged` routes pool-backed layers through the per-sequence `update_and_fetch` gather+SDPA intercept (`src/models/llama3.rs` returns `None` for `is_paged_backed()` caches), and `mlxcel-bench-decode` uses dense single-stream caches. Neither reaches the fused kernel, so `MLXCEL_PAGED_ATTENTION_NATIVE` does not change their output. The single-slab constraint (#235) also caps the servable context at 1024 tokens; beyond that the kernel declines and falls back to gather, which bounds the achievable long-context win independently of this port. This PR delivers the CUDA kernel and gate; wiring the fused kernel onto the server decode path would reverse #720's deliberate library-only decision and is out of scope. ## Verification (GB10, CUDA) Numerical parity: `test_fused_paged_decode_native_vs_fallback_matrix` passes with worst RMS 5.45e-5, worst max-abs 2.35e-4 over 60 configs (well under the 5e-3 fp16 threshold). All 11 `ffi_tests` paged-decode tests and 53 `layers::tests` pass under `--features cuda`. Sanity A/B (`mlxcel-bench-decode`, llama-3.1-8b-4bit, 2048-token prompt, 32 decode tokens): `MLXCEL_PAGED_ATTENTION_NATIVE=0` decode 50.93 tok/s, `=1` decode 50.94 tok/s. Identical because bench_decode does not route through the fused kernel (dense single-stream caches), confirming no regression and no accidental path divergence. Kernel A/B (`examples/paged_attention_kernel_bench`, the path that does reach the fused kernel): `select=native` now fires on CUDA for single-slab shapes; fused vs gather speedup 1.31x (batch 1, ctx 512), 1.49x (batch 1, ctx 1024), and 2.55x / 6.20x / 1.27x on the single-slab batched island. Multi-slab shapes correctly decline and read gather. ## Test plan - [x] `cargo test --release --features cuda -p mlxcel-core --lib ffi_tests::test_fused_paged_decode_native_vs_fallback_matrix -- --test-threads=1` (worst RMS 5.45e-5) - [x] `cargo test --release --features cuda -p mlxcel-core --lib paged_decode -- --test-threads=1` (11 passed) - [x] `cargo test --release --features cuda -p mlxcel-core --lib layers::tests -- --test-threads=1` (53 passed) - [x] `cargo run --release --features cuda --example paged_attention_kernel_bench` (native fires on CUDA, 1.31x-6.20x) - [x] `MLXCEL_PAGED_ATTENTION_NATIVE=0 vs =1` bench_decode A/B (no crash, no regression) Closes #634
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Resolves the PR #709 review finding that
paged_decode_attention_pooledhas nomlxcel-servercaller. The decision recorded here is to retire the pooled decode entry point (and itsselect_pooled_paged_dispatchselector) to a library-only API rather than wire it into the scheduler decode.Why retire (occupancy evidence, Apple M1 Ultra)
The fused kernel's winning island is single-slab only. A slab is
POOL_SLAB_BLOCKS = 32block rows (src/lib/mlxcel-core/src/cache/paged.rs) and a block holdsDEFAULT_PAGED_BLOCK_SIZE = 32tokens (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_countreturns the layer's whole slab-list length). The selector requiresslab_count <= 1andbatch_size >= 4, so the pool stays native only whileB * ceil(len / 32) <= 32. At the shipped--parallel 4default (#714,n_parallel = 4) that islen <= 256total tokens per sequence, counting prompt, chat template, and generated tokens.That island is negligible in production. The #714 serving-throughput bench drives 512-token prompts, which at
B = 4need4 x 16 = 64block rows = 2 slabs, so every layer is multi-slab from the first decode step (the #331 bench table in ADR 0001 shows everyB >= 4row atctx >= 512readingdeclined). The pool only appends slabs (#235: existing slabs are never freed), so a request that starts inside the island leaves it permanently within its first few dozen generated tokens. Wiring the pooled path in would thread a new batched-decode arm through every transformer family (high blast radius, jitter-class parity risk) for that transient sliver. Since the pooled path has no server caller, its measured production occupancy is definitionally zero today, so the closed-form geometry is stronger evidence than a thermally-noisy instrumented run.What changed
src/lib/mlxcel-core/src/layers.rs: library-only doc comments onpaged_decode_attention_pooled,select_pooled_paged_dispatch, theMLXCEL_PAGED_ATTENTION_NATIVEoverride (NativePagedOverride), and the per-shape dispatch memo (PagedDispatchCache), stating they are not on themlxcel-serverdecode path and are kept for external mlxcel-core consumers andexamples/paged_attention_kernel_bench.rs.docs/adr/0001-paged-attention-gather-vs-fused-kernel.md: new decision record with the occupancy derivation, the geometry math, the chosen exit, and what would reopen the question.docs/environment-variables.md,docs/turbo-kv-cache.md:MLXCEL_PAGED_ATTENTION_NATIVEis now described as a library-consumer control and A/B pin for the bench, not a server knob.What was kept vs removed
PagedBlockPool::paged_decode_fused), the selector, the memo, the env override, andexamples/paged_attention_kernel_bench.rs(tested, benchmarked library surface).use_native_paged_kernelscheduler request. It also gates the livepaged_decode_attention_dense_compatblock-table decode path (src/models/llama3.rsand the other families), so it is shared plumbing, not dead code. No plumbing was removed.Test plan
cargo fmt --all -- --checkcargo check --lib --tests --features metal,acceleratecargo clippy --lib --tests --features metal,accelerate -- -D warningscargo test --release --features metal,accelerate -p mlxcel-core layers::tests::(53 passed: selector island, boundaries, override pinning)cargo test --release --features metal,accelerate -p mlxcel-core ffi_tests::test_fused_paged_decode(2 passed: gather parity over 200 steps, GQA + batched)Hardware: Apple M1 Ultra.
Closes #710