Skip to content

refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803

Description

@inureyes

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

  1. 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).
  2. Convert the ten sites to use it. Rocm and None take the existing graph fallback wherever one exists.
  3. 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.
  4. 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.
  5. 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 the MLXCEL_FUSED_MOE=0 run 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_matmul in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificpriority:mediumMedium prioritystatus:backlogIn the backlog, not yet readytype:refactorCode restructuring without changing functionality

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions