fix

.mlx_typings/mlx/core/__init__.pyi

@@ -1767,12 +1767,12 @@ def clip(
         array: The clipped array.
     """
 
-def compile(
-    fun: Callable,
+def compile[F: Callable[..., object]](
+    fun: F,
     inputs: object | None = ...,
     outputs: object | None = ...,
     shapeless: bool = ...,
-) -> Callable:
+) -> F:
     """
     Returns a compiled function which produces the same output as ``fun``.
 
@@ -2915,8 +2915,8 @@ def gather_mm(
     a: array,
     b: array,
     /,
-    lhs_indices: array,
-    rhs_indices: array,
+    lhs_indices: array | None = ...,
+    rhs_indices: array | None = ...,
     *,
     sorted_indices: bool = ...,
     stream: Stream | Device | None = ...,
@@ -3683,7 +3683,14 @@ def logsumexp(
         array: The output array with the corresponding axes reduced.
     """
 
-def matmul(a: array, b: array, /, *, stream: Stream | Device | None = ...) -> array:
+def matmul(
+    a: array,
+    b: array,
+    /,
+    *,
+    output_dtype: Dtype | None = ...,
+    stream: Stream | Device | None = ...,
+) -> array:
     """
     Matrix multiplication.
 
@@ -4707,6 +4714,7 @@ def softmax(
     /,
     axis: int | Sequence[int] | None = ...,
     *,
+    precise: bool = ...,
     stream: Stream | Device | None = ...,
 ) -> array:
     """
@@ -5431,6 +5439,12 @@ def zeros_like(a: array, /, *, stream: Stream | Device | None = ...) -> array:
         array: The output array filled with zeros.
     """
 
+def compute_splitk_partitions(m: int, n: int, k: int) -> int:
+    """Return the splitk partition count mlx would use for an (m x k) @ (k x n) matmul."""
+
+def set_splitk_partitions_override(n: int) -> None:
+    """Override the splitk partition count for subsequent matmul dispatches. Pass 0 to clear."""
+
 scalar: TypeAlias = int | float | bool
 list_or_scalar: TypeAlias = scalar | list["list_or_scalar"]
 bool_: Dtype = ...

.mlx_typings/mlx/nn/layers/base.pyi

@@ -57,6 +57,10 @@ class Module(dict):
     def __init__(self) -> None:
         """Should be called by the subclasses of ``Module``."""
 
+    def __getitem__(self, key: str) -> mx.array | Module: ...
+    def get(
+        self, key: str, default: mx.array | Module | None = ...
+    ) -> mx.array | Module | None: ...
     @property
     def training(self):  # -> bool:
         """Boolean indicating if the model is in training mode."""

.mlx_typings/mlx/nn/layers/distributed.pyi

@@ -70,6 +70,10 @@ def shard_linear(
     """
 
 class AllToShardedLinear(Module):
+    weight: mx.array
+    bias: mx.array | None
+    group: mx.distributed.Group
+
     """Each member of the group applies part of the affine transformation such
     that the result is sharded across the group.
 
@@ -102,6 +106,10 @@ class AllToShardedLinear(Module):
     ) -> AllToShardedLinear: ...
 
 class ShardedToAllLinear(Module):
+    weight: mx.array
+    bias: mx.array | None
+    group: mx.distributed.Group
+
     """Each member of the group applies part of the affine transformation and
     then aggregates the results.
 

.mlx_typings/mlx/nn/layers/quantized.pyi

@@ -85,27 +85,14 @@ class QuantizedEmbedding(Module):
         """Create a :obj:`QuantizedEmbedding` layer from an :obj:`Embedding` layer."""
 
 class QuantizedLinear(Module):
-    """Applies an affine transformation to the input using a quantized weight matrix.
+    weight: mx.array
+    scales: mx.array
+    biases: mx.array
+    bias: mx.array | None
+    group_size: int
+    bits: int
+    mode: str
 
-    It is the quantized equivalent of :class:`Linear`. For now its
-    parameters are frozen and will not be included in any gradient computation
-    but this will probably change in the future.
-
-    :obj:`QuantizedLinear` also provides a classmethod :meth:`from_linear` to
-    convert linear layers to :obj:`QuantizedLinear` layers.
-
-    Args:
-        input_dims (int): The dimensionality of the input features.
-        output_dims (int): The dimensionality of the output features.
-        bias (bool, optional): If set to ``False`` then the layer will not use
-            a bias. Default: ``True``.
-        group_size (int, optional): The group size to use for the quantized
-            weight. See :func:`~mlx.core.quantize`. Default: ``64``.
-        bits (int, optional): The bit width to use for the quantized weight.
-            See :func:`~mlx.core.quantize`. Default: ``4``.
-        mode (str): The quantization method to use (see
-           :func:`mlx.core.quantize`). Default: ``"affine"``.
-    """
     def __init__(
         self,
         input_dims: int,

.mlx_typings/mlx_lm/models/base.pyi

@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
 """
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional
 
 import mlx.core as mx
 
@@ -37,10 +37,10 @@ def quantized_scaled_dot_product_attention(
     bits: int = ...,
 ) -> mx.array: ...
 def scaled_dot_product_attention(
-    queries,
-    keys,
-    values,
-    cache,
+    queries: mx.array,
+    keys: mx.array,
+    values: mx.array,
+    cache: Optional[Any],
     scale: float,
     mask: Optional[mx.array],
     sinks: Optional[mx.array] = ...,

.mlx_typings/mlx_lm/models/gpt_oss.pyi

@@ -0,0 +1,103 @@
+"""Type stubs for mlx_lm.models.gpt_oss"""
+
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import KVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    vocab_size: int
+    rms_norm_eps: float
+    sliding_window: int
+    layer_types: Optional[List[str]]
+
+def mlx_topk(a: mx.array, k: int, axis: int = -1) -> tuple[mx.array, mx.array]: ...
+
+class AttentionBlock(nn.Module):
+    head_dim: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_key_value_groups: int
+    sinks: mx.array
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    sm_scale: float
+    rope: nn.Module
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class TransformerBlock(nn.Module):
+    self_attn: AttentionBlock
+    mlp: MLPBlock
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MLPBlock(nn.Module):
+    hidden_size: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    experts: SwitchGLU
+    router: nn.Linear
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class GptOssMoeModel(nn.Module):
+    embed_tokens: nn.Embedding
+    norm: nn.RMSNorm
+    layer_types: List[str]
+    layers: list[TransformerBlock]
+    window_size: int
+    swa_idx: int
+    ga_idx: int
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: GptOssMoeModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[nn.Module]: ...
+    def make_cache(self) -> list[KVCache]: ...

.mlx_typings/mlx_lm/models/minimax.pyi

@@ -0,0 +1,94 @@
+"""Type stubs for mlx_lm.models.minimax"""
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    max_position_embeddings: int
+
+class MiniMaxAttention(nn.Module):
+    num_heads: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    q_norm: nn.Module
+    k_norm: nn.Module
+    rope: nn.Module
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxSparseMoeBlock(nn.Module):
+    num_experts_per_tok: int
+    gate: nn.Linear
+    switch_mlp: SwitchGLU
+    e_score_correction_bias: mx.array
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class MiniMaxDecoderLayer(nn.Module):
+    self_attn: MiniMaxAttention
+    block_sparse_moe: MiniMaxSparseMoeBlock
+    input_layernorm: nn.RMSNorm
+    post_attention_layernorm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxModel(nn.Module):
+    embed_tokens: nn.Embedding
+    layers: list[MiniMaxDecoderLayer]
+    norm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: MiniMaxModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[MiniMaxDecoderLayer]: ...

.mlx_typings/mlx_lm/models/nemotron_h.pyi

@@ -92,6 +92,15 @@ class NemotronHAttention(nn.Module):
         cache: Optional[KVCache] = None,
     ) -> mx.array: ...
 
+class MoEGate(nn.Module):
+    config: ModelArgs
+    top_k: int
+    norm_topk_prob: bool
+    weight: mx.array
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
 class NemotronHMLP(nn.Module):
     up_proj: nn.Linear
     down_proj: nn.Linear
@@ -102,9 +111,14 @@ class NemotronHMLP(nn.Module):
     def __call__(self, x: mx.array) -> mx.array: ...
 
 class NemotronHMoE(nn.Module):
+    config: ModelArgs
     num_experts_per_tok: int
+    moe_latent_size: Optional[int]
     switch_mlp: SwitchMLP
+    gate: MoEGate
     shared_experts: NemotronHMLP
+    fc1_latent_proj: nn.Linear
+    fc2_latent_proj: nn.Linear
 
     def __init__(self, config: ModelArgs) -> None: ...
     def __call__(self, x: mx.array) -> mx.array: ...

.mlx_typings/mlx_lm/models/qwen3_next.pyi

@@ -71,6 +71,7 @@ class Qwen3NextAttention(nn.Module):
 class Qwen3NextSparseMoeBlock(nn.Module):
     norm_topk_prob: bool
     num_experts: int
+    num_experts_per_tok: int
     top_k: int
     gate: nn.Linear
     switch_mlp: SwitchGLU

.mlx_typings/mlx_lm/models/switch_layers.pyi

@@ -3,11 +3,20 @@ This type stub file was generated by pyright.
 """
 
 from functools import partial
+from typing import Callable
 
 import mlx.core as mx
 import mlx.nn as nn
 
 class QuantizedSwitchLinear(nn.Module):
+    weight: mx.array
+    scales: mx.array
+    biases: mx.array
+    bias: mx.array | None
+    group_size: int
+    bits: int
+    mode: str
+
     def __init__(
         self,
         input_dims: int,
@@ -19,42 +28,47 @@ class QuantizedSwitchLinear(nn.Module):
         mode: str = ...,
     ) -> None: ...
     @property
-    def input_dims(self):  # -> int:
-        ...
+    def input_dims(self) -> int: ...
     @property
-    def output_dims(self):  # -> int:
-        ...
+    def output_dims(self) -> int: ...
     @property
-    def num_experts(self):  # -> int:
-        ...
-    def __call__(self, x, indices, sorted_indices=...):  # -> array:
-        ...
+    def num_experts(self) -> int: ...
+    def __call__(
+        self, x: mx.array, indices: mx.array, sorted_indices: bool = ...
+    ) -> mx.array: ...
 
 class SwitchLinear(nn.Module):
+    weight: mx.array
+    bias: mx.array | None
+
     def __init__(
         self, input_dims: int, output_dims: int, num_experts: int, bias: bool = ...
     ) -> None: ...
     @property
-    def input_dims(self):  # -> int:
-        ...
+    def input_dims(self) -> int: ...
     @property
-    def output_dims(self):  # -> int:
-        ...
+    def output_dims(self) -> int: ...
     @property
-    def num_experts(self):  # -> int:
-        ...
-    def __call__(self, x, indices, sorted_indices=...): ...
+    def num_experts(self) -> int: ...
+    def __call__(
+        self, x: mx.array, indices: mx.array, sorted_indices: bool = ...
+    ) -> mx.array: ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ..., mode: str = ...
-    ):  # -> QuantizedSwitchLinear:
-        ...
+    ) -> QuantizedSwitchLinear: ...
 
 @partial(mx.compile, shapeless=True)
-def swiglu(x, gate): ...
+def swiglu(x: mx.array, gate: mx.array) -> mx.array: ...
+def _gather_sort(
+    x: mx.array, indices: mx.array
+) -> tuple[mx.array, mx.array, mx.array]: ...
+def _scatter_unsort(
+    x: mx.array, inv_order: mx.array, shape: tuple[int, ...]
+) -> mx.array: ...
 
 class SwiGLU(nn.Module):
     def __init__(self) -> None: ...
-    def __call__(self, x, gate): ...
+    def __call__(self, x: mx.array, gate: mx.array) -> mx.array: ...
 
 class SwitchGLU(nn.Module):
     gate_proj: SwitchLinear
@@ -67,21 +81,22 @@ class SwitchGLU(nn.Module):
         input_dims: int,
         hidden_dims: int,
         num_experts: int,
-        activation=...,
+        activation: object = ...,
         bias: bool = ...,
     ) -> None: ...
-    def __call__(self, x, indices) -> mx.array: ...
+    def __call__(self, x: mx.array, indices: mx.array) -> mx.array: ...
 
 class SwitchMLP(nn.Module):
     fc1: SwitchLinear
     fc2: SwitchLinear
+    activation: Callable[[mx.array], mx.array]
 
     def __init__(
         self,
         input_dims: int,
         hidden_dims: int,
         num_experts: int,
-        activation=...,
+        activation: object = ...,
         bias: bool = ...,
     ) -> None: ...
-    def __call__(self, x, indices) -> mx.array: ...
+    def __call__(self, x: mx.array, indices: mx.array) -> mx.array: ...

bench/eval_tool_calls.py

@@ -3,11 +3,13 @@ from __future__ import annotations
 
 import argparse
 import contextlib
+import io
 import json
 import os
 import sys
 import time
 import tomllib
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Literal
@@ -209,7 +211,7 @@ def _openai_build_request(
         "model": model,
         "messages": messages,
         "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     return "/v1/chat/completions", body
@@ -276,7 +278,7 @@ def _openai_build_followup(
         "model": model,
         "messages": followup_messages,
         "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     return "/v1/chat/completions", body
@@ -379,7 +381,7 @@ def _claude_build_request(
         "model": model,
         "messages": claude_messages,
         "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     if system_content is not None:
@@ -489,7 +491,7 @@ def _claude_build_followup(
         "model": model,
         "messages": claude_messages,
         "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     if system_content is not None:
@@ -913,6 +915,12 @@ Examples:
         default=1,
         help="Repeat each scenario N times (default: 1)",
     )
+    parser.add_argument(
+        "--concurrency",
+        type=int,
+        default=1,
+        help="Run up to N scenarios in parallel against the same instance (default: 1)",
+    )
     parser.add_argument(
         "--scenarios",
         nargs="*",
@@ -935,6 +943,13 @@ Examples:
     )
     args = parser.parse_args()
 
+    if args.concurrency < 1:
+        print(
+            f"--concurrency must be >= 1 (got {args.concurrency})",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
     all_scenarios = load_scenarios(SCENARIOS_PATH)
     if args.scenarios:
         scenarios = [s for s in all_scenarios if s.name in args.scenarios]
@@ -1010,42 +1025,72 @@ Examples:
     cluster_snapshot = capture_cluster_snapshot(exo)
     all_results: list[ScenarioResult] = []
 
+    tasks: list[tuple[int, Scenario, ApiName]] = [
+        (run_idx, scenario, api_name)
+        for run_idx in range(args.repeat)
+        for scenario in scenarios
+        for api_name in api_names
+    ]
+
+    def _run_one(
+        http_client: httpx.Client,
+        task: tuple[int, Scenario, ApiName],
+    ) -> tuple[tuple[int, Scenario, ApiName], list[ScenarioResult], str]:
+        run_idx, scenario, api_name = task
+        buf = io.StringIO()
+        run_tag = f"[run {run_idx + 1}/{args.repeat}]" if args.repeat > 1 else ""
+        print(
+            f"\n  {run_tag}[{api_name:>9}] {scenario.name}: {scenario.description}",
+            file=buf,
+        )
+        scenario_results = run_scenario(
+            http_client,
+            args.host,
+            args.port,
+            full_model_id,
+            scenario,
+            api_name,
+            args.timeout,
+            args.verbose,
+        )
+        for r in scenario_results:
+            status = "PASS" if r.passed else "FAIL"
+            print(
+                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
+                file=buf,
+            )
+            for check_name, check_ok in r.checks.items():
+                mark = "+" if check_ok else "-"
+                print(f"      {mark} {check_name}", file=buf)
+            if r.error:
+                print(f"      ! {r.error}", file=buf)
+        return task, scenario_results, buf.getvalue()
+
     try:
         with httpx.Client() as http_client:
-            for run_idx in range(args.repeat):
-                if args.repeat > 1:
-                    print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
-
-                for scenario in scenarios:
-                    for api_name in api_names:
-                        print(
-                            f"\n  [{api_name:>9}] {scenario.name}: {scenario.description}",
-                            file=log,
-                        )
-
-                        scenario_results = run_scenario(
-                            http_client,
-                            args.host,
-                            args.port,
-                            full_model_id,
-                            scenario,
-                            api_name,
-                            args.timeout,
-                            args.verbose,
-                        )
+            if args.concurrency == 1:
+                current_run = -1
+                for task in tasks:
+                    run_idx = task[0]
+                    if args.repeat > 1 and run_idx != current_run:
+                        print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
+                        current_run = run_idx
+                    _, scenario_results, buffered = _run_one(http_client, task)
+                    all_results.extend(scenario_results)
+                    log.write(buffered)
+                    log.flush()
+            else:
+                print(
+                    f"Running {len(tasks)} tasks with concurrency={args.concurrency}",
+                    file=log,
+                )
+                with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
+                    futures = [pool.submit(_run_one, http_client, t) for t in tasks]
+                    for fut in as_completed(futures):
+                        _, scenario_results, buffered = fut.result()
                         all_results.extend(scenario_results)
-
-                        for r in scenario_results:
-                            status = "PASS" if r.passed else "FAIL"
-                            print(
-                                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
-                                file=log,
-                            )
-                            for check_name, check_ok in r.checks.items():
-                                mark = "+" if check_ok else "-"
-                                print(f"      {mark} {check_name}", file=log)
-                            if r.error:
-                                print(f"      ! {r.error}", file=log)
+                        log.write(buffered)
+                        log.flush()
     finally:
         try:
             exo.request_json("DELETE", f"/instance/{instance_id}")

pyproject.toml

@@ -70,8 +70,8 @@ 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-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arrayscache-leak" }
+mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "splitk", marker = "sys_platform == 'darwin'" }
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/add-tp-helpers" }
 torch = [
   { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
   { index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },

resources/inference_model_cards/moonshotai--Kimi-K2.6.toml

@@ -0,0 +1,32 @@
+model_id = "moonshotai/Kimi-K2.6"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "kimi"
+quantization = ""
+base_model = "Kimi K2.6"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 595148192736
+
+[vision]
+image_token_id = 163605
+model_type = "kimi_vl"
+processor_repo = "moonshotai/Kimi-K2.6"
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+min_p = 0.01
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults.non_thinking]
+temperature = 0.6
+top_p = 0.95
+min_p = 0.01

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

@@ -7,12 +7,13 @@ from typing import TYPE_CHECKING, Literal, Protocol, cast
 import mlx.core as mx
 import mlx.nn as nn
 from mlx.nn.layers.distributed import (
+    AllToShardedLinear,
     shard_inplace,
     shard_linear,
     sum_gradients,
 )
 from mlx_lm.models.base import (
-    scaled_dot_product_attention,  # pyright: ignore[reportUnknownVariableType]
+    scaled_dot_product_attention,
 )
 from mlx_lm.models.cache import ArraysCache, KVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
@@ -65,6 +66,32 @@ from exo.worker.runner.bootstrap import logger
 if TYPE_CHECKING:
     from mlx_lm.models.cache import Cache
 
+
+def _splitk_override_for_unsharded(m: int, n_full: int, k: int) -> int:
+    n = mx.compute_splitk_partitions(m, n_full, k)
+    return n if n > 0 else -1
+
+
+def _splitk_override_all_to_sharded_call(
+    self: AllToShardedLinear, x: mx.array
+) -> mx.array:
+    """All-to-sharded matmul matching the unsharded kernel's K-reduction."""
+    x = sum_gradients(self.group)(x)
+    weight = cast(mx.array, self["weight"])
+    per_rank_n, k = weight.shape
+    n_full = per_rank_n * self.group.size()
+    m = x.shape[-2] if x.ndim >= 2 else 1
+    mx.set_splitk_partitions_override(_splitk_override_for_unsharded(m, n_full, k))
+    if "bias" in self:
+        y = mx.addmm(cast(mx.array, self["bias"]), x, weight.T)
+    else:
+        y = mx.matmul(x, weight.T)
+    return y
+
+
+AllToShardedLinear.__call__ = _splitk_override_all_to_sharded_call
+
+
 LayerLoadedCallback = Callable[[int, int], None]  # (layers_loaded, total_layers)
 
 
@@ -99,7 +126,7 @@ class CustomMlxLayer(nn.Module):
 
     def __init__(self, original_layer: _LayerCallable):
         super().__init__()
-        dict.__setitem__(self, "_original_layer", original_layer)  # pyright: ignore[reportUnknownMemberType]
+        dict.__setitem__(self, "_original_layer", original_layer)  # type: ignore
 
     @property
     def original_layer(self) -> _LayerCallable:
@@ -178,7 +205,7 @@ class PipelineLastLayer(CustomMlxLayer):
                 # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA)
                 # doesn't have .keys directly; access via first sub-cache.
                 _cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
-                if hasattr(_cache, "keys"):  # pyright: ignore[reportAny]
+                if hasattr(_cache, "keys"):  # type: ignore
                     _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
             mx.eval(output)
             if cache is not None and hasattr(_cache, "keys"):  # type: ignore
@@ -309,24 +336,20 @@ def pipeline_auto_parallel(
     )
 
     if isinstance(inner_model_instance, GptOssMoeModel):
-        inner_model_instance.layer_types = inner_model_instance.layer_types[  # type: ignore
+        inner_model_instance.layer_types = inner_model_instance.layer_types[
             start_layer:end_layer
         ]
         # We can assume the model has at least one layer thanks to placement.
         # If a layer type doesn't exist, we can set it to 0.
         inner_model_instance.swa_idx = (
             0
-            if "sliding_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "sliding_attention"
-            )
+            if "sliding_attention" not in inner_model_instance.layer_types
+            else inner_model_instance.layer_types.index("sliding_attention")
         )
         inner_model_instance.ga_idx = (
             0
-            if "full_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "full_attention"
-            )
+            if "full_attention" not in inner_model_instance.layer_types
+            else inner_model_instance.layer_types.index("full_attention")
         )
 
     if isinstance(inner_model_instance, Step35InnerModel):
@@ -439,17 +462,17 @@ def patch_tensor_model[T](model: T) -> T:
         *args: object,
         **kwargs: object,
     ) -> mx.array:
-        logits: mx.array = original_call(self, *args, **kwargs)  # pyright: ignore[reportAny]
+        logits: mx.array = original_call(self, *args, **kwargs)  # type: ignore
         cache = call_signature.bind_partial(self, *args, **kwargs).arguments.get(
             "cache", None
         )
 
         # Add dependency to last cache entry to ensure distributed ops are evaluated
-        if cache is not None and len(cache) > 0:  # pyright: ignore[reportAny]
-            last = cache[-1]  # pyright: ignore[reportAny]
-            dep_cache = last[0] if hasattr(last, "caches") else last  # pyright: ignore[reportAny]
+        if cache is not None and len(cache) > 0:  # type: ignore
+            last = cache[-1]  # type: ignore
+            dep_cache = last[0] if hasattr(last, "caches") else last  # type: ignore
             if hasattr(dep_cache, "keys"):  # type: ignore
-                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny]
+                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
 
         return logits
 
