lablup / lablup/mlxcel

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

Open
#1,885 0 comments 0 reactions 0 assignees View on GitHub
status:backlog type:bug
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

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 ? : `. 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;` 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>` 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>`. 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

```bash
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:

```bash
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 #1801 (AMD GPU (ROCm) backend on Linux via mlxcelverse), phase 1
- #1805 and PR #1883, where this was found
- #1803, which introduced `GpuKernelBackend` and the `paged_attention_decode` refusal precedent
- #1814, the follow-up home for HIP sampler kernel ports
- #1862, the ROCm BitLinear kernel port, as the model for option (a)
- #900 (Gumbel-max sampling), #901 (dual-pivot rejection sampling)

Contributor guide

Open the contributing guide

Research direction

Start with the launcher sites in src/lib/mlx-cpp/turbo/sampling.cpp and sampling_rejection.cpp, then read the paged_attention.cpp refusal precedent and the cxx declarations in src/lib/mlxcel-core/src/lib.rs. Trace Rust callers and direct calls in the two sampling kill-switch tests and benches. Done means ROCm returns typed errors without aborting, tests assert refusal or preserve supported behavior, and the listed ROCm/CUDA checks pass.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, rust
Domain
backend, testing-qa
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.