Repository navigation
refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803
Copy link
Copy link
Closed
Labels
area:coremlxcel-core: MLX FFI, primitives, KV cache, layersmlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificLinux (CUDA / packaging) specificpriority:mediumMedium priorityMedium prioritystatus:backlogIn the backlog, not yet readyIn the backlog, not yet readytype:refactorCode restructuring without changing functionalityCode restructuring without changing functionality
Description
Activity
- addedstatus:backlogIn the backlog, not yet readyIn the backlog, not yet readytype:refactorCode restructuring without changing functionalityCode restructuring without changing functionalitypriority:mediumMedium priorityMedium priorityarea:coremlxcel-core: MLX FFI, primitives, KV cache, layersmlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificLinux (CUDA / packaging) specific
on Sep 11, 2026 - added a parent issue
on Sep 11, 2026 - added 6 commits that reference this issue
on Sep 12, 2026 - added 7 commits that reference this issue
on Sep 13, 2026 - added 7 commits that reference this issue
on Sep 26, 2026
Metadata
Metadata
Assignees
Labels
area:coremlxcel-core: MLX FFI, primitives, KV cache, layersmlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificLinux (CUDA / packaging) specificpriority:mediumMedium priorityMedium prioritystatus:backlogIn the backlog, not yet readyIn the backlog, not yet readytype:refactorCode restructuring without changing functionalityCode restructuring without changing functionality
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.