diff --git a/pyproject.toml b/pyproject.toml
index ecb66133..f3a9e677 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -62,7 +62,7 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
-mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
+mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "leo/add-uneven-sharding", marker = "sys_platform == 'darwin'" }
 mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arrayscache-leak" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 0085e2f5..af5405f5 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Literal, Protocol, cast
 import mlx.core as mx
 import mlx.nn as nn
 from mlx.nn.layers.distributed import (
+    compute_shard_sizes,
     shard_inplace,
     shard_linear,
     sum_gradients,
@@ -657,13 +658,15 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
         for i, layer in enumerate(model.layers):
             # Force load weights before sharding to avoid FAST_SYNCH deadlock
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
-            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
-            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
-            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)
-            layer.self_attn.n_heads //= self.N
-            if layer.self_attn.n_kv_heads is not None:
-                layer.self_attn.n_kv_heads //= self.N
+            head_dim = layer.self_attn.head_dim
+            n_kv = layer.self_attn.n_kv_heads or layer.self_attn.n_heads
+            gqa_unit = head_dim * (layer.self_attn.n_heads // n_kv)
+            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj, unit=gqa_unit)
+            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj, unit=head_dim)
+            layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj, unit=head_dim)
+            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=gqa_unit)
+            layer.self_attn.n_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
+            layer.self_attn.n_kv_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
 
             layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
             layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
@@ -715,22 +718,24 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
 
             # Shard the self attention
+            original_num_heads = layer.self_attn.num_heads
+            q_head_dim = layer.self_attn.q_b_proj.weight.shape[0] // original_num_heads if layer.self_attn.q_lora_rank is not None else layer.self_attn.q_proj.weight.shape[0] // original_num_heads
             if layer.self_attn.q_lora_rank is None:
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_proj
+                    layer.self_attn.q_proj, unit=q_head_dim
                 )
             else:
                 layer.self_attn.q_b_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_b_proj
+                    layer.self_attn.q_b_proj, unit=q_head_dim
                 )
 
-            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
-            layer.self_attn.num_heads //= self.N
+            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=q_head_dim)
+            head_sizes = compute_shard_sizes(original_num_heads, self.N)
+            layer.self_attn.num_heads = head_sizes[self.group.rank()]
 
             # Logic from upstream mlx
-            num_heads = layer.self_attn.num_heads
-            sh = self.group.rank() * num_heads
-            eh = sh + num_heads
+            sh = sum(head_sizes[:self.group.rank()])
+            eh = sh + head_sizes[self.group.rank()]
 
             def shard_heads(w: mx.array, sh: int = sh, eh: int = eh) -> mx.array:
                 return w[sh:eh]
@@ -802,22 +807,24 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
                 timeout_seconds / total,
                 on_timeout,
             )
+            original_num_heads = layer.self_attn.num_heads  # type: ignore
+            q_head_dim = layer.self_attn.q_b_proj.weight.shape[0] // original_num_heads if layer.self_attn.q_lora_rank is not None else layer.self_attn.q_proj.weight.shape[0] // original_num_heads  # type: ignore
             if layer.self_attn.q_lora_rank is None:  # type: ignore
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_proj
+                    layer.self_attn.q_proj, unit=q_head_dim
                 )
             else:
                 layer.self_attn.q_b_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_b_proj
+                    layer.self_attn.q_b_proj, unit=q_head_dim
                 )
 
-            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
-            layer.self_attn.num_heads //= self.N
+            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=q_head_dim)
+            head_sizes = compute_shard_sizes(original_num_heads, self.N)
+            layer.self_attn.num_heads = head_sizes[self.group.rank()]
 
             # Logic from upstream mlx
-            num_heads = layer.self_attn.num_heads
-            sh = self.group.rank() * num_heads
-            eh = sh + num_heads
+            sh = sum(head_sizes[:self.group.rank()])
+            eh = sh + head_sizes[self.group.rank()]
 
             def shard_heads(w: mx.array, sh: int = sh, eh: int = eh) -> mx.array:
                 return w[sh:eh]
@@ -944,13 +951,15 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
         for i, layer in enumerate(model.layers):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
             # Shard the self attention
-            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
-            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
-            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)
+            head_dim = layer.self_attn.head_dim
+            gqa_unit = head_dim * (layer.self_attn.num_attention_heads // layer.self_attn.num_key_value_heads)
+            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj, unit=gqa_unit)
+            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj, unit=head_dim)
+            layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj, unit=head_dim)
+            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=gqa_unit)
 
-            layer.self_attn.num_attention_heads //= self.N
-            layer.self_attn.num_key_value_heads //= self.N
+            layer.self_attn.num_attention_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
+            layer.self_attn.num_key_value_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
 
             layer.self_attn = WrappedMiniMaxAttention(layer.self_attn, self.group)  # pyright: ignore[reportAttributeAccessIssue,reportArgumentType]
 
@@ -988,20 +997,22 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
             # Shard the self attention
             if isinstance(layer, Qwen3MoeDecoderLayer):