@@ -462,6 +485,10 @@ def tensor_auto_parallel(
     group: mx.distributed.Group,
     on_layer_loaded: LayerLoadedCallback | None,
 ) -> nn.Module:
+    if not hasattr(mx, "set_splitk_partitions_override"):
+        raise RuntimeError(
+            "TP sharding requires the exo MLX fork (mx.set_splitk_partitions_override missing)"
+        )
     all_to_sharded_linear = partial(
         shard_linear,
         sharding="all-to-sharded",
@@ -637,13 +664,17 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
             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.o_proj = NShardedLinear.from_linear(  # type: ignore
+                layer.self_attn.o_proj, self.group
+            )
             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
 
             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)
+            layer.mlp.down_proj = NShardedLinear.from_linear(  # type: ignore
+                layer.mlp.down_proj, self.group
+            )
             layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
             mx.eval(layer)
             if on_layer_loaded is not None:
@@ -699,7 +730,9 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
                     layer.self_attn.q_b_proj
                 )
 
-            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
+            layer.self_attn.o_proj = NShardedLinear.from_linear(  # type: ignore
+                layer.self_attn.o_proj, self.group
+            )
             layer.self_attn.num_heads //= self.N
 
             # Logic from upstream mlx
@@ -716,24 +749,26 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
             # Shard the MLP
             if isinstance(layer.mlp, (DeepseekV3MLP, DeepseekV32MLP)):
                 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)
