Document where a Materialize node belongs, with the measurements - #225
Merged
Conversation
cboulay
force-pushed
the
docs/mlx-async-eval-placement
branch
from
August 22, 2026 15:53
c819bc6 to
b16c835
Compare
mx.async_eval never blocks -- not on a dependency still in flight, and not when the queue is already deep -- which invites the conclusion that an evaluation point can go anywhere, or after every MLX node. It cannot. Each call costs ~17-21 us of host time to submit a command buffer, against ~1 us to build the op graph it replaces, and in a live pipeline that host time belongs to the executor rather than the device. Across a three-node MLX segment on an M4 Pro, evaluating per node is 1.3-1.7x worse than one evaluation at the tail at streaming chunk sizes, and only reaches break-even at chunks far larger than a real-time graph carries. Fan-out is not a special case: evaluating each branch as it finishes measured 1.3x worse than evaluating once at the join. What does pay is placing that single evaluation point at the end of the MLX segment, before whatever else the executor has to do -- downstream units, or the NumPy conversion at a process boundary that forces completion anyway. That is worth 1.3-1.7x, scaling with how much other host work it overlaps, and exactly nothing when there is none. Also puts a size on the run-ahead warning ASYNC already carried: 2083 MiB of in-flight intermediates against 33 MiB for SYNC, for 1.6x the throughput, replaying 400 messages of 512x1024 as fast as the host could build them. The new benchmark reproduces all six results.
cboulay
force-pushed
the
docs/mlx-async-eval-placement
branch
from
August 22, 2026 15:57
b16c835 to
d951d34
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follows up a question that came out of the MLX performance work: if
mx.async_evaldoesn't block, should every MLX-capable_processcall it — so the GPU starts immediately and the executor moves on?The two premises check out. The conclusion does not.
What was measured
async_evalnever blocks on an in-flight dependency. Four chained heavy stages, each consuming the previous one's output: host returned at 2.4 ms against 71 ms of device work.There is no back-pressure. 24 independent chunks enqueued in 5.7 ms against 201 ms of device work, per-chunk host cost flat at 0.09–0.36 ms. Work queues and runs back-to-back with no host involvement — so yes, chunk 2 starts its GPU work without waiting for chunk 1 to reach the end of the pipeline.
But each evaluation point costs ~17–21 µs of host time to submit a command buffer, against ~1 µs to build the op graph it replaces. In a live graph that is the executor's time, not the device's.
Three-node MLX segment (Metal SOS filter, scaler, abs), µs/message:
Per-node is 1.3–1.7x worse at streaming chunk sizes, and only reaches break-even at chunks far above what a real-time graph carries. Fan-out is not a special case either — evaluating each branch as it finishes measured 1.3x worse than evaluating once at the join.
What does pay is putting the single evaluation point at the end of the MLX segment, before whatever else the executor has to do — downstream units, or the NumPy conversion at a process boundary (shmem can't carry MLX arrays, so that conversion forces completion regardless):
The gain is entirely overlap, so it is worth exactly as much as the host has to get on with — which is the argument for placing the node earlier than the point that would force evaluation anyway.
Changes
materialize.pymodule docstring gains a Where to place it section carrying the above. This is where someone would look before making the per-node change, so the negative result belongs here rather than only in a benchmark.MaterializeMode.ASYNCgets a size on the run-ahead warning it already carried: replaying 400 messages of 512x1024 float32 as fast as the host could build them holds 2083 MiB in flight against 33 MiB forSYNC, for 1.6x the throughput. Fine for a live source that paces itself; a hazard for file replay with no downstream backpressure.benchmarks/benchmark_mlx_async_eval_placement.pyreproduces all six results. Self-contained, so it runs on this branch without themain-based stack.Notes
Docs-and-benchmark only — no behavior change, no source logic touched. 4071 tests pass. Both new RST tables validated through docutils.
Two measurement bugs were caught and fixed while writing the benchmark, both of which had made the early-
async_evalcase look falsely neutral: anelifthat appliedasync_evalto the "deferred" arm as well, and a conversion at the process boundary that was skipped when the intervening-work count was zero.