feat(mlx-distributed): add new MLX-distributed backend

.github/workflows/backend.yml

@@ -157,6 +157,19 @@ jobs:
             dockerfile: "./backend/Dockerfile.python"
             context: "./"
             ubuntu-version: '2404'
+          - build-type: ''
+            cuda-major-version: ""
+            cuda-minor-version: ""
+            platforms: 'linux/amd64'
+            tag-latest: 'auto'
+            tag-suffix: '-cpu-mlx-distributed'
+            runs-on: 'ubuntu-latest'
+            base-image: "ubuntu:24.04"
+            skip-drivers: 'true'
+            backend: "mlx-distributed"
+            dockerfile: "./backend/Dockerfile.python"
+            context: "./"
+            ubuntu-version: '2404'
           # CUDA 12 builds
           - build-type: 'cublas'
             cuda-major-version: "12"
@@ -470,6 +483,19 @@ jobs:
             dockerfile: "./backend/Dockerfile.python"
             context: "./"
             ubuntu-version: '2404'
+          - build-type: 'cublas'
+            cuda-major-version: "12"
+            cuda-minor-version: "8"
+            platforms: 'linux/amd64'
+            tag-latest: 'auto'
+            tag-suffix: '-gpu-nvidia-cuda-12-mlx-distributed'
+            runs-on: 'ubuntu-latest'
+            base-image: "ubuntu:24.04"
+            skip-drivers: 'false'
+            backend: "mlx-distributed"
+            dockerfile: "./backend/Dockerfile.python"
+            context: "./"
+            ubuntu-version: '2404'
           - build-type: 'cublas'
             cuda-major-version: "12"
             cuda-minor-version: "8"
@@ -822,6 +848,19 @@ jobs:
             backend: "mlx-audio"
             dockerfile: "./backend/Dockerfile.python"
             context: "./"
+          - build-type: 'l4t'
+            cuda-major-version: "13"
+            cuda-minor-version: "0"
+            platforms: 'linux/arm64'
+            tag-latest: 'auto'
+            tag-suffix: '-nvidia-l4t-cuda-13-arm64-mlx-distributed'
+            runs-on: 'ubuntu-24.04-arm'
+            base-image: "ubuntu:24.04"
+            skip-drivers: 'false'
+            ubuntu-version: '2404'
+            backend: "mlx-distributed"
+            dockerfile: "./backend/Dockerfile.python"
+            context: "./"
           - build-type: 'cublas'
             cuda-major-version: "13"
             cuda-minor-version: "0"
@@ -926,6 +965,19 @@ jobs:
             dockerfile: "./backend/Dockerfile.python"
             context: "./"
             ubuntu-version: '2404'
+          - build-type: 'cublas'
+            cuda-major-version: "13"
+            cuda-minor-version: "0"
+            platforms: 'linux/amd64'
+            tag-latest: 'auto'
+            tag-suffix: '-gpu-nvidia-cuda-13-mlx-distributed'
+            runs-on: 'ubuntu-latest'
+            base-image: "ubuntu:24.04"
+            skip-drivers: 'false'
+            backend: "mlx-distributed"
+            dockerfile: "./backend/Dockerfile.python"
+            context: "./"
+            ubuntu-version: '2404'
           - build-type: 'cublas'
             cuda-major-version: "13"
             cuda-minor-version: "0"
@@ -1423,6 +1475,19 @@ jobs:
             dockerfile: "./backend/Dockerfile.python"
             context: "./"
             ubuntu-version: '2204'
+          - build-type: 'l4t'
+            cuda-major-version: "12"
+            cuda-minor-version: "0"
+            platforms: 'linux/arm64'
+            tag-latest: 'auto'
+            tag-suffix: '-nvidia-l4t-mlx-distributed'
+            runs-on: 'ubuntu-24.04-arm'
+            base-image: "nvcr.io/nvidia/l4t-jetpack:r36.4.0"
+            skip-drivers: 'true'
+            backend: "mlx-distributed"
+            dockerfile: "./backend/Dockerfile.python"
+            context: "./"
+            ubuntu-version: '2204'
           # SYCL additional backends
           - build-type: 'intel'
             cuda-major-version: ""
@@ -2016,6 +2081,9 @@ jobs:
           - backend: "mlx-audio"
             tag-suffix: "-metal-darwin-arm64-mlx-audio"
             build-type: "mps"
+          - backend: "mlx-distributed"
+            tag-suffix: "-metal-darwin-arm64-mlx-distributed"
+            build-type: "mps"
           - backend: "stablediffusion-ggml"
             tag-suffix: "-metal-darwin-arm64-stablediffusion-ggml"
             build-type: "metal"

Makefile

@@ -1,5 +1,5 @@
 # Disable parallel execution for backend builds
