Skip to content

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

Open
tivris wants to merge 1 commit into
ml-explore:mainfrom
tivris:fix-concatenate-large-offsets
Open

tivris wants to merge 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

This branch has not been deployed

No deployments
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.

1 participant