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

Part of #623.

Context

The native fused paged-attention decode kernel exists only as a Metal kernel: src/lib/mlx-cpp/turbo/paged_attention.cpp uses mlx::core::fast::metal_kernel (line ~211) with no cuda_kernel counterpart. The runtime gate (src/lib/mlxcel-core/src/layers.rs:2815-2854, env MLXCEL_PAGED_ATTENTION_NATIVE) refuses the native path off-Metal, so CUDA always takes paged_decode_attention_pooled_fallback: gather the paged KV blocks into contiguous tensors, then run regular SDPA. That adds a gather pass plus extra memory traffic per step, precisely where long-context decode is bandwidth-bound.

This matters for the server's paged decode storage (--decode-storage paged, DecodeStorageBackend::Paged, block budget --kv-cache-budget), which is the memory-efficient multi-tenant configuration.

A proven porting pattern exists in this codebase: the fused decode-MoE kernel was ported from metal_kernel to fast::cuda_kernel in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp (selection via use_cuda = !metal::is_available(), simd_sum -> __shfl_down_sync, header-source split; see lines ~103-160 and ~1400-1650).

Scope

CUDA port of the paged-attention decode kernel, gated the same way as Metal, with the gather fallback retained.

Implementation plan

1. Read the Metal kernel in src/lib/mlx-cpp/turbo/paged_attention.cpp: block table indexing, head layout, GQA handling, softmax accumulation scheme, and the pooled-fallback contract (layers.rs fallback call) that defines expected outputs.
2. Port to fast::cuda_kernel following the perf(moe): implement the fused decode-MoE kernel on CUDA #319 MoE pattern: one thread block per (sequence, head-group), warp-parallel over KV block pages, online-softmax accumulation in registers/shared memory, __shfl_down_sync reductions replacing simd_sum. Respect the existing quant modes if the Metal kernel supports quantized KV pages; if it is fp16-only, keep the port fp16-only and note the interaction with perf(cuda/kv): validate and enable quantized KV cache modes on CUDA #635.
3. Wire the gate: layers.rs:2849-2854 currently allows native only on Metal; extend to CUDA when the kernel is compiled in, still honoring MLXCEL_PAGED_ATTENTION_NATIVE (default: enabled where available, env=0 forces fallback).
4. Correctness: exact-match (or fp16-tolerance) comparison of native vs fallback outputs across: head dims used by the model zoo (64/80/96/128), GQA ratios, page-boundary cases (sequence length exactly at block edge, single block, many blocks), and batch>1.
5. Performance: scripts/benchmark_paged_decode_rollout.sh plus a long-context measurement (8k and 32k context decode on llama-3.1-8b-4bit with --decode-storage paged) comparing native vs fallback on GB10.

Acceptance criteria

• Native paged-attention runs on CUDA behind the existing env gate; fallback preserved and selectable.
• Numerical agreement with the fallback within fp16 tolerance on the case matrix above (add a unit/integration test that runs both paths and compares).
• Measured decode tok/s at 8k+ context with paged storage improves vs fallback on GB10 (record the table; expected gain grows with context length as the gather traffic dominates).
• Server E2E: mlxcel-server --decode-storage paged on GB10 exercises the native kernel (verify via log/metric), long-context requests stream correctly.
• cargo test --features cuda passes; Metal path untouched.

Validation

cargo build --release --features cuda
scripts/benchmark_paged_decode_rollout.sh
MLXCEL_PAGED_ATTENTION_NATIVE=0 ./target/release/mlxcel-bench-decode --model ./models/llama-3.1-8b-4bit --prompt-tokens 8192   # fallback
MLXCEL_PAGED_ATTENTION_NATIVE=1 ./target/release/mlxcel-bench-decode --model ./models/llama-3.1-8b-4bit --prompt-tokens 8192   # native

References

• Metal kernel: src/lib/mlx-cpp/turbo/paged_attention.cpp:211 area.
• Gate and fallback: src/lib/mlxcel-core/src/layers.rs:2815-2854.
• Port pattern: src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp (fused MoE cuda_kernel port).
• Paged cache: src/lib/mlxcel-core/src/cache/paged.rs, update_and_fetch_paged (cache.rs:2653).