-.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/outetts backends/piper backends/stablediffusion-ggml backends/whisper backends/faster-whisper backends/silero-vad backends/local-store backends/huggingface backends/rfdetr backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/voxtral
+.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/outetts backends/piper backends/stablediffusion-ggml backends/whisper backends/faster-whisper backends/silero-vad backends/local-store backends/huggingface backends/rfdetr backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/voxtral
 
 GOCMD=go
 GOTEST=$(GOCMD) test
@@ -451,6 +451,10 @@ backends/mlx-audio:
 	BACKEND=mlx-audio $(MAKE) build-darwin-python-backend
 	./local-ai backends install "ocifile://$(abspath ./backend-images/mlx-audio.tar)"
 
+backends/mlx-distributed:
+	BACKEND=mlx-distributed $(MAKE) build-darwin-python-backend
+	./local-ai backends install "ocifile://$(abspath ./backend-images/mlx-distributed.tar)"
+
 backends/stablediffusion-ggml-darwin:
 	BACKEND=stablediffusion-ggml BUILD_TYPE=metal $(MAKE) build-darwin-go-backend
 	./local-ai backends install "ocifile://$(abspath ./backend-images/stablediffusion-ggml.tar)"
@@ -495,6 +499,7 @@ BACKEND_NEMO = nemo|python|.|false|true
 BACKEND_VOXCPM = voxcpm|python|.|false|true
 BACKEND_WHISPERX = whisperx|python|.|false|true
 BACKEND_ACE_STEP = ace-step|python|.|false|true
+BACKEND_MLX_DISTRIBUTED = mlx-distributed|python|./|false|true
 
 # Helper function to build docker image for a backend
 # Usage: $(call docker-build-backend,BACKEND_NAME,DOCKERFILE_TYPE,BUILD_CONTEXT,PROGRESS_FLAG,NEEDS_BACKEND_ARG)
@@ -548,12 +553,13 @@ $(eval $(call generate-docker-build-target,$(BACKEND_NEMO)))
 $(eval $(call generate-docker-build-target,$(BACKEND_VOXCPM)))
 $(eval $(call generate-docker-build-target,$(BACKEND_WHISPERX)))
 $(eval $(call generate-docker-build-target,$(BACKEND_ACE_STEP)))
+$(eval $(call generate-docker-build-target,$(BACKEND_MLX_DISTRIBUTED)))
 
 # Pattern rule for docker-save targets
 docker-save-%: backend-images
 	docker save local-ai-backend:$* -o backend-images/$*.tar
 
-docker-build-backends: docker-build-llama-cpp docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-voxtral
+docker-build-backends: docker-build-llama-cpp docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-voxtral docker-build-mlx-distributed
 
 ########################################################
 ### Mock Backend for E2E Tests

README.md

@@ -434,6 +434,7 @@ LocalAI couldn't have been built without the help of great software already avai
 - https://github.com/EdVince/Stable-Diffusion-NCNN
 - https://github.com/ggerganov/whisper.cpp
 - https://github.com/rhasspy/piper
