chore: import upstream snapshot with attribution

backend/python/mlx-distributed/coordinator.py

@@ -0,0 +1,104 @@
+"""
+Distributed coordination using MLX distributed primitives.
+
+Rank 0 broadcasts commands and tokens to all ranks via all_sum/all_gather.
+Worker ranks wait in a loop for commands from rank 0.
+"""
+import json
+import struct
+
+import mlx.core as mx
+
+
+CMD_IDLE = 0
+CMD_GENERATE = 1
+CMD_LOAD_MODEL = 2
+CMD_SHUTDOWN = -1
+
+
+class DistributedCoordinator:
+    def __init__(self, group):
+        self.group = group
+        self.rank = group.rank()
+        self.world_size = group.size()
+
+    def broadcast_command(self, cmd, payload_size=0):
+        """Rank 0 broadcasts a command to all ranks.
+
+        Uses all_sum with only rank 0 providing non-zero values so every
+        rank receives the same command array.
+        """
+        if self.rank == 0:
+            cmd_array = mx.array([cmd, payload_size], dtype=mx.int32)
+        else:
+            cmd_array = mx.zeros((2,), dtype=mx.int32)
+        result = mx.distributed.all_sum(cmd_array, group=self.group)
+        mx.eval(result)
+        return int(result[0].item()), int(result[1].item())
+
+    def broadcast_tokens(self, tokens):
+        """Broadcast input token ids from rank 0 to all ranks.
+
+        Rank 0 provides the real token array; other ranks provide zeros of the
+        same shape.  ``all_sum`` ensures every rank ends up with identical data.
+        """
+        if self.rank == 0:
+            token_array = mx.array(tokens, dtype=mx.int32)
+        else:
+            token_array = mx.zeros((len(tokens),), dtype=mx.int32)
+        result = mx.distributed.all_sum(token_array, group=self.group)
+        mx.eval(result)
+        return result
+
+    def broadcast_token_count(self, count):
+        """Broadcast the number of tokens so workers can prepare a buffer."""
+        if self.rank == 0:
+            count_array = mx.array([count], dtype=mx.int32)
+        else:
+            count_array = mx.zeros((1,), dtype=mx.int32)
+        result = mx.distributed.all_sum(count_array, group=self.group)
+        mx.eval(result)
+        return int(result[0].item())
+
+    def broadcast_generation_params(self, max_tokens=200, temperature=0.6, top_p=1.0):
+        """Broadcast generation parameters from rank 0."""
+        if self.rank == 0:
+            params = mx.array([max_tokens, temperature, top_p], dtype=mx.float32)
+        else:
+            params = mx.zeros((3,), dtype=mx.float32)
+        result = mx.distributed.all_sum(params, group=self.group)
+        mx.eval(result)
+        return {
+            "max_tokens": int(result[0].item()),
+            "temperature": float(result[1].item()),
+            "top_p": float(result[2].item()),
+        }
+
+    def wait_for_command(self):
+        """Worker ranks block here until rank 0 broadcasts a command."""
+        return self.broadcast_command(CMD_IDLE, 0)
+
+    def broadcast_model_name(self, model_name=""):
+        """Broadcast model name string from rank 0 to all ranks.
+
+        Encodes the model name as int32 codepoints so it can travel via
+        all_sum.
+        """
+        if self.rank == 0:
+            encoded = [ord(c) for c in model_name]
+            # First broadcast the length
+            length = self.broadcast_token_count(len(encoded))
+            if length > 0:
+                name_array = mx.array(encoded, dtype=mx.int32)
+                result = mx.distributed.all_sum(name_array, group=self.group)
+                mx.eval(result)
+                return model_name
+            return ""
+        else:
+            length = self.broadcast_token_count(0)
+            if length > 0:
+                name_array = mx.zeros((length,), dtype=mx.int32)
+                result = mx.distributed.all_sum(name_array, group=self.group)
+                mx.eval(result)
+                return "".join(chr(int(c.item())) for c in result)
+            return ""