fix(metal): use-after-free in gpu::eval for primitives that synchronize mid-eval

[Gajesh2007]

Problem

gpu::eval (mlx/backend/metal/eval.cpp) captures the MTL::CommandBuffer* before eval_gpu runs, then attaches the buffer-liveness completion handler to that pointer afterwards. Any primitive that calls CommandEncoder::synchronize() mid-eval — e.g. the expert-tile route's sortedness-retract check (Gemma E=128 runs this in production today) — commits and replaces the encoder's buffer, leaving the captured pointer dangling. The handler is then attached to a freed object: objc_msgSend/addCompletedHandler on a stale MTL::CommandBuffer*.

Production Gemma has been surviving this by allocator luck. The Qwen E=256 expert-tile shapes (#6) crash deterministically — SIGSEGV in objc_msgSend/addCompletedHandler (crash logs ~/Library/Logs/DiagnosticReports/xctest-2026-08-13-1430*.ips).

Change

Fetch the command buffer from the encoder after eval_gpu. The current buffer holds the tail of the primitive's work, which is exactly what the input buffers must outlive. The debug label call likewise re-fetches instead of using the stale capture.

Behavior

flowchart LR
  subgraph Before
    A1["capture cbuf ptr"] --> B1[eval_gpu] --> C1{"primitive syncs mid-eval?"}
    C1 -->|no| D1["handler on cbuf — OK"]
    C1 -->|yes| E1["encoder committed + replaced buffer<br/>handler on DANGLING ptr → UAF<br/>Gemma: latent · Qwen E=256: deterministic SIGSEGV"]
  end
  subgraph After
    A2[eval_gpu] --> B2["fetch encoder's CURRENT buffer"] --> C2["handler always on live tail buffer<br/>inputs outlive exactly the right work"]
  end

Loading

Code

flowchart LR
  subgraph Before
    F1["eval(): auto* command_buffer = encoder.get_command_buffer()"] --> G1[eval_gpu] --> H1["command_buffer->addCompletedHandler(...)  ← stale after mid-eval commit"]
  end
  subgraph After
    F2["eval(): (no early capture)"] --> G2[eval_gpu] --> H2["encoder.get_command_buffer()->addCompletedHandler(...)  ← always current"]
  end

Loading

Verification

Same diff running on the codex/gemma4-autoresearch-v0.8.2 line (ab5e1dd): Qwen E=256 route went from deterministic SIGSEGV to 10/10 green GPU kernel tests (SortedGatherQuantizedMMTests via mlx-swift) and clean release-server e2e runs; Gemma E=128 suites intact. Cherry-pick onto main is clean (this PR); libmlx compiles on this branch (Release, Metal JIT).

Stack

Merge order (top to bottom):

1. fix(metal): use-after-free in gpu::eval for primitives that synchronize mid-eval #5 — gpu::eval UAF fix (base main) — prod-relevant crash fix, merge first ← this PR
2. perf(metal): E=256 expert-tile route for Qwen 3.5/3.6 MoE prefill + trust mode (unit 1.22-1.32x, e2e +15.2% @8k) #6 — E=256 expert tiles + trust (base codex/gemma4-autoresearch-v0.8.2, the line mlx-swift pins Source/Cmlx/mlx to) · perf(metal): E=256 expert-tile route + trust + gpu::eval UAF fix — darkbloom-base mirror #7 — same commits mirrored onto darkbloom-base (the line d-inference pins libs/mlx to)
3. perf(mlx-swift): Cmlx bump to E=256 expert-tile route (Qwen 3.6) + regenerated sources + Qwen kernel tests mlx-swift#12 — Cmlx bump to 76663df + Qwen kernel tests (base main)
4. perf(cbv2-mtp): Qwen3.6 GDN capture-verify + target-prefix acceptance (1.53x vs serial verify) mlx-swift-lm#106 — GDN capture-verify — merged (squash d06da67)
5. perf(qwen35): fuse routed-expert gate_up + weightedExpertSum — E=256 tile-route consumer mlx-swift-lm#107 — fuseGateUp consumer, +15.2% prefill @8k (base main)
6. perf(qwen36): bump mlx-swift-lm to GDN capture-verify + target-prefix MTP (mlx-swift-lm#106) d-inference#616 — capture-verify pin bump (base master)
7. perf(qwen36): E=256 expert-tile prefill + fused gate_up + trust d-inference#617 — prefill submodule pins + metallib contract + trust env projection (base perf/qwen36-mtp-capture-verify = Fix typo in CMakeLists.txt ml-explore/mlx#616's head)

Chain: mlx UAF fix → mlx tiles (+mirror) → mlx-swift → mlx-swift-lm#106 (merged) → mlx-swift-lm prefill → d-inference#616 → d-inference prefill.