lablup / lablup/mlxcel

fix(rocm): get_launch_args clamps the grid without a grid-stride contract

Open
#1,874 0 comments 0 reactions 0 assignees View on GitHub
area:core platform:linux priority:medium status:ready type:bug
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

Part of #1801

## Problem / Background

`get_launch_args` in `src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/kernel_utils.hpp:227` computes `num_blocks` and then caps it at `src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/kernel_utils.hpp:236` with `num_blocks = std::min(num_blocks, 65535)`. It returns that capped grid with nothing in its name, signature, or comment saying the caller's kernel must be grid-stride to cover the rest. A kernel written the ordinary way, one index per thread with an early return past the end, silently leaves every element past `65535 * 256` unwritten, and the output buffer keeps whatever the allocation held. That is the failure shape of lablup/mlxcel#1823, where a branch returned without launching and `quantized_matmul` returned the contents of a fresh allocation.

Nothing calls the helper today, which is why this has not bitten. The kernels in this backend write their own grid-stride loops (`binary.hip:21`, `unary.hip:24`). The two kernels added in lablup/mlxcel#1856 compute their own clamped geometry instead of calling it, and both carry a comment saying why: `hadamard.hip:112` and `sort.hip:418-420` ("the launch geometry below caps the grid, so a one-index-per-thread kernel would silently leave the tail of a large output unwritten"). The next person who reaches for the obvious-looking helper gets the bug.

## Current Behavior

The fork's helper is not upstream's. Upstream CUDA's definition at `mlx/backend/cuda/kernel_utils.cu:33-50` does not clamp x at all (`num_blocks.x = cuda::ceil_div(nthreads, block_dim)`; CUDA's x limit is 2^31 - 1), takes a `max_block_dim` parameter, and honours the `large` flag through `get_2d_grid_dims`. The 65535 constant appears there only as `max_grid_yz_dim` in `get_launch_args_general` (`kernel_utils.cu:53-62`), which folds the y overflow into z rather than dropping it. The fork's version ignores `shape`, `strides`, and `large`, fixes `block_size` at 256, and drops the overflow. Upstream's `hadamard.cu:150` does clamp to 65535, but its kernel is grid-stride (`mlx/backend/cuda/device/hadamard.cuh:61`, `transform += gridDim.x`), so clamp and loop are a matched pair. Metal has no equivalent helper.

## Scope

**In scope:** `src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/kernel_utils.hpp` and an entry in `src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md`.

**Out of scope:** changing any existing kernel's launch geometry. `binary.hip`, `unary.hip`, `hadamard.hip`, and `sort.hip` are already correct and must not be touched.

## Proposed Solution

Make the helper safe to use, or remove it. Three options to weigh rather than a prescription, since the file's future depends on how closely the overlay tracks upstream:

1. Delete both overloads. Nothing calls them, and every kernel here writes its own loop. Cheapest, and it cannot regress.
2. Keep them and make the contract impossible to miss: rename to something that names the requirement (for example `get_grid_stride_launch_args`) and add a comment stating that the returned grid is capped and the kernel must be grid-stride.
3. Fold the overflow into y the way `get_launch_args_general` does, so a non-grid-stride kernel is also correct. Closest to upstream, most code.

Record the decision and its reason so it is not relitigated.

## Implementation Notes

- **Reuse**: whichever option is taken, follow the pattern already used by `hadamard.hip:110-115` and `sort.hip:726-728`, where the clamp and the grid-stride loop sit next to each other with a comment tying them together.
- **Constraints**: this file is a fork overlay copied over the ROCm backend, so a rename must not collide with an upstream symbol the fork may later introduce. Metal and CUDA builds never copy `patches-rocm/`.
- **Edge cases**: `size == 0` yields `num_blocks == 0` and a zero-extent launch; whichever option is taken must not make that throw. `work_per_thread > size` already collapses `adjusted_size` to 1.

## Acceptance Criteria

- [ ] No caller can obtain a silently truncated grid from `kernel_utils.hpp`: either the helper is gone, or its name and comment state the grid-stride requirement, or it covers the full range.
- [ ] `grep -rn "get_launch_args" src/lib/mlx-cpp/patches-rocm/` shows no call site that pairs a capped grid with a non-grid-stride kernel.
- [ ] The decision and its reason are recorded as a new numbered entry in `src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md` (entries currently run to 18), ending with the file's standard "Applies to the fork; to be proposed there" marker.
- [ ] The entry is listed as an upstreaming candidate under #1813, since the defect is fork-side and not present in upstream CUDA.
- [ ] The ROCm build still compiles and the gfx1151 correctness matrix is unchanged.

## Verification

```
cargo build --release --features rocm
make verify
```

Then run the gfx1151 correctness gate added in lablup/mlxcel#1826 and confirm no regression. A pass is a clean build plus an unchanged matrix result; there is no behavior change to observe, because the helper has no callers.

## Technical Considerations

The defect is fork-side, so it belongs in the upstreaming list maintained by #1813 rather than in an upstream MLX report. Related: #1823 (same silent-uninitialized-output failure shape), #1856 (the two kernels that deliberately avoided this helper).

Contributor guide

Open the contributing guide

Research direction

Start with src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/kernel_utils.hpp:227-236, then compare the overflow handling in get_launch_args_general and the grid-stride patterns in hadamard.hip:110-115 and sort.hip:726-728. Record the chosen approach in src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md, run cargo build --release --features rocm and make verify, and confirm the gfx1151 correctness matrix is unchanged.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.