+                head_dim = layer.self_attn.q_proj.weight.shape[0] // layer.self_attn.n_heads
+                gqa_unit = head_dim * (layer.self_attn.n_heads // layer.self_attn.n_kv_heads)
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_proj
+                    layer.self_attn.q_proj, unit=gqa_unit
                 )
                 layer.self_attn.k_proj = self.all_to_sharded_linear(
-                    layer.self_attn.k_proj
+                    layer.self_attn.k_proj, unit=head_dim
                 )
                 layer.self_attn.v_proj = self.all_to_sharded_linear(
-                    layer.self_attn.v_proj
+                    layer.self_attn.v_proj, unit=head_dim
                 )
                 layer.self_attn.o_proj = self.sharded_to_all_linear(
-                    layer.self_attn.o_proj
+                    layer.self_attn.o_proj, unit=gqa_unit
                 )
-                layer.self_attn.n_heads //= self.N
-                layer.self_attn.n_kv_heads //= self.N
+                layer.self_attn.n_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
+                layer.self_attn.n_kv_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
             else:
                 assert isinstance(layer, (Qwen3NextDecoderLayer, Qwen3_5DecoderLayer))
                 if hasattr(layer, "linear_attn"):
@@ -1019,16 +1030,19 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                         # Qwen3.5: separate projections
                         # in_proj_qkv has sections [q(key_dim), k(key_dim), v(value_dim)]
                         # that must be split section-aware, not as a contiguous block
+                        head_k_dim = linear_attn.head_k_dim
+                        head_v_dim = linear_attn.head_v_dim
                         key_dim = linear_attn.key_dim
                         value_dim = linear_attn.value_dim
                         linear_attn.in_proj_qkv = shard_linear(
                             linear_attn.in_proj_qkv,
                             "all-to-sharded",
                             segments=[key_dim, key_dim + key_dim],
+                            unit=head_k_dim,
                             group=self.group,
                         )
                         linear_attn.in_proj_z = self.all_to_sharded_linear(
-                            linear_attn.in_proj_z
+                            linear_attn.in_proj_z, unit=head_v_dim
                         )
                         linear_attn.in_proj_b = self.all_to_sharded_linear(
                             linear_attn.in_proj_b
@@ -1037,7 +1051,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                             linear_attn.in_proj_a
                         )
                     linear_attn.out_proj = self.sharded_to_all_linear(
-                        linear_attn.out_proj
+                        linear_attn.out_proj, unit=linear_attn.head_v_dim
                     )
 
                     # Shard conv1d: depthwise conv with non-contiguous channel slicing.
@@ -1046,31 +1060,37 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                     rank = self.group.rank()
                     key_dim = linear_attn.key_dim
                     value_dim = linear_attn.value_dim
-                    key_dim_shard = key_dim // self.N
-                    value_dim_shard = value_dim // self.N
+                    head_k_dim = linear_attn.head_k_dim
+                    head_v_dim = linear_attn.head_v_dim
+                    key_shard_sizes = compute_shard_sizes(key_dim, self.N, unit=head_k_dim)
+                    value_shard_sizes = compute_shard_sizes(value_dim, self.N, unit=head_v_dim)
+                    key_dim_shard = key_shard_sizes[rank]
+                    value_dim_shard = value_shard_sizes[rank]
+                    key_dim_offset = sum(key_shard_sizes[:rank])
+                    value_dim_offset = sum(value_shard_sizes[:rank])
 
-                    q_idx = mx.arange(rank * key_dim_shard, (rank + 1) * key_dim_shard)
+                    q_idx = mx.arange(key_dim_offset, key_dim_offset + key_dim_shard)
                     k_idx = mx.arange(
-                        key_dim + rank * key_dim_shard,
-                        key_dim + (rank + 1) * key_dim_shard,
+                        key_dim + key_dim_offset,
+                        key_dim + key_dim_offset + key_dim_shard,
                     )
                     v_idx = mx.arange(
-                        2 * key_dim + rank * value_dim_shard,
-                        2 * key_dim + (rank + 1) * value_dim_shard,
+                        2 * key_dim + value_dim_offset,
+                        2 * key_dim + value_dim_offset + value_dim_shard,
                     )
                     conv_indices = mx.concatenate([q_idx, k_idx, v_idx])
                     linear_attn.conv1d.weight = linear_attn.conv1d.weight[conv_indices]
                     new_conv_dim = key_dim_shard * 2 + value_dim_shard
                     linear_attn.conv1d.groups = new_conv_dim
 
-                    num_v_shard = linear_attn.num_v_heads // self.N
-                    v_start = rank * num_v_shard
-                    v_end = v_start + num_v_shard
-                    linear_attn.A_log = linear_attn.A_log[v_start:v_end]
-                    linear_attn.dt_bias = linear_attn.dt_bias[v_start:v_end]
+                    num_k_per_rank = key_dim_shard // head_k_dim
+                    num_v_per_rank = value_dim_shard // head_v_dim
+                    v_offset = value_dim_offset // head_v_dim
+                    linear_attn.A_log = linear_attn.A_log[v_offset:v_offset + num_v_per_rank]
+                    linear_attn.dt_bias = linear_attn.dt_bias[v_offset:v_offset + num_v_per_rank]
 
-                    linear_attn.num_k_heads //= self.N
-                    linear_attn.num_v_heads //= self.N
+                    linear_attn.num_k_heads = num_k_per_rank
+                    linear_attn.num_v_heads = num_v_per_rank
                     linear_attn.key_dim = (
                         linear_attn.head_k_dim * linear_attn.num_k_heads
                     )
@@ -1081,20 +1101,22 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                         linear_attn.key_dim * 2 + linear_attn.value_dim
                     )
                 else:
