huggingface
diff --git a/‎docs/source/openvino/models.mdx‎
Lines changed: 2 additions & 0 deletions b/‎docs/source/openvino/models.mdx‎
Lines changed: 2 additions & 0 deletions
diff --git a/‎optimum/exporters/openvino/model_configs.py‎
Lines changed: 126 additions & 0 deletions b/‎optimum/exporters/openvino/model_configs.py‎
Lines changed: 126 additions & 0 deletions
@@ -52,6 +52,7 @@ Here is the list of the supported architectures :
 - ESM
 - Exaone
 - Falcon
+- Falcon-Mamba
 - Flaubert
 - GLM-4
 - GLM-Edge
@@ -82,6 +83,7 @@ Here is the list of the supported architectures :
 - Llava-Next-Video
 - M2-M100
 - MAIRA-2
+- Mamba
 - MBart
 - MPNet
 - MPT
 
@@ -117,6 +117,7 @@
     LlavaNextVideoImageEmbeddingModelPatcher,
     LlavaQwen2ImageEmbeddingsModelPatcher,
     MairaImageEmbeddingModelPatcher,
+    MambaPatcher,
     MarianModelPatcher,
     MarianStatefulSeq2SeqDecoderPatcher,
     MiniCPM3Patcher,
@@ -4382,3 +4383,128 @@ def patch_model_for_export(
         if self._behavior != VLMConfigBehavior.VISION_EMBEDDINGS:
             return super().patch_model_for_export(model, model_kwargs)
         return Llama4ImageEmbeddingsModelPatcher(self, model, model_kwargs)
+
+
+class MambaCacheDummyInputGenerator(DummyInputGenerator):
+    """
+    Generates dummy past_ssm_states, past_conv_states and cache_position inputs for Mamba architectures.
+    """
+
+    SUPPORTED_INPUT_NAMES = ("cache_params", "cache_position")
+
+    def __init__(
+        self,
+        task: str,
+        normalized_config,
+        batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"],
+        sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"],
+        **kwargs,
+    ):
+        self.normalized_config = normalized_config
+        self.batch_size = batch_size
+        self.sequence_length = sequence_length
+        self.intermediate_size = self.normalized_config.config.intermediate_size
+        self.ssm_state_size = self.normalized_config.config.state_size
+        self.conv_kernel_size = self.normalized_config.config.conv_kernel
+
+    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
+        if input_name == "cache_params":
+            ssm_shape = [self.batch_size, self.intermediate_size, self.ssm_state_size]
+            conv_shape = [self.batch_size, self.intermediate_size, self.conv_kernel_size]
+            return [
+                (
+                    self.random_float_tensor(ssm_shape, framework=framework, dtype=float_dtype),
+                    self.random_float_tensor(conv_shape, framework=framework, dtype=float_dtype),
+                )
+                for _ in range(self.normalized_config.num_layers)
+            ]
+        elif input_name == "cache_position":
+            return self.random_int_tensor(
+                shape=[self.conv_kernel_size],
+                max_value=self.sequence_length,
+                framework=framework,
+                dtype=int_dtype,
+            )
+
+        raise ValueError(f"Unsupported input name {input_name}")
+
+
+@register_in_tasks_manager(
+    "falcon_mamba", *["text-generation", "text-generation-with-past"], library_name="transformers"
+)
+@register_in_tasks_manager("mamba", *["text-generation", "text-generation-with-past"], library_name="transformers")
+class MambaOpenVINOConfig(TextDecoderOnnxConfig):
+    DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MambaCacheDummyInputGenerator)
+    DUMMY_PKV_GENERATOR_CLASS = MambaCacheDummyInputGenerator
+    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
+    MIN_TRANSFORMERS_VERSION = version.parse("4.43.0")
+
+    @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"},
+            "cache_position": {0: "cache_sequence_length"},
+        }
+        if self.use_past_in_inputs:
+            self.add_past_key_values(common_inputs, direction="inputs")
+        return common_inputs
+
+    def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str):
+        """
+        Fills `input_or_outputs` mapping with past_key_values dynamic axes considering the direction.
+
+        Args:
+            inputs_or_outputs (`Dict[str, Dict[int, str]]`):
+                The mapping to fill.
+            direction (`str`):
+                either "inputs" or "outputs", it specifies whether `input_or_outputs` is the input mapping or the
+                output mapping, this is important for axes naming.
+        """
+        if direction not in ["inputs", "outputs"]:
+            raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given')
+
+        if direction == "inputs":
+            ssm_conv_states_name = "cache_params.past"
+        else:
+            ssm_conv_states_name = "cache_params.present"
+
+        for i in range(self._normalized_config.num_layers):
+            # [batch_size, d_state, d_model]
+            inputs_or_outputs[f"{ssm_conv_states_name}.ssm.{i}"] = {0: "batch_size"}
+            # [batch_size, conv_kernel_size - 1, d_model]
+            inputs_or_outputs[f"{ssm_conv_states_name}.conv.{i}"] = {0: "batch_size"}
+
+    def patch_model_for_export(
+        self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
+    ):
+        return MambaPatcher(self, model, model_kwargs)
+
+    def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
+        # need to override `generate_dummy_inputs` since mamba model has other states: ssm_states and conv_states
+        # which we separate and call them as past_ssm_states and past_conv_states
+        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 ONNX config.'
+                )
+
+        return dummy_inputs