Skip to content

[CUDA] Wait for tasks in cuda#2636

Merged
awni merged 1 commit into
ml-explore:mainfrom
awni:cuda_wait_tasks
Sep 30, 2025
Merged

[CUDA] Wait for tasks in cuda#2636
awni merged 1 commit into
ml-explore:mainfrom
awni:cuda_wait_tasks

Conversation

@awni

@awni awni commented Sep 30, 2025

Copy link
Copy Markdown
Member

This helps keep the work submission from running to far ahead of complete work for the CUDA back-end. I'ts virtually identical to what we do for Metal.

Some benchmarks:

Model Pre It/sec Post it/sec Pre Peak Mem GB Post Peak Mem GB
0.6 B 6.45 6.44 54.3 28.9
0.86 B 3.59 5.82 64.79 36.47

Generation speed is unaffected.

@awni
awni requested review from angeloskath and zcbenz September 30, 2025 20:14
@awni

awni commented Sep 30, 2025

Copy link
Copy Markdown
Member Author

Thanks @nastya236 for diagnosing the issue here.

@angeloskath angeloskath left a comment

Copy link
Copy Markdown
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 bbf1423 into ml-explore:main Sep 30, 2025
7 checks passed
@awni
awni deleted the cuda_wait_tasks branch October 13, 2025 14:55
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