Skip to content

fix(rocm): let strided copies run past 2^32 elements - #2237

Merged
inureyes merged 2 commits into
mainfrom
fix/issue-2184-strided-copy-grid
Oct 8, 2026
Merged

inureyes merged 2 commits into
mainfrom
fix/issue-2184-strided-copy-grid

Conversation

@inureyes

@inureyes inureyes commented Oct 8, 2026

Copy link
Copy Markdown
Member

Summary

  • The ROCm general copy kernels (copy_g_byval, copy_gg_byval, copy_gg_dynamic_nd, copy_gg_dynamic) now cap the grid at 65535 blocks (rocm::copy_grid_blocks, 64-bit block count) and loop past the cap, so strided copies of 2^32 elements or more no longer fail with hipErrorInvalidConfiguration.
  • copy_general_dynamic sizes its launch by the copied shape instead of out.size(). DynamicSliceUpdate passes the whole destination as out, so a one-block update into a 2^32-element slab used to launch 2^32 threads.
  • Each kernel takes a kLoop template argument. Below the cap the instantiation is the old one-thread-per-element body; past it the threads stride on a 64-bit counter. A single kernel holding both paths, and a loop for every size, each measured 1.2 to 1.7% slower on a 4M-element transpose, so the split is what keeps typical copies unchanged. It doubles those four kernels' instantiations.
  • copy_contiguous uses the same helper and switches to 64-bit indices above INT32_MAX: its uint32_t counter adds up to 2^26 per pass and would wrap just under 2^32. This was found by reading the code and was not reproduced.
  • Recorded as LOCAL_FIXES.md item 42. It stays in mlxcelverse under the fork policy.

Other ROCm launchers with the same uncapped grid (binary, unary/ternary 2-D, reduce init/col, arange, gather/scatter, quantize, sort iota, copy_col_row) are not on the copy path. They are filed as #2234.

Verification (gfx1151)

  • tests/rocm_strided_copy_grid.rs, fast tests: 42.6M elements, about 2.5 passes of the capped grid, through contiguous, concatenate and slice_update_dynamic, compared with a host reference. All 3 pass. With the loop cut to one pass, all 3 fail.
  • --ignored tests (8 GiB each, 2^32 + 32,768 f16 elements, rows checked on both sides of 2^31 and 2^32 and at the end): all 4 failed on main with "invalid configuration argument" and pass with the fix. Each ran inside scripts/rocm_gpu_guard.sh. AMD_LOG_LEVEL=3 confirms each op reaches the kernel it targets.
  • Microbench under the guard, before vs after, six interleaved rounds, median of per-round medians, f16 contiguous of a transposed view:
    • [1, 8, 4096, 128]: 147.82 vs 147.99 us
    • [1, 2048, 32, 128]: 388.85 vs 387.40 us
    • concatenate of two [1, 32, 2048, 128] views: 836.75 vs 833.55 us
    • [1, 1, 32, 128] (12 to 14 us) and [1, 8192, 32, 128] (past the cap; best run 5455 vs 5423 us) vary more between rounds than between builds.
  • make verify-rocm on e5f4419 (this branch rebased on 2aa5211): OK; 163 test suites, 12194 passed, 0 failed, 403 ignored, smoke included.
  • Fast gates pass: verify-versions, verify-kernel-dtype-keys, verify-kernel-port-dispatch, verify-llama-compat, verify-fmt, verify-rocm-overlay. Clippy on the new test target passes.

Not verified: Metal and CUDA are not available on this host. The change touches only patches-rocm/ and a test that is gated on feature = "rocm", so the Metal and CUDA paths are not built from it.

Closes #2184

The ROCm general copy kernels (copy_g_byval, copy_gg_byval, copy_gg_dynamic_nd, copy_gg_dynamic) launched one thread per element on an uncapped int grid, so HIP rejected any strided copy of 2^32 elements or more with hipErrorInvalidConfiguration. copy_general_dynamic also sized its launch by the whole destination, so a one-block DynamicSliceUpdate into a slab that large failed the same way.

The grid is now capped at 65535 blocks (rocm::copy_grid_blocks, computed in 64 bits), the dynamic launch covers the copied shape, and each kernel takes a kLoop template argument: below the cap it is the old body, past it a grid-stride loop on a 64-bit counter. One kernel with both paths, or a loop for every size, measured 1.2 to 1.7% slower on a 4M-element transpose, so the split keeps typical copies unchanged within noise. copy_contiguous uses the helper and moves to 64-bit indices above INT32_MAX, so its uint32 counter cannot wrap.

tests/rocm_strided_copy_grid.rs adds fast tests past the cap and four ignored 8 GiB tests that failed before this change and pass after it on gfx1151. Recorded as LOCAL_FIXES.md item 42.

Closes #2184
@inureyes inureyes added status:done Completed type:bug Bug fixes, error corrections, or issue resolutions priority:low Low priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers labels Oct 8, 2026
@inureyes
inureyes merged commit a18d3d7 into main Oct 8, 2026
27 checks passed
@inureyes
inureyes deleted the fix/issue-2184-strided-copy-grid branch October 8, 2026 14:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers priority:low Low priority status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(rocm): strided copy kernels refuse arrays of 2^32 elements or more

1 participant