fix(rocm): the Gumbel-max and rejection sampler entry points have no ROCm arm and abort on metal_kernel #1885

Phase 1 ROCm work under epic #1801. Found while validating PR #1883 (issue #1805) on the ROCm spike host: AMD Ryzen AI MAX+ 395 / Radeon 8060S, gfx1151, RDNA 3.5, wave32, ROCm 10.0.0, Debian 13, branch feature/issue-1805-gpu-vendor. Pre-existing defect, not a regression from that PR.

Context

cargo test --workspace --profile test-fast --features rocm fails two targets, -p mlxcel --test sampling_gumbel_kill_switch and -p mlxcel --test sampling_rejection_kill_switch. The failure is a process abort, not an assertion:

running 1 test
test falsy_env_restores_the_categorical_sampling_path ... [mlx-rocm] bound HIP device 0: gfx1151 (AMD Radeon 8060S Graphics) cus=20 warp=32 lds=64KB
terminate called after throwing an instance of 'std::runtime_error'
  what():  [metal_kernel] No Metal back-end.

Both sampler launchers dispatch on a two-valued boolean whose false arm still means "Metal" rather than "no port for this backend". src/lib/mlx-cpp/turbo/sampling.cpp:412-415 and src/lib/mlx-cpp/turbo/sampling_rejection.cpp:779-783 are both const bool use_cuda = mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Cuda; followed by use_cuda ? <cuda kernel> : <metal kernel>. Issue #1803 replaced the old !metal::is_available() idiom at these sites, so the CUDA arm is now selected correctly, but it did not add a refusal for GpuKernelBackend::Rocm or None. On ROCm the false arm is taken and mlx::core::fast::metal_kernel throws.

The routing gate is correct and is not the defect. gumbel_max_sample_supported() at src/lib/mlx-cpp/turbo/sampling.cpp:355-360 and its rejection counterpart at src/lib/mlx-cpp/turbo/sampling_rejection.cpp:750 both return mlxcel::custom_kernels_available(), which is false on ROCm per custom_kernels_available_for in src/lib/mlx-cpp/turbo/gpu_backend.h:50-53, so production sampling takes the random::categorical fallback and never reaches the kernel. What is unguarded is the DIRECT entry point. tests/sampling_gumbel_kill_switch.rs:87 calls gumbel_max_sample(&batched, 1.0) directly and tests/sampling_rejection_kill_switch.rs:118 calls fused_sample_rejection(&batched, 1.0, 40, 0.9, 0.0, 32) directly, both under the comment "The kernel entry point itself stays callable: the switch gates routing, not the kernel, so an explicit caller and the benchmark still work." That premise holds only on a backend that has a port, which today is Metal and CUDA. Benchmarks that call either entry point directly have the same exposure.

The abort rather than an error is a second, independent defect: src/lib/mlxcel-core/src/lib.rs:2220 declares fn gumbel_max_sample(logits: &MlxArray, temperature: f32) -> UniquePtr<MlxArray>; and src/lib/mlxcel-core/src/lib.rs:2237-2244 declares fused_sample_rejection the same way, neither as Result. A C++ throw crossing a noexcept cxx extern ends in std::terminate, so even a deliberate refusal at these sites would kill the process as written.

Scope

In scope: src/lib/mlx-cpp/turbo/sampling.cpp (the gumbel_max_sample launcher), src/lib/mlx-cpp/turbo/sampling_rejection.cpp (the fused_sample_rejection launcher), the corresponding cxx declarations in src/lib/mlxcel-core/src/lib.rs, their Rust callers, and the two kill-switch tests plus any benchmark that calls either entry point directly.

