feat(models): vendor Nemotron-H Puzzle checkpoint support

tests/test_nemotron_puzzle_vendored.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Focused coverage for the Nemotron Puzzle heterogeneous-MoE vendor."""
+
+import importlib
+import sys
+
+import mlx.core as mx
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def _restore_nemotron_modules():
+    """Registration replaces an old native module; restore test isolation."""
+    from vllm_mlx.utils.tokenizer import _VENDORED_MODEL_TYPES
+
+    original = sys.modules.get("mlx_lm.models.nemotron_h")
+    sys.modules.pop("mlx_lm.models.nemotron_h_puzzle", None)
+    _VENDORED_MODEL_TYPES.difference_update({"nemotron_h", "nemotron_h_puzzle"})
+    yield
+    sys.modules.pop("mlx_lm.models.nemotron_h_puzzle", None)
+    if original is not None:
+        sys.modules["mlx_lm.models.nemotron_h"] = original
+    else:
+        sys.modules.pop("mlx_lm.models.nemotron_h", None)
+    _VENDORED_MODEL_TYPES.difference_update({"nemotron_h", "nemotron_h_puzzle"})
+
+
+def _args(model_type="nemotron_h_puzzle"):
+    from vllm_mlx.models.nemotron_h import ModelArgs
+
+    return ModelArgs(
+        model_type=model_type,
+        vocab_size=128,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=4,
+        max_position_embeddings=64,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        attention_bias=False,
+        mamba_num_heads=4,
+        mamba_head_dim=8,
+        mamba_proj_bias=False,
+        ssm_state_size=8,
+        conv_kernel=4,
+        n_groups=1,
+        mlp_bias=False,
+        layer_norm_epsilon=1e-5,
+        use_bias=False,
+        use_conv_bias=False,
+        layers_block_type=["moe", "mamba", "attention", "moe"],
+        moe_intermediate_size=48,
+        moe_latent_size=16,
+        n_routed_experts=4,
+        num_experts_per_tok=2,
+        norm_topk_prob=True,
+        block_configs=[
+            {"moe_intermediate_size": 64, "num_experts_per_tok": 2},
+            {},
+            {},
+            {"moe_intermediate_size": 96, "num_experts_per_tok": 4},
+        ],
+        num_nextn_predict_layers=1,
+        mtp_layers_block_type=["attention"],
+        mtp_block_configs=[{"num_experts_per_tok": 4}],
+    )
+
+
+def test_heterogeneous_blocks_forward_and_cache():
+    from vllm_mlx.models.nemotron_h import Model
+
+    args = _args()
+    model = Model(args)
+    assert not hasattr(model, "mtp")
+    first, _, _, last = model.layers
+    assert first.mixer.num_experts_per_tok == 2
+    assert last.mixer.num_experts_per_tok == 4
+    assert first.mixer.switch_mlp.fc1.weight.shape[-2] == 64
+    assert last.mixer.switch_mlp.fc1.weight.shape[-2] == 96
+
+    cache = model.make_cache()
+    assert len(cache) == 2
+    logits = model(mx.array([[1, 2, 3]], dtype=mx.int32), cache=cache)
+    mx.eval(logits)
+    assert cache[1].offset == 3
+    next_logits = model(mx.array([[4]], dtype=mx.int32), cache=cache)
+    mx.eval(next_logits)
+    assert next_logits.shape == (1, 1, args.vocab_size)
+    assert cache[1].offset == 4
+
+
+def test_loader_registration_dispatches_both_puzzle_config_names():
+    from mlx_lm.utils import _get_classes
+
+    from vllm_mlx.utils import tokenizer
+
+    tokenizer._register_vendored_archs()
+    for model_type in ("nemotron_h", "nemotron_h_puzzle"):
+        module = importlib.import_module(f"mlx_lm.models.{model_type}")
+        assert module.__name__ == "vllm_mlx.models.nemotron_h"
+        model_class, args_class = _get_classes({"model_type": model_type})
+        assert model_class is module.Model
+        assert args_class is module.ModelArgs
+    assert {"nemotron_h", "nemotron_h_puzzle"} <= tokenizer._VENDORED_MODEL_TYPES
+
+
+def test_puzzle_quantization_keeps_lm_head_unquantized():
+    from vllm_mlx.models.nemotron_h import Model
+
+    puzzle = Model(_args("nemotron_h_puzzle"))
+    assert not puzzle.quant_predicate("lm_head", puzzle.lm_head)
+    assert puzzle.quant_predicate("backbone.layers.0.mixer", puzzle.layers[0].mixer)
+
+    uniform = Model(_args("nemotron_h"))
+    assert uniform.quant_predicate("lm_head", uniform.lm_head)
+
+
+def test_profile_uses_nemotron_xml_parser_and_disables_speculation():
+    from vllm_mlx.model_aliases import resolve_profile
+
+    profile = resolve_profile("nemotron-puzzle-75b-a9b-6bit")
+    assert profile.tool_call_parser == "nemotron"
+    assert profile.reasoning_parser == "qwen3"
+    assert profile.is_hybrid and profile.is_moe
+    assert not profile.supports_spec_decode

vllm_mlx/aliases.json

@@ -1302,6 +1302,15 @@
     "supports_spec_decode": false,
     "is_moe": true
   },
+  "nemotron-puzzle-75b-a9b-6bit": {
+    "hf_path": "georgeis55/Nemotron-Labs-3-Puzzle-75B-A9B-MLX-6bit",
+    "tool_call_parser": "nemotron",
+    "reasoning_parser": "qwen3",
+    "is_hybrid": true,
+    "is_moe": true,
+    "supports_spec_decode": false,
+    "pflash_tier": "unknown"
+  },
   "bonsai-1.7b-2bit": {
     "hf_path": "prism-ml/Ternary-Bonsai-1.7B-mlx-2bit",
     "tool_call_parser": "hermes",

vllm_mlx/models/nemotron_h.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright © 2026 Apple Inc.
+"""Compatibility vendor for heterogeneous Nemotron-H Puzzle checkpoints.
+
+The installed mlx-lm 0.31.x ``nemotron_h`` implementation assumes every MoE
+block has the same intermediate width and top-k.  NVIDIA Nemotron Puzzle
+instead supplies those values in per-layer ``block_configs``.  This module
+keeps the upstream Nemotron-H implementation intact for uniform checkpoints
+and adds only the Puzzle plumbing from ml-explore/mlx-lm#1536.
+
+It is registered only while the installed mlx-lm lacks ``block_configs``;
+newer native implementations take precedence.  The implementation deliberately
+does not include MTP: the commonly available 6-bit Puzzle artifact omits its
+``mtp.*`` tensors, and self-speculative rollback for the hybrid SSM backbone is
+a separate follow-up.
+
+Upstream provenance: ml-explore/mlx-lm#1536, commits
+aec9d8ecf90f3fc600be4b8f76b816a52c6a3944 and
+77c06c84592b60015ac5ba17dd636b8fd2887746.
+"""
+
+import copy
+from dataclasses import dataclass
+
+import mlx.nn as nn
+
+from .. import _mlx_compat as _mlx_compat
+
+_mlx_compat.install()
+
+from mlx_lm.models import nemotron_h as _native  # noqa: E402
+
+
+@dataclass
+class ModelArgs(_native.ModelArgs):
+    """Nemotron-H arguments plus Puzzle's per-layer MoE descriptions."""
+
+    block_configs: list[dict] | None = None
+    # Parse-only compatibility for Puzzle configs whose MTP tensors were
+    # stripped during quantization. This vendor deliberately does not build
+    # an MTP module; inherited sanitize() drops any stray ``mtp.*`` weights.
+    num_nextn_predict_layers: int = 0
+    mtp_layers_block_type: list[str] | None = None
+    mtp_hybrid_override_pattern: list[str] | None = None
+    mtp_block_configs: list[dict] | None = None
+
+
+def _moe_layer_args(args: ModelArgs, block_cfg: dict | None) -> ModelArgs:
+    """Return a shallow per-layer view with Puzzle's MoE dimensions."""
+    if not block_cfg:
+        return args
+    layer_args = copy.copy(args)
+    if block_cfg.get("moe_intermediate_size") is not None:
+        layer_args.moe_intermediate_size = block_cfg["moe_intermediate_size"]
+    if block_cfg.get("num_experts_per_tok") is not None:
+        layer_args.num_experts_per_tok = block_cfg["num_experts_per_tok"]
+    return layer_args
+
+
+class MoEGate(_native.MoEGate):
+    """Puzzle omits optional group-routing values; use identity defaults."""
+
+    def __init__(self, config: ModelArgs):
+        nn.Module.__init__(self)
+        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 or 1.0
+        self.n_group = config.n_group or 1
+        self.topk_group = config.topk_group or 1
+        self.weight = _native.mx.zeros((self.n_routed_experts, config.hidden_size))
+        self.e_score_correction_bias = _native.mx.zeros((self.n_routed_experts,))
+
+
+class NemotronHMoE(nn.Module):
+    """Native MoE with the Puzzle-safe gate; all other behavior is unchanged."""
+
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.moe_latent_size = config.moe_latent_size
+        expert_input_dim = config.moe_latent_size or config.hidden_size
+        self.switch_mlp = _native.SwitchMLP(
+            expert_input_dim,
+            config.moe_intermediate_size,
+            config.n_routed_experts,
+            activation=nn.ReLU2(),
+        )
+        self.gate = MoEGate(config)
+        if config.n_shared_experts is not None:
+            self.shared_experts = _native.NemotronHMLP(
+                config, intermediate_size=config.moe_shared_expert_intermediate_size
+            )
+        if config.moe_latent_size is not None:
+            self.fc1_latent_proj = nn.Linear(
+                config.hidden_size, config.moe_latent_size, bias=config.mlp_bias
+            )
+            self.fc2_latent_proj = nn.Linear(
+                config.moe_latent_size, config.hidden_size, bias=config.mlp_bias
+            )
+
+    def __call__(self, x):
+        residuals = x
+        inds, scores = self.gate(x)
+        if self.moe_latent_size is not None:
+            x = self.fc1_latent_proj(x)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+        if self.moe_latent_size is not None:
+            y = self.fc2_latent_proj(y)
+        if self.config.n_shared_experts is not None:
+            y = y + self.shared_experts(residuals)
+        return y
+
+
+class NemotronHBlock(_native.NemotronHBlock):
+    """Native block constructor with Puzzle's MoE implementation for E blocks."""
+
+    def __init__(self, args: ModelArgs, block_type: str):
+        nn.Module.__init__(self)
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
+        self.block_type = block_type
+        if block_type == "M":
+            self.mixer = _native.NemotronHMamba2Mixer(args)
+        elif block_type == "*":
+            self.mixer = _native.NemotronHAttention(args)
+        elif block_type == "-":
+            self.mixer = _native.NemotronHMLP(args)
+        elif block_type == "E":
+            self.mixer = NemotronHMoE(args)
+
+
+class NemotronHModel(_native.NemotronHModel):
+    """Native backbone with heterogeneous construction for MoE layers."""
+
+    def __init__(self, args: ModelArgs):
+        nn.Module.__init__(self)
+        self.embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+        pattern = args.hybrid_override_pattern
+        block_configs = args.block_configs or [None] * len(pattern)
+        if len(block_configs) != len(pattern):
+            raise ValueError("block_configs must align with hybrid_override_pattern")
+        self.layers = [
+            NemotronHBlock(_moe_layer_args(args, cfg) if kind == "E" else args, kind)
+            for kind, cfg in zip(pattern, block_configs)
+        ]
+        self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
+        self.fa_idx = 0
+        self.ssm_idx = 0
+        for kind in pattern:
+            if kind == "*":
+                break
+            if kind == "M":
+                self.fa_idx += 1
+        for kind in pattern:
+            if kind == "*":
+                self.ssm_idx += 1
+            elif kind == "M":
+                break
+
+
+class Model(_native.Model):
+    """Puzzle-aware model wrapper retaining native cache and sanitize paths."""
+
+    def __init__(self, args: ModelArgs):
+        nn.Module.__init__(self)
+        self.args = args
+        self.backbone = NemotronHModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.model_type = args.model_type
+
+    @property
+    def quant_predicate(self):
+        if self.model_type != "nemotron_h_puzzle":
+            return lambda _path, _module: True
+        # Upstream #1536: Puzzle's output projection is too sensitive to
+        # affine low-bit quantization; retain its checkpoint precision.
+        return lambda path, _module: path != "lm_head"

