ollama source for Momentry Core verification

2026-05-22 17:19:10 +08:00
commit 0b31ff9135
2020 changed files with 1413145 additions and 0 deletions
--- a/x/imagegen/mlx/CMakeLists.txt
+++ b/x/imagegen/mlx/CMakeLists.txt
@@ -0,0 +1,196 @@
+include(FetchContent)
+
+# Read MLX-C version from top-level file (shared with Dockerfile)
+file(READ "${CMAKE_SOURCE_DIR}/MLX_C_VERSION" MLX_C_GIT_TAG)
+string(STRIP "${MLX_C_GIT_TAG}" MLX_C_GIT_TAG)
+
+# Read MLX version from top-level file
+file(READ "${CMAKE_SOURCE_DIR}/MLX_VERSION" MLX_GIT_TAG)
+string(STRIP "${MLX_GIT_TAG}" MLX_GIT_TAG)
+
+set(MLX_C_BUILD_EXAMPLES OFF)
+
+set(MLX_BUILD_GGUF OFF)
+set(MLX_BUILD_SAFETENSORS ON)
+
+function(set_target_output_directory _target)
+    if(TARGET ${_target})
+        set_target_properties(${_target} PROPERTIES
+            RUNTIME_OUTPUT_DIRECTORY ${OLLAMA_BUILD_DIR}
+            LIBRARY_OUTPUT_DIRECTORY ${OLLAMA_BUILD_DIR}
+            ARCHIVE_OUTPUT_DIRECTORY ${OLLAMA_BUILD_DIR}
+        )
+    endif()
+endfunction()
+
+# Check for Metal support (macOS only)
+if(CMAKE_SYSTEM_NAME MATCHES "Darwin")
+    execute_process(
+      COMMAND
+        zsh "-c"
+        "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'"
+      OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY)
+
+    if(NOT MLX_METAL_VERSION)
+        message(STATUS "`xcrun metal` error. Setting MLX_BUILD_METAL=OFF")
+        set(MLX_BUILD_METAL OFF)
+    endif()
+else()
+    # On Linux, disable Metal backend
+    message(STATUS "Non-macOS platform detected. Setting MLX_BUILD_METAL=OFF")
+    set(MLX_BUILD_METAL OFF)
+endif()
+
+# Map CMAKE_CUDA_ARCHITECTURES to MLX_CUDA_ARCHITECTURES if not explicitly set
+if(NOT MLX_CUDA_ARCHITECTURES AND CMAKE_CUDA_ARCHITECTURES)
+    set(MLX_CUDA_ARCHITECTURES ${CMAKE_CUDA_ARCHITECTURES})
+    message(STATUS "Using CMAKE_CUDA_ARCHITECTURES for MLX: ${MLX_CUDA_ARCHITECTURES}")
+endif()
+
+# Forward cuDNN environment variables to cmake variables so MLX's FindCUDNN.cmake
+# can find them via HINTS ${CUDNN_INCLUDE_PATH} / ${CUDNN_LIBRARY_PATH}.
+if(DEFINED ENV{CUDNN_INCLUDE_PATH} AND NOT CUDNN_INCLUDE_PATH)
+    set(CUDNN_INCLUDE_PATH "$ENV{CUDNN_INCLUDE_PATH}" CACHE PATH "cuDNN include path")
+    message(STATUS "Using CUDNN_INCLUDE_PATH from environment: ${CUDNN_INCLUDE_PATH}")
+endif()
+if(DEFINED ENV{CUDNN_LIBRARY_PATH} AND NOT CUDNN_LIBRARY_PATH)
+    set(CUDNN_LIBRARY_PATH "$ENV{CUDNN_LIBRARY_PATH}" CACHE PATH "cuDNN library path")
+    message(STATUS "Using CUDNN_LIBRARY_PATH from environment: ${CUDNN_LIBRARY_PATH}")
+endif()
+
+# Enable CUDA backend if CUDA architectures are specified and CUDA compiler is available
+if(MLX_CUDA_ARCHITECTURES AND CMAKE_CUDA_COMPILER)
+    set(MLX_BUILD_CUDA ON CACHE BOOL "Build CUDA backend for MLX" FORCE)
+    message(STATUS "Enabling MLX CUDA backend with architectures: ${MLX_CUDA_ARCHITECTURES}")
+elseif(MLX_CUDA_ARCHITECTURES)
+    message(WARNING "MLX_CUDA_ARCHITECTURES specified but CUDA compiler not found, CUDA backend will be disabled")
+endif()
+
+# Allow local source overrides via environment variables
+# Resolve to absolute paths so FetchContent doesn't break on relative dirs.
+if(DEFINED ENV{OLLAMA_MLX_SOURCE})
+    get_filename_component(_mlx_src "$ENV{OLLAMA_MLX_SOURCE}" ABSOLUTE BASE_DIR ${CMAKE_SOURCE_DIR})
+    set(FETCHCONTENT_SOURCE_DIR_MLX "${_mlx_src}" CACHE PATH "" FORCE)
+    message(STATUS "Using local MLX source: ${_mlx_src}")
+endif()
+if(DEFINED ENV{OLLAMA_MLX_C_SOURCE})
+    get_filename_component(_mlx_c_src "$ENV{OLLAMA_MLX_C_SOURCE}" ABSOLUTE BASE_DIR ${CMAKE_SOURCE_DIR})
+    set(FETCHCONTENT_SOURCE_DIR_MLX-C "${_mlx_c_src}" CACHE PATH "" FORCE)
+    message(STATUS "Using local MLX-C source: ${_mlx_c_src}")
+endif()
+
+# Pre-declare mlx so our pinned version takes precedence over the one
+# hardcoded in mlx-c's CMakeLists.txt (first FetchContent_Declare wins).
+FetchContent_Declare(
+    mlx
+    GIT_REPOSITORY "https://github.com/ml-explore/mlx.git"
+    GIT_TAG ${MLX_GIT_TAG}
+)
+
+FetchContent_Declare(
+    mlx-c
+    GIT_REPOSITORY "https://github.com/ml-explore/mlx-c.git"
+    GIT_TAG ${MLX_C_GIT_TAG}
+)
+FetchContent_MakeAvailable(mlx-c)
+
+# To avoid a "long tail" when building MLX with a large set of GPU
+# architectures, utilize a higher --threads (-t) setting. At high -t
+# every .cu spawns concurrent cicc instances; each cicc can consume several GB
+# compiling MLX's CUTLASS-using kernels.  This in turn can cause OOMs.
+#
+# We use a pool to cover all MLX CUDA sources. Pool size is derived from total 
+# host RAM via a per-file memory budget.
+#
+# This was calibrated with `-t 6`.  Higher -t may require overriding
+# MLX_CUDA_RAM_MB
+if(CMAKE_GENERATOR STREQUAL "Ninja")
+    file(GLOB_RECURSE _mlx_cu
+        "${mlx_SOURCE_DIR}/mlx/backend/cuda/*.cu"
+        "${mlx_BINARY_DIR}/mlx/backend/cuda/*.cu"
+    )
+    if(_mlx_cu)
+        set(MLX_CUDA_RAM_MB 22000 CACHE STRING
+            "Per-file memory budget (MB) for the cuda_compile JOB_POOL.  Override for higher -t.")
+        cmake_host_system_information(RESULT _ram_mb QUERY TOTAL_PHYSICAL_MEMORY)
+        math(EXPR _cuda_pool "${_ram_mb} / ${MLX_CUDA_RAM_MB}")
+        if(_cuda_pool LESS 2)
+            set(_cuda_pool 2)
+        endif()
+        set_property(GLOBAL APPEND PROPERTY JOB_POOLS cuda_compile=${_cuda_pool})
+        list(LENGTH _mlx_cu _cu_count)
+        # SOURCE properties default to directory-scoped, which means a plain
+        # set_property(SOURCE ...) here would NOT affect the build rules
+        # generated for the mlx target (defined in mlx_SOURCE_DIR after
+        # FetchContent).  TARGET_DIRECTORY mlx puts the property in the
+        # directory where mlx was defined, so it actually applies.
+        foreach(f ${_mlx_cu})
+            set_property(SOURCE "${f}"
+                TARGET_DIRECTORY mlx
+                PROPERTY JOB_POOL_COMPILE cuda_compile)
+        endforeach()
+        message(STATUS "MLX cuda_compile JOB_POOL: ${_cu_count} files, pool size ${_cuda_pool} (host RAM ${_ram_mb} MB / ${MLX_CUDA_RAM_MB} MB per file)")
+    else()
+        message(WARNING "MLX cuda_compile JOB_POOL: no .cu files found under mlx/backend/cuda/ - check MLX layout")
+    endif()
+endif()
+
+# Sync vendored headers with fetched version
+file(GLOB _mlx_c_hdrs "${mlx-c_SOURCE_DIR}/mlx/c/*.h")
+file(COPY ${_mlx_c_hdrs} DESTINATION "${CMAKE_SOURCE_DIR}/x/mlxrunner/mlx/include/mlx/c/")
+
+# Regenerate Go/C shim wrappers from the (possibly updated) headers.
+# Skip during cross-compilation — the generated files are arch-independent.
+if(CMAKE_SYSTEM_PROCESSOR STREQUAL CMAKE_HOST_SYSTEM_PROCESSOR OR NOT APPLE)
+    find_program(GO_EXECUTABLE go REQUIRED)
+    message(STATUS "Regenerating MLX Go wrappers")
+
+    # CGo's probe compilation is sensitive to CGO_CFLAGS/CGO_CXXFLAGS and CC.
+    # Clear them so go generate uses default compiler settings:
+    # - On Windows, CC may contain spaces (e.g., "C:/Program Files/.../cl.exe")
+    #   which breaks CGo's CC parsing.
+    # - On macOS, CGO_CFLAGS with -mmacosx-version-min breaks header search
+    #   when cmake also sets CMAKE_OSX_DEPLOYMENT_TARGET.
+    set(_SAVE_CC "$ENV{CC}")
+    set(_SAVE_CGO_CFLAGS "$ENV{CGO_CFLAGS}")
+    set(_SAVE_CGO_CXXFLAGS "$ENV{CGO_CXXFLAGS}")
+    set(ENV{CC} "")
+    set(ENV{CGO_CFLAGS} "")
+    set(ENV{CGO_CXXFLAGS} "")
+
+    execute_process(
+        COMMAND ${GO_EXECUTABLE} generate ./x/...
+        WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
+        COMMAND_ERROR_IS_FATAL ANY
+    )
+
+    set(ENV{CC} "${_SAVE_CC}")
+    set(ENV{CGO_CFLAGS} "${_SAVE_CGO_CFLAGS}")
+    set(ENV{CGO_CXXFLAGS} "${_SAVE_CGO_CXXFLAGS}")
+else()
+    message(STATUS "Skipping MLX Go wrapper generation (cross-compiling)")
+endif()
+
+# For local dev builds, override MLX_VERSION with git describe output
+if(TARGET mlx_version AND DEFINED FETCHCONTENT_SOURCE_DIR_MLX)
+    execute_process(
+        COMMAND git describe --tags --first-parent --abbrev=7 --long --dirty --always
+        WORKING_DIRECTORY ${mlx_SOURCE_DIR}
+        OUTPUT_VARIABLE _mlx_git_version
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        ERROR_QUIET
+        RESULT_VARIABLE _mlx_git_result
+    )
+    if(_mlx_git_result EQUAL 0 AND _mlx_git_version)
+        # Strip leading "v" prefix for consistency
+        string(REGEX REPLACE "^v" "" _mlx_git_version "${_mlx_git_version}")
+        get_target_property(_mlx_defs mlx_version COMPILE_DEFINITIONS)
+        list(FILTER _mlx_defs EXCLUDE REGEX "^MLX_VERSION=")
+        set_target_properties(mlx_version PROPERTIES COMPILE_DEFINITIONS "${_mlx_defs}")
+        target_compile_definitions(mlx_version PRIVATE "MLX_VERSION=\"${_mlx_git_version}\"")
+        message(STATUS "MLX version (local dev): ${_mlx_git_version}")
+    endif()
+endif()
+
+set_target_output_directory(mlx)
+set_target_output_directory(mlxc)
--- a/x/imagegen/mlx/README.md
+++ b/x/imagegen/mlx/README.md
@@ -0,0 +1,46 @@
+# MLX Memory Management
+
+| This package will get consolidated with `x/ml/backend/mlx` in the future.
+
+## Automatic Tracking
+
+All arrays are automatically tracked when created. On `Eval()`, non-kept arrays are freed.
+
+### API
+
+```go
+result := mlx.Matmul(x, w) // arrays automatically tracked
+mlx.Eval(result)           // free non-kept, eval result (auto-kept)
+```
+
+### Key Functions
+
+- `mlx.Eval(outputs...)` - free non-kept arrays, then evaluate (outputs auto-kept)
+- `mlx.AsyncEval(outputs...)` - async version of Eval (outputs auto-kept)
+- `mlx.Keep(arrays...)` - mark arrays to survive cleanup (for weights, caches)
+- `array.Free()` - mark array for cleanup on next Eval
+
+### Loop Pattern
+
+```go
+for step := 0; step < maxTokens; step++ {
+    logits := model.Forward(token, caches)
+    oldToken := token
+    token = sample(logits)
+
+    // Keep cache state across iterations
+    for _, c := range caches {
+        mlx.Keep(c.State()...)
+    }
+
+    oldToken.Free()       // mark for cleanup
+    mlx.AsyncEval(token)  // frees old, evals new
+}
+```
+
+### Notes
+
+- `Eval()` and `AsyncEval()` auto-keep their outputs
+- `Free()` marks for cleanup - actual free happens during next Eval
+- Use `Keep()` for weights and cache state that must survive multiple Eval cycles
+- Arrays created inside compiled closures are managed by MLX, not tracked
--- a/x/imagegen/mlx/compile.go
+++ b/x/imagegen/mlx/compile.go
@@ -0,0 +1,171 @@
+package mlx
+
+/*
+#include "mlx.h"
+#include <stdlib.h>
+
+// Forward declaration for Go callback
+extern int goClosureCallback(mlx_vector_array* res, mlx_vector_array input, void* payload);
+
+// Destructor for payload (Go handle)
+extern void goClosureDestructor(void* payload);
+*/
+import "C"
+import (
+	"runtime/cgo"
+	"sync"
+	"unsafe"
+)
+
+// inClosureCallback is set to true during closure callback execution.
+var inClosureCallback bool
+var closureCallbackMu sync.Mutex
+
+// InClosureCallback returns true if we're currently executing inside a closure callback.
+func InClosureCallback() bool {
+	closureCallbackMu.Lock()
+	defer closureCallbackMu.Unlock()
+	return inClosureCallback
+}
+
+// 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.
+type CompiledFunc struct {
+	closure  C.mlx_closure
+	compiled C.mlx_closure
+}
+
+// ClosureFunc is the signature for functions that can be compiled.
+// It takes a slice of input arrays and returns a slice of output arrays.
+type ClosureFunc func(inputs []*Array) []*Array
+
+// Compile compiles a Go function into an optimized MLX closure.
+// The function is traced once during compilation, then subsequent calls
+// run the optimized graph without creating Go intermediate arrays.
+//
+// Example:
+//
+//	compiled := mlx.Compile(func(inputs []*mlx.Array) []*mlx.Array {
+//	    a, b := inputs[0], inputs[1]
+//	    c := mlx.Add(a, b)
+//	    d := mlx.Mul(c, c)
+//	    return []*mlx.Array{d}
+//	})
+//	defer compiled.Free()
+//
+//	result := compiled.Call(x, y)[0]
+func Compile(fn ClosureFunc) *CompiledFunc {
+	return CompileShapeless(fn, false)
+}
+
+// CompileShapeless compiles with optional shapeless mode.
+// If shapeless=true, the function works for any input shape after tracing.
+func CompileShapeless(fn ClosureFunc, shapeless bool) *CompiledFunc {
+	// Create a cgo.Handle to prevent the Go function from being GC'd
+	handle := cgo.NewHandle(fn)
+
+	// Create the closure from the Go callback
+	closure := C.mlx_closure_new_func_payload(
+		(*[0]byte)(C.goClosureCallback),
+		unsafe.Pointer(handle),
+		(*[0]byte)(C.goClosureDestructor),
+	)
+
+	// Compile the closure
+	compiled := C.mlx_closure_new()
+	C.mlx_compile(&compiled, closure, C.bool(shapeless))
+
+	return &CompiledFunc{
+		closure:  closure,
+		compiled: compiled,
+	}
+}
+
+// Call invokes the compiled function with the given inputs.
+func (cf *CompiledFunc) Call(inputs ...*Array) []*Array {
+	// Pack inputs into vector
+	inputVec := C.mlx_vector_array_new()
+	for _, arr := range inputs {
+		C.mlx_vector_array_append_value(inputVec, arr.c)
+	}
+
+	// Apply compiled closure
+	outputVec := C.mlx_vector_array_new()
+	C.mlx_closure_apply(&outputVec, cf.compiled, inputVec)
+	C.mlx_vector_array_free(inputVec)
+
+	// Unpack outputs
+	numOutputs := int(C.mlx_vector_array_size(outputVec))
+	outputs := make([]*Array, numOutputs)
+	for i := 0; i < numOutputs; i++ {
+		var arr C.mlx_array
+		C.mlx_vector_array_get(&arr, outputVec, C.size_t(i))
+		outputs[i] = newArray(arr)
+	}
+	C.mlx_vector_array_free(outputVec)
+
+	return outputs
+}
+
+// CallEval invokes the compiled function and evaluates the results.
+func (cf *CompiledFunc) CallEval(inputs ...*Array) []*Array {
+	outputs := cf.Call(inputs...)
+	Eval(outputs...)
+	return outputs
+}
+
+// Free releases the compiled function resources.
+func (cf *CompiledFunc) Free() {
+	C.mlx_closure_free(cf.compiled)
+	C.mlx_closure_free(cf.closure)
+}
+
+// borrowArray wraps a C array WITHOUT setting up GC cleanup.
+// Use this for arrays we don't own (e.g., borrowed references in callbacks).
+func borrowArray(array C.mlx_array) *Array {
+	return &Array{c: 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
+	closureCallbackMu.Lock()
+	inClosureCallback = true
+	closureCallbackMu.Unlock()
+	defer func() {
+		closureCallbackMu.Lock()
+		inClosureCallback = false
+		closureCallbackMu.Unlock()
+	}()
+
+	// Recover the Go function from the handle
+	handle := cgo.Handle(payload)
+	fn := handle.Value().(ClosureFunc)
+
+	// Convert input vector to Go slice - use borrowArray since MLX owns these
+	numInputs := int(C.mlx_vector_array_size(input))
+	inputs := make([]*Array, numInputs)
+	for i := 0; i < numInputs; i++ {
+		var arr C.mlx_array
+		C.mlx_vector_array_get(&arr, input, C.size_t(i))
+		inputs[i] = borrowArray(arr) // Don't set up cleanup - MLX owns these
+	}
+
+	// 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)
+	}
+
+	return 0
+}
+
+//export goClosureDestructor
+func goClosureDestructor(payload unsafe.Pointer) {
+	handle := cgo.Handle(payload)
+	handle.Delete()
+}
--- a/x/imagegen/mlx/doc.go
+++ b/x/imagegen/mlx/doc.go
@@ -0,0 +1,4 @@
+// Package mlx provides Go bindings for the MLX-C library with dynamic loading support.
+//
+//go:generate go run generate_wrappers.go ../../mlxrunner/mlx/include/mlx/c mlx.h mlx.c
+package mlx
--- a/x/imagegen/mlx/generate_wrappers.go
+++ b/x/imagegen/mlx/generate_wrappers.go
@@ -0,0 +1,451 @@
+//go:build ignore
+
+// This tool generates MLX-C dynamic loading wrappers.
+// Usage: go run generate_wrappers.go <mlx-c-include-dir> <output-header> [output-impl]
+package main
+
+import (
+	"bytes"
+	"flag"
+	"fmt"
+	"io/fs"
+	"os"
+	"path/filepath"
+	"regexp"
+	"strings"
+)
+
+type Function struct {
+	Name            string
+	ReturnType      string
+	Params          string
+	ParamNames      []string
+	NeedsARM64Guard bool
+}
+
+func findHeaders(directory string) ([]string, error) {
+	var headers []string
+	err := filepath.WalkDir(directory, func(path string, d fs.DirEntry, err error) error {
+		if err != nil {
+			return err
+		}
+		// Private headers contain C++ implementation helpers and are not part of
+		// the C API surface; parsing them can produce invalid wrapper signatures.
+		if d.IsDir() && d.Name() == "private" {
+			return fs.SkipDir
+		}
+		if !d.IsDir() && strings.HasSuffix(path, ".h") {
+			headers = append(headers, path)
+		}
+		return nil
+	})
+	return headers, err
+}
+
+func cleanContent(content string) string {
+	// Remove single-line comments
+	re := regexp.MustCompile(`//.*?\n`)
+	content = re.ReplaceAllString(content, "\n")
+
+	// Remove multi-line comments
+	re = regexp.MustCompile(`/\*.*?\*/`)
+	content = re.ReplaceAllString(content, "")
+
+	// Remove preprocessor directives (lines starting with #) - use multiline mode
+	re = regexp.MustCompile(`(?m)^\s*#.*?$`)
+	content = re.ReplaceAllString(content, "")
+
+	// Remove extern "C" { and } blocks more conservatively
+	// Only remove the extern "C" { line, not the content inside
+	re = regexp.MustCompile(`extern\s+"C"\s*\{\s*?\n`)
+	content = re.ReplaceAllString(content, "\n")
+	// Remove standalone closing braces that are not part of function declarations
+	re = regexp.MustCompile(`\n\s*\}\s*\n`)
+	content = re.ReplaceAllString(content, "\n")
+
+	// Collapse whitespace and newlines
+	re = regexp.MustCompile(`\s+`)
+	content = re.ReplaceAllString(content, " ")
+
+	return content
+}
+
+func extractParamNames(params string) []string {
+	if params == "" || strings.TrimSpace(params) == "void" {
+		return []string{}
+	}
+
+	var names []string
+
+	// Split by comma, but respect parentheses (for function pointers)
+	parts := splitParams(params)
+
+	// Remove array brackets
+	arrayBrackets := regexp.MustCompile(`\[.*?\]`)
+
+	// Function pointer pattern
+	funcPtrPattern := regexp.MustCompile(`\(\s*\*\s*(\w+)\s*\)`)
+
+	// Type keywords to skip
+	typeKeywords := map[string]bool{
+		"const":     true,
+		"struct":    true,
+		"unsigned":  true,
+		"signed":    true,
+		"long":      true,
+		"short":     true,
+		"int":       true,
+		"char":      true,
+		"float":     true,
+		"double":    true,
+		"void":      true,
+		"size_t":    true,
+		"uint8_t":   true,
+		"uint16_t":  true,
+		"uint32_t":  true,
+		"uint64_t":  true,
+		"int8_t":    true,
+		"int16_t":   true,
+		"int32_t":   true,
+		"int64_t":   true,
+		"intptr_t":  true,
+		"uintptr_t": true,
+	}
+
+	for _, part := range parts {
+		if part == "" {
+			continue
+		}
+
+		// Remove array brackets
+		part = arrayBrackets.ReplaceAllString(part, "")
+
+		// For function pointers like "void (*callback)(int)"
+		if matches := funcPtrPattern.FindStringSubmatch(part); len(matches) > 1 {
+			names = append(names, matches[1])
+			continue
+		}
+
+		// Regular parameter: last identifier
+		tokens := regexp.MustCompile(`\w+`).FindAllString(part, -1)
+		if len(tokens) > 0 {
+			// The last token is usually the parameter name
+			// Skip type keywords
+			for i := len(tokens) - 1; i >= 0; i-- {
+				if !typeKeywords[tokens[i]] {
+					names = append(names, tokens[i])
+					break
+				}
+			}
+		}
+	}
+
+	return names
+}
+
+func splitParams(params string) []string {
+	var parts []string
+	var current bytes.Buffer
+	depth := 0
+
+	for _, char := range params + "," {
+		switch char {
+		case '(':
+			depth++
+			current.WriteRune(char)
+		case ')':
+			depth--
+			current.WriteRune(char)
+		case ',':
+			if depth == 0 {
+				parts = append(parts, strings.TrimSpace(current.String()))
+				current.Reset()
+			} else {
+				current.WriteRune(char)
+			}
+		default:
+			current.WriteRune(char)
+		}
+	}
+
+	return parts
+}
+
+func parseFunctions(content string) []Function {
+	var functions []Function
+
+	// Match function declarations: return_type function_name(params);
+	// Matches both mlx_* and _mlx_* functions
+	pattern := regexp.MustCompile(`\b((?:const\s+)?(?:struct\s+)?[\w\s]+?[\*\s]*)\s+(_?mlx_\w+)\s*\(([^)]*(?:\([^)]*\)[^)]*)*)\)\s*;`)
+
+	matches := pattern.FindAllStringSubmatch(content, -1)
+	for _, match := range matches {
+		returnType := strings.TrimSpace(match[1])
+		funcName := strings.TrimSpace(match[2])
+		params := strings.TrimSpace(match[3])
+
+		// Skip if this looks like a variable declaration
+		if params == "" || strings.Contains(params, "{") {
+			continue
+		}
+
+		// Clean up return type
+		returnType = strings.Join(strings.Fields(returnType), " ")
+
+		// Extract parameter names
+		paramNames := extractParamNames(params)
+
+		// Check if ARM64 guard is needed
+		needsGuard := needsARM64Guard(funcName, returnType, params)
+
+		functions = append(functions, Function{
+			Name:            funcName,
+			ReturnType:      returnType,
+			Params:          params,
+			ParamNames:      paramNames,
+			NeedsARM64Guard: needsGuard,
+		})
+	}
+
+	return functions
+}
+
+func needsARM64Guard(name, retType, params string) bool {
+	return strings.Contains(name, "float16") ||
+		strings.Contains(name, "bfloat16") ||
+		strings.Contains(retType, "float16_t") ||
+		strings.Contains(retType, "bfloat16_t") ||
+		strings.Contains(params, "float16_t") ||
+		strings.Contains(params, "bfloat16_t")
+}
+
+func generateWrapperFiles(functions []Function, headerPath, implPath string) error {
+	// Generate header file
+	var headerBuf bytes.Buffer
+
+	headerBuf.WriteString("// AUTO-GENERATED by generate_wrappers.go - DO NOT EDIT\n")
+	headerBuf.WriteString("// This file provides wrapper declarations for MLX-C functions that use dlopen/dlsym\n")
+	headerBuf.WriteString("//\n")
+	headerBuf.WriteString("// Strategy: Include MLX-C headers for type definitions, then provide wrapper\n")
+	headerBuf.WriteString("// functions that shadow the originals, allowing Go code to call them directly (e.g., C.mlx_add).\n")
+	headerBuf.WriteString("// Function pointers are defined in mlx.c (single compilation unit).\n\n")
+	headerBuf.WriteString("#ifndef MLX_WRAPPERS_H\n")
+	headerBuf.WriteString("#define MLX_WRAPPERS_H\n\n")
+
+	headerBuf.WriteString("// Include MLX headers for type definitions and original declarations\n")
+	headerBuf.WriteString("#include \"mlx/c/mlx.h\"\n")
+	headerBuf.WriteString("#include \"mlx_dynamic.h\"\n")
+	headerBuf.WriteString("#include <stdio.h>\n\n")
+
+	// Undef all MLX functions to avoid conflicts
+	headerBuf.WriteString("// Undefine any existing MLX function macros\n")
+	for _, fn := range functions {
+		headerBuf.WriteString(fmt.Sprintf("#undef %s\n", fn.Name))
+	}
+	headerBuf.WriteString("\n")
+
+	// Function pointer extern declarations
+	headerBuf.WriteString("// Function pointer declarations (defined in mlx.c, loaded via dlsym)\n")
+	for _, fn := range functions {
+		if fn.NeedsARM64Guard {
+			headerBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
+		}
+		headerBuf.WriteString(fmt.Sprintf("extern %s (*%s_ptr)(%s);\n", fn.ReturnType, fn.Name, fn.Params))
+		if fn.NeedsARM64Guard {
+			headerBuf.WriteString("#endif\n")
+		}
+	}
+	headerBuf.WriteString("\n")
+
+	// Initialization function declaration
+	headerBuf.WriteString("// Initialize all function pointers via dlsym (defined in mlx.c)\n")
+	headerBuf.WriteString("int mlx_load_functions(void* handle);\n\n")
+
+	// Wrapper function declarations
+	headerBuf.WriteString("// Wrapper function declarations that call through function pointers\n")
+	headerBuf.WriteString("// Go code calls these directly as C.mlx_* (no #define redirection needed)\n")
+	for _, fn := range functions {
+		if fn.NeedsARM64Guard {
+			headerBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
+		}
+		headerBuf.WriteString(fmt.Sprintf("%s %s(%s);\n", fn.ReturnType, fn.Name, fn.Params))
+		if fn.NeedsARM64Guard {
+			headerBuf.WriteString("#endif\n")
+		}
+		headerBuf.WriteString("\n")
+	}
+
+	headerBuf.WriteString("#endif // MLX_WRAPPERS_H\n")
+
+	// Write header file
+	if err := os.WriteFile(headerPath, headerBuf.Bytes(), 0644); err != nil {
+		return fmt.Errorf("failed to write header file: %w", err)
+	}
+
+	// Generate implementation file
+	var implBuf bytes.Buffer
+
+	implBuf.WriteString("// AUTO-GENERATED by generate_wrappers.go - DO NOT EDIT\n")
+	implBuf.WriteString("// This file contains the function pointer definitions and initialization\n")
+	implBuf.WriteString("// All function pointers are in a single compilation unit to avoid duplication\n\n")
+
+	implBuf.WriteString("#include \"mlx/c/mlx.h\"\n")
+	implBuf.WriteString("#include \"mlx_dynamic.h\"\n")
+	implBuf.WriteString("#include <stdio.h>\n\n")
+	implBuf.WriteString("// Platform-specific dynamic loading\n")
+	implBuf.WriteString("#ifdef _WIN32\n")
+	implBuf.WriteString("#include <windows.h>\n")
+	implBuf.WriteString("#define GET_SYM(handle, name) (void*)GetProcAddress((HMODULE)(handle), name)\n")
+	implBuf.WriteString("#else\n")
+	implBuf.WriteString("#include <dlfcn.h>\n")
+	implBuf.WriteString("#define GET_SYM(handle, name) dlsym(handle, name)\n")
+	implBuf.WriteString("#endif\n\n")
+
+	// Function pointer definitions
+	implBuf.WriteString("// Function pointer definitions\n")
+	for _, fn := range functions {
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
+		}
+		implBuf.WriteString(fmt.Sprintf("%s (*%s_ptr)(%s) = NULL;\n", fn.ReturnType, fn.Name, fn.Params))
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#endif\n")
+		}
+	}
+	implBuf.WriteString("\n")
+
+	// Initialization function
+	implBuf.WriteString("// Initialize all function pointers\n")
+	implBuf.WriteString("int mlx_load_functions(void* handle) {\n")
+	implBuf.WriteString("    if (handle == NULL) {\n")
+	implBuf.WriteString("        fprintf(stderr, \"MLX: Invalid library handle\\n\");\n")
+	implBuf.WriteString("        return -1;\n")
+	implBuf.WriteString("    }\n\n")
+
+	for _, fn := range functions {
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
+		}
+		implBuf.WriteString(fmt.Sprintf("    %s_ptr = GET_SYM(handle, \"%s\");\n", fn.Name, fn.Name))
+		implBuf.WriteString(fmt.Sprintf("    if (%s_ptr == NULL) {\n", fn.Name))
+		implBuf.WriteString(fmt.Sprintf("        fprintf(stderr, \"MLX: Failed to load symbol: %s\\n\");\n", fn.Name))
+		implBuf.WriteString("        return -1;\n")
+		implBuf.WriteString("    }\n")
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#endif\n")
+		}
+	}
+
+	implBuf.WriteString("    return 0;\n")
+	implBuf.WriteString("}\n\n")
+
+	// Wrapper function implementations
+	implBuf.WriteString("// Wrapper function implementations that call through function pointers\n")
+	for _, fn := range functions {
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
+		}
+		implBuf.WriteString(fmt.Sprintf("%s %s(%s) {\n", fn.ReturnType, fn.Name, fn.Params))
+
+		// Call through function pointer
+		if fn.ReturnType != "void" {
+			implBuf.WriteString(fmt.Sprintf("    return %s_ptr(", fn.Name))
+		} else {
+			implBuf.WriteString(fmt.Sprintf("    %s_ptr(", fn.Name))
+		}
+
+		// Pass parameters
+		implBuf.WriteString(strings.Join(fn.ParamNames, ", "))
+		implBuf.WriteString(");\n")
+		implBuf.WriteString("}\n")
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#endif\n")
+		}
+		implBuf.WriteString("\n")
+	}
+
+	// Write implementation file
+	if err := os.WriteFile(implPath, implBuf.Bytes(), 0644); err != nil {
+		return fmt.Errorf("failed to write implementation file: %w", err)
+	}
+
+	return nil
+}
+
+func main() {
+	flag.Usage = func() {
+		fmt.Fprintf(flag.CommandLine.Output(), "Usage: go run generate_wrappers.go <mlx-c-include-dir> <output-header> [output-impl]\n")
+		fmt.Fprintf(flag.CommandLine.Output(), "Generate MLX-C dynamic loading wrappers.\n\n")
+		flag.PrintDefaults()
+	}
+	flag.Parse()
+
+	args := flag.Args()
+	if len(args) < 2 {
+		fmt.Fprintf(flag.CommandLine.Output(), "ERROR: Missing required arguments\n\n")
+		flag.Usage()
+		os.Exit(1)
+	}
+
+	headerDir := args[0]
+	outputHeader := args[1]
+	// Default implementation file is same name with .c extension
+	outputImpl := outputHeader
+	if len(args) > 2 {
+		outputImpl = args[2]
+	} else if strings.HasSuffix(outputHeader, ".h") {
+		outputImpl = outputHeader[:len(outputHeader)-2] + ".c"
+	}
+
+	// Check if header directory exists
+	if _, err := os.Stat(headerDir); os.IsNotExist(err) {
+		fmt.Fprintf(os.Stderr, "ERROR: MLX-C headers directory not found at: %s\n\n", headerDir)
+		fmt.Fprintf(os.Stderr, "Please run CMake first to download MLX-C dependencies:\n")
+		fmt.Fprintf(os.Stderr, "  cmake -B build\n\n")
+		fmt.Fprintf(os.Stderr, "The CMake build will download and extract MLX-C headers needed for wrapper generation.\n")
+		os.Exit(1)
+	}
+
+	fmt.Fprintf(os.Stderr, "Parsing MLX-C headers from: %s\n", headerDir)
+
+	// Find all headers
+	headers, err := findHeaders(headerDir)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "ERROR: Failed to find header files: %v\n", err)
+		os.Exit(1)
+	}
+	fmt.Fprintf(os.Stderr, "Found %d header files\n", len(headers))
+
+	// Parse all headers
+	var allFunctions []Function
+	seen := make(map[string]bool)
+
+	for _, header := range headers {
+		content, err := os.ReadFile(header)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Error reading %s: %v\n", header, err)
+			continue
+		}
+
+		cleaned := cleanContent(string(content))
+		functions := parseFunctions(cleaned)
+
+		// Deduplicate
+		for _, fn := range functions {
+			if !seen[fn.Name] {
+				seen[fn.Name] = true
+				allFunctions = append(allFunctions, fn)
+			}
+		}
+	}
+
+	fmt.Fprintf(os.Stderr, "Found %d unique function declarations\n", len(allFunctions))
+
+	// Generate wrapper files
+	if err := generateWrapperFiles(allFunctions, outputHeader, outputImpl); err != nil {
+		fmt.Fprintf(os.Stderr, "ERROR: Failed to generate wrapper files: %v\n", err)
+		os.Exit(1)
+	}
+
+	fmt.Fprintf(os.Stderr, "Generated %s and %s successfully\n", outputHeader, outputImpl)
+}
--- a/x/imagegen/mlx/mlx.c
+++ b/x/imagegen/mlx/mlx.c
--- a/x/imagegen/mlx/mlx.go
+++ b/x/imagegen/mlx/mlx.go
--- a/x/imagegen/mlx/mlx.h
+++ b/x/imagegen/mlx/mlx.h
--- a/x/imagegen/mlx/mlx_dynamic.c
+++ b/x/imagegen/mlx/mlx_dynamic.c
@@ -0,0 +1,93 @@
+// mlx_dynamic.c - Dynamic loading wrapper for MLX-C library
+// This file provides runtime dynamic loading of libmlxc instead of link-time binding
+
+#include "mlx_dynamic.h"
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#ifdef _WIN32
+#include <windows.h>
+typedef HMODULE lib_handle_t;
+static char win_error_buffer[256] = {0};
+static const char* get_win_error(void) {
+    DWORD err = GetLastError();
+    snprintf(win_error_buffer, sizeof(win_error_buffer), "error code %lu", err);
+    return win_error_buffer;
+}
+#define LIB_ERROR() get_win_error()
+#else
+#include <dlfcn.h>
+typedef void* lib_handle_t;
+#define LIB_ERROR() dlerror()
+#endif
+
+static lib_handle_t mlx_handle = NULL;
+static int mlx_initialized = 0;
+static char mlx_error_buffer[512] = {0};
+
+#ifdef _WIN32
+// Windows: Load library from a path with dependency resolution.
+// Temporarily adds the library's directory to the DLL search path
+// so that dependencies (like mlx.dll) in the same directory are found.
+static int try_load_win(const char* path) {
+    if (!path) return 0;
+
+    // Extract directory and add to DLL search path for dependency resolution
+    char dir_path[MAX_PATH];
+    strncpy(dir_path, path, MAX_PATH - 1);
+    dir_path[MAX_PATH - 1] = '\0';
+    char* last_slash = strrchr(dir_path, '\\');
+    if (!last_slash) last_slash = strrchr(dir_path, '/');
+    if (last_slash) {
+        *last_slash = '\0';
+        SetDllDirectoryA(dir_path);
+    }
+
+    mlx_handle = LoadLibraryA(path);
+    SetDllDirectoryA(NULL);
+    return mlx_handle != NULL;
+}
+#endif
+
+// Try to load library from a specific path
+static int try_load_lib(const char* path) {
+    if (!path) return 0;
+#ifdef _WIN32
+    return try_load_win(path);
+#else
+    mlx_handle = dlopen(path, RTLD_LAZY | RTLD_GLOBAL);
+    return mlx_handle != NULL;
+#endif
+}
+
+// Initialize the MLX dynamic library from a specific path.
+// Returns 0 on success, -1 on failure.
+int mlx_dynamic_init_path(const char* path) {
+    if (mlx_initialized) {
+        return 0;
+    }
+
+    if (try_load_lib(path)) {
+        mlx_initialized = 1;
+        snprintf(mlx_error_buffer, sizeof(mlx_error_buffer),
+                 "MLX: Successfully loaded %s", path ? path : "library");
+        return 0;
+    }
+
+    const char* err = LIB_ERROR();
+    snprintf(mlx_error_buffer, sizeof(mlx_error_buffer),
+             "MLX: Failed to load %s: %s", path ? path : "(null)", err ? err : "unknown error");
+    return -1;
+}
+
+// Get the last error message
+const char* mlx_dynamic_error(void) {
+    return mlx_error_buffer;
+}
+
+// Get the library handle (for use by generated wrappers)
+void* mlx_get_handle(void) {
+    return mlx_handle;
+}
+
--- a/x/imagegen/mlx/mlx_dynamic.h
+++ b/x/imagegen/mlx/mlx_dynamic.h
@@ -0,0 +1,23 @@
+// mlx_dynamic.h - Dynamic loading interface for MLX-C library
+#ifndef MLX_DYNAMIC_H
+#define MLX_DYNAMIC_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Initialize the MLX dynamic library from a specific path
+// Returns 0 on success, -1 on failure
+int mlx_dynamic_init_path(const char* path);
+
+// Get the last error message from dynamic loading
+const char* mlx_dynamic_error(void);
+
+// Get the library handle (for use by generated wrappers)
+void* mlx_get_handle(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLX_DYNAMIC_H
--- a/x/imagegen/mlx/mlx_error_handler.c
+++ b/x/imagegen/mlx/mlx_error_handler.c
@@ -0,0 +1,38 @@
+// mlx_error_handler.c - Safe error handling for MLX initialization
+// Provides a non-fatal error handler for use during init(), so that
+// GPU failures are captured instead of calling exit(-1).
+
+#include "mlx_error_handler.h"
+#include "mlx.h"
+#include <string.h>
+
+static char mlx_init_error_msg[1024] = {0};
+static int  mlx_init_error_flag = 0;
+
+// Error handler that silently stores the error message.
+// The error is surfaced on the Go side via mlxInitError / GetMLXInitError()
+// only when MLX is actually needed.
+static void mlx_silent_error_handler(const char* msg, void* data) {
+    (void)data;
+    strncpy(mlx_init_error_msg, msg, sizeof(mlx_init_error_msg) - 1);
+    mlx_init_error_msg[sizeof(mlx_init_error_msg) - 1] = '\0';
+    mlx_init_error_flag = 1;
+}
+
+void mlx_set_safe_init_mode(void) {
+    mlx_init_error_flag = 0;
+    mlx_init_error_msg[0] = '\0';
+    mlx_set_error_handler(mlx_silent_error_handler, NULL, NULL);
+}
+
+void mlx_set_default_error_mode(void) {
+    mlx_set_error_handler(NULL, NULL, NULL);
+}
+
+int mlx_had_init_error(void) {
+    return mlx_init_error_flag;
+}
+
+const char* mlx_get_init_error(void) {
+    return mlx_init_error_flag ? mlx_init_error_msg : NULL;
+}
--- a/x/imagegen/mlx/mlx_error_handler.h
+++ b/x/imagegen/mlx/mlx_error_handler.h
@@ -0,0 +1,22 @@
+// mlx_error_handler.h - Safe error handling for MLX initialization
+// This replaces the default exit(-1) MLX error handler during init()
+// so that GPU failures don't kill the process.
+
+#ifndef MLX_ERROR_HANDLER_H
+#define MLX_ERROR_HANDLER_H
+
+// Enter safe mode before any MLX compute calls during init().
+// Replaces the default exit(-1) handler with one that silently stores errors.
+void mlx_set_safe_init_mode(void);
+
+// Restore the default MLX error handler (exit on error).
+// Call from runner entry points after confirming MLX is available.
+void mlx_set_default_error_mode(void);
+
+// Check whether an error occurred while in safe init mode.
+int mlx_had_init_error(void);
+
+// Get the error message from the last init error, or NULL if none.
+const char* mlx_get_init_error(void);
+
+#endif // MLX_ERROR_HANDLER_H
--- a/x/imagegen/mlx/mlx_test.go
+++ b/x/imagegen/mlx/mlx_test.go