Allow small num_heads

Don't reject uneven placements in placement
Add advanced dashboard elements for tensor sharding and update exo bench
2026-04-01 16:55:38 +01:00 · 2026-04-01 11:46:53 +01:00 · 2026-04-01 11:46:53 +01:00 · 2026-04-01 11:46:53 +01:00 · 2026-04-01 11:46:53 +01:00
18 changed files with 2065 additions and 250 deletions
@@ -9,6 +9,12 @@ import mlx.core as mx
 from base import Module
 from mlx.nn.layers.linear import Linear

+def compute_shard_sizes(
+    dim: int,
+    N: int,
+    unit: int = ...,
+    weights: list[float] | None = ...,
+) -> list[int]: ...
@lru_cache
 def sum_gradients(
    group: mx.distributed.Group,
@@ -20,6 +26,8 @@ def shard_inplace(
    *,
    segments: Union[int, list[int]] = ...,
    group: Optional[mx.distributed.Group] = ...,
+    unit: int = ...,
+    weights: list[float] | None = ...,
 ) -> None:
    """Shard a module in-place by updating its parameter dictionary with the
    sharded parameter dictionary.
@@ -51,6 +59,8 @@ def shard_linear(
    *,
    segments: Union[int, list[int]] = ...,
    group: Optional[mx.distributed.Group] = ...,
+    unit: int = ...,
+    weights: list[float] | None = ...,
 ) -> Linear:
    """Create a new linear layer that has its parameters sharded and also
    performs distributed communication either in the forward or backward
@@ -8,6 +8,9 @@ import mlx.core as mx
 import mlx.nn as nn

 class QuantizedSwitchLinear(nn.Module):
+    weight: mx.array
+    scales: mx.array
+
    def __init__(
        self,
        input_dims: int,
@@ -31,6 +34,8 @@ class QuantizedSwitchLinear(nn.Module):
        ...

 class SwitchLinear(nn.Module):
+    weight: mx.array
+
    def __init__(
        self, input_dims: int, output_dims: int, num_experts: int, bias: bool = ...
    ) -> None: ...
@@ -35,14 +35,17 @@ from harness import (
    ExoHttpError,
    add_common_instance_args,
    capture_cluster_snapshot,
+    get_all_instance_ids,
    instance_id_from_instance,
    node_ids_from_instance,
    nodes_used_in_instance,
    resolve_model_short_id,
    run_planning_phase,
    settle_and_fetch_placements,
+    shard_split_summary,
    wait_for_instance_gone,
    wait_for_instance_ready,
+    wait_for_new_instance,
 )
 from loguru import logger
 from transformers import AutoTokenizer
@@ -378,8 +381,21 @@ def main() -> int:
        default=1.0,
        help="System metrics polling interval in seconds (default: 1.0).",
    )
+    ap.add_argument(
+        "--tensor-strategies",
+        nargs="+",
+        default=["Naive"],
+        help="Tensor shard strategies to benchmark. Choices: Naive, Memory, Compute, Bandwidth, all. Default: Naive.",
+    )
    args = ap.parse_args()

+    tensor_strategies = []
+    for s in args.tensor_strategies:
+        if s.lower() == "all":
+            tensor_strategies = ["Naive", "Memory", "Compute", "Bandwidth"]
+            break
+        tensor_strategies.append(s)
+
    pp_list = parse_int_list(args.pp)
    tg_list = parse_int_list(args.tg)
    if not pp_list or not tg_list:
@@ -467,20 +483,51 @@ def main() -> int:
    all_rows: list[dict[str, Any]] = []
    all_system_metrics: dict[str, dict[str, dict[str, float]]] = {}

+    placement_runs: list[tuple[dict[str, Any], str]] = []
    for preview in selected:
+        sharding = str(preview["sharding"])
+        if sharding == "Tensor":
+            for ts in tensor_strategies:
+                placement_runs.append((preview, ts))
+        else:
+            placement_runs.append((preview, "Naive"))
+
+    for preview, tensor_strategy in placement_runs:
        instance = preview["instance"]
        instance_id = instance_id_from_instance(instance)

        sharding = str(preview["sharding"])
        instance_meta = str(preview["instance_meta"])
        n_nodes = nodes_used_in_instance(instance)
+        strategy_label = (
+            f" / strategy={tensor_strategy}" if sharding == "Tensor" else ""
+        )

        logger.info("=" * 80)
        logger.info(
-            f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
+            f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes}{strategy_label} / instance_id={instance_id}"
        )

-        client.request_json("POST", "/instance", body={"instance": instance})
+        if sharding == "Tensor" and tensor_strategy != "Naive":
+            before_ids = get_all_instance_ids(client)
+            client.request_json(
+                "POST",
+                "/place_instance",
+                body={
+                    "model_id": full_model_id,
+                    "sharding": sharding,
+                    "instance_meta": instance_meta,
+                    "min_nodes": n_nodes,
+                    "tensor_strategy": tensor_strategy,
+                },
+            )
+            instance_id = wait_for_new_instance(client, before_ids)
+            instance = client.get_instance(instance_id)
+            logger.info(f"place_instance created instance_id={instance_id}")
+            logger.info(f"Shard split: {shard_split_summary(instance)}")
+        else:
+            client.request_json("POST", "/instance", body={"instance": instance})
+            logger.info(f"Shard split: {shard_split_summary(instance)}")
        try:
            wait_for_instance_ready(client, instance_id)
        except (RuntimeError, TimeoutError) as e:
@@ -541,6 +588,7 @@ def main() -> int:
                                    "placement_sharding": sharding,
                                    "placement_instance_meta": instance_meta,
                                    "placement_nodes": n_nodes,
+                                    "tensor_strategy": tensor_strategy,
                                    "instance_id": instance_id,
                                    "pp_tokens": actual_pp_tokens,
                                    "tg": tg,
@@ -595,6 +643,7 @@ def main() -> int:
                                        "placement_sharding": sharding,
                                        "placement_instance_meta": instance_meta,
                                        "placement_nodes": n_nodes,
+                                        "tensor_strategy": tensor_strategy,
                                        "instance_id": instance_id,
                                        "pp_tokens": actual_pp_tokens,
                                        "tg": tg,
@@ -656,7 +705,9 @@ def main() -> int:
        finally:
            if sampler:
                sampler.stop()
-                placement_label = f"{sharding}/{instance_meta}/{n_nodes} nodes"
+                placement_label = (
+                    f"{sharding}/{instance_meta}/{n_nodes} nodes{strategy_label}"
+                )
                sampler.print_summary(placement_label)
                placement_metrics = sampler.summarize()
                if placement_metrics:
@@ -131,6 +131,24 @@ def node_ids_from_instance(instance: dict[str, Any]) -> list[str]:
    return list(inner["shardAssignments"]["nodeToRunner"].keys())


+def shard_split_summary(instance: dict[str, Any]) -> str:
+    inner = unwrap_instance(instance)
+    runner_to_shard = inner["shardAssignments"]["runnerToShard"]
+    parts: list[str] = []
+    for rid, shard_wrapper in runner_to_shard.items():
+        shard = next(iter(shard_wrapper.values()))
+        start = shard.get("startLayer", "?")
+        end = shard.get("endLayer", "?")
+        n = shard.get("nLayers", "?")
+        weights = shard.get("shardWeights")
+        rank = shard.get("deviceRank", "?")
+        part = f"rank {rank}: layers [{start},{end})/{n}"
+        if weights:
+            part += f" weights={[round(w, 3) for w in weights]}"
+        parts.append(part)
+    return " | ".join(parts)
+
+
 def runner_ready(runner: dict[str, Any]) -> bool:
    return "RunnerReady" in runner

