[BUG] Unable to build metal library from source with uv

[Wysie]

Describe the bug

Getting the following error when loading models (across 2 nodes) with commit 7117d74.

[ 04:42:23.8983PM | WARNING ] Runner f191e512-da70-460e-a765-22ec2c3e617b crashed with critical exception [metal::Device] Unable to build metal library from source
mlx/backend/metal/kernels/utils.h:64:25: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
instantiate_float_limit(bfloat16_t);
                        ^~~~~~~~~~
                        float16_t
mlx/backend/metal/kernels/utils.h:51:17: note: expanded from macro 'instantiate_float_limit'
  struct Limits<type> {                           \
                ^
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/utils.h:64:1: error: redefinition of 'Limits<half>'
instantiate_float_limit(bfloat16_t);
^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:51:10: note: expanded from macro 'instantiate_float_limit'
  struct Limits<type> {                           \
         ^~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:62:1: note: previous definition is here
instantiate_float_limit(half);
^
mlx/backend/metal/kernels/utils.h:51:10: note: expanded from macro 'instantiate_float_limit'
  struct Limits<type> {                           \
         ^
mlx/backend/metal/kernels/utils.h:73:15: error: use of undeclared identifier 'complex64_t'
struct Limits<complex64_t> {
              ^
mlx/backend/metal/kernels/utils.h:144:12: error: no template named 'vec'; did you mean 'metal::vec'?
METAL_FUNC vec<IdxT, 2> elem_to_loc_2_nd(
           ^~~
           metal::vec
/System/Library/PrivateFrameworks/GPUCompiler.framework/Versions/32023/Libraries/lib/clang/32023.864/include/metal/metal_extended_vector:447:1: note: 'metal::vec' declared here
using vec = T __attribute__((__ext_vector_type__(N)));
^
mlx/backend/metal/kernels/utils.h:150:3: error: no template named 'vec'; did you mean 'metal::vec'?
  vec<IdxT, 2> loc = {
  ^~~
  metal::vec
/System/Library/PrivateFrameworks/GPUCompiler.framework/Versions/32023/Libraries/lib/clang/32023.864/include/metal/metal_extended_vector:447:1: note: 'metal::vec' declared here
using vec = T __attribute__((__ext_vector_type__(N)));
^
mlx/backend/metal/kernels/utils.h:167:12: error: no template named 'vec'; did you mean 'metal::vec'?
METAL_FUNC vec<IdxT, 3> elem_to_loc_3_nd(
           ^~~
           metal::vec
/System/Library/PrivateFrameworks/GPUCompiler.framework/Versions/32023/Libraries/lib/clang/32023.864/include/metal/metal_extended_vector:447:1: note: 'metal::vec' declared here
using vec = T __attribute__((__ext_vector_type__(N)));
^
mlx/backend/metal/kernels/utils.h:174:3: error: no template named 'vec'; did you mean 'metal::vec'?
  vec<IdxT, 3> loc = {
  ^~~
  metal::vec
/System/Library/PrivateFrameworks/GPUCompiler.framework/Versions/32023/Libraries/lib/clang/32023.864/include/metal/metal_extended_vector:447:1: note: 'metal::vec' declared here
using vec = T __attribute__((__ext_vector_type__(N)));
^
mlx/backend/metal/kernels/utils.h:318:8: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
inline bfloat16_t log1p(bfloat16_t x) {
       ^~~~~~~~~~
       float16_t
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/utils.h:318:25: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
inline bfloat16_t log1p(bfloat16_t x) {
                        ^~~~~~~~~~
                        float16_t
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/utils.h:321:19: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
    return Limits<bfloat16_t>::max;
                  ^~~~~~~~~~
                  float16_t
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/utils.h:327:10: error: use of undeclared identifier 'bfloat16_t'; did you mean 'float16_t'?
  return bfloat16_t(x * (metal::log(xp1) / (xp1 - 1.0f)));
         ^
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/utils.h:330:8: error: unknown type name 'complex64_t'
inline complex64_t log1p(complex64_t in) {
       ^
mlx/backend/metal/kernels/utils.h:330:26: error: unknown type name 'complex64_t'
inline complex64_t log1p(complex64_t in) {
                         ^
mlx/backend/metal/kernels/utils.h:338:14: error: excess elements in scalar initializer
      return {x, theta};
             ^ ~~~~~~~~
mlx/backend/metal/kernels/utils.h:340:12: error: excess elements in scalar initializer
    return {0.5f * log1p(r), theta};
           ^               ~~~~~~~~
mlx/backend/metal/kernels/utils.h:343:12: error: excess elements in scalar initializer
    return {metal::log(z0), theta};
           ^              ~~~~~~~~
mlx/backend/metal/kernels/utils.h:362:10: error: call to 'simd_shuffle_down' is ambiguous
  return simd_shuffle_down(static_cast<uint32_t>(data), delta);
         ^~~~~~~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:351:17: note: candidate function
inline uint64_t simd_shuffle_down(uint64_t data, uint16_t delta) {
                ^
mlx/backend/metal/kernels/utils.h:356:16: note: candidate function
inline int64_t simd_shuffle_down(int64_t data, uint16_t delta) {
               ^
mlx/backend/metal/kernels/utils.h:361:13: note: candidate function
inline bool simd_shuffle_down(bool data, uint16_t delta) {
            ^
mlx/backend/metal/kernels/utils.h:365:8: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle_down(complex64_t data, uint16_t delta) {
       ^
mlx/backend/metal/kernels/utils.h:365:38: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle_down(complex64_t data, uint16_t delta) {
                                     ^
mlx/backend/metal/kernels/utils.h:379:10: error: call to 'simd_shuffle_up' is ambiguous
  return simd_shuffle_up(static_cast<uint32_t>(data), delta);
         ^~~~~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:370:17: note: candidate function
inline uint64_t simd_shuffle_up(uint64_t data, uint16_t delta) {
                ^
mlx/backend/metal/kernels/utils.h:374:16: note: candidate function
inline int64_t simd_shuffle_up(int64_t data, uint16_t delta) {
               ^
mlx/backend/metal/kernels/utils.h:378:13: note: candidate function
inline bool simd_shuffle_up(bool data, uint16_t delta) {
            ^
mlx/backend/metal/kernels/utils.h:382:8: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle_up(complex64_t data, uint16_t delta) {
       ^
mlx/backend/metal/kernels/utils.h:382:36: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle_up(complex64_t data, uint16_t delta) {
                                   ^
mlx/backend/metal/kernels/utils.h:400:10: error: call to 'simd_shuffle_and_fill_up' is ambiguous
  return simd_shuffle_and_fill_up(
         ^~~~~~~~~~~~~~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:388:1: note: candidate function
simd_shuffle_and_fill_up(uint64_t data, uint64_t filling, uint16_t delta) {
^
mlx/backend/metal/kernels/utils.h:394:1: note: candidate function
simd_shuffle_and_fill_up(int64_t data, int64_t filling, uint16_t delta) {
^
mlx/backend/metal/kernels/utils.h:399:13: note: candidate function
inline bool simd_shuffle_and_fill_up(bool data, bool filling, uint16_t delta) {
            ^
mlx/backend/metal/kernels/utils.h:404:8: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle_and_fill_up(
       ^
mlx/backend/metal/kernels/utils.h:405:5: error: unknown type name 'complex64_t'
    complex64_t data,
    ^
mlx/backend/metal/kernels/utils.h:406:5: error: unknown type name 'complex64_t'
    complex64_t filling,
    ^
mlx/backend/metal/kernels/utils.h:422:10: error: call to 'simd_shuffle' is ambiguous
  return simd_shuffle(static_cast<uint32_t>(data), lane);
         ^~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:413:17: note: candidate function
inline uint64_t simd_shuffle(uint64_t data, uint16_t lane) {
                ^
mlx/backend/metal/kernels/utils.h:417:16: note: candidate function
inline int64_t simd_shuffle(int64_t data, uint16_t lane) {
               ^
mlx/backend/metal/kernels/utils.h:421:13: note: candidate function
inline bool simd_shuffle(bool data, uint16_t lane) {
            ^
mlx/backend/metal/kernels/utils.h:425:8: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle(complex64_t data, uint16_t lane) {
       ^
mlx/backend/metal/kernels/utils.h:425:33: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle(complex64_t data, uint16_t lane) {
                                ^
mlx/backend/metal/kernels/indexing/gather_front.h:26:93: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
template [[host_name("gather_frontbfloat16_int32_int_2")]] [[kernel]] decltype(gather_front<bfloat16_t, int32_t, int, 2>) gather_front<bfloat16_t, int32_t, int, 2>;
                                                                                            ^~~~~~~~~~
                                                                                            float16_t
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/indexing/gather_front.h:26:136: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
template [[host_name("gather_frontbfloat16_int32_int_2")]] [[kernel]] decltype(gather_front<bfloat16_t, int32_t, int, 2>) gather_front<bfloat16_t, int32_t, int, 2>;
                                                                                                                                       ^~~~~~~~~~
                                                                                                                                       float16_t
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/indexing/gather_front.h:14:14: error: use of undeclared identifier 'offset_neg_idx'
  auto idx = offset_neg_idx(indices[index.y], size);
             ^
mlx/backend/metal/kernels/indexing/gather_front.h:26:123: note: in instantiation of function template specialization 'gather_front<half, int, int, 2>' requested here
template [[host_name("gather_frontbfloat16_int32_int_2")]] [[kernel]] decltype(gather_front<bfloat16_t, int32_t, int, 2>) gather_front<bfloat16_t, int32_t, int, 2>;
                                                                                                                          ^


Traceback (most recent call last):

  File "<string>", line 1, in <module>
  File "/Users/wysie/.local/share/uv/python/cpython-3.13.12-macos-aarch64-none/lib/python3.13/multiprocessing/spawn.py", line 122, in spawn_main
    exitcode = _main(fd, parent_sentinel)
               │     │   └ 3
               │     └ 88
               └ <function _main at 0x1030467a0>
  File "/Users/wysie/.local/share/uv/python/cpython-3.13.12-macos-aarch64-none/lib/python3.13/multiprocessing/spawn.py", line 135, in _main
    return self._bootstrap(parent_sentinel)
           │    │          └ 3
           │    └ <function BaseProcess._bootstrap at 0x102f41760>
           └ <Process name='Process-1' parent=46494 started daemon>
  File "/Users/wysie/.local/share/uv/python/cpython-3.13.12-macos-aarch64-none/lib/python3.13/multiprocessing/process.py", line 313, in _bootstrap
    self.run()
    │    └ <function BaseProcess.run at 0x102f40cc0>
    └ <Process name='Process-1' parent=46494 started daemon>
  File "/Users/wysie/.local/share/uv/python/cpython-3.13.12-macos-aarch64-none/lib/python3.13/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
    │    │        │    │        │    └ {}
    │    │        │    │        └ <Process name='Process-1' parent=46494 started daemon>
    │    │        │    └ (BoundInstance(instance=MlxJacclInstance(instance_id='60536771-cde9-4ed8-b610-56442ea1a915', shard_assignments=ShardAssignmen...
    │    │        └ <Process name='Process-1' parent=46494 started daemon>
    │    └ <function entrypoint at 0x10af7ff60>
    └ <Process name='Process-1' parent=46494 started daemon>

> File "/Users/wysie/exo/src/exo/worker/runner/bootstrap.py", line 54, in entrypoint
    runner.main()
    │      └ <function Runner.main at 0x11a61bce0>
    └ <exo.worker.runner.llm_inference.runner.Runner object at 0x11a4d34d0>

  File "/Users/wysie/exo/src/exo/worker/runner/llm_inference/runner.py", line 141, in main
    self.handle_first_task(task)
    │    │                 └ StartWarmup(task_id='c2a4d968-f67d-4656-a766-3809deea23b7', task_status=<TaskStatus.Pending: 'Pending'>, instance_id='6053677...
    │    └ <function Runner.handle_first_task at 0x11a61bd80>
    └ <exo.worker.runner.llm_inference.runner.Runner object at 0x11a4d34d0>

  File "/Users/wysie/exo/src/exo/worker/runner/llm_inference/runner.py", line 220, in handle_first_task
    self.generator.warmup()
    │    │         └ <function BatchGenerator.warmup at 0x11a61b4c0>
    │    └ BatchGenerator(model=Model(
    │        (model): MiniMaxModel(
    │          (embed_tokens): QuantizedEmbedding(200064, 3072, group_size=64, bits...
    └ <exo.worker.runner.llm_inference.runner.Runner object at 0x11a4d34d0>

  File "/Users/wysie/exo/src/exo/worker/runner/llm_inference/batch_generator.py", line 342, in warmup
    self.check_for_cancel_every = warmup_inference(
    │    │                        └ <function warmup_inference at 0x11a5d5800>
    │    └ 50
    └ BatchGenerator(model=Model(
        (model): MiniMaxModel(
          (embed_tokens): QuantizedEmbedding(200064, 3072, group_size=64, bits...

  File "/Users/wysie/exo/src/exo/worker/engines/mlx/generator/generate.py", line 344, in warmup_inference
    for _r in stream_generate(
              └ <function stream_generate at 0x11a09f060>

  File "/Users/wysie/exo/.venv/lib/python3.13/site-packages/mlx_lm/generate.py", line 705, in stream_generate
    for n, (token, logprobs, from_draft) in enumerate(token_generator):
                                                      └ <generator object stream_generate.<locals>.<genexpr> at 0x11a70e420>
  File "/Users/wysie/exo/.venv/lib/python3.13/site-packages/mlx_lm/generate.py", line 695, in <genexpr>
    (token, logprobs, False) for token, logprobs in token_generator
  File "/Users/wysie/exo/.venv/lib/python3.13/site-packages/mlx_lm/generate.py", line 438, in generate_step
    mx.eval([c.state for c in prompt_cache])
    │  │                      └ [<mlx_lm.models.cache.KVCache object at 0x11a602a50>, <mlx_lm.models.cache.KVCache object at 0x11a573750>, <mlx_lm.models.cac...
    │  └ <nanobind.nb_func object at 0x10b023ab0>
    └ <module 'mlx.core' from '/Users/wysie/exo/.venv/lib/python3.13/site-packages/mlx/core.cpython-313-darwin.so'>

RuntimeError: [metal::Device] Unable to build metal library from source
mlx/backend/metal/kernels/utils.h:64:25: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
instantiate_float_limit(bfloat16_t);
                        ^~~~~~~~~~
                        float16_t
mlx/backend/metal/kernels/utils.h:51:17: note: expanded from macro 'instantiate_float_limit'
  struct Limits<type> {                           \
                ^
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/utils.h:64:1: error: redefinition of 'Limits<half>'
instantiate_float_limit(bfloat16_t);
^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:51:10: note: expanded from macro 'instantiate_float_limit'
  struct Limits<type> {                           \
         ^~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:62:1: note: previous definition is here
instantiate_float_limit(half);
^
mlx/backend/metal/kernels/utils.h:51:10: note: expanded from macro 'instantiate_float_limit'
  struct Limits<type> {                           \
         ^
mlx/backend/metal/kernels/utils.h:73:15: error: use of undeclared identifier 'complex64_t'
struct Limits<complex64_t> {
              ^
mlx/backend/metal/kernels/utils.h:144:12: error: no template named 'vec'; did you mean 'metal::vec'?
METAL_FUNC vec<IdxT, 2> elem_to_loc_2_nd(
           ^~~
           metal::vec
/System/Library/PrivateFrameworks/GPUCompiler.framework/Versions/32023/Libraries/lib/clang/32023.864/include/metal/metal_extended_vector:447:1: note: 'metal::vec' declared here
using vec = T __attribute__((__ext_vector_type__(N)));
^
mlx/backend/metal/kernels/utils.h:150:3: error: no template named 'vec'; did you mean 'metal::vec'?
  vec<IdxT, 2> loc = {
  ^~~
  metal::vec
/System/Library/PrivateFrameworks/GPUCompiler.framework/Versions/32023/Libraries/lib/clang/32023.864/include/metal/metal_extended_vector:447:1: note: 'metal::vec' declared here
using vec = T __attribute__((__ext_vector_type__(N)));
^
mlx/backend/metal/kernels/utils.h:167:12: error: no template named 'vec'; did you mean 'metal::vec'?
METAL_FUNC vec<IdxT, 3> elem_to_loc_3_nd(
           ^~~
           metal::vec
/System/Library/PrivateFrameworks/GPUCompiler.framework/Versions/32023/Libraries/lib/clang/32023.864/include/metal/metal_extended_vector:447:1: note: 'metal::vec' declared here
using vec = T __attribute__((__ext_vector_type__(N)));
^
mlx/backend/metal/kernels/utils.h:174:3: error: no template named 'vec'; did you mean 'metal::vec'?
  vec<IdxT, 3> loc = {
  ^~~
  metal::vec
/System/Library/PrivateFrameworks/GPUCompiler.framework/Versions/32023/Libraries/lib/clang/32023.864/include/metal/metal_extended_vector:447:1: note: 'metal::vec' declared here
using vec = T __attribute__((__ext_vector_type__(N)));
^
mlx/backend/metal/kernels/utils.h:318:8: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
inline bfloat16_t log1p(bfloat16_t x) {
       ^~~~~~~~~~
       float16_t
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/utils.h:318:25: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
inline bfloat16_t log1p(bfloat16_t x) {
                        ^~~~~~~~~~
                        float16_t
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/utils.h:321:19: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
    return Limits<bfloat16_t>::max;
                  ^~~~~~~~~~
                  float16_t
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/utils.h:327:10: error: use of undeclared identifier 'bfloat16_t'; did you mean 'float16_t'?
  return bfloat16_t(x * (metal::log(xp1) / (xp1 - 1.0f)));
         ^
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/utils.h:330:8: error: unknown type name 'complex64_t'
inline complex64_t log1p(complex64_t in) {
       ^
mlx/backend/metal/kernels/utils.h:330:26: error: unknown type name 'complex64_t'
inline complex64_t log1p(complex64_t in) {
                         ^
mlx/backend/metal/kernels/utils.h:338:14: error: excess elements in scalar initializer
      return {x, theta};
             ^ ~~~~~~~~
mlx/backend/metal/kernels/utils.h:340:12: error: excess elements in scalar initializer
    return {0.5f * log1p(r), theta};
           ^               ~~~~~~~~
mlx/backend/metal/kernels/utils.h:343:12: error: excess elements in scalar initializer
    return {metal::log(z0), theta};
           ^              ~~~~~~~~
mlx/backend/metal/kernels/utils.h:362:10: error: call to 'simd_shuffle_down' is ambiguous
  return simd_shuffle_down(static_cast<uint32_t>(data), delta);
         ^~~~~~~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:351:17: note: candidate function
inline uint64_t simd_shuffle_down(uint64_t data, uint16_t delta) {
                ^
mlx/backend/metal/kernels/utils.h:356:16: note: candidate function
inline int64_t simd_shuffle_down(int64_t data, uint16_t delta) {
               ^
mlx/backend/metal/kernels/utils.h:361:13: note: candidate function
inline bool simd_shuffle_down(bool data, uint16_t delta) {
            ^
mlx/backend/metal/kernels/utils.h:365:8: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle_down(complex64_t data, uint16_t delta) {
       ^
mlx/backend/metal/kernels/utils.h:365:38: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle_down(complex64_t data, uint16_t delta) {
                                     ^
mlx/backend/metal/kernels/utils.h:379:10: error: call to 'simd_shuffle_up' is ambiguous
  return simd_shuffle_up(static_cast<uint32_t>(data), delta);
         ^~~~~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:370:17: note: candidate function
inline uint64_t simd_shuffle_up(uint64_t data, uint16_t delta) {
                ^
mlx/backend/metal/kernels/utils.h:374:16: note: candidate function
inline int64_t simd_shuffle_up(int64_t data, uint16_t delta) {
               ^
mlx/backend/metal/kernels/utils.h:378:13: note: candidate function
inline bool simd_shuffle_up(bool data, uint16_t delta) {
            ^
mlx/backend/metal/kernels/utils.h:382:8: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle_up(complex64_t data, uint16_t delta) {
       ^
mlx/backend/metal/kernels/utils.h:382:36: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle_up(complex64_t data, uint16_t delta) {
                                   ^
mlx/backend/metal/kernels/utils.h:400:10: error: call to 'simd_shuffle_and_fill_up' is ambiguous
  return simd_shuffle_and_fill_up(
         ^~~~~~~~~~~~~~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:388:1: note: candidate function
simd_shuffle_and_fill_up(uint64_t data, uint64_t filling, uint16_t delta) {
^
mlx/backend/metal/kernels/utils.h:394:1: note: candidate function
simd_shuffle_and_fill_up(int64_t data, int64_t filling, uint16_t delta) {
^
mlx/backend/metal/kernels/utils.h:399:13: note: candidate function
inline bool simd_shuffle_and_fill_up(bool data, bool filling, uint16_t delta) {
            ^
mlx/backend/metal/kernels/utils.h:404:8: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle_and_fill_up(
       ^
mlx/backend/metal/kernels/utils.h:405:5: error: unknown type name 'complex64_t'
    complex64_t data,
    ^
mlx/backend/metal/kernels/utils.h:406:5: error: unknown type name 'complex64_t'
    complex64_t filling,
    ^
mlx/backend/metal/kernels/utils.h:422:10: error: call to 'simd_shuffle' is ambiguous
  return simd_shuffle(static_cast<uint32_t>(data), lane);
         ^~~~~~~~~~~~
mlx/backend/metal/kernels/utils.h:413:17: note: candidate function
inline uint64_t simd_shuffle(uint64_t data, uint16_t lane) {
                ^
mlx/backend/metal/kernels/utils.h:417:16: note: candidate function
inline int64_t simd_shuffle(int64_t data, uint16_t lane) {
               ^
mlx/backend/metal/kernels/utils.h:421:13: note: candidate function
inline bool simd_shuffle(bool data, uint16_t lane) {
            ^
mlx/backend/metal/kernels/utils.h:425:8: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle(complex64_t data, uint16_t lane) {
       ^
mlx/backend/metal/kernels/utils.h:425:33: error: unknown type name 'complex64_t'
inline complex64_t simd_shuffle(complex64_t data, uint16_t lane) {
                                ^
mlx/backend/metal/kernels/indexing/gather_front.h:26:93: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
template [[host_name("gather_frontbfloat16_int32_int_2")]] [[kernel]] decltype(gather_front<bfloat16_t, int32_t, int, 2>) gather_front<bfloat16_t, int32_t, int, 2>;
                                                                                            ^~~~~~~~~~
                                                                                            float16_t
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/indexing/gather_front.h:26:136: error: unknown type name 'bfloat16_t'; did you mean 'float16_t'?
template [[host_name("gather_frontbfloat16_int32_int_2")]] [[kernel]] decltype(gather_front<bfloat16_t, int32_t, int, 2>) gather_front<bfloat16_t, int32_t, int, 2>;
                                                                                                                                       ^~~~~~~~~~
                                                                                                                                       float16_t
mlx/backend/metal/kernels/utils.h:7:14: note: 'float16_t' declared here
typedef half float16_t;
             ^
mlx/backend/metal/kernels/indexing/gather_front.h:14:14: error: use of undeclared identifier 'offset_neg_idx'
  auto idx = offset_neg_idx(indices[index.y], size);
             ^
mlx/backend/metal/kernels/indexing/gather_front.h:26:123: note: in instantiation of function template specialization 'gather_front<half, int, int, 2>' requested here
template [[host_name("gather_frontbfloat16_int32_int_2")]] [[kernel]] decltype(gather_front<bfloat16_t, int32_t, int, 2>) gather_front<bfloat16_t, int32_t, int, 2>;
                                                                                                                          ^

To Reproduce

Steps to reproduce the behavior:

1. checkout 7117d74
2. run exo and load a model that needs > 2 nodes (e.g MiniMax 2.5 6-bit)

Expected behavior

Model should load and run.

Actual behavior

Throws the error shown above.

Environment

• macOS Version:
• EXO Version:
• Hardware:
  • Device 1: Mac Studio M4 Max 128GB
  • Device 2: Mac Studio M4 Max 128GB
  • Additional devices:
• Interconnection:
  • Thunderbolt

Additional context

This error does not happen if you go just one commit back (178c617)

[Evanev7]

can i confirm that you need two nodes to hit this crash? also are you running exo with uv or nix? the mlx build with uv can be a bit temperamental about metal versions

[Wysie]

I'm running with uv. Let me test with a single model as well and report back.

[Wysie]

@Evanev7 One node crashes as well. I tried loading a model with just one node (closed exo on the other node completely) and the same error occurs.

[Wysie]

@Evanev7 Tried running the same thing with nix, no errors (as you've suspected)

[Evanev7]

sweet. i do want to get down to the bottom of this, but for now the nix build is what we're testing so is the 'correct' way to run exo from source to avoid the os updating underneath us. thanks for testing!

[Evanev7]

do you run into the same issue with uv after cleaning the uv cache? (just sync-clean or uv sync --force-reinstall --all-packages --no-cache)

[Wysie]

Hm I did uv sync and rm -rf ~/Library/Caches/mlx but not the two commands you recommended. Let me try and report back later.

[Wysie]

@Evanev7 Did with uv sync --force-reinstall --all-packages --no-cache and it still fails. It's very weird because I get some Metal compilation errors, and when I try to fix them they said to run xcodebuild -downloadComponent MetalToolchain again (which I already did). Anyway I ran it, and then I could run the uv sync command without errors. However, when loading models it crashes again with the same error.

[Evanev7]

ahhh ok-
current metal toolchain version is 17C7003j, but our nix build uses 17C48. i'm not sure why mlx's cmake doesn't correctly adapt to the latest metal version, but this is an upstream issue at this point if --force-reinstall --no-cache didn't work.

[thelamedia]

Hey, ran into this same issue on my two-machine setup (M4 Max 128GB + M5 Max 128GB, both on macOS 26.4) and got it working using the Nix build. Sharing the steps in case it helps others.

The issue is that the MLX fork's mlx.metallib is missing some kernels, so MLX tries to JIT compile them at runtime. macOS 26's Metal compiler doesn't support bfloat in JIT mode, so it blows up. The Nix build avoids this by pre-compiling everything at build time with a pinned Metal Toolchain.

Here's what I did on each machine:

• Install Nix:
• Make sure the Metal Toolchain is downloaded: xcodebuild -downloadComponent MetalToolchain
• Extract it and add to the Nix store:

hdiutil attach "$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' | head -1)" -mountpoint /private/tmp/metal-dmg
cp -R /private/tmp/metal-dmg/Metal.xctoolchain /private/tmp/metal-export
hdiutil detach /private/tmp/metal-dmg
nix nar pack /private/tmp/metal-export > /private/tmp/metal-toolchain-17C48.nar
nix store add --mode flat /private/tmp/metal-toolchain-17C48.nar

Update the hash in nix/metal-toolchain.nix:

The Nix flake expects a specific hash that didn't match my local toolchain, so I had to get the actual hash:
nix hash file --type sha256 --sri /private/tmp/metal-toolchain-17C48.nar

Then swap it into the file:
sed -i '' 's|hash = "sha256-.*";|hash = "sha256-YOUR_HASH_HERE=";|' nix/metal-toolchain.nix

One thing that tripped me up — my two machines produced different hashes even though they're on the same macOS version. So you need to do this on each machine separately.

Run exo using nix run .

A couple of gotchas:

• Use /private/tmp not /tmp in the Nix commands — Nix doesn't like the symlink.
• If find returns multiple .dmg files, I used the one with 058e1b31 / 022-21521-014.dmg.

After all this, both machines are running models properly across the cluster with no Metal JIT errors. The real fix would be getting the MLX fork to pre-compile all the kernels into mlx.metallib so uv run exo just works, but this got things running for me in the meantime.