MLX on CUDA

[awni]

Quick Start

The first version of the MLX CUDA back-end is up on PyPi and ready to use. To get started:

pip install "mlx[cuda]"

The CUDA back-end works well for many common cases with notable exceptions listed below. For example, you can use mlx-lm to generate text and train or fine-tune LLMs.

To do so, first install mlx-lm:

pip install mlx-lm

Then generate text with:

mlx_lm.generate \
  --model meta-llama/Llama-3.2-3B-Instruct \
  --prompt "Write a story about Einstein" \
  -m 512

And LoRA fine-tune with:

mlx_lm.lora \
  --model meta-llama/Llama-3.2-3B-Instruct \
  --data mlx-community/WikiSQL \
  --train \
  --iters 100

Dependencies

Note the current minimum requirements on the CUDA toolkit and driver.

Dependency	Minimum Version
CUDA Toolkit	12.0
Nvidia Driver	550.54.14

Missing Operations

• Quantized matrix-vector and matrix-matrix operations
• Gather matrix-matrix multiplication and matrix-vector multiplication (used by MoEs)
• FFTs
• LAPACK-style operations (mx.linalg.svd, mx.linalg.eig, etc)

As always if you run into any problems, like missing operations, bugs, or performance cliffs please open an issue with steps to reproduce.

[lightandshadow68]

Does MLX fall back to CPU for the missing operations?

[awni]

No it will crash for missing ops. You can use the CPU fallback explicitly by passing mx.cpu to the stream kwarg of an operation.

[Velnbur]

What's the general usage of mlx with CUDA support? Macs with external GPUs?

[lightandshadow68]

This does not add CUDA support on Mac. Rather, it adds CUDA as a backend to run MLX on Linux. So you can author and test locally on a Mac, then run on the server.

However, some MLX operations are not yet implemented via CUDA.

[Velnbur]

Got it! Thanks!

[6thday]

Hello, thx for the handworks)) any chance it will work on arm cpu soon? (nvidia jetson/Ampere Altra Max etc)

[awni]

Would be great if you could make a separate issue for that to help us gauge priority

[6thday]

thx, will do that)

[samair]

This is great, was waiting for a long time!

[samair]

@awni I see that I need to install sentencepiece in addition to installing mlx* (I tried on two MacBooks). Is this expected?

[awni]

It depends what you are doing. Sentencepiece is not a dependency of mlx directly.