From 0f7ab8407fcc94217808535f3165c2e84ebdf3a4 Mon Sep 17 00:00:00 2001
From: Aegis AI Assistant <simba@aegis-ai.dev>
Date: Thu, 24 Sep 2026 10:49:03 -0700
Subject: [PATCH] fix: run compiled functions on the task-scoped default device

The backend keys its compile cache on, and builds the Compiled primitive
for, the C++ default stream. Device.withDefaultDevice only sets a Swift
TaskLocal, so a function first traced on the GPU was reused inside a
CPU-scoped call and ran its fused kernel on the GPU against CPU-stream
inputs. Set the C++ default device for the duration of the call.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 Source/MLX/Transforms+Compile.swift | 13 +++++++++++++
 Tests/MLXTests/TransformTests.swift | 30 +++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/Source/MLX/Transforms+Compile.swift b/Source/MLX/Transforms+Compile.swift
index 1423f848f..eb8b8030e 100644
--- a/Source/MLX/Transforms+Compile.swift
+++ b/Source/MLX/Transforms+Compile.swift
@@ -94,6 +94,19 @@ final class CompiledFunction: @unchecked (Sendable) {
             evalLock.unlock()
         }
 
+        // The backend keys its compile cache on, and builds the compiled
+        // primitive for, the C++ default stream, which never sees the
+        // task-scoped device from `Device.withDefaultDevice`. Without this, a
+        // function first traced on the GPU is reused inside a CPU-scoped call
+        // and runs its fused kernel on the GPU against CPU-stream inputs.
+        var previousDevice = mlx_device_new()
+        mlx_get_default_device(&previousDevice)
+        mlx_set_default_device(Device.defaultDevice().ctx)
+        defer {
+            mlx_set_default_device(previousDevice)
+            mlx_device_free(previousDevice)
+        }
+
         let innerInputs = arguments + stateInputs
         let innerInputsVector = new_mlx_vector_array(innerInputs)
         defer { mlx_vector_array_free(innerInputsVector) }
diff --git a/Tests/MLXTests/TransformTests.swift b/Tests/MLXTests/TransformTests.swift
index fbfbaac5e..6c74f6def 100644
--- a/Tests/MLXTests/TransformTests.swift
+++ b/Tests/MLXTests/TransformTests.swift
@@ -133,6 +133,36 @@ class TransformTests: XCTestCase {
         assertEqual(r1, r3)
     }
 
+    func testCompileHonorsScopedDefaultDevice() {
+        // A function traced on the GPU must not be reused inside a CPU-scoped
+        // call: the backend caches compiled graphs per default stream, so the
+        // CPU call has to retrace for its own stream.
+        var traces = 0
+        let compiled = compile { (inputs: [MLXArray]) -> [MLXArray] in
+            traces += 1
+            return [square(inputs[0] * inputs[1])]
+        }
+
+        let i1 = MLXRandom.normal([20, 20])
+        let i2 = MLXRandom.normal([20, 20])
+        let expected = square(i1 * i2)
+
+        assertEqual(compiled([i1, i2])[0], expected)
+        XCTAssertEqual(traces, 1)
+
+        let cpuResult = Device.withDefaultDevice(.cpu) {
+            let r = compiled([i1, i2])[0]
+            eval(r)
+            return r
+        }
+        assertEqual(cpuResult, expected)
+        XCTAssertEqual(traces, 2)
+
+        // Back on the GPU, the original entry is reused.
+        assertEqual(compiled([i1, i2])[0], expected)
+        XCTAssertEqual(traces, 2)
+    }
+
     class CompileTestState: CustomStringConvertible, Updatable {
         var y: MLXArray
         var o: MLXArray?