Skip to content

Document where a Materialize node belongs, with the measurements - #225

Merged
cboulay merged 1 commit into
devfrom
docs/mlx-async-eval-placement
Aug 22, 2026
Merged

Document where a Materialize node belongs, with the measurements#225
cboulay merged 1 commit into
devfrom
docs/mlx-async-eval-placement

Conversation

@cboulay

@cboulay cboulay commented Aug 22, 2026

Copy link
Copy Markdown
Member

Stacked on #221 (which is stacked on #220dev). Needs MaterializeMode from #220, so it cannot sit on the main-based MLX stack (#222#224).

Follows up a question that came out of the MLX performance work: if mx.async_eval doesn't block, should every MLX-capable _process call it — so the GPU starts immediately and the executor moves on?

The two premises check out. The conclusion does not.

What was measured

async_eval never 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:

chunk SYNC at tail ASYNC tail every node
30–64 x 256 ch 300–334 139–172 201–223
128 x 1024 ch 594–641 338–342 328–331
1024 x 2048 ch 1961–2009 1602–1650 1502–1593

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):

other host work in between gain from ASYNC at the end
none 1.00x
~0.5 ms 1.30–1.41x
~2 ms 1.63–1.69x
~6 ms 1.39–1.42x

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.py module 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.ASYNC gets 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 for SYNC, 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.py reproduces all six results. Self-contained, so it runs on this branch without the main-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_eval case look falsely neutral: an elif that applied async_eval to the "deferred" arm as well, and a conversion at the process boundary that was skipped when the intervening-work count was zero.

@cboulay
cboulay force-pushed the docs/mlx-async-eval-placement branch from c819bc6 to b16c835 Compare August 22, 2026 15:53
Base automatically changed from refactor/no-implicit-mlx-sync to dev August 22, 2026 15:57
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
cboulay force-pushed the docs/mlx-async-eval-placement branch from b16c835 to d951d34 Compare August 22, 2026 15:57
@cboulay
cboulay merged commit 7cc0ea8 into dev Aug 22, 2026
13 checks passed
@cboulay
cboulay deleted the docs/mlx-async-eval-placement branch August 22, 2026 15:58
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