[michaeltrefry]

Story

sc-8529 — A. mlx-rs — expose mx.fast.metal_kernel JIT binding (Phase 0 enabler). Epic sc-8528.

Validation finding: fresh wrapper, not a new C binding, not a cherry-pick

• No existing Rust binding: searched mlx-rs/mlx-sys for metal_kernel/MetalKernel/fast_metal — none.
• C API already present: the pinned mlx-c submodule exposes the full mlx_fast_metal_kernel* family in mlx/c/fast.h (config new/free, add_output_arg, set_grid, set_thread_group, set_init_value, set_verbose, add_template_arg_{dtype,int,bool}, kernel new/free/apply).
• Already surfaced by mlx-sys: fast.h is included by mlx.h, and mlx-sys runs bindgen with no allowlist, so all mlx_fast_metal_kernel* symbols are already generated. So this is a fresh safe wrapper over existing C symbols — no new bindgen/C work, and no upstream cherry-pick needed (upstream mlx-rs has no safe wrapper either).

What changed

Added a safe MetalKernel API to mlx-rs/src/fast.rs, mirroring the existing fast:: wrapping conventions (Array in/out, try_from_op/VectorArray guards, errors via the global mlx error handler):

• MetalKernel::new(name, input_names, output_names, source) and with_options(.., header, ensure_row_contiguous, atomic_outputs) — builds the kernel; MLX JIT-compiles + caches it on first dispatch.
• MetalKernel::apply() → MetalKernelDispatch builder: .input()/.inputs(), .output()/.output_shape(shape, dtype), .grid(x,y,z), .thread_group(x,y,z), .template_arg(name, TemplateArg::{Dtype,Int,Bool}), .init_value(), .verbose(), then .run() (default stream) / .run_device(stream) → Result<Vec<Array>>.
• All C handles (mlx_vector_string, config, kernel) are freed; null kernel ctx is turned into an Err.

This is ergonomic enough to later port _scalar_quant.py, _bit_packing.py, _rvq_attend.py, _vecinfer.py, _rabitq.py, _comm_vq.py, _qjl.py, fused_sdpa.py.

Tests

Gated on the metal feature (Apple Silicon), both green:

• test_metal_kernel_elementwise_scale — trivial MSL out = a * 2.0f round-trip vs pure-ops a * 2.0 within 1e-5.
• test_metal_kernel_two_inputs_with_template — 2 inputs + a dtype template arg, MSL out = a + b vs a + b within 1e-5.

cargo build, cargo test fast::tests::test_metal_kernel, cargo clippy --all-features, and cargo fmt --check all pass. Single-mlx-rs pin invariant intact (no crate-graph changes).

Follow-ups

• None required for this enabler. Porting the individual VeloxQuant MSL kernels (Phase 2) is tracked separately under epic sc-8528.

🤖 Generated with Claude Code