Skip to content

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

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

citruslab wants to merge 1 commit into
ml-explore:mainfrom
Sogni-AI:fix-concatenate-large-offsets

Conversation

@citruslab

Copy link
Copy Markdown

Proposed changes

concatenate_gpu (Metal and CUDA) builds each destination view with copy_shared_buffer and a data_size equal to the view's element count. The copy dispatch selects 32-bit indexing when data_size fits in INT32_MAX, so a destination whose rows start beyond 2^31 elements is written with wrapped offsets and everything past the boundary is garbage, while the output as a whole can still be well under the 2^32 size limit.

This builds the destination view with slice() instead, so its data_size is the span it reaches, exactly as #4744 did for pad. Same class of issue as the split fix in #4540.

Reproduction (before this change, on an M2 Max): concatenate([q, k, v], axis=-1) with three [105000, 7168] bfloat16 inputs returns correct rows up to 99,864 and wrong rows after that (2^31 / 21504 = 99,864). We hit it in a video transformer's fused Q/K/V projection at long sequence lengths.

The added test mirrors test_pad_large_offsets: row 2^27 of a uint8 (2^27 + 1, 16) output starts at offset 2^31.

Checklist

  • I have read the CONTRIBUTING document
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

concatenate_gpu built each destination view with copy_shared_buffer and a
data_size equal to the view's element count. The copy dispatch picks 32-bit
indexing when data_size fits in INT32_MAX, so a destination whose rows start
beyond 2^31 elements was written with wrapped offsets (garbage past the
boundary). Build the view with slice() so data_size is the reachable span,
exactly as the pad fix does. Regression test mirrors test_pad_large_offsets.

(cherry picked from commit b1f132a)

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