+- [exo](https://github.com/exo-explore/exo) for the MLX distributed auto-parallel sharding implementation
 
 ## 🤗 Contributors
 

backend/index.yaml

@@ -118,7 +118,7 @@
   capabilities:
     nvidia: "cuda12-rfdetr"
     intel: "intel-rfdetr"
     #amd: "rocm-rfdetr"
     nvidia-l4t: "nvidia-l4t-arm64-rfdetr"
     metal: "metal-rfdetr"
     default: "cpu-rfdetr"
@@ -259,6 +259,31 @@
     nvidia-l4t: "nvidia-l4t-mlx-audio"
     nvidia-l4t-cuda-12: "nvidia-l4t-mlx-audio"
     nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-mlx-audio"
+- &mlx-distributed
+  name: "mlx-distributed"
+  uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-mlx-distributed"
+  icon: https://avatars.githubusercontent.com/u/102832242?s=200&v=4
+  urls:
+    - https://github.com/ml-explore/mlx-lm
+  mirrors:
+    - localai/localai-backends:latest-metal-darwin-arm64-mlx-distributed
+  license: MIT
+  description: |
+      Run distributed LLM inference with MLX across multiple Apple Silicon Macs
+  tags:
+    - text-to-text
+    - LLM
+    - MLX
+    - distributed
+  capabilities:
+    default: "cpu-mlx-distributed"
+    nvidia: "cuda12-mlx-distributed"
+    metal: "metal-mlx-distributed"
+    nvidia-cuda-12: "cuda12-mlx-distributed"
+    nvidia-cuda-13: "cuda13-mlx-distributed"
+    nvidia-l4t: "nvidia-l4t-mlx-distributed"
+    nvidia-l4t-cuda-12: "nvidia-l4t-mlx-distributed"
+    nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-mlx-distributed"
 - &rerankers
   name: "rerankers"
   alias: "rerankers"
@@ -791,6 +816,11 @@
   uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-mlx-audio"
   mirrors:
     - localai/localai-backends:master-metal-darwin-arm64-mlx-audio
+- !!merge <<: *mlx-distributed
+  name: "mlx-distributed-development"
+  uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-mlx-distributed"
+  mirrors:
+    - localai/localai-backends:master-metal-darwin-arm64-mlx-distributed
 ## mlx
 - !!merge <<: *mlx
   name: "cpu-mlx"
@@ -944,6 +974,57 @@
   uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-cuda-13-arm64-mlx-audio"
   mirrors:
     - localai/localai-backends:master-nvidia-l4t-cuda-13-arm64-mlx-audio
+## mlx-distributed
+- !!merge <<: *mlx-distributed
+  name: "cpu-mlx-distributed"
+  uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-mlx-distributed"
+  mirrors:
+    - localai/localai-backends:latest-cpu-mlx-distributed
+- !!merge <<: *mlx-distributed
+  name: "cpu-mlx-distributed-development"
+  uri: "quay.io/go-skynet/local-ai-backends:master-cpu-mlx-distributed"
+  mirrors:
+    - localai/localai-backends:master-cpu-mlx-distributed
+- !!merge <<: *mlx-distributed
+  name: "cuda12-mlx-distributed"
+  uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-mlx-distributed"
+  mirrors:
+    - localai/localai-backends:latest-gpu-nvidia-cuda-12-mlx-distributed
+- !!merge <<: *mlx-distributed
+  name: "cuda12-mlx-distributed-development"
+  uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-mlx-distributed"
+  mirrors:
+    - localai/localai-backends:master-gpu-nvidia-cuda-12-mlx-distributed
+- !!merge <<: *mlx-distributed
+  name: "cuda13-mlx-distributed"
+  uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-13-mlx-distributed"
+  mirrors:
+    - localai/localai-backends:latest-gpu-nvidia-cuda-13-mlx-distributed
+- !!merge <<: *mlx-distributed
+  name: "cuda13-mlx-distributed-development"
+  uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-13-mlx-distributed"
+  mirrors:
+    - localai/localai-backends:master-gpu-nvidia-cuda-13-mlx-distributed
+- !!merge <<: *mlx-distributed
+  name: "nvidia-l4t-mlx-distributed"
+  uri: "quay.io/go-skynet/local-ai-backends:latest-nvidia-l4t-mlx-distributed"
+  mirrors:
+    - localai/localai-backends:latest-nvidia-l4t-mlx-distributed
+- !!merge <<: *mlx-distributed
+  name: "nvidia-l4t-mlx-distributed-development"
+  uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-mlx-distributed"
+  mirrors:
+    - localai/localai-backends:master-nvidia-l4t-mlx-distributed
+- !!merge <<: *mlx-distributed
+  name: "cuda13-nvidia-l4t-arm64-mlx-distributed"
+  uri: "quay.io/go-skynet/local-ai-backends:latest-nvidia-l4t-cuda-13-arm64-mlx-distributed"
+  mirrors:
+    - localai/localai-backends:latest-nvidia-l4t-cuda-13-arm64-mlx-distributed
+- !!merge <<: *mlx-distributed
+  name: "cuda13-nvidia-l4t-arm64-mlx-distributed-development"
+  uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-cuda-13-arm64-mlx-distributed"
+  mirrors:
+    - localai/localai-backends:master-nvidia-l4t-cuda-13-arm64-mlx-distributed
 - !!merge <<: *kitten-tts
   name: "kitten-tts-development"
   uri: "quay.io/go-skynet/local-ai-backends:master-kitten-tts"
@@ -1390,7 +1471,7 @@
   capabilities:
     nvidia: "cuda12-rfdetr-development"
     intel: "intel-rfdetr-development"
     #amd: "rocm-rfdetr-development"
     nvidia-l4t: "nvidia-l4t-arm64-rfdetr-development"
     metal: "metal-rfdetr-development"
     default: "cpu-rfdetr-development"

backend/python/mlx-distributed/Makefile

@@ -0,0 +1,23 @@
+.PHONY: mlx-distributed
+mlx-distributed:
+	bash install.sh
+
+.PHONY: run
+run:
+	@echo "Running mlx-distributed..."
+	bash run.sh
+	@echo "mlx-distributed run."
+
+.PHONY: test
+test:
+	@echo "Testing mlx-distributed..."
+	bash test.sh
+	@echo "mlx-distributed tested."
+
+.PHONY: protogen-clean
+protogen-clean:
+	$(RM) backend_pb2_grpc.py backend_pb2.py
+
+.PHONY: clean
+clean: protogen-clean
+	rm -rf venv __pycache__