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

[inureyes]

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

• 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)
• cargo test --release --features cuda -p mlxcel-core --lib paged_decode -- --test-threads=1 (11 passed)
• cargo test --release --features cuda -p mlxcel-core --lib layers::tests -- --test-threads=1 (53 passed)
• cargo run --release --features cuda --example paged_attention_kernel_bench (native fires on CUDA, 1.31x-6.20x)
• MLXCEL_PAGED_ATTENTION_NATIVE=0 vs =1 bench_decode A/B (no crash, no regression)

Closes #634