Skip to content

docs(core): retire pooled paged-attention decode to library-only API#720

Merged
inureyes merged 1 commit into
mainfrom
perf/710-pooled-paged-attn-retire
Jul 9, 2026
Merged

docs(core): retire pooled paged-attention decode to library-only API#720
inureyes merged 1 commit into
mainfrom
perf/710-pooled-paged-attn-retire

Conversation

@inureyes

@inureyes inureyes commented Jul 9, 2026

Copy link
Copy Markdown
Member

Summary

Resolves the PR #709 review finding that paged_decode_attention_pooled has no mlxcel-server caller. The decision recorded here is to retire the pooled decode entry point (and its select_pooled_paged_dispatch selector) 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 = 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). The selector requires slab_count <= 1 and batch_size >= 4, so the pool stays native only while B * ceil(len / 32) <= 32. At the shipped --parallel 4 default (#714, n_parallel = 4) that is len <= 256 total 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 = 4 need 4 x 16 = 64 block rows = 2 slabs, so every layer is multi-slab from the first decode step (the #331 bench table in ADR 0001 shows every B >= 4 row at ctx >= 512 reading declined). 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 on paged_decode_attention_pooled, select_pooled_paged_dispatch, the MLXCEL_PAGED_ATTENTION_NATIVE override (NativePagedOverride), and the per-shape dispatch memo (PagedDispatchCache), stating they are not on the mlxcel-server decode path and are kept for external mlxcel-core consumers and examples/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_NATIVE is now described as a library-consumer control and A/B pin for the bench, not a server knob.

What was kept vs removed

  • Kept: the fused kernel (PagedBlockPool::paged_decode_fused), the selector, the memo, the env override, and examples/paged_attention_kernel_bench.rs (tested, benchmarked library surface).
  • Kept: use_native_paged_kernel scheduler request. 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. No plumbing was removed.
  • Removed: nothing. The pooled function, selector, kernel, and bench all remain.

Test plan

  • cargo fmt --all -- --check
  • cargo check --lib --tests --features metal,accelerate
  • cargo clippy --lib --tests --features metal,accelerate -- -D warnings
  • cargo 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

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
@inureyes inureyes added type:performance Performance improvements priority:low Low priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:inference Generation, sampling, decoding (incl. speculative, DRY) status:review Under review status:done Completed and removed status:review Under review labels Jul 9, 2026
@inureyes inureyes merged commit 8b2329c into main Jul 9, 2026
5 checks passed
@inureyes inureyes deleted the perf/710-pooled-paged-attn-retire branch July 9, 2026 16:15
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
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:inference Generation, sampling, decoding (incl. speculative, DRY) priority:low Low priority status:done Completed type:performance Performance improvements

Projects

None yet

Development

Successfully merging this pull request may close these issues.

perf(core): wire the pooled paged-attention decode entry point into the scheduler or retire it

1 participant