Use different minimax sharding

src/exo/worker/engines/mlx/auto_parallel.py

@@ -652,18 +652,11 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
             layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
 
-            # Shard qk_norm weights if present (must match sharded head count)
-            if getattr(layer.self_attn, "use_qk_norm", False):
-                layer.self_attn.q_norm.weight = layer.self_attn.q_norm.weight.split(  # type: ignore
-                    self.N, axis=-1
-                )[rank]
-                layer.self_attn.k_norm.weight = layer.self_attn.k_norm.weight.split(  # type: ignore
-                    self.N, axis=-1
-                )[rank]
-
             layer.self_attn.num_attention_heads //= self.N
             layer.self_attn.num_key_value_heads //= self.N
 
+            layer.self_attn = WrappedMinimaxAttention(layer.self_attn, self.group)
+
             # Shard the MoE. Shard in place since the MoE should be responsible
             # for aggregating the results.
             self.all_to_sharded_linear_in_place(
@@ -680,6 +673,65 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             mx.eval(layer)
         return model
 
+    class WrappedMiniMaxAttention(CustomMlxLayer):
+        def __init__(self, layer: _LayerCallable, group: mx.distributed.Group):
+            super().__init__(layer)
+            self.group = group
+
+        def __call__(
+            self,
+            x: mx.array,
+            mask: Optional[mx.array] = None,
+            cache: Optional[Any] = None,
+        ) -> mx.array:
+            B, L, D = x.shape
+
+            queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+            q_dim = queries.shape[-1]
+            k_dim = keys.shape[-1]
+
+            qk = mx.concatenate([queries, keys], axis=-1)
+            qk = mx.distributed.all_gather(qk, group=self.group)
+
+
+            if getattr(self.self_attn, "use_qk_norm", False):
+                queries, keys = qk[..., :q_dim], qk[..., q_dim: q_dim + k_dim]
+
+                queries = self.q_norm(queries)
+                keys = self.k_norm(keys)
+
+                queries = queries.split(  # type: ignore
+                    self.group.size(), axis=-1
+                )[self.group.rank()]
+                keys = keys.split(self.group.size(), axis=-1)[self.group.rank()]  # type: ignore
+
+            queries = queries.reshape(B, L, self.num_attention_heads, -1).transpose(
+                0, 2, 1, 3
+            )
+            keys = keys.reshape(B, L, self.num_key_value_heads, -1).transpose(
+                0, 2, 1, 3
+            )
+            values = values.reshape(B, L, self.num_key_value_heads, -1).transpose(
+                0, 2, 1, 3
+            )
+
+            if cache is not None:
+                queries = self.rope(queries, offset=cache.offset)
+                keys = self.rope(keys, offset=cache.offset)
+                keys, values = cache.update_and_fetch(keys, values)
+            else:
+                queries = self.rope(queries)
+                keys = self.rope(keys)
+
+            output = scaled_dot_product_attention(
+                queries, keys, values, cache=cache, scale=self.scale, mask=mask
+            )
+
+            output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+
+            return self.o_proj(output)
+
 
 class QwenShardingStrategy(TensorParallelShardingStrategy):
     def shard_model(