+                layer.mlp.down_proj = NShardedLinear.from_linear(  # type: ignore
+                    layer.mlp.down_proj, self.group
+                )
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
 
-            # Shard the MoE.
+            # Shard the MoE with column-sharded down_proj for bit-exactness.
             else:
                 if getattr(layer.mlp, "shared_experts", None) is not None:
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.gate_proj
                     )
-                    self.sharded_to_all_linear_in_place(
-                        layer.mlp.shared_experts.down_proj
-                    )
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.up_proj
                     )
+                    self.all_to_sharded_linear_in_place(
+                        layer.mlp.shared_experts.down_proj
+                    )
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
-                self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
                 layer.mlp.sharding_group = self.group
 
@@ -744,20 +779,155 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
         return model
 
 
+class NShardedLinear(nn.Module):
+    group: mx.distributed.Group
+    quantized: bool
+    weight: mx.array
+    bias: mx.array
+    scales: mx.array
+    biases: mx.array
+    group_size: int
+    bits: int
+    mode: str
+
+    def __init__(
+        self, in_dims: int, out_dims: int, bias: bool, group: mx.distributed.Group
+    ):
+        super().__init__()
+        n = group.size()
+        self.group = group
+        self.quantized = False
+        self.weight = mx.zeros((out_dims // n, in_dims))
+        if bias:
+            self.bias = mx.zeros((out_dims // n,))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x_full = _all_gather_last(x, self.group)
+        if self.quantized:
+            y = mx.quantized_matmul(
+                x_full,
+                cast(mx.array, self["weight"]),
+                cast(mx.array, self["scales"]),
+                cast(mx.array | None, self.get("biases")),
+                transpose=True,
+                group_size=self.group_size,
+                bits=self.bits,
+                mode=self.mode,
+            )
+            if "bias" in self:
+                y = y + cast(mx.array, self["bias"])
+            return _all_gather_last(y, self.group)
+        weight = cast(mx.array, self["weight"])
+        m = x_full.shape[-2] if x_full.ndim >= 2 else 1
+        per_rank_n, k = weight.shape
+        n_full = per_rank_n * self.group.size()
+        mx.set_splitk_partitions_override(_splitk_override_for_unsharded(m, n_full, k))
+        try:
+            if "bias" in self:
+                y = mx.addmm(cast(mx.array, self["bias"]), x_full, weight.T)
+            else:
+                y = mx.matmul(x_full, weight.T)
+        finally:
+            mx.set_splitk_partitions_override(0)
+        return _all_gather_last(y, self.group)
+
+    @classmethod
+    def from_linear(
+        cls, linear: nn.Linear, group: mx.distributed.Group
+    ) -> "NShardedLinear":
+        if isinstance(linear, nn.QuantizedLinear):
+            return cls._from_quantized(linear, group)
+        out_dims, in_dims = linear.weight.shape
+        n = group.size()
+        rank = group.rank()
+        per_rank = out_dims // n
+        instance = cls(in_dims, out_dims, hasattr(linear, "bias"), group)
+        new_weight = cast(mx.array, linear["weight"])[
+            rank * per_rank : (rank + 1) * per_rank
+        ]
+        instance.update({"weight": new_weight})
+        if hasattr(linear, "bias"):
+            new_bias = cast(mx.array, linear["bias"])[
+                rank * per_rank : (rank + 1) * per_rank
+            ]
+            instance.update({"bias": new_bias})
+        return instance
+
+    @classmethod
+    def _from_quantized(
+        cls, linear: nn.QuantizedLinear, group: mx.distributed.Group
+    ) -> "NShardedLinear":
+        out_dims = linear.weight.shape[0]
+        n = group.size()
+        rank = group.rank()
+        per_rank = out_dims // n
+        in_dims = linear.scales.shape[1] * linear.group_size
+        instance = cls(in_dims, out_dims, hasattr(linear, "bias"), group)
+        # Replace the empty bf16 stub weight with the sharded quantized tensors.
+        sl = slice(rank * per_rank, (rank + 1) * per_rank)
+        update: dict[str, mx.array] = {
+            "weight": cast(mx.array, linear["weight"])[sl],
+            "scales": cast(mx.array, linear["scales"])[sl],
+        }
+        if "biases" in linear:
+            update["biases"] = cast(mx.array, linear["biases"])[sl]
+        if "bias" in linear:
+            update["bias"] = cast(mx.array, linear["bias"])[sl]
+        del instance["weight"]
+        if "bias" in instance:
+            del instance["bias"]
+        for k, v in update.items():
+            setattr(instance, k, v)
+        instance.quantized = True
+        instance.group_size = linear.group_size
+        instance.bits = linear.bits
+        instance.mode = linear.mode
+        return instance
+
+
+def _all_gather_last(x: mx.array, group: mx.distributed.Group) -> mx.array:
+    """Fast all_gather over the last axis."""
+    leading = x.shape[:-1]
+    last = x.shape[-1]
+    x2 = x.reshape(-1, last)
+    xt = mx.contiguous(x2.T)
+    g = mx.distributed.all_gather(xt, group=group)
+    return mx.contiguous(g.T).reshape(*leading, last * group.size())
+
+
+class ShardedInputNorm(CustomMlxLayer):
+    def __init__(self, norm: _LayerCallable, group: mx.distributed.Group):
+        super().__init__(norm)
+        self.group = group
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.original_layer(_all_gather_last(x, self.group))
+
+
+class ShardedEmbedding(CustomMlxLayer):
+    def __init__(self, embed: _LayerCallable, group: mx.distributed.Group):
+        super().__init__(embed)
+        self.group = group
+
+    def __call__(self, ids: mx.array) -> mx.array:
+        y = self.original_layer(ids)
+        n = self.group.size()
+        per_rank = y.shape[-1] // n
+        rank = self.group.rank()
+        return y[..., rank * per_rank : (rank + 1) * per_rank]
+
+
 class ShardedMoE(CustomMlxLayer):
-    """Wraps any MoE layer with distributed sum_gradients / all_sum."""
+    """Delegates to the wrapped MoE block's ``call_sharded`` when a group is set."""
 
     def __init__(self, layer: _LayerCallable):
         super().__init__(layer)
         self.sharding_group: mx.distributed.Group | None = None
 
     def __call__(self, x: mx.array) -> mx.array:
-        if self.sharding_group is not None:
-            x = sum_gradients(self.sharding_group)(x)
-        y = self.original_layer.__call__(x)
-        if self.sharding_group is not None:
-            y = mx.distributed.all_sum(y, group=self.sharding_group)
-        return y
+        if self.sharding_group is None:
+            return self.original_layer.__call__(x)
+        return self.original_layer.call_sharded(x, self.sharding_group)  # type: ignore
 
 
 class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
@@ -780,7 +950,9 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
                     layer.self_attn.q_b_proj
                 )
 
-            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
+            layer.self_attn.o_proj = NShardedLinear.from_linear(  # type: ignore
+                layer.self_attn.o_proj, self.group
+            )
             layer.self_attn.num_heads //= self.N
 
             # Logic from upstream mlx
@@ -796,7 +968,9 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
 
             if isinstance(layer.mlp, Glm4MoeLiteMLP):
                 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)
+                layer.mlp.down_proj = NShardedLinear.from_linear(  # type: ignore
+                    layer.mlp.down_proj, self.group
+                )
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
 
             else:
@@ -804,15 +978,15 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.gate_proj
                     )
-                    self.sharded_to_all_linear_in_place(
-                        layer.mlp.shared_experts.down_proj
-                    )
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.up_proj
                     )
+                    self.all_to_sharded_linear_in_place(
+                        layer.mlp.shared_experts.down_proj
+                    )
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
-                self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
                 layer.mlp.sharding_group = self.group  # type: ignore
             mx.eval(layer)
@@ -891,7 +1065,7 @@ class WrappedMiniMaxAttention(CustomMlxLayer):
             keys,
             values,
             cache=cache,
-            scale=self._original_layer.scale,  # type: ignore
+            scale=self._original_layer.scale,
             mask=mask,
         )
 
@@ -914,25 +1088,27 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             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.o_proj = NShardedLinear.from_linear(  # type: ignore
+                layer.self_attn.o_proj, self.group
+            )
 
             layer.self_attn.num_attention_heads //= self.N
             layer.self_attn.num_key_value_heads //= self.N
 
-            layer.self_attn = WrappedMiniMaxAttention(layer.self_attn, self.group)  # pyright: ignore[reportAttributeAccessIssue,reportArgumentType]
+            layer.self_attn = WrappedMiniMaxAttention(layer.self_attn, self.group)  # type: ignore
 
             # Shard the MoE.
             self.all_to_sharded_linear_in_place(
                 layer.block_sparse_moe.switch_mlp.gate_proj
             )
-            self.sharded_to_all_linear_in_place(
-                layer.block_sparse_moe.switch_mlp.down_proj
-            )
             self.all_to_sharded_linear_in_place(
                 layer.block_sparse_moe.switch_mlp.up_proj
             )
