Faster synchronization Fence primitive - #1773
Conversation
move event fixes update bench fix fix
| if (stream.encoder == nullptr) { | ||
| stream.encoder = std::make_unique<CommandEncoder>(stream.buffer); | ||
| stream.fence = std::make_shared<Fence>(device_->newFence()); | ||
| stream.encoder->wait_for_fence(stream.event_fence); |
There was a problem hiding this comment.
I am kinda confused by this wait here. In the case where we don't use the Fence at all when would the event_fence be updated?
There was a problem hiding this comment.
I'm trying an alternative that doesn't require this fence cause I don't like it. But basically you can always wait on a fence. The wait will wait for any preceding calls to update. So if you never update it the wait it is essentially a no-op.
This fence is used to ensure no kernels start before the GPU signal kernel is done. So we update this fence when we signal from the GPU and then any command encoder that waits after that will wait for that update to finish.
There was a problem hiding this comment.
Ok, I pushed a change to get rid of this that should (and seems) to work, so you can disregard all the previous stuff with event_fence.
Basically it requires modifying the call to wait_gpu to take an array that we want to be sure is ready before any future kernels that depend on it run. It reuses our existing synchronization machinery (barriers + fences) and is nice in that it only encoders which actually depend on the output will wait for it.
I had to add a way to register_output_array since it's not actually part of a kernel.. but I think it's cleaner/ more efficient / and doesn't require this random stream_event which was very icky.
There was a problem hiding this comment.
So if you never update it the wait it is essentially a no-op.
What a horrible hidden API. I assumed wait without update is a deadlock hence the confusion before.
I had to add a way to
register_output_array.
Yeah, so much better! Also only waiting on things that matter rather than everything on the whole stream. Plus avoid waiting on two fences which could have been the case before.
|
|
||
| #include <metal_atomic> | ||
|
|
||
| [[kernel]] void input_coherent( |
There was a problem hiding this comment.
What does the input_coherent do? Does it mean to ensure the data gets fully transferred from cpu to gpu?
There was a problem hiding this comment.
It's to ensure the data is safe to be read by the CPU.
There was a problem hiding this comment.
Is this a metal specific thing?
I have a weird flaky test in cuda backend that an array loaded from cpu sometimes ends up with wrong value in gpu:
Line 36 in 9daa6b0
I wonder if I need to do something similar here, though it is completely different hardware.
There was a problem hiding this comment.
It's very metal specific. In fact.. most of the time this kerne is not used. It's only enabled for an experimental polling based synchronization when you set MLX_FAST_METAL_SYNCH=1.
If you are using standard cuda synchronization primitives you should need anything like that.
There was a problem hiding this comment.
Cuda does not have a standard synchronization primitive for signaling a gpu wait from cpu, I was using cuda::atomic but its wait/notify APIs had some problems with CUDA 11 (#2137) so I'm currently using busy wait for fence which is almost identical to metal backend's fast fence.
There was a problem hiding this comment.
so I'm currently using busy wait for fence which is almost identical to metal backend's fast fence.
Oh right. Well in that case.. we are in pretty uncharted territory. Do you know of examples that have used a similar concept in CUDA? That would be a good place to look and see if they have to do anything fancy to ensure the input safe to read.
Metal we need this because even if the fence value is written by the GPU kernel (which certainly implies and the input is done being computed) it can still be in cache and not necessarily "coherent" with the rest of the device (and system in unified memory).
There was a problem hiding this comment.
No I couldn't find any open source code synchronizes gpu with cpu work, and there is actually no framework using cuda's unified memory APIs. I also searched nvidia forum and there was probably 2 posts related to this topic in past 10 years.
The CPU->GPU release half of a tensor-parallel gate blocks the command processor on an MTLSharedEvent. Resuming from that is the single largest term in a TP decode token: DS4_TP_GATE_PROFILE on a 2x M2 Ultra pair reports a steady-state gpu-wait of 470us and an exchange of 38us, and a Metal System Trace puts the GPU's own compute at 284us, so the resume is ~186us of a 508us gate. Across 86 gates that is 16ms of a 43.6ms token -- 37% of decode spent stopping and restarting the GPU. The peer is not the problem: tp_rdma_gate_exchange already blocks on recv_done, and the 38us it costs contains the entire rendezvous. Add an opt-in release fence, after MLX (ml-explore/mlx#1773): the GPU spins on a word the service thread writes instead of blocking, so the command processor never stops. A harness mirroring the gate measures the release at 6.0us against 101us for encodeWaitForEvent. Both halves of the handshake have to change together. The existing arrival publish, kernel_dsv4_tp_flag_set, does a relaxed store to an ordinary device atomic, which is sufficient only because the command processor blocks on an event immediately afterwards -- that stall is what makes the write land. Under the fence the GPU never stalls, so arrival needs system scope of its own, or the service thread spins on a stale word for a full coherency window every gate. Hence kernel_dsv4_tp_flag_set_coherent, selected only on the fast path. Off by default; DS4_METAL_FAST_SYNC=1 enables it, and it falls back to the shared event if the buffer or pipeline is unavailable. Row gates only -- batch and big gates keep the event path, and both sides key the choice off the same conditions, fixed at init. Three properties are load-bearing and were established by A/B: coherent(system) is mandatory. With a plain volatile device pointer the GPU never observes the CPU store (71 of 120 releases missed); with coherent(device) it misses roughly one in a few hundred, which at 86 gates per token is a hang every few tokens. Reaching that qualifier needs #pragma METAL internals, scoped here to just these two kernels. One threadgroup of one thread. Against a saturated GPU, one threadgroup costs 4%, two cost 21%, four cost 61%; threads within a threadgroup are free. The shape is deliberately not configurable. No backoff, no chunking. An fma backoff between polls takes the release from 6.0us to 13.1us (256 iters) and 33.0us (2048); splitting the wait into 400 bounded dispatches costs 1265us, because Metal cannot cancel dispatches already encoded. The spin is iteration-bounded (DS4_TP_FENCE_MAX_ITERS) so a peer dying mid-gate cannot wedge the command processor into a GPU watchdog kill; the service thread already writes the release on the failure path, so the bound is a backstop.
The GPU waits on a word the host stores, so a producer/consumer handoff does not need the host to wait for a command buffer and then submit the next one. Needs coherent(system), which is undocumented, so the kernels live in their own runtime-compiled library and init returns NULL when it is unavailable. After MLX ml-explore/mlx#1773. Not used yet; exposed by name for backend-agnostic callers. Assisted-by: Claude Opus
3 participants
Adds a
Fencesynchronization primitive for polling-based fast CPU / GPU synchronizationConditionally builds the above for Metal 3.2 and up.
Adds an environment variable
MLX_METAL_FAST_SYNCHto toggle using the fast primitive or not (default0).Updates distributed Metal ops to use the
Fenceso we have the option to make them fast.Distributed LLM tested on M1 Max, M2 Ultra, M3 Max and working.
Benchmarks suggests 10x reduction in synchronization latency for all-reduce (15-20 microseconds depending on the machine, compared to 150-200 with shared events).
Microbenchmarks on M3 Max:
All Reduce: time per iteration 0.149781 (ms)
All gather: time per iteration 0.136384 (ms)
All Reduce: time per iteration 0.016508 (ms)
All gather: time per iteration 0.015927 (ms)