Support an MLX CUDA backend to improve Exo inference on NVIDIA GPUs

[satoutahhaithem]

MLX just landed a native CUDA backend - see PR #1983
[Daily.dev](https://app.daily.dev/posts/wip-cuda-backend-by-zcbenz-pull-request-1983-ml-explore-mlx-ergfm6crk , ml-explore/mlx#1983)
; installing pip install mlx-cuda gives matmul and other kernels running directly on NVIDIA GPUs. My earlier inference runs on Exo with PyTorch/tinygrad hit ~5 tokens/s even while the Jetson GPU sat at 100 % utilisation—details in issue #834
GitHub
—so the slowdown is clearly in Exo, not the hardware. Switching Exo to the new MLX-CUDA path on Linux machines could remove that bottleneck, simplify deps, and bring us up to BLAS-level performance. Would love guidance on wiring an Exo backend to MLX or any other ideas the team has.

[satoutahhaithem]

Mlx-cuda is still under development
#1983 (comment)

[fakebizprez]

🙌

[satoutahhaithem]

✅ First results: Llama inference on mlx-cuda

I pulled the Llama example from ml-explore/mlx-examples and ran it with the brand-new CUDA wheels. Environment and raw numbers:

Component	Value
Host	DGX-Station (4 × V100-16 GB)
Driver	570.158.01
CUDA	12.8
Python	3.10.13
mlx-cuda	0.0.9.dev20250731
Command	python lms/llama/llama.py --model-path mlx_model/weights.npz --gpu --prompt "hello"

Stage	Wall-time	Notes
Prompt processing	5.94 s	includes weight load
First token	7.33 s	end-to-end
Sustained throughput	≈ 48 tok/s	after warm-up
GPU-0 utilisation	~59 %	6.6 GiB VRAM
GPU 1-3	idle	61 °C / 65 W

---

❓ Where Exo comes in: sharding & multi-GPU

Single-GPU inference is now sorted; the next goal is to split the workload across the other three V100s. A few open questions for the Exo maintainers:

1. Backend plumbing
   Is it enough to pip install mlx-cuda and set backend="mlx" in Exo, or do we need a dedicated MlxCudaBackend that implements sharding helpers (shard_tensor, gather_tensor, NCCL collectives, etc.)?

❓ Exo integration

Do we need to rewrite any inference code in Exo to support mlx-cuda, or is it sufficient to simply replace the original mlx installation with pip install mlx-cuda?

[satoutahhaithem]

As an update, one of the MLX contributors informed me that mlx.distributed has not yet been integrated with mlx-cuda, but they are working on it (see: ml-explore/mlx-examples#1378 (comment)). I will provide further updates on this issue as they become available.

[fakebizprez]

Component Value
Host DGX-Station (4 × V100-16 GB)
Driver 570.158.01
CUDA 12.8
Python 3.10.13
mlx-cuda 0.0.9.dev20250731

You're doing god's work, but i am curious after seeing your gear -- why are you running MLX and not CUDA?

I'm genuinely curious; I have the new Mac Studio that I run MLX on, but my Threadripper station I'm running the same DGX GPUs as you are, except only two of the 32GB versions with no NVLink/Bridge.

[satoutahhaithem]

MLX was originally developed for Apple machines, and it works perfectly on them. However, because the MLX team has added support for a CUDA backend, we tried running only the inference on Ubuntu machines. After that, we plan to integrate it into EXO.

As you can see, it worked perfectly on a single node. But to run it in a distributed system, we need to distribute the inference using MLX’s mx.distributed package. Unfortunately, the MLX team has not added the support of cuda in this package, as noted in this issue on their GitHub: mlx-examples #1378.

[fakebizprez]

Ahh, totally get it. A lot of progress.

Are you going to continue contributing to this repo? I heard the Exo team have already went closed source, which is their prerogative, but are actively ignoring their own community.

[peteygao]

@fakebizprez Are you sure they have gone closed source? (was there an announcement I missed somewhere?)

Exo going closed source is completely antithetical to their whole mission statement:

EXO LABS

Humanity is currently headed towards a world controlled by a cartel of a handful of companies that want to control the AI models that everyone will use.

   exo's mission is to democratize access to AI.

WHAT ARE WE GOING TO DO ABOUT IT?

Build open infrastructure to train frontier models and enable any human to run them anywhere.

exo was started by a team of AI researchers and engineers from Oxford University singularly focused on this.

Quoted directly from https://exolabs.net/

[fakebizprez]

@peteygao can't say for sure, but actions speak louder than words.