-            layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
-            layer.block_sparse_moe.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
+            self.all_to_sharded_linear_in_place(
+                layer.block_sparse_moe.switch_mlp.down_proj
+            )
+            layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe)  # type: ignore
+            layer.block_sparse_moe.sharding_group = self.group
             mx.eval(layer)
             if on_layer_loaded is not None:
                 on_layer_loaded(i, total)
@@ -968,8 +1144,8 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.self_attn.v_proj = self.all_to_sharded_linear(
                     layer.self_attn.v_proj
                 )
-                layer.self_attn.o_proj = self.sharded_to_all_linear(
-                    layer.self_attn.o_proj
+                layer.self_attn.o_proj = NShardedLinear.from_linear(  # type: ignore
+                    layer.self_attn.o_proj, self.group
                 )
                 layer.self_attn.n_heads //= self.N
                 layer.self_attn.n_kv_heads //= self.N
@@ -1007,8 +1183,8 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                         linear_attn.in_proj_a = self.all_to_sharded_linear(
                             linear_attn.in_proj_a
                         )
-                    linear_attn.out_proj = self.sharded_to_all_linear(
-                        linear_attn.out_proj
+                    linear_attn.out_proj = NShardedLinear.from_linear(  # type: ignore
+                        linear_attn.out_proj, self.group
                     )
 
                     # Shard conv1d: depthwise conv with non-contiguous channel slicing.
@@ -1061,13 +1237,17 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                     layer.self_attn.v_proj = self.all_to_sharded_linear(
                         layer.self_attn.v_proj
                     )
-                    layer.self_attn.o_proj = self.sharded_to_all_linear(
-                        layer.self_attn.o_proj
+                    layer.self_attn.o_proj = NShardedLinear.from_linear(  # type: ignore
+                        layer.self_attn.o_proj, self.group
                     )
                     layer.self_attn.num_attention_heads //= self.N
                     layer.self_attn.num_key_value_heads //= self.N
 
-            # Shard the MoE.
+            # Shard the MoE. Down_proj is column-sharded (output dim) so each
+            # per-rank matmul runs the full K-reduction and its bf16 output is
+            # bit-exact per column to tp=1. ShardedMoE.__call__ inserts the
+            # required all_gather of the intermediate before down_proj and of
+            # the output after.
             if isinstance(
                 layer.mlp,
                 (
@@ -1077,25 +1257,27 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 ),
             ):
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
-                self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 if isinstance(
                     layer.mlp, (Qwen3NextSparseMoeBlock, Qwen3_5SparseMoeBlock)
                 ):
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_expert.gate_proj
                     )
-                    self.sharded_to_all_linear_in_place(
+                    self.all_to_sharded_linear_in_place(layer.mlp.shared_expert.up_proj)
+                    self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_expert.down_proj
                     )
-                    self.all_to_sharded_linear_in_place(layer.mlp.shared_expert.up_proj)
-                layer.mlp = ShardedMoE(layer.mlp)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
+                layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
                 layer.mlp.sharding_group = self.group
 
             # Shard the MLP
             else:
                 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)
+                layer.mlp.down_proj = NShardedLinear.from_linear(  # type: ignore
+                    layer.mlp.down_proj, self.group
+                )
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
 
             mx.eval(layer)
@@ -1118,30 +1300,33 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
             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.o_proj = NShardedLinear.from_linear(  # type: ignore
+                layer.self_attn.o_proj, self.group
+            )
             layer.self_attn.n_heads //= self.N
             layer.self_attn.n_kv_heads //= self.N
 
             if isinstance(layer.mlp, MoE):
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
-                self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 if getattr(layer.mlp, "shared_experts", None) is not None:
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.gate_proj
                     )
-                    self.sharded_to_all_linear_in_place(
-                        layer.mlp.shared_experts.down_proj
-                    )
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.up_proj
                     )
-                layer.mlp = ShardedMoE(layer.mlp)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
+                    self.all_to_sharded_linear_in_place(
+                        layer.mlp.shared_experts.down_proj
+                    )
                 layer.mlp.sharding_group = self.group
 
             else:
                 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)
+                layer.mlp.down_proj = NShardedLinear.from_linear(  # type: ignore
+                    layer.mlp.down_proj, self.group
+                )
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
 
             mx.eval(layer)
@@ -1164,7 +1349,9 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
             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.o_proj = NShardedLinear.from_linear(  # type: ignore
+                layer.self_attn.o_proj, self.group
+            )
 
             layer.self_attn.num_attention_heads //= self.N
             layer.self_attn.num_key_value_heads //= self.N
@@ -1180,11 +1367,11 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
             ]
 
             self.all_to_sharded_linear_in_place(layer.mlp.experts.gate_proj)
-            self.sharded_to_all_linear_in_place(layer.mlp.experts.down_proj)
             self.all_to_sharded_linear_in_place(layer.mlp.experts.up_proj)
+            self.all_to_sharded_linear_in_place(layer.mlp.experts.down_proj)
 
             layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
-            layer.mlp.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
+            layer.mlp.sharding_group = self.group
             mx.eval(layer)
             if on_layer_loaded is not None:
                 on_layer_loaded(i, total)
@@ -1205,7 +1392,9 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
             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.o_proj = NShardedLinear.from_linear(  # type: ignore
+                layer.self_attn.o_proj, self.group
+            )
 
             layer.self_attn.num_heads //= self.N
             layer.self_attn.num_kv_heads //= self.N
@@ -1218,15 +1407,18 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
             if isinstance(layer.mlp, Step35MLP):
                 layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
-                layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
+                layer.mlp.down_proj = NShardedLinear.from_linear(  # type: ignore
+                    layer.mlp.down_proj, self.group
+                )
             else:
-                layer.mlp.sharding_group = self.group
                 self.all_to_sharded_linear_in_place(layer.mlp.share_expert.gate_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.share_expert.up_proj)
-                self.sharded_to_all_linear_in_place(layer.mlp.share_expert.down_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.share_expert.down_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
-                self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.down_proj)
+                layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
+                layer.mlp.sharding_group = self.group
 
             mx.eval(layer)
             if on_layer_loaded is not None:
@@ -1252,7 +1444,7 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
                 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.o_proj = NShardedLinear.from_linear(mixer.o_proj, self.group)  # type: ignore
                 mixer.num_heads //= self.N
                 mixer.num_key_value_heads //= self.N
 
@@ -1260,16 +1452,19 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
                 self._shard_mamba2_mixer(mixer, rank)
 
             elif isinstance(mixer, NemotronHMoE):
-                # Shard routed experts (SwitchMLP uses fc1/fc2)
+                # N-shard both fc1 and fc2 so each per-rank matmul runs the
+                # full K reduction (bit-exact per column). ShardedMoE does the
+                # all_gather of the intermediate between fc1 and fc2 and of the
+                # fc2 output.
                 self.all_to_sharded_linear_in_place(mixer.switch_mlp.fc1)
-                self.sharded_to_all_linear_in_place(mixer.switch_mlp.fc2)
-                # Shard shared expert in-place (no all-reduce — ShardedMoE handles that)
+                self.all_to_sharded_linear_in_place(mixer.switch_mlp.fc2)
                 if hasattr(mixer, "shared_experts"):
+                    self.all_to_sharded_linear_in_place(mixer.shared_experts.gate_proj)  # type: ignore
                     self.all_to_sharded_linear_in_place(mixer.shared_experts.up_proj)
-                    self.sharded_to_all_linear_in_place(mixer.shared_experts.down_proj)
-                mixer = ShardedMoE(mixer)  # pyright: ignore[reportArgumentType]
+                    self.all_to_sharded_linear_in_place(mixer.shared_experts.down_proj)
+                mixer = ShardedMoE(mixer)  # type: ignore
                 mixer.sharding_group = self.group
-                layer.mixer = mixer  # pyright: ignore[reportAttributeAccessIssue]
+                layer.mixer = mixer  # type: ignore
 
             mx.eval(layer)
             if on_layer_loaded is not None:
@@ -1320,7 +1515,7 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
         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 = NShardedLinear.from_linear(mixer.out_proj, self.group)  # type: ignore
 
         # === conv1d: depthwise conv on conv_dim channels ===
         # conv_dim layout: [ssm_hidden:IS | B:NG*SS | C:NG*SS]
@@ -1368,12 +1563,11 @@ class WrappedGemma4Experts(CustomMlxLayer):
     def __call__(
         self, x: mx.array, top_k_indices: mx.array, top_k_weights: mx.array
     ) -> mx.array:
-        if self.sharding_group is not None:
-            x = sum_gradients(self.sharding_group)(x)
-        y: mx.array = self.original_layer(x, top_k_indices, top_k_weights)
-        if self.sharding_group is not None:
-            y = mx.distributed.all_sum(y, group=self.sharding_group)
-        return y
+        if self.sharding_group is None:
+            return self.original_layer(x, top_k_indices, top_k_weights)
+        return self.original_layer.call_sharded(  # type: ignore
+            x, top_k_indices, top_k_weights, self.sharding_group
+        )
 
 
 class Gemma4ShardingStrategy(TensorParallelShardingStrategy):
@@ -1393,19 +1587,21 @@ class Gemma4ShardingStrategy(TensorParallelShardingStrategy):
             attn.k_proj = self.all_to_sharded_linear(attn.k_proj)
             if not attn.use_k_eq_v:
                 attn.v_proj = self.all_to_sharded_linear(attn.v_proj)
-            attn.o_proj = self.sharded_to_all_linear(attn.o_proj)
+            attn.o_proj = NShardedLinear.from_linear(attn.o_proj, self.group)  # type: ignore
             attn.n_heads //= self.N
             attn.n_kv_heads //= self.N
 
             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)
+            layer.mlp.down_proj = NShardedLinear.from_linear(  # type: ignore
+                layer.mlp.down_proj, self.group
+            )
             layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
 
             if layer.enable_moe:
                 self.all_to_sharded_linear_in_place(layer.experts.switch_glu.gate_proj)
-                self.sharded_to_all_linear_in_place(layer.experts.switch_glu.down_proj)
                 self.all_to_sharded_linear_in_place(layer.experts.switch_glu.up_proj)
-                layer.experts = WrappedGemma4Experts(layer.experts)  # pyright: ignore[reportAttributeAccessIssue,reportArgumentType]
+                self.all_to_sharded_linear_in_place(layer.experts.switch_glu.down_proj)
+                layer.experts = WrappedGemma4Experts(layer.experts)  # type: ignore
                 layer.experts.sharding_group = self.group
 
             mx.eval(layer)

