refactor(core): route custom kernels by GPU backend kind

[inureyes]

Ten sites picked a fused kernel with const bool use_cuda = !mlx::core::metal::is_available(), reading "not Metal" as "CUDA". That held while Metal and CUDA were the only GPU backends. With ROCm (#1802) it stopped holding: Metal is unavailable there too, so every site took the CUDA branch and called fast::cuda_kernel, whose no-CUDA stub throws. Most callers catch that, but the throw crosses the cxx bridge into a noexcept extern wherever they do not, which ends the process in std::terminate. On the ROCm gate that was 53 aborts.

src/lib/mlx-cpp/turbo/gpu_backend.h replaces the boolean with GpuKernelBackend (None, Metal, Cuda, Rocm) and custom_kernels_available(). On Metal and CUDA builds the enum resolves to exactly what the boolean did, which is the point: only a ROCm build sees a new value. Four availability predicates spelled metal::is_available() || cu::is_available() become the same call.

Each family already had a fallback; what was missing was a way to reach it. split_kv_enabled(), the paged_decode_batched planner and fused_moe_enabled() now decline on a backend without ports, so the MLA absorbed path, the gather-then-SDPA path and SwitchGLU serve those steps. An eleventh site, paged_decode_backend() in layers.rs, made the same decision in Rust where the issue's grep would not have found it; it now asks the predicate first, so correctness no longer depends on cuda_is_available() staying false on ROCm, and its Other variant stops describing ROCm as a CPU.

Tests reach three launchers directly, past every gate, so the launchers refuse too. paged_attention_decode, paged_attention_decode_v2_partial and paged_attention_merge_states throw a message naming the backend rather than falling through to whichever port was tried, and their bridge declarations return Result so cxx turns the throw into an Err. The C++ signatures are unchanged. paged.rs already carried a comment from the #634 review saying a bare UniquePtr return meant a failed launch would abort rather than error; that is the condition this removes.

bitlinear_matmul is the one fused kernel with no graph fallback at all, so a CPU-only build terminated on a BitNet checkpoint just as ROCm did. BitNet checkpoints are now refused at load with a message naming the reason, the op returns Result, and its test skips with a printed reason where no port exists. The ROCm kernel port is #1862.

Measured on gfx1151: make verify-test-rocm goes from 53 aborts to 2, neither from this issue (nvfp4 group 16 is #1806, and one sampling kill-switch test predates this). 9359 tests pass. paged_v2, mla::split_kv and cache::paged_batch_decode report ordinary failures naming the backend instead of killing the binary.

Not verified here: that Metal and CUDA select the same kernels as before. The enum is equivalent by construction on those backends, but the runners were down and neither host is reachable from this one. The A/B is set up on the Metal side and the CUDA arm is described in a PR comment.

Closes #1803

Validated on gfx1151: cargo fmt --all -- --check, scripts/ci/check_cross_repo_refs.py, scripts/insert_apache_header.py --check, scripts/ci/check_kernel_dtype_keys.py, and make verify-test-rocm.

On that last checker, one thing is worth writing down because the next refactor here will meet it. It scopes itself by whether a file contains the literal cuda_kernel(, so moving a launch call out of these files would take all eight out of scope silently. This PR moves only the selection boolean and leaves the launches where they were, so the scope is unchanged (verified: the eight files still carry the token, and it reports OK). Its glob is *.cpp only, so a launch that ends up in a header is invisible to it whether or not the token is there.

[inureyes]

Running the CUDA arm on a GB10 node

This is the half of the last acceptance criterion that cannot be checked from the machine this branch was written on. It needs a CUDA host, and neither that host nor a GB10 node is reachable from here or from the Metal host running the other arm. Everything below is self-contained: it assumes only a GB10 node with a CUDA toolchain and a clone of this repository, and no context from this PR.

What it is checking: that this branch selects exactly the same kernels on CUDA as main does. The change replaces const bool use_cuda = !mlx::core::metal::is_available() with an enum that resolves to Cuda on a CUDA build, so the two should be equivalent by construction. "Should be" is the reason for the check.

1. Build both arms

main is the baseline, this branch is the arm. The merge base is 9ead2b7d, so main at or after that commit is the right baseline.

git fetch origin
git checkout 9ead2b7d
cargo build --release --features cuda --bin mlxcel
/bin/cp target/release/mlxcel target/release/mlxcel.before

git checkout feature/issue-1803-kernel-backend-kind
cargo build --release --features cuda --bin mlxcel

Check the second build's log before going on. It must show mlx_cxx_kernels.cpp and the files under src/lib/mlx-cpp/turbo/ actually recompiling. A cargo build that reports a fraction of a second and rebuilds nothing after a bridge .cpp edit has been observed on this project, and it makes the arm binary secretly identical to the baseline, at which point every comparison below passes for the wrong reason.

2. The test gate

This is the stronger of the two checks, because the tests that exercise these launchers call them directly.

git checkout 9ead2b7d && make verify-test-cuda 2>&1 | tee /tmp/cuda-before.log
git checkout feature/issue-1803-kernel-backend-kind && make verify-test-cuda 2>&1 | tee /tmp/cuda-after.log
grep '^test result:' /tmp/cuda-before.log | tail -5
grep '^test result:' /tmp/cuda-after.log  | tail -5

Pass condition: the same tests pass in both, and no test that passed before fails after. A new failure in paged_v2::*, mla::split_kv or cache::paged_batch_decode is the signal that matters; those are the paths this PR changed.

3. Greedy output equality

./scripts/ab_output_equality.sh \
    --baseline target/release/mlxcel.before \
    --arm target/release/mlxcel \
    --model models/mlx/Qwen3-30B-A3B-4bit \
    --model models/mlx/Meta-Llama-3.1-8B-Instruct-4bit \
    --prompt "Explain mixture-of-experts routing in two sentences." \
    -n 64

The script pins --temp 0 and --show-reasoning on both arms and runs the baseline twice as a self-consistency control, so it reports INCONCLUSIVE rather than FAIL when the baseline does not reproduce itself. Exit 0 means every pair matched.

If it refuses to start saying the two binaries are byte-identical, that is the no-op build from step 1, not a script problem. Rebuild and check the log.

4. The paged launchers, if you have time for one more

Steps 2 and 3 leave a gap: mlxcel generate has no paged-cache flags, and the fused paged decode kernel is off by default, so the generate arm never reaches paged_attention.cpp or paged_attention_v2*.cpp. The test gate in step 2 does cover them, which is why it is listed first. To exercise them through a server instead:

MLXCEL_PAGED_ATTENTION_NATIVE=1 ./target/release/mlxcel-server \
    -m models/mlx/Meta-Llama-3.1-8B-Instruct-4bit --kv-unified --port 8080

Then send a few chat completions long enough to build context, and compare against the same run on the baseline binary. MLXCEL_PAGED_ATTENTION_NATIVE=1 forces the fused path past the token floors that would otherwise decline it.

Worth watching rather than assuming: that the fused kernel actually ran. This PR makes the paged path ask a new predicate before dispatching, and if that predicate ever answered wrongly on a host that has the kernels, the path would fall back to gather. The output would be identical and only the speed would change, so the equality check in step 3 cannot see it. /slots and the decode-path counters report which path served each step.

What to report back

The two test result: summaries from step 2, the script's exit status and verdict lines from step 3, and the cargo build timing from step 1 that shows the C++ actually recompiled. A plain "all equal" is not enough on its own, because the most likely failure mode here produces exactly that.

[inureyes]

Superseded by #1877. GitHub closed this automatically when the FFT merge deleted its base branch, and a closed PR whose base is gone cannot be reopened or retargeted. Same branch, rebased onto the new main so it carries only the routing commit.