mlx: update and fix CUDA JIT packaging (#16871)

Bump MLX to the latest selected upstream ref and update the MLX/imagegen
wrappers and tests for the new API behavior.

Fix the CUDA MLX archive so runtime NVRTC kernels work after deployment:
package CUTE/CUTLASS headers, include the CUDA runtime header closure, and
stage a coherent CUDA-toolkit-matched CCCL tree instead of MLX's fetched CCCL
for CUDA payloads. The previous archive could build successfully but crash at
runtime due to missing or incompatible JIT headers.

x/imagegen/mlx/compile.go

@@ -13,6 +13,7 @@ extern void goClosureDestructor(void* payload);
 import "C"
 
 import (
+	"log/slog"
 	"runtime/cgo"
 	"sync"
 	"unsafe"
@@ -21,6 +22,7 @@ import (
 // inClosureCallback is set to true during closure callback execution.
 var (
 	inClosureCallback bool
+	closureScratch    []*Array
 	closureCallbackMu sync.Mutex
 )
 
@@ -31,6 +33,14 @@ func InClosureCallback() bool {
 	return inClosureCallback
 }
 
+func trackClosureArray(a *Array) {
+	closureCallbackMu.Lock()
+	defer closureCallbackMu.Unlock()
+	if inClosureCallback {
+		closureScratch = append(closureScratch, a)
+	}
+}
+
 // CompiledFunc is a compiled MLX function that can be called efficiently.
 // All intermediate arrays during execution stay inside MLX - only inputs
 // and outputs cross the Go boundary.
@@ -131,15 +141,31 @@ func borrowArray(array C.mlx_array) *Array {
 }
 
 //export goClosureCallback
-func goClosureCallback(res *C.mlx_vector_array, input C.mlx_vector_array, payload unsafe.Pointer) C.int {
-	// Set flag to disable AddCleanup during callback
+func goClosureCallback(res *C.mlx_vector_array, input C.mlx_vector_array, payload unsafe.Pointer) (rc C.int) {
+	defer func() {
+		if r := recover(); r != nil {
+			slog.Error("mlx closure callback panicked", "panic", r)
+			rc = 1
+		}
+	}()
+
 	closureCallbackMu.Lock()
 	inClosureCallback = true
+	closureScratch = nil
 	closureCallbackMu.Unlock()
 	defer func() {
 		closureCallbackMu.Lock()
+		scratch := closureScratch
+		closureScratch = nil
 		inClosureCallback = false
 		closureCallbackMu.Unlock()
+
+		for _, a := range scratch {
+			if a != nil && a.Valid() {
+				C.mlx_array_free(a.c)
+				a.c.ctx = nil
+			}
+		}
 	}()
 
 	// Recover the Go function from the handle
@@ -158,13 +184,16 @@ func goClosureCallback(res *C.mlx_vector_array, input C.mlx_vector_array, payloa
 	// Call the Go function
 	outputs := fn(inputs)
 
-	// Build output vector
-	*res = C.mlx_vector_array_new()
-	for _, arr := range outputs {
-		C.mlx_vector_array_append_value(*res, arr.c)
+	var arrPtr *C.mlx_array
+	if len(outputs) > 0 {
+		handles := make([]C.mlx_array, len(outputs))
+		for i, arr := range outputs {
+			handles[i] = arr.c
+		}
+		arrPtr = &handles[0]
 	}
 
-	return 0
+	return C.mlx_vector_array_set_data(res, arrPtr, C.size_t(len(outputs)))
 }
 
 //export goClosureDestructor

x/imagegen/mlx/mlx.go

@@ -149,7 +149,9 @@ var arrayPool = sync.Pool{
 func newArray(array C.mlx_array) *Array {
 	// In compiled closures, MLX manages memory - skip Go tracking
 	if InClosureCallback() {
-		return &Array{c: array}
+		a := &Array{c: array}
+		trackClosureArray(a)
+		return a
 	}
 
 	// Use pooled Array struct for efficiency
@@ -245,6 +247,18 @@ func FreeStruct(v any) {
 	}
 }
 
+// ReleaseAll releases all arrays tracked by this package.
+// Use this when tearing down a whole MLX context; normal model code should
+// prefer FreeStruct or explicit Array.Free calls for owned state.
+func ReleaseAll() {
+	for _, a := range arrays {
+		if a != nil {
+			a.kept = false
+		}
+	}
+	cleanup()
+}
+
 // Keep marks arrays to persist across Eval() cleanup.
 // Kept arrays will NOT be freed when Eval() runs cleanup.
 func Keep(arrays ...*Array) {
@@ -275,6 +289,27 @@ func cleanup() int {
 	return freed
 }
 
+func keepDuringRead(readArrays ...*Array) func() {
+	type state struct {
+		array *Array
+		kept  bool
+	}
+
+	states := make([]state, 0, len(readArrays))
+	for _, a := range readArrays {
+		if a != nil {
+			states = append(states, state{array: a, kept: a.kept})
+			a.kept = true
+		}
+	}
+
+	return func() {
+		for _, s := range states {
+			s.array.kept = s.kept
+		}
+	}
+}
+
 // DebugArrays prints summary info about all tracked arrays.
 func DebugArrays() {
 	var totalBytes int64
@@ -1277,17 +1312,38 @@ func (d Dtype) ItemSize() int64 {
 // Note: Arrays of other dtypes (bf16, f16, etc) are automatically converted to float32.
 // Note: Triggers cleanup of non-kept arrays.
 func (a *Array) Data() []float32 {
-	cleanup()
+	if a == nil || !a.Valid() {
+		return nil
+	}
+	restore := keepDuringRead(a)
+	defer func() {
+		restore()
+		cleanup()
+	}()
+
 	size := a.Size()
 	if size == 0 {
 		return nil
 	}
 
 	arr := a
+	var restoreCast func()
 	if a.Dtype() != DtypeFloat32 {
 		arr = AsType(a, DtypeFloat32)
+		restoreCast = keepDuringRead(arr)
 		arr.Eval()
-		// Cast array will be cleaned up on next Eval
+		defer func() {
+			restoreCast()
+		}()
+	}
+	var restoreContiguous func()
+	if !arr.IsContiguous() {
+		arr = Contiguous(arr)
+		restoreContiguous = keepDuringRead(arr)
+		arr.Eval()
+		defer func() {
+			restoreContiguous()
+		}()
 	}
 
 	ptr := C.mlx_array_data_float32(arr.c)
@@ -1313,7 +1369,15 @@ func (a *Array) Item() float32 {
 // Note: For non-contiguous arrays (e.g., from SliceStride), call Contiguous() first.
 // Note: Triggers cleanup of non-kept arrays.
 func (a *Array) DataInt32() []int32 {
-	cleanup()
+	if a == nil || !a.Valid() {
+		return nil
+	}
+	restore := keepDuringRead(a)
+	defer func() {
+		restore()
+		cleanup()
+	}()
+
 	size := a.Size()
 	if size == 0 {
 		return nil
@@ -1330,7 +1394,15 @@ func (a *Array) DataInt32() []int32 {
 // ItemInt32 gets a single scalar value efficiently (no array copy).
 // Note: Triggers cleanup of non-kept arrays.
 func (a *Array) ItemInt32() int32 {
-	cleanup()
+	if a == nil || !a.Valid() {
+		return 0
+	}
+	restore := keepDuringRead(a)
+	defer func() {
+		restore()
+		cleanup()
+	}()
+
 	var val C.int32_t
 	C.mlx_array_item_int32(&val, a.c)
 	return int32(val)
@@ -1341,7 +1413,15 @@ func (a *Array) ItemInt32() int32 {
 // For non-contiguous arrays, call Contiguous() first.
 // Note: Triggers cleanup of non-kept arrays.
 func (a *Array) Bytes() []byte {
-	cleanup()
+	if a == nil || !a.Valid() {
+		return nil
+	}
+	restore := keepDuringRead(a)
+	defer func() {
+		restore()
+		cleanup()
+	}()
+
 	nbytes := a.Nbytes()
 	if nbytes == 0 {
 		return nil
@@ -1361,7 +1441,9 @@ func (a *Array) Bytes() []byte {
 	default:
 		// For other types (bf16, f16, etc), convert to float32
 		arr := AsType(a, DtypeFloat32)
+		restoreCast := keepDuringRead(arr)
 		arr.Eval()
+		defer restoreCast()
 		ptr = unsafe.Pointer(C.mlx_array_data_float32(arr.c))
 		nbytes = arr.Nbytes()
 	}

x/imagegen/mlx/mlx_test.go

@@ -26,8 +26,20 @@ func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
 
+func useMLXTestThread(t *testing.T) {
+	t.Helper()
+	runtime.LockOSThread()
+	t.Cleanup(func() {
+		ReleaseAll()
+		ClearCache()
+		runtime.UnlockOSThread()
+	})
+}
+
 // TestBasicCleanup verifies non-kept arrays are freed and kept arrays survive.
 func TestBasicCleanup(t *testing.T) {
+	useMLXTestThread(t)
+
 	weight := NewArrayFloat32([]float32{1, 2, 3, 4}, []int32{2, 2})
 	Keep(weight)
 	weight.Eval()
@@ -62,6 +74,8 @@ func TestBasicCleanup(t *testing.T) {
 
 // TestKeptSurvives verifies kept arrays are not freed.
 func TestKeptSurvives(t *testing.T) {
+	useMLXTestThread(t)
+
 	a := NewArrayFloat32([]float32{1, 2}, []int32{2})
 	b := NewArrayFloat32([]float32{3, 4}, []int32{2})
 	result := Add(a, b)
@@ -81,6 +95,8 @@ func TestKeptSurvives(t *testing.T) {
 
 // TestEvalAutoKeeps verifies Eval automatically keeps its outputs.
 func TestEvalAutoKeeps(t *testing.T) {
+	useMLXTestThread(t)
+
 	a := NewArrayFloat32([]float32{1, 2}, []int32{2})
 	b := NewArrayFloat32([]float32{3, 4}, []int32{2})
 	result := Add(a, b)
@@ -110,6 +126,8 @@ func TestEvalAutoKeeps(t *testing.T) {
 
 // TestWeightsSurvive verifies kept arrays survive multiple Eval cycles.
 func TestWeightsSurvive(t *testing.T) {
+	useMLXTestThread(t)
+
 	weight := NewArrayFloat32([]float32{1, 2, 3, 4}, []int32{2, 2})
 	Keep(weight)
 	weight.Eval()
@@ -128,6 +146,8 @@ func TestWeightsSurvive(t *testing.T) {
 
 // TestAsyncEvalCleanup verifies AsyncEval cleans up and dispatches.
 func TestAsyncEvalCleanup(t *testing.T) {
+	useMLXTestThread(t)
+
 	weight := NewArrayFloat32([]float32{1, 0, 0, 1}, []int32{2, 2}) // Identity matrix
 	Keep(weight)
 	weight.Eval()
@@ -164,6 +184,8 @@ func TestAsyncEvalCleanup(t *testing.T) {
 
 // TestMultiOutput verifies multiple kept arrays survive.
 func TestMultiOutput(t *testing.T) {
+	useMLXTestThread(t)
+
 	a := NewArrayFloat32([]float32{1, 2, 3, 4}, []int32{2, 2})
 	sum := Add(a, a)
 	prod := Mul(a, a)
@@ -186,6 +208,8 @@ func TestMultiOutput(t *testing.T) {
 
 // TestChaining verifies output from one step can be used in next.
 func TestChaining(t *testing.T) {
+	useMLXTestThread(t)
+
 	weight := NewArrayFloat32([]float32{1, 0, 0, 1}, []int32{2, 2})
 	Keep(weight)
 	weight.Eval()
@@ -215,6 +239,8 @@ func TestChaining(t *testing.T) {
 
 // TestGenerationLoop simulates the LLM generation pattern with cache.
 func TestGenerationLoop(t *testing.T) {
+	useMLXTestThread(t)
+
 	weight := NewArrayFloat32([]float32{1, 0, 0, 1}, []int32{2, 2})
 	Keep(weight)
 	weight.Eval()
@@ -429,6 +455,8 @@ func gelu(x *Array) *Array {
 
 // TestCompileBasic verifies compiled function produces correct output.
 func TestCompileBasic(t *testing.T) {
+	useMLXTestThread(t)
+
 	x := NewArrayFloat32([]float32{-1, 0, 1, 2}, []int32{4})
 	Keep(x)
 	x.Eval()
@@ -464,6 +492,8 @@ func TestCompileBasic(t *testing.T) {
 
 // TestCompileMultipleInputs verifies compiled function with multiple inputs.
 func TestCompileMultipleInputs(t *testing.T) {
+	useMLXTestThread(t)
+
 	a := NewArrayFloat32([]float32{1, 2, 3, 4}, []int32{4})
 	b := NewArrayFloat32([]float32{5, 6, 7, 8}, []int32{4})
 	Keep(a, b)
@@ -489,6 +519,8 @@ func TestCompileMultipleInputs(t *testing.T) {
 
 // TestCompileReuse verifies compiled function can be called multiple times.
 func TestCompileReuse(t *testing.T) {
+	useMLXTestThread(t)
+
 	compiled := Compile(func(inputs []*Array) []*Array {
 		return []*Array{Add(inputs[0], inputs[0])}
 	})
@@ -551,6 +583,8 @@ func BenchmarkGELUCompiled(b *testing.B) {
 
 // TestCompileNoMemoryLeak verifies compiled functions don't leak memory.
 func TestCompileNoMemoryLeak(t *testing.T) {
+	useMLXTestThread(t)
+
 	x := RandomNormal([]int32{100, 100}, 42)
 	Keep(x)
 	x.Eval()
@@ -598,6 +632,8 @@ func TestCompileNoMemoryLeak(t *testing.T) {
 
 // TestCompileWithRandomState verifies compiled function can capture and update random state.
 func TestCompileWithRandomState(t *testing.T) {
+	useMLXTestThread(t)
+
 	// Simulate logits for sampling
 	logits := NewArrayFloat32([]float32{0.1, 0.2, 0.3, 0.4}, []int32{1, 4})
 	Keep(logits)
@@ -666,6 +702,8 @@ func swiGLU(gate, up *Array, alpha, limit float32) *Array {
 
 // TestCompileSwiGLU verifies compiled SwiGLU produces correct output.
 func TestCompileSwiGLU(t *testing.T) {
+	useMLXTestThread(t)
+
 	gate := NewArrayFloat32([]float32{-1, 0, 1, 2, 5, 10}, []int32{6})
 	up := NewArrayFloat32([]float32{-5, -1, 0, 1, 5, 10}, []int32{6})
 	Keep(gate, up)
@@ -900,6 +938,8 @@ func BenchmarkSampleTopPCompiled(b *testing.B) {
 
 // TestCompiledSamplerMemoryStable verifies compiled samplers don't leak memory.
 func TestCompiledSamplerMemoryStable(t *testing.T) {
+	useMLXTestThread(t)
+
 	vocabSize := int32(32000)
 	logits := RandomNormal([]int32{vocabSize}, 42)
 	key := RandomKey(42)
@@ -1096,6 +1136,8 @@ func BenchmarkCleanupIsolated(b *testing.B) {
 
 // TestMemoryStable verifies that cleanup doesn't cause unbounded memory growth.
 func TestMemoryStable(t *testing.T) {
+	useMLXTestThread(t)
+
 	if testing.Short() {
 		t.Skip("skipping memory test in short mode")
 	}

x/imagegen/nn/nn_test.go

@@ -29,8 +29,20 @@ func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
 
+func useMLXTestThread(t *testing.T) {
+	t.Helper()
+	runtime.LockOSThread()
+	t.Cleanup(func() {
+		mlx.ReleaseAll()
+		mlx.ClearCache()
+		runtime.UnlockOSThread()
+	})
+}
+
 // TestLinearNoBias verifies Linear without bias computes x @ w.T correctly.
 func TestLinearNoBias(t *testing.T) {
+	useMLXTestThread(t)
+
 	// Weight: [out=2, in=3] -> transposed at forward time
 	weight := mlx.NewArrayFloat32([]float32{
 		1, 2, 3, // row 0
@@ -56,6 +68,8 @@ func TestLinearNoBias(t *testing.T) {
 
 // TestLinearWithBias verifies Linear with bias computes x @ w.T + b correctly.
 func TestLinearWithBias(t *testing.T) {
+	useMLXTestThread(t)
+
 	weight := mlx.NewArrayFloat32([]float32{
 		1, 2, 3,
 		4, 5, 6,
@@ -80,6 +94,8 @@ func TestLinearWithBias(t *testing.T) {
 
 // TestLinearBatched verifies Linear works with batched input.
 func TestLinearBatched(t *testing.T) {
+	useMLXTestThread(t)
+
 	weight := mlx.NewArrayFloat32([]float32{
 		1, 0,
 		0, 1,
@@ -111,6 +127,8 @@ func TestLinearBatched(t *testing.T) {
 
 // TestRMSNorm verifies RMSNorm computation.
 func TestRMSNorm(t *testing.T) {
+	useMLXTestThread(t)
+
 	weight := mlx.NewArrayFloat32([]float32{1, 1, 1, 1}, []int32{4})
 	mlx.Eval(weight)
 
@@ -134,6 +152,8 @@ func TestRMSNorm(t *testing.T) {
 
 // TestRMSNormWithScale verifies RMSNorm applies weight scaling.
 func TestRMSNormWithScale(t *testing.T) {
+	useMLXTestThread(t)
+
 	weight := mlx.NewArrayFloat32([]float32{2, 2, 2, 2}, []int32{4})
 	mlx.Eval(weight)
 
@@ -156,6 +176,8 @@ func TestRMSNormWithScale(t *testing.T) {
 
 // TestEmbedding verifies embedding lookup.
 func TestEmbedding(t *testing.T) {
+	useMLXTestThread(t)
+
 	// Embedding table: 4 tokens, dim 3
 	weight := mlx.NewArrayFloat32([]float32{
 		0, 0, 0, // token 0
@@ -185,6 +207,8 @@ func TestEmbedding(t *testing.T) {
 
 // TestRepeatKV verifies K/V repetition for GQA.
 func TestRepeatKV(t *testing.T) {
+	useMLXTestThread(t)
+
 	// [B=1, num_kv_heads=2, S=2, head_dim=2]
 	x := mlx.NewArrayFloat32([]float32{
 		// head 0
@@ -222,6 +246,8 @@ func TestRepeatKV(t *testing.T) {
 
 // TestRepeatKVNoOp verifies RepeatKV with factor 1 returns input unchanged.
 func TestRepeatKVNoOp(t *testing.T) {
+	useMLXTestThread(t)
+
 	x := mlx.NewArrayFloat32([]float32{1, 2, 3, 4}, []int32{1, 1, 2, 2})
 	mlx.Eval(x)
 
@@ -234,6 +260,8 @@ func TestRepeatKVNoOp(t *testing.T) {
 
 // TestApplyCausalMask verifies causal masking.
 func TestApplyCausalMask(t *testing.T) {
+	useMLXTestThread(t)
+
 	// [B=1, heads=1, S=3, S=3] - all ones
 	scores := mlx.Ones(1, 1, 3, 3)
 	mlx.Eval(scores)
@@ -259,6 +287,8 @@ func TestApplyCausalMask(t *testing.T) {
 
 // TestApplyCausalMaskWithOffset verifies causal masking with cache offset.
 func TestApplyCausalMaskWithOffset(t *testing.T) {
+	useMLXTestThread(t)
+
 	// Simulating: cache has 2 tokens, adding 1 new query
 	// scores: [B=1, heads=1, queryLen=1, keyLen=3]
 	scores := mlx.Ones(1, 1, 1, 3)
@@ -276,6 +306,8 @@ func TestApplyCausalMaskWithOffset(t *testing.T) {
 
 // TestApplyCausalMaskWithOffsetZero verifies offset=0 falls back to regular causal.
 func TestApplyCausalMaskWithOffsetZero(t *testing.T) {
+	useMLXTestThread(t)
+
 	scores := mlx.Ones(1, 1, 2, 2)
 	mlx.Eval(scores)