src/exo/worker/tests/unittests/test_mlx/test_tp_bit_exact.py

@@ -0,0 +1,394 @@
+# type: ignore
+"""uv run pytest -v -m "" src/exo/worker/tests/unittests/test_mlx/test_tp_bit_exact.py"""
+
+import importlib
+import json
+import multiprocessing as mp
+import os
+import sys
+import tempfile
+import traceback
+
+import numpy as np
+import pytest
+
+MODEL_CONFIGS = {
+    "llama": dict(
+        module="mlx_lm.models.llama",
+        args=dict(
+            model_type="llama",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            rms_norm_eps=1e-6,
+            vocab_size=512,
+            max_position_embeddings=128,
+            head_dim=32,
+            rope_theta=10000.0,
+        ),
+    ),
+    "qwen3_5_moe": dict(
+        module="mlx_lm.models.qwen3_5_moe",
+        args=dict(
+            model_type="qwen3_5_moe",
+            text_config=dict(
+                model_type="qwen3_5_moe",
+                vocab_size=512,
+                hidden_size=512,
+                intermediate_size=1024,
+                num_hidden_layers=4,
+                num_attention_heads=16,
+                num_key_value_heads=4,
+                head_dim=32,
+                max_position_embeddings=128,
+                rms_norm_eps=1e-6,
+                tie_word_embeddings=False,
+                attention_bias=False,
+                full_attention_interval=2,
+                linear_num_value_heads=32,
+                linear_num_key_heads=16,
+                linear_key_head_dim=32,
+                linear_value_head_dim=32,
+                linear_conv_kernel_dim=4,
+                num_experts=16,
+                num_experts_per_tok=2,
+                decoder_sparse_step=1,
+                shared_expert_intermediate_size=256,
+                moe_intermediate_size=256,
+                norm_topk_prob=True,
+                rope_parameters={
+                    "type": "default",
+                    "rope_theta": 10000.0,
+                    "partial_rotary_factor": 0.25,
+                    "mrope_section": [11, 11, 10],
+                },
+            ),
+        ),
+    ),
+    "qwen3_next": dict(
+        module="mlx_lm.models.qwen3_next",
+        args=dict(
+            model_type="qwen3_next",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=4,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            head_dim=32,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            vocab_size=512,
+            attention_bias=False,
+            full_attention_interval=2,
+            linear_num_value_heads=32,
+            linear_num_key_heads=16,
+            linear_key_head_dim=32,
+            linear_value_head_dim=32,
+            linear_conv_kernel_dim=4,
+            num_experts=16,
+            num_experts_per_tok=2,
+            decoder_sparse_step=1,
+            shared_expert_intermediate_size=256,
+            moe_intermediate_size=256,
+            norm_topk_prob=True,
+            mlp_only_layers=[],
+            rope_theta=10000.0,
+            partial_rotary_factor=0.25,
+        ),
+    ),
+    "deepseek_v3": dict(
+        module="mlx_lm.models.deepseek_v3",
+        args=dict(
+            model_type="deepseek_v3",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=16,
+            vocab_size=512,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            n_routed_experts=8,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            moe_intermediate_size=256,
+            moe_layer_freq=1,
+            first_k_dense_replace=0,
+            n_group=1,
+            topk_group=1,
+            routed_scaling_factor=1.0,
+            q_lora_rank=None,
+            kv_lora_rank=16,
+            qk_nope_head_dim=16,
+            qk_rope_head_dim=16,
+            v_head_dim=32,
+            rope_theta=10000.0,
+            rope_scaling={},
+            attention_bias=False,
+            norm_topk_prob=True,
+            scoring_func="sigmoid",
+            topk_method="noaux_tc",
+        ),
+    ),
+    "deepseek_v3_q4": dict(
+        module="mlx_lm.models.deepseek_v3",
+        quantize=dict(group_size=32, bits=4, mode="affine"),
+        args=dict(
+            model_type="deepseek_v3",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=16,
+            vocab_size=512,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            n_routed_experts=8,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            moe_intermediate_size=256,
+            moe_layer_freq=1,
+            first_k_dense_replace=0,
+            n_group=1,
+            topk_group=1,
+            routed_scaling_factor=1.0,
+            q_lora_rank=None,
+            kv_lora_rank=64,
+            qk_nope_head_dim=32,
+            qk_rope_head_dim=32,
+            v_head_dim=32,
+            rope_theta=10000.0,
+            rope_scaling={},
+            attention_bias=False,
+            norm_topk_prob=True,
+            scoring_func="sigmoid",
+            topk_method="noaux_tc",
+        ),
+    ),
+    "glm4_moe_lite": dict(
+        module="mlx_lm.models.glm4_moe_lite",
+        args=dict(
+            model_type="glm4_moe_lite",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=16,
+            vocab_size=512,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            n_routed_experts=8,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            moe_intermediate_size=256,
+            first_k_dense_replace=1,
+            n_group=1,
+            topk_group=1,
+            routed_scaling_factor=1.0,
+            rope_theta=10000.0,
+            attention_bias=False,
+            q_lora_rank=None,
+            kv_lora_rank=16,
+            qk_rope_head_dim=16,
+            qk_nope_head_dim=16,
+            v_head_dim=32,
+        ),
+    ),
+    "minimax": dict(
+        module="mlx_lm.models.minimax",
+        args=dict(
+            model_type="minimax",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            max_position_embeddings=128,
+            num_experts_per_tok=2,
+            num_local_experts=8,
+            shared_intermediate_size=256,
+            num_hidden_layers=2,
+            rms_norm_eps=1e-6,
+            rope_theta=10000.0,
+            rotary_dim=32,
+            vocab_size=512,
+        ),
+    ),
+    "gpt_oss": dict(
+        module="mlx_lm.models.gpt_oss",
+        args=dict(
+            model_type="gpt_oss",
+            hidden_size=512,
+            intermediate_size=256,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            vocab_size=512,
+            head_dim=32,
+            rms_norm_eps=1e-6,
+            num_local_experts=8,
+            num_experts_per_tok=2,
+            layer_types=["sliding_attention", "full_attention"],
+            sliding_window=64,
+            rope_theta=10000.0,
+        ),
+    ),
+    "gemma4": dict(
+        module="mlx_lm.models.gemma4",
+        args=dict(
+            model_type="gemma4",
+            vocab_size=512,
+            text_config=dict(
+                vocab_size=512,
+                hidden_size=512,
+                intermediate_size=1024,
+                num_hidden_layers=4,
+                num_attention_heads=16,
+                num_key_value_heads=4,
+                head_dim=32,
+                global_head_dim=32,
+                num_kv_shared_layers=0,
+                vocab_size_per_layer_input=512,
+                hidden_size_per_layer_input=512,
+                rms_norm_eps=1e-6,
+                max_position_embeddings=128,
+                sliding_window=64,
+                sliding_window_pattern=2,
+                layer_types=[
+                    "sliding_attention",
+                    "full_attention",
+                    "sliding_attention",
+                    "full_attention",
+                ],
+                enable_moe_block=True,
+                num_experts=8,
+                top_k_experts=2,
+                moe_intermediate_size=256,
+            ),
+        ),
+    ),
+}
+
+_PROMPT = [[1, 23, 45, 67, 89, 12, 34, 56]]
+
+
+def _build(name):
+    import mlx.core as mx
+    import mlx.nn as nn
+    from mlx.utils import tree_map_with_path
+
+    import exo.worker.engines.mlx.auto_parallel  # noqa: F401
+
+    cfg = MODEL_CONFIGS[name]
+    module = importlib.import_module(cfg["module"])
+    model_cls = module.Model
+    model_args_cls = module.ModelArgs
+
+    mx.random.seed(0)
+    args = model_args_cls(**cfg["args"])
+    m = model_cls(args)
+
+    def _to_bf16(_p, v):
+        if hasattr(v, "dtype") and v.dtype in (mx.float16, mx.float32, mx.bfloat16):
+            return v.astype(mx.bfloat16)
+        return v
+
+    m.update(tree_map_with_path(_to_bf16, m.parameters()))
+    if "quantize" in cfg:
+        nn.quantize(m, **cfg["quantize"])
+    mx.eval(m.parameters())
+    return mx, m
+
+
+def _run(name, out_path, shard):
+    import mlx.core as mx
+
+    if shard:
+        g = mx.distributed.init(backend="ring", strict=True)
+    mx_, m = _build(name)
+    if shard:
+        from exo.worker.engines.mlx.auto_parallel import tensor_auto_parallel
+
+        m = tensor_auto_parallel(m, g, on_layer_loaded=None)
+        mx_.eval(m.parameters())
+    inputs = mx_.array(_PROMPT, dtype=mx_.int32)
+    logits = m(inputs)
+    mx_.eval(logits)
+    np.savez(out_path, logits=np.asarray(logits.astype(mx_.float32)))
+
+
+def _ref_worker(name, out_path, q):
+    try:
+        _run(name, out_path, shard=False)
+        q.put(True)
+    except BaseException as e:
+        q.put(f"{e}\n{traceback.format_exc()}")
+
+
+def _tp_worker(name, rank, hf, out_path, q):
+    os.environ["MLX_HOSTFILE"] = hf
+    os.environ["MLX_RANK"] = str(rank)
+    try:
+        path = out_path if rank == 0 else out_path + f".r{rank}"
+        _run(name, path, shard=True)
+        q.put((rank, True, None))
+    except BaseException as e:
+        q.put((rank, False, f"{e}\n{traceback.format_exc()}"))
+
+
+def _run_compare(name, world_size, port_base):
+    d = tempfile.mkdtemp()
+    ref_path = f"{d}/ref.npz"
+    tp_path = f"{d}/tp.npz"
+    ctx = mp.get_context("spawn")
+    q = ctx.Queue()
+
+    p = ctx.Process(target=_ref_worker, args=(name, ref_path, q))
+    p.start()
+    p.join(300)
+    r = q.get(timeout=10)
+    if r is not True:
+        pytest.fail(f"[{name}] ref FAIL: {str(r)[:500]}")
+
+    hosts = [f"127.0.0.1:{port_base + i}" for i in range(world_size)]
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(hosts, f)
+        hf = f.name
+    ps = [
+        ctx.Process(target=_tp_worker, args=(name, rank, hf, tp_path, q))
+        for rank in range(world_size)
+    ]
+    for pp in ps:
+        pp.start()
+    results = [q.get(timeout=300) for _ in range(world_size)]
+    for pp in ps:
+        pp.join(60)
+    for rank, ok, payload in results:
+        if not ok:
+            pytest.fail(f"[{name}] rank {rank} FAIL: {payload[:500]}")
+
+    ref = np.load(ref_path)["logits"]
+    tp = np.load(tp_path)["logits"]
+    diff = np.abs(ref - tp)
+    max_diff = float(diff.max())
+    mean_diff = float(diff.mean())
+    assert max_diff == 0.0, (
+        f"[{name} TP={world_size}] not bit-exact: max={max_diff} mean={mean_diff}"
+    )
+
+
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.skipif(
+        sys.platform != "darwin", reason="MLX distributed requires Metal"
+    ),
+]
+
+
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("name", list(MODEL_CONFIGS))
+def test_tp_bit_exact(name, world_size):
+    name_idx = list(MODEL_CONFIGS).index(name)
+    port = 32000 + name_idx * 20 + world_size
+    _run_compare(name, world_size, port)

