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

[inureyes]

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 #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 #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 chore(rocm): mlxcelverse ROCm fork sync script, MLX pin-bump procedure, and upstreaming local fixes #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 #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).