[EXPERIMENT][WIP][OpenVINO] Add NemotronH (Nemotron-3 Mamba2) support

docs/source/openvino/models.mdx

@@ -115,6 +115,7 @@ Here is the list of the supported architectures :
 - MobileNet v1
 - MobileNet v2
 - MobileViT
+- Nemotron-H (NVIDIA-Nemotron-3-Nano)
 - Nystromformer
 - OLMo
 - OLMo 2

optimum/exporters/openvino/_ov_ops.py

@@ -111,3 +111,107 @@ def convert_recurrent_attention_cell(context):
     final_output = ops.concat([core_attn_out_new, last_recurrent_state_new], 0)
 
     return [final_output.output(0)]
+
+
+
+
+# Conversion rule for the `Mamba2RecurrentCellOp` operation in a Torch graph.
+#
+# This generalizes the recurrent-cell-to-`ov::Loop` approach (originally introduced for the
+# GatedDeltaNet block, see `convert_recurrent_attention_cell` above) to the Mamba2 selective
+# state-space recurrence used by hybrid Mamba2 models such as NemotronH.
+#
+# The Mamba2 single-step recurrence over the SSM state is a linear recurrence:
+#       state_t = state_{t-1} * dA_t + dBx_t            # [B, H, P, N]
+#       y_t     = reduce_sum(state_t * C_t, axis=N)      # [B, H, P]
+# where `dA` (discretized A), `dBx` (discretized B * x) and `C` are precomputed and vectorized
+# over the sequence in the patched mixer forward. The skip connection `x_t * D` does not depend
+# on the recurrent state and is therefore added outside the loop.
+#
+# The `Mamba2RecurrentCellOp` appears in the Torch graph as a result of replacing the
+# `Mamba2RecurrentCell` `torch.nn.Module` via a registered `ModuleExtension` in the OpenVINO
+# PyTorch frontend; OpenVINO then applies this conversion rule to the resulting operation.
+def convert_recurrent_mamba2_cell(context):
+    # Inputs match the forward signature of `Mamba2RecurrentCell`.
+    # `dA` is broadcastable ([B, H, T, 1, 1]) while `dBx` carries the full state shape
+    # ([B, H, T, P, N]); shapes for the accumulator and the trip count are therefore
+    # derived from `dBx`.
+    dA = context.get_input(0)  # [B, H, T, 1, 1]
+    dBx = context.get_input(1)  # [B, H, T, P, N]
+    C = context.get_input(2)  # [B, H, T, N]
+    last_state_old = context.get_input(3)  # [B, H, P, N]
+
+    const_zero_axis = ops.constant(0, dtype=np.int32)
+    const_two = ops.constant(2, dtype=np.int32)
+    const_minus_one = ops.constant(-1, dtype=np.int32)
+
+    # Build the zero-initialized output accumulator with shape [B, H, T, P].
+    dBx_shape = ops.shape_of(dBx)
+    core_shape = ops.gather(dBx_shape, ops.constant([0, 1, 2, 3], dtype=np.int32), const_zero_axis)
+    const_zero_f32 = ops.constant(0, dtype=np.float32)
+    core_out = ops.broadcast(const_zero_f32, core_shape)
+
+    # Trip count for the loop equals the sequence length (dim 2).
+    seq_len = ops.gather(dBx_shape, const_two, const_zero_axis)
+    seq_len = ops.convert(seq_len, "i32")
+
+    # Body parameters (one timestep slice each).
+    timestep_param = ops.parameter([], np.int32, "timestep")
+    dA_t_param = ops.parameter([-1, -1, 1, -1, -1], np.float32, "dA_t")
+    dBx_t_param = ops.parameter([-1, -1, 1, -1, -1], np.float32, "dBx_t")
+    C_t_param = ops.parameter([-1, -1, 1, -1], np.float32, "C_t")
+    last_state_t = ops.parameter([-1, -1, -1, -1], np.float32, "last_state_t")
+    core_out_t = ops.parameter([-1, -1, -1, -1], np.float32, "core_out_t")
+
+    # Drop the singleton sequence dimension introduced by slicing.
+    dA_t = ops.squeeze(dA_t_param, const_two)  # [B, H, 1, 1] (broadcastable to [B, H, P, N])
+    dBx_t = ops.squeeze(dBx_t_param, const_two)  # [B, H, P, N]
+    C_t = ops.squeeze(C_t_param, const_two)  # [B, H, N]
+
+    # state_t = state_{t-1} * dA_t + dBx_t
+    last_state_new = ops.add(ops.multiply(last_state_t, dA_t), dBx_t)  # [B, H, P, N]
+
+    # y_t = reduce_sum(state_t * C_t, axis=N) -> [B, H, P]
+    const_minus_two = ops.constant(-2, dtype=np.int32)
+    y_t = ops.multiply(last_state_new, ops.unsqueeze(C_t, const_minus_two))  # [B, H, P, N]
+    y_t = ops.reduce_sum(y_t, const_minus_one, False)  # [B, H, P]
+    y_t = ops.unsqueeze(y_t, const_two)  # [B, H, 1, P]
+
+    timestep = ops.unsqueeze(timestep_param, const_zero_axis)
+    core_out_res = ops.scatter_update(core_out_t, timestep, y_t, const_two)
+    last_state_res = last_state_new
+
+    body_cond = ops.constant([True], dtype=bool)
+    body_model = ov.Model(
+        [body_cond, last_state_res, core_out_res],
+        [
+            timestep_param,
+            dA_t_param,
+            dBx_t_param,
+            C_t_param,
+            last_state_t,
+            core_out_t,
+        ],
+        "mamba2_body_model",
+    )
+
+    loop = ops.loop(seq_len, ops.constant(True, dtype="bool"))
+    loop.set_function(body_model)
+
+    loop.set_sliced_input(dA_t_param, dA, 0, 1, 1, -1, 2)
+    loop.set_sliced_input(dBx_t_param, dBx, 0, 1, 1, -1, 2)
+    loop.set_sliced_input(C_t_param, C, 0, 1, 1, -1, 2)
+    loop.set_merged_input(last_state_t, last_state_old, last_state_res.output(0))
+    loop.set_merged_input(core_out_t, core_out.output(0), core_out_res.output(0))
+    loop.set_special_body_ports([0, 0])
+
+    core_out_new = loop.get_iter_value(core_out_res.output(0), -1)
+    last_state_new = loop.get_iter_value(last_state_res.output(0), -1)
+
+    flatten_shape = ops.constant([-1], dtype=np.int32)
+    core_out_new = ops.reshape(core_out_new, flatten_shape, False)
+    last_state_new = ops.reshape(last_state_new, flatten_shape, False)
+
+    final_output = ops.concat([core_out_new, last_state_new], 0)
+
+    return [final_output.output(0)]

