optimum-rbln 0.9.3rc0__py3-none-any.whl → 0.9.5a4__py3-none-any.whl

optimum/rbln/transformers/models/qwen3/modeling_qwen3.py

@@ -12,24 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING
-
-from transformers import PretrainedConfig
 
 from ....utils import logging
 from ...models.decoderonly import (
     RBLNDecoderOnlyModel,
     RBLNDecoderOnlyModelForCausalLM,
-    RBLNDecoderOnlyModelForCausalLMConfig,
 )
 from .qwen3_architecture import Qwen3Wrapper
 
 
 logger = logging.get_logger(__name__)
 
-if TYPE_CHECKING:
-    from transformers import PretrainedConfig
-
 
 class RBLNQwen3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
@@ -84,19 +77,6 @@ class RBLNQwen3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     _decoder_wrapper_cls = Qwen3Wrapper
 
-    @classmethod
-    def _update_sliding_window_config(
-        cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
-    ):
-        # https://github.com/huggingface/transformers/issues/35896
-        # There seems to be a bug in transformers(v4.52.4). Therefore, similar to when attn_implementation is eager,
-        # we set all layers to use sliding window in this version. This should be updated once the bug is fixed.
-
-        rbln_config.cache_impl = "sliding_window"
-        rbln_config.sliding_window = model_config.sliding_window
-        rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
-        return rbln_config
-
     def forward(self, *args, **kwargs):
         kwargs["return_dict"] = True
         return super().forward(*args, **kwargs)

optimum/rbln/transformers/models/qwen3/qwen3_architecture.py

@@ -22,10 +22,10 @@ class Qwen3Wrapper(DecoderOnlyWrapper):
 
 
 class Qwen3Attention(DecoderOnlyAttention):
-    def __post_init__(self):
-        self.k_proj = self._original_mod.k_proj
-        self.v_proj = self._original_mod.v_proj
-        self.q_proj = self._original_mod.q_proj
-        self.o_proj = self._original_mod.o_proj
-        self.q_norm = self._original_mod.q_norm
-        self.k_norm = self._original_mod.k_norm
+    def __post_init__(self, self_attn):
+        self.q_proj = self_attn.q_proj
+        self.k_proj = self_attn.k_proj
+        self.v_proj = self_attn.v_proj
+        self.o_proj = self_attn.o_proj
+        self.q_norm = self_attn.q_norm
+        self.k_norm = self_attn.k_norm

optimum/rbln/transformers/models/qwen3_moe/__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_qwen3_moe import RBLNQwen3MoeForCausalLMConfig
+from .modeling_qwen3_moe import RBLNQwen3MoeForCausalLM

optimum/rbln/transformers/models/qwen3_moe/configuration_qwen3_moe.py

