Integrating Custom Kernels with MLX's Lazy Compute Mechanism

[Zijie-Tian]

I’m developing a custom kernel in MLX framework and need it to integrate seamlessly with MLX’s lazy evaluation system. MLX’s built-in operators work with lazy computation (e.g., deferred execution until eval()), but my custom kernel requires explicit input evaluation (x.eval()) to function correctly, which seems inconsistent with native operators.

# Built-in operator (works with lazy tensors)  
y = mlx.add(a, b)  # No eval needed  
y = y.eval()  # Triggers computation  

# Custom kernel (fails without eval)  
z = my_custom_kernel(a, b)  # Requires a.eval(), b.eval()  

So, How do native operators avoid explicit eval() on inputs?

[awni]

You have two options:

1. Simple and preferred: use mx.metal.custom_kernel
2. Write a custom primitive.

[Zijie-Tian]

Thank you! @awni . However, after reading the implementations of eval and scheduler, I am quite confused:

1. Threadpool Usage: Since MLX maintains its own threadpool(in scheduler), should I avoid creating additional threadpools inside CPU operators? Would using multiple streams be a better approach for parallelism?
2. Cross-Device Stream Scheduling: What happens if I call CPU op in eval_gpu or call GPU op in eval_cpu, or How can I use both of them simultaneously?
3. I notice the MLX CPU operator cannot achieve the performance of other frameworks such as llama.cpp. Why?