Skip to content

[CUDA] Use native events for GPU fence waits - #4401

Open
strayberry wants to merge 1 commit into
ml-explore:mainfrom
strayberry:fix/cuda-fence-native-gpu-event
Open

[CUDA] Use native events for GPU fence waits#4401
strayberry wants to merge 1 commit into
ml-explore:mainfrom
strayberry:fix/cuda-fence-native-gpu-event

Conversation

@strayberry

Copy link
Copy Markdown
Contributor

#4400
Use native CUDA events for GPU-to-GPU dependencies to avoid host polling, while retaining counted AtomicEvent synchronization for CPU dependencies.

Fence simultaneously maintains:

  • producer stream;
  • MLX Event with value fixed to 1;
  • AtomicEvent maintaining an incrementing count.

Synchronization rules:

Producer Consumer Implementation
GPU GPU producer Event::set_value(1), consumer stream asynchronous wait
GPU CPU counted AtomicEvent
CPU GPU counted AtomicEvent
CPU CPU counted AtomicEvent

Fixing the value to 1 makes CUDA Event select native cudaEvent_t, avoiding falling back to counting device-memory events. The consumer GPU stream establishes order via cudaStreamWaitEvent without blocking host.

Constraints

  • CPU side cannot be changed to a value-one event, otherwise it breaks repeated update/wait semantics;

  • GPU→GPU must use the passed-in consumer stream;

  • Cross-device synchronization or dependencies involving CPU continue using AtomicEvent;

  • Modifications are strictly limited to MLX CUDA provider, without altering Metal, MCI format, or model weights;

  • ☑️ I understand it is strictly prohibited to use AI to write PR description

  • AI usage disclosure:
    analysis and review code

Use native CUDA events for GPU-to-GPU dependencies to avoid host polling, while retaining counted AtomicEvent synchronization for CPU dependencies.
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.

1 participant