+                    kv_head_dim = layer.self_attn.k_proj.weight.shape[0] // layer.self_attn.num_key_value_heads
+                    gqa_repeat = layer.self_attn.num_attention_heads // layer.self_attn.num_key_value_heads
                     layer.self_attn.q_proj = self.all_to_sharded_linear(
-                        layer.self_attn.q_proj
+                        layer.self_attn.q_proj, unit=kv_head_dim * 2 * gqa_repeat
                     )
                     layer.self_attn.k_proj = self.all_to_sharded_linear(
-                        layer.self_attn.k_proj
+                        layer.self_attn.k_proj, unit=kv_head_dim
                     )
                     layer.self_attn.v_proj = self.all_to_sharded_linear(
-                        layer.self_attn.v_proj
+                        layer.self_attn.v_proj, unit=kv_head_dim
                     )
                     layer.self_attn.o_proj = self.sharded_to_all_linear(
-                        layer.self_attn.o_proj
+                        layer.self_attn.o_proj, unit=kv_head_dim * gqa_repeat
                     )
-                    layer.self_attn.num_attention_heads //= self.N
-                    layer.self_attn.num_key_value_heads //= self.N
+                    layer.self_attn.num_attention_heads = layer.self_attn.q_proj.weight.shape[0] // (kv_head_dim * 2)
+                    layer.self_attn.num_key_value_heads = layer.self_attn.k_proj.weight.shape[0] // kv_head_dim
 
             # Shard the MoE.
             if isinstance(
@@ -1146,12 +1168,14 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
         for i, layer in enumerate(model.layers):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
 
-            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
-            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
-            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)
-            layer.self_attn.n_heads //= self.N
-            layer.self_attn.n_kv_heads //= self.N
+            head_dim = layer.self_attn.q_proj.weight.shape[0] // layer.self_attn.n_heads
+            gqa_unit = head_dim * (layer.self_attn.n_heads // layer.self_attn.n_kv_heads)
+            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj, unit=gqa_unit)
+            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj, unit=head_dim)
+            layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj, unit=head_dim)
+            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=gqa_unit)
+            layer.self_attn.n_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
+            layer.self_attn.n_kv_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
 
             if isinstance(layer.mlp, MoE):
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
@@ -1194,23 +1218,26 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
 
         for i, layer in enumerate(model.layers):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
-            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
-            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
-            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)
+            head_dim = layer.self_attn.head_dim
+            original_num_heads = layer.self_attn.num_attention_heads
+            gqa_unit = head_dim * (layer.self_attn.num_attention_heads // layer.self_attn.num_key_value_heads)
+            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj, unit=gqa_unit)
+            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj, unit=head_dim)
+            layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj, unit=head_dim)
+            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=gqa_unit)
 
-            layer.self_attn.num_attention_heads //= self.N
-            layer.self_attn.num_key_value_heads //= self.N
+            layer.self_attn.num_attention_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
+            layer.self_attn.num_key_value_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
             layer.self_attn.num_key_value_groups = (
                 layer.self_attn.num_attention_heads
                 // layer.self_attn.num_key_value_heads
             )
 
-            layer.self_attn.sinks = layer.self_attn.sinks[
-                layer.self_attn.num_attention_heads
-                * self.group.rank() : layer.self_attn.num_attention_heads
-                * (self.group.rank() + 1)
-            ]
+            rank = self.group.rank()
+            q_head_sizes = compute_shard_sizes(original_num_heads, self.N, unit=gqa_unit // head_dim)
+            sink_start = sum(q_head_sizes[:rank])
+            sink_end = sink_start + q_head_sizes[rank]
+            layer.self_attn.sinks = layer.self_attn.sinks[sink_start:sink_end]
 
             self.all_to_sharded_linear_in_place(layer.mlp.experts.gate_proj)
             self.sharded_to_all_linear_in_place(layer.mlp.experts.down_proj)
@@ -1237,17 +1264,19 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
 
         for i, layer in enumerate(model.layers):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
-            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
-            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
-            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)
+            head_dim = layer.self_attn.head_dim
+            gqa_unit = head_dim * (layer.self_attn.num_heads // layer.self_attn.num_kv_heads)
+            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj, unit=gqa_unit)
+            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj, unit=head_dim)
+            layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj, unit=head_dim)
+            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=gqa_unit)
 
-            layer.self_attn.num_heads //= self.N
-            layer.self_attn.num_kv_heads //= self.N
+            layer.self_attn.num_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
+            layer.self_attn.num_kv_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
 
             if getattr(layer.self_attn, "use_head_wise_attn_gate", False):
                 layer.self_attn.g_proj = self.all_to_sharded_linear(
-                    layer.self_attn.g_proj
+                    layer.self_attn.g_proj, unit=gqa_unit // head_dim
                 )
 
             if isinstance(layer.mlp, Step35MLP):
@@ -1286,12 +1315,14 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
             mixer = layer.mixer
 
             if isinstance(mixer, NemotronHAttention):
