* mlx: port the Metal custom kernels to CUDA
mamba2_scan, depthwise_conv_silu and gated_delta (plus their _states/_bias
variants) had a Metal source only, so on CUDA each fell back to graph ops.
Deviations from the Metal sources:
- The scan returns early past D: Metal dispatches an exact thread count,
CUDA rounds the grid up to whole blocks.
- T is dereferenced; CUDA passes inputs by address, Metal scalarizes
zero-dimensional ones.
- simd_sum becomes a shfl butterfly, and the threadgroup q/k caches become
per-warp registers: each warp normalizes the row its own lanes consume,
costing a redundant norm but no shared memory and no barriers.
- The gates and SiLU call MLX's Sigmoid/LogAddExp/Exp instead of restating
their formulas. Metal cannot: its custom-kernel preamble is metal::utils()
alone, while CUDA embeds the op headers for nvrtc.
The parity tests now run on Metal and CUDA hosts. The fused gated-delta
comparison takes a tolerance: the kernel and the graph path reduce in a
different order, so the bf16 output can land one ULP away.
* address comments
The MLX runner is the only Go inference runner left and is no longer
experimental, so its packages leave x/. The bindings become a top-level
mlx package beside the carried patches in mlx/compat, mirroring how
llama/ holds the llama.cpp integration, and the runner becomes mlxrunner
with the architectures nested under the package they implement.
Subpackages move with their parent unless listed.
x/mlxrunner/mlx mlx
x/internal/mlxthread mlx/mlxthread
x/internal/mlxthreadtest mlx/mlxthread/mlxthreadtest
x/internal/mlxtest mlx/mlxtest
x/quant mlx/quant
mlx/compat/*.patch mlx/compat/mlx-c (MLX patches go in mlx/compat/mlx)
x/mlxrunner mlxrunner
x/models/nn mlxrunner/nn
x/models/<arch> mlxrunner/model/<arch>
x/mlxrunner/imports.go mlxrunner/model/architectures (new package)
x/create create
x/safetensors fs/safetensors
x/tokenizer mlxrunner/tokenizer
Every package keeps its name, so the Go changes are the import path
rewrites the moves force, and the CMake, Dockerfile, CI cache keys, drift
check and Darwin payload script follow the new paths. Four edits are not
paths: the runner's blank architecture imports become the package
mlxrunner/model/architectures, so the list to extend for a new model sits
beside the architecture directories; a depguard rule keeps the two test
harnesses out of non-test code, as the x/internal placement used to; the
CI change filter's two entries for the long-deleted x/imagegen/mlx now
name the bindings' CMake project and the carried patches, so a change to
either builds the payload; and the tokenizer parity test reads its
fixtures from its own testdata instead of walking out of x/.
x/server and x/imagegen/manifest stay for the next two commits.
Daniel HiltgenandGitHub
Jesse Gross