`--gpu-layers N` crashes the server with a Metal GPU timeout on the first request

[solderzzc]

With --gpu-layers set below the model's layer count, the first request aborts the whole server process:

libc++abi: terminating due to uncaught exception of type std::runtime_error: [METAL] Command buffer execution failed: Caused GPU Timeout Error (00000002:kIOGPUCommandBufferCallbackErrorTimeout)

The client only sees the connection drop.

Repro

• SwiftLM main @ 318f712, release build
• Runs on M5 Pro (64 GB, macOS 26.6.2) and on M6
• SwiftLM --model mlx-community/gemma-4-26b-a4b-it-4bit --gpu-layers 23 --port 5415
  • The ready event reports gpu_layers: 23, cpu_layers: 7, total_layers: 30.
• Any request of a few hundred tokens triggers it. With a 490-token prompt, the M5 aborts about 105 s into prefill; on the M6 a 521-token request dies.

Root cause

partitionedLayerCall runs CPU layers under Device.withDefaultDevice(.cpu), which only sets a Swift TaskLocal. The C++ default device stays on the GPU.

MLX's compile keys its cache on the C++ default stream and builds the Compiled primitive for that stream. Gemma 4's MoE calls compiled helpers such as weightedExpertSum and the GELU product. Those helpers are first traced in a GPU layer. When a CPU layer calls them, they hit the same cache entry and run on the GPU.

Their inputs come from GatherQMM on the CPU stream, so the GPU command buffer waits on a CPU fence. That CPU-side MoE matmul for a whole prefill chunk takes far longer than the Metal watchdog allows.

A process sample during the stall shows:

• the MLX CPU StreamThread inside GatherQMM::eval_cpu;
• the Metal command-queue thread stuck submitting.

Fix

The fix is in mlx-swift (SharpAI/mlx-swift#17). CompiledFunction sets the C++ default device to Device.defaultDevice() for the duration of the compiled call and restores it afterwards. Upstream ml-explore/mlx-swift has the same behavior.

With the fix, the same repro no longer crashes. The 490-token request completes with a correct answer ("The provided list consists of forty repetitive placeholder sentences regarding weather and travel.").

Unit tests: 306 executed, 0 failures.

The fix PR adds a regression test, testCompileHonorsScopedDefaultDevice. It traces a compiled function on the GPU, then calls it under Device.withDefaultDevice(.cpu) and asserts that the call retraces. The test fails without the fix and passes with it.

Remaining problem: CPU layers are very slow

After the fix, the CPU layers run single-threaded. GatherQMM / QuantizedMatmul eval_cpu on bf16 run on one core at ~99% CPU. The 490-token prefill took 1,243 s (0.4 tok/s), and decode took about 3 s per token.

For MoE models, --gpu-layers is therefore only useful to avoid OOM, not as a practical speed/memory trade-off. Options:

• warn at startup when CPU layers contain MoE blocks;
• have the README recommend --stream-experts first for MoE models on memory-constrained machines. For comparison, M6 measured Qwen3.6-35B at 13.2 tok/s with --stream-experts.
• look at multithreading the CPU quantized matmul.