@@ -0,0 +1,38 @@
+# 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 ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+
+
+class RBLNQwen3MoeForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    """
+    Configuration class for RBLN Qwen3 Moe models.
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNQwen3MoeForCausalLM, RBLNQwen3MoeForCausalLMConfig
+    # Create a configuration object
+    config = RBLNQwen3MoeForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=262144,
+        tensor_parallel_size=4
+    )
+    # Use the configuration with from_pretrained
+    model = RBLNQwen3MoeForCausalLM.from_pretrained(
+        "Qwen/Qwen3-30B-A3B-Thinking-2507",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

optimum/rbln/transformers/models/qwen3_moe/modeling_qwen3_moe.py

@@ -0,0 +1,68 @@
+# 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 ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from .qwen3_moe_architecture import Qwen3MoeWrapper
+
+
+class RBLNQwen3MoeForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The Qwen3 Moe is a Mixture-of-Experts (MoE) variant of Qwen3, 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 Qwen3MoeForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers Qwen3MoeForCausalLM 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 [`RBLNQwen3MoeForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNQwen3MoeForCausalLMConfig`] or a dictionary conforming to its structure.
+    See the [`RBLNQwen3MoeForCausalLMConfig`] class for all available configuration options.
+    Examples:
+        ```python
+        from optimum.rbln import RBLNQwen3MoeForCausalLM
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNQwen3MoeForCausalLM.from_pretrained(
+            "Qwen/Qwen3-30B-A3B-Thinking-2507",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 262144,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNQwen3MoeForCausalLM.from_pretrained(
+            "Qwen/Qwen3-30B-A3B-Thinking-2507",
+            export=True,
+            rbln_config=rbln_config
+        )
+        # Using a RBLNQwen3ForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNQwen3MoeForCausalLMConfig
+        config = RBLNQwen3MoeForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=262144,
+            tensor_parallel_size=4
+        )
+        model = RBLNQwen3MoeForCausalLM.from_pretrained(
+            "Qwen/Qwen3-30B-A3B-Thinking-2507",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    _decoder_wrapper_cls = Qwen3MoeWrapper

optimum/rbln/transformers/models/qwen3_moe/qwen3_moe_architecture.py

@@ -0,0 +1,100 @@
+# 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
+from torch import nn
+
+from ..decoderonly.configuration_decoderonly import RBLNLoRAConfig
+from ..decoderonly.decoderonly_architecture import DecoderOnlyAttention, DecoderOnlyLayer, DecoderOnlyWrapper
+
+
+class Qwen3MoeWrapper(DecoderOnlyWrapper):
+    def get_rbln_layer_class(self):
+        return Qwen3MoeLayer
+
+    def get_rbln_attn_class(self):
+        return Qwen3MoeAttention
+
+
+class Qwen3MoeAttention(DecoderOnlyAttention):
+    def __post_init__(self, self_attn):
+        self.q_proj = self_attn.q_proj
+        self.k_proj = self_attn.k_proj
+        self.v_proj = self_attn.v_proj
+        self.o_proj = self_attn.o_proj
+        self.q_norm = self_attn.q_norm
+        self.k_norm = self_attn.k_norm
+
+
+class Qwen3MoeLayer(DecoderOnlyLayer):
+    def __init__(self, layer, self_attn: DecoderOnlyAttention, lora_config: Optional[RBLNLoRAConfig] = None):
+        super().__init__(layer, self_attn, lora_config)
+        self.mlp = (
+            Qwen3MoeSparseMoeBlock(layer.mlp)
+            if layer.mlp.__class__.__name__ == "Qwen3MoeSparseMoeBlock"
+            else layer.mlp
+        )
+
+    def get_mlp(self) -> nn.Module:
+        return self.mlp
+
+
+class Qwen3MoeSparseMoeBlock(nn.Module):
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        self.num_experts = model.num_experts
+        self.top_k = model.top_k
+        self.norm_topk_prob = model.norm_topk_prob
+        self.gate = model.gate
+        self.experts = Qwen3MoeMLP(model.experts, self.top_k, self.norm_topk_prob)
+
+    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 Qwen3MoeMLP(nn.Module):
+    def __init__(self, expert_list, top_k, norm_topk_prob):
+        super().__init__()
+        self.hidden_size = expert_list[0].hidden_size
+        self.intermediate_size = expert_list[0].intermediate_size
+        self.top_k = top_k
+        self.norm_topk_prob = norm_topk_prob
+        self.num_experts = len(expert_list)
+        self.gate_proj = nn.Linear(self.hidden_size, self.num_experts * self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.num_experts * self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.num_experts * self.intermediate_size, self.hidden_size, bias=False)
+        self.gate_proj.weight.data = torch.stack([expert.gate_proj.weight.data for expert in expert_list], dim=0)
+        self.up_proj.weight.data = torch.stack([expert.up_proj.weight.data for expert in expert_list], dim=0)
+        self.down_proj.weight.data = torch.stack([expert.down_proj.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.gate_proj.weight,
+            self.up_proj.weight,
+            self.down_proj.weight,
+            router_logits,
+            self.top_k,
+            self.norm_topk_prob,
+        )

optimum/rbln/transformers/models/resnet/configuration_resnet.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 
+from typing import Optional
+
 from ...configuration_generic import RBLNModelForImageClassificationConfig
 
 
@@ -23,3 +25,18 @@ class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConf
     This configuration class stores the configuration parameters specific to
     RBLN-optimized ResNet models for image classification tasks.
     """
+
+    def __init__(self, output_hidden_states: Optional[bool] = None, **kwargs):
+        """
+        Args:
+            image_size (Optional[Union[int, Tuple[int, int]]]): The size of input images.
+                Can be an integer for square images or a tuple (height, width).
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            output_hidden_states (bool, optional) — Whether or not to return the hidden states of all layers.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.output_hidden_states = output_hidden_states

optimum/rbln/transformers/models/resnet/modeling_resnet.py

@@ -13,7 +13,17 @@
 # limitations under the License.
 
 
+from typing import TYPE_CHECKING, Optional, Tuple, Union
+
+import torch
+from transformers.modeling_outputs import ImageClassifierOutputWithNoAttention
+
 from ...modeling_generic import RBLNModelForImageClassification
+from .configuration_resnet import RBLNResNetForImageClassificationConfig
+
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig, PreTrainedModel
 
 
 class RBLNResNetForImageClassification(RBLNModelForImageClassification):
@@ -24,3 +34,66 @@ class RBLNResNetForImageClassification(RBLNModelForImageClassification):
     on RBLN devices, supporting image classification with convolutional neural networks
     designed for computer vision tasks.
     """
