[jonathan308]

Summary

Adds a per-thread override for when Metal command buffers commit. mlx_metal_set_thread_command_buffer_limits(max_ops, max_mb, &prev_ops, &prev_mb) makes command buffers encoded on the calling thread commit after max_ops operations or max_mb megabytes of distinct inputs, instead of the device defaults. Zero keeps a default, and the previous values are returned so the caller can restore them.

Evaluation encodes on the thread that calls eval/async_eval. A caller can therefore give single-token decode steps larger buffers, so the CPU encodes further ahead of the GPU, while prefill keeps the defaults and the memory bound they give. On Ultra parts the defaults are 50 ops / 50 MB, and the MB counter includes weight inputs, so multi-GB MoE expert banks force a commit after almost every expert operation.

Partitioning work into command buffers does not change results. The device defaults and the MLX_MAX_OPS_PER_BUFFER / MLX_MAX_MB_PER_BUFFER environment overrides behave as before when no thread limit is set.

fork.yaml describes the change, and scripts/check_forkdiff.py passes against upstream main.

Why not the environment variables

MLX_MAX_OPS_PER_BUFFER=1000 MLX_MAX_MB_PER_BUFFER=4000 speeds up decode just as much. But the setting is process-wide: temporaries live until their command buffer completes, so an 8K-token prefill of Nemotron 3.5 Super peaked about 75 GB higher.

Layers

Layer	PR
Layr-Labs/mlx	this PR
Layr-Labs/mlx-swift	Layr-Labs/mlx-swift#57 (GPU.withCommandBufferLimits)
Layr-Labs/mlx-swift-lm	Nemotron Super decode scope, to be merged into Layr-Labs/mlx-swift-lm#315
Layr-Labs/d-inference	pins via Layr-Labs/d-inference#1517

Test

Full Nemotron 3.5 Super NVFP4 checkpoint, M3 Ultra (60-core GPU, 256 GB), release build. Greedy decoding, 27-token prompt, 64 tokens. The engine scopes only single-token plain decode steps, at 100 ops with no MB limit. Prefill, mixed steps and speculative rounds keep the defaults.

Commit limit	plain, 1 row	plain, 8 rows	MTP, 1 row	MTP, 8 rows
Defaults (50 ops / 50 MB)	34.2 tok/s	81.9 tok/s	47.3 tok/s	61.3 tok/s
Scoped: 100 ops, no MB limit	43.0 tok/s	91.4 tok/s	47.9 tok/s	61.3 tok/s

• Limits of 200, 300 and 1000 ops gave the same rates as 100.
• The MLX peak during decode is unchanged: 75.8 GB at 1 row, 79.2 GB at 8. The 8K-prefill peak is unchanged at 81.07 GB.
• Outputs:
  • Under every limit, the greedy loop reproduces all 60 golden tokens.
  • Speculative (MTP) tokens equal the plain tokens.
  • All 8 batched rows match their 1-row tokens.
• The SDK debug suite on this MLX build ran 249 tests: 0 failed, 21 skipped.
• The fork's own C++ unit tests were not run.

Merge Danger

Door: two-way. This is an additive C entry point and a thread-local check in Device::get_max_ops_mb_per_buffer. With no thread limit set, behaviour is unchanged.

🤖 Generated with Claude Code

[Gajesh2007]

Merge-readiness review (independent reviewer agents + local validation) before merging:

• Verified thread-locality: CommandEncoder::needs_commit / get_max_ops_mb_per_buffer run only inside gpu::eval on the thread that calls eval/async_eval; every eval commits its buffers before returning, so no buffer mixes two threads' limits. Partitioning does not change results.
• Fixed (f7fb1563, in the squash): the C entry point is declared with MLX_API like the sibling gemma4 entry points (libmlx builds with hidden visibility, so a shared build would not have exported it); the header is included from device.cpp only; the C++ helper moved into the file's anonymous namespace.
• clang-format 21.1.8 clean; compiled through mlx-swift's Cmlx; Swift-side coverage added in GPU: thread-scoped command-buffer commit limits mlx-swift#57 (CommandBufferLimitsTests). Squash 71271f08 has the same tree as the reviewed head.