diff --git a/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md b/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md index 8569de99..ff10c0c5 100644 --- a/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md +++ b/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. diff --git a/docs/benchmark_results/paged-attention-cuda-port-gb10-2026-07-10.md b/docs/benchmark_results/paged-attention-cuda-port-gb10-2026-07-10.md new file mode 100644 index 00000000..37995906 --- /dev/null +++ b/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). diff --git a/docs/environment-variables.md b/docs/environment-variables.md index a049260a..0a1fcac3 100644 --- a/docs/environment-variables.md +++ b/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. | diff --git a/docs/turbo-kv-cache.md b/docs/turbo-kv-cache.md index 5a438c2b..9e55b5ff 100644 --- a/docs/turbo-kv-cache.md +++ b/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 diff --git a/examples/paged_attention_kernel_bench.rs b/examples/paged_attention_kernel_bench.rs index 5cdd215a..44220f0a 100644 --- a/examples/paged_attention_kernel_bench.rs +++ b/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 { @@ -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", diff --git a/src/lib/mlx-cpp/turbo/paged_attention.cpp b/src/lib/mlx-cpp/turbo/paged_attention.cpp index 99f859c8..a0f31737 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention.cpp @@ -17,6 +17,13 @@ #include #include +// Backend availability probes for the Metal-vs-CUDA kernel gate (#634). Both +// headers are backend-agnostic: an Apple build links the CUDA no_cuda stub, a +// CUDA build links the Metal no_metal stub, so `is_available()` is always +// resolvable and returns false for the absent backend. +#include +#include + #include #include #include @@ -194,8 +201,169 @@ constexpr const char* PAGED_ATTENTION_DECODE_SOURCE = R"( } )"; +// CUDA port of the fused paged-attention decode kernel (#634). The Metal source +// above is `mx.fast.metal_kernel`, which throws "[metal_kernel] No Metal +// back-end" on the CUDA backend, so before this port every CUDA paged decode +// fell back to the gather-then-SDPA path (`paged_decode_attention_pooled_fallback`): +// a separate gather pass plus extra KV-pool traffic per step, exactly where +// long-context decode is bandwidth-bound. +// +// Same split-K flash-decoding scheme, thread mapping, and online-softmax +// accumulation as the Metal kernel: one CUDA block per (batch * query head), a +// `(32, NumSplits)` thread layout where the 32 lanes of a warp partition the +// head dimension and the `NumSplits` warps sweep strided token stripes. The +// per-token QK dot is warp-reduced with a butterfly `__shfl_xor_sync` (an +// all-reduce so every lane sees the score, matching Metal's `simd_sum`), and a +// `__syncthreads()` combine over shared memory merges the `NumSplits` partial +// softmaxes. +// +// Grid mapping (MLX passes Metal-style total threads and ceil-divides by the +// threadgroup tuple, cf. backend/cuda/custom_kernel.cpp): grid +// `(32, NumSplits, B*Hq)` over threadgroup `(32, NumSplits, 1)` yields blocks +// `(1, 1, B*Hq)` with `blockDim = (32, NumSplits, 1)`. Every field is an exact +// multiple, so no padded threads are launched and no guards are needed for +// grid padding. `vlen` is block-uniform (a function of the batch index only), +// so the empty-window `return` is taken by the whole block together and never +// strands a warp at the later `__syncthreads()`. +// +// `k_pool`/`v_pool` are f16 (`const __half*` after MLX type substitution); the +// `(float)` casts use `__half`'s implicit float conversion. Buffers, template +// constants, and the `*_shape` metadata match the Metal launcher exactly. +constexpr const char* PAGED_ATTENTION_DECODE_CUDA_SOURCE = R"( + uint32_t lane = threadIdx.x; // 0 .. 31 (within warp) + uint32_t sg = threadIdx.y; // 0 .. NumSplits-1 + uint32_t bhq = blockIdx.z; // 0 .. B*Hq-1 + + uint32_t hq_count = (uint32_t)q_shape[1]; // Hq + uint32_t block_size = (uint32_t)k_pool_shape[1]; // tokens per block + uint32_t hkv_count = (uint32_t)k_pool_shape[2]; // Hkv + uint32_t dim = (uint32_t)Dim; + uint32_t dpt = (uint32_t)DimsPerThread; // dims this lane owns + uint32_t d0 = lane * dpt; // first dim of this lane + + uint32_t b = bhq / hq_count; // batch index + uint32_t h = bhq % hq_count; // query head + uint32_t kv_head = h / (uint32_t)NRep; // grouped-query KV head + if (kv_head >= hkv_count) { + kv_head = 0; // defensive + } + + int vlen_i = visible_lens[b]; + uint32_t vlen = vlen_i > 0 ? (uint32_t)vlen_i : 0u; + uint32_t logical_start = (uint32_t)logical_starts[b]; + uint32_t row_off = (uint32_t)row_offsets[b]; + + // Stage this lane's Q slice in registers. + float q_reg[DimsPerThread]; + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + q_reg[j] = (d < dim) ? q[bhq * dim + d] : 0.0f; + } + + // Shared scratch for the cross-warp flash combine. + __shared__ float tg_m[NumSplits]; + __shared__ float tg_l[NumSplits]; + __shared__ float tg_acc[NumSplits * Dim]; + + // Empty window: only warp 0 emits zeros (uniform across the whole block). + if (vlen == 0u) { + if (sg == 0u) { + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + if (d < dim) { + out[bhq * dim + d] = 0.0f; + } + } + } + return; + } + + float scale_v = scale[0]; + + // This warp's online softmax over its strided token stripe. + float m = -INFINITY; + float l = 0.0f; + float acc[DimsPerThread]; + for (uint32_t j = 0; j < dpt; j++) { + acc[j] = 0.0f; + } + + uint32_t stride_kv = hkv_count * dim; // elements per (block,slot) + for (uint32_t t = sg; t < vlen; t += (uint32_t)NumSplits) { + uint32_t abs_pos = logical_start + t; + uint32_t block_idx = abs_pos / block_size; + uint32_t slot = abs_pos - block_idx * block_size; + uint32_t row = (uint32_t)rows[row_off + block_idx]; + uint32_t base = (row * block_size + slot) * stride_kv + kv_head * dim; + + float partial = 0.0f; + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + float kd = (d < dim) ? (float)k_pool[base + d] : 0.0f; + partial += q_reg[j] * kd; + } + // Butterfly all-reduce over the 32 lanes: full q . k_t in every lane + // (Metal simd_sum is an all-reduce; every lane needs `score`). The loop + // trip count is warp-uniform, so all 32 lanes reach each shuffle. + #pragma unroll + for (int o = 16; o > 0; o >>= 1) { + partial += __shfl_xor_sync(0xffffffffu, partial, o); + } + float score = partial * scale_v; + + float m_new = fmaxf(m, score); + float corr = expf(m - m_new); + float p = expf(score - m_new); + l = l * corr + p; + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + float vd = (d < dim) ? (float)v_pool[base + d] : 0.0f; + acc[j] = acc[j] * corr + p * vd; + } + m = m_new; + } + + // Publish this warp's partial. Every lane stores its dim slice; lane 0 + // stores the scalar (max, denominator). + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + if (d < dim) { + tg_acc[sg * dim + d] = acc[j]; + } + } + if (lane == 0u) { + tg_m[sg] = m; + tg_l[sg] = l; + } + __syncthreads(); + + // Warp 0 merges the NumSplits partials (flash rescale) and writes out. + if (sg == 0u) { + float m_g = tg_m[0]; + for (uint32_t s = 1; s < (uint32_t)NumSplits; s++) { + m_g = fmaxf(m_g, tg_m[s]); + } + float l_g = 0.0f; + for (uint32_t s = 0; s < (uint32_t)NumSplits; s++) { + l_g += tg_l[s] * expf(tg_m[s] - m_g); + } + float inv_l = l_g > 0.0f ? (1.0f / l_g) : 0.0f; + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + if (d < dim) { + float a = 0.0f; + for (uint32_t s = 0; s < (uint32_t)NumSplits; s++) { + a += tg_acc[s * dim + d] * expf(tg_m[s] - m_g); + } + out[bhq * dim + d] = a * inv_l; + } + } + } +)"; + // Apple Silicon SIMD width. Each SIMD group is 32 lanes that partition the head -// dimension; `NumSplits` SIMD groups split the token range. +// dimension; `NumSplits` SIMD groups split the token range. On CUDA the same +// constant is the warp width. constexpr int PAGED_ATTENTION_SIMD_WIDTH = 32; // Thread-safe lazy-initialised holder for the JIT-compiled kernel. Mirrors the @@ -224,6 +392,31 @@ inline PagedAttentionKernelHolder& get_paged_attention_kernel() { return holder; } +// CUDA counterpart of `PagedAttentionKernelHolder`. Lazily JIT-compiles the +// `cuda_kernel` port (#634) under the same `call_once` contract, reached only on +// a CUDA backend where `metal::is_available()` is false. +struct PagedAttentionKernelHolderCuda { + std::optional kernel; + std::once_flag init_flag; + + mlx::core::fast::CustomKernelFunction& get() { + std::call_once(init_flag, [this] { + kernel = mlx::core::fast::cuda_kernel( + "mlxcel_paged_attention_decode", + {"q", "k_pool", "v_pool", "rows", "row_offsets", "logical_starts", + "visible_lens", "scale"}, + {"out"}, + std::string(PAGED_ATTENTION_DECODE_CUDA_SOURCE)); + }); + return *kernel; + } +}; + +inline PagedAttentionKernelHolderCuda& get_paged_attention_kernel_cuda() { + static PagedAttentionKernelHolderCuda holder; + return holder; +} + } // namespace mlx::core::array paged_attention_decode( @@ -251,7 +444,14 @@ mlx::core::array paged_attention_decode( n_rep = 1; } - auto& kernel = get_paged_attention_kernel().get(); + // Metal kernel on Apple, CUDA port elsewhere. `mx.fast.metal_kernel` throws + // "[metal_kernel] No Metal back-end" on the CUDA backend, so dispatch the + // `cuda_kernel` port there; `metal::is_available()` is false on a CUDA-only + // build (#634). Both kernels share the template args, grid, and buffer + // contract below, so only the JIT-compiled body differs. + const bool use_cuda = !mlx::core::metal::is_available(); + auto& kernel = use_cuda ? get_paged_attention_kernel_cuda().get() + : get_paged_attention_kernel().get(); // Each of the 32 lanes owns a ceil(Dim/32)-wide slice of the head. int dims_per_thread = (dim + PAGED_ATTENTION_SIMD_WIDTH - 1) / PAGED_ATTENTION_SIMD_WIDTH; diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 1b7e7b11..5bb88ced 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1638,6 +1638,11 @@ bool gated_delta_kernel_available(); // metal::is_available() gate used to pick the Metal vs CUDA kernel port. #626. bool metal_is_available(); +// True when the MLX CUDA backend is available at runtime; false on Metal-only +// and CPU-only builds. Backend-agnostic (`no_cuda` stub off CUDA). Used by the +// paged-attention decode gate to allow the fused native kernel on CUDA. #634. +bool cuda_is_available(); + // Start/stop Metal GPU capture. Requires the process to run under // `MTL_CAPTURE_ENABLED=1`; otherwise Metal drops the capture silently. void metal_start_capture(rust::Str path); diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp index 717c1fd3..a7bf9b04 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp @@ -862,6 +862,15 @@ bool metal_is_available() { #endif } +// True when the MLX CUDA backend is available at runtime. Backend-agnostic: the +// `mlx/backend/cuda/cuda.h` probe links a `no_cuda` stub (returning false) on +// builds without CUDA, so Rust callers can pick a backend-specific default +// without a compile-time cfg. Used by the paged-attention decode gate to allow +// the fused native kernel on CUDA (#634). +bool cuda_is_available() { + return mlx::core::cu::is_available(); +} + // Start a GPU trace capture and stop an active one. Mirrors the Python // `mx.metal.start_capture` / `mx.metal.stop_capture` API so a mlxcel // decode run can emit the same `.gputrace` files that `mlx-lm` produces diff --git a/src/lib/mlxcel-core/src/cache/paged.rs b/src/lib/mlxcel-core/src/cache/paged.rs index 701857fe..7c8b6c32 100644 --- a/src/lib/mlxcel-core/src/cache/paged.rs +++ b/src/lib/mlxcel-core/src/cache/paged.rs @@ -1607,6 +1607,21 @@ impl PagedBlockPool { (Some([k]), Some([v])) => (k, v), _ => return Ok(None), }; + + // The CUDA launch places batch * query_heads in gridDim.z, which CUDA + // caps at 65535 (Metal has no such cap). The bridge entry returns a + // bare UniquePtr, so a failed launch would abort rather than error; + // decline here and let the caller take the gather fallback (#634 + // security review). Unreachable for the single-slab island's small + // batches, kept as a guard for external library callers. + if crate::layers::paged_decode_backend() == crate::layers::PagedDecodeBackend::Cuda { + let q_shape = ffi::array_shape(q); + let grid_z = i64::from(q_shape.first().copied().unwrap_or(0).max(0)) + * i64::from(q_shape.get(1).copied().unwrap_or(0).max(0)); + if grid_z > 65535 { + return Ok(None); + } + } let block_size = self.layout.block_size as i32; // Flatten every sequence's physical pool rows (block-table order) into diff --git a/src/lib/mlxcel-core/src/ffi_tests.rs b/src/lib/mlxcel-core/src/ffi_tests.rs index e7cfd339..15ffab16 100644 --- a/src/lib/mlxcel-core/src/ffi_tests.rs +++ b/src/lib/mlxcel-core/src/ffi_tests.rs @@ -736,11 +736,10 @@ fn test_fused_paged_decode_matches_gather_over_200_steps() { fn test_fused_paged_decode_gqa_and_batched() { use crate::cache::{PagedBlockPool, PagedSequenceState}; - // The fused paged-decode kernel is Metal-only; on a CUDA build it throws - // "No Metal back-end" and aborts the process. CUDA uses the gather-then-SDPA - // fallback instead (native CUDA paged attention is issue #634), so skip here - // on non-Metal backends. - if !crate::metal_is_available() { + // The fused paged-decode kernel dispatches a Metal JIT body on Apple and a + // CUDA JIT body on NVIDIA (#634); on a CPU-only build neither backend can + // launch it, so skip there rather than aborting the process. + if !crate::metal_is_available() && !crate::cuda_is_available() { return; } @@ -824,6 +823,168 @@ fn test_fused_paged_decode_gqa_and_batched() { println!("test_fused_paged_decode_gqa_and_batched: RMS = {rms:e}"); } +/// #634 native-vs-fallback parity matrix for the fused paged-attention decode +/// kernel. This is the CUDA-port acceptance test: it runs the native fused +/// kernel ([`crate::cache::PagedBlockPool::paged_decode_fused`]) and the +/// gather-then-SDPA reference +/// ([`crate::layers::paged_decode_attention_pooled_fallback`]) on the same +/// synthetic pool and asserts they agree within fp16 tolerance across the case +/// matrix the issue calls out: head dims 64/80/96/128, GQA ratios 1:1 / 4:1 / +/// 8:1, page-boundary cases (a single partial block, a context exactly at a +/// block edge, many blocks), and batch > 1. +/// +/// The kernel is now dispatched on both Metal (original JIT body) and CUDA (the +/// #634 port), so this runs on either GPU backend and is skipped only on a +/// CPU-only build where neither backend can launch it. `paged_decode_fused` is +/// called directly rather than through `paged_decode_attention_pooled` so the +/// adaptive selector cannot silently route a case to gather and turn the +/// comparison into a gather-vs-gather tautology. +#[test] +fn test_fused_paged_decode_native_vs_fallback_matrix() { + use crate::cache::{PagedBlockPool, PagedSequenceState}; + + if !crate::metal_is_available() && !crate::cuda_is_available() { + // No GPU backend can launch the fused kernel; the CPU build has no + // native path to compare (it always gathers). Nothing to verify. + return; + } + + let backend = if crate::metal_is_available() { + "metal" + } else { + "cuda" + }; + let block_size = 16usize; + let layer_idx = 0usize; + + // One parity run for a single (head layout, batch shape) configuration. + // `seq_lens[b]` is sequence b's visible length; sequences are grown with + // interleaved single-token writes so their physical pool rows fragment, and + // every visible token gets distinct pseudo-random K/V. Returns the RMS and + // max-abs deviation of native vs gather over the flattened `[B, Hq, 1, D]` + // outputs. + let run_case = |n_kv_heads: i32, n_q_heads: i32, head_dim: i32, seq_lens: &[i32]| { + let scale = 1.0f32 / (head_dim as f32).sqrt(); + let mut pool = PagedBlockPool::new(pooled_layout(block_size, 1, n_kv_heads, head_dim)); + let batch = seq_lens.len(); + let mut states: Vec = (0..batch) + .map(|_| PagedSequenceState::new(pool.layout())) + .collect(); + + let max_len = *seq_lens.iter().max().unwrap(); + for t in 0..max_len { + for (bi, state) in states.iter_mut().enumerate() { + if t >= seq_lens[bi] { + continue; + } + pool.append_tokens(state, layer_idx, 1).unwrap(); + let ids = state.layer(layer_idx).unwrap().block_ids.clone(); + let seed = (bi as u64 + 1) + .wrapping_mul(0x1_0001) + .wrapping_add(t as u64 + 1); + let k = from_slice_f32( + &pooled_pseudo_f32(seed, (n_kv_heads * head_dim) as usize), + &[1, n_kv_heads, 1, head_dim], + ); + let v = from_slice_f32( + &pooled_pseudo_f32(seed.wrapping_mul(7), (n_kv_heads * head_dim) as usize), + &[1, n_kv_heads, 1, head_dim], + ); + pool.write_block( + ids[(t as usize) / block_size], + layer_idx, + (t as usize) % block_size, + &k, + &v, + ) + .unwrap(); + } + } + + // Q: [B, Hq, 1, D], one distinct pseudo-random vector per sequence. + let mut q = None::>; + for bi in 0..batch { + let qb = from_slice_f32( + &pooled_pseudo_f32(0xB0B0 + bi as u64, (n_q_heads * head_dim) as usize), + &[1, n_q_heads, 1, head_dim], + ); + q = Some(match q { + None => qb, + Some(prev) => concatenate(&prev, &qb, 0), + }); + } + let q = q.unwrap(); + + let state_refs: Vec<&PagedSequenceState> = states.iter().collect(); + let native = pool + .paged_decode_fused(&q, &state_refs, layer_idx, scale) + .unwrap() + .expect("single-slab layer: fused kernel must serve it"); + let gather = crate::layers::paged_decode_attention_pooled_fallback( + &q, + &pool, + &state_refs, + layer_idx, + scale, + ) + .unwrap(); + + assert_eq!( + array_shape(&native), + vec![batch as i32, n_q_heads, 1, head_dim], + "native output shape mismatch (kv={n_kv_heads} q={n_q_heads} d={head_dim})" + ); + + let f = flatten_f32_local(&native); + let g = flatten_f32_local(&gather); + let rms = pooled_rms(&f, &g); + let max_abs = f + .iter() + .zip(g.iter()) + .map(|(x, y)| (x - y).abs()) + .fold(0.0f32, f32::max); + (rms, max_abs) + }; + + // Head dims from the model zoo, GQA ratios 1:1 / 4:1 / 8:1, and page-boundary + // shapes: a single partial block (10), a context exactly at a block edge + // (16, 32), many blocks (112 = 7 * 16), and a mixed batch > 1. All row + // counts stay within one pool slab (POOL_SLAB_BLOCKS = 32). + let head_dims = [64i32, 80, 96, 128]; + let gqa: [(i32, i32); 3] = [(4, 4), (2, 8), (1, 8)]; // 1:1, 4:1, 8:1 + let scenarios: [&[i32]; 5] = [ + &[10], // single partial block, batch 1 + &[16], // exactly one full block (page edge) + &[112], // many blocks (7), batch 1 + &[16, 48, 80], // batch 3: block edges of varying length + &[10, 33, 64, 17], // batch 4: partial + past-edge mix + ]; + + let mut worst_rms = 0.0f32; + let mut worst_max = 0.0f32; + let tol = 5e-3f32; + for &head_dim in &head_dims { + for &(n_kv_heads, n_q_heads) in &gqa { + for scenario in &scenarios { + let (rms, max_abs) = run_case(n_kv_heads, n_q_heads, head_dim, scenario); + assert!( + rms < tol, + "[{backend}] native vs gather RMS {rms:e} exceeded {tol:e} \ + (kv={n_kv_heads} q={n_q_heads} d={head_dim} lens={scenario:?})" + ); + worst_rms = worst_rms.max(rms); + worst_max = worst_max.max(max_abs); + } + } + } + println!( + "test_fused_paged_decode_native_vs_fallback_matrix [{backend}]: \ + worst RMS = {worst_rms:e}, worst max-abs = {worst_max:e} over \ + {} configs", + head_dims.len() * gqa.len() * scenarios.len() + ); +} + /// Sliding-window parity: a sequence with `logical_start > 0` (post-trim) must /// gather/decode exactly the visible window and match the dense reference of /// just that window. diff --git a/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index 229b8053..33fb2934 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/src/lib/mlxcel-core/src/layers.rs @@ -3275,14 +3275,18 @@ pub enum PagedDecodeDispatch { Gather, } -/// Compute backend the pooled decode runs on. The fused kernel is a Metal JIT -/// kernel and ADR 0001's regime table was measured on Apple Silicon Metal, so -/// only [`PagedDecodeBackend::Metal`] is ever a native candidate. +/// Compute backend the pooled decode runs on. The fused kernel has a Metal JIT +/// body (ADR 0001, measured on Apple Silicon) and a CUDA JIT body (#634). Both +/// are native candidates; the selector applies backend-specific thresholds. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PagedDecodeBackend { - /// Apple Silicon Metal (the fused kernel's home). + /// Apple Silicon Metal (the fused kernel's original home). Metal, - /// Anything else (CUDA/CPU): no fused-kernel evidence, always gather. + /// NVIDIA CUDA (the ported fused kernel, #634). The kernel reads scattered + /// pool blocks with no gather pass, so its advantage grows with context; + /// the Metal-measured batch/context ceilings do not apply. + Cuda, + /// Anything else (CPU): no fused kernel, always gather. Other, } @@ -3314,9 +3318,9 @@ const NATIVE_MAX_SLABS: usize = 1; /// Pure regime selector for the pooled paged-attention decode (issue #331). /// /// Library-only: this selector governs [`paged_decode_attention_pooled`], which -/// #710 retired to a library-only API. It is not on the `mlxcel-server` decode -/// path; it is kept for external mlxcel-core consumers and -/// `examples/paged_attention_kernel_bench.rs`. See ADR 0001 (#710). +/// #720 retired to a library-only API (ADR 0001, resolving issue #710). It is +/// not on the `mlxcel-server` decode path; it is kept for external mlxcel-core +/// consumers and `examples/paged_attention_kernel_bench.rs`. See ADR 0001. /// /// Returns [`PagedDecodeDispatch::Native`] only inside the island where ADR 0001 /// Phase 6 measured the fused kernel winning: Apple Silicon Metal, batched @@ -3335,10 +3339,23 @@ pub fn select_pooled_paged_dispatch( slab_count: usize, backend: PagedDecodeBackend, ) -> PagedDecodeDispatch { - let native = backend == PagedDecodeBackend::Metal - && slab_count <= NATIVE_MAX_SLABS - && batch_size >= NATIVE_MIN_BATCH - && visible_len <= NATIVE_MAX_VISIBLE_LEN; + let native = match backend { + // Apple Silicon Metal: only the ADR 0001 Phase 6 win island (batched, + // moderate context, single-slab). + PagedDecodeBackend::Metal => { + slab_count <= NATIVE_MAX_SLABS + && batch_size >= NATIVE_MIN_BATCH + && visible_len <= NATIVE_MAX_VISIBLE_LEN + } + // CUDA (#634): the fused kernel reads the scattered pool blocks in-place + // with no gather copy, so it is the default wherever it can serve the + // layer. The kernel reads one contiguous pool buffer per side, so it + // still declines multi-slab layers ([`NATIVE_MAX_SLABS`]); the + // Metal-measured batch and context ceilings do not apply because the + // CUDA win grows with context rather than eroding at long context. + PagedDecodeBackend::Cuda => slab_count <= NATIVE_MAX_SLABS, + PagedDecodeBackend::Other => false, + }; if native { PagedDecodeDispatch::Native } else { @@ -3422,18 +3439,21 @@ fn resolve_dispatch_decision( /// The backend the fused kernel would run on, cached for the decode hot path. /// -/// The fused kernel is a Metal JIT kernel, so it is a candidate only when Metal -/// is available at runtime ([`crate::metal_is_available`], the same gate the -/// model dispatch paths use). This correctly falls to gather for a non-Metal -/// build (CUDA/CPU), a macOS build without the `metal` feature, and a machine -/// with no usable Metal device, none of which a compile-time `target_os` check -/// would catch. Detection is process-static, so it is read once. -fn paged_decode_backend() -> PagedDecodeBackend { +/// The fused kernel has a Metal JIT body and a CUDA JIT body (#634), so it is a +/// native candidate whenever either backend is available at runtime +/// ([`crate::metal_is_available`] / [`crate::cuda_is_available`], the same gates +/// the model dispatch and C++ kernel selection use). This correctly falls to +/// gather only for a CPU-only build or a machine with no usable GPU, neither of +/// which a compile-time `target_os` check would catch. Metal is probed first so +/// the Apple path is unchanged. Detection is process-static, so it is read once. +pub(crate) fn paged_decode_backend() -> PagedDecodeBackend { use std::sync::OnceLock; static BACKEND: OnceLock = OnceLock::new(); *BACKEND.get_or_init(|| { if crate::metal_is_available() { PagedDecodeBackend::Metal + } else if crate::cuda_is_available() { + PagedDecodeBackend::Cuda } else { PagedDecodeBackend::Other } @@ -3459,17 +3479,18 @@ struct PagedDispatchCache { } /// Sentinel value that never collides with a real packed cell: real cells pack -/// three `u16`-bounded fields plus the backend bit into bits `0..=48` (see -/// [`PagedDispatchCache::pack_key`]) and the decision into bit `49`, so bits -/// `50..=63` stay zero and the all-ones sentinel is unambiguous. +/// three `u16`-bounded fields plus a 2-bit backend tag into bits `0..=49` (see +/// [`PagedDispatchCache::pack_key`]) and the decision into bit `50`, so bits +/// `51..=63` stay zero and the all-ones sentinel is unambiguous. const PAGED_DISPATCH_CACHE_EMPTY: u64 = u64::MAX; impl PagedDispatchCache { /// Bit carrying the decision alongside the packed key: set means /// [`PagedDecodeDispatch::Native`], clear means [`PagedDecodeDispatch::Gather`]. - /// The key uses bits `0..=48` ([`Self::pack_key`]), so bit `49` is free. - const DECISION_BIT: u64 = 1u64 << 49; - /// Mask covering the packed-key bits (`0..=48`), used to compare a cell's + /// The key uses bits `0..=49` ([`Self::pack_key`]: the backend tag now needs + /// two bits for the Metal/Cuda/Other trichotomy), so bit `50` is free. + const DECISION_BIT: u64 = 1u64 << 50; + /// Mask covering the packed-key bits (`0..=49`), used to compare a cell's /// key half against a freshly packed key while ignoring the decision bit. const KEY_MASK: u64 = Self::DECISION_BIT - 1; @@ -3495,7 +3516,8 @@ impl PagedDispatchCache { let s = slab_count.min(u16::MAX as usize) as u64; let k = match backend { PagedDecodeBackend::Metal => 0u64, - PagedDecodeBackend::Other => 1u64, + PagedDecodeBackend::Cuda => 1u64, + PagedDecodeBackend::Other => 2u64, }; b | (v << 16) | (s << 32) | (k << 48) } @@ -3556,27 +3578,30 @@ fn resolve_pooled_paged_dispatch( /// Adaptive pooled paged decode attention (epic #116 Phase 6, #123; adaptive /// selector #331). /// -/// **Library-only (not on the `mlxcel-server` decode path).** #710 retired this -/// entry point to a library API. The in-server decode routes pool-backed layers -/// through the per-sequence `update_and_fetch` pool intercept (gather + standard -/// SDPA, `src/models/llama3.rs` `is_paged_backed()`), and production models call +/// **Library-only (not on the `mlxcel-server` decode path).** #720 retired this +/// entry point to a library API (ADR 0001, resolving issue #710). The in-server +/// decode routes pool-backed layers through the per-sequence `update_and_fetch` +/// pool intercept (gather + standard SDPA, `src/models/llama3.rs` +/// `is_paged_backed()`), and production models call /// [`paged_decode_attention_dense_compat`] / [`paged_decode_attention_dense_fallback`], /// never this function. It is kept and tested for external mlxcel-core consumers -/// and `examples/paged_attention_kernel_bench.rs`. ADR 0001 (#710) records the +/// and `examples/paged_attention_kernel_bench.rs`. ADR 0001 records the /// occupancy derivation: the selector's single-slab batched island (`B>=4`, at /// most ~256 tokens per sequence at `block_size` 32) is unreachable under the /// shipped serving defaults and transient even when entered, so a scheduler /// wire-in is not worth its cross-family blast radius. /// /// The per-config `use_native_paged_kernel` flag marks the caller as *willing* -/// to run the fused Metal kernel +/// to run the fused kernel /// ([`crate::cache::PagedBlockPool::paged_decode_fused`], ADR 0001 strategy B). -/// The actual dispatch is then decided by [`select_pooled_paged_dispatch`], -/// which only picks the kernel inside the regime ADR 0001 Phase 6 measured it -/// winning (Metal, `batch >= 4`, `visible_len <= 4096`, single-slab layer) and -/// otherwise uses [`paged_decode_attention_pooled_fallback`], the -/// gather-then-SDPA reference (strategy A). This keeps the long-context and -/// `b=1` regimes on gather, where the ADR shows the kernel loses. +/// The actual dispatch is then decided by [`select_pooled_paged_dispatch`]. On +/// Metal it picks the kernel only inside the regime ADR 0001 Phase 6 measured it +/// winning (`batch >= 4`, `visible_len <= 4096`, single-slab layer), keeping the +/// long-context and `b=1` regimes on gather where the ADR shows Metal loses. On +/// CUDA (#634) the ported kernel reads scattered pool blocks with no gather pass, +/// so it is the default for any single-slab layer regardless of batch/context. +/// Everything else uses [`paged_decode_attention_pooled_fallback`], the +/// gather-then-SDPA reference (strategy A). /// /// `MLXCEL_PAGED_ATTENTION_NATIVE` overrides the selector both ways: the /// original force-on values (`1`/`true`/`on`/`yes`) still pin the kernel, and @@ -5160,8 +5185,9 @@ mod tests { // single-slab, and loses at b=1, ctx 16384, or a multi-slab layer. use super::{ - NativePagedOverride, PagedDecodeBackend, PagedDecodeDispatch, PagedDispatchCache, - parse_native_paged_override, resolve_dispatch_decision, select_pooled_paged_dispatch, + NativePagedOverride, PAGED_DISPATCH_CACHE_EMPTY, PagedDecodeBackend, PagedDecodeDispatch, + PagedDispatchCache, parse_native_paged_override, resolve_dispatch_decision, + select_pooled_paged_dispatch, }; #[test] @@ -5235,6 +5261,30 @@ mod tests { } } + #[test] + fn selector_cuda_backend_native_on_single_slab_only() { + // CUDA (#634): the ported kernel is the default wherever it can serve + // the layer, independent of the Metal-measured batch/context ceilings. + // Single-slab -> native across batch and context, including the b=1 and + // long-context regimes Metal routes to gather. + for &(b, ctx) in &[(1usize, 128usize), (1, 65536), (4, 4096), (16, 131072)] { + assert_eq!( + select_pooled_paged_dispatch(b, ctx, 1, PagedDecodeBackend::Cuda), + PagedDecodeDispatch::Native, + "CUDA single-slab b={b} ctx={ctx} must select native" + ); + } + // Multi-slab -> gather: the kernel reads one contiguous pool buffer per + // side and declines past NATIVE_MAX_SLABS regardless of batch/context. + for &slabs in &[2usize, 3, 8] { + assert_eq!( + select_pooled_paged_dispatch(1, 128, slabs, PagedDecodeBackend::Cuda), + PagedDecodeDispatch::Gather, + "CUDA multi-slab (slabs={slabs}) must decline to gather" + ); + } + } + #[test] fn selector_batch_and_context_boundaries() { // batch boundary: 3 -> gather, 4 -> native (NATIVE_MIN_BATCH). @@ -5379,4 +5429,69 @@ mod tests { PagedDecodeDispatch::Gather ); } + + #[test] + fn dispatch_cache_cuda_tag_does_not_alias_metal_or_other() { + // The backend tag widened to two bits (#634). A single-slab shape that + // differs only by backend must key three distinct cells: Metal at b=1 + // gathers (needs b>=4), CUDA natives (single-slab default), Other + // gathers (no kernel). If the CUDA tag aliased Metal's or Other's cell, + // the second/third lookup would return the first's cached decision. + let cache = PagedDispatchCache::new(); + + // Prime CUDA (native), then query the same shape on Metal (gather) and + // Other (gather). No aliasing means each returns its own decision. + assert_eq!( + cache.select(1, 512, 1, PagedDecodeBackend::Cuda), + PagedDecodeDispatch::Native, + "CUDA single-slab b=1 must be native" + ); + assert_eq!( + cache.select(1, 512, 1, PagedDecodeBackend::Metal), + PagedDecodeDispatch::Gather, + "Metal b=1 must be gather (not the aliased CUDA native)" + ); + assert_eq!( + cache.select(1, 512, 1, PagedDecodeBackend::Other), + PagedDecodeDispatch::Gather, + "Other must be gather (no kernel)" + ); + // Re-query CUDA: still native, proving the Metal/Other writes did not + // clobber the CUDA cell (each backend keeps its own last-key cell only + // for the most recent distinct key, so this also confirms recompute + // agrees with the pure selector rather than returning a stale hit). + assert_eq!( + cache.select(1, 512, 1, PagedDecodeBackend::Cuda), + PagedDecodeDispatch::Native, + "CUDA cell must recompute to native after Metal/Other queries" + ); + } + + #[test] + fn dispatch_cache_cuda_pack_key_roundtrips_without_collision() { + // The 2-bit backend tag (Metal=0, Cuda=1, Other=2) must produce three + // distinct packed keys for one identical shape, and each must stay clear + // of the decision bit and the empty sentinel. + let shape = (4usize, 4096usize, 1usize); + let k_metal = + PagedDispatchCache::pack_key(shape.0, shape.1, shape.2, PagedDecodeBackend::Metal); + let k_cuda = + PagedDispatchCache::pack_key(shape.0, shape.1, shape.2, PagedDecodeBackend::Cuda); + let k_other = + PagedDispatchCache::pack_key(shape.0, shape.1, shape.2, PagedDecodeBackend::Other); + assert_ne!(k_metal, k_cuda, "Metal and CUDA keys must differ"); + assert_ne!(k_cuda, k_other, "CUDA and Other keys must differ"); + assert_ne!(k_metal, k_other, "Metal and Other keys must differ"); + for k in [k_metal, k_cuda, k_other] { + assert_eq!( + k & PagedDispatchCache::DECISION_BIT, + 0, + "packed key must not touch the decision bit" + ); + assert_ne!( + k, PAGED_DISPATCH_CACHE_EMPTY, + "packed key must never equal the empty sentinel" + ); + } + } } diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index c835fcc1..c83d6b93 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -1871,6 +1871,13 @@ mod ffi { /// without a compile-time cfg. See issue #626. fn metal_is_available() -> bool; + /// True when the MLX CUDA backend is available at runtime. False on + /// Metal-only and CPU-only builds. Backend-agnostic (`no_cuda` stub off + /// CUDA), so Rust callers can choose a backend-specific default without + /// a compile-time cfg. Used by the paged-attention decode gate to allow + /// the fused native kernel on CUDA. See issue #634. + fn cuda_is_available() -> bool; + /// Fused sampling: temperature + top-k + top-p + min-p + categorical /// in a single C++ call to minimize FFI round-trips. /// Input: 2D logits [batch, vocab] (already sliced, penalties applied)