Mirror_Repos/ollama_ollama
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
Daniel Hiltgen
2026-06-24 10:36:02 -07:00
committed by GitHub
parent 89a171cc70
commit 570679c9e0
7 changed files with 315 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
MLX_VERSION
+1 -1
View File
@@ -1 +1 @@
2165dc08d7b33258260aa849d39f087d50e62962
51b2768da7e1897d3c4258f7ddbb47083d1eef01
cmake/mlx/CMakeLists.txt
+114 -31
View File
@@ -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\")
scripts/build_linux.sh
+2
View File
@@ -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 \
x/imagegen/mlx/compile.go
+36 -7
View File
@@ -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
+88 -6
View File
@@ -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
+42
View File
@@ -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
+32
View File
@@ -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)