From 9ac7dbe877fb4e929b78232a36ed95fb62ed49e1 Mon Sep 17 00:00:00 2001
From: Cheng <zcbenz@gmail.com>
Date: Sun, 16 Nov 2025 07:12:18 +0900
Subject: [PATCH] Fix MPI distributed tests with CUDA backend (#2775)

---
 .github/actions/test-linux/action.yml | 13 +++++++++----
 mlx/backend/cuda/utils.cpp            |  2 +-
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/.github/actions/test-linux/action.yml b/.github/actions/test-linux/action.yml
index 9f318c8c..9a0f49b5 100644
--- a/.github/actions/test-linux/action.yml
+++ b/.github/actions/test-linux/action.yml
@@ -9,13 +9,18 @@ inputs:
 runs:
   using: "composite"
   steps:
+    - name: Run MPI tests
+      shell: bash
+      run: |
+        echo "::group::MPI tests"
+        mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
+        echo "::endgroup::"
+
     - name: Run distributed tests
-      # FIXME: This test fails with CUDA build.
       if: ${{ inputs.cpu-only == 'true' }}
       shell: bash
       run: |
         echo "::group::Distributed tests"
-        mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
         mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
         if grep -Fq '[WARN]' stderr.log ; then
           grep -F '[WARN]' stderr.log
@@ -34,7 +39,7 @@ runs:
         echo "::endgroup::"
 
     - name: Run Python tests - GPU
-      if: ${{ !inputs.cpu-only }}
+      if: ${{ inputs.cpu-only == 'false' }}
       shell: bash
       env:
         DEVICE: gpu
@@ -53,7 +58,7 @@ runs:
         echo "::endgroup::"
 
     - name: Run CPP tests - GPU
-      if: ${{ !inputs.cpu-only }}
+      if: ${{ inputs.cpu-only == 'false' }}
       shell: bash
       env:
         DEVICE: gpu
diff --git a/mlx/backend/cuda/utils.cpp b/mlx/backend/cuda/utils.cpp
index 09894d4c..3ab1a05b 100644
--- a/mlx/backend/cuda/utils.cpp
+++ b/mlx/backend/cuda/utils.cpp
@@ -60,7 +60,7 @@ const char* dtype_to_cuda_type(const Dtype& dtype) {
     case float64:
       return "double";
     case complex64:
-      return "complex64_t";
+      return "mlx::core::cu::complex64_t";
     default:
       return "unknown";
   }