From 1f3f11cae6cb22fbfb2cefcb3d89f6e524f6f5d7 Mon Sep 17 00:00:00 2001
From: Gajesh Naik <26431906+Gajesh2007@users.noreply.github.com>
Date: Thu, 13 Aug 2026 14:38:07 -0700
Subject: [PATCH] fix(metal): use-after-free in gpu::eval for primitives that
 synchronize mid-eval
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

gpu::eval captured the MTL::CommandBuffer* BEFORE eval_gpu, but a primitive
that calls CommandEncoder::synchronize() inside eval_gpu (the expert-tile
route's sortedness-retract check) commits and REPLACES the encoder's buffer;
eval then attached its buffer-liveness completion handler to the dangling
pointer (objc_msgSend on freed object). The Gemma E=128 route has been
surviving this by allocator luck; the Qwen E=256 shapes crash deterministically.
Fetch the buffer from the encoder AFTER eval_gpu — the current buffer holds
the tail of the primitive's work, which is exactly what the inputs must
outlive.
---
 mlx/backend/metal/eval.cpp | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/mlx/backend/metal/eval.cpp b/mlx/backend/metal/eval.cpp
index e6826253de..fcc572b476 100644
--- a/mlx/backend/metal/eval.cpp
+++ b/mlx/backend/metal/eval.cpp
@@ -30,7 +30,6 @@ void eval(array& arr) {
   auto pool = metal::new_scoped_memory_pool();
   auto s = arr.primitive().stream();
   auto& encoder = metal::get_command_encoder(s);
-  auto* command_buffer = encoder.get_command_buffer();
 
   auto outputs = arr.outputs();
   {
@@ -41,7 +40,8 @@ void eval(array& arr) {
       inputs = arr.inputs();
     }
 
-    debug_set_primitive_buffer_label(command_buffer, arr.primitive());
+    debug_set_primitive_buffer_label(
+        encoder.get_command_buffer(), arr.primitive());
     arr.primitive().eval_gpu(arr.inputs(), outputs);
   }
   std::unordered_set<std::shared_ptr<array::Data>> buffers;
@@ -63,7 +63,13 @@ void eval(array& arr) {
       scheduler::notify_task_completion(s);
     });
   } else {
-    command_buffer->addCompletedHandler(
+    // Fetch the command buffer AFTER eval_gpu: primitives that synchronize
+    // mid-eval (e.g. the expert-tile route's descriptor retract check)
+    // commit and REPLACE the encoder's buffer, so a pointer captured before
+    // eval_gpu would be stale here — attaching the buffer-liveness handler
+    // to it is a use-after-free. The current buffer holds the tail of this
+    // primitive's work, which is exactly what the inputs must outlive.
+    encoder.get_command_buffer()->addCompletedHandler(
         [buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {});
   }
 }