mirror of
https://github.com/ollama/ollama.git
synced 2026-10-21 19:11:17 +00:00
mlx: factor custom GPU kernel scaffolding into helpers
Each custom kernel repeated the same host-side creation and launch boilerplate plus a CUDA-then-Metal-then-graph dispatch at every call site. gpuKernel declares the sources (either backend may be absent) and a graph fallback; run executes the first that works.
This commit is contained in:
Jesse Gross
parent
bd3f22e2f7
commit
accd6d656a
3 changed files with 425 additions and 379 deletions
|
|
@ -1,23 +1,17 @@
|
|||
package mlx
|
||||
|
||||
// #include <stdlib.h>
|
||||
// #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)")
|
||||
|
|
|
|||
359
x/mlxrunner/mlx/gpu_kernel.go
Normal file
359
x/mlxrunner/mlx/gpu_kernel.go
Normal file
|
|
@ -0,0 +1,359 @@
|
|||
package mlx
|
||||
|
||||
// #include <stdlib.h>
|
||||
// #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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue