From c36adebc2095d1b3cfdf52c7ee6f1a17c662f7f4 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Wed, 2 Sep 2026 10:07:52 -0700 Subject: [PATCH] mlxrunner: capture MLX errors in a single buffer read after every call MLX runs on one goroutine locked to its OS thread, so the thread-local error buffers and closure-based check helpers defended against a calling pattern that is already invalid. Replace them with a single buffer that the handler fills and Go reads after every call. mlxError returns the captured message; mlxCheck panics on it and passes the call's result through, so a checked call is one expression. Only an int status carries a failure signal, which lets a message next to a zero status be reported as an earlier unchecked call. Fix two tests that relied on errors being dropped: the laguna mixed-precision fixture used an unsupported quantization group size, and the compile callback test expected the callback's own panic. --- x/mlxrunner/mlx/compile.go | 8 +-- x/mlxrunner/mlx/compile_test.go | 4 +- x/mlxrunner/mlx/memory.go | 20 +++--- x/mlxrunner/mlx/mlx.go | 113 +++++++++++++++++++------------- x/models/laguna/laguna_test.go | 58 ++++++++-------- 5 files changed, 109 insertions(+), 94 deletions(-) diff --git a/x/mlxrunner/mlx/compile.go b/x/mlxrunner/mlx/compile.go index 91e564ba0..847ed2d21 100644 --- a/x/mlxrunner/mlx/compile.go +++ b/x/mlxrunner/mlx/compile.go @@ -63,9 +63,7 @@ func Compile(name string, fn CompileFunc, opts ...CompileOption) CompileFunc { defer C.mlx_closure_free(src) closure = C.mlx_closure_new() - mlxCheck(name+": compile failed", func() C.int { - return C.mlx_compile(&closure, src, C.bool(cfg.shapeless)) - }) + mlxCheck(C.mlx_compile(&closure, src, C.bool(cfg.shapeless))) }) inVec := C.mlx_vector_array_new() @@ -76,9 +74,7 @@ func Compile(name string, fn CompileFunc, opts ...CompileOption) CompileFunc { outVec := C.mlx_vector_array_new() defer C.mlx_vector_array_free(outVec) - mlxCheck(name+": closure apply failed", func() C.int { - return C.mlx_closure_apply(&outVec, closure, inVec) - }) + mlxCheck(C.mlx_closure_apply(&outVec, closure, inVec)) n := int(C.mlx_vector_array_size(outVec)) outputs := make([]*Array, n) diff --git a/x/mlxrunner/mlx/compile_test.go b/x/mlxrunner/mlx/compile_test.go index 9978cb739..cd264c69f 100644 --- a/x/mlxrunner/mlx/compile_test.go +++ b/x/mlxrunner/mlx/compile_test.go @@ -124,8 +124,8 @@ func testCompileCallbackPanicRecovers(t *mlxthreadtest.T) { if r == nil { t.Fatal("expected panic from Call, got none") } - if _, ok := r.(string); !ok { - t.Fatalf("expected string panic, got %T: %v", r, r) + if _, ok := r.(error); !ok { + t.Fatalf("expected error panic, got %T: %v", r, r) } }() boom(x) diff --git a/x/mlxrunner/mlx/memory.go b/x/mlxrunner/mlx/memory.go index f75937bad..8da461498 100644 --- a/x/mlxrunner/mlx/memory.go +++ b/x/mlxrunner/mlx/memory.go @@ -72,23 +72,23 @@ func ResetPeakMemory() { // for resident Metal allocations. func MaxRecommendedWorkingSetSize() (int, error) { info := C.mlx_device_info_new() - if err := mlxCall("get device info failed", func() C.int { - return C.mlx_device_info_get(&info, DefaultDevice().ctx) - }); err != nil { - C.mlx_device_info_free(info) + defer C.mlx_device_info_free(info) + if err := mlxError(C.mlx_device_info_get(&info, DefaultDevice().ctx)); err != nil { return 0, err } - defer C.mlx_device_info_free(info) key := C.CString("max_recommended_working_set_size") defer C.free(unsafe.Pointer(key)) var size C.size_t - if err := mlxCall("max recommended working set size unavailable", func() C.int { - return C.mlx_device_info_get_size(&size, info, key) - }); err != nil { + rc := C.mlx_device_info_get_size(&size, info, key) + if err := lastError(); err != nil { return 0, err } + if rc != 0 { + // mlx-c reports a missing key with a non-zero return and no message. + return 0, fmt.Errorf("mlx: no max_recommended_working_set_size in device info") + } return int(size), nil } @@ -100,9 +100,7 @@ func SetWiredLimit(limit int) (int, error) { } var previous C.size_t - if err := mlxCall("set wired limit failed", func() C.int { - return C.mlx_set_wired_limit(&previous, C.size_t(limit)) - }); err != nil { + if err := mlxError(C.mlx_set_wired_limit(&previous, C.size_t(limit))); err != nil { return 0, err } return int(previous), nil diff --git a/x/mlxrunner/mlx/mlx.go b/x/mlxrunner/mlx/mlx.go index 99d579cad..71ce4dae7 100644 --- a/x/mlxrunner/mlx/mlx.go +++ b/x/mlxrunner/mlx/mlx.go @@ -1,3 +1,8 @@ +// Package mlx wraps the MLX C API. +// +// MLX keeps stream and backend state in thread-locals, so all calls into this +// package must come from a single goroutine locked to its OS thread (see +// x/internal/mlxthread). package mlx //go:generate go run generator/main.go -output=. ./include/mlx/c/*.h @@ -9,35 +14,27 @@ package mlx // #include "generated.h" // #include // -// static __thread char _mlx_last_error_msg[1024] = {0}; -// static __thread int _mlx_last_error_flag = 0; +// static char _mlx_last_error[1024]; // -// static void _mlx_capture_error_handler(const char* msg, void* data) { +// static void _mlx_capture_error(const char* msg, void* data) { // (void)data; -// strncpy(_mlx_last_error_msg, msg, sizeof(_mlx_last_error_msg) - 1); -// _mlx_last_error_msg[sizeof(_mlx_last_error_msg) - 1] = '\0'; -// _mlx_last_error_flag = 1; +// strncpy(_mlx_last_error, msg, sizeof(_mlx_last_error) - 1); // } // // static void mlx_install_capture_handler(void) { // if (mlx_set_error_handler_) { -// mlx_set_error_handler_(_mlx_capture_error_handler, NULL, NULL); +// mlx_set_error_handler_(_mlx_capture_error, NULL, NULL); // } // } // -// static void mlx_clear_last_error(void) { -// _mlx_last_error_flag = 0; -// _mlx_last_error_msg[0] = '\0'; -// } -// -// static const char* mlx_get_last_error(void) { -// return _mlx_last_error_flag ? _mlx_last_error_msg : ""; +// static char* mlx_last_error(void) { +// return _mlx_last_error; // } import "C" import ( + "errors" "fmt" - "runtime" ) func init() { @@ -46,6 +43,56 @@ func init() { C.mlx_install_capture_handler() } +var errBuf = C.mlx_last_error() + +// lastError consumes the captured MLX error, or returns nil when none is +// pending. +func lastError() error { + if *errBuf == 0 { + return nil + } + err := fmt.Errorf("mlx: %s", C.GoString(errBuf)) + *errBuf = 0 + return err +} + +// mlxError returns the MLX error captured by the call that produced v. mlx-c +// signals failure with a non-zero int status; a message next to a zero +// status came from an earlier unchecked call. +func mlxError[T comparable](v T) error { + var zero T + var failed, signaled bool + switch any(zero).(type) { + case C.int: + failed, signaled = v != zero, true + default: + // Only an int status signals failure. Handles, pointers, sizes, and + // dtypes are all valid at zero: a null handle is what the out-param + // constructors return, and an empty array has no data. + } + if *errBuf != 0 { + err := lastError() + if signaled && !failed { + return fmt.Errorf("mlx: unchecked error from an earlier call: %w", err) + } + return err + } + if failed { + return errors.New("mlx: call failed without an error message") + } + return nil +} + +// mlxCheck panics on a failed call and otherwise passes its result through. +// Most array operations cannot recover from a failed graph construction or +// evaluation. +func mlxCheck[T comparable](v T) T { + if err := mlxError(v); err != nil { + panic(err) + } + return v +} + // Version returns the MLX core library version string. func Version() string { str := C.mlx_string_new() @@ -54,31 +101,6 @@ func Version() string { return C.GoString(C.mlx_string_data(str)) } -// mlxCall locks the goroutine to its OS thread so the thread-local error state -// is read from the same thread that executed fn. -func mlxCall(fallback string, fn func() C.int) error { - runtime.LockOSThread() - defer runtime.UnlockOSThread() - - C.mlx_clear_last_error() - if fn() != 0 { - msg := C.GoString(C.mlx_get_last_error()) - if msg == "" { - msg = fallback - } - return fmt.Errorf("mlx: %s", msg) - } - return nil -} - -// mlxCheck panics with the captured MLX error. Most array operations cannot -// recover from a failed graph construction or evaluation. -func mlxCheck(fallback string, fn func() C.int) { - if err := mlxCall(fallback, fn); err != nil { - panic(err.Error()) - } -} - func doEval(outputs []*Array, async bool) { if len(outputs) == 0 { return @@ -93,12 +115,11 @@ func doEval(outputs []*Array, async bool) { } } - mlxCheck("eval failed", func() C.int { - if async { - return C.mlx_async_eval(vector) - } - return C.mlx_eval(vector) - }) + if async { + mlxCheck(C.mlx_async_eval(vector)) + } else { + mlxCheck(C.mlx_eval(vector)) + } } func AsyncEval(outputs ...*Array) { diff --git a/x/models/laguna/laguna_test.go b/x/models/laguna/laguna_test.go index ba6e61457..4310a1430 100644 --- a/x/models/laguna/laguna_test.go +++ b/x/models/laguna/laguna_test.go @@ -197,7 +197,7 @@ func TestTinyLagunaLoadAndForward(t *testing.T) { mlxtest.Run(t, func(t *mlxtest.T) { cfg, err := parseConfig([]byte(`{ "model_type": "laguna", - "hidden_size": 8, + "hidden_size": 32, "intermediate_size": 12, "moe_intermediate_size": 4, "shared_expert_intermediate_size": 4, @@ -264,8 +264,8 @@ func TestTinyLagunaLoadAndForward(t *testing.T) { SeqQueryLens: []int32{int32(tokens.Dim(1))}, }, caches) mlx.Eval(hidden) - if got := hidden.Dims(); len(got) != 3 || got[0] != 1 || got[1] != 3 || got[2] != 8 { - t.Fatalf("hidden shape = %v, want [1 3 8]", got) + if got := hidden.Dims(); len(got) != 3 || got[0] != 1 || got[1] != 3 || got[2] != 32 { + t.Fatalf("hidden shape = %v, want [1 3 32]", got) } logits := m.Unembed(hidden) @@ -285,7 +285,7 @@ func TestTinyLagunaLoadWeightsFusesDenseGateUp(t *testing.T) { mlxtest.Run(t, func(t *mlxtest.T) { cfg, err := parseConfig([]byte(`{ "model_type": "laguna", - "hidden_size": 8, + "hidden_size": 32, "intermediate_size": 12, "moe_intermediate_size": 4, "shared_expert_intermediate_size": 4, @@ -329,7 +329,7 @@ func TestTinyLagunaLoadWeightsFusesDenseGateUp(t *testing.T) { if moe.SwitchMLP.GateUpWeight == nil { t.Fatal("expected fused GateUpWeight to be populated") } - if got, want := moe.SwitchMLP.GateUpWeight.Dims(), []int{2, 8, 8}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] { + if got, want := moe.SwitchMLP.GateUpWeight.Dims(), []int{2, 32, 8}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] { t.Fatalf("GateUpWeight dims = %v, want %v", got, want) } }) @@ -339,7 +339,7 @@ func TestTinyLagunaLoadWeightsKeepsBF16SourceLayout(t *testing.T) { mlxtest.Run(t, func(t *mlxtest.T) { cfg, err := parseConfig([]byte(`{ "model_type": "laguna", - "hidden_size": 8, + "hidden_size": 32, "intermediate_size": 12, "moe_intermediate_size": 4, "shared_expert_intermediate_size": 4, @@ -392,7 +392,7 @@ func TestTinyLagunaLoadWeightsKeepsBF16SourceLayout(t *testing.T) { if moe.SwitchMLP.GateUpWeight != nil { t.Fatal("expected BF16 source-layout SwitchMLP to avoid pre-fused gate/up weights") } - if got, want := moe.SwitchMLP.GateWeight.Dims(), []int{2, 4, 8}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] { + if got, want := moe.SwitchMLP.GateWeight.Dims(), []int{2, 4, 32}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] { t.Fatalf("GateWeight dims = %v, want %v", got, want) } }) @@ -401,7 +401,7 @@ func TestTinyLagunaLoadWeightsKeepsBF16SourceLayout(t *testing.T) { func TestTinyLagunaLoadWeightsKeepsMixedExpertPrecision(t *testing.T) { mlxtest.Run(t, func(t *mlxtest.T) { cfg := &Config{ - HiddenSize: 8, + HiddenSize: 32, IntermediateSize: 12, MoeIntermediateSize: 4, SharedExpertIntermediate: 4, @@ -418,7 +418,7 @@ func TestTinyLagunaLoadWeightsKeepsMixedExpertPrecision(t *testing.T) { NumExpertsPerTok: 1, MoeRoutedScalingFactor: 2.5, RMSNormEps: 1e-5, - QuantGroupSize: 4, + QuantGroupSize: 32, QuantBits: 4, QuantMode: "affine", } @@ -739,38 +739,38 @@ func TestCombinedTensorGlobalScaleIgnoresInputGlobalScale(t *testing.T) { func tinyLagunaTensors() map[string]*mlx.Array { tensors := map[string]*mlx.Array{ - "model.embed_tokens.weight": weights(16, 8), - "model.norm.weight": ones(8), - "lm_head.weight": weights(16, 8), + "model.embed_tokens.weight": weights(16, 32), + "model.norm.weight": ones(32), + "lm_head.weight": weights(16, 32), } for layer := range 2 { prefix := "model.layers." + string(rune('0'+layer)) - tensors[prefix+".input_layernorm.weight"] = ones(8) - tensors[prefix+".post_attention_layernorm.weight"] = ones(8) - tensors[prefix+".self_attn.q_proj.weight"] = weights(8, 8) - tensors[prefix+".self_attn.k_proj.weight"] = weights(4, 8) - tensors[prefix+".self_attn.v_proj.weight"] = weights(4, 8) - tensors[prefix+".self_attn.o_proj.weight"] = weights(8, 8) - tensors[prefix+".self_attn.g_proj.weight"] = weights(2, 8) + tensors[prefix+".input_layernorm.weight"] = ones(32) + tensors[prefix+".post_attention_layernorm.weight"] = ones(32) + tensors[prefix+".self_attn.q_proj.weight"] = weights(8, 32) + tensors[prefix+".self_attn.k_proj.weight"] = weights(4, 32) + tensors[prefix+".self_attn.v_proj.weight"] = weights(4, 32) + tensors[prefix+".self_attn.o_proj.weight"] = weights(32, 8) + tensors[prefix+".self_attn.g_proj.weight"] = weights(2, 32) tensors[prefix+".self_attn.q_norm.weight"] = ones(4) tensors[prefix+".self_attn.k_norm.weight"] = ones(4) } - tensors["model.layers.0.mlp.gate_proj.weight"] = weights(12, 8) - tensors["model.layers.0.mlp.up_proj.weight"] = weights(12, 8) - tensors["model.layers.0.mlp.down_proj.weight"] = weights(8, 12) + tensors["model.layers.0.mlp.gate_proj.weight"] = weights(12, 32) + tensors["model.layers.0.mlp.up_proj.weight"] = weights(12, 32) + tensors["model.layers.0.mlp.down_proj.weight"] = weights(32, 12) - tensors["model.layers.1.mlp.gate.weight"] = weights(2, 8) + tensors["model.layers.1.mlp.gate.weight"] = weights(2, 32) tensors["model.layers.1.mlp.experts.e_score_correction_bias"] = mlx.FromValues([]float32{0.1, -0.1}, 2) for expert := range 2 { prefix := "model.layers.1.mlp.experts." + string(rune('0'+expert)) - tensors[prefix+".gate_proj.weight"] = weights(4, 8) - tensors[prefix+".up_proj.weight"] = weights(4, 8) - tensors[prefix+".down_proj.weight"] = weights(8, 4) + tensors[prefix+".gate_proj.weight"] = weights(4, 32) + tensors[prefix+".up_proj.weight"] = weights(4, 32) + tensors[prefix+".down_proj.weight"] = weights(32, 4) } - tensors["model.layers.1.mlp.shared_expert.gate_proj.weight"] = weights(4, 8) - tensors["model.layers.1.mlp.shared_expert.up_proj.weight"] = weights(4, 8) - tensors["model.layers.1.mlp.shared_expert.down_proj.weight"] = weights(8, 4) + tensors["model.layers.1.mlp.shared_expert.gate_proj.weight"] = weights(4, 32) + tensors["model.layers.1.mlp.shared_expert.up_proj.weight"] = weights(4, 32) + tensors["model.layers.1.mlp.shared_expert.down_proj.weight"] = weights(32, 4) return tensors }