Part of #1801. Phase 1. Depends on #1802.
Context
mlxcel's fused kernels pick an implementation with the same idiom in ten places:
const bool use_cuda = !mlx::core::metal::is_available();
Sites: src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:166,1483,1991 and src/lib/mlx-cpp/turbo/fused_rope_append.cpp:421, sampling.cpp:412, fused_norm.cpp:370, paged_attention.cpp:467, paged_attention_v2.cpp:600, paged_attention_v2_merge.cpp:203, sampling_rejection.cpp:779.
On a ROCm build Metal is unavailable, so every site takes the CUDA branch and calls fast::cuda_kernel. The upstream no-CUDA stub throws [cuda_kernel] No CUDA back-end. Some callers are gated and fall back to an MLX graph (sampling, rejection sampling, fused add+RMSNorm, fused RoPE+append, SSM kernels via ssm_kernel_available, paged decode). Two are not:
- fused MoE (
MLXCEL_FUSED_MOE, on by default for mixtral.rs, qwen2_moe.rs, qwen3_moe.rs and others), whose bridge function does not return Result, so the C++ exception is likely to reach std::terminate;
- BitNet
bitlinear_matmul, which has no fallback.
Scope
Replace the Metal-or-CUDA boolean with an explicit backend kind and make every custom-kernel site correct on ROCm, using the existing graph fallbacks. Porting kernels to fast::hip_kernel is out of scope (#1814).
Implementation plan
- Add a single bridge helper, e.g.
enum class GpuKernelBackend { Metal, Cuda, Rocm, None }, resolved from the build (MLX_USE_ROCM / CUDA / Metal compile definitions) and the runtime device. Expose it to Rust through the cxx bridge next to gpu_backend_available (src/lib/mlxcel-core/src/lib.rs:2078).
- Convert the ten sites to use it.
Rocm and None take the existing graph fallback wherever one exists.
- Fused MoE: on
Rocm, default to the SwitchGLU path (the same path MLXCEL_FUSED_MOE=0 selects) and make the fused entry point return an error instead of throwing if it is reached on an unsupported backend.
- BitNet: add a graph fallback for
bitlinear_matmul (dequantize-and-matmul is acceptable), or reject BitNet checkpoints at load on backends without the kernel, with a clear message.
- Log the resolved kernel backend once at startup under the existing diagnostics flag.
Acceptance criteria
Validation
cargo build --release --features rocm
./target/release/mlxcel generate -m models/mlx/Qwen3-30B-A3B-4bit -p "Explain MoE routing." -n 64 --temp 0
MLXCEL_FUSED_MOE=0 ./target/release/mlxcel generate -m models/mlx/Qwen3-30B-A3B-4bit -p "Explain MoE routing." -n 64 --temp 0
References
- Kernel sites listed above.
- Fused MoE opt-out comments:
src/models/mixtral.rs:127, src/models/qwen2_moe.rs:241, src/models/qwen3_moe.rs:469.
- BitNet:
src/models/bitnet.rs, bitlinear_matmul in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp.
Part of #1801. Phase 1. Depends on #1802.
Context
mlxcel's fused kernels pick an implementation with the same idiom in ten places:
Sites:
src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:166,1483,1991andsrc/lib/mlx-cpp/turbo/fused_rope_append.cpp:421,sampling.cpp:412,fused_norm.cpp:370,paged_attention.cpp:467,paged_attention_v2.cpp:600,paged_attention_v2_merge.cpp:203,sampling_rejection.cpp:779.On a ROCm build Metal is unavailable, so every site takes the CUDA branch and calls
fast::cuda_kernel. The upstream no-CUDA stub throws[cuda_kernel] No CUDA back-end.Some callers are gated and fall back to an MLX graph (sampling, rejection sampling, fused add+RMSNorm, fused RoPE+append, SSM kernels viassm_kernel_available, paged decode). Two are not:MLXCEL_FUSED_MOE, on by default formixtral.rs,qwen2_moe.rs,qwen3_moe.rsand others), whose bridge function does not returnResult, so the C++ exception is likely to reachstd::terminate;bitlinear_matmul, which has no fallback.Scope
Replace the Metal-or-CUDA boolean with an explicit backend kind and make every custom-kernel site correct on ROCm, using the existing graph fallbacks. Porting kernels to
fast::hip_kernelis out of scope (#1814).Implementation plan
enum class GpuKernelBackend { Metal, Cuda, Rocm, None }, resolved from the build (MLX_USE_ROCM/ CUDA / Metal compile definitions) and the runtime device. Expose it to Rust through the cxx bridge next togpu_backend_available(src/lib/mlxcel-core/src/lib.rs:2078).RocmandNonetake the existing graph fallback wherever one exists.Rocm, default to the SwitchGLU path (the same pathMLXCEL_FUSED_MOE=0selects) and make the fused entry point return an error instead of throwing if it is reached on an unsupported backend.bitlinear_matmul(dequantize-and-matmul is acceptable), or reject BitNet checkpoints at load on backends without the kernel, with a clear message.Acceptance criteria
turbo/decides CUDA from!metal::is_available()(grep in CI or a unit test).MLXCEL_FUSED_MOE=0), output matching theMLXCEL_FUSED_MOE=0run under greedy decoding.Validation
References
src/models/mixtral.rs:127,src/models/qwen2_moe.rs:241,src/models/qwen3_moe.rs:469.src/models/bitnet.rs,bitlinear_matmulinsrc/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp.