perf(cuda/attn): port native paged-attention decode kernel to CUDA

docs/adr/0001-paged-attention-gather-vs-fused-kernel.md

@@ -174,6 +174,10 @@ Retire `paged_decode_attention_pooled` and `select_pooled_paged_dispatch` (with
 - 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.
 
+## Addendum (2026-07-10): CUDA port of the fused kernel (#634)
+
+#634 ports the fused kernel (`src/lib/mlx-cpp/turbo/paged_attention.cpp`) from Metal-only `mx.fast.metal_kernel` to a `mx.fast.cuda_kernel` body with the same split-K flash-decoding scheme and grid geometry; the reduction changes from `simd_sum` to a `__shfl_xor_sync` butterfly all-reduce. The Metal-measured batch>=4/ctx<=4096 gate above does not carry over: on CUDA the fused kernel's win grows with context rather than losing it, so `select_pooled_paged_dispatch` dispatches native for any single-slab layer (any batch, any context) and still declines multi-slab layers per the #235 constraint. The entry point stays the library-only surface this ADR retired above; #634 does not reopen server wiring. See `docs/benchmark_results/paged-attention-cuda-port-gb10-2026-07-10.md` for the GB10 kernel A/B and parity numbers.
+
 ## References
 
 - Epic #116, unified KV cache.

docs/benchmark_results/paged-attention-cuda-port-gb10-2026-07-10.md

@@ -0,0 +1,42 @@
+# Paged-attention decode: CUDA port of the fused kernel (issue #634, GB10)
+
+Date: 2026-07-10. Host: NVIDIA GB10 (Grace-Blackwell, sm_121 / cc 12.1), CUDA 13.0, MLX 0.32.1.
+Binary: `cargo build --release --features cuda`. Bench: `examples/paged_attention_kernel_bench.rs` (kernel A/B), `mlxcel-bench-decode` (sanity A/B), `ffi_tests::test_fused_paged_decode_native_vs_fallback_matrix` (numerical parity).
+
+## Summary
+
+#634 ports the fused split-K paged-attention decode kernel (`src/lib/mlx-cpp/turbo/paged_attention.cpp`) from Metal-only `mx.fast.metal_kernel` to a `mx.fast.cuda_kernel` body, so CUDA no longer falls back to gather-then-SDPA on every step of the pooled paged decode path. The thread mapping and grid geometry carry over unchanged from Metal (one block per (batch, query head), `(32, NumSplits)` threads, online-softmax accumulation); the reduction changes from `simd_sum` to a `__shfl_xor_sync` butterfly all-reduce. `select_pooled_paged_dispatch` makes native the default on CUDA for any single-slab layer, since the Metal-measured batch>=4/ctx<=4096 ceiling (ADR 0001) does not apply on CUDA, where the fused win grows with context instead of losing it. Multi-slab layers still decline to gather, same as on Metal.
+
+## Numerical parity
+
+`test_fused_paged_decode_native_vs_fallback_matrix` runs the native kernel against the gather fallback over 60 configs (head dims 64/80/96/128, GQA ratios 1:1/4:1/8:1, page-boundary cases, batch > 1): worst RMS 5.45e-5, worst max-abs 2.35e-4, both well under the 5e-3 fp16 parity threshold.
+
+## Kernel A/B (`examples/paged_attention_kernel_bench`)
+
+This is the only path that reaches the fused kernel; `select=native` fires on CUDA once a layer is single-slab.
+
+| shape | select | fused vs gather |
+|---|---|---:|
+| batch 1, ctx 512 | native | 1.31x |
+| batch 1, ctx 1024 | native | 1.49x |
+| single-slab batched island | native | 2.55x / 6.20x / 1.27x |
+| multi-slab shapes | gather (declined) | - |
+
+Run-to-run variance: these kernels sit in the 200-900 us range where GB10 timing jitter is material. An independent re-run of the same sweep measured b1/512 at 0.88x (gather 266 us vs fused 303 us) and b1/1024 at 1.97x, so treat the single-batch ctx-512 case as noise-dominated around parity; the batched island and the b1/1024 win reproduce across runs.
+
+## 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. The two runs land within noise of each other because `mlxcel-bench-decode` uses dense single-stream caches and never reaches the fused kernel either way; this A/B confirms no accidental path divergence rather than a performance win.
+
+## Reachability caveat
+
+The fused kernel and its selector are library-only surface. #720 retired the pooled decode entry point (`paged_decode_attention_pooled` -> `PagedBlockPool::paged_decode_fused`) from the server: `mlxcel-server --decode-storage paged` routes pool-backed layers through the per-sequence `update_and_fetch` gather-then-SDPA intercept, and `mlxcel-bench-decode` uses dense single-stream caches, so neither reaches this kernel and `MLXCEL_PAGED_ATTENTION_NATIVE` does not change their output. The single-slab constraint (#235) also caps the servable context before the kernel declines to gather, which bounds the achievable long-context win independently of this port. Wiring the fused kernel onto the server decode path would reverse #720's deliberate library-only decision and stays out of scope for #634.
+
+## Reproduce
+
+```bash
+cargo run --release --features cuda --example paged_attention_kernel_bench
+cargo test --release --features cuda -p mlxcel-core --lib ffi_tests::test_fused_paged_decode_native_vs_fallback_matrix -- --test-threads=1
+```
+
+Refs #634, #720, #235, [ADR 0001](../adr/0001-paged-attention-gather-vs-fused-kernel.md).

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 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_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 (Metal: native only for batch>=4 + ctx<=4096 + single-slab; CUDA: native for any single-slab layer since #634; gather otherwise) | Overrides the fused split-K paged-attention decode kernel (Metal since epic #116 Phase 6/#123, and since #634 also CUDA via `mx.fast.cuda_kernel`) behind the library-only `paged_decode_attention_pooled` entry point. 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 on Metal, and the same island (any batch/context, single-slab only) on CUDA; 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. |

docs/turbo-kv-cache.md

@@ -293,13 +293,16 @@ batch 4 the native kernel runs at 276 tok/s for a 512-token prompt and 84 tok/s
 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
+the native kernel. The separate fused split-K kernel (Metal, and since #634
+also CUDA via `mx.fast.cuda_kernel`)
 (`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
+longer a plain on/off switch: an adaptive selector dispatches it on Metal only
+inside the batch>=4 / ctx<=4096 / single-slab island ADR 0001 measured it
+winning; on CUDA (#634) it dispatches on any single-slab layer regardless of
+batch or context, since the Metal ceilings do not apply there. Gather runs
+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. #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

examples/paged_attention_kernel_bench.rs

@@ -15,12 +15,14 @@
 //! Fused paged-attention decode kernel throughput bench (epic #116 Phase 6,
 //! #123; adaptive selector #331).
 //!
-//! 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.
+//! Compares the fused 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. The fused kernel runs its Metal JIT body
+//! on Apple and the #634 CUDA port on NVIDIA, so this bench exercises the native
+//! arm on either GPU backend.
 //!
 //! This is the #123 counterpart to `examples/page_gather_microbench.rs` (the
 //! ADR 0001 spike): the spike measured gather overhead against a contiguous
@@ -34,11 +36,13 @@
 //! reads `declined`. The selector encodes exactly that (plus the ADR 0001
 //! batch/context regime), so the two are cross-checked here.
 //!
-//! Run:
+//! Run (Apple):
 //!   cargo run --release --features metal,accelerate \
 //!     --example paged_attention_kernel_bench
-//! Run under `caffeinate -i` and let the machine cool between sweeps; Apple
-//! Silicon down-clocks under sustained load.
+//! Run (CUDA):
+//!   cargo run --release --features cuda --example paged_attention_kernel_bench
+//! On Apple Silicon run under `caffeinate -i` and let the machine cool between
+//! sweeps; it down-clocks under sustained load.
 
 use std::time::{Duration, Instant};
 
@@ -57,6 +61,19 @@ const LAYER: usize = 0;
 const WARMUP: usize = 20;
 const ITERS: usize = 50;
 
+/// Backend the fused kernel actually dispatches on: Metal on Apple, the #634
+/// CUDA port on NVIDIA, else CPU (no native kernel). Drives the selector label
+/// so the printed `select=` column matches production dispatch on this host.
+fn bench_backend() -> PagedDecodeBackend {
+    if mlxcel_core::metal_is_available() {
+        PagedDecodeBackend::Metal
+    } else if mlxcel_core::cuda_is_available() {
+        PagedDecodeBackend::Cuda
+    } else {
+        PagedDecodeBackend::Other
+    }
+}
+
 /// Deterministic pseudo-random f32 fill (a cheap LCG; only the timing matters,
 /// but distinct values keep softmax non-degenerate).
 fn pseudo_f32(seed: u64, n: usize) -> Vec<f32> {
@@ -157,16 +174,15 @@ 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.
+    // Selector decision for this shape on this host's backend (Metal on Apple,
+    // the #634 CUDA port on NVIDIA). 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 decision = select_pooled_paged_dispatch(batch, visible_len, slabs, bench_backend());
     let pick = match decision {
         PagedDecodeDispatch::Native => "native",
         PagedDecodeDispatch::Gather => "gather",