MLX matmul (and therefore tensordot) 2,000+ times faster than Numpy on AArch64 CPU. How?

[jli05]

With MLX v0.29 (latest build as of writing), I timed just the matmul call on AArch64 CPU (Graviton3). MLX is 2,000+ times faster than NumPy. Likewise tensordot.

It looks:

• Both NumPy and MLX arrays are row-major
• Both translate to the CBLAS gemm calls, although NumPy C API does more work checking the dimensions etc. mlx/backend/cpu/gemms/cblas.cpp

Why such drastic difference?

[awni]

It's probably an issue with the timing. MLX is lazy so unless you mx.eval before you record the time.

You can read more about lazy evaluation here.