-                mixer.q_proj = self.all_to_sharded_linear(mixer.q_proj)
-                mixer.k_proj = self.all_to_sharded_linear(mixer.k_proj)
-                mixer.v_proj = self.all_to_sharded_linear(mixer.v_proj)
-                mixer.o_proj = self.sharded_to_all_linear(mixer.o_proj)
-                mixer.num_heads //= self.N
-                mixer.num_key_value_heads //= self.N
+                attn_head_dim = mixer.head_dim
+                gqa_unit = attn_head_dim * (mixer.num_heads // mixer.num_key_value_heads)
+                mixer.q_proj = self.all_to_sharded_linear(mixer.q_proj, unit=gqa_unit)
+                mixer.k_proj = self.all_to_sharded_linear(mixer.k_proj, unit=attn_head_dim)
+                mixer.v_proj = self.all_to_sharded_linear(mixer.v_proj, unit=attn_head_dim)
+                mixer.o_proj = self.sharded_to_all_linear(mixer.o_proj, unit=gqa_unit)
+                mixer.num_heads = mixer.q_proj.weight.shape[0] // attn_head_dim
+                mixer.num_key_value_heads = mixer.k_proj.weight.shape[0] // attn_head_dim
 
             elif isinstance(mixer, NemotronHMamba2Mixer):
                 self._shard_mamba2_mixer(mixer, rank)
@@ -1322,12 +1353,22 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
         ssm_state_size = mixer.ssm_state_size
         intermediate_size = mixer.intermediate_size  # = num_heads * head_dim
 
-        # Per-rank sizes
-        heads_per_rank = num_heads // world_size
-        groups_per_rank = n_groups // world_size
+        # Distribute groups first, derive heads from groups
+        heads_per_group = num_heads // n_groups
+        group_sizes = compute_shard_sizes(n_groups, world_size)
+        head_sizes = [g * heads_per_group for g in group_sizes]
+
+        # Per-rank sizes from uneven distribution
+        groups_per_rank = group_sizes[rank]
+        heads_per_rank = head_sizes[rank]
         is_per_rank = heads_per_rank * head_dim
         bc_per_rank = groups_per_rank * ssm_state_size
 
+        # Cumulative offsets
+        is_offset = sum(head_sizes[:rank]) * head_dim
+        bc_offset = sum(group_sizes[:rank]) * ssm_state_size
+        head_offset = sum(head_sizes[:rank])
+
         # === in_proj: output layout is [gate:IS | conv_ssm:IS | B:NG*SS | C:NG*SS | dt:NH] ===
         gate_start = 0
         conv_ssm_start = intermediate_size
@@ -1337,38 +1378,38 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
 
         # Build index tensor for this rank's slice of each section
         gate_idx = mx.arange(
-            gate_start + rank * is_per_rank, gate_start + (rank + 1) * is_per_rank
+            gate_start + is_offset, gate_start + is_offset + is_per_rank
         )
         conv_ssm_idx = mx.arange(
-            conv_ssm_start + rank * is_per_rank,
-            conv_ssm_start + (rank + 1) * is_per_rank,
+            conv_ssm_start + is_offset,
+            conv_ssm_start + is_offset + is_per_rank,
         )
         b_idx = mx.arange(
-            b_start + rank * bc_per_rank, b_start + (rank + 1) * bc_per_rank
+            b_start + bc_offset, b_start + bc_offset + bc_per_rank
         )
         c_idx = mx.arange(
-            c_start + rank * bc_per_rank, c_start + (rank + 1) * bc_per_rank
+            c_start + bc_offset, c_start + bc_offset + bc_per_rank
         )
         dt_idx = mx.arange(
-            dt_start + rank * heads_per_rank, dt_start + (rank + 1) * heads_per_rank
+            dt_start + head_offset, dt_start + head_offset + heads_per_rank
         )
 
         indices = mx.concatenate([gate_idx, conv_ssm_idx, b_idx, c_idx, dt_idx])
         mixer.in_proj.weight = mixer.in_proj.weight[indices]
 
         # === out_proj: input is intermediate_size (sharded) → hidden_size (reduce) ===
-        mixer.out_proj = self.sharded_to_all_linear(mixer.out_proj)
+        mixer.out_proj = self.sharded_to_all_linear(mixer.out_proj, unit=heads_per_group * head_dim)
 
         # === conv1d: depthwise conv on conv_dim channels ===
         # conv_dim layout: [ssm_hidden:IS | B:NG*SS | C:NG*SS]
-        conv_ssm_idx_local = mx.arange(rank * is_per_rank, (rank + 1) * is_per_rank)
+        conv_ssm_idx_local = mx.arange(is_offset, is_offset + is_per_rank)
         conv_b_idx = mx.arange(
-            intermediate_size + rank * bc_per_rank,
-            intermediate_size + (rank + 1) * bc_per_rank,
+            intermediate_size + bc_offset,
+            intermediate_size + bc_offset + bc_per_rank,
         )
         conv_c_idx = mx.arange(
-            intermediate_size + n_groups * ssm_state_size + rank * bc_per_rank,
-            intermediate_size + n_groups * ssm_state_size + (rank + 1) * bc_per_rank,
+            intermediate_size + n_groups * ssm_state_size + bc_offset,
+            intermediate_size + n_groups * ssm_state_size + bc_offset + bc_per_rank,
         )
         conv_indices = mx.concatenate([conv_ssm_idx_local, conv_b_idx, conv_c_idx])
         mixer.conv1d.weight = mixer.conv1d.weight[conv_indices]
@@ -1378,16 +1419,15 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
             mixer.conv1d.bias = mixer.conv1d.bias[conv_indices]
 
         # === Per-head parameters ===
-        h_start = rank * heads_per_rank
+        h_start = head_offset
         h_end = h_start + heads_per_rank
         mixer.dt_bias = mixer.dt_bias[h_start:h_end]
         mixer.A_log = mixer.A_log[h_start:h_end]
         mixer.D = mixer.D[h_start:h_end]
 
         # === Norm: weight is intermediate_size ===
-        mixer.norm.weight = mixer.norm.weight[
-            rank * is_per_rank : (rank + 1) * is_per_rank
-        ]
+        mixer.norm.weight = mixer.norm.weight[is_offset : is_offset + is_per_rank]
+        mixer.norm.group_size = is_per_rank // groups_per_rank
 
         # === Update dimensions ===
         mixer.num_heads = heads_per_rank
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_uneven_sharding.py b/src/exo/worker/tests/unittests/test_mlx/test_uneven_sharding.py
new file mode 100644
index 00000000..6d400a9b
--- /dev/null
+++ b/src/exo/worker/tests/unittests/test_mlx/test_uneven_sharding.py
@@ -0,0 +1,517 @@
+# type: ignore
+import importlib
+import itertools
+import json
+import multiprocessing as mp
+import os
+import tempfile
+import traceback
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+import pytest
+from mlx.nn.layers.distributed import compute_shard_sizes
+
+RANDOM_SEED = 42
+INPUT_TOKENS = [1, 100, 200, 300]
+
+REDUCED_CONFIGS = {
+    "llama": {
+        "model_type": "llama",
+        "num_hidden_layers": 2,
+        "hidden_size": 256,
+        "num_attention_heads": 8,
+        "num_key_value_heads": 4,
+        "vocab_size": 1024,
+        "intermediate_size": 512,
+        "rms_norm_eps": 1e-5,
+        "rope_theta": 10000.0,
+        "tie_word_embeddings": True,
+    },
+    "gpt_oss": {
+        "model_type": "gpt_oss",
+        "num_hidden_layers": 4,
+        "hidden_size": 256,
+        "num_attention_heads": 8,
+        "num_key_value_heads": 4,
+        "head_dim": 64,
+        "vocab_size": 1024,
+        "intermediate_size": 256,
+        "num_local_experts": 4,
+        "num_experts_per_tok": 2,
+        "sliding_window": 64,
+        "rope_theta": 10000.0,
+        "tie_word_embeddings": True,
+        "rms_norm_eps": 1e-5,
+    },
+    "deepseek_v3": {
+        "model_type": "deepseek_v3",
+        "num_hidden_layers": 4,
+        "hidden_size": 256,
+        "num_attention_heads": 8,
+        "num_key_value_heads": 4,
+        "vocab_size": 1024,
+        "intermediate_size": 512,
+        "moe_intermediate_size": 128,
+        "n_routed_experts": 4,
+        "n_shared_experts": 1,
+        "num_experts_per_tok": 2,
+        "moe_layer_freq": 2,
+        "qk_rope_head_dim": 32,
+        "qk_nope_head_dim": 32,
+        "v_head_dim": 64,
+        "q_lora_rank": None,
+        "kv_lora_rank": 64,
+        "rope_theta": 10000.0,
+        "rms_norm_eps": 1e-5,
+        "tie_word_embeddings": True,
+    },
+    "step3p5": {
+        "model_type": "step3p5",
+        "num_hidden_layers": 4,
+        "hidden_size": 256,
+        "num_attention_heads": 8,
+        "num_attention_groups": 4,
+        "head_dim": 32,
+        "vocab_size": 1024,
+        "intermediate_size": 512,
+        "moe_num_experts": 4,
+        "moe_top_k": 2,
+        "moe_intermediate_size": 128,
+        "share_expert_dim": 128,
+        "sliding_window": 64,
+        "rms_norm_eps": 1e-5,
+        "rope_theta": 10000.0,
+        "tie_word_embeddings": True,
+    },
+    "minimax": {
+        "model_type": "minimax",
+        "num_hidden_layers": 2,
+        "hidden_size": 256,
+        "num_attention_heads": 8,
+        "num_key_value_heads": 4,
+        "vocab_size": 1024,
+        "intermediate_size": 512,
+        "num_local_experts": 4,
+        "num_experts_per_tok": 2,
+        "shared_intermediate_size": 256,
+        "max_position_embeddings": 1024,
+        "rms_norm_eps": 1e-5,
+        "rope_theta": 10000.0,
+        "rotary_dim": 32,
+        # QK norm's all_gather pattern requires equal shard sizes across ranks,
+        # incompatible with uneven tp — needs separate fix for padded all_gather
+        "use_qk_norm": False,
+    },
+    "qwen3_moe": {
+        "model_type": "qwen3_moe",
+        "num_hidden_layers": 4,
+        "hidden_size": 256,
+        "num_attention_heads": 8,
+        "num_key_value_heads": 4,
+        "head_dim": 32,
+        "vocab_size": 1024,
+        "intermediate_size": 512,
+        "num_experts": 4,
+        "num_experts_per_tok": 2,
+        "decoder_sparse_step": 2,
+        "mlp_only_layers": [],
+        "moe_intermediate_size": 128,
+        "rms_norm_eps": 1e-5,
+        "rope_theta": 10000.0,
+        "tie_word_embeddings": True,
+        "max_position_embeddings": 1024,
+        "norm_topk_prob": True,
+    },
+    "qwen3_5": {
+        "model_type": "qwen3_5",
+        "text_config": {
+            "model_type": "qwen3_5",
+            "num_hidden_layers": 4,
+            "hidden_size": 256,
+            "num_attention_heads": 8,
+            "num_key_value_heads": 4,
+            "vocab_size": 1024,
+            "intermediate_size": 512,
+            "num_experts": 4,
+            "num_experts_per_tok": 2,
+            "shared_expert_intermediate_size": 256,
+            "moe_intermediate_size": 128,
+            "linear_num_key_heads": 4,
+            "linear_num_value_heads": 4,
+            "linear_key_head_dim": 32,
+            "linear_value_head_dim": 32,
+            "linear_conv_kernel_dim": 4,
+            "full_attention_interval": 2,
+            "rms_norm_eps": 1e-5,
+            "rope_theta": 10000.0,
+            "tie_word_embeddings": True,
+            "max_position_embeddings": 1024,
+            "head_dim": 32,
+        },
+    },
+    "glm4_moe": {
+        "model_type": "glm4_moe",
+        "num_hidden_layers": 4,
+        "hidden_size": 256,
+        "num_attention_heads": 8,
+        "num_key_value_heads": 4,
+        "head_dim": 32,
+        "vocab_size": 1024,
+        "intermediate_size": 512,
+        "moe_intermediate_size": 128,
+        "n_routed_experts": 4,
+        "n_shared_experts": 1,
+        "n_group": 1,
+        "topk_group": 1,
+        "num_experts_per_tok": 2,
+        "first_k_dense_replace": 1,
+        "routed_scaling_factor": 1.0,
+        "max_position_embeddings": 1024,
+        "rms_norm_eps": 1e-5,
+        "rope_theta": 10000.0,
+        "rope_scaling": None,
+        "use_qk_norm": False,
+        "tie_word_embeddings": True,
+        "attention_bias": False,
+        "partial_rotary_factor": 1.0,
+        "norm_topk_prob": True,
+    },
+    "nemotron_h": {
+        "model_type": "nemotron_h",
+        "num_hidden_layers": 4,
+        "hidden_size": 256,
+        "num_attention_heads": 8,
+        "num_key_value_heads": 4,
+        "vocab_size": 1024,
+        "intermediate_size": 512,
+        "max_position_embeddings": 1024,
+        "attention_bias": False,
+        "mamba_num_heads": 8,
+        "mamba_head_dim": 32,
+        "mamba_proj_bias": False,
+        "ssm_state_size": 16,
+        "conv_kernel": 4,
+        "n_groups": 4,
+        "mlp_bias": False,
+        "layer_norm_epsilon": 1e-5,
+        "use_bias": False,
+        "use_conv_bias": True,
+        "head_dim": 32,
+        "hybrid_override_pattern": "M-*E",
+        "n_routed_experts": 4,
+        "num_experts_per_tok": 2,
+        "moe_intermediate_size": 128,
+        "n_group": 1,
+        "topk_group": 1,
+        "norm_topk_prob": True,
+        "routed_scaling_factor": 1.0,
+    },
+    "glm4_moe_lite": {
+        "model_type": "glm4_moe_lite",
+        "num_hidden_layers": 4,
+        "hidden_size": 256,
+        "num_attention_heads": 4,
+        "num_key_value_heads": 4,
+        "vocab_size": 1024,
+        "intermediate_size": 512,
+        "moe_intermediate_size": 128,
+        "n_routed_experts": 4,
+        "n_shared_experts": 1,
+        "num_experts_per_tok": 2,
+        "moe_layer_freq": 2,
+        "qk_rope_head_dim": 32,
+        "qk_nope_head_dim": 32,
+        "v_head_dim": 64,
+        "q_lora_rank": None,
+        "kv_lora_rank": 64,
+        "rope_theta": 10000.0,
+        "rms_norm_eps": 1e-5,
+        "tie_word_embeddings": True,
+    },
+}
+
+
+def _build_model(config):
+    mx.random.seed(RANDOM_SEED)
+    mod = importlib.import_module(f"mlx_lm.models.{config['model_type']}")
+    args = mod.ModelArgs.from_dict(config)
+    model = mod.Model(args)
+    mx.eval(model.parameters())
+    return model
+
+
+def _forward(model, tokens):
+    x = mx.array([tokens])
+    logits = model(x)
+    mx.eval(logits)
+    return np.array(logits[0, -1, :])
+
+
+def _create_hostfile(world_size, base_port):
+    hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)]
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False)
+    json.dump(hosts, f)
+    f.close()
+    return f.name
+
+
+def _run_single_device(config, result_queue):
+    try:
+        model = _build_model(config)
+        logits = _forward(model, INPUT_TOKENS)
+        result_queue.put((0, True, logits))
+    except Exception as e:
+        result_queue.put((0, False, f"{e}\n{traceback.format_exc()}"))
+
+
+def _run_tensor_device(rank, world_size, hostfile_path, config, result_queue):
+    os.environ["MLX_HOSTFILE"] = hostfile_path
+    os.environ["MLX_RANK"] = str(rank)
+
+    try:
+        group = mx.distributed.init(backend="ring", strict=True)
+
+        model = _build_model(config)
+
+        from exo.worker.engines.mlx.auto_parallel import tensor_auto_parallel
+
+        model = tensor_auto_parallel(
+            model, group, timeout_seconds=60.0, on_timeout=None, on_layer_loaded=None
+        )
+
+        logits = _forward(model, INPUT_TOKENS)
+        result_queue.put((rank, True, logits))
+    except Exception as e:
+        result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}"))
+
+
+def _run_single(config):
+    ctx = mp.get_context("spawn")
+    result_queue = ctx.Queue()
+    p = ctx.Process(target=_run_single_device, args=(config, result_queue))
+    p.start()
+    p.join(timeout=60)
+    if p.is_alive():
+        p.terminate()
+        p.join(timeout=5)
+        raise TimeoutError("Single device timed out")
+    rank, success, value = result_queue.get()
+    assert success, f"Single device failed: {value}"
+    return value
+
+
+def _run_tensor(config, world_size, base_port):
+    ctx = mp.get_context("spawn")
+    hostfile_path = _create_hostfile(world_size, base_port)
+    try:
+        result_queue = ctx.Queue()
+        processes = []
+        for rank in range(world_size):
+            p = ctx.Process(
+                target=_run_tensor_device,
+                args=(rank, world_size, hostfile_path, config, result_queue),
+            )
+            p.start()
+            processes.append(p)
+
+        for p in processes:
+            p.join(timeout=120)
+
+        timed_out = any(p.is_alive() for p in processes)
+        for p in processes:
+            if p.is_alive():
+                p.terminate()
+                p.join(timeout=5)
+
+        assert not timed_out, "Tensor parallel timed out"
+
+        results = {}
+        while not result_queue.empty():
+            rank, success, value = result_queue.get()
+            results[rank] = (success, value)
+
+        assert len(results) == world_size, f"Missing results: got {list(results.keys())}"
+        for rank, (success, value) in results.items():
+            assert success, f"Rank {rank} failed: {value}"
+
+        return results[0][1]
+    finally:
+        os.unlink(hostfile_path)
+
+
+class TestComputeShardSizes:
+    def test_even_division(self):
+        assert compute_shard_sizes(64, 2) == [32, 32]
+        assert compute_shard_sizes(64, 4) == [16, 16, 16, 16]
+
+    def test_uneven_division(self):
+        assert compute_shard_sizes(8, 3) == [3, 3, 2]
+        assert compute_shard_sizes(64, 3) == [22, 21, 21]
+        assert compute_shard_sizes(10, 3) == [4, 3, 3]
+
+    def test_sum_invariant(self):
+        for total in [7, 8, 64, 100, 255, 2880]:
+            for n in [2, 3, 5, 7]:
+                sizes = compute_shard_sizes(total, n)
+                assert sum(sizes) == total, f"sum({sizes}) != {total}"
+
+
+class TestWeightSplitMath:
+    def test_all_to_sharded_unquantized(self):
+        mx.random.seed(RANDOM_SEED)
+        weight = mx.random.normal((64, 256))
+        x = mx.random.normal((1, 4, 256))
+
+        full_output = x @ weight.T
+        mx.eval(full_output)
+
+        for n in [2, 3, 5, 7]:
+            sizes = compute_shard_sizes(64, n)
+            indices = list(itertools.accumulate(sizes[:-1]))
+            shards = mx.split(weight, indices, axis=0)
+
+            reconstructed = mx.concatenate([x @ s.T for s in shards], axis=-1)
+            mx.eval(reconstructed)
+
+            diff = float(mx.max(mx.abs(full_output - reconstructed)))
+            assert diff < 1e-5, f"all-to-sharded N={n}: diff={diff}"
+
+    def test_sharded_to_all_unquantized(self):
+        mx.random.seed(RANDOM_SEED)
+        weight = mx.random.normal((128, 256))
+        x = mx.random.normal((1, 4, 256))
+
+        full_output = x @ weight.T
+        mx.eval(full_output)
+
+        for n in [2, 3, 5, 7]:
+            sizes = compute_shard_sizes(256, n)
+            w_indices = list(itertools.accumulate(sizes[:-1]))
+            x_indices = list(itertools.accumulate(sizes[:-1]))
+
+            w_shards = mx.split(weight, w_indices, axis=-1)
+            x_shards = mx.split(x, x_indices, axis=-1)
+
+            partial_outputs = [xs @ ws.T for xs, ws in zip(x_shards, w_shards)]
+            reconstructed = sum(partial_outputs)
+            mx.eval(reconstructed)
+
+            diff = float(mx.max(mx.abs(full_output - reconstructed)))
+            assert diff < 5e-5, f"sharded-to-all N={n}: diff={diff}"
+
+    def test_all_to_sharded_quantized(self):
+        mx.random.seed(RANDOM_SEED)
+        weight = mx.random.normal((64, 256))
+        group_size = 32
+        bits = 4
+        qw, scales, biases = mx.quantize(weight, group_size=group_size, bits=bits)
+        x = mx.random.normal((1, 4, 256))
+
+        full_output = mx.quantized_matmul(
+            x, qw, scales=scales, biases=biases, transpose=True, group_size=group_size, bits=bits
+        )
+        mx.eval(full_output)
+
+        for n in [2, 3]:
+            sizes = compute_shard_sizes(64, n)
+            indices = list(itertools.accumulate(sizes[:-1]))
+
+            qw_shards = mx.split(qw, indices, axis=0)
+            scales_shards = mx.split(scales, indices, axis=0)
+            biases_shards = mx.split(biases, indices, axis=0)
+
+            partial = [
+                mx.quantized_matmul(
+                    x, qw_s, scales=sc_s, biases=bi_s,
+                    transpose=True, group_size=group_size, bits=bits,
+                )
+                for qw_s, sc_s, bi_s in zip(qw_shards, scales_shards, biases_shards)
+            ]
+            reconstructed = mx.concatenate(partial, axis=-1)
+            mx.eval(reconstructed)
+
+            diff = float(mx.max(mx.abs(full_output - reconstructed)))
+            assert diff < 1e-5, f"quantized all-to-sharded N={n}: diff={diff}"
+
+    def test_sharded_to_all_quantized(self):
+        mx.random.seed(RANDOM_SEED)
+        weight = mx.random.normal((128, 256))
+        group_size = 32
+        bits = 4
+        qw, scales, biases = mx.quantize(weight, group_size=group_size, bits=bits)
+        x = mx.random.normal((1, 4, 256))
+
+        full_output = mx.quantized_matmul(
+            x, qw, scales=scales, biases=biases, transpose=True, group_size=group_size, bits=bits
+        )
+        mx.eval(full_output)
+
+        pack_factor = 32 // bits
+        num_quant_groups = scales.shape[-1]
+        for n in [2]:
+            # Split in quantization-group space (same as _shard_quantized_s2a)
+            group_counts = compute_shard_sizes(num_quant_groups, n)
+            weight_ppg = group_size * bits // 32
+
+            packed_sizes = [gc * weight_ppg for gc in group_counts]
+            packed_indices = list(itertools.accumulate(packed_sizes[:-1]))
+            qw_shards = mx.split(qw, packed_indices, axis=-1)
+
+            scale_indices = list(itertools.accumulate(group_counts[:-1]))
+            scales_shards = mx.split(scales, scale_indices, axis=-1)
+            biases_shards = mx.split(biases, scale_indices, axis=-1)
+
+            logical_sizes = [gc * group_size for gc in group_counts]
+            x_indices = list(itertools.accumulate(logical_sizes[:-1]))
+            x_shards = mx.split(x, x_indices, axis=-1)
+
+            partial = [
+                mx.quantized_matmul(
+                    xs, qw_s, scales=sc_s, biases=bi_s,
+                    transpose=True, group_size=group_size, bits=bits,
+                )
+                for xs, qw_s, sc_s, bi_s in zip(x_shards, qw_shards, scales_shards, biases_shards)
+            ]
+            reconstructed = sum(partial)
+            mx.eval(reconstructed)
+
+            diff = float(mx.max(mx.abs(full_output - reconstructed)))
+            assert diff < 1e-4, f"quantized sharded-to-all N={n}: diff={diff}"
+
+
+# Port allocation: 31200-31999 (non-colliding with conftest 29600-29800 and qwen35 29950-31100)
+_BASE_PORT = 31200
+_port_counter = 0
+
+
+def _next_port_block():
+    global _port_counter
+    port = _BASE_PORT + _port_counter * 100
+    _port_counter += 1
+    return port
+
+
+@pytest.mark.slow
+class TestTensorParallelTP2:
+    @pytest.mark.parametrize("model_name", list(REDUCED_CONFIGS.keys()))
+    def test_tp2_matches_single(self, model_name):
+        config = REDUCED_CONFIGS[model_name]
+        single_logits = _run_single(config)
+        tp2_logits = _run_tensor(config, world_size=2, base_port=_next_port_block())
+
+        diff = float(np.max(np.abs(single_logits - tp2_logits)))
+        assert diff < 3e-6, f"{model_name} tp=2 logit diff: {diff}"
+
+
+@pytest.mark.slow
+class TestTensorParallelTP3:
+    @pytest.mark.parametrize("model_name", list(REDUCED_CONFIGS.keys()))
+    def test_tp3_matches_single(self, model_name):
+        config = REDUCED_CONFIGS[model_name]
+        single_logits = _run_single(config)
+        tp3_logits = _run_tensor(config, world_size=3, base_port=_next_port_block())
+
+        diff = float(np.max(np.abs(single_logits - tp3_logits)))
+        assert diff < 3e-6, f"{model_name} tp=3 logit diff: {diff}"
diff --git a/uv.lock b/uv.lock
index 133393ef..a063f5ae 100644
--- a/uv.lock
+++ b/uv.lock
@@ -558,7 +558,7 @@ dependencies = [
     { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx-vlm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -597,7 +597,7 @@ requires-dist = [
     { name = "hypercorn", specifier = ">=0.18.0" },
     { name = "loguru", specifier = ">=0.7.3" },
     { name = "mflux", specifier = "==0.17.2" },
-    { name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
+    { name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding" },
     { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" },
     { name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak" },
     { name = "mlx-vlm", specifier = ">=0.3.11" },
@@ -1436,7 +1436,7 @@ dependencies = [
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1494,8 +1494,8 @@ cuda13 = [
 
 [[package]]
 name = "mlx"
-version = "0.31.2.dev20260324+e5e64331"
-source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }
+version = "0.31.2.dev20260331+71bcd7a2"
+source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'darwin'",
     "python_full_version < '3.14' and sys_platform == 'darwin'",
@@ -1531,7 +1531,7 @@ version = "0.31.2"
 source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#d36e9b661e55a5fc0f77fb6f17ea643aa2dc87aa" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1548,7 +1548,7 @@ dependencies = [
     { name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "miniaudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },