search

ml-explore / mlx

MLX: An array framework for Apple silicon
https://ml-explore.github.io/mlx/
MIT License
17.48k stars 1.01k forks source link

Dispatch bf16 at run time when using the JIT #1584

Closed awni closed 1 week ago

awni commented 1 week ago

This dispatches based on metal version at runtime when using the JIT.

The idea is to simplify distributing a single binary while supporting older OS yet getting newer features when available.

In general I think these are good guidelines to follow that give the most flexibility in terms of building applications with MLX:

  1. As much as possible minimize where we conditionally compile based on the metal version
  2. When using the JIT we can easily choose the right branch at run-time with negligible cost. As much as possible we should put things which dispatch in Metal in the JIT (when MLX_METAL_JIT=True) so that we at least have this option.
  3. When not using the JIT you have two options: either deploy a single binary for the minimal target without using newer features OR ship multiple binaries (like we do in PyPi).

CC @davidkoski