optimum-rbln 0.9.4a2__py3-none-any.whl → 0.10.0.post1__py3-none-any.whl

optimum/rbln/transformers/models/gpt_oss/gpt_oss_architecture.py

@@ -0,0 +1,122 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ..decoderonly.configuration_decoderonly import RBLNLoRAConfig
+from ..decoderonly.decoderonly_architecture import (
+    DecoderOnlyAttention,
+    DecoderOnlyLayer,
+    DecoderOnlyWrapper,
+)
+
+
+class RBLNGptOssWrapper(DecoderOnlyWrapper):
+    def get_rbln_layer_class(self):
+        return RBLNGptOssLayer
+
+
+class RBLNGptOssLayer(DecoderOnlyLayer):
+    def __init__(self, layer, self_attn: DecoderOnlyAttention, lora_config: Optional[RBLNLoRAConfig] = None):
+        super().__init__(layer, self_attn, lora_config)
+        self.mlp = RBLNGptOssMLP(layer.mlp)
+
+    def get_mlp(self) -> nn.Module:
+        return self.mlp
+
+
+class RBLNGptOssTopKRouter(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.weight = model.weight
+        self.bias = model.bias
+
+    def forward(self, hidden_states):
+        return F.linear(hidden_states, self.weight, self.bias)  # (seq_len, num_experts)
+
+
+class RBLNGptOssExperts(nn.Module):
+    def __init__(self, model, top_k: Optional[int] = None):
+        super().__init__()
+        self.intermediate_size = model.intermediate_size
+        self.num_experts = model.num_experts
+        self.hidden_size = model.hidden_size
+
+        self.register_buffer(
+            "gate_proj_blocks",
+            model.gate_up_proj_blocks.data[:, ::2, :, :].reshape(self.num_experts, self.intermediate_size, -1),
+        )
+        self.register_buffer("gate_proj_scales", model.gate_up_proj_scales.data[:, ::2, :])
+        self.register_buffer(
+            "gate_proj_bias",
+            model.gate_up_proj_bias.data[:, ::2].reshape(self.num_experts, self.intermediate_size),
+        )
+
+        self.register_buffer(
+            "up_proj_blocks",
+            model.gate_up_proj_blocks.data[:, 1::2, :, :].reshape(self.num_experts, self.intermediate_size, -1),
+        )
+        self.register_buffer("up_proj_scales", model.gate_up_proj_scales.data[:, 1::2, :])
+        self.register_buffer(
+            "up_proj_bias", model.gate_up_proj_bias.data[:, 1::2].reshape(self.num_experts, self.intermediate_size)
+        )
+
+        self.register_buffer(
+            "down_proj_blocks", model.down_proj_blocks.data.reshape(self.num_experts, self.hidden_size, -1)
+        )
+        self.register_buffer("down_proj_scales", model.down_proj_scales.data)
+        self.register_buffer("down_proj_bias", model.down_proj_bias.data)
+
+        self.alpha = model.alpha  # 1.702
+        self.limit = model.limit  # 7.0
+        self.top_k = top_k
+
+    def forward(self, hidden_states: torch.Tensor, router_logits: torch.Tensor) -> torch.Tensor:
+        return torch.ops.rbln_custom_ops.custom_moe_glu_mxfp4(
+            hidden_states,
+            self.gate_proj_blocks,
+            self.gate_proj_scales,
+            self.gate_proj_bias,
+            self.up_proj_blocks,
+            self.up_proj_scales,
+            self.up_proj_bias,
+            self.down_proj_blocks,
+            self.down_proj_scales,
+            self.down_proj_bias,
+            router_logits,
+            torch.tensor(self.alpha, dtype=hidden_states.dtype),
+            torch.tensor(self.limit, dtype=hidden_states.dtype),
+            k=self.top_k,
+            post_norm=True,
+        )
+
+
+class RBLNGptOssMLP(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.router = RBLNGptOssTopKRouter(model.router)
+        self.experts = RBLNGptOssExperts(model.experts, top_k=model.router.top_k)
+
+    def forward(self, hidden_states):
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        router_logits = self.router(hidden_states)
+        routed_out = self.experts(hidden_states, router_logits=router_logits)
+        routed_out = routed_out.reshape(batch_size, sequence_length, hidden_dim)
+        return routed_out

optimum/rbln/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -0,0 +1,168 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+from safetensors.torch import load_file
+from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig
+from transformers.integrations.mxfp4 import Mxfp4GptOssExperts
+from transformers.modeling_utils import PreTrainedModel, no_init_weights
+
+from ....utils.logging import get_logger
+from ...models.decoderonly import (
+    RBLNDecoderOnlyModelConfig,
+    RBLNDecoderOnlyModelForCausalLM,
+    RBLNDecoderOnlyModelForCausalLMConfig,
+)
+from ...utils.rbln_quantization import load_weight_files
+from .gpt_oss_architecture import RBLNGptOssWrapper
+
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+logger = get_logger(__name__)
+
+
+class RBLNGptOssForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The GPT-OSS Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based GptOssForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers GptOssForCausalLM model into a RBLN transformer model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    **Configuration:**
+    This model uses [`RBLNGptOssForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNGptOssForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNGptOssForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNGptOssForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNGptOssForCausalLM.from_pretrained(
+            "openai/gpt-oss-20b",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=8,
+            rbln_kvcache_partition_len=8192,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "tensor_parallel_size": 8,
+            "kvcache_partition_len": 8192,
+        }
+        model = RBLNGptOssForCausalLM.from_pretrained(
+            "openai/gpt-oss-20b",
+            export=True,
+            rbln_config=rbln_config,
+        )
+
+
+        # Using a RBLNGptOssForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNGptOssForCausalLMConfig
+
+        config = RBLNGptOssForCausalLMConfig(
+            batch_size=1,
+            tensor_parallel_size=8,
+            kvcache_partition_len=8192,
+        )
+        model = RBLNGptOssForCausalLM.from_pretrained(
+            "openai/gpt-oss-20b",
+            export=True,
+            rbln_config=config,
+        )
+        ```
+    """
+
+    _decoder_wrapper_cls = RBLNGptOssWrapper
+
+    @staticmethod
+    def _get_dtype(dtype: Union[str, torch.dtype] = None, torch_dtype: Union[str, torch.dtype] = None):
+        # For BC on torch_dtype argument
+        if torch_dtype is not None:
+            logger.warning_once("`torch_dtype` is deprecated! Use `dtype` instead!")
+            # If both kwargs are provided, use `dtype`
+            dtype = dtype if dtype is not None else torch_dtype
+
+        # As mxfp4_quantizer's default dtype
+        if dtype is None or dtype == "auto":
+            dtype = torch.bfloat16
+
+        return dtype
+
+    @classmethod
+    def get_pytorch_model(
+        cls,
+        model_id: str,
+        *args,
+        rbln_config: Optional[RBLNDecoderOnlyModelConfig] = None,
+        dtype: Union[str, torch.dtype] = None,
+        torch_dtype: Union[str, torch.dtype] = None,
+        config: Optional[PretrainedConfig] = None,
+        **kwargs,
+    ) -> PreTrainedModel:
+        safetensor_files = load_weight_files(model_id, exception_keywords=["original"])
+        state_dict = {k: v for f in safetensor_files for k, v in load_file(f).items()}
+
+        if config is None:
+            config, kwargs = AutoConfig.from_pretrained(model_id, return_unused_kwargs=True)
+
+        dtype = cls._get_dtype(dtype, torch_dtype)
+
+        with no_init_weights():
+            model = AutoModelForCausalLM.from_config(config, dtype=dtype, **kwargs)
+
+        _replace_with_mxfp4_linear(model, config)
+        model.load_state_dict(state_dict, strict=False)
+
+        return model
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
+        model: Optional["PreTrainedModel"] = None,
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_config: Optional[RBLNDecoderOnlyModelForCausalLMConfig] = None,
+    ) -> RBLNDecoderOnlyModelForCausalLMConfig:
+        rbln_config = super()._update_rbln_config(preprocessors, model, model_config, rbln_config)
+
+        if rbln_config.use_attention_mask:
+            raise ValueError(
+                "use_attention_mask is not supported for GPT-OSS because custom attention does not support attention sink for masked attention"
+            )
+
+        return rbln_config
+
+
+def _replace_with_mxfp4_linear(
+    model,
+    config,
+):
+    for name, module in model.named_children():
+        if module.__class__.__name__ == "GptOssExperts":
+            model._modules[name] = Mxfp4GptOssExperts(config)
+        if len(list(module.children())) > 0:
+            _replace_with_mxfp4_linear(module, config)

optimum/rbln/transformers/models/grounding_dino/configuration_grounding_dino.py

@@ -50,11 +50,14 @@ class RBLNGroundingDinoForObjectDetectionConfig(RBLNImageModelConfig):
         Raises:
             ValueError: If batch_size is not a positive integer.
         """
-        super().__init__(**kwargs)
-        self.encoder = encoder
-        self.decoder = decoder
-        self.text_backbone = text_backbone
-        self.backbone = backbone
+
+        super().__init__(batch_size=batch_size, **kwargs)
+        self.encoder = self.initialize_submodule_config(submodule_config=encoder, batch_size=self.batch_size)
+        self.decoder = self.initialize_submodule_config(submodule_config=decoder, batch_size=self.batch_size)
+        self.text_backbone = self.initialize_submodule_config(
+            submodule_config=text_backbone, batch_size=self.batch_size
+        )
+        self.backbone = self.initialize_submodule_config(submodule_config=backbone, batch_size=self.batch_size)
         self.output_attentions = output_attentions if output_attentions is not None else False
         self.output_hidden_states = output_hidden_states if output_hidden_states is not None else False
 

optimum/rbln/transformers/models/grounding_dino/grounding_dino_architecture.py

@@ -509,10 +509,12 @@ class _GroundingDinoBiMultiHeadAttention(torch.nn.Module):
 
         # mask vision for language
         if vision_attention_mask is not None:
-            # RBLN FIX: bool tensor to float tensor
-            mask = vision_attention_mask * torch.finfo(torch.float16).min
-            text_attn_weights = text_attn_weights.transpose(1, 2) + mask
-            text_attn_weights = text_attn_weights.transpose(1, 2)
+            # RBLN FIX: bool tensor to float tensor, broadcast across heads and src_len
+            mask = vision_attention_mask
+            if mask.dim() == 3:
+                mask = mask[..., 0]
+            mask = mask[:, None, None, :].repeat(1, self.num_heads, 1, 1).flatten(0, 1)
+            text_attn_weights = text_attn_weights + mask * torch.finfo(text_attn_weights.dtype).min
 
         text_attn_weights = text_attn_weights.softmax(dim=-1)
 

optimum/rbln/transformers/models/llava/modeling_llava.py

@@ -116,7 +116,6 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
     RBLNLlavaForConditionalGeneration is a multi-modal model that combines vision and language processing capabilities,
     optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
     This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
-
     Important Note:
         This model includes a Large Language Model (LLM) as a submodule. For optimal performance, it is highly recommended to use
         tensor parallelism for the language model. This can be achieved by using the `rbln_config` parameter in the

optimum/rbln/transformers/models/midm/midm_architecture.py

@@ -71,6 +71,12 @@ class MidmLMHeadModelWrapper(DecoderOnlyWrapper):
 
 
 class MidmModel(DecoderOnlyModel):
+    def __init__(self, model, layers, rbln_config, use_learned_pos_emb=None, use_rotary_emb=True):
+        super().__init__(
+            model, layers, rbln_config, use_learned_pos_emb=use_learned_pos_emb, use_rotary_emb=use_rotary_emb
+        )
+        self.use_layernorm1p = getattr(model, "use_layernorm1p", False)
+
     def get_layernorm1p(self, module: nn.LayerNorm):
         def layernorm1p(input: torch.Tensor):
             """Applies Layer Normalization with a slight modification on the weights."""
@@ -81,19 +87,22 @@ class MidmModel(DecoderOnlyModel):
         return layernorm1p
 
     def get_last_layernorm(self) -> nn.LayerNorm:
-        if self._original_mod.use_layernorm1p:
-            return self.get_layernorm1p(self._original_mod.ln_f)
-        else:
-            return self._original_mod.ln_f
+        if self.use_layernorm1p:
+            return self.get_layernorm1p(self.norm)
+        return self.norm
 
     def get_embedding(self) -> nn.Embedding:
-        return self._original_mod.wte
+        return self.embed_tokens
 
     def get_pos_embedding(self) -> nn.Embedding:
-        return self._original_mod.wpe
+        return self.embed_positions
 
 
 class MidmLayer(DecoderOnlyLayer):
+    def __init__(self, layer, self_attn: DecoderOnlyAttention, lora_config=None):
+        super().__init__(layer, self_attn, lora_config)
+        self.use_layernorm1p = getattr(layer, "use_layernorm1p", False)
+
     def get_layernorm1p(self, module: nn.LayerNorm):
         def layernorm1p(input: torch.Tensor):
             """Applies Layer Normalization with a slight modification on the weights."""
@@ -104,24 +113,22 @@ class MidmLayer(DecoderOnlyLayer):
         return layernorm1p
 
     def get_pre_attention_layernorm(self) -> nn.LayerNorm:
-        if self._original_mod.use_layernorm1p:
-            return self.get_layernorm1p(self._original_mod.ln_1)
-        else:
-            return self._original_mod.ln_1
+        if self.use_layernorm1p:
+            return self.get_layernorm1p(self.pre_attention_layernorm)
+        return self.pre_attention_layernorm
 
     def get_post_attention_layernorm(self) -> nn.LayerNorm:
-        if self._original_mod.use_layernorm1p:
-            return self.get_layernorm1p(self._original_mod.ln_2)
-        else:
-            return self._original_mod.ln_2
+        if self.use_layernorm1p:
+            return self.get_layernorm1p(self.post_attention_layernorm)
+        return self.post_attention_layernorm
 
 
 class MidmAttention(DecoderOnlyAttention):
-    def __post_init__(self):
-        self.c_attn = self._original_mod.c_attn
-        self.o_proj = self._original_mod.c_proj
-        self.split_size = self._original_mod.split_size
-        self.num_key_value_heads = self._original_mod.num_heads
+    def __post_init__(self, self_attn):
+        self.c_attn = self_attn.c_attn
+        self.o_proj = self_attn.c_proj
+        self.split_size = self_attn.split_size
+        self.num_key_value_heads = self_attn.num_heads
 
     def projection(self, hidden_states, lora_int_id) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         if lora_int_id is not None:
@@ -130,12 +137,12 @@ class MidmAttention(DecoderOnlyAttention):
         query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
         return query_states, key_states, value_states
 
-    def get_attn_scale(self):
+    def get_attn_scale(self, self_attn):
         scale = 1.0
-        if self._original_mod.scale_attn_weights:
+        if self_attn.scale_attn_weights:
             scale /= math.sqrt(self.head_dim)
 
-        if self._original_mod.scale_attn_by_inverse_layer_idx and not self._original_mod.scale_qk_by_inverse_layer_idx:
+        if self_attn.scale_attn_by_inverse_layer_idx:
             scale /= 1 + self.layer_idx
 
         return scale

optimum/rbln/transformers/models/mixtral/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .configuration_mixtral import RBLNMixtralForCausalLMConfig
+from .modeling_mixtral import RBLNMixtralForCausalLM

optimum/rbln/transformers/models/mixtral/configuration_mixtral.py

@@ -0,0 +1,38 @@
+# Copyright 2026 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+
+
+class RBLNMixtralForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    """
+    Configuration class for RBLN Mixtral models.
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNMixtralForCausalLM, RBLNMixtralForCausalLMConfig
+    # Create a configuration object
+    config = RBLNMixtralForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=32768,
+        tensor_parallel_size=4
+    )
+    # Use the configuration with from_pretrained
+    model = RBLNMixtralForCausalLM.from_pretrained(
+        "mistralai/Mixtral-8x7B-Instruct-v0.1",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

optimum/rbln/transformers/models/mixtral/mixtral_architecture.py

@@ -0,0 +1,76 @@
+# Copyright 2026 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+import torch
+from torch import nn
+
+from ..decoderonly.configuration_decoderonly import RBLNLoRAConfig
+from ..decoderonly.decoderonly_architecture import DecoderOnlyAttention, DecoderOnlyLayer, DecoderOnlyWrapper
+
+
+class MixtralWrapper(DecoderOnlyWrapper):
+    def get_rbln_layer_class(self):
+        return MixtralLayer
+
+
+class MixtralLayer(DecoderOnlyLayer):
+    _MLP_ATTR = ("block_sparse_moe",)
+
+    def __init__(self, layer, self_attn: DecoderOnlyAttention, lora_config: Optional[RBLNLoRAConfig] = None):
+        super().__init__(layer, self_attn, lora_config)
+        self.block_sparse_moe = MixtralSparseMoeBlock(self.mlp)
+
+    def get_mlp(self) -> nn.Module:
+        return self.block_sparse_moe
+
+
+class MixtralSparseMoeBlock(nn.Module):
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        # self.num_experts = model.num_experts
+        self.top_k = model.top_k
+        self.gate = model.gate
+        self.experts = MixtralBlockSparseTop2MLP(model.experts, self.top_k)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        # router_logits: (batch * sequence_length, n_experts)
+        router_logits = self.gate(hidden_states)
+        final_hidden_states = self.experts(hidden_states, router_logits)
+        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+        return final_hidden_states
+
+
+class MixtralBlockSparseTop2MLP(nn.Module):
+    def __init__(self, expert_list, top_k):
+        super().__init__()
+        self.hidden_dim = expert_list[0].hidden_dim
+        self.ffn_dim = expert_list[0].ffn_dim
+        self.top_k = top_k
+
+        self.num_experts = len(expert_list)
+        self.w1 = nn.Linear(self.hidden_dim, self.num_experts * self.ffn_dim, bias=False)
+        self.w2 = nn.Linear(self.num_experts * self.ffn_dim, self.hidden_dim, bias=False)
+        self.w3 = nn.Linear(self.hidden_dim, self.num_experts * self.ffn_dim, bias=False)
+        self.w1.weight.data = torch.stack([expert.w1.weight.data for expert in expert_list], dim=0)
+        self.w2.weight.data = torch.stack([expert.w2.weight.data for expert in expert_list], dim=0)
+        self.w3.weight.data = torch.stack([expert.w3.weight.data for expert in expert_list], dim=0)
+
+    def forward(self, x, router_logits):
+        return torch.ops.rbln_custom_ops.custom_moe_glu(
+            x, self.w1.weight, self.w3.weight, self.w2.weight, router_logits, self.top_k, True
+        )

optimum/rbln/transformers/models/mixtral/modeling_mixtral.py

@@ -0,0 +1,68 @@
+# Copyright 2026 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from .mixtral_architecture import MixtralWrapper
+
+
+class RBLNMixtralForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The Mixtral is a Mixture-of-Experts (MoE) variant of Mixtral, available as a base model and an aligned chat model.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    A class to convert and run pre-trained transformers based MixtralForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers MixtralForCausalLM model into a RBLN transformer model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+    **Configuration:**
+    This model uses [`RBLNMixtralForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNMixtralForCausalLMConfig`] or a dictionary conforming to its structure.
+    See the [`RBLNMixtralForCausalLMConfig`] class for all available configuration options.
+    Examples:
+        ```python
+        from optimum.rbln import RBLNMixtralForCausalLM
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNMixtralForCausalLM.from_pretrained(
+            "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 32768,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNMixtralForCausalLM.from_pretrained(
+            "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            export=True,
+            rbln_config=rbln_config
+        )
+        # Using a RBLNMixtralForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNMixtralForCausalLMConfig
+        config = RBLNMixtralForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=32768,
+            tensor_parallel_size=4
+        )
+        model = RBLNMixtralForCausalLM.from_pretrained(
+            "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    _decoder_wrapper_cls = MixtralWrapper