Out of scope: HIP ports of the two sampler kernels (see the recommendation below; that work belongs to #1814). Any change to gumbel_max_sample_supported() / the rejection support predicate, which are already correct. Any change to the random::categorical and argpartition fallback paths.

Proposed solution

Two candidates were considered.

(a) Add HIP ports of both sampler kernels, the way #1862 added the ROCm BitLinear kernel, so the direct entry point works on every backend. This removes the abort as a side effect but is a full kernel-porting effort with its own numerical-equivalence burden.

(b) Make the direct entry point refuse cleanly on a backend with no port, and have the tests and benchmarks skip the direct-call assertion when custom_kernels_available() is false.

Recommend (b) as the immediate fix, because it removes a process abort, is small, and matches an existing in-tree precedent; keep (a) as follow-up performance work under #1814. The precedent is paged_attention_decode, fixed under #1803: src/lib/mlx-cpp/turbo/paged_attention.cpp:470-479 refuses with 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"); } placed BEFORE the port is selected, so the message names the real reason rather than the port that happened to be tried, and src/lib/mlxcel-core/src/lib.rs:1346-1356 declares that bridge function -> Result<UniquePtr<MlxArray>> so the throw becomes an Err. src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp:207-213 follows the same shape. Reuse it verbatim rather than inventing a second refusal idiom.

The refusal must be a typed Rust-side error, never a bare C++ throw across a noexcept extern. Changing the launcher without also changing the cxx declaration to -> Result<...> converts one abort message into a different abort message and fixes nothing.

Implementation plan

1. In src/lib/mlx-cpp/turbo/sampling.cpp, insert a custom_kernels_available() guard immediately before the use_cuda computation at line 412, throwing std::runtime_error("[gumbel_max_sample] no custom kernel port for this GPU backend; mlxcel's callers take the categorical fallback instead"). Match the wording and placement of paged_attention.cpp:470-479.
2. Do the same in src/lib/mlx-cpp/turbo/sampling_rejection.cpp before line 779, with [fused_sample_rejection] as the prefix and "take the argpartition fallback instead" as the tail.
3. Change the cxx declarations at src/lib/mlxcel-core/src/lib.rs:2220 (gumbel_max_sample) and src/lib/mlxcel-core/src/lib.rs:2237-2244 (fused_sample_rejection) to -> Result<UniquePtr<MlxArray>>. Check whether fused_sample_rejection_deferred (src/lib/mlxcel-core/src/lib.rs:2251) shares the launcher and needs the same treatment.
4. Update every Rust caller of the two now-fallible functions. Production callers already gate on the support predicate, so the expected shape is an expect with a message naming the gate, or propagation where the caller returns Result. Do not add a silent unwrap_or_else fallback that would mask a real gating bug on Metal or CUDA.
5. Gate the direct-call sections of tests/sampling_gumbel_kill_switch.rs:93-96 and tests/sampling_rejection_kill_switch.rs:115-124 on the support predicate: keep the existing assertions when a port exists, and assert the typed Err (not a skip that asserts nothing) when custom_kernels_available() is false, so the refusal itself is covered.
6. Audit benches/ for direct calls to either entry point and apply the same gate. Grep for gumbel_max_sample and fused_sample_rejection across the workspace to find every direct caller.
7. Update the stale dispatch comments at sampling.cpp:408-411 and sampling_rejection.cpp:775-778, both of which still say "Metal kernel on Apple, CUDA port elsewhere" and describe a two-backend world.

Acceptance criteria

• cargo test --workspace --profile test-fast --features rocm passes -p mlxcel --test sampling_gumbel_kill_switch and -p mlxcel --test sampling_rejection_kill_switch on the gfx1151 host, with no terminate called after throwing in the output.
• Calling gumbel_max_sample or fused_sample_rejection directly on a ROCm build returns a typed Err whose message names the missing port, and does not abort the process.
• Both cxx declarations are -> Result<...>; no refusal path throws across a noexcept extern.
• A test asserts the Err on a backend without a port, rather than only skipping.
• The kill-switch tests still assert the full direct-call behavior on Metal and CUDA, unchanged.
• CUDA behavior is byte-identical: the same kernel is selected and no new error path is reachable there.
• The refusal is wired into the real call path (the cxx bridge and its Rust callers), not left as an unreferenced helper.
• The stale two-backend dispatch comments at both sites are corrected.

Validation

cargo test --workspace --profile test-fast --features rocm -p mlxcel --test sampling_gumbel_kill_switch
cargo test --workspace --profile test-fast --features rocm -p mlxcel --test sampling_rejection_kill_switch
cargo test --workspace --profile test-fast --features rocm
cargo clippy --workspace --all-targets --features rocm -- -D warnings
cargo fmt --all -- --check

Regression guard on a CUDA host, which must stay unchanged per #1805:

cargo test --workspace --profile test-fast --features cuda -p mlxcel --test sampling_gumbel_kill_switch
cargo test --workspace --profile test-fast --features cuda -p mlxcel --test sampling_rejection_kill_switch

A pass is both ROCm targets green with no terminate called after throwing an instance of 'std::runtime_error' line anywhere in the output.

References

• Epic epic: AMD GPU (ROCm) backend on Linux via mlxcelverse #1801 (AMD GPU (ROCm) backend on Linux via mlxcelverse), phase 1
• feat(rocm): memory estimation, device info and hardware detection on AMD UMA hosts #1805 and PR feat(rocm): add a GPU vendor concept and AMD device reporting #1883, where this was found
• refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803, which introduced GpuKernelBackend and the paged_attention_decode refusal precedent
• perf(rocm): allocator footprint and ROCm ports of mlxcel fused kernels #1814, the follow-up home for HIP sampler kernel ports
• feat(rocm): port the BitNet bitlinear_matmul kernel to mlxcelverse #1862, the ROCm BitLinear kernel port, as the model for option (a)
• Softmax-free GPU sampling: Gumbel-max categorical sampling kernel #900 (Gumbel-max sampling), Sorting-free top-k/top-p/min-p sampling with dual-pivot rejection kernels #901 (dual-pivot rejection sampling)