Thanks to visit codestin.com
Credit goes to github.com

Skip to content

No copy batch rope#2595

Merged
awni merged 1 commit intomainfrom
no_copy_batch_rope
Sep 15, 2025
Merged

No copy batch rope#2595
awni merged 1 commit intomainfrom
no_copy_batch_rope

Conversation

@awni
Copy link
Member

@awni awni commented Sep 15, 2025

  • Avoid an extra non-contiguous copy in rope when batching
  • Use smaller type for indexing math when possible.

Gives a small win in batch generation:

Pre:

Averages: prompt_tps=7043.840, generation_tps=825.242, peak_memory=4.667

Post

Averages: prompt_tps=7129.770, generation_tps=828.059, peak_memory=4.667

@awni awni requested a review from angeloskath September 15, 2025 18:45
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.

Love it, less kernels overall but more options 🥳


constant bool forward [[function_constant(1)]];
constant bool traditional [[function_constant(2)]];
constant bool hs_transpose [[function_constant(3)]];
Copy link
Member

Choose a reason for hiding this comment

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

❤️

@awni awni merged commit caecbe8 into main Sep 15, 2025
6 checks passed
@awni awni deleted the no_copy_batch_rope branch September 15, 2025 21:23
faisalmemon pushed a commit to faisalmemon/mlx that referenced this pull request Oct 30, 2025
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