uv.lock

@@ -485,7 +485,7 @@ dependencies = [
     { name = "hypercorn", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mflux", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx", version = "0.31.2.dev20260420+553a7adb", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#553a7adbb20ed1b71fe643f4075982a639aaef1b" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx", version = "0.31.2.dev20260421+b27d5e35", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=splitk#b27d5e35da9d0fae4f70918830a7f52b2b81d131" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx-vlm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -538,11 +538,11 @@ requires-dist = [
     { name = "hypercorn", specifier = ">=0.18.0" },
     { name = "loguru", specifier = ">=0.7.3" },
     { name = "mflux", marker = "sys_platform == 'darwin'", 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=splitk" },
     { name = "mlx-cpu", marker = "sys_platform == 'linux' and extra == 'cpu'", specifier = "==0.31.1" },
     { name = "mlx-cuda-12", marker = "sys_platform == 'linux' and extra == 'cuda12'", specifier = "==0.31.1" },
     { name = "mlx-cuda-13", marker = "sys_platform == 'linux' and extra == 'cuda13'", specifier = "==0.31.1" },
-    { name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak" },
+    { name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fadd-tp-helpers" },
     { name = "mlx-vlm", specifier = ">=0.3.11" },
     { name = "msgspec", specifier = ">=0.19.0" },
     { name = "nanobind", marker = "extra == 'build'" },
@@ -1437,7 +1437,7 @@ dependencies = [
     { name = "hf-transfer", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "matplotlib", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx", version = "0.31.2.dev20260420+553a7adb", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#553a7adbb20ed1b71fe643f4075982a639aaef1b" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx", version = "0.31.2.dev20260421+b27d5e35", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=splitk#b27d5e35da9d0fae4f70918830a7f52b2b81d131" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "numpy", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "piexif", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -1493,8 +1493,8 @@ wheels = [
 
 [[package]]
 name = "mlx"
-version = "0.31.2.dev20260420+553a7adb"
-source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#553a7adbb20ed1b71fe643f4075982a639aaef1b" }
+version = "0.31.2.dev20260421+b27d5e35"
+source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=splitk#b27d5e35da9d0fae4f70918830a7f52b2b81d131" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'darwin'",
     "python_full_version < '3.14' and sys_platform == 'darwin'",
@@ -1542,10 +1542,10 @@ wheels = [
 [[package]]
 name = "mlx-lm"
 version = "0.31.3"
-source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#f5b9c9f42ef7577c73c7f6eeeb15b35a6682ff57" }
+source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fadd-tp-helpers#19ab177bbe26f5990efabbe2125cdb173494b0e5" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx", version = "0.31.2.dev20260420+553a7adb", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#553a7adbb20ed1b71fe643f4075982a639aaef1b" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx", version = "0.31.2.dev20260421+b27d5e35", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=splitk#b27d5e35da9d0fae4f70918830a7f52b2b81d131" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -1562,7 +1562,7 @@ dependencies = [
     { name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "miniaudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx", version = "0.31.2.dev20260420+553a7adb", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#553a7adbb20ed1b71fe643f4075982a639aaef1b" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx", version = "0.31.2.dev20260421+b27d5e35", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=splitk#b27d5e35da9d0fae4f70918830a7f52b2b81d131" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },