Repository navigation
refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803
Description
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 formixtral.rs,qwen2_moe.rs,qwen3_moe.rsand others), whose bridge function does not returnResult, so the C++ exception is likely to reachstd::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 togpu_backend_available(src/lib/mlxcel-core/src/lib.rs:2078). - Convert the ten sites to use it.
RocmandNonetake the existing graph fallback wherever one exists. - Fused MoE: on
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. - 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
- No site in the bridge or
turbo/decides CUDA from!metal::is_available()(grep in CI or a unit test). - On ROCm, Qwen3-30B-A3B-4bit generates with default settings (no
MLXCEL_FUSED_MOE=0), output matching theMLXCEL_FUSED_MOE=0run under greedy decoding. - On ROCm, a BitNet checkpoint either generates through the fallback or is rejected at load with an actionable message. It never aborts the process.
- Metal and CUDA select exactly the same kernels as before (existing gates green; greedy output byte-identical on one Metal and one CUDA host for a fixed prompt).
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_matmulinsrc/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp.
Activity
test(rocm): add the gfx1151 correctness matrix and its gate
test(rocm): verify mlxcel-server chat completions on gfx1151
feat(rocm): implement Hadamard and SearchSorted, fix narrow index dtypes
fix: refuse byte-identical arms in the A/B output gate (#1867)
feat(rocm): implement FFT through hipFFT
refactor(core): route custom kernels by GPU backend kind
refactor(core): route custom kernels by GPU backend kind
feat(rocm): port the BitNet bitlinear_matmul kernel
feat(rocm): port the BitNet bitlinear_matmul kernel
docs: record #1803 Metal gate, add paged decode counter A/B (#1881)
refactor(moe): migrate qwen3_moe onto the shared SwitchGLU (#1888)
feat(rocm): add a GPU vendor concept and AMD device reporting
feat(rocm): add a GPU vendor concept and AMD device reporting (#1883)
ci(rocm): build, link and generate on a self-hosted gfx1151 runner
ci(rocm): build, link and generate on a self-hosted gfx1151 runner
ci(rocm): build, link and generate on a self-hosted gfx1151 runner (#…
fix(rocm): refuse instead of aborting at the sampler entry points (#2010
fix(rocm): guard the four remaining custom-kernel launchers (#2018)
refactor(core): choose custom-kernel ports through one helper (#2026)
refactor: finish the kernel port standardization (#2029)