Pipeline parallelism support for qwen3_moe

[fretnoyz]

Summary

qwen3_moe (used by Qwen3-235B-A22B and similar) does not implement pipeline() or tensor parallelism, so running a 2-node distributed benchmark with mlx.launch --pipeline raises:

ValueError: The model does not support pipelining but a pipeline_group was provided

This patch adds PipelineMixin to Qwen3MoeModel, following the same pattern already used in glm4_moe.py.

What the patch does

Four changes, all mechanical:

1. from .pipeline import PipelineMixin import
2. class Qwen3MoeModel(PipelineMixin, nn.Module) — adds pipeline() method and pipeline_layers property via MRO
3. __call__ updated with recv/send around the layer loop (identical pattern to glm4_moe.py)
4. Model.layers returns self.model.pipeline_layers instead of self.model.layers so make_prompt_cache creates the right number of KV cache entries per rank (not all 94 — only the slice owned by this rank)

The Model.layers fix was the non-obvious bug: without it, make_prompt_cache creates 94 caches for a 94-layer model, but rank 0 (47 layers) only touches 47 of them. c.state then fails on the uninitialized entries with AttributeError: 'NoneType' object has no attribute 'shape'.

Hardware / measured result

5 × M3 Ultra Mac Studio, 96 GB each, mlx-lm 0.31.3, mlx 0.32.0.

mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit (132 GB, 94 layers), 2 nodes, pipeline parallel, tested over Tailscale (18.6 ms RTT) while waiting for a firewall fix on the Thunderbolt fabric:

Timing with prompt_tokens=512, generation_tokens=128, batch_size=1.
Trial 1:  prompt_tps=505.706, generation_tps=64.332, peak_memory=67.055
Trial 2:  prompt_tps=494.689, generation_tps=64.273, peak_memory=67.055
Trial 3:  prompt_tps=498.561, generation_tps=64.362, peak_memory=67.055
Averages: prompt_tps=499.652, generation_tps=64.322, peak_memory=67.055

4k-prompt pass:

Averages: prompt_tps=508.649, generation_tps=55.435, peak_memory=67.938

Peak memory per node is 67 GB — the 132 GB model fits cleanly across 2 × 96 GB nodes.

Patch

--- a/mlx_lm/models/qwen3_moe.py
+++ b/mlx_lm/models/qwen3_moe.py
@@ -8,6 +8,7 @@
 
 from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .pipeline import PipelineMixin
 from .switch_layers import SwitchGLU
 
 
@@ -171,7 +172,7 @@
         return out
 
 
-class Qwen3MoeModel(nn.Module):
+class Qwen3MoeModel(PipelineMixin, nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
         self.args = args
@@ -191,19 +192,32 @@
         cache=None,
         input_embeddings: Optional[mx.array] = None,
     ) -> mx.array:
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
+
         if input_embeddings is not None:
             h = input_embeddings
         else:
             h = self.embed_tokens(inputs)
 
         if cache is None:
-            cache = [None] * len(self.layers)
+            cache = [None] * len(self.pipeline_layers)
 
         mask = create_attention_mask(h, cache[0])
 
-        for layer, c in zip(self.layers, cache):
+        # Receive from the next rank in the pipeline (higher rank = earlier stage)
+        if pipeline_rank < pipeline_size - 1:
+            h = mx.distributed.recv_like(h, pipeline_rank + 1)
+
+        for layer, c in zip(self.pipeline_layers, cache):
             h = layer(h, mask, c)
 
+        # Send to the previous rank in the pipeline (lower rank = later stage)
+        if pipeline_rank != 0:
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache and cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
+
         return self.norm(h)
 
 
@@ -256,4 +270,8 @@
 
     @property
     def layers(self):
-        return self.model.layers
+        return self.model.pipeline_layers
+
+    @property
+    def pipeline_layers(self):
+        return self.model.pipeline_layers

Happy to submit as a PR if that's preferred. The patch also applies to any other qwen3_moe-based models (e.g. qwen3_5_moe may need the same treatment).