Skip to content

Fix GPU concatenate for outputs larger than 2^31 elements - #4758

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
tivris:fix-concatenate-large-offsets
Oct 28, 2026
Merged

zcbenz merged 1 commit into
ml-explore:mainfrom
tivris:fix-concatenate-large-offsets

Conversation

@tivris

@tivris tivris commented Oct 24, 2026

Copy link
Copy Markdown
Contributor

Follow-up to #4744. concatenate_gpu built destination views with copy_shared_buffer(..., out_slice.size(), data_offset), so data_size was the input element count rather than its span in the output. When concatenating along axis 1 or later into an output past 2^31 elements (mx.stack along axis 1 or later uses the same path), the copy picks 32-bit indices and wraps: Metal silently leaves stale memory where late rows should be, and CUDA is expected to fault.

In both Metal and CUDA backends, views are now built using slice_gpu, which computes the true span. The CPU backend is already correct and untouched.

test_concatenate_large_offsets concatenates along axis 1 to produce a (2^27 + 1, 16) uint8 array and checks the last row. It fails on main on GPU and passes with the fix on GPU and CPU. Neighboring ops tests and pre-commit pass. CUDA was compile-checked only, with warnings as errors on CUDA 12.9 and 13.0.

  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: debugging, code and tests

@zcbenz
zcbenz force-pushed the fix-concatenate-large-offsets branch from aa19390 to 71eb60f Compare October 27, 2026 09:23
@zcbenz
zcbenz merged commit df46f58 into ml-explore:main Oct 28, 2026
30 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants