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

  1. added
    status:backlogIn the backlog, not yet ready
    type:refactorCode restructuring without changing functionality
    area:coremlxcel-core: MLX FFI, primitives, KV cache, layers
    platform:linuxLinux (CUDA / packaging) specific
    on Sep 11, 2026
  2. added 7 commits that reference this issue on Sep 13, 2026
    d492c81
    7d03956
    d67b64d
    3d8fd7b
    c1234f3
    c683b91
    9457ec5
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