Skip to content

Fix silent Metal OOM error by sticky stream - #4134

Closed
jasonge27 wants to merge 5 commits into
ml-explore:mainfrom
jasonge27:fix/sticky-stream-metal-error
Closed

Fix silent Metal OOM error by sticky stream#4134
jasonge27 wants to merge 5 commits into
ml-explore:mainfrom
jasonge27:fix/sticky-stream-metal-error

Conversation

@jasonge27

@jasonge27 jasonge27 commented Aug 10, 2026

Copy link
Copy Markdown

Proposed changes

Addressing issue #3979, where the MLX drops OOM errors silently in an edge case.

import mlx.core as mx, numpy as np

a = mx.full((3100, 1048576), 3.0)
b = mx.full((3100, 1048576), 4.0)
mx.eval(a, b)
c = a + b
mx.eval(c)
print(float(np.array(c[:1, :1])[0, 0]))   # returns 0.0 (M4 Pro, 48GB)

What goes wrong today:

  1. Mid-eval, MLX hits needs_commit() and submits a buffer with no signal events (eval.cpp around the mid-eval commit).
  2. That buffer fails → error is parked on the encoder.
  3. Nothing is poisoned, because poison only happens for events in signal_events.
  4. A later healthy buffer signals the eval synchronizer → mx.eval returns “success.”
  5. Output was never written → we get 0.0.

Later, get_command_encoder() / error_.reset() discards the error. In the end, GPU failed, user got silence + wrong data.. I'm also linking a cpp gist to reproduce the error (without having to install python library).

Changes:
Handling error in lazy-eval / async systems is hard. CUDA uses sticky last-error and PyTorch surfaces async errors at the next sync. I think following industry best practices, we should park the error per stream rather than per encoder, and report it from synchronization points rather than from events:

  • StreamError keyed by stream index, so it outlives the encode cycle. get_command_encoder() no longer resets it.
  • commit() captures the holder by shared_ptr instead of &error_, which also removes a reference that could outlive the encoder.
  • Event::wait checks the stream after waiting; EventImpl::wait no longer throws, leaving the event error to propagate across streams via wait_events as before.
  • array::wait stops taking theis_available() shortcut, which detached the event as soon as the GPU signaled and skipped the check.
  • The cpu-stream branch of Event::wait(Stream) reports nothing. It runs on a scheduler worker with no handler above it, so throwing there terminated the process rather than reaching the user.

With our changes, now the test script above would throw error [METAL] Command buffer execution failed: Insufficient Memory (00000008:kIOGPUCommandBufferCallbackErrorOutOfMemory).

Testing
I added four cases intotests/gpu_tests.cpp. It injects the error into the completion handler, so that we can reproduce the real one needs more memory than a test machine can assume. Reverting either source change fails the matching test. Full suite passes (267 cases).

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)

cc @zcbenz

…l-explore#3979)

Command-buffer failures could be discarded when get_command_encoder()
reset encoder-local error state after a mid-eval commit with no signal
event. Park the error per stream until synchronize()/Event::wait().
Report only from the stream: EventImpl::wait no longer throws, array::wait
stops taking the shortcut, and the cpu-stream wait leaves the error parked
for the next synchronization point.
@jasonge27
jasonge27 marked this pull request as draft August 10, 2026 08:47
@jasonge27
jasonge27 marked this pull request as ready for review August 10, 2026 08:50
@zcbenz

zcbenz commented Aug 11, 2026

Copy link
Copy Markdown
Member

Thanks for the PR, while I think your solution works, we are in the middle of refactoring the event handling also for the cpu side (#3742) and I do need to have a minimal fix that would not block too much things, so I'm closing this in favor of #4174.

@zcbenz zcbenz closed this Aug 11, 2026
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