Skip to content

Validate dims in rope#3230

Merged
angeloskath merged 1 commit intoml-explore:mainfrom
MillaFleurs:pr-3218-clean
Mar 10, 2026
Merged

Validate dims in rope#3230
angeloskath merged 1 commit intoml-explore:mainfrom
MillaFleurs:pr-3218-clean

Conversation

@MillaFleurs
Copy link
Contributor

Proposed changes

Updates for #3218 . RoPE missing dims validation can trigger invalid GPU launch / out-of-bounds indexing.

Bug summary:

fast::rope accepts arbitrary dims and forwards it into RoPE primitive state.
There is no guard for invalid values such as dims <= 0, dims > feature_dim, or odd dims.
GPU kernels derive launch/indexing from dims_/2, which can become unsafe or semantically incorrect.
Primary locations:

mlx/fast.cpp:367-415, 417-512, 533-540
No checks enforce: dims > 0, dims <= x.shape(-1), dims even

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

Copy link
Member

@angeloskath angeloskath left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@angeloskath
Copy link
Member

Could you run the linter? Then we can merge after the tests pass.

@MillaFleurs
Copy link
Contributor Author

Same comment here I thought the pre-commit was running. give me a minute...

@zcbenz zcbenz changed the title PR 3218 Validate dims in rope Mar 10, 2026
@angeloskath angeloskath merged commit 572e0a4 into ml-explore:main Mar 10, 2026
16 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.

3 participants