+
+    @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["RBLNResNetForImageClassificationConfig"] = None,
+    ) -> "RBLNResNetForImageClassificationConfig":
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
+
+        rbln_config = super()._update_rbln_config(
+            preprocessors=preprocessors,
+            model=model,
+            model_config=model_config,
+            rbln_config=rbln_config,
+        )
+
+        return rbln_config
+
+    @classmethod
+    def _wrap_model_if_needed(
+        cls, model: torch.nn.Module, rbln_config: "RBLNResNetForImageClassificationConfig"
+    ) -> torch.nn.Module:
+        class _ResNetForImageClassification(torch.nn.Module):
+            def __init__(self, model: torch.nn.Module, output_hidden_states: bool):
+                super().__init__()
+                self.model = model
+                self.output_hidden_states = output_hidden_states
+
+            def forward(self, *args, **kwargs):
+                output = self.model(*args, output_hidden_states=self.output_hidden_states, **kwargs)
+                return output
+
+        return _ResNetForImageClassification(model, rbln_config.output_hidden_states)
+
+    def forward(
+        self, pixel_values: torch.Tensor, output_hidden_states: bool = None, return_dict: bool = None, **kwargs
+    ) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:
+        """
+        Foward pass for the RBLN-optimized ResNet model for image classification.
+
+        Args:
+            pixel_values (torch.FloatTensor of shape (batch_size, channels, height, width)): The tensors corresponding to the input images.
+            output_hidden_states (bool, *optional*, defaults to False): Whether or not to return the hidden states of all layers.
+                See hidden_states under returned tensors for more details.
+            return_dict (bool, *optional*, defaults to True): Whether to return a dictionary of outputs.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a ImageClassifierOutputWithNoAttention object.
+        """
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        if output_hidden_states != self.rbln_config.output_hidden_states:
+            raise ValueError(
+                f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
+                f"Please compile again with the correct argument."
+            )
+
+        return super().forward(pixel_values=pixel_values, return_dict=return_dict, **kwargs)

optimum/rbln/transformers/models/roberta/modeling_roberta.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Tuple, Union
+
+import torch
+from transformers.modeling_outputs import MaskedLMOutput, SequenceClassifierOutput
+
 from ...modeling_generic import RBLNModelForMaskedLM, RBLNModelForSequenceClassification
 
 
@@ -26,6 +31,19 @@ class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
 
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> Union[Tuple, MaskedLMOutput]:
+        """
+        Forward pass for the RBLN-optimized RoBERTa model for masked language modeling tasks.
+
+        Args:
+            input_ids (torch.LongTensor of shape (batch_size, sequence_length), optional): Indices of input sequence tokens in the vocabulary.
+            attention_mask (torch.FloatTensor of shape (batch_size, sequence_length), optional): Mask to avoid performing attention on padding token indices.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a MaskedLMOutput object.
+        """
+        return super().forward(input_ids, attention_mask, **kwargs)
+
 
 class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     """
@@ -37,3 +55,18 @@ class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     """
 
     rbln_model_input_names = ["input_ids", "attention_mask"]
+
+    def forward(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs
+    ) -> Union[Tuple, SequenceClassifierOutput]:
+        """
+        Forward pass for the RBLN-optimized RoBERTa model for sequence classification tasks.
+
+        Args:
+            input_ids (torch.LongTensor of shape (batch_size, sequence_length), optional): Indices of input sequence tokens in the vocabulary.
+            attention_mask (torch.FloatTensor of shape (batch_size, sequence_length), optional): Mask to avoid performing attention on padding token indices.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a SequenceClassifierOutput object.
+        """
+        return super().forward(input_ids, attention_mask, **kwargs)

optimum/rbln/transformers/models/seq2seq/configuration_seq2seq.py

@@ -15,6 +15,7 @@
 from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.deprecation import deprecate_kwarg
 from ....utils.logging import get_logger
 
 
@@ -24,13 +25,13 @@ logger = get_logger()
 class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
     support_paged_attention = None
 
+    @deprecate_kwarg(old_name="pad_token_id", version="0.10.0")
     def __init__(
         self,
         batch_size: Optional[int] = None,
         enc_max_seq_len: Optional[int] = None,
         dec_max_seq_len: Optional[int] = None,
         use_attention_mask: Optional[bool] = None,
-        pad_token_id: Optional[int] = None,
         kvcache_num_blocks: Optional[int] = None,
         kvcache_block_size: Optional[int] = None,
         **kwargs: Any,
@@ -41,7 +42,6 @@ class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
             enc_max_seq_len (Optional[int]): Maximum sequence length for the encoder.
             dec_max_seq_len (Optional[int]): Maximum sequence length for the decoder.
             use_attention_mask (Optional[bool]): Whether to use attention masks during inference.
-            pad_token_id (Optional[int]): The ID of the padding token in the vocabulary.
             kvcache_num_blocks (Optional[int]): The total number of blocks to allocate for the
                 PagedAttention KV cache for the SelfAttention. Defaults to batch_size.
             kvcache_block_size (Optional[int]): Sets the size (in number of tokens) of each block
@@ -61,8 +61,6 @@ class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
 
         self.use_attention_mask = use_attention_mask
 
-        self.pad_token_id = pad_token_id
-
         if self.support_paged_attention:
             self.kvcache_num_blocks = kvcache_num_blocks
             self.kvcache_block_size = kvcache_block_size

optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py

@@ -20,8 +20,9 @@ import rebel
 import torch
 from rebel.compile_context import CompileContext
 from transformers import AutoModelForSeq2SeqLM, PretrainedConfig, PreTrainedModel
+from transformers.generation.configuration_utils import GenerationConfig
 from transformers.generation.utils import GenerationMixin
-from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
+from transformers.modeling_outputs import BaseModelOutput, ModelOutput, Seq2SeqLMOutput
 
 from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
@@ -33,7 +34,7 @@ from .configuration_seq2seq import RBLNModelForSeq2SeqLMConfig
 logger = get_logger(__name__)
 
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, GenerationConfig, PretrainedConfig
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
 
 
 class RBLNRuntimeEncoder(RBLNPytorchRuntime):
@@ -140,7 +141,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, GenerationMixin, ABC):
     @classmethod
     @torch.inference_mode()
     def get_compiled_model(cls, model: PreTrainedModel, rbln_config: RBLNModelForSeq2SeqLMConfig):
-        wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
+        wrapped_model = cls._wrap_model_if_needed(model, rbln_config)
 
         enc_compile_config = rbln_config.compile_cfgs[0]
         dec_compile_config = rbln_config.compile_cfgs[1]
@@ -209,8 +210,8 @@ class RBLNModelForSeq2SeqLM(RBLNModel, GenerationMixin, ABC):
         if not cls.support_causal_attn:
             rbln_config.use_attention_mask = True
 
-        n_layer = getattr(model_config, "decoder_layers", None) or getattr(model_config, "num_layers")
-        n_head = getattr(model_config, "decoder_attention_heads", None) or getattr(model_config, "num_heads")
+        n_layer = getattr(model_config, "decoder_layers", None) or model_config.num_layers
+        n_head = getattr(model_config, "decoder_attention_heads", None) or model_config.num_heads
         d_kv = (
             model_config.d_kv
             if hasattr(model_config, "d_kv")
@@ -221,12 +222,6 @@ class RBLNModelForSeq2SeqLM(RBLNModel, GenerationMixin, ABC):
             model_config, "max_position_embeddings", None
         )
 
-        pad_token_id = getattr(model_config, "pad_token_id", None)
-        pad_token_id = pad_token_id or getattr(model_config, "bos_token_id", None)
-        pad_token_id = pad_token_id or getattr(model_config, "eos_token_id", None)
-        pad_token_id = pad_token_id or -1
-        rbln_config.pad_token_id = pad_token_id
-
         if rbln_config.enc_max_seq_len is None:
             enc_max_seq_len = max_position_embeddings
             for tokenizer in preprocessors:
@@ -432,7 +427,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, GenerationMixin, ABC):
         inputs_tensor = torch.nn.functional.pad(
             inputs_tensor,
             (0, self.rbln_config.enc_max_seq_len - input_len),
-            value=self.rbln_config.pad_token_id,
+            value=self.config.pad_token_id,
         )
         model_kwargs["attention_mask"] = torch.nn.functional.pad(
             model_kwargs["attention_mask"], (0, self.rbln_config.enc_max_seq_len - input_len)
@@ -451,3 +446,32 @@ class RBLNModelForSeq2SeqLM(RBLNModel, GenerationMixin, ABC):
             model_kwargs["encoder_outputs"] = encoder(**encoder_kwargs, block_tables=block_tables)
 
         return model_kwargs
+
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        **kwargs,
+    ) -> Union[ModelOutput, torch.LongTensor]:
+        """
+        The generate function is utilized in its standard form as in the HuggingFace transformers library. User can use this function to generate text from the model.
+        Check the [HuggingFace transformers documentation](https://huggingface.co/docs/transformers/v4.57.1/en/main_classes/text_generation#transformers.GenerationMixin.generate) for more details.
+
+        Args:
+            input_ids (torch.LongTensor): The input ids to the model.
+            attention_mask (torch.LongTensor, optional): The attention mask to the model.
+            generation_config (GenerationConfig, optional): The generation configuration to be used as base parametrization for the generation call. **kwargs passed to generate matching the attributes of generation_config will override them.
+                If generation_config is not provided, the default will be used, which had the following loading priority: 1) from the generation_config.json model file, if it exists; 2) from the model configuration.
+                Please note that unspecified parameters will inherit [GenerationConfig](https://huggingface.co/docs/transformers/v4.57.1/en/main_classes/text_generation#transformers.GenerationConfig)’s default values.
+            kwargs (dict[str, Any], optional): Additional arguments passed to the generate function. See the HuggingFace transformers documentation for more details.
+
+        Returns:
+            Generates sequences of token ids for models with a language modeling head.
+        """
+        if generation_config is not None:
+            kwargs["generation_config"] = generation_config
+        if attention_mask is not None:
+            kwargs["attention_mask"] = attention_mask
+
+        return super().generate(input_ids, **kwargs)

optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py

@@ -268,13 +268,12 @@ class Seq2SeqDecoder(torch.nn.Module):
 
     def __init__(self, model, layers, **kwargs):
         super().__init__()
-        self._original_mod = model
         self.layers = nn.ModuleList(layers)
         self.embed_tokens = model.embed_tokens
-        self.final_layer_norm = getattr(model, "final_layer_norm", None)
-        self.__post_init__(**kwargs)
+        self.final_layer_norm = getattr(model, "final_layer_norm", None) or getattr(model, "layer_norm", None)
+        self.__post_init__(model, **kwargs)
 
-    def __post_init__(self, **kwargs):
+    def __post_init__(self, model: nn.Module, **kwargs):
         """
         Abstract method intended to be overridden by subclasses to modify or override
         the attributes of the original model after initialization.
@@ -344,12 +343,11 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
 
     def __init__(self, decoder_layer, self_attn, cross_attn):
         super().__init__()
-        self._original_mod = decoder_layer
         self.self_attn = self_attn
         self.cross_attn = cross_attn
-        self.__post_init__()
+        self.__post_init__(decoder_layer)
 
-    def __post_init__(self, **kwargs):
+    def __post_init__(self, decoder_layer: nn.Module, **kwargs):
         """
         Abstract method intended to be overridden by subclasses to modify or override
         the attributes of the original model after initialization.
@@ -423,10 +421,9 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
 class Seq2SeqSelfAttention(nn.Module):
     def __init__(self, attn, **kwargs):
         super().__init__()
-        self._original_mod = attn
-        self.__post_init__(**kwargs)
+        self.__post_init__(attn, **kwargs)
 
-    def __post_init__(self, **kwargs):
+    def __post_init__(self, attn: nn.Module, **kwargs):
         """
         Abstract method intended to be overridden by subclasses to modify or override
         the attributes of the original model after initialization.
@@ -495,8 +492,13 @@ class Seq2SeqSelfAttention(nn.Module):
 class Seq2SeqCrossAttention(nn.Module):
     def __init__(self, attn, **kwargs):
         super().__init__()
-        self._original_mod = attn
-        self.__post_init__(**kwargs)
+        self.__post_init__(attn, **kwargs)
+
+    def __post_init__(self, attn: nn.Module, **kwargs):
+        """
+        Optional post-init hook for subclasses (e.g., to register q/k/v/out projections).
+        """
+        pass
 
     def forward(
         self,