diff --git a/x/mlxrunner/mlx/gated_delta.go b/x/mlxrunner/mlx/gated_delta.go index 13d034931..1e9107c4d 100644 --- a/x/mlxrunner/mlx/gated_delta.go +++ b/x/mlxrunner/mlx/gated_delta.go @@ -1,23 +1,17 @@ package mlx -// #include -// #include "generated.h" -import "C" - -import ( - "sync" - "unsafe" -) - -var ( - gatedDeltaMetalKernelOnce sync.Once - gatedDeltaMetalKernel C.mlx_fast_metal_kernel - gatedDeltaMetalDisabled bool - - gatedDeltaCUDAKernelOnce sync.Once - gatedDeltaCUDAKernel C.mlx_fast_cuda_kernel - gatedDeltaCUDADisabled bool -) +var gatedDelta = &gpuKernel{ + name: "gated_delta_step", + inputs: []string{"q", "k", "v", "g", "beta", "state_in", "T"}, + outputs: []string{"y", "state_out"}, + metal: gpuSource{source: gatedDeltaMetalKernelSource}, + cuda: gpuSource{source: gatedDeltaCUDAKernelSource}, + fallback: func(launch gpuLaunch) []*Array { + in := launch.inputs + y, state := gatedDeltaFallback(in[0], in[1], in[2], in[3], in[4], in[5]) + return []*Array{y, state} + }, +} const gatedDeltaMetalKernelSource = ` auto n = thread_position_in_grid.z; @@ -167,196 +161,49 @@ for (int i = 0; i < n_per_t; ++i) { } ` -func cStringVector(values []string) (C.mlx_vector_string, func(), bool) { - vec := C.mlx_vector_string_new() - ok := true - for _, s := range values { - cs := C.CString(s) - if C.mlx_vector_string_append_value(vec, cs) != 0 { - ok = false - } - C.free(unsafe.Pointer(cs)) - if !ok { - break - } - } - cleanup := func() { - C.mlx_vector_string_free(vec) - } - return vec, cleanup, ok +// gatedDeltaDims are the batch and head geometry of one scan, recovered +// from the input shapes. +type gatedDeltaDims struct { + B, T, Hk, Dk, Hv, Dv int } -func initGatedDeltaMetalKernel() { - inputs, freeInputs, ok := cStringVector([]string{"q", "k", "v", "g", "beta", "state_in", "T"}) - if !ok { - gatedDeltaMetalDisabled = true - freeInputs() - return - } - defer freeInputs() - - outputs, freeOutputs, ok := cStringVector([]string{"y", "state_out"}) - if !ok { - gatedDeltaMetalDisabled = true - freeOutputs() - return - } - defer freeOutputs() - - cName := C.CString("gated_delta_step") - defer C.free(unsafe.Pointer(cName)) - cSource := C.CString(gatedDeltaMetalKernelSource) - defer C.free(unsafe.Pointer(cSource)) - cHeader := C.CString("") - defer C.free(unsafe.Pointer(cHeader)) - - gatedDeltaMetalKernel = C.mlx_fast_metal_kernel_new( - cName, - inputs, - outputs, - cSource, - cHeader, - C.bool(true), - C.bool(false), - ) -} - -// gatedDeltaKernel runs a fused Metal kernel for the qwen3.5 recurrent update. -// It returns ok=false on unsupported shapes/devices or kernel setup/apply failure. -func gatedDeltaKernel(q, k, v, g, beta, state *Array) (y, nextState *Array, ok bool) { - if gatedDeltaMetalDisabled { - return nil, nil, false - } +// resolveGatedDeltaDims validates the inputs against the GPU kernels' +// contract and recovers the launch geometry. ok=false routes to the graph +// fallback: shapes that disagree, Dk not a multiple of the 32-lane simd +// width, or mixed input dtypes. +func resolveGatedDeltaDims(q, k, v, g, beta, state *Array) (gatedDeltaDims, bool) { + var dims gatedDeltaDims if q == nil || k == nil || v == nil || g == nil || beta == nil || state == nil { - return nil, nil, false + return dims, false } - - qd := q.Dims() - kd := k.Dims() - vd := v.Dims() - gd := g.Dims() - bd := beta.Dims() - sd := state.Dims() + qd, kd, vd, gd, bd, sd := q.Dims(), k.Dims(), v.Dims(), g.Dims(), beta.Dims(), state.Dims() if len(qd) != 4 || len(kd) != 4 || len(vd) != 4 || len(gd) != 3 || len(bd) != 3 || len(sd) != 4 { - return nil, nil, false + return dims, false } - - B, T, Hk, Dk := qd[0], qd[1], qd[2], qd[3] - if T <= 0 || Hk <= 0 || Dk <= 0 || Dk%32 != 0 { - return nil, nil, false + dims.B, dims.T, dims.Hk, dims.Dk = qd[0], qd[1], qd[2], qd[3] + if dims.T <= 0 || dims.Hk <= 0 || dims.Dk <= 0 || dims.Dk%32 != 0 { + return dims, false } - if kd[0] != B || kd[1] != T || kd[2] != Hk || kd[3] != Dk { - return nil, nil, false + if kd[0] != dims.B || kd[1] != dims.T || kd[2] != dims.Hk || kd[3] != dims.Dk { + return dims, false } - Hv, Dv := vd[2], vd[3] - if vd[0] != B || vd[1] != T || Hv <= 0 || Dv <= 0 || Hv%Hk != 0 { - return nil, nil, false + dims.Hv, dims.Dv = vd[2], vd[3] + if vd[0] != dims.B || vd[1] != dims.T || dims.Hv <= 0 || dims.Dv <= 0 || dims.Hv%dims.Hk != 0 { + return dims, false } - if gd[0] != B || gd[1] != T || gd[2] != Hv { - return nil, nil, false + if gd[0] != dims.B || gd[1] != dims.T || gd[2] != dims.Hv { + return dims, false } - if bd[0] != B || bd[1] != T || bd[2] != Hv { - return nil, nil, false + if bd[0] != dims.B || bd[1] != dims.T || bd[2] != dims.Hv { + return dims, false } - if sd[0] != B || sd[1] != Hv || sd[2] != Dv || sd[3] != Dk { - return nil, nil, false + if sd[0] != dims.B || sd[1] != dims.Hv || sd[2] != dims.Dv || sd[3] != dims.Dk { + return dims, false } - - inputDType := q.DType() - stateDType := state.DType() - if k.DType() != inputDType || v.DType() != inputDType || g.DType() != inputDType || beta.DType() != inputDType { - return nil, nil, false + if k.DType() != q.DType() || v.DType() != q.DType() || g.DType() != q.DType() || beta.DType() != q.DType() { + return dims, false } - - gatedDeltaMetalKernelOnce.Do(initGatedDeltaMetalKernel) - if gatedDeltaMetalDisabled { - return nil, nil, false - } - - cfg := C.mlx_fast_metal_kernel_config_new() - defer C.mlx_fast_metal_kernel_config_free(cfg) - - cInT := C.CString("InT") - defer C.free(unsafe.Pointer(cInT)) - if C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, cInT, C.mlx_dtype(inputDType)) != 0 { - gatedDeltaMetalDisabled = true - return nil, nil, false - } - cStT := C.CString("StT") - defer C.free(unsafe.Pointer(cStT)) - if C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, cStT, C.mlx_dtype(stateDType)) != 0 { - gatedDeltaMetalDisabled = true - return nil, nil, false - } - for _, tpl := range []struct { - name string - value int - }{ - {name: "Dk", value: Dk}, - {name: "Dv", value: Dv}, - {name: "Hk", value: Hk}, - {name: "Hv", value: Hv}, - } { - cn := C.CString(tpl.name) - rc := C.mlx_fast_metal_kernel_config_add_template_arg_int(cfg, cn, C.int(tpl.value)) - C.free(unsafe.Pointer(cn)) - if rc != 0 { - gatedDeltaMetalDisabled = true - return nil, nil, false - } - } - - yShape := []C.int{C.int(B), C.int(T), C.int(Hv), C.int(Dv)} - stateShape := []C.int{C.int(B), C.int(Hv), C.int(Dv), C.int(Dk)} - if C.mlx_fast_metal_kernel_config_add_output_arg(cfg, unsafe.SliceData(yShape), C.size_t(len(yShape)), C.mlx_dtype(inputDType)) != 0 { - gatedDeltaMetalDisabled = true - return nil, nil, false - } - if C.mlx_fast_metal_kernel_config_add_output_arg(cfg, unsafe.SliceData(stateShape), C.size_t(len(stateShape)), C.mlx_dtype(stateDType)) != 0 { - gatedDeltaMetalDisabled = true - return nil, nil, false - } - if C.mlx_fast_metal_kernel_config_set_grid(cfg, 32, C.int(Dv), C.int(B*Hv)) != 0 { - gatedDeltaMetalDisabled = true - return nil, nil, false - } - threadY := Dv - if threadY > 4 { - threadY = 4 - } - if C.mlx_fast_metal_kernel_config_set_thread_group(cfg, 32, C.int(threadY), 1) != 0 { - gatedDeltaMetalDisabled = true - return nil, nil, false - } - - tScalar := FromValue(T) - inputs := []C.mlx_array{ - q.ctx, - k.ctx, - v.ctx, - g.ctx, - beta.ctx, - state.ctx, - tScalar.ctx, - } - inVec := C.mlx_vector_array_new_data(unsafe.SliceData(inputs), C.size_t(len(inputs))) - defer C.mlx_vector_array_free(inVec) - - outVec := C.mlx_vector_array_new() - defer C.mlx_vector_array_free(outVec) - if C.mlx_fast_metal_kernel_apply(&outVec, gatedDeltaMetalKernel, inVec, cfg, DefaultStream().ctx) != 0 { - gatedDeltaMetalDisabled = true - return nil, nil, false - } - if int(C.mlx_vector_array_size(outVec)) < 2 { - return nil, nil, false - } - - y = New("GATED_DELTA_METAL_Y") - nextState = New("GATED_DELTA_METAL_STATE") - C.mlx_vector_array_get(&y.ctx, outVec, 0) - C.mlx_vector_array_get(&nextState.ctx, outVec, 1) - return y, nextState, true + return dims, true } func repeatHeadsForGatedDelta(x *Array, repeatFactor int) *Array { @@ -443,183 +290,6 @@ func gatedDeltaFallback(q, k, v, g, beta, state *Array) (y, nextState *Array) { return Concatenate(outs, 1), nextState } -func initGatedDeltaCUDAKernel() { - var cudaAvail C.bool - if C.mlx_cuda_is_available(&cudaAvail) != 0 || !bool(cudaAvail) { - gatedDeltaCUDADisabled = true - return - } - - inputs, freeInputs, ok := cStringVector([]string{"q", "k", "v", "g", "beta", "state_in", "T"}) - if !ok { - gatedDeltaCUDADisabled = true - freeInputs() - return - } - defer freeInputs() - - outputs, freeOutputs, ok := cStringVector([]string{"y", "state_out"}) - if !ok { - gatedDeltaCUDADisabled = true - freeOutputs() - return - } - defer freeOutputs() - - cName := C.CString("gated_delta_step") - defer C.free(unsafe.Pointer(cName)) - cSource := C.CString(gatedDeltaCUDAKernelSource) - defer C.free(unsafe.Pointer(cSource)) - cHeader := C.CString("") - defer C.free(unsafe.Pointer(cHeader)) - - gatedDeltaCUDAKernel = C.mlx_fast_cuda_kernel_new( - cName, - inputs, - outputs, - cSource, - cHeader, - C.bool(true), - C.int(0), - ) -} - -func gatedDeltaCUDAKernelApply(q, k, v, g, beta, state *Array) (y, nextState *Array, ok bool) { - if gatedDeltaCUDADisabled { - return nil, nil, false - } - if q == nil || k == nil || v == nil || g == nil || beta == nil || state == nil { - return nil, nil, false - } - - qd := q.Dims() - kd := k.Dims() - vd := v.Dims() - gd := g.Dims() - bd := beta.Dims() - sd := state.Dims() - if len(qd) != 4 || len(kd) != 4 || len(vd) != 4 || len(gd) != 3 || len(bd) != 3 || len(sd) != 4 { - return nil, nil, false - } - - B, T, Hk, Dk := qd[0], qd[1], qd[2], qd[3] - if T <= 0 || Hk <= 0 || Dk <= 0 || Dk%32 != 0 { - return nil, nil, false - } - if kd[0] != B || kd[1] != T || kd[2] != Hk || kd[3] != Dk { - return nil, nil, false - } - Hv, Dv := vd[2], vd[3] - if vd[0] != B || vd[1] != T || Hv <= 0 || Dv <= 0 || Hv%Hk != 0 { - return nil, nil, false - } - if gd[0] != B || gd[1] != T || gd[2] != Hv { - return nil, nil, false - } - if bd[0] != B || bd[1] != T || bd[2] != Hv { - return nil, nil, false - } - if sd[0] != B || sd[1] != Hv || sd[2] != Dv || sd[3] != Dk { - return nil, nil, false - } - - inputDType := q.DType() - stateDType := state.DType() - if k.DType() != inputDType || v.DType() != inputDType || g.DType() != inputDType || beta.DType() != inputDType { - return nil, nil, false - } - - gatedDeltaCUDAKernelOnce.Do(initGatedDeltaCUDAKernel) - if gatedDeltaCUDADisabled { - return nil, nil, false - } - - cfg := C.mlx_fast_cuda_kernel_config_new() - defer C.mlx_fast_cuda_kernel_config_free(cfg) - - cInT := C.CString("InT") - defer C.free(unsafe.Pointer(cInT)) - if C.mlx_fast_cuda_kernel_config_add_template_arg_dtype(cfg, cInT, C.mlx_dtype(inputDType)) != 0 { - gatedDeltaCUDADisabled = true - return nil, nil, false - } - cStT := C.CString("StT") - defer C.free(unsafe.Pointer(cStT)) - if C.mlx_fast_cuda_kernel_config_add_template_arg_dtype(cfg, cStT, C.mlx_dtype(stateDType)) != 0 { - gatedDeltaCUDADisabled = true - return nil, nil, false - } - for _, tpl := range []struct { - name string - value int - }{ - {name: "Dk", value: Dk}, - {name: "Dv", value: Dv}, - {name: "Hk", value: Hk}, - {name: "Hv", value: Hv}, - } { - cn := C.CString(tpl.name) - rc := C.mlx_fast_cuda_kernel_config_add_template_arg_int(cfg, cn, C.int(tpl.value)) - C.free(unsafe.Pointer(cn)) - if rc != 0 { - gatedDeltaCUDADisabled = true - return nil, nil, false - } - } - - yShape := []C.int{C.int(B), C.int(T), C.int(Hv), C.int(Dv)} - stateShape := []C.int{C.int(B), C.int(Hv), C.int(Dv), C.int(Dk)} - if C.mlx_fast_cuda_kernel_config_add_output_arg(cfg, unsafe.SliceData(yShape), C.size_t(len(yShape)), C.mlx_dtype(inputDType)) != 0 { - gatedDeltaCUDADisabled = true - return nil, nil, false - } - if C.mlx_fast_cuda_kernel_config_add_output_arg(cfg, unsafe.SliceData(stateShape), C.size_t(len(stateShape)), C.mlx_dtype(stateDType)) != 0 { - gatedDeltaCUDADisabled = true - return nil, nil, false - } - if C.mlx_fast_cuda_kernel_config_set_grid(cfg, 32, C.int(Dv), C.int(B*Hv)) != 0 { - gatedDeltaCUDADisabled = true - return nil, nil, false - } - threadY := Dv - if threadY > 4 { - threadY = 4 - } - if C.mlx_fast_cuda_kernel_config_set_thread_group(cfg, 32, C.int(threadY), 1) != 0 { - gatedDeltaCUDADisabled = true - return nil, nil, false - } - - tScalar := FromValue(T) - inputs := []C.mlx_array{ - q.ctx, - k.ctx, - v.ctx, - g.ctx, - beta.ctx, - state.ctx, - tScalar.ctx, - } - inVec := C.mlx_vector_array_new_data(unsafe.SliceData(inputs), C.size_t(len(inputs))) - defer C.mlx_vector_array_free(inVec) - - outVec := C.mlx_vector_array_new() - defer C.mlx_vector_array_free(outVec) - if C.mlx_fast_cuda_kernel_apply(&outVec, gatedDeltaCUDAKernel, inVec, cfg, DefaultStream().ctx) != 0 { - gatedDeltaCUDADisabled = true - return nil, nil, false - } - if int(C.mlx_vector_array_size(outVec)) < 2 { - return nil, nil, false - } - - y = New("GATED_DELTA_CUDA_Y") - nextState = New("GATED_DELTA_CUDA_STATE") - C.mlx_vector_array_get(&y.ctx, outVec, 0) - C.mlx_vector_array_get(&nextState.ctx, outVec, 1) - return y, nextState, true -} - // FastGatedDelta runs the recurrent update operation. // // When mask is non-nil, it must be a [B, T] bool tensor identifying real @@ -628,8 +298,9 @@ func gatedDeltaCUDAKernelApply(q, k, v, g, beta, state *Array) (y, nextState *Ar // kernel iteration is a no-op — state passes through unchanged and the // final state equals the state after the last real token of each row. // -// It tries the fused CUDA kernel first, then Metal, then falls back to a -// backend-agnostic MLX implementation with identical inputs/outputs. +// Inputs that fit the GPU kernels' contract run there (CUDA or Metal, with +// the graph implementation covering boxes where neither can run); anything +// else runs the graph implementation directly. func FastGatedDelta(q, k, v, g, beta, state, mask *Array) (y, nextState *Array) { // TODO: handle this more efficiently with a masked kernel (MLX-LM has one). if mask != nil { @@ -649,12 +320,21 @@ func FastGatedDelta(q, k, v, g, beta, state, mask *Array) (y, nextState *Array) g = Where(m3, g, oneG) } - if y, nextState, ok := gatedDeltaCUDAKernelApply(q, k, v, g, beta, state); ok { - return y, nextState - } - if y, nextState, ok := gatedDeltaKernel(q, k, v, g, beta, state); ok { - return y, nextState + if dims, ok := resolveGatedDeltaDims(q, k, v, g, beta, state); ok { + outs := gatedDelta.run(gpuLaunch{ + dtypes: []gpuDTypeArg{{"InT", q.DType()}, {"StT", state.DType()}}, + ints: []gpuIntArg{{"Dk", dims.Dk}, {"Dv", dims.Dv}, {"Hk", dims.Hk}, {"Hv", dims.Hv}}, + outputs: []gpuOutputSpec{ + {"GATED_DELTA_Y", []int32{int32(dims.B), int32(dims.T), int32(dims.Hv), int32(dims.Dv)}, q.DType()}, + {"GATED_DELTA_STATE", []int32{int32(dims.B), int32(dims.Hv), int32(dims.Dv), int32(dims.Dk)}, state.DType()}, + }, + grid: [3]int{32, dims.Dv, dims.B * dims.Hv}, + threadGroup: [3]int{32, min(dims.Dv, 4), 1}, + inputs: []*Array{q, k, v, g, beta, state, FromValue(dims.T)}, + }) + return outs[0], outs[1] } + y, nextState = gatedDeltaFallback(q, k, v, g, beta, state) if y == nil || nextState == nil { panic("mlx.FastGatedDelta: fallback failed (invalid inputs or unsupported shapes)") diff --git a/x/mlxrunner/mlx/gpu_kernel.go b/x/mlxrunner/mlx/gpu_kernel.go new file mode 100644 index 000000000..b8d055384 --- /dev/null +++ b/x/mlxrunner/mlx/gpu_kernel.go @@ -0,0 +1,359 @@ +package mlx + +// #include +// #include "generated.h" +import "C" + +import ( + "log/slog" + "sync" + "unsafe" +) + +// gpuSource is one backend's implementation of a kernel. +type gpuSource struct { + source string + header string +} + +// gpuKernel is a custom kernel with per-backend sources and a graph +// fallback. Either backend may be absent; a backend that cannot be created +// or launched disables itself permanently. Contract checks belong to the +// caller, before run: run itself cannot fail. +type gpuKernel struct { + name string + inputs []string + outputs []string + metal gpuSource + cuda gpuSource + + // fallback computes the same outputs with graph ops when no GPU + // backend can run the launch. + fallback func(launch gpuLaunch) []*Array + + metalOnce sync.Once + metalKernel C.mlx_fast_metal_kernel + metalDisabled bool + + cudaOnce sync.Once + cudaKernel C.mlx_fast_cuda_kernel + cudaDisabled bool +} + +// gpuDTypeArg and gpuIntArg name template arguments for one launch. +type gpuDTypeArg struct { + name string + dtype DType +} + +type gpuIntArg struct { + name string + value int +} + +// gpuOutputSpec declares one kernel output buffer. +type gpuOutputSpec struct { + name string + shape []int32 + dtype DType +} + +// gpuLaunch is the per-call configuration for gpuKernel.run. Grid and +// thread-group units are shared across backends. +type gpuLaunch struct { + dtypes []gpuDTypeArg + ints []gpuIntArg + outputs []gpuOutputSpec + grid [3]int + threadGroup [3]int + inputs []*Array +} + +func cStringVector(values []string) (C.mlx_vector_string, func(), bool) { + vec := C.mlx_vector_string_new() + ok := true + for _, s := range values { + cs := C.CString(s) + if C.mlx_vector_string_append_value(vec, cs) != 0 { + ok = false + } + C.free(unsafe.Pointer(cs)) + if !ok { + break + } + } + cleanup := func() { + C.mlx_vector_string_free(vec) + } + return vec, cleanup, ok +} + +// run executes the kernel with the first backend that works, in CUDA, +// Metal, fallback order. It panics if no variant can run the launch. +func (k *gpuKernel) run(launch gpuLaunch) []*Array { + if outs, ok := k.applyCUDA(launch); ok { + return outs + } + if outs, ok := k.applyMetal(launch); ok { + return outs + } + if k.fallback == nil { + panic("mlx: kernel " + k.name + " has no usable implementation") + } + outs := k.fallback(launch) + if len(outs) != len(k.outputs) { + panic("mlx: kernel " + k.name + " fallback returned wrong output count") + } + return outs +} + +func (k *gpuKernel) disableMetal(reason string) { + k.metalDisabled = true + slog.Warn("custom GPU kernel backend disabled", "kernel", k.name, "backend", "metal", "reason", reason) +} + +func (k *gpuKernel) disableCUDA(reason string) { + k.cudaDisabled = true + slog.Warn("custom GPU kernel backend disabled", "kernel", k.name, "backend", "cuda", "reason", reason) +} + +func (k *gpuKernel) getMetal() (C.mlx_fast_metal_kernel, bool) { + k.metalOnce.Do(func() { + if !MetalIsAvailable() { + k.metalDisabled = true + return + } + + if k.metal.source == "" { + k.disableMetal("no source") + return + } + + inputs, freeInputs, ok := cStringVector(k.inputs) + if !ok { + freeInputs() + k.disableMetal("creating input names failed") + return + } + defer freeInputs() + + outputs, freeOutputs, ok := cStringVector(k.outputs) + if !ok { + freeOutputs() + k.disableMetal("creating output names failed") + return + } + defer freeOutputs() + + cName := C.CString(k.name) + defer C.free(unsafe.Pointer(cName)) + cSource := C.CString(k.metal.source) + defer C.free(unsafe.Pointer(cSource)) + cHeader := C.CString(k.metal.header) + defer C.free(unsafe.Pointer(cHeader)) + + k.metalKernel = C.mlx_fast_metal_kernel_new( + cName, + inputs, + outputs, + cSource, + cHeader, + // ensure_row_contiguous, so kernels can index inputs linearly. + C.bool(true), + C.bool(false), + ) + if k.metalKernel.ctx == nil { + k.disableMetal("creating kernel failed") + } + }) + return k.metalKernel, !k.metalDisabled +} + +func (k *gpuKernel) applyMetal(launch gpuLaunch) ([]*Array, bool) { + if k.metalDisabled { + return nil, false + } + kernel, ok := k.getMetal() + if !ok { + return nil, false + } + + cfg := C.mlx_fast_metal_kernel_config_new() + defer C.mlx_fast_metal_kernel_config_free(cfg) + for _, arg := range launch.dtypes { + name := C.CString(arg.name) + rc := C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, name, C.mlx_dtype(arg.dtype)) + C.free(unsafe.Pointer(name)) + if rc != 0 { + k.disableMetal("setting dtype template arg failed") + return nil, false + } + } + for _, arg := range launch.ints { + name := C.CString(arg.name) + rc := C.mlx_fast_metal_kernel_config_add_template_arg_int(cfg, name, C.int(arg.value)) + C.free(unsafe.Pointer(name)) + if rc != 0 { + k.disableMetal("setting int template arg failed") + return nil, false + } + } + for _, out := range launch.outputs { + shape := make([]C.int, len(out.shape)) + for i, d := range out.shape { + shape[i] = C.int(d) + } + if C.mlx_fast_metal_kernel_config_add_output_arg(cfg, unsafe.SliceData(shape), C.size_t(len(shape)), C.mlx_dtype(out.dtype)) != 0 { + k.disableMetal("adding output failed") + return nil, false + } + } + if C.mlx_fast_metal_kernel_config_set_grid(cfg, C.int(launch.grid[0]), C.int(launch.grid[1]), C.int(launch.grid[2])) != 0 || + C.mlx_fast_metal_kernel_config_set_thread_group(cfg, C.int(launch.threadGroup[0]), C.int(launch.threadGroup[1]), C.int(launch.threadGroup[2])) != 0 { + k.disableMetal("setting grid failed") + return nil, false + } + + inputs := make([]C.mlx_array, len(launch.inputs)) + for i, in := range launch.inputs { + inputs[i] = in.ctx + } + inVec := C.mlx_vector_array_new_data(unsafe.SliceData(inputs), C.size_t(len(inputs))) + defer C.mlx_vector_array_free(inVec) + outVec := C.mlx_vector_array_new() + defer C.mlx_vector_array_free(outVec) + if C.mlx_fast_metal_kernel_apply(&outVec, kernel, inVec, cfg, DefaultStream().ctx) != 0 { + k.disableMetal("launching failed") + return nil, false + } + if int(C.mlx_vector_array_size(outVec)) < len(launch.outputs) { + return nil, false + } + + outs := make([]*Array, len(launch.outputs)) + for i, out := range launch.outputs { + outs[i] = New(out.name) + C.mlx_vector_array_get(&outs[i].ctx, outVec, C.size_t(i)) + } + return outs, true +} + +func (k *gpuKernel) getCUDA() (C.mlx_fast_cuda_kernel, bool) { + k.cudaOnce.Do(func() { + if !CUDAIsAvailable() { + k.cudaDisabled = true + return + } + + if k.cuda.source == "" { + k.disableCUDA("no source") + return + } + + inputs, freeInputs, ok := cStringVector(k.inputs) + if !ok { + freeInputs() + k.disableCUDA("creating input names failed") + return + } + defer freeInputs() + + outputs, freeOutputs, ok := cStringVector(k.outputs) + if !ok { + freeOutputs() + k.disableCUDA("creating output names failed") + return + } + defer freeOutputs() + + cName := C.CString(k.name) + defer C.free(unsafe.Pointer(cName)) + cSource := C.CString(k.cuda.source) + defer C.free(unsafe.Pointer(cSource)) + cHeader := C.CString(k.cuda.header) + defer C.free(unsafe.Pointer(cHeader)) + + k.cudaKernel = C.mlx_fast_cuda_kernel_new( + cName, + inputs, + outputs, + cSource, + cHeader, + C.bool(true), + C.int(0), + ) + if k.cudaKernel.ctx == nil { + k.disableCUDA("creating kernel failed") + } + }) + return k.cudaKernel, !k.cudaDisabled +} + +func (k *gpuKernel) applyCUDA(launch gpuLaunch) ([]*Array, bool) { + if k.cudaDisabled { + return nil, false + } + kernel, ok := k.getCUDA() + if !ok { + return nil, false + } + + cfg := C.mlx_fast_cuda_kernel_config_new() + defer C.mlx_fast_cuda_kernel_config_free(cfg) + for _, arg := range launch.dtypes { + name := C.CString(arg.name) + rc := C.mlx_fast_cuda_kernel_config_add_template_arg_dtype(cfg, name, C.mlx_dtype(arg.dtype)) + C.free(unsafe.Pointer(name)) + if rc != 0 { + k.disableCUDA("setting dtype template arg failed") + return nil, false + } + } + for _, arg := range launch.ints { + name := C.CString(arg.name) + rc := C.mlx_fast_cuda_kernel_config_add_template_arg_int(cfg, name, C.int(arg.value)) + C.free(unsafe.Pointer(name)) + if rc != 0 { + k.disableCUDA("setting int template arg failed") + return nil, false + } + } + for _, out := range launch.outputs { + shape := make([]C.int, len(out.shape)) + for i, d := range out.shape { + shape[i] = C.int(d) + } + if C.mlx_fast_cuda_kernel_config_add_output_arg(cfg, unsafe.SliceData(shape), C.size_t(len(shape)), C.mlx_dtype(out.dtype)) != 0 { + k.disableCUDA("adding output failed") + return nil, false + } + } + if C.mlx_fast_cuda_kernel_config_set_grid(cfg, C.int(launch.grid[0]), C.int(launch.grid[1]), C.int(launch.grid[2])) != 0 || + C.mlx_fast_cuda_kernel_config_set_thread_group(cfg, C.int(launch.threadGroup[0]), C.int(launch.threadGroup[1]), C.int(launch.threadGroup[2])) != 0 { + k.disableCUDA("setting grid failed") + return nil, false + } + + inputs := make([]C.mlx_array, len(launch.inputs)) + for i, in := range launch.inputs { + inputs[i] = in.ctx + } + inVec := C.mlx_vector_array_new_data(unsafe.SliceData(inputs), C.size_t(len(inputs))) + defer C.mlx_vector_array_free(inVec) + outVec := C.mlx_vector_array_new() + defer C.mlx_vector_array_free(outVec) + if C.mlx_fast_cuda_kernel_apply(&outVec, kernel, inVec, cfg, DefaultStream().ctx) != 0 { + k.disableCUDA("launching failed") + return nil, false + } + if int(C.mlx_vector_array_size(outVec)) < len(launch.outputs) { + return nil, false + } + + outs := make([]*Array, len(launch.outputs)) + for i, out := range launch.outputs { + outs[i] = New(out.name) + C.mlx_vector_array_get(&outs[i].ctx, outVec, C.size_t(i)) + } + return outs, true +} diff --git a/x/mlxrunner/mlx/mlx.go b/x/mlxrunner/mlx/mlx.go index c863b9f3e..99d579cad 100644 --- a/x/mlxrunner/mlx/mlx.go +++ b/x/mlxrunner/mlx/mlx.go @@ -115,3 +115,10 @@ func MetalIsAvailable() bool { C.mlx_metal_is_available(&available) return bool(available) } + +// CUDAIsAvailable returns true if a CUDA GPU is available. +func CUDAIsAvailable() bool { + var available C._Bool + C.mlx_cuda_is_available(&available) + return bool(available) +}