support nemotron 3 (#678)

* support nemotron 3

* fix

* bump version

.github/workflows/pull_request.yml

@@ -38,4 +38,6 @@ jobs:
       - name: Run tests
         shell: bash -l {0}
         run: |
-          python -m xmlrunner discover -v tests -o test-results/
+          curl -o test_data.zip -L https://github.com/ml-explore/mlx-lm/releases/download/test_data/test_data.zip
+          unzip test_data.zip
+          HF_HOME="." python -m xmlrunner discover -v tests -o test-results/

mlx_lm/_version.py

@@ -1,3 +1,3 @@
 # Copyright © 2023-2025 Apple Inc.
 
-__version__ = "0.29.0"
+__version__ = "0.29.1"

mlx_lm/fuse.py

@@ -76,8 +76,9 @@ def main() -> None:
 
     if args.dequantize:
         print("Dequantizing model")
-        model = dequantize(model)
+        model = dequantize_model(model)
         config.pop("quantization", None)
+        config.pop("quantization_config", None)
 
     save_path = Path(args.save_path)
     save(

mlx_lm/models/deepseek_v32.py

@@ -222,6 +222,11 @@ class DeepseekV32Attention(nn.Module):
             if mask is not None:
                 sparse_mask = sparse_mask & mask
             mask = sparse_mask
+        # Ensure the indexer cache is evaluated even if the topk_indices are unused
+        # to keep the graph from getting too large
+        if cache is not None and cache[0] is not None:
+            cache[0].keys = mx.depends(cache[0].keys, (cache[1].keys, cache[1].values))
+
         output = scaled_dot_product_attention(
             queries, keys, values, cache=cache[0], scale=self.scale, mask=mask
         )

mlx_lm/models/gemma3_text.py

@@ -54,13 +54,20 @@ class Attention(nn.Module):
         self.k_norm = RMSNorm(dims=head_dim, eps=args.rms_norm_eps)
         self.is_sliding = (layer_idx + 1) % args.sliding_window_pattern != 0
 
-        self.rope = initialize_rope(
-            dims=head_dim,
-            base=(args.rope_local_base_freq if self.is_sliding else args.rope_theta),
-            traditional=False,
-            max_position_embeddings=args.max_position_embeddings,
-            scaling_config=args.rope_scaling,
-        )
+        if self.is_sliding:
+            self.rope = initialize_rope(
+                dims=head_dim,
+                base=args.rope_local_base_freq,
+                traditional=False,
+            )
+        else:
+            self.rope = initialize_rope(
+                dims=head_dim,
+                base=args.rope_theta,
+                traditional=False,
+                max_position_embeddings=args.max_position_embeddings,
+                scaling_config=args.rope_scaling,
+            )
 
     def __call__(
         self,

mlx_lm/models/ministral3.py

@@ -36,13 +36,17 @@ class ModelArgs(BaseModelArgs):
             self.layer_types = ["full_attention"] * self.num_hidden_layers
 
 
-def _get_llama_4_attn_scale(
-    start: int, stop: int, beta: float, max_position_embeddings: int
-):
+def _get_llama_4_attn_scale(size, offset, beta: float, max_position_embeddings: int):
+    if isinstance(offset, mx.array) and offset.ndim > 0:
+        offset = offset[:, None]
+
     scaling = 1 + beta * mx.log(
-        1 + mx.floor(mx.arange(start, stop) / max_position_embeddings)
+        1 + mx.floor((mx.arange(size) + offset) / max_position_embeddings)
     )
-    return scaling[:, None]
+    if scaling.ndim == 2:
+        return scaling[:, None, :, None]
+    else:
+        return scaling[:, None]
 
 
 class Attention(nn.Module):
@@ -191,8 +195,8 @@ class LanguageModel(nn.Module):
             )
 
         attn_scale = _get_llama_4_attn_scale(
+            inputs.shape[1],
             offset,
-            offset + inputs.shape[1],
             self.args.rope_parameters["llama_4_scaling_beta"],
             self.args.rope_parameters["original_max_position_embeddings"],
         ).astype(h.dtype)

mlx_lm/models/nemotron_h.py

@@ -15,6 +15,7 @@ from .base import (
 )
 from .cache import KVCache, MambaCache
 from .ssm import ssm_update
+from .switch_layers import SwitchMLP
 
 
 @dataclass()
@@ -37,24 +38,34 @@ class ModelArgs(BaseModelArgs):
     time_step_limit: Tuple[float, float]
     mlp_bias: bool
     layer_norm_epsilon: float
-    rms_norm_eps: float
     use_bias: bool
     use_conv_bias: bool
-    residual_in_fp32: bool
     hybrid_override_pattern: List[str]
     head_dim: Optional[int] = None
+    moe_intermediate_size: Optional[int] = None
+    moe_shared_expert_intermediate_size: Optional[int] = None
+    n_group: Optional[int] = None
+    n_routed_experts: Optional[int] = None
+    n_shared_experts: Optional[int] = None
+    topk_group: Optional[int] = None
+    num_experts_per_tok: Optional[int] = None
+    norm_topk_prob: Optional[bool] = None
+    routed_scaling_factor: Optional[float] = None
 
 
 class MambaRMSNormGated(nn.Module):
-    def __init__(self, hidden_size: int, eps: float = 1e-6):
+    def __init__(self, hidden_size: int, eps: float, group_size: int):
         super().__init__()
         self.eps = eps
         self.weight = mx.ones(hidden_size)
+        self.group_size = group_size
 
-    def __call__(self, hidden_states: mx.array, gate: mx.array = None) -> mx.array:
+    def __call__(self, x: mx.array, gate: mx.array = None) -> mx.array:
         if gate is not None:
-            hidden_states = hidden_states * nn.silu(gate)
-        return mx.fast.rms_norm(hidden_states, self.weight, self.eps)
+            x = x * nn.silu(gate)
+        x = mx.unflatten(x, axis=-1, shape=(-1, self.group_size))
+        x = mx.fast.rms_norm(x, weight=None, eps=self.eps)
+        return self.weight * x.flatten(-2)
 
 
 class NemotronHMamba2Mixer(nn.Module):
@@ -90,8 +101,11 @@ class NemotronHMamba2Mixer(nn.Module):
         self.A_log = mx.log(mx.arange(1, self.num_heads + 1, dtype=mx.float32))
         self.D = mx.ones(self.num_heads)
 
+        group_size = self.intermediate_size // self.n_groups
         self.norm = MambaRMSNormGated(
-            self.intermediate_size, eps=args.layer_norm_epsilon
+            self.intermediate_size,
+            eps=args.layer_norm_epsilon,
+            group_size=group_size,
         )
         self.out_proj = nn.Linear(
             self.intermediate_size, self.hidden_size, bias=args.mamba_proj_bias
@@ -139,7 +153,7 @@ class NemotronHMamba2Mixer(nn.Module):
             self.A_log,
             B,
             C,
-            self.D,
+            self.D.astype(hidden_states.dtype),
             dt,
             self.dt_bias,
             state,
@@ -245,24 +259,113 @@ class NemotronHAttention(nn.Module):
 
 
 class NemotronHMLP(nn.Module):
-    def __init__(self, args: ModelArgs):
+    def __init__(self, args: ModelArgs, intermediate_size=None):
         super().__init__()
+        intermediate_size = intermediate_size or args.intermediate_size
+
         self.up_proj = nn.Linear(
-            args.hidden_size, args.intermediate_size, bias=args.mlp_bias
+            args.hidden_size, intermediate_size, bias=args.mlp_bias
         )
         self.down_proj = nn.Linear(
-            args.intermediate_size, args.hidden_size, bias=args.mlp_bias
+            intermediate_size, args.hidden_size, bias=args.mlp_bias
         )
 
     def __call__(self, x):
         return self.down_proj(nn.relu2(self.up_proj(x)))
 
 
+@mx.compile
+def group_expert_select(
+    gates,
+    e_score_correction_bias,
+    top_k,
+    n_group,
+    topk_group,
+    routed_scaling_factor,
+    norm_topk_prob,
+):
+
+    orig_scores = scores = mx.sigmoid(gates.astype(mx.float32))
+    scores = scores + e_score_correction_bias
+    if n_group > 1:
+        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
+        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
+        k = n_group - topk_group
+        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
+        scores = mx.put_along_axis(
+            scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+        )
+        scores = mx.flatten(scores, -2, -1)
+
+    k = top_k
+    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
+    if top_k > 1 and norm_topk_prob:
+        denominator = scores.sum(axis=-1, keepdims=True)
+        scores = scores / (denominator + 1e-20)
+    scores = scores * routed_scaling_factor
+
+    return inds, scores
+
+
+class MoEGate(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        self.n_routed_experts = config.n_routed_experts
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
+        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
+
+    def __call__(self, x):
+        return group_expert_select(
+            x @ self.weight.T,
+            self.e_score_correction_bias,
+            self.top_k,
+            self.n_group,
+            self.topk_group,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+        )
+
+
+class NemotronHMoE(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.switch_mlp = SwitchMLP(
+            config.hidden_size,
+            config.moe_intermediate_size,
+            config.n_routed_experts,
+            activation=nn.ReLU2(),
+        )
+
+        self.gate = MoEGate(config)
+        if config.n_shared_experts is not None:
+            intermediate_size = config.moe_shared_expert_intermediate_size
+            self.shared_experts = NemotronHMLP(
+                config, intermediate_size=intermediate_size
+            )
+
+    def __call__(self, x):
+        inds, scores = self.gate(x)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+        if self.config.n_shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
 class NemotronHBlock(nn.Module):
     def __init__(self, args: ModelArgs, block_type: str):
         super().__init__()
-        self.residual_in_fp32 = args.residual_in_fp32
-        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
 
         self.block_type = block_type
 
@@ -272,6 +375,8 @@ class NemotronHBlock(nn.Module):
             self.mixer = NemotronHAttention(args)
         elif self.block_type == "-":
             self.mixer = NemotronHMLP(args)
+        elif self.block_type == "E":
+            self.mixer = NemotronHMoE(args)
 
     def __call__(
         self,
@@ -296,7 +401,7 @@ class NemotronHModel(nn.Module):
             NemotronHBlock(args, block_type)
             for block_type in args.hybrid_override_pattern
         ]
-        self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
         self.fa_idx = 0
         self.ssm_idx = 0
         for b in args.hybrid_override_pattern:
@@ -372,4 +477,23 @@ class Model(nn.Module):
         for k, v in weights.items():
             if "conv1d.weight" in k and v.shape[-1] != 1:
                 weights[k] = v.moveaxis(2, 1)
+
+        # Stack experts
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"backbone.layers.{l}.mixer"
+            for m, n in [("down_proj", "fc2"), ("up_proj", "fc1")]:
+                if f"{prefix}.experts.0.{m}.weight" in weights:
+                    to_join = [
+                        weights.pop(f"{prefix}.experts.{e}.{m}.weight")
+                        for e in range(self.args.n_routed_experts)
+                    ]
+                    weights[f"{prefix}.switch_mlp.{n}.weight"] = mx.stack(to_join)
+
         return weights
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k and "A_log" not in k
+
+        return predicate

mlx_lm/models/rope_utils.py

@@ -43,18 +43,36 @@ class SuScaledRoPE(nn.Module):
             long_mscale (float, optional): Scale the input prior to embedding.
         """
         super().__init__()
-        freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
-        self._freqs = mx.array(long_factor, dtype=mx.float32) * freqs
         self.original_max_position_embeddings = original_max_position_embeddings
-        self.scale = long_mscale or math.sqrt(
-            1
-            + math.log(max_position_embeddings / original_max_position_embeddings)
-            / math.log(original_max_position_embeddings)
-        )
         self.dim = dims
 
+        freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
+        self._short_freqs = mx.array(short_factor, dtype=mx.float32) * freqs
+        self._long_freqs = mx.array(long_factor, dtype=mx.float32) * freqs
+
+        def default_scale(factor):
+            return math.sqrt(
+                1 + math.log(factor) / math.log(original_max_position_embeddings)
+            )
+
+        factor = max_position_embeddings / original_max_position_embeddings
+        self._short_scale = short_mscale or (
+            1.0 if factor <= 1.0 else default_scale(factor)
+        )
+        self._long_scale = long_mscale or (
+            1.0 if factor <= 1.0 else default_scale(factor)
+        )
+
     def __call__(self, x, offset: int = 0):
-        x[..., : self.dim] = self.scale * x[..., : self.dim]
+        seq_len = offset + x.shape[-2]
+        if seq_len > self.original_max_position_embeddings:
+            freqs = self._long_freqs
+            scale = self._long_scale
+        else:
+            freqs = self._short_freqs
+            scale = self._short_scale
+
+        x[..., : self.dim] = scale * x[..., : self.dim]
         return mx.fast.rope(
             x,
             self.dim,
@@ -62,7 +80,7 @@ class SuScaledRoPE(nn.Module):
             base=None,
             scale=1.0,
             offset=offset,
-            freqs=self._freqs,
+            freqs=freqs,
         )
 
 

mlx_lm/models/ssm.py

@@ -139,7 +139,7 @@ def ssm_attn(
 
     dt = compute_dt(dt, dt_bias, time_step_limit)
     repeats = h // g
-    A = -mx.exp(A_log)
+    A = -mx.exp(A_log).astype(dt.dtype)
     dtA = dt * A.reshape(1, 1, -1)
     dtx = dt.reshape(b, l, h, 1) * x
 

mlx_lm/server.py

@@ -34,7 +34,13 @@ from huggingface_hub import scan_cache_dir
 
 from ._version import __version__
 from .generate import BatchGenerator, stream_generate
-from .models.cache import can_trim_prompt_cache, make_prompt_cache, trim_prompt_cache
+from .models.cache import (
+    KVCache,
+    RotatingKVCache,
+    can_trim_prompt_cache,
+    make_prompt_cache,
+    trim_prompt_cache,
+)
 from .sample_utils import make_logits_processors, make_sampler
 from .utils import load
 
@@ -378,6 +384,7 @@ class ModelProvider:
         self.model = None
         self.tokenizer = None
         self.draft_model = None
+        self.cache_types = set()
 
         # Preload the default model if it is provided
         self.default_model_map = {}
@@ -448,6 +455,15 @@ class ModelProvider:
         elif draft_model_path is not None and draft_model_path != "default_model":
             self.draft_model, draft_tokenizer = load(draft_model_path)
             validate_draft_tokenizer(draft_tokenizer)
+
+        # Figure out the cache types and save them in a set for anybody that
+        # wants to make a decision based on those.
+        for c in make_prompt_cache(self.model):
+            self.cache_types.add(type(c))
+        if self.draft_model is not None:
+            for c in make_prompt_cache(self.draft_model):
+                self.cache_types.add(type(c))
+
         return self.model, self.tokenizer
 
 
@@ -477,6 +493,7 @@ class ResponseGenerator:
                     messages,
                     tools,
                     add_generation_prompt=True,
+                    tokenize=True,
                     **self.model_provider.cli_args.chat_template_args,
                 )
             else:
@@ -490,6 +507,9 @@ class ResponseGenerator:
             or self.model_provider.cli_args.draft_model is not None
         ):
             return False
+        for c in self.model_provider.cache_types:
+            if c not in (KVCache, RotatingKVCache):
+                return False
         if args.logits.logit_bias is not None:
             return False
         if args.logits.repetition_penalty != 0:
@@ -921,7 +941,7 @@ class APIHandler(BaseHTTPRequestHandler):
         self.top_p = self.body.get("top_p", self.response_generator.cli_args.top_p)
         self.top_k = self.body.get("top_k", self.response_generator.cli_args.top_k)
         self.min_p = self.body.get("min_p", self.response_generator.cli_args.min_p)
-        self.repetition_penalty = self.body.get("repetition_penalty", 1.0)
+        self.repetition_penalty = self.body.get("repetition_penalty", 0.0)
         self.repetition_context_size = self.body.get("repetition_context_size", 20)
         self.xtc_probability = self.body.get("xtc_probability", 0.0)
         self.xtc_threshold = self.body.get("xtc_threshold", 0.0)

tests/test_datsets.py

@@ -21,6 +21,8 @@ class TestDatasets(unittest.TestCase):
         cls.test_dir = cls.test_dir_fid.name
         if not os.path.isdir(cls.test_dir):
             os.mkdir(cls.test_dir_fid.name)
+        # Only one HF request
+        AutoTokenizer.from_pretrained(HF_MODEL_PATH)
 
     @classmethod
     def tearDownClass(cls):
@@ -37,7 +39,7 @@ class TestDatasets(unittest.TestCase):
         data = {"text": "This is an example for the model."}
         self.save_data(4 * [data])
         args = types.SimpleNamespace(train=True, test=False, data=self.test_dir)
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH, local_files_only=True)
         train, valid, test = datasets.load_dataset(args, tokenizer)
         self.assertEqual(len(train), 4)
         self.assertEqual(len(valid), 4)
@@ -50,7 +52,7 @@ class TestDatasets(unittest.TestCase):
         data = {"prompt": "What is the capital of France?", "completion": "Paris."}
         self.save_data(4 * [data])
         args = types.SimpleNamespace(train=True, test=False, data=self.test_dir)
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH, local_files_only=True)
         train, valid, test = datasets.load_dataset(args, tokenizer)
         self.assertEqual(len(train), 4)
         self.assertEqual(len(valid), 4)
@@ -69,7 +71,7 @@ class TestDatasets(unittest.TestCase):
         }
         self.save_data(4 * [data])
         args = types.SimpleNamespace(train=True, test=False, data=self.test_dir)
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH, local_files_only=True)
         train, valid, test = datasets.load_dataset(args, tokenizer)
         self.assertEqual(len(train), 4)
         self.assertEqual(len(valid), 4)
@@ -91,7 +93,7 @@ class TestDatasets(unittest.TestCase):
             test=False,
             train=True,
         )
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH, local_files_only=True)
         train, valid, test = datasets.load_dataset(args, tokenizer)
         self.assertTrue(len(train) > 0)
         self.assertTrue(len(train[0]) > 0)

tests/test_generate.py

@@ -81,7 +81,7 @@ class TestGenerate(unittest.TestCase):
 
     def test_stream_generate_speculative(self):
         # Use same model as draft model, this is not a speed test
-        draft_model, _ = load(self.HF_MODEL_PATH)
+        draft_model = self.model
 
         results: List[GenerationResponse] = []
         drafted: List[bool] = []

tests/test_prompt_cache.py

@@ -34,6 +34,7 @@ class TestPromptCache(unittest.TestCase):
     def setUpClass(cls):
         cls.test_dir_fid = tempfile.TemporaryDirectory()
         cls.test_dir = cls.test_dir_fid.name
+        cls.model, cls.tokenizer = load(HF_MODEL_PATH)
 
     @classmethod
     def tearDownClass(cls):
@@ -132,7 +133,7 @@ class TestPromptCache(unittest.TestCase):
                 self.assertTrue(mx.array_equal(v, lv))
 
     def test_cache_with_generate(self):
-        model, tokenizer = load(HF_MODEL_PATH)
+        model, tokenizer = self.model, self.tokenizer
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]
         results = list(generate_step(prompt, model, max_tokens=4))
         toks, all_logits = zip(*results)
@@ -212,7 +213,7 @@ class TestPromptCache(unittest.TestCase):
         self.assertEqual(num_trimmed, 3)
 
     def test_trim_cache_with_generate(self):
-        model, tokenizer = load(HF_MODEL_PATH)
+        model, tokenizer = self.model, self.tokenizer
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]
 
         prompt_cache = make_prompt_cache(model)
@@ -289,7 +290,7 @@ class TestPromptCache(unittest.TestCase):
         self.assertEqual(metadata, loaded_metadata)
 
     def test_cache_to_quantized(self):
-        model, tokenizer = load(HF_MODEL_PATH)
+        model, tokenizer = self.model, self.tokenizer
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]
         results = zip(range(4), generate_step(prompt, model))
         toks, all_logits = zip(*(r[1] for r in results))

tests/test_server.py

@@ -19,6 +19,7 @@ class DummyModelProvider:
         HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
         self.model, self.tokenizer = load(HF_MODEL_PATH)
         self.model_key = (HF_MODEL_PATH, None)
+        self.cache_types = set([KVCache])
 
         # Add draft model support
         self.draft_model = None