diff --git a/MLX_VERSION b/MLX_VERSION index 806a0f6c..64bc1aa8 100644 --- a/MLX_VERSION +++ b/MLX_VERSION @@ -1 +1 @@ -2165dc08d7b33258260aa849d39f087d50e62962 +51b2768da7e1897d3c4258f7ddbb47083d1eef01 diff --git a/cmake/mlx/CMakeLists.txt b/cmake/mlx/CMakeLists.txt index a12a2151..123845f4 100644 --- a/cmake/mlx/CMakeLists.txt +++ b/cmake/mlx/CMakeLists.txt @@ -102,7 +102,8 @@ install(RUNTIME_DEPENDENCY_SET mlx_runtime_deps LIBRARY DESTINATION ${OLLAMA_INSTALL_DIR} COMPONENT MLX_VENDOR ) -if(TARGET jaccl) +get_target_property(_MLX_LINK_LIBRARIES mlx LINK_LIBRARIES) +if(TARGET jaccl AND "jaccl" IN_LIST _MLX_LINK_LIBRARIES) install(TARGETS jaccl RUNTIME DESTINATION ${OLLAMA_INSTALL_DIR} COMPONENT MLX LIBRARY DESTINATION ${OLLAMA_INSTALL_DIR} COMPONENT MLX @@ -123,29 +124,53 @@ endif() # --component MLX. Headers are installed alongside libmlx in OLLAMA_INSTALL_DIR. # # Layout: -# ${OLLAMA_INSTALL_DIR}/include/cccl/{cuda,nv}/ - CCCL headers -# ${OLLAMA_INSTALL_DIR}/include/*.h - CUDA toolkit headers +# ${OLLAMA_INSTALL_DIR}/include/cccl/ - CCCL headers +# ${OLLAMA_INSTALL_DIR}/include/{cute,cutlass}/ - CUTLASS/CUTE headers +# ${OLLAMA_INSTALL_DIR}/include/ - CUDA runtime/core headers # -# MLX's jit_module.cpp resolves CCCL via -# current_binary_dir()[.parent_path()] / "include" / "cccl" -# On Linux, MLX's jit_module.cpp resolves CCCL via -# current_binary_dir().parent_path() / "include" / "cccl", so we create a -# symlink from lib/ollama/include -> ${OLLAMA_RUNNER_DIR}/include. +# MLX's jit_module.cpp resolves JIT support headers from the backend-local +# include directory. On Linux it also probes current_binary_dir().parent_path() +# / "include", so we create a symlink from lib/ollama/include to the backend +# include directory for archive packaging. # This will need refinement if we add multiple CUDA versions for MLX in the future. # CUDA runtime headers are found via CUDA_PATH env var (set by mlxrunner). -if(EXISTS ${CMAKE_BINARY_DIR}/_deps/cccl-src/include/cuda) - install(DIRECTORY ${CMAKE_BINARY_DIR}/_deps/cccl-src/include/cuda - DESTINATION ${OLLAMA_INSTALL_DIR}/include/cccl +set(_mlx_jit_cccl_include_dir "") +if(CUDAToolkit_FOUND) + foreach(_dir ${CUDAToolkit_INCLUDE_DIRS}) + if(EXISTS "${_dir}/cccl/cuda/std") + set(_mlx_jit_cccl_include_dir "${_dir}/cccl") + break() + endif() + endforeach() +endif() +if(NOT _mlx_jit_cccl_include_dir AND EXISTS ${CMAKE_BINARY_DIR}/_deps/cccl-src/include/cuda) + set(_mlx_jit_cccl_include_dir "${CMAKE_BINARY_DIR}/_deps/cccl-src/include") +endif() +if(_mlx_jit_cccl_include_dir) + foreach(_cccl_dir cuda nv cub thrust) + if(EXISTS "${_mlx_jit_cccl_include_dir}/${_cccl_dir}") + install(DIRECTORY "${_mlx_jit_cccl_include_dir}/${_cccl_dir}" + DESTINATION ${OLLAMA_INSTALL_DIR}/include/cccl + COMPONENT MLX) + endif() + endforeach() +endif() +if(EXISTS ${CMAKE_BINARY_DIR}/_deps/cutlass-src/include/cute) + install(DIRECTORY ${CMAKE_BINARY_DIR}/_deps/cutlass-src/include/cute + DESTINATION ${OLLAMA_INSTALL_DIR}/include COMPONENT MLX) - install(DIRECTORY ${CMAKE_BINARY_DIR}/_deps/cccl-src/include/nv - DESTINATION ${OLLAMA_INSTALL_DIR}/include/cccl + install(DIRECTORY ${CMAKE_BINARY_DIR}/_deps/cutlass-src/include/cutlass + DESTINATION ${OLLAMA_INSTALL_DIR}/include COMPONENT MLX) endif() -# Install minimal CUDA toolkit headers needed by MLX JIT kernels. -# These are the transitive closure of includes from mlx/backend/cuda/device/*.cuh. +# Install CUDA runtime/core headers needed by MLX JIT kernels. +# NVIDIA's NVRTC bundled-header model is CUDA Runtime + CCCL, not the entire +# toolkit include tree. Keep CCCL coherent above, include CUTLASS/CUTE above, +# and avoid shipping unrelated SDK headers such as NPP, CUPTI, cuRAND, NVML, +# cuBLAS, cuSPARSE, and cuSOLVER. # The Go mlxrunner sets CUDA_PATH to OLLAMA_INSTALL_DIR so MLX finds them at -# $CUDA_PATH/include/*.h via NVRTC --include-path. +# $CUDA_PATH/include via NVRTC --include-path. if(CUDAToolkit_FOUND) # CUDAToolkit_INCLUDE_DIRS may be a semicolon-separated list # (e.g. ".../include;.../include/cccl"). Find the entry that @@ -161,39 +186,97 @@ if(CUDAToolkit_FOUND) message(WARNING "Could not find cuda_runtime_api.h in CUDAToolkit_INCLUDE_DIRS: ${CUDAToolkit_INCLUDE_DIRS}") else() set(_dst "${OLLAMA_INSTALL_DIR}/include") - set(_MLX_JIT_CUDA_HEADERS + + set(_mlx_jit_cuda_headers builtin_types.h + channel_descriptor.h + common_functions.h cooperative_groups.h + cuComplex.h + cuda.h + cudaTypedefs.h + cuda_awbarrier.h + cuda_awbarrier_helpers.h + cuda_awbarrier_primitives.h cuda_bf16.h cuda_bf16.hpp cuda_device_runtime_api.h - cuda_fp16.h - cuda_fp16.hpp + cuda_fp4.h + cuda_fp4.hpp + cuda_fp6.h + cuda_fp6.hpp cuda_fp8.h cuda_fp8.hpp + cuda_fp16.h + cuda_fp16.hpp + cuda_occupancy.h + cuda_pipeline.h + cuda_pipeline_helpers.h + cuda_pipeline_primitives.h + cuda_runtime.h cuda_runtime_api.h + cuda_stdint.h + cudart_platform.h + device_atomic_functions.h + device_atomic_functions.hpp + device_double_functions.h + device_functions.h + device_launch_parameters.h device_types.h + driver_functions.h driver_types.h + fatbinary_section.h + host_config.h + host_defines.h + library_types.h math_constants.h + math_functions.h + mma.h + nvrtc_device_runtime.h + sm_20_atomic_functions.h + sm_20_atomic_functions.hpp + sm_20_intrinsics.h + sm_20_intrinsics.hpp + sm_30_intrinsics.h + sm_30_intrinsics.hpp + sm_32_atomic_functions.h + sm_32_atomic_functions.hpp + sm_32_intrinsics.h + sm_32_intrinsics.hpp + sm_35_atomic_functions.h + sm_35_intrinsics.h + sm_60_atomic_functions.h + sm_60_atomic_functions.hpp + sm_61_intrinsics.h + sm_61_intrinsics.hpp + surface_indirect_functions.h surface_types.h + target + texture_indirect_functions.h texture_types.h vector_functions.h vector_functions.hpp - vector_types.h - ) - foreach(_hdr ${_MLX_JIT_CUDA_HEADERS}) - install(FILES "${_cuda_inc}/${_hdr}" + vector_types.h) + set(_mlx_jit_cuda_header_paths "") + foreach(_header IN LISTS _mlx_jit_cuda_headers) + if(EXISTS "${_cuda_inc}/${_header}") + list(APPEND _mlx_jit_cuda_header_paths "${_cuda_inc}/${_header}") + endif() + endforeach() + if(_mlx_jit_cuda_header_paths) + install(FILES ${_mlx_jit_cuda_header_paths} DESTINATION ${_dst} COMPONENT MLX) + endif() + + foreach(_runtime_dir cooperative_groups crt) + if(EXISTS "${_cuda_inc}/${_runtime_dir}") + install(DIRECTORY "${_cuda_inc}/${_runtime_dir}" + DESTINATION ${_dst} + COMPONENT MLX) + endif() endforeach() - # Subdirectory headers. - install(DIRECTORY "${_cuda_inc}/cooperative_groups" - DESTINATION ${_dst} - COMPONENT MLX - FILES_MATCHING PATTERN "*.h") - install(FILES "${_cuda_inc}/crt/host_defines.h" - DESTINATION "${_dst}/crt" - COMPONENT MLX) + if(NOT WIN32 AND NOT APPLE) install(CODE " set(_link \"${CMAKE_INSTALL_PREFIX}/${OLLAMA_LIB_DIR}/include\") diff --git a/scripts/build_linux.sh b/scripts/build_linux.sh index c3d083a9..b4485e75 100755 --- a/scripts/build_linux.sh +++ b/scripts/build_linux.sh @@ -28,6 +28,8 @@ if ! command -v zstd >/dev/null 2>&1; then exit 1 fi +rm -rf dist/bin dist/lib dist/linux_amd64 dist/linux_arm64 +rm -f dist/ollama-linux-*.tar.zst mkdir -p dist docker buildx build \ diff --git a/x/imagegen/mlx/compile.go b/x/imagegen/mlx/compile.go index 4d0587da..e152c305 100644 --- a/x/imagegen/mlx/compile.go +++ b/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 diff --git a/x/imagegen/mlx/mlx.go b/x/imagegen/mlx/mlx.go index 944b6a5b..3c3d354d 100644 --- a/x/imagegen/mlx/mlx.go +++ b/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() } diff --git a/x/imagegen/mlx/mlx_test.go b/x/imagegen/mlx/mlx_test.go index 939652db..06855e41 100644 --- a/x/imagegen/mlx/mlx_test.go +++ b/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") } diff --git a/x/imagegen/nn/nn_test.go b/x/imagegen/nn/nn_test.go index 32433dba..0fc433ef 100644 --- a/x/imagegen/nn/nn_test.go +++ b/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)