Skip to content

leak global CommandEncoder -- avoid cuda synchronize on process shutdown - #4480

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
davidkoski:cuda-shutdown
Sep 9, 2026
Merged

leak global CommandEncoder -- avoid cuda synchronize on process shutdown#4480
zcbenz merged 1 commit into
ml-explore:mainfrom
davidkoski:cuda-shutdown

Conversation

@davidkoski

Copy link
Copy Markdown
Member
  • used in mlx-swift
  • global CommandEncoder are still alive after exit even if no work is happening
  • the destructor tries to call cuda api, which causes a fatal error
terminate called after throwing an instance of 'std::runtime_error'
 what(): cudaStreamSynchronize(stream_) failed: driver shutting down

from

frame #11: 0x0000560c45de753f example1`mlx::core::cu::CommandEncoder::synchronize() + 73
frame #12: 0x0000560c45de5c8a example1`mlx::core::cu::CommandEncoder::~CommandEncoder() + 28
frame #13: 0x0000560c45df3ee4 example1`std::pair<int const, mlx::core::cu::CommandEncoder>::~pair() + 32
  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: none used

- used in mlx-swift
- global CommandEncoder are still alive after exit even if no work is happening
- the destructor tries to call cuda api, which causes a fatal error
Comment on lines +620 to +623
// encoders are leaked intentionally as they would synchronize on process
// shutdown
static auto* encoders = new std::unordered_map<int, CommandEncoder>();
return *encoders;

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Per internal discussion we decided leaking these on process exit was ok. I was able to test this on a cuda device with mlx-swift.

Comment on lines 615 to 616
static thread_local std::unordered_map<int, CommandEncoder> encoders;
return encoders;

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I don't know if this has the same issue -- mlx-swift doesn't use the TLS Streams, but since the python code does, perhaps it is ok.

@zcbenz
zcbenz merged commit 24c699e into ml-explore:main Sep 9, 2026
29 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.

2 participants