vllm_mlx/utils/tokenizer.py

@@ -662,6 +662,36 @@ def _register_vendored_archs() -> None:
             # still doesn't know the arch, native mlx-lm module or not).
             _VENDORED_MODEL_TYPES.add("hy_v3")
 
+    # mlx-lm <=0.31.x has ``nemotron_h`` but assumes uniform MoE widths.
+    # Puzzle's ``block_configs`` needs the small compatibility vendor above.
+    # Once #1536 (or equivalent native support) is installed, leave native
+    # Nemotron-H untouched and let its own remapping handle Puzzle configs.
+    _native_nemotron_h = sys.modules.get("mlx_lm.models.nemotron_h")
+    if _native_nemotron_h is None:
+        try:
+            import importlib
+
+            _native_nemotron_h = importlib.import_module("mlx_lm.models.nemotron_h")
+        except ImportError:
+            _native_nemotron_h = None
+    _native_args = getattr(_native_nemotron_h, "ModelArgs", None)
+    _native_fields = getattr(_native_args, "__dataclass_fields__", {})
+    if "block_configs" not in _native_fields:
+        try:
+            from ..models import nemotron_h as _puzzle_nemotron_h
+
+            # The community 6-bit conversion identifies as ``nemotron_h``;
+            # NVIDIA's original config identifies as ``nemotron_h_puzzle``.
+            # Register both names without changing mlx-lm's general loader.
+            sys.modules["mlx_lm.models.nemotron_h"] = _puzzle_nemotron_h
+            sys.modules.setdefault(
+                "mlx_lm.models.nemotron_h_puzzle", _puzzle_nemotron_h
+            )
+        except Exception as e:
+            logger.warning("Nemotron Puzzle vendor failed to register: %s", e)
+        else:
+            _VENDORED_MODEL_TYPES.update({"nemotron_h", "nemotron_h_puzzle"})
+
 
 def _is_vendored_arch_model(model_name: str) -> bool:
     """Return True if model's config.json declares a model_type we vendor."""