Add support for Nemotron sharding (#1693)

### Automated Testing
tested logits match

.mlx_typings/mlx/nn/layers/convolution.pyi

@@ -32,6 +32,7 @@ class Conv1d(Module):
     """
 
     weight: mx.array
+    bias: mx.array | None
     groups: int
     def __init__(
         self,

.mlx_typings/mlx/nn/layers/linear.pyi

@@ -40,6 +40,10 @@ class Linear(Module):
         bias (bool, optional): If set to ``False`` then the layer will
           not use a bias. Default is ``True``.
     """
+
+    weight: mx.array
+    bias: mx.array | None
+
     def __init__(self, input_dims: int, output_dims: int, bias: bool = ...) -> None: ...
     def __call__(self, x: mx.array) -> mx.array: ...
     def to_quantized(

.mlx_typings/mlx/nn/layers/normalization.pyi

@@ -88,6 +88,9 @@ class RMSNorm(Module):
         dims (int): The feature dimension of the input to normalize over
         eps (float): A small additive constant for numerical stability
     """
+
+    weight: mx.array
+
     def __init__(self, dims: int, eps: float = ...) -> None: ...
     def __call__(self, x) -> mx.array: ...
 

.mlx_typings/mlx_lm/models/nemotron_h.pyi

@@ -0,0 +1,154 @@
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .cache import ArraysCache, KVCache
+from .switch_layers import SwitchMLP
+
+@dataclass
+class ModelArgs:
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    max_position_embeddings: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    attention_bias: bool
+    mamba_num_heads: int
+    mamba_head_dim: int
+    mamba_proj_bias: bool
+    ssm_state_size: int
+    conv_kernel: int
+    n_groups: int
+    mlp_bias: bool
+    layer_norm_epsilon: float
+    use_bias: bool
+    use_conv_bias: bool
+    hybrid_override_pattern: List[str]
+    head_dim: Optional[int]
+    moe_intermediate_size: Optional[int]
+    moe_shared_expert_intermediate_size: Optional[int]
+    n_group: Optional[int]
+    n_routed_experts: Optional[int]
+    n_shared_experts: Optional[int]
+    topk_group: Optional[int]
+    num_experts_per_tok: Optional[int]
+    norm_topk_prob: Optional[bool]
+    routed_scaling_factor: Optional[float]
+    time_step_limit: Optional[Tuple[float, float]]
+    time_step_min: Optional[float]
+    time_step_max: Optional[float]
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
+    def __post_init__(self) -> None: ...
+
+class NemotronHMamba2Mixer(nn.Module):
+    num_heads: int
+    hidden_size: int
+    ssm_state_size: int
+    conv_kernel_size: int
+    intermediate_size: int
+    n_groups: int
+    head_dim: int
+    conv_dim: int
+    conv1d: nn.Conv1d
+    in_proj: nn.Linear
+    dt_bias: mx.array
+    A_log: mx.array
+    D: mx.array
+    norm: nn.RMSNorm
+    heads_per_group: int
+    out_proj: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        mask: Optional[mx.array],
+        cache: Optional[ArraysCache] = None,
+    ) -> mx.array: ...
+
+class NemotronHAttention(nn.Module):
+    hidden_size: int
+    num_heads: int
+    head_dim: int
+    num_key_value_heads: int
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[KVCache] = None,
+    ) -> mx.array: ...
+
+class NemotronHMLP(nn.Module):
+    up_proj: nn.Linear
+    down_proj: nn.Linear
+
+    def __init__(
+        self, args: ModelArgs, intermediate_size: Optional[int] = None
+    ) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class NemotronHMoE(nn.Module):
+    num_experts_per_tok: int
+    switch_mlp: SwitchMLP
+    shared_experts: NemotronHMLP
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class NemotronHBlock(nn.Module):
+    block_type: str
+    norm: nn.RMSNorm
+    mixer: NemotronHMamba2Mixer | NemotronHAttention | NemotronHMLP | NemotronHMoE
+
+    def __init__(self, args: ModelArgs, block_type: str) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class NemotronHModel(nn.Module):
+    embeddings: nn.Embedding
+    layers: list[NemotronHBlock]
+    norm_f: nn.RMSNorm
+    fa_idx: int
+    ssm_idx: int
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    backbone: NemotronHModel
+    lm_head: nn.Linear
+    model_type: str
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[NemotronHBlock]: ...
+    def make_cache(self) -> list[ArraysCache | KVCache]: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...

.mlx_typings/mlx_lm/models/qwen3_next.pyi

@@ -5,6 +5,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .cache import ArraysCache, KVCache
 from .switch_layers import SwitchGLU
 
 class Qwen3NextRMSNormGated(nn.Module):
@@ -99,6 +100,8 @@ class Qwen3NextModel(nn.Module):
     embed_tokens: nn.Embedding
     layers: list[Qwen3NextDecoderLayer]
     norm: nn.RMSNorm
+    ssm_idx: int
+    fa_idx: int
 
     def __init__(self, args: Any) -> None: ...
     def __call__(
@@ -121,3 +124,4 @@ class Model(nn.Module):
     def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
     @property
     def layers(self) -> list[Qwen3NextDecoderLayer]: ...
+    def make_cache(self) -> list[ArraysCache | KVCache]: ...

.mlx_typings/mlx_lm/models/switch_layers.pyi

@@ -73,6 +73,9 @@ class SwitchGLU(nn.Module):
     def __call__(self, x, indices) -> mx.array: ...
 
 class SwitchMLP(nn.Module):
+    fc1: SwitchLinear
+    fc2: SwitchLinear
+
     def __init__(
         self,
         input_dims: int,

resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-70B-Instruct-HF-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Llama-3.1-Nemotron-70B-Instruct-HF-4bit"
+n_layers = 80
+hidden_size = 8192
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "llama"
+quantization = "4bit"
+base_model = "NVIDIA Llama-3.1-Nemotron-70B-Instruct"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 39688355840

resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-70B-Instruct-HF-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Llama-3.1-Nemotron-70B-Instruct-HF-8bit"
+n_layers = 80
+hidden_size = 8192
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "llama"
+quantization = "8bit"
+base_model = "NVIDIA Llama-3.1-Nemotron-70B-Instruct"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 74964549632

resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-70B-Instruct-HF-bf16.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Llama-3.1-Nemotron-70B-Instruct-HF-bf16"
+n_layers = 80
+hidden_size = 8192
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "llama"
+quantization = "bf16"
+base_model = "NVIDIA Llama-3.1-Nemotron-70B-Instruct"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 141107412992

resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-Nano-4B-v1.1-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Llama-3.1-Nemotron-Nano-4B-v1.1-4bit"
+n_layers = 32
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "llama"
+quantization = "4bit"
+base_model = "NVIDIA Llama-3.1-Nemotron-Nano-4B-v1.1"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 2538706944

resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-Nano-4B-v1.1-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Llama-3.1-Nemotron-Nano-4B-v1.1-8bit"
+n_layers = 32
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "llama"
+quantization = "8bit"
+base_model = "NVIDIA Llama-3.1-Nemotron-Nano-4B-v1.1"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 4794980352

resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-Nano-4B-v1.1-bf16.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Llama-3.1-Nemotron-Nano-4B-v1.1-bf16"
+n_layers = 32
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "llama"
+quantization = "bf16"
+base_model = "NVIDIA Llama-3.1-Nemotron-Nano-4B-v1.1"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 9025492992

resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-4Bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-4Bit"
+n_layers = 52
+hidden_size = 2688
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "nemotron"
+quantization = "4bit"
+base_model = "NVIDIA Nemotron-3-Nano-30B-A3B"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 17775342336

resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-5Bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-5Bit"
+n_layers = 52
+hidden_size = 2688
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "nemotron"
+quantization = "5bit"
+base_model = "NVIDIA Nemotron-3-Nano-30B-A3B"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 21721476864

resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-6Bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-6Bit"
+n_layers = 52
+hidden_size = 2688
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "nemotron"
+quantization = "6bit"
+base_model = "NVIDIA Nemotron-3-Nano-30B-A3B"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 25667611392

resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-8Bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-8Bit"
+n_layers = 52
+hidden_size = 2688
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "nemotron"
+quantization = "8bit"
+base_model = "NVIDIA Nemotron-3-Nano-30B-A3B"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 33559880448

resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16"
+n_layers = 52
+hidden_size = 2688
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "nemotron"
+quantization = "bf16"
+base_model = "NVIDIA Nemotron-3-Nano-30B-A3B"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 63155889408

resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-MXFP4.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-MXFP4"
+n_layers = 52
+hidden_size = 2688
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "nemotron"
+quantization = "4bit"
+base_model = "NVIDIA Nemotron-3-Nano-30B-A3B"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 16788808704

resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4"
+n_layers = 52
+hidden_size = 2688
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "nemotron"
+quantization = "4bit"
+base_model = "NVIDIA Nemotron-3-Nano-30B-A3B"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 19323906944

resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-Nano-9B-v2-4bits.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/NVIDIA-Nemotron-Nano-9B-v2-4bits"
+n_layers = 56
+hidden_size = 4480
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "nemotron"
+quantization = "4bit"
+base_model = "NVIDIA Nemotron-Nano-9B-v2"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 5002791936

resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-Nano-9B-v2-6bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/NVIDIA-Nemotron-Nano-9B-v2-6bit"
+n_layers = 56
+hidden_size = 4480
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "nemotron"
+quantization = "6bit"
+base_model = "NVIDIA Nemotron-Nano-9B-v2"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 7224298496

src/exo/shared/models/model_cards.py

@@ -196,6 +196,7 @@ class ConfigData(BaseModel):
             ["LlamaForCausalLM"],
             ["GptOssForCausalLM"],
             ["Step3p5ForCausalLM"],
+            ["NemotronHForCausalLM"],
         ]
 
     @model_validator(mode="before")

src/exo/worker/engines/image/models/flux/adapter.py

@@ -93,7 +93,7 @@ class FluxModelAdapter(ModelAdapter[Flux1, Transformer]):
 
     @property
     def hidden_dim(self) -> int:
-        return self._transformer.x_embedder.weight.shape[0]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        return self._transformer.x_embedder.weight.shape[0]
 
     @property
     def needs_cfg(self) -> bool:

src/exo/worker/engines/image/models/flux/kontext_adapter.py

@@ -133,7 +133,7 @@ class FluxKontextModelAdapter(ModelAdapter[Flux1Kontext, Transformer]):
 
     @property
     def hidden_dim(self) -> int:
-        return self._transformer.x_embedder.weight.shape[0]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        return self._transformer.x_embedder.weight.shape[0]
 
     @property
     def needs_cfg(self) -> bool:

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

@@ -4,7 +4,7 @@ from abc import ABC, abstractmethod
 from collections.abc import Callable
 from functools import partial
 from inspect import signature
-from typing import TYPE_CHECKING, Any, Protocol, cast
+from typing import TYPE_CHECKING, Any, Literal, Protocol, cast
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -32,6 +32,13 @@ from mlx_lm.models.llama import Model as LlamaModel
 from mlx_lm.models.minimax import MiniMaxAttention
 from mlx_lm.models.minimax import Model as MiniMaxModel
 from mlx_lm.models.ministral3 import Model as Ministral3Model
+from mlx_lm.models.nemotron_h import Model as NemotronHModel
+from mlx_lm.models.nemotron_h import (
+    NemotronHAttention,
+    NemotronHMamba2Mixer,
+    NemotronHMoE,
+)
+from mlx_lm.models.nemotron_h import NemotronHModel as NemotronHInnerModel
 from mlx_lm.models.qwen3_5 import DecoderLayer as Qwen3_5DecoderLayer
 from mlx_lm.models.qwen3_5 import Model as Qwen3_5TextModel
 from mlx_lm.models.qwen3_5 import Qwen3_5TextModel as Qwen3_5TextModelInner
@@ -45,6 +52,7 @@ from mlx_lm.models.qwen3_next import (
     Qwen3NextGatedDeltaNet,
     Qwen3NextSparseMoeBlock,
 )
+from mlx_lm.models.qwen3_next import Qwen3NextModel as Qwen3NextInnerModel
 from mlx_lm.models.step3p5 import Model as Step35Model
 from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
 from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel
@@ -243,7 +251,13 @@ def get_inner_model(model: nn.Module) -> nn.Module:
         if isinstance(inner_inner, nn.Module):
             return inner_inner
 
-    raise ValueError("Model must either have a 'model' or 'transformer' attribute")
+    inner = getattr(model, "backbone", None)
+    if isinstance(inner, nn.Module):
+        return inner
+
+    raise ValueError(
+        "Model must either have a 'model', 'transformer', or 'backbone' attribute"
+    )
 
 
 def get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]:
@@ -259,8 +273,8 @@ def get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]:
     return layers
 
 
-def _patch_qwen35_cache(
-    model: Qwen3_5TextModel,
+def _patch_hybrid_cache(
+    model: Qwen3_5TextModel | Qwen3NextModel | NemotronHModel,
     fa_idx: int,
     has_full_attn: bool,
     ssm_idx: int,
@@ -270,16 +284,20 @@ def _patch_qwen35_cache(
     original = model.make_cache
 
     def patched() -> list[ArraysCache | KVCache]:
-        cache: list[ArraysCache | KVCache] = original()
+        cache = original()
         if not has_full_attn:
             entry = cache[fa_idx]
             orig_make_mask = entry.make_mask
             entry.make_mask = lambda n, **_kw: orig_make_mask(n)  # type: ignore
         if not has_linear:
             orig_ssm_make_mask = cache[ssm_idx].make_mask
-            cache[ssm_idx].make_mask = (  # type: ignore
-                lambda n, **kw: orig_ssm_make_mask(n, **kw) if kw else None  # type: ignore
-            )
+
+            def _ssm_mask(
+                n: int, **kw: bool | int | None
+            ) -> mx.array | Literal["causal"] | None:
+                return orig_ssm_make_mask(n, **kw) if kw else None
+
+            cache[ssm_idx].make_mask = _ssm_mask  # type: ignore
         return cache
 
     model.make_cache = patched
@@ -355,7 +373,7 @@ def pipeline_auto_parallel(
         inner_model_instance._swa_idx = 0 if not sliding_layers else sliding_layers[0]
         inner_model_instance._full_idx = 0 if not full_layers else full_layers[0]
 
-    if isinstance(inner_model_instance, Qwen3_5TextModelInner):
+    if isinstance(inner_model_instance, (Qwen3_5TextModelInner, Qwen3NextInnerModel)):
         full_attn_layers = [
             i for i, layer in enumerate(layers) if not getattr(layer, "is_linear", True)
         ]
@@ -365,14 +383,44 @@ def pipeline_auto_parallel(
         inner_model_instance.fa_idx = full_attn_layers[0] if full_attn_layers else 0
         inner_model_instance.ssm_idx = linear_layers[0] if linear_layers else 0
         if not full_attn_layers or not linear_layers:
-            _patch_qwen35_cache(
-                cast(Qwen3_5TextModel, model),
+            _patch_hybrid_cache(
+                cast(Qwen3_5TextModel | Qwen3NextModel, model),
                 fa_idx=inner_model_instance.fa_idx,
                 has_full_attn=bool(full_attn_layers),
                 ssm_idx=inner_model_instance.ssm_idx,
                 has_linear=bool(linear_layers),
             )
 
+    if isinstance(inner_model_instance, NemotronHInnerModel):
+        # NemotronH uses block_type: "M" (Mamba/SSM), "*" (Attention), "E" (MoE), "-" (MLP)
+        # Only "M" and "*" blocks have cache entries.
+        # Recompute fa_idx and ssm_idx as cache-array indices for the shard's layers.
+        cache_idx = 0
+        fa_idx: int | None = None
+        ssm_idx: int | None = None
+        for layer in layers:
+            block_type = getattr(layer, "block_type", None)
+            if block_type == "*":
+                if fa_idx is None:
+                    fa_idx = cache_idx
+                cache_idx += 1
+            elif block_type == "M":
+                if ssm_idx is None:
+                    ssm_idx = cache_idx
+                cache_idx += 1
+        has_attn = fa_idx is not None
+        has_mamba = ssm_idx is not None
+        inner_model_instance.fa_idx = fa_idx if fa_idx is not None else 0
+        inner_model_instance.ssm_idx = ssm_idx if ssm_idx is not None else 0
+        if not has_attn or not has_mamba:
+            _patch_hybrid_cache(
+                cast(NemotronHModel, model),
+                fa_idx=inner_model_instance.fa_idx,
+                has_full_attn=has_attn,
+                ssm_idx=inner_model_instance.ssm_idx,
+                has_linear=has_mamba,
+            )
+
     _set_layers(model, layers)
 
     assert isinstance(layers, list), (
@@ -431,7 +479,8 @@ def patch_tensor_model[T](model: T) -> T:
         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]
-            dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
+            if hasattr(dep_cache, "keys"):  # type: ignore
+                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
 
         return logits
 
@@ -552,6 +601,14 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
+    elif isinstance(model, NemotronHModel):
+        tensor_parallel_sharding_strategy = NemotronHShardingStrategy(
+            group,
+            all_to_sharded_linear,
+            sharded_to_all_linear,
+            all_to_sharded_linear_in_place,
+            sharded_to_all_linear_in_place,
+        )
     else:
         raise ValueError(f"Unsupported model type: {type(model)}")
 
@@ -1210,3 +1267,131 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
             if on_layer_loaded is not None:
                 on_layer_loaded(i, total)
         return model
+
+
+class NemotronHShardingStrategy(TensorParallelShardingStrategy):
+    def shard_model(
+        self,
+        model: nn.Module,
+        timeout_seconds: float,
+        on_timeout: TimeoutCallback | None,
+        on_layer_loaded: LayerLoadedCallback | None,
+    ) -> nn.Module:
+        model = cast(NemotronHModel, model)
+        rank = self.group.rank()
+        total = len(model.layers)
+        for i, layer in enumerate(model.layers):
+            eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
+
+            mixer = layer.mixer
+
+            if isinstance(mixer, NemotronHAttention):
+                mixer.q_proj = self.all_to_sharded_linear(mixer.q_proj)
+                mixer.k_proj = self.all_to_sharded_linear(mixer.k_proj)
+                mixer.v_proj = self.all_to_sharded_linear(mixer.v_proj)
+                mixer.o_proj = self.sharded_to_all_linear(mixer.o_proj)
+                mixer.num_heads //= self.N
+                mixer.num_key_value_heads //= self.N
+
+            elif isinstance(mixer, NemotronHMamba2Mixer):
+                self._shard_mamba2_mixer(mixer, rank)
+
+            elif isinstance(mixer, NemotronHMoE):
+                # Shard routed experts (SwitchMLP uses fc1/fc2)
+                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)
+                if hasattr(mixer, "shared_experts"):
+                    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]
+                mixer.sharding_group = self.group
+                layer.mixer = mixer  # pyright: ignore[reportAttributeAccessIssue]
+
+            mx.eval(layer)
+            if on_layer_loaded is not None:
+                on_layer_loaded(i, total)
+        return model
+
+    def _shard_mamba2_mixer(self, mixer: NemotronHMamba2Mixer, rank: int) -> None:
+        """Shard the Mamba2 mixer along the head dimension."""
+        world_size = self.N
+        num_heads = mixer.num_heads
+        head_dim = mixer.head_dim
+        n_groups = mixer.n_groups
+        ssm_state_size = mixer.ssm_state_size
+        intermediate_size = mixer.intermediate_size  # = num_heads * head_dim
+
+        # Per-rank sizes
+        heads_per_rank = num_heads // world_size
+        groups_per_rank = n_groups // world_size
+        is_per_rank = heads_per_rank * head_dim
+        bc_per_rank = groups_per_rank * ssm_state_size
+
+        # === in_proj: output layout is [gate:IS | conv_ssm:IS | B:NG*SS | C:NG*SS | dt:NH] ===
+        gate_start = 0
+        conv_ssm_start = intermediate_size
+        b_start = 2 * intermediate_size
+        c_start = b_start + n_groups * ssm_state_size
+        dt_start = c_start + n_groups * ssm_state_size
+
+        # Build index tensor for this rank's slice of each section
+        gate_idx = mx.arange(
+            gate_start + rank * is_per_rank, gate_start + (rank + 1) * is_per_rank
+        )
+        conv_ssm_idx = mx.arange(
+            conv_ssm_start + rank * is_per_rank,
+            conv_ssm_start + (rank + 1) * is_per_rank,
+        )
+        b_idx = mx.arange(
+            b_start + rank * bc_per_rank, b_start + (rank + 1) * bc_per_rank
+        )
+        c_idx = mx.arange(
+            c_start + rank * bc_per_rank, c_start + (rank + 1) * bc_per_rank
+        )
+        dt_idx = mx.arange(
+            dt_start + rank * heads_per_rank, dt_start + (rank + 1) * heads_per_rank
+        )
+
+        indices = mx.concatenate([gate_idx, conv_ssm_idx, b_idx, c_idx, dt_idx])
+        mixer.in_proj.weight = mixer.in_proj.weight[indices]
+
+        # === out_proj: input is intermediate_size (sharded) → hidden_size (reduce) ===
+        mixer.out_proj = self.sharded_to_all_linear(mixer.out_proj)
+
+        # === conv1d: depthwise conv on conv_dim channels ===
+        # conv_dim layout: [ssm_hidden:IS | B:NG*SS | C:NG*SS]
+        conv_ssm_idx_local = mx.arange(rank * is_per_rank, (rank + 1) * is_per_rank)
+        conv_b_idx = mx.arange(
+            intermediate_size + rank * bc_per_rank,
+            intermediate_size + (rank + 1) * bc_per_rank,
+        )
+        conv_c_idx = mx.arange(
+            intermediate_size + n_groups * ssm_state_size + rank * bc_per_rank,
+            intermediate_size + n_groups * ssm_state_size + (rank + 1) * bc_per_rank,
+        )
+        conv_indices = mx.concatenate([conv_ssm_idx_local, conv_b_idx, conv_c_idx])
+        mixer.conv1d.weight = mixer.conv1d.weight[conv_indices]
+        new_conv_dim = is_per_rank + 2 * bc_per_rank
+        mixer.conv1d.groups = new_conv_dim
+        if mixer.conv1d.bias is not None:
+            mixer.conv1d.bias = mixer.conv1d.bias[conv_indices]
+
+        # === Per-head parameters ===
+        h_start = rank * heads_per_rank
+        h_end = h_start + heads_per_rank
+        mixer.dt_bias = mixer.dt_bias[h_start:h_end]
+        mixer.A_log = mixer.A_log[h_start:h_end]
+        mixer.D = mixer.D[h_start:h_end]
+
+        # === Norm: weight is intermediate_size ===
+        mixer.norm.weight = mixer.norm.weight[
+            rank * is_per_rank : (rank + 1) * is_per_rank
+        ]
+
+        # === Update dimensions ===
+        mixer.num_heads = heads_per_rank
+        mixer.n_groups = groups_per_rank
+        mixer.intermediate_size = is_per_rank
+        mixer.conv_dim = new_conv_dim
+        mixer.heads_per_group = heads_per_rank // groups_per_rank