diff --git a/examples/page_gather_microbench.rs b/examples/page_gather_microbench.rs index 8389def0d..e028c49a0 100644 --- a/examples/page_gather_microbench.rs +++ b/examples/page_gather_microbench.rs @@ -669,6 +669,7 @@ fn run_config( scale, 0, ) + .expect("paged_attention_decode: this microbench needs a backend with the fused kernel") }); let mut rot = Rotation::new(rot_count); let fused_v2 = time_body(warmup, iters, || { diff --git a/src/lib/mlx-cpp/turbo/fused_norm.cpp b/src/lib/mlx-cpp/turbo/fused_norm.cpp index ea96579f3..56fccabe5 100644 --- a/src/lib/mlx-cpp/turbo/fused_norm.cpp +++ b/src/lib/mlx-cpp/turbo/fused_norm.cpp @@ -13,6 +13,7 @@ // limitations under the License. #include "fused_norm.h" +#include "gpu_backend.h" #include #include @@ -325,7 +326,7 @@ inline FusedNormKernelHolderCuda& get_fused_norm_kernel_cuda() { } // namespace bool fused_add_rms_norm_available() { - return mlx::core::metal::is_available() || mlx::core::cu::is_available(); + return mlxcel::custom_kernels_available(); } std::vector fused_add_rms_norm( @@ -367,7 +368,8 @@ std::vector fused_add_rms_norm( // Metal kernel on Apple, CUDA port elsewhere. `fast::metal_kernel` throws // "[metal_kernel] No Metal back-end" on the CUDA backend and vice versa; // `metal::is_available()` is false on a CUDA-only build. - const bool use_cuda = !mlx::core::metal::is_available(); + const bool use_cuda = + mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Cuda; auto& kernel = use_cuda ? get_fused_norm_kernel_cuda().get() : get_fused_norm_kernel().get(); diff --git a/src/lib/mlx-cpp/turbo/fused_rope_append.cpp b/src/lib/mlx-cpp/turbo/fused_rope_append.cpp index 401aa3e9b..f01fd0e3e 100644 --- a/src/lib/mlx-cpp/turbo/fused_rope_append.cpp +++ b/src/lib/mlx-cpp/turbo/fused_rope_append.cpp @@ -13,6 +13,7 @@ // limitations under the License. #include "fused_rope_append.h" +#include "gpu_backend.h" #include #include @@ -372,7 +373,7 @@ inline FusedRopeKernelHolderCuda& get_fused_rope_kernel_cuda() { } // namespace bool fused_rope_qk_append_available() { - return mlx::core::metal::is_available() || mlx::core::cu::is_available(); + return mlxcel::custom_kernels_available(); } std::vector fused_rope_qk_append( @@ -418,7 +419,8 @@ std::vector fused_rope_qk_append( "[fused_rope_qk_append] qkv trailing dim does not match the head geometry."); } - const bool use_cuda = !mlx::core::metal::is_available(); + const bool use_cuda = + mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Cuda; auto& kernel = use_cuda ? get_fused_rope_kernel_cuda().get() : get_fused_rope_kernel().get(); diff --git a/src/lib/mlx-cpp/turbo/gpu_backend.cpp b/src/lib/mlx-cpp/turbo/gpu_backend.cpp new file mode 100644 index 000000000..7fdf11d7d --- /dev/null +++ b/src/lib/mlx-cpp/turbo/gpu_backend.cpp @@ -0,0 +1,91 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "gpu_backend.h" + +#include +#include + +// `mlx/backend/rocm/rocm.h` ships in the ROCm overlay +// (`src/lib/mlx-cpp/patches-rocm/`), which CMake copies into the MLX source +// tree only for a ROCm build, so the include has to be gated the same way the +// Metal-backend-only symbols are. build.rs defines this alongside +// MLXCEL_BRIDGE_METAL_BACKEND. +#ifdef MLXCEL_BRIDGE_ROCM_BACKEND +#include +#endif + +#include +#include + +namespace mlxcel { + +namespace { + +GpuKernelBackend resolve_gpu_kernel_backend() { + if (mlx::core::metal::is_available()) { + return GpuKernelBackend::Metal; + } +#ifdef MLXCEL_BRIDGE_ROCM_BACKEND + // Checked before CUDA: on a ROCm build `cu::is_available()` is the no-CUDA + // stub and returns false, but ordering it this way keeps the answer right if + // a future build ever links both. + if (mlx::core::rocm::is_available()) { + return GpuKernelBackend::Rocm; + } +#endif + if (mlx::core::cu::is_available()) { + return GpuKernelBackend::Cuda; + } + return GpuKernelBackend::None; +} + +const char* backend_name(GpuKernelBackend backend) { + switch (backend) { + case GpuKernelBackend::Metal: + return "metal"; + case GpuKernelBackend::Cuda: + return "cuda"; + case GpuKernelBackend::Rocm: + return "rocm"; + case GpuKernelBackend::None: + return "none"; + } + return "unknown"; +} + +} // namespace + +GpuKernelBackend gpu_kernel_backend() { + static const GpuKernelBackend backend = [] { + GpuKernelBackend resolved = resolve_gpu_kernel_backend(); + if (std::getenv("MLXCEL_DEBUG_KERNEL_BACKEND") != nullptr) { + std::fprintf( + stderr, + "[mlxcel] custom kernel backend: %s%s\n", + backend_name(resolved), + custom_kernels_available_for(resolved) + ? "" + : " (no custom kernel ports; using MLX graph fallbacks)"); + } + return resolved; + }(); + return backend; +} + +bool custom_kernels_available() { + return custom_kernels_available_for(gpu_kernel_backend()); +} + +} // namespace mlxcel diff --git a/src/lib/mlx-cpp/turbo/gpu_backend.h b/src/lib/mlx-cpp/turbo/gpu_backend.h new file mode 100644 index 000000000..6d4f4fb31 --- /dev/null +++ b/src/lib/mlx-cpp/turbo/gpu_backend.h @@ -0,0 +1,62 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +// Which GPU backend mlxcel's custom kernels should target (issue #1803). +// +// Every fused-kernel launcher used to decide with `!metal::is_available()`, +// reading "not Metal" as "CUDA". That was true while Metal and CUDA were the +// only GPU backends. It stopped being true with the ROCm backend (issue #1802): +// there Metal is unavailable too, so every site took the CUDA branch and called +// `fast::cuda_kernel`, whose no-CUDA stub throws `[cuda_kernel] No CUDA +// back-end.` Most callers catch that and fall back to an MLX graph, but the +// throw crosses the cxx bridge into a `noexcept` extern wherever they do not, +// which ends the process in `std::terminate`. +// +// The replacement names the backend instead of negating one. On Metal and CUDA +// builds it resolves to exactly what the old boolean did, which is the point: +// `Metal` where `metal::is_available()` was true, `Cuda` otherwise on a build +// that has CUDA. Only a ROCm build sees a new value. + +namespace mlxcel { + +enum class GpuKernelBackend { + // No GPU backend, or one with no custom-kernel port. Callers take their MLX + // graph fallback. + None = 0, + Metal = 1, + Cuda = 2, + // ROCm is a real GPU backend with no `fast::hip_kernel` ports yet (issue + // #1814). Callers treat it like `None` and take the graph fallback; they must + // not call `fast::cuda_kernel`. + Rocm = 3, +}; + +// True when this backend has custom kernel ports, that is Metal or CUDA. This +// is the condition the launchers gate on; it is deliberately not "a GPU is +// present", which is what the old idiom conflated it with. +constexpr bool custom_kernels_available_for(GpuKernelBackend backend) { + return backend == GpuKernelBackend::Metal || + backend == GpuKernelBackend::Cuda; +} + +// Resolved once from the compiled-in backends and the runtime device. Set +// MLXCEL_DEBUG_KERNEL_BACKEND to have the resolved value printed once. +GpuKernelBackend gpu_kernel_backend(); + +// custom_kernels_available_for(gpu_kernel_backend()). +bool custom_kernels_available(); + +} // namespace mlxcel diff --git a/src/lib/mlx-cpp/turbo/paged_attention.cpp b/src/lib/mlx-cpp/turbo/paged_attention.cpp index 83b9dd967..96b832888 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention.cpp @@ -13,6 +13,8 @@ // limitations under the License. #include "paged_attention.h" +#include +#include "gpu_backend.h" #include #include @@ -464,7 +466,20 @@ mlx::core::array paged_attention_decode( // `cuda_kernel` port there; `metal::is_available()` is false on a CUDA-only // build (#634). Both kernels share the template args, grid, and buffer // contract below, so only the JIT-compiled body differs. - const bool use_cuda = !mlx::core::metal::is_available(); + // Refuse before selecting a port, so the message names the real reason + // rather than the port that happened to be tried. mlxcel's Rust callers + // gate on `custom_kernels_available()` and take a graph fallback, so + // reaching this means a direct call; the bridge declares this function + // `Result`, so the throw becomes an `Err` instead of ending the process + // (issue #1803). + if (!mlxcel::custom_kernels_available()) { + throw std::runtime_error( + "[paged_attention_decode] no custom kernel port for this GPU backend; " + "mlxcel's callers take the graph fallback instead"); + } + + const bool use_cuda = + mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Cuda; auto& kernel = use_cuda ? get_paged_attention_kernel_cuda().get() : get_paged_attention_kernel().get(); diff --git a/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp b/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp index b9cf02196..57c3c48d7 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp @@ -13,6 +13,8 @@ // limitations under the License. #include "paged_attention_v2.h" +#include +#include "gpu_backend.h" #include #include @@ -597,7 +599,20 @@ std::vector paged_attention_decode_v2_partial( int dims_per_thread = (dim + PAGED_V2_SIMD_WIDTH - 1) / PAGED_V2_SIMD_WIDTH; int num_warps = paged_attention_v2_num_warps(dim, q_heads); - const bool use_cuda = !mlx::core::metal::is_available(); + // Refuse before selecting a port, so the message names the real reason + // rather than the port that happened to be tried. mlxcel's Rust callers + // gate on `custom_kernels_available()` and take a graph fallback, so + // reaching this means a direct call; the bridge declares this function + // `Result`, so the throw becomes an `Err` instead of ending the process + // (issue #1803). + if (!mlxcel::custom_kernels_available()) { + throw std::runtime_error( + "[paged_attention_decode_v2_partial] no custom kernel port for this GPU backend; " + "mlxcel's callers take the graph fallback instead"); + } + + const bool use_cuda = + mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Cuda; auto& kernel = get_partial_kernel(use_cuda).get(); // `QType`/`KVType` are load-bearing even though the body never names them: diff --git a/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp b/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp index 80fad5bda..bffa52331 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp @@ -21,6 +21,8 @@ // the full contract, including the log2 LSE units. #include "paged_attention_v2.h" +#include +#include "gpu_backend.h" #include #include @@ -200,7 +202,20 @@ std::vector paged_attention_merge_states( num_outputs = 0; } - const bool use_cuda = !mlx::core::metal::is_available(); + // Refuse before selecting a port, so the message names the real reason + // rather than the port that happened to be tried. mlxcel's Rust callers + // gate on `custom_kernels_available()` and take a graph fallback, so + // reaching this means a direct call; the bridge declares this function + // `Result`, so the throw becomes an `Err` instead of ending the process + // (issue #1803). + if (!mlxcel::custom_kernels_available()) { + throw std::runtime_error( + "[paged_attention_merge_states] no custom kernel port for this GPU backend; " + "mlxcel's callers take the graph fallback instead"); + } + + const bool use_cuda = + mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Cuda; auto& kernel = get_merge_kernel(use_cuda).get(); // `VType`/`LseType` key the JIT cache on the input dtypes; see the longer diff --git a/src/lib/mlx-cpp/turbo/sampling.cpp b/src/lib/mlx-cpp/turbo/sampling.cpp index 42f6f6afe..134819ecc 100644 --- a/src/lib/mlx-cpp/turbo/sampling.cpp +++ b/src/lib/mlx-cpp/turbo/sampling.cpp @@ -13,6 +13,7 @@ // limitations under the License. #include "sampling.h" +#include "gpu_backend.h" #include #include @@ -355,7 +356,7 @@ bool gumbel_max_sample_supported() { if (mlx::core::default_device() != mlx::core::Device::gpu) { return false; } - return mlx::core::metal::is_available() || mlx::core::cu::is_available(); + return mlxcel::custom_kernels_available(); } bool gumbel_max_sample_accepts(const mlx::core::array& logits) { @@ -409,7 +410,8 @@ mlx::core::array gumbel_max_sample( // "[metal_kernel] No Metal back-end" on the CUDA backend, so dispatch the // `cuda_kernel` port there; `metal::is_available()` is false on a CUDA-only // build. Both kernels share the template args, grid, and buffer contract. - const bool use_cuda = !mlx::core::metal::is_available(); + const bool use_cuda = + mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Cuda; auto& kernel = use_cuda ? get_gumbel_kernel_cuda().get() : get_gumbel_kernel().get(); diff --git a/src/lib/mlx-cpp/turbo/sampling_rejection.cpp b/src/lib/mlx-cpp/turbo/sampling_rejection.cpp index bfe5f6397..c3f8025e1 100644 --- a/src/lib/mlx-cpp/turbo/sampling_rejection.cpp +++ b/src/lib/mlx-cpp/turbo/sampling_rejection.cpp @@ -13,6 +13,7 @@ // limitations under the License. #include "sampling_rejection.h" +#include "gpu_backend.h" #include #include @@ -746,7 +747,7 @@ bool rejection_sample_supported() { if (mlx::core::default_device() != mlx::core::Device::gpu) { return false; } - return mlx::core::metal::is_available() || mlx::core::cu::is_available(); + return mlxcel::custom_kernels_available(); } bool rejection_sample_accepts(const mlx::core::array& probs) { @@ -776,7 +777,8 @@ RejectionSampleResult rejection_sample( // "[metal_kernel] No Metal back-end" on the CUDA backend, so dispatch the // `cuda_kernel` port there; both share the template args, grid, and buffer // contract. - const bool use_cuda = !mlx::core::metal::is_available(); + const bool use_cuda = + mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Cuda; auto& kernel = use_cuda ? get_rejection_kernel_cuda().get() : get_rejection_kernel().get(); diff --git a/src/lib/mlxcel-core/build.rs b/src/lib/mlxcel-core/build.rs index e45a898d2..f43ae6242 100644 --- a/src/lib/mlxcel-core/build.rs +++ b/src/lib/mlxcel-core/build.rs @@ -85,6 +85,9 @@ fn main() { // `src/lib/mlx-cpp/turbo/` so the MLX-upstream-commit upgrade // checklist ("Bumping the MLX upstream pin" in CONTRIBUTING.md) // treats this directory as in-scope. + // Resolves which GPU backend the custom kernels target (issue + // #1803). Shared by the bridge and every turbo/ launcher. + .file("../mlx-cpp/turbo/gpu_backend.cpp") .file("../mlx-cpp/turbo/sparse_v_sdpa.cpp") // Fused Turbo4Delegated cold-V weighted-sum kernel // launcher. Reads the packed cold V directly so the dequantised @@ -154,6 +157,14 @@ fn main() { bridge.define("MLXCEL_BRIDGE_METAL_BACKEND", None); } + // Same shape for ROCm: `mlx/backend/rocm/rocm.h` ships in the ROCm overlay + // and CMake copies it into the MLX source tree only for a ROCm build, so + // the bridge can only include it when that build is the one running + // (issue #1803). + if std::env::var("CARGO_FEATURE_ROCM").is_ok() { + bridge.define("MLXCEL_BRIDGE_ROCM_BACKEND", None); + } + // Add optimization flags for release builds #[cfg(not(debug_assertions))] { diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index dbfb97b6f..d93c48546 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -2,6 +2,7 @@ // Direct C++ bridge implementation for MLX via cxx #include "mlx_cxx_internal.h" +#include "../../mlx-cpp/turbo/gpu_backend.h" #include "mlx/primitives.h" @@ -4830,6 +4831,14 @@ bool gpu_backend_available() { return mlx::core::device_count(mlx::core::Device::gpu) > 0; } +// True when the resolved GPU backend has mlxcel's fused kernel ports, that is +// Metal or CUDA (issue #1803). ROCm answers false until #1814 ports them, so +// the Rust callers take the graph fallback each family already has instead of +// reaching a `fast::cuda_kernel` that throws. +bool custom_kernels_available() { + return mlxcel::custom_kernels_available(); +} + // Top-p (nucleus) filtering. // Not compiled: MLX v0.31.x Scan primitive (cumsum) lacks output_shapes, // which causes "CumSum cannot infer output shapes" when used inside diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 2dd75d171..60afa3bc4 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1313,6 +1313,11 @@ bool default_device_is_gpu(); // default device; never moves it or creates a stream. Issue #1421. bool gpu_backend_available(); +// True when the resolved GPU backend has mlxcel's fused kernel ports, that is +// Metal or CUDA (issue #1803). Narrower than `gpu_backend_available`, which +// only says a GPU exists. +bool custom_kernels_available(); + // Fused sampling: top-k + top-p + min-p on the untempered distribution, then // one temperature scaling, then the categorical draw, in a single function // call to minimize FFI round-trips (chain order per issue #1379). diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp index 798a1ecfe..d7e2aa939 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp @@ -4,6 +4,7 @@ // mlx_cxx_bridge.cpp; see mlx_cxx_internal.h for the shared helpers. #include "mlx_cxx_internal.h" +#include "../../mlx-cpp/turbo/gpu_backend.h" // CUDA backend availability probe for the fused-kernel gates (#631). The // header is backend-agnostic: builds without CUDA link the no_cuda stub. @@ -142,6 +143,7 @@ namespace { static BitlinearKernelHolderCuda holder; return holder; } + } std::unique_ptr bitlinear_matmul( @@ -163,9 +165,25 @@ std::unique_ptr bitlinear_matmul( // mx.fast.metal_kernel throws on CUDA, so dispatch the cuda_kernel port // there. metal::is_available() is false on a CUDA-only build. - const bool use_cuda = !mlx::core::metal::is_available(); - auto& kernel = use_cuda ? get_bitlinear_kernel_cuda().get() - : get_bitlinear_kernel().get(); + // Named backends rather than a negation of Metal (issue #1803). Backends + // with no port throw here, and the bridge declares this function `Result`, + // so that reaches the caller as an error instead of ending the process: + // this op is the one fused kernel with no graph fallback, so before this a + // CPU-only build and a ROCm build both terminated on a BitNet checkpoint. + auto& kernel = [&]() -> mlx::core::fast::CustomKernelFunction& { + switch (mlxcel::gpu_kernel_backend()) { + case mlxcel::GpuKernelBackend::Metal: + return get_bitlinear_kernel().get(); + case mlxcel::GpuKernelBackend::Cuda: + return get_bitlinear_kernel_cuda().get(); + case mlxcel::GpuKernelBackend::Rocm: + case mlxcel::GpuKernelBackend::None: + break; + } + throw std::runtime_error( + "[bitlinear_matmul] no BitLinear kernel port for this GPU backend; " + "BitNet needs Metal or CUDA (ROCm port: lablup/mlxcel#1862)"); + }(); std::vector> ta = { {"T", T}, {"in_features", in_features}, @@ -1480,7 +1498,8 @@ void ssm_update_kernel( // Metal kernel on Apple, CUDA port elsewhere (mx.fast.metal_kernel throws // "[metal_kernel] No Metal back-end" on the CUDA backend), cf. #631. - const bool use_cuda = !mlx::core::metal::is_available(); + const bool use_cuda = + mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Cuda; auto& kernel = use_cuda ? get_ssm_kernel_cuda().get() : get_ssm_kernel().get(); // CustomKernelFunction signature: @@ -1988,7 +2007,8 @@ std::unique_ptr run_fused_moe_two_kernel( // mx.fast.metal_kernel throws on CUDA ("No Metal back-end"), so dispatch the // cuda_kernel port there. metal::is_available() is false on a CUDA-only build. - const bool use_cuda = !mlx::core::metal::is_available(); + const bool use_cuda = + mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Cuda; // A) gate/up + activation -> act_g[K, Dff] (f32 for the down GEMV). auto& kA = use_cuda ? get_moe_gateup_kernel_cuda().get() diff --git a/src/lib/mlxcel-core/src/autotune/ops/paged_decode_splits.rs b/src/lib/mlxcel-core/src/autotune/ops/paged_decode_splits.rs index 6fa2a078d..b50d55bf3 100644 --- a/src/lib/mlxcel-core/src/autotune/ops/paged_decode_splits.rs +++ b/src/lib/mlxcel-core/src/autotune/ops/paged_decode_splits.rs @@ -259,7 +259,8 @@ impl TunableOp for PagedDecodeSplitsOp<'_> { self.visible_lens, self.scale, splits, - ); + ) + .map_err(|e| TuneError::infeasible(tactic, format!("launch failed: {e}")))?; crate::eval(&out); Ok(()) } diff --git a/src/lib/mlxcel-core/src/cache/paged.rs b/src/lib/mlxcel-core/src/cache/paged.rs index 02113c662..3efa2fe0c 100644 --- a/src/lib/mlxcel-core/src/cache/paged.rs +++ b/src/lib/mlxcel-core/src/cache/paged.rs @@ -1972,6 +1972,21 @@ impl PagedBlockPool { select_paged_v2_dispatch, }; + // Both fused paths end in custom kernels with Metal and CUDA ports + // only. On a backend without them the launcher would reach + // `fast::cuda_kernel`, whose throw crosses the cxx bridge into a + // `noexcept` extern and ends the process, so decline before planning + // and let the caller's gather-then-SDPA fallback serve the step + // (issue #1803). + if !crate::ffi::custom_kernels_available() { + return Ok(( + None, + PagedDecodeOutcome::NotServable( + "the GPU backend has no fused paged-attention kernel port", + ), + )); + } + // Same single-slab restriction as v1 and as `paged_decode_fused_v2`: // both kernels read one contiguous pool buffer per side, so a layer // grown past one slab (#235) is declined rather than stitched. This is @@ -2288,7 +2303,8 @@ impl PagedBlockPool { let out_f32 = ffi::paged_attention_decode( q_in, pool_k, pool_v, &rows_arr, &off_arr, &ls_arr, &vl_arr, scale, num_splits, - ); + ) + .map_err(|e| format!("paged_attention_decode: {e}"))?; let out = if q_dtype == crate::dtype::FLOAT32 { out_f32 } else { diff --git a/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index b0bdd7849..28c8f4241 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/src/lib/mlxcel-core/src/layers.rs @@ -5032,7 +5032,11 @@ pub enum PagedDecodeBackend { /// pool blocks with no gather pass, so its advantage grows with context; /// the Metal-measured batch/context ceilings do not apply. Cuda, - /// Anything else (CPU): no fused kernel, always gather. + /// Any backend with no fused kernel port, which today means CPU-only + /// builds and ROCm (issue #1803). Always gather. ROCm is a real GPU here, + /// not a CPU, and gets its own variant when its kernels are ported + /// (issue #1814); until then the answer is the same, so it shares this one + /// rather than pretending the ports exist. Other, } @@ -5233,7 +5237,14 @@ pub(crate) fn paged_decode_backend() -> PagedDecodeBackend { use std::sync::OnceLock; static BACKEND: OnceLock = OnceLock::new(); *BACKEND.get_or_init(|| { - if crate::metal_is_available() { + // Asked first, so a backend with no fused kernel port can never be + // read as one that has it. `cuda_is_available()` is false on a ROCm + // build today, so this is belt and braces rather than a live fix, but + // the old order made correctness depend on that staying true + // (issue #1803). + if !crate::custom_kernels_available() { + PagedDecodeBackend::Other + } else if crate::metal_is_available() { PagedDecodeBackend::Metal } else if crate::cuda_is_available() { PagedDecodeBackend::Cuda diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index 8d55e0590..70d55b13d 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -1336,6 +1336,13 @@ mod ffi { /// ceiling, which is exactly the pre-#906 behavior; out-of-range values /// clamp back to the ceiling, so a stale cached tactic can never /// produce an infeasible launch. + /// Paged-attention decode launcher. + /// + /// Returns `Err` instead of ending the process when the GPU backend has + /// no custom kernel port (issue #1803): the C++ launcher throws and cxx + /// turns that into an `Err` here, which a `noexcept` extern could not + /// do. mlxcel's own callers gate on `custom_kernels_available()` and + /// never reach it. fn paged_attention_decode( q: &MlxArray, k_pool: &MlxArray, @@ -1346,7 +1353,7 @@ mod ffi { visible_lens: &MlxArray, scale: f32, num_splits_override: i32, - ) -> UniquePtr; + ) -> Result>; /// Largest feasible `NumSplits` for a head dimension (issue #906). /// @@ -1383,7 +1390,7 @@ mod ffi { scale: f32, partial_v_out: &mut UniquePtr, lse_out: &mut UniquePtr, - ); + ) -> Result<()>; /// Variable-length attention-state merge kernel (issue #898). /// @@ -1392,13 +1399,20 @@ mod ffi { /// into `[M, H, D]` / `[M, H]`, where output row `o` covers partial /// rows `[o_indptr[o], o_indptr[o + 1])`. Deliberately paging-agnostic: /// the cascade-attention issue #903 reuses it unchanged. + /// Merge partial paged-attention states. + /// + /// Returns `Err` instead of ending the process when the GPU backend has + /// no custom kernel port (issue #1803): the C++ launcher throws and cxx + /// turns that into an `Err` here, which a `noexcept` extern could not + /// do. mlxcel's own callers gate on `custom_kernels_available()` and + /// never reach it. fn paged_attention_merge_states( v_in: &MlxArray, lse_in: &MlxArray, o_indptr: &MlxArray, v_out: &mut UniquePtr, lse_out: &mut UniquePtr, - ); + ) -> Result<()>; /// Query heads one v2 CTA processes together (issue #898). Always /// divides `n_rep`, so the plan's CTA count and the launcher's grid @@ -1780,6 +1794,13 @@ mod ffi { /// [out_features/4, in_features] uint8 (2-bit ternary, 4 rows/byte), /// scaled by `weight_scale[0]` (inverted unless linear_class is /// autobitlinear). + /// BitLinear ternary matmul, with Metal and CUDA kernel ports. + /// + /// Returns `Err` on a backend with neither, which today means ROCm and + /// CPU-only builds, instead of ending the process: the op has no graph + /// fallback, so before issue #1803 the `fast::*_kernel` throw crossed a + /// `noexcept` extern and terminated. Callers that want to refuse early + /// should ask [`custom_kernels_available`]. fn bitlinear_matmul( x: &MlxArray, packed_weights: &MlxArray, @@ -1787,7 +1808,7 @@ mod ffi { in_features: i32, out_features: i32, invert_weight_scales: bool, - ) -> UniquePtr; + ) -> Result>; /// Fused gated-delta single-token decode step. /// Combines: decay → kv_mem → delta → state_update → output into one C++ call. @@ -2077,6 +2098,15 @@ mod ffi { /// before `initialize_runtime` finishes. fn gpu_backend_available() -> bool; + /// True when the resolved GPU backend has mlxcel's fused kernel ports, + /// that is Metal or CUDA (issue #1803). ROCm has a GPU but no ports + /// yet, so it answers `false` and callers take the MLX graph fallback + /// their family already has. This is deliberately narrower than + /// `gpu_backend_available`, which only says a GPU exists: conflating + /// the two is what sent ROCm into `fast::cuda_kernel` and aborted the + /// process. + fn custom_kernels_available() -> bool; + /// True when the MLX Metal backend is available at runtime (macOS /// Apple Silicon). False on CUDA-only and CPU-only builds. Mirrors the /// `mlx::core::metal::is_available()` gate that picks the Metal vs CUDA diff --git a/src/lib/mlxcel-core/src/mla/mod.rs b/src/lib/mlxcel-core/src/mla/mod.rs index 163235e00..eaf5bdb22 100644 --- a/src/lib/mlxcel-core/src/mla/mod.rs +++ b/src/lib/mlxcel-core/src/mla/mod.rs @@ -139,6 +139,13 @@ pub fn split_kv_enabled() -> bool { static ENABLED: OnceLock = OnceLock::new(); *ENABLED.get_or_init(|| parse_enabled(std::env::var(SPLIT_KV_ENV).ok().as_deref())) && absorbed_enabled() + // The split path ends in `paged_attention_merge_states`, a custom + // kernel with Metal and CUDA ports only. On a backend without them the + // launcher would reach `fast::cuda_kernel`, whose throw crosses the + // cxx bridge into a `noexcept` extern and ends the process. Declining + // here takes the `absorbed_decode` fallback the caller already has + // for a plan that declines (issue #1803). + && crate::ffi::custom_kernels_available() } /// The per-layer MLA shape the absorbed path needs. diff --git a/src/lib/mlxcel-core/src/mla/split_kv.rs b/src/lib/mlxcel-core/src/mla/split_kv.rs index 63f314508..cdc1ce838 100644 --- a/src/lib/mlxcel-core/src/mla/split_kv.rs +++ b/src/lib/mlxcel-core/src/mla/split_kv.rs @@ -302,7 +302,8 @@ pub fn absorbed_decode_split_kv( let mut v_out = UniquePtr::null(); let mut lse_out = UniquePtr::null(); - ffi::paged_attention_merge_states(&v_in, &lse_in, &o_indptr, &mut v_out, &mut lse_out); + ffi::paged_attention_merge_states(&v_in, &lse_in, &o_indptr, &mut v_out, &mut lse_out) + .map_err(|e| format!("mla split: merge launch failed: {e}"))?; v_out } else { // One chunk per request: the partial is the answer, no merge launch. diff --git a/src/lib/mlxcel-core/src/mla/split_kv_tests.rs b/src/lib/mlxcel-core/src/mla/split_kv_tests.rs index 8f3065a77..b858f946d 100644 --- a/src/lib/mlxcel-core/src/mla/split_kv_tests.rs +++ b/src/lib/mlxcel-core/src/mla/split_kv_tests.rs @@ -187,13 +187,8 @@ fn merge_rejects_natural_log_lse_units() { let lse_arr = crate::ffi::from_slice_f32(lse, &[2, HEADS as i32]); let mut out_v = cxx::UniquePtr::null(); let mut out_lse = cxx::UniquePtr::null(); - crate::ffi::paged_attention_merge_states( - &v_arr, - &lse_arr, - &indptr, - &mut out_v, - &mut out_lse, - ); + ffi::paged_attention_merge_states(&v_arr, &lse_arr, &indptr, &mut out_v, &mut out_lse) + .expect("merge_states"); to_vec_f32(&out_v) }; diff --git a/src/lib/mlxcel-core/src/paged_v2/cascade_launch.rs b/src/lib/mlxcel-core/src/paged_v2/cascade_launch.rs index b6e9523d0..e74ab3171 100644 --- a/src/lib/mlxcel-core/src/paged_v2/cascade_launch.rs +++ b/src/lib/mlxcel-core/src/paged_v2/cascade_launch.rs @@ -247,7 +247,8 @@ pub fn run_cascade_decode( let mut merged_v = UniquePtr::null(); let mut merged_lse = UniquePtr::null(); - ffi::paged_attention_merge_states(&v_in, &lse_in, &o_indptr, &mut merged_v, &mut merged_lse); + ffi::paged_attention_merge_states(&v_in, &lse_in, &o_indptr, &mut merged_v, &mut merged_lse) + .map_err(|e| format!("cascade merge launch failed: {e}"))?; let stats = CascadeLaunchStats { prefix_chunks: prefix_plan.num_chunks, diff --git a/src/lib/mlxcel-core/src/paged_v2/cascade_launch_tests.rs b/src/lib/mlxcel-core/src/paged_v2/cascade_launch_tests.rs index b11522718..be3f2c099 100644 --- a/src/lib/mlxcel-core/src/paged_v2/cascade_launch_tests.rs +++ b/src/lib/mlxcel-core/src/paged_v2/cascade_launch_tests.rs @@ -442,7 +442,8 @@ fn run_cascade_variant( let o_indptr = ffi::from_slice_i32(&plan.o_indptr, &[b + 1]); let mut merged_v = UniquePtr::null(); let mut merged_lse = UniquePtr::null(); - ffi::paged_attention_merge_states(&v_in, &lse_in, &o_indptr, &mut merged_v, &mut merged_lse); + ffi::paged_attention_merge_states(&v_in, &lse_in, &o_indptr, &mut merged_v, &mut merged_lse) + .expect("merge_states"); to_vec_f32(&merged_v) } diff --git a/src/lib/mlxcel-core/src/paged_v2/launch.rs b/src/lib/mlxcel-core/src/paged_v2/launch.rs index 7d1d5b852..d78c8ff27 100644 --- a/src/lib/mlxcel-core/src/paged_v2/launch.rs +++ b/src/lib/mlxcel-core/src/paged_v2/launch.rs @@ -134,7 +134,8 @@ impl<'a> V2Context<'a> { self.scale, &mut partial_v, &mut lse, - ); + ) + .map_err(|e| format!("paged decode v2: partial launch failed: {e}"))?; let batch = i32::try_from(plan.batch) .map_err(|_| format!("paged decode v2: batch {} overflows i32", plan.batch))?; @@ -154,7 +155,8 @@ impl<'a> V2Context<'a> { &o_indptr, &mut merged_v, &mut merged_lse, - ); + ) + .map_err(|e| format!("paged decode v2: merge launch failed: {e}"))?; Ok((merged_v, merged_lse)) } } diff --git a/src/lib/mlxcel-core/src/paged_v2/launch_tests.rs b/src/lib/mlxcel-core/src/paged_v2/launch_tests.rs index f35c020f1..8989b1f07 100644 --- a/src/lib/mlxcel-core/src/paged_v2/launch_tests.rs +++ b/src/lib/mlxcel-core/src/paged_v2/launch_tests.rs @@ -522,7 +522,8 @@ fn merge_kernel_matches_the_closed_form() { let indptr_arr = ffi::from_slice_i32(&indptr, &[indptr.len() as i32]); let mut out_v = UniquePtr::null(); let mut out_lse = UniquePtr::null(); - ffi::paged_attention_merge_states(&v_arr, &lse_arr, &indptr_arr, &mut out_v, &mut out_lse); + ffi::paged_attention_merge_states(&v_arr, &lse_arr, &indptr_arr, &mut out_v, &mut out_lse) + .expect("merge_states"); let got_v = to_vec_f32(&out_v); let got_lse = to_vec_f32(&out_lse); @@ -580,7 +581,8 @@ fn merge_kernel_returns_zeros_for_an_all_empty_row() { let indptr_arr = ffi::from_slice_i32(&indptr, &[2]); let mut out_v = UniquePtr::null(); let mut out_lse = UniquePtr::null(); - ffi::paged_attention_merge_states(&v_arr, &lse_arr, &indptr_arr, &mut out_v, &mut out_lse); + ffi::paged_attention_merge_states(&v_arr, &lse_arr, &indptr_arr, &mut out_v, &mut out_lse) + .expect("merge_states"); assert!(to_vec_f32(&out_v).iter().all(|&x| x == 0.0)); assert!( to_vec_f32(&out_lse) @@ -606,7 +608,8 @@ fn merge_is_associative_across_regroupings() { let arr = ffi::from_slice_i32(ptr, &[ptr.len() as i32]); let mut ov = UniquePtr::null(); let mut ol = UniquePtr::null(); - ffi::paged_attention_merge_states(v_in, lse_in, &arr, &mut ov, &mut ol); + ffi::paged_attention_merge_states(v_in, lse_in, &arr, &mut ov, &mut ol) + .expect("merge_states"); (ov, ol) }; diff --git a/src/models/bitnet.rs b/src/models/bitnet.rs index dc371a206..c951a06ed 100644 --- a/src/models/bitnet.rs +++ b/src/models/bitnet.rs @@ -32,6 +32,27 @@ use mlxcel_core::{MlxArray, UniquePtr}; use serde::Deserialize; use std::path::Path; +/// Refuse a BitNet checkpoint on a backend with no `bitlinear_matmul` port. +/// +/// Every other fused kernel in mlxcel has an MLX graph fallback; this one does +/// not, so before issue #1803 the launcher reached `fast::cuda_kernel` on any +/// non-Metal backend and its throw crossed the cxx bridge into a `noexcept` +/// extern, ending the process. That was never ROCm-only: a CPU-only build took +/// the same path. Metal, CUDA and ROCm all have ports now (issue #1862), so +/// this refuses only where no kernel exists at all, and it refuses at load +/// rather than mid-request. +fn reject_without_bitlinear_kernel() -> Result<(), String> { + if mlxcel_core::custom_kernels_available() { + return Ok(()); + } + Err( + "BitNet checkpoints need the bitlinear_matmul kernel, which has Metal and CUDA ports only. \ + This backend has neither and the op has no graph fallback, so the model cannot run here. \ + The ROCm port is tracked as lablup/mlxcel#1862." + .to_string(), + ) +} + // Config. #[derive(Debug, Clone, Deserialize)] pub struct BitNetConfig { @@ -133,7 +154,12 @@ impl BitLinear { self.in_features, self.out_features, self.invert_weight_scales, - ); + ) + // `reject_without_bitlinear_kernel` refuses the checkpoint at load on a + // backend with no port, so reaching here means one exists. A panic + // unwinds and fails the request; the pre-#1803 behaviour was a C++ + // throw across a `noexcept` extern, which ended the process. + .expect("bitlinear_matmul: no kernel port, which load should have refused"); match &self.bias { Some(b) => mlxcel_core::add(&y, b), None => y, @@ -417,6 +443,7 @@ impl BitNetModel { } pub fn load>(model_dir: P) -> Result<(Self, BitNetConfig), String> { + reject_without_bitlinear_kernel()?; let model_dir = model_dir.as_ref(); let config_str = std::fs::read_to_string(model_dir.join("config.json")) .map_err(|e| format!("Failed to read config.json: {}", e))?; @@ -428,6 +455,7 @@ impl BitNetModel { } pub fn from_weights(weights: &WeightMap, config: &BitNetConfig) -> Result { + reject_without_bitlinear_kernel()?; // Embedding / lm_head are full precision (not BitLinear); group_size/bits // only matter for the affine-quantized -4/6/8bit variants. let embed_tokens = UnifiedEmbedding::from_weights(weights, "model.embed_tokens", 64, 4)?; @@ -512,7 +540,17 @@ mod tests { let packed = mlxcel_core::from_bytes(&[134u8, 137, 100, 97], &[1, 4], mlxcel_core::dtype::UINT8); let scale = mlxcel_core::from_slice_f32(&[2.0], &[1]); - let y = mlxcel_core::bitlinear_matmul(&x, &packed, &scale, 4, 4, false); + if !mlxcel_core::custom_kernels_available() { + // No BitLinear port on this backend, so there is nothing to check. + // Printed rather than silent: a gate run that skips this should say + // so, not look like it passed (issue #1803). + eprintln!( + "skipped bitlinear_matmul_known_ternary_case: this backend has no BitLinear kernel port" + ); + return; + } + let y = mlxcel_core::bitlinear_matmul(&x, &packed, &scale, 4, 4, false) + .expect("bitlinear_matmul"); let expected = mlxcel_core::from_slice_f32(&[-4.0, -4.0, 8.0, 6.0], &[1, 4]); let diff = mlxcel_core::abs(&mlxcel_core::subtract(&y, &expected)); let max_abs = mlxcel_core::item_f32(&mlxcel_core::max_axis( diff --git a/src/models/switch_layers.rs b/src/models/switch_layers.rs index 92d05bf17..cdf50ca6b 100644 --- a/src/models/switch_layers.rs +++ b/src/models/switch_layers.rs @@ -87,6 +87,12 @@ pub fn fused_moe_enabled() -> bool { static ENABLED: std::sync::OnceLock = std::sync::OnceLock::new(); *ENABLED .get_or_init(|| fused_moe_enabled_from(std::env::var("MLXCEL_FUSED_MOE").ok().as_deref())) + // The fused MoE launcher has Metal and CUDA ports only, and its bridge + // function does not return `Result`, so on a backend without a port the + // `fast::cuda_kernel` throw ends the process rather than failing the + // call. Answering false here selects the same SwitchGLU path + // `MLXCEL_FUSED_MOE=0` selects (issue #1803). + && mlxcel_core::custom_kernels_available() } /// Pure decision behind [`fused_moe_enabled`], split out so it can be unit-tested