@@ -151,6 +169,7 @@ def wait_for_instance_ready(
    start_time = time.time()
    instance_existed = False
    last_loaded: dict[str, int] = {}
+    last_states: dict[str, str] = {}
    while time.time() - start_time < timeout:
        instance = client.get_instance(instance_id)

@@ -166,21 +185,29 @@ def wait_for_instance_ready(
        rids = runner_ids_from_instance(instance)

        all_ready = True
+        runner_states: dict[str, str] = {}
        for rid in rids:
            runner = client.get_runner(rid) or {}
            if runner_failed(runner):
                error_msg = get_runner_failed_message(runner) or "Unknown error"
                raise RuntimeError(f"Runner {rid} failed: {error_msg}")
+            state_tag = next(iter(runner), "Unknown")
            if "RunnerLoading" in runner:
                loading = runner["RunnerLoading"]
                loaded = loading.get("layersLoaded", 0)
                total = loading.get("totalLayers", 0)
-                if total > 0 and last_loaded.get(rid) != loaded:
+                if total > 0:
                    last_loaded[rid] = loaded
-                    logger.debug(f"Runner {rid}: loading layers {loaded}/{total}")
+                    state_tag = f"Loading {loaded}/{total}"
+            runner_states[rid] = state_tag
            if not runner_ready(runner):
                all_ready = False

+        if runner_states != last_states:
+            last_states = runner_states
+            parts = [f"{rid[:8]}: {state}" for rid, state in runner_states.items()]
+            logger.debug(f"Runners: {', '.join(parts)}")
+
        if all_ready:
            return

@@ -189,6 +216,24 @@ def wait_for_instance_ready(
    raise TimeoutError(f"Instance {instance_id} did not become ready within {timeout=}")


+def get_all_instance_ids(client: ExoClient) -> set[str]:
+    instances = client.get_state_path("instances") or {}
+    return set(instances.keys())
+
+
+def wait_for_new_instance(
+    client: ExoClient, before_ids: set[str], timeout: float = 60.0
+) -> str:
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        current_ids = get_all_instance_ids(client)
+        new_ids = current_ids - before_ids
+        if new_ids:
+            return next(iter(new_ids))
+        time.sleep(0.2)
+    raise TimeoutError(f"No new instance appeared within {timeout}s")
+
+
 def wait_for_instance_gone(
    client: ExoClient, instance_id: str, timeout: float = 3.0
 ) -> None:
@@ -886,6 +886,8 @@
  }

  let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
+  let selectedTensorStrategy = $state<string>("Naive");
+  let showAdvancedSettings = $state(false);
  type InstanceMeta = "MlxRing" | "MlxJaccl";

  // Launch defaults persistence
@@ -1446,6 +1448,8 @@
            sharding: selectedSharding,
            instance_meta: selectedInstanceType,
            min_nodes: 1,
+            tensor_strategy:
+              selectedSharding === "Tensor" ? selectedTensorStrategy : "Naive",
          }),
        });
      }
@@ -5874,6 +5878,61 @@
                    </div>
                  </div>
                </div>
+
+                <!-- Advanced Settings -->
+                {#if selectedSharding === "Tensor"}
+                  <div class="border-t border-exo-medium-gray/30 pt-3 mt-3">
+                    <button
+                      onclick={() =>
+                        (showAdvancedSettings = !showAdvancedSettings)}
+                      class="flex items-center gap-2 text-xs text-white/40 font-mono hover:text-white/60 transition-colors cursor-pointer"
+                    >
+                      <span
+                        class="transition-transform duration-200 {showAdvancedSettings
+                          ? 'rotate-90'
+                          : ''}">▶</span
+                      >
+                      Advanced
+                    </button>
+                    {#if showAdvancedSettings}
+                      <div class="mt-3 space-y-3">
+                        <div>
+                          <div class="text-xs text-white/50 font-mono mb-2">
+                            Tensor Strategy:
+                          </div>
+                          <div class="flex gap-2 flex-wrap">
+                            {#each ["Naive", "Memory", "Compute", "Bandwidth"] as strategy}
+                              <button
+                                onclick={() => {
+                                  selectedTensorStrategy = strategy;
+                                  saveLaunchDefaults();
+                                }}
+                                class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedTensorStrategy ===
+                                strategy
+                                  ? 'bg-transparent text-exo-yellow border-exo-yellow'
+                                  : 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
+                              >
+                                <span
+                                  class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedTensorStrategy ===
+                                  strategy
+                                    ? 'border-exo-yellow'
+                                    : 'border-exo-medium-gray'}"
+                                >
+                                  {#if selectedTensorStrategy === strategy}
+                                    <span
+                                      class="w-1.5 h-1.5 rounded-full bg-exo-yellow"
+                                    ></span>
+                                  {/if}
+                                </span>
+                                {strategy}
+                              </button>
+                            {/each}
+                          </div>
+                        </div>
+                      </div>
+                    {/if}
+                  </div>
+                {/if}
              {/if}
            </div>

@@ -49,7 +49,7 @@ let
      owner = "rltakashige";
      repo = "mlx-jaccl-fix-small-recv";
      rev = uvLockMlxRev;
-      hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
+      hash = "sha256-WZGQKGcKQR9uyXf5X/a1+79ycPdbcs/spfTykDUjLE4=";
    };

    patches = [
@@ -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 }
@@ -383,6 +383,7 @@ class API:
            sharding=payload.sharding,
            instance_meta=payload.instance_meta,
            min_nodes=payload.min_nodes,
+            tensor_strategy=payload.tensor_strategy,
        )
        await self._send(command)

@@ -10,7 +10,7 @@ from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.text_generation import ReasoningEffort
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding, ShardMetadata
+from exo.shared.types.worker.shards import Sharding, ShardMetadata, TensorShardStrategy
 from exo.utils.pydantic_ext import CamelCaseModel

 FinishReason = Literal[
@@ -253,6 +253,7 @@ class PlaceInstanceParams(BaseModel):
    sharding: Sharding = Sharding.Pipeline
    instance_meta: InstanceMeta = InstanceMeta.MlxRing
    min_nodes: int = 1
+    tensor_strategy: TensorShardStrategy = TensorShardStrategy.Naive


 class CreateInstanceParams(BaseModel):
@@ -298,6 +298,7 @@ class Master:
                                self.state.instances,
                                self.state.node_memory,
                                self.state.node_network,
+                                node_identities=self.state.node_identities,
                                download_status=self.state.downloads,
                            )
                            transition_events = get_transition_events(
@@ -29,7 +29,7 @@ from exo.shared.types.events import (
    TaskStatusUpdated,
 )
 from exo.shared.types.memory import Memory
-from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
+from exo.shared.types.profiling import MemoryUsage, NodeIdentity, NodeNetworkInfo
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
    DownloadCompleted,
@@ -108,6 +108,7 @@ def place_instance(
    current_instances: Mapping[InstanceId, Instance],
    node_memory: Mapping[NodeId, MemoryUsage],
    node_network: Mapping[NodeId, NodeNetworkInfo],
+    node_identities: Mapping[NodeId, NodeIdentity] | None = None,
    required_nodes: set[NodeId] | None = None,
    download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
 ) -> dict[InstanceId, Instance]:
@@ -127,26 +128,12 @@ def place_instance(
    if len(cycles_with_sufficient_memory) == 0:
        raise ValueError("No cycles found with sufficient memory")

-    if command.sharding == Sharding.Tensor:
-        if not command.model_card.supports_tensor:
-            raise ValueError(
-                f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_card.model_id}"
-            )
-        # TODO: the condition here for tensor parallel is not correct, but it works good enough for now.
-        kv_heads = command.model_card.num_key_value_heads
-        cycles_with_sufficient_memory = [
-            cycle
-            for cycle in cycles_with_sufficient_memory
-            if command.model_card.hidden_size % len(cycle) == 0
-            and (kv_heads is None or kv_heads % len(cycle) == 0)
-        ]
-        if not cycles_with_sufficient_memory:
-            raise ValueError(
-                f"No tensor sharding found for model with "
-                f"hidden_size={command.model_card.hidden_size}"
-                f"{f', num_key_value_heads={kv_heads}' if kv_heads is not None else ''}"
-                f" across candidate cycles"
-            )
+    if command.sharding == Sharding.Tensor and not command.model_card.supports_tensor:
+        raise ValueError(
+            f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_card.model_id}"
+        )
+
+    # Uneven tensor sharding handles arbitrary world sizes — no divisibility check needed
    if command.sharding == Sharding.Pipeline and command.model_card.model_id == ModelId(
        "mlx-community/DeepSeek-V3.1-8bit"
    ):
@@ -197,7 +184,12 @@ def place_instance(
        command.sharding = Sharding.Pipeline

    shard_assignments = get_shard_assignments(
-        command.model_card, selected_cycle, command.sharding, node_memory
+        command.model_card,
+        selected_cycle,
+        command.sharding,
+        node_memory,
+        tensor_strategy=command.tensor_strategy,
+        node_identities=node_identities,
    )

    cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle.node_ids)
@@ -6,7 +6,7 @@ from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Topology
 from exo.shared.types.common import Host, NodeId
 from exo.shared.types.memory import Memory
-from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
+from exo.shared.types.profiling import MemoryUsage, NodeIdentity, NodeNetworkInfo
 from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments
 from exo.shared.types.worker.shards import (
@@ -15,8 +15,54 @@ from exo.shared.types.worker.shards import (
    Sharding,
    ShardMetadata,
    TensorShardMetadata,
+    TensorShardMode,
+    TensorShardStrategy,
 )

+# FP32 TFLOPS by chip_id (used for Compute strategy)
+APPLE_SILICON_FLOPS: dict[str, float] = {
+    "Apple M1": 2.6,
+    "Apple M1 Pro": 5.3,
+    "Apple M1 Max": 10.4,
+    "Apple M1 Ultra": 21.0,
+    "Apple M2": 3.6,
+    "Apple M2 Pro": 6.8,
+    "Apple M2 Max": 13.6,
+    "Apple M2 Ultra": 27.2,
+    "Apple M3": 3.5,
+    "Apple M3 Pro": 5.0,
+    "Apple M3 Max": 14.2,
+    "Apple M3 Ultra": 28.4,
+    "Apple M4": 4.3,
+    "Apple M4 Pro": 9.2,
+    "Apple M4 Max": 18.4,
+    "Apple M5": 4.2,
+    "Apple M5 Pro": 8.3,
+    "Apple M5 Max": 19.9,
+}
+
+# Memory bandwidth in GB/s by chip_id (used for Bandwidth strategy)
+APPLE_SILICON_BANDWIDTH: dict[str, float] = {
+    "Apple M1": 68,
+    "Apple M1 Pro": 200,
+    "Apple M1 Max": 400,
+    "Apple M1 Ultra": 800,
+    "Apple M2": 100,
+    "Apple M2 Pro": 200,
+    "Apple M2 Max": 400,
+    "Apple M2 Ultra": 800,
+    "Apple M3": 100,
+    "Apple M3 Pro": 150,
+    "Apple M3 Max": 400,
+    "Apple M3 Ultra": 800,
+    "Apple M4": 120,
+    "Apple M4 Pro": 273,
+    "Apple M4 Max": 546,
+    "Apple M5": 154,
+    "Apple M5 Pro": 307,
+    "Apple M5 Max": 614,
+}
+

 def filter_cycles_by_memory(
    cycles: list[Cycle],
@@ -243,12 +289,39 @@ def _get_shard_assignments_for_pure_pipeline(
 def get_shard_assignments_for_tensor_parallel(
    model_card: ModelCard,
    cycle: Cycle,
+    node_memory: Mapping[NodeId, MemoryUsage] | None = None,
+    strategy: TensorShardStrategy = TensorShardStrategy.Naive,
+    node_identities: Mapping[NodeId, NodeIdentity] | None = None,
 ):
    total_layers = model_card.n_layers
    world_size = len(cycle)
    runner_to_shard: dict[RunnerId, ShardMetadata] = {}
    node_to_runner: dict[NodeId, RunnerId] = {}

+    shard_weights: list[float] | None = None
+    shard_mode = TensorShardMode.Constant
+    match strategy:
+        case TensorShardStrategy.Naive:
+            pass
+        case TensorShardStrategy.Memory:
+            if node_memory is not None:
+                shard_weights = [
+                    node_memory[node_id].ram_available.in_gb for node_id in cycle
+                ]
+            shard_mode = TensorShardMode.Greedy
+        case TensorShardStrategy.Compute:
+            if node_identities is not None:
+                shard_weights = [
+                    APPLE_SILICON_FLOPS.get(node_identities[nid].chip_id, 1.0)
+                    for nid in cycle
+                ]
+        case TensorShardStrategy.Bandwidth:
+            if node_identities is not None:
+                shard_weights = [
+                    APPLE_SILICON_BANDWIDTH.get(node_identities[nid].chip_id, 1.0)
+                    for nid in cycle
+                ]
+
    for i, node_id in enumerate(cycle):
        shard = TensorShardMetadata(
            model_card=model_card,
@@ -257,6 +330,8 @@ def get_shard_assignments_for_tensor_parallel(
            start_layer=0,
            end_layer=total_layers,
            n_layers=total_layers,
+            shard_weights=shard_weights,
+            shard_mode=shard_mode,
        )

        runner_id = RunnerId()
@@ -278,6 +353,8 @@ def get_shard_assignments(
    cycle: Cycle,
    sharding: Sharding,
    node_memory: Mapping[NodeId, MemoryUsage],
+    tensor_strategy: TensorShardStrategy = TensorShardStrategy.Naive,
+    node_identities: Mapping[NodeId, NodeIdentity] | None = None,
 ) -> ShardAssignments:
    match sharding:
        case Sharding.Pipeline:
@@ -290,6 +367,9 @@ def get_shard_assignments(
            return get_shard_assignments_for_tensor_parallel(
                model_card=model_card,
                cycle=cycle,
+                node_memory=node_memory,
+                strategy=tensor_strategy,
+                node_identities=node_identities,
            )


@@ -9,7 +9,7 @@ from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding, ShardMetadata
+from exo.shared.types.worker.shards import Sharding, ShardMetadata, TensorShardStrategy
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel


@@ -38,6 +38,7 @@ class PlaceInstance(BaseCommand):
    sharding: Sharding
    instance_meta: InstanceMeta
    min_nodes: int
+    tensor_strategy: TensorShardStrategy = TensorShardStrategy.Naive


 class CreateInstance(BaseCommand):
@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import TypeAlias, final

-from pydantic import Field
+from pydantic import Field, field_validator

 from exo.shared.models.model_cards import ModelCard
 from exo.utils.pydantic_ext import TaggedModel
@@ -12,6 +12,18 @@ class Sharding(str, Enum):
    Pipeline = "Pipeline"


+class TensorShardMode(str, Enum):
+    Greedy = "Greedy"
+    Constant = "Constant"
+
+
+class TensorShardStrategy(str, Enum):
+    Naive = "Naive"
+    Memory = "Memory"
+    Compute = "Compute"
+    Bandwidth = "Bandwidth"
+
+
 class BaseShardMetadata(TaggedModel):
    """
    Defines a specific shard of the model that is ready to be run on a device.
@@ -76,7 +88,17 @@ class CfgShardMetadata(BaseShardMetadata):

@final
 class TensorShardMetadata(BaseShardMetadata):
-    pass
+    shard_weights: list[float] | None = None
+    shard_mode: TensorShardMode = TensorShardMode.Constant
+
+    @field_validator("shard_mode", mode="before")
+    @classmethod
+    def _coerce_shard_mode(cls, v: object) -> TensorShardMode:
+        if isinstance(v, str):
+            return TensorShardMode(v)
+        if isinstance(v, TensorShardMode):
+            return v
+        raise ValueError(f"expected TensorShardMode or str, got {type(v).__name__}")


 ShardMetadata: TypeAlias = (
@@ -271,9 +271,18 @@ def shard_and_load(

    match shard_metadata:
        case TensorShardMetadata():
-            logger.info(f"loading model from {model_path} with tensor parallelism")
+            logger.info(
+                f"loading model from {model_path} with tensor parallelism "
+                f"(weights={shard_metadata.shard_weights}, mode={shard_metadata.shard_mode})"
+            )
            model = tensor_auto_parallel(
-                model, group, timeout_seconds, on_timeout, on_layer_loaded
+                model,
+                group,
+                timeout_seconds,
+                on_timeout,
+                on_layer_loaded,
+                shard_weights=shard_metadata.shard_weights,
+                shard_mode=shard_metadata.shard_mode,
            )
        case PipelineShardMetadata():
            logger.info(f"loading model from {model_path} with pipeline parallelism")
@@ -0,0 +1,617 @@
+# type: ignore
+import importlib
+import itertools
+import json
+import multiprocessing as mp
+import os
+import tempfile
+import traceback
+
+import mlx.core as mx
+import numpy as np
+import pytest
+from mlx.nn.layers.distributed import compute_shard_sizes
+
+from exo.shared.types.worker.shards import TensorShardMode
+
+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": 2,
+        "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": 2,
+            "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)]
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(hosts, f)
+    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,
+    shard_weights=None,
+    shard_mode=None,
+):
+    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,
+            shard_weights=shard_weights,
+            shard_mode=shard_mode,
+        )
+
+        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, shard_weights=None, shard_mode=None):
+    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,
+                    shard_weights,
+                    shard_mode,
+                ),
+            )
+            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, strict=True)
+            ]
+            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, strict=True
+                )
+            ]
+            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)
+
+        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, strict=True
+                )
+            ]
+            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 = 40000
+_port_counter = 0
+
+
+def _next_port_block():
+    global _port_counter
+    port = _BASE_PORT + _port_counter * 10
+    _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}"
+
+
+@pytest.mark.slow
+class TestWeightedShardingTP2:
+    @pytest.mark.parametrize("model_name", list(REDUCED_CONFIGS.keys()))
+    def test_weighted_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(), shard_weights=[2.0, 1.0]
+        )
+
+        diff = float(np.max(np.abs(single_logits - tp2_logits)))
+        assert diff < 3e-6, f"{model_name} weighted tp=2 logit diff: {diff}"
+
+
+@pytest.mark.slow
+class TestWeightedShardingTP3:
+    @pytest.mark.parametrize("model_name", list(REDUCED_CONFIGS.keys()))
+    def test_weighted_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(),
+            shard_weights=[3.0, 2.0, 1.0],
+        )
+
+        diff = float(np.max(np.abs(single_logits - tp3_logits)))
+        assert diff < 3e-6, f"{model_name} weighted tp=3 logit diff: {diff}"
+
+
+@pytest.mark.slow
+class TestGreedyShardingTP2:
+    @pytest.mark.parametrize("model_name", list(REDUCED_CONFIGS.keys()))
+    def test_greedy_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(),
+            shard_weights=[2.0, 1.0],
+            shard_mode=TensorShardMode.Greedy,
+        )
+
+        diff = float(np.max(np.abs(single_logits - tp2_logits)))
+        assert diff < 3e-6, f"{model_name} greedy tp=2 logit diff: {diff}"
@@ -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.dev20260401+fbfe79c7", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#fbfe79c7b1238273969328abc5720ba18f7265d5" }, 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.dev20260401+fbfe79c7", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#fbfe79c7b1238273969328abc5720ba18f7265d5" }, 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.dev20260401+fbfe79c7"
+source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#fbfe79c7b1238273969328abc5720ba18f7265d5" }
 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.dev20260401+fbfe79c7", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#fbfe79c7b1238273969328abc5720ba18f7265d5" }, 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.dev20260401+fbfe79c7", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#fbfe79c7b1238273969328abc5720ba18f7265d5" }, 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'" },
Author	SHA1	Message	Date
Ryuichi Leo Takashige	03c7b7627d	Allow small num_heads	2026-04-01 16:55:38 +01:00
Ryuichi Leo Takashige	e78454de41	Don't reject uneven placements in placement	2026-04-01 11:46:53 +01:00
Ryuichi Leo Takashige	b30ee156c4	Add advanced dashboard elements for tensor sharding and update exo bench	2026-04-01 11:46:53 +01:00
Ryuichi Leo Takashige	ae7ba5d054	Messy POC	2026-04-01 11:46:53 +01:00
Ryuichi Leo Takashige	55fa5362bb	Add uneven sharding	2026-04-01 11:46:53 +01:00