Integrate MLX-c 0.31.1

Author: ronaldmannak

DISTRIBUTED-LM-INTEGRATION.md

@@ -894,7 +894,7 @@ Implement in this order. Each step produces a testable, shippable increment:
 | Limitation | Impact | Workaround |
 |-----------|--------|------------|
 | All distributed ops are CPU-only | Must use `Device.withDefaultDevice(.cpu)` | Wrap model loading and generation in CPU scope |
-| MLX-C has no backend selection parameter | Cannot programmatically choose ring vs JACCL | MLX-C tries JACCL first, then ring — usually correct |
+| No backend introspection API | Cannot query which backend was initialized for an existing group | Use `isAvailable(backend:)` to check before init |
 | `mlx_distributed_group_free()` not in public C API | Group deallocation relies on C++ shared_ptr | No action needed — works via ref counting |
 | `group.split()` unsupported by ring/JACCL | Cannot create subgroups | Not needed for tensor parallelism |
 | `sumScatter` not implemented in ring backend | Cannot use reduce-scatter collective | Use allSum instead (slightly more bandwidth) |

Source/Examples/DistributedWorker.swift

@@ -42,8 +42,8 @@ struct DistributedWorker {
     }
 
     static func runWorker(rank: Int, testOp: String) {
-        // Initialize distributed with strict=true (ring backend must be available)
-        guard let group = MLXDistributed.`init`(strict: true) else {
+        // Initialize distributed with strict=true using the ring backend
+        guard let group = MLXDistributed.`init`(strict: true, backend: .ring) else {
             fputs("ERROR: Failed to initialize distributed group (strict=true)\n", stderr)
             exit(1)
         }

Source/MLX/Distributed.swift

@@ -6,12 +6,12 @@ import Foundation
 /// Wrapper around the MLX C distributed group handle.
 ///
 /// A `DistributedGroup` represents a group of independent MLX processes
-/// that can communicate using collective operations. Use ``MLXDistributed/init(strict:)``
+/// that can communicate using collective operations. Use ``MLXDistributed/init(strict:backend:)``
 /// to create the initial group, then ``split(color:key:)`` to create sub-groups.
 ///
 /// ### See Also
 /// - ``MLXDistributed``
-/// - ``MLXDistributed/init(strict:)``
+/// - ``MLXDistributed/init(strict:backend:)``
 public final class DistributedGroup: @unchecked Sendable {
 
     let ctx: mlx_distributed_group
@@ -69,6 +69,23 @@ public final class DistributedGroup: @unchecked Sendable {
     }
 }
 
+/// The distributed communication backend to use.
+///
+/// When ``DistributedBackend/any`` is specified, MLX chooses the best available
+/// backend automatically. Use a specific case to force a particular backend.
+public enum DistributedBackend: String, CaseIterable, Sendable {
+    /// Let MLX choose the best available backend automatically.
+    case any
+    /// TCP socket-based ring backend.
+    case ring
+    /// Joint Accelerator Communication Library (Thunderbolt 5 RDMA).
+    case jaccl
+    /// Message Passing Interface backend.
+    case mpi
+    /// NVIDIA Collective Communications Library backend.
+    case nccl
+}
+
 /// Collection of distributed communication operations.
 ///
 /// Use ``MLXDistributed`` to check for distributed backend availability,
@@ -77,7 +94,7 @@ public final class DistributedGroup: @unchecked Sendable {
 ///
 /// ```swift
 /// // Initialize distributed communication
-/// let group = MLXDistributed.init()
+/// let group = MLXDistributed.`init`()
 /// print("Rank \(group.rank) of \(group.size)")
 ///
 /// // Perform an all-sum reduction
@@ -91,10 +108,10 @@ public enum MLXDistributed {
 
     /// Check if a distributed communication backend is available.
     ///
-    /// Returns `true` when the ring backend (or another backend) is compiled and
-    /// available for use.
-    public static func isAvailable() -> Bool {
-        mlx_distributed_is_available()
+    /// - Parameter backend: the backend to check (default: `.any`, checks all)
+    /// - Returns: `true` when the specified backend is available
+    public static func isAvailable(backend: DistributedBackend = .any) -> Bool {
+        backend.rawValue.withCString { mlx_distributed_is_available($0) }
     }
 
     /// Initialize the distributed backend and return the group containing
@@ -105,16 +122,21 @@ public enum MLXDistributed {
     /// When `strict` is `true`, returns `nil` if initialization fails
     /// (e.g., no hostfile configured).
     ///
-    /// > Note: MLX-C does not currently expose a backend selection parameter.
-    /// > The C layer tries backends in priority order (JACCL first, then ring).
-    /// > Track upstream mlx-c for a future `backend` parameter.
+    /// ```swift
+    /// // Use a specific backend
+    /// let group = MLXDistributed.`init`(strict: true, backend: .ring)
+    /// ```
     ///
-    /// - Parameter strict: if `true`, return `nil` on initialization failure
-    ///   instead of falling back to a singleton group
+    /// - Parameters:
+    ///   - strict: if `true`, return `nil` on initialization failure
+    ///     instead of falling back to a singleton group
+    ///   - backend: the backend to use (default: `.any`, let MLX choose)
     /// - Returns: the ``DistributedGroup`` for this process, or `nil` if
     ///   `strict` is `true` and initialization failed
-    public static func `init`(strict: Bool = false) -> DistributedGroup? {
-        let group = mlx_distributed_init(strict)
+    public static func `init`(strict: Bool = false, backend: DistributedBackend = .any)
+        -> DistributedGroup?
+    {
+        let group = backend.rawValue.withCString { mlx_distributed_init(strict, $0) }
         if group.ctx == nil {
             return nil
         }

Tests/MLXTests/DistributedTests.swift

@@ -67,6 +67,11 @@ class DistributedTests: XCTestCase {
     func testIsAvailable() {
         // Ring backend is compiled in, so isAvailable should return true
         XCTAssertTrue(MLXDistributed.isAvailable())
+
+        // Verify backend-specific availability check
+        XCTAssertTrue(
+            MLXDistributed.isAvailable(backend: .ring),
+            "Ring backend should always be available")
     }
 
     // MARK: - (2b) JACCL availability check
@@ -86,10 +91,9 @@ class DistributedTests: XCTestCase {
         // 2. The ring backend is available (true)
         // 3. On this hardware, the overall availability is true (ring)
         //
-        // NOTE: We cannot directly query which backend (ring vs JACCL) was
-        // selected because MLX-C does not expose a backend-name API. The
-        // isAvailable() call returns true if ANY backend is available. On
-        // machines without RDMA/TB5, this is the ring backend.
+        // NOTE: Backend selection is supported (e.g., .ring, .jaccl), but
+        // MLX-C does not expose a backend introspection API — there is no way
+        // to query which backend was actually initialized for an existing group.
 
         // (1) Verify isAvailable() returns a Bool
         let available = MLXDistributed.isAvailable()

skills/mlx-distributed/SKILL.md

@@ -403,7 +403,7 @@ let avgGrads3 = averageGradients(
 
 | Limitation | Impact |
 |------------|--------|
-| MLX-C doesn't expose backend selection parameter | Cannot choose between JACCL and ring; tries JACCL first, falls back to ring |
+| No backend introspection API | Cannot query which backend was initialized for an existing group; use `isAvailable(backend:)` to check before init |
 | `mlx_distributed_group_free()` not exposed in public C API | Groups leak small amounts of memory on deallocation (minimal practical impact) |
 | `group.split()` unsupported by ring and JACCL backends | Only MPI (not available on macOS) supports sub-group creation |
 | `sumScatter`/`reduceScatter` not implemented in ring backend | Use allSum + manual slicing as a workaround |

skills/mlx-distributed/references/multi-process.md

@@ -25,7 +25,7 @@ JACCL (Joint Accelerator Communication Library) uses RDMA over Thunderbolt 5 for
 - RDMA explicitly enabled in Recovery Mode (`csrutil`)
 - Physical Thunderbolt 5 cable between nodes
 
-> **Note:** MLX-C does not expose a backend selection parameter. You cannot force one backend over the other. If JACCL hardware is present, it will be preferred.
+> **Note:** You can select a specific backend using the `backend` parameter (e.g., `MLXDistributed.\`init\`(backend: .jaccl)`). Use `.any` (the default) to let MLX choose automatically.
 
 ---
 

skills/mlx-distributed/references/primitives.md

@@ -81,36 +81,46 @@ public enum MLXDistributed
 
 ### Static Methods
 
-#### isAvailable()
+#### isAvailable(backend:)
 
 Check if a distributed communication backend is available.
 
 ```swift
-public static func isAvailable() -> Bool
+public static func isAvailable(backend: DistributedBackend = .any) -> Bool
 ```
 
-**Returns:** `true` when the ring backend (or another backend) is compiled and available.
+**Parameters:**
+- `backend`: The backend to check. Default is `.any`, which checks if any backend is available.
+
+**Returns:** `true` when the specified backend is available.
 
 ```swift
+// Check if any backend is available
 if MLXDistributed.isAvailable() {
     print("Distributed backend ready")
 }
+
+// Check a specific backend
+if MLXDistributed.isAvailable(backend: .ring) {
+    print("Ring backend ready")
+}
 ```
 
-#### init(strict:)
+#### init(strict:backend:)
 
 Initialize the distributed backend and return the group containing all discoverable processes.
 
 ```swift
-public static func `init`(strict: Bool = false) -> DistributedGroup?
+public static func `init`(strict: Bool = false, backend: DistributedBackend = .any) -> DistributedGroup?
 ```
 
 **Parameters:**
 - `strict`: If `true`, returns `nil` on initialization failure instead of falling back to a singleton group. Default is `false`.
+- `backend`: The backend to use. Default is `.any`, which lets MLX choose automatically.
 
 **Returns:** The `DistributedGroup` for this process, or `nil` if `strict` is `true` and initialization failed.
 
-When `strict` is `false` (default), returns a singleton group (rank 0, size 1) if no distributed backend can be initialized. MLX-C does not expose a backend selection parameter — it tries JACCL first, then ring.
+When `strict` is `false` (default), returns a singleton group (rank 0, size 1) if no distributed backend can be initialized.
 
 ```swift
 // Non-strict: always returns a group (size-1 fallback)