optimum/exporters/openvino/input_generators.py

@@ -1703,6 +1703,77 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int
         return cache_params
 
 
+
+class NemotronHDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator):
+    """
+    Generates dummy cache_params inputs for NemotronH (hybrid Mamba2 + Attention + MoE) architectures.
+
+    The cache is a flat list of tensors ordered as:
+      * for each mamba layer: conv_state, recurrent_state
+      * for each attention layer: key cache, value cache
+    """
+
+    SUPPORTED_INPUT_NAMES = ("cache_params",)
+
+    def __init__(
+        self,
+        task: str,
+        normalized_config,
+        batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"],
+        sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"],
+        **kwargs,
+    ):
+        super().__init__(
+            task=task,
+            normalized_config=normalized_config,
+            batch_size=batch_size,
+            sequence_length=sequence_length,
+            **kwargs,
+        )
+
+        config = normalized_config.config
+        self.layers_block_type = config.layers_block_type
+        self.num_mamba_layers = config.layers_block_type.count("mamba")
+        self.num_attn_layers = config.layers_block_type.count("attention")
+
+        self.conv_kernel_size = config.conv_kernel
+        self.ssm_state_size = config.ssm_state_size
+        self.n_groups = config.n_groups
+        self.mamba_num_heads = config.mamba_num_heads
+        self.mamba_head_dim = config.mamba_head_dim
+        self.intermediate_size = config.mamba_num_heads * config.mamba_head_dim
+        self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+
+        self.num_key_value_heads = config.num_key_value_heads
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+
+    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
+        cache_params = []
+
+        for _ in range(self.num_mamba_layers):
+            conv_state_shape = (self.batch_size, self.conv_dim, self.conv_kernel_size)
+            conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype)
+            cache_params.append(conv_state)
+            recurrent_state_shape = (
+                self.batch_size,
+                self.mamba_num_heads,
+                self.mamba_head_dim,
+                self.ssm_state_size,
+            )
+            recurrent_state = self.random_float_tensor(recurrent_state_shape, framework=framework, dtype=float_dtype)
+            cache_params.append(recurrent_state)
+
+        for _ in range(self.num_attn_layers):
+            kv_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.head_dim)
+            k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype)
+            v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype)
+            cache_params.append(k)
+            cache_params.append(v)
+
+        return cache_params
+
+
+
 class Qwen3_5DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator):
     """
     Generates dummy cache_params inputs for Qwen3.5 architectures.

optimum/exporters/openvino/model_configs.py

@@ -75,6 +75,7 @@
     LTXTransformerDummyInputGenerator,
     LTXVaeDummyInputGenerator,
     MambaCacheDummyInputGenerator,
+    NemotronHDummyPastKeyValuesGenerator,
     OVFalconDummyPastKeyValuesGenerator,
     OVMiniCPM3DummyPastKeyValuesGenerator,
     PooledProjectionsDummyInputGenerator,
@@ -141,6 +142,7 @@
     MixtralModelPatcher,
     ModelPatcher,
     MPTModelPatcher,
+    NemotronHModelPatcher,
     OVDecoderModelPatcher,
     OVSeq2SeqModelPatcher,
     PegasusModelPatcher,
@@ -6101,3 +6103,79 @@ class TrOCROpenVINOConfig(TextSeq2SeqOpenVINOConfig):
         decoder_num_attention_heads="decoder_attention_heads",
         hidden_size="hidden_size",
     )
+
+
+@register_in_tasks_manager(
+    "nemotron_h",
+    *["text-generation", "text-generation-with-past"],
+    library_name="transformers",
+)
+class NemotronHOpenVINOConfig(TextDecoderOpenVINOConfig):
+    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
+    MIN_TRANSFORMERS_VERSION = "5.0.0"
+
+    def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str):
+        if direction not in ["inputs", "outputs"]:
+            raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given')
+
+        if direction == "inputs":
+            decoder_sequence_name = "past_sequence_length"
+            cache_name_prefix = "cache_params.past"
+        else:
+            decoder_sequence_name = "past_sequence_length + sequence_length"
+            cache_name_prefix = "cache_params.present"
+
+        layers_block_type = self._normalized_config.layers_block_type
+        self.num_mamba_layers = layers_block_type.count("mamba")
+        self.num_attn_layers = layers_block_type.count("attention")
+
+        for i in range(self.num_mamba_layers):
+            # [batch_size, conv_dim, conv_kernel_size]
+            inputs_or_outputs[f"{cache_name_prefix}.conv.{i}"] = {0: "batch_size"}
+            # [batch_size, mamba_num_heads, mamba_head_dim, ssm_state_size]
+            inputs_or_outputs[f"{cache_name_prefix}.ssm.{i}"] = {0: "batch_size"}
+
+        for i in range(self.num_attn_layers):
+            inputs_or_outputs[f"{cache_name_prefix}.key.{i}"] = {0: "batch_size", 2: decoder_sequence_name}
+            inputs_or_outputs[f"{cache_name_prefix}.value.{i}"] = {0: "batch_size", 2: decoder_sequence_name}
+
+    @property
+    def inputs(self) -> Dict[str, Dict[int, str]]:
+        common_inputs = {
+            "input_ids": {0: "batch_size", 1: "sequence_length"},
+            "attention_mask": {0: "batch_size", 1: "sequence_length"},
+        }
+        if self.use_past_in_inputs:
+            self.add_past_key_values(common_inputs, direction="inputs")
+        return common_inputs
+
+    def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
+        # need to override `generate_dummy_inputs` since the hybrid Mamba2 model exposes
+        # conv/recurrent states (for mamba layers) and key/value caches (for attention layers)
+        # bundled together under a single flat `cache_params` input.
+        dummy_inputs_generators = self._create_dummy_input_generator_classes(**kwargs)
+
+        dummy_inputs = {}
+        input_names = [key for key in self.inputs.keys() if not key.startswith("cache_params")]
+        if self.use_past_in_inputs:
+            input_names.extend(["cache_params"])
+
+        for input_name in input_names:
+            input_was_inserted = False
+            for dummy_input_gen in dummy_inputs_generators:
+                if dummy_input_gen.supports_input(input_name):
+                    dummy_inputs[input_name] = self.overwrite_shape_and_generate_input(
+                        dummy_input_gen,
+                        input_name,
+                        framework,
+                        input_shapes=kwargs,
+                    )
+                    input_was_inserted = True
+                    break
+            if not input_was_inserted:
+                raise RuntimeError(
+                    f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.'
+                )
+
+        return dummy_inputs
+