optimum-rbln 0.8.0.post2__py3-none-any.whl → 0.8.1__py3-none-any.whl

optimum/rbln/transformers/models/qwen2/modeling_qwen2.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from transformers import PretrainedConfig
+
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyModelForCausalLMConfig
 from .qwen2_architecture import QWEN2Wrapper
 
 
@@ -22,13 +24,74 @@ logger = logging.get_logger(__name__)
 
 class RBLNQwen2ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
-    The Llama Model transformer with a language modeling head (linear layer) on top.
+    The Qwen2 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 LlamaForCausalLM model on RBLN devices.
-    It implements the methods to convert a pre-trained transformers LlamaForCausalLM model into a RBLN transformer model by:
+    A class to convert and run pre-trained transformers based Qwen2ForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers Qwen2ForCausalLM 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 [`RBLNQwen2ForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNQwen2ForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNQwen2ForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNQwen2ForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNQwen2ForCausalLM.from_pretrained(
+            "Qwen/Qwen2-7B-Instruct",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNQwen2ForCausalLM.from_pretrained(
+            "Qwen/Qwen2-7B-Instruct",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+
+        # Using a RBLNQwen2ForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNQwen2ForCausalLMConfig
+
+        config = RBLNQwen2ForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNQwen2ForCausalLM.from_pretrained(
+            "Qwen/Qwen2-7B-Instruct",
+            export=True,
+            rbln_config=config
+        )
+        ```
     """
 
     _decoder_wrapper_cls = QWEN2Wrapper
+
+    @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

optimum/rbln/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py

@@ -12,20 +12,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from ....configuration_utils import RBLNModelConfig
 from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNQwen2_5_VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    """
+    Configuration class for RBLNQwen2_5_VLForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Qwen2.5-VL models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
     submodules = ["visual"]
 
     def __init__(
         self,
         visual: Optional[RBLNModelConfig] = None,
         use_inputs_embeds: bool = True,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         super().__init__(use_inputs_embeds=use_inputs_embeds, **kwargs)
         if not self.use_inputs_embeds:
@@ -37,7 +45,15 @@ class RBLNQwen2_5_VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausal
 
 
 class RBLNQwen2_5_VisionTransformerPretrainedModelConfig(RBLNModelConfig):
-    def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs):
+    """
+    Configuration class for RBLNQwen2_5_VisionTransformerPretrainedModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Qwen2.5-VL vision transformer models with window-based attention
+    mechanisms for processing images and videos.
+    """
+
+    def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs: Dict[str, Any]):
         """
         Args:
             max_seq_lens (Optional[Union[int, List[int]]]): Maximum sequence lengths for Vision
@@ -54,6 +70,18 @@ class RBLNQwen2_5_VisionTransformerPretrainedModelConfig(RBLNModelConfig):
 
         Raises:
             ValueError: If batch_size is not a positive integer.
+
+        Max Seq Lens:
+            Since `Qwen2_5_VLForConditionalGeneration` performs inference on a per-image or per-frame basis,
+            `max_seq_lens` should be set based on the maximum expected resolution of the input images or video frames,
+            according to the following guidelines:
+
+            1. **Minimum Value**: `max_seq_lens` must be greater than or equal to the number of patches generated from the input image.
+                For example, a 224x224 image with a patch size of 14 results in (224 / 14) * (224 / 14) = 256 patches.
+                Therefore, `max_seq_lens` must be at least 256.
+            2. **Alignment Requirement**: `max_seq_lens` must be a multiple of `(window_size / patch_size)^2` due to the requirements
+                of the window-based attention mechanism. For instance, if `window_size` is 112 and `patch_size` is 14, then
+                `(112 / 14)^2 = 64`, meaning valid values for `max_seq_lens` include 64, 128, 192, 256, etc.
         """
         super().__init__(**kwargs)
 

optimum/rbln/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py

@@ -37,6 +37,7 @@ from ....utils.logging import get_logger
 from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyOutput
 from .configuration_qwen2_5_vl import (
     RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
+    RBLNQwen2_5_VLForConditionalGenerationConfig,
 )
 from .qwen2_5_vl_architecture import Qwen2_5_VisionTransformerWrapper, Qwen2_5_VL_LanguageModelWrapper
 
@@ -53,6 +54,14 @@ if TYPE_CHECKING:
 
 
 class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
+    """
+    RBLN optimized Qwen2.5-VL vision transformer model.
+
+    This class provides hardware-accelerated inference for Qwen2.5-VL vision transformers
+    on RBLN devices, supporting image and video encoding for multimodal vision-language tasks
+    with window-based attention mechanisms.
+    """
+
     auto_model_class = None
 
     def __post_init__(self, **kwargs):
@@ -338,6 +347,40 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
 
 
 class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
+    """
+    RBLNQwen2_5_VLForConditionalGeneration is a multi-modal model that integrates 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 [`RBLNDecoderOnlyModelForCausalLM`]. 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). 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
+        `from_pretrained` method. Refer to the `from_pretrained` documentation and the RBLNQwen2_5_VLForConditionalGenerationConfig class for details.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNQwen2_5_VLForConditionalGeneration
+
+        model = RBLNQwen2_5_VLForConditionalGeneration.from_pretrained(
+            "Qwen/Qwen2.5-VL-7B-Instruct",
+            export=True,
+            rbln_config={
+                "visual": {
+                    "max_seq_lens": 6400,
+                    "device": 0,
+                },
+                "tensor_parallel_size": 8,
+                "kvcache_partition_len": 16_384,
+                "max_seq_len": 114_688,
+                "device": [0, 1, 2, 3, 4, 5, 6, 7],
+            },
+        )
+
+        model.save_pretrained("compiled-qwen2.5-vl-7b-instruct")
+        ```
+    """
+
     auto_model_class = AutoModelForVision2Seq
     _rbln_submodules = [
         {"name": "visual"},
@@ -369,33 +412,19 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         cls,
         batch_size: int,
         query_length: int,
-        use_inputs_embeds: bool,
-        use_attention_mask: bool,
-        use_position_ids: bool,
-        max_seq_len: int,
-        kvcache_block_size: int,
-        kvcache_num_blocks: int,
-        num_key_value_heads: int,
-        num_hidden_layers: int,
-        hidden_size: int,
-        head_dim: int,
+        rbln_config: RBLNQwen2_5_VLForConditionalGenerationConfig,
+        model_config: PretrainedConfig,
     ):
-        input_info = super().get_input_info(
-            batch_size,
-            query_length,
-            use_inputs_embeds,
-            use_attention_mask,
-            use_position_ids,
-            max_seq_len,
-            kvcache_block_size,
-            kvcache_num_blocks,
-            num_key_value_heads,
-            num_hidden_layers,
-            hidden_size,
-            head_dim,
-        )
+        input_info = super().get_input_info(batch_size, query_length, rbln_config, model_config)
         pos_idx = 3
-        input_info.insert(pos_idx, ("position_emb", [2, batch_size, 1, query_length, head_dim], "float32"))
+        input_info.insert(
+            pos_idx,
+            (
+                "position_emb",
+                [2, batch_size, 1, query_length, model_config.hidden_size // model_config.num_attention_heads],
+                "float32",
+            ),
+        )
 
         return input_info
 

optimum/rbln/transformers/models/qwen2_5_vl/qwen2_5_vl_architecture.py

@@ -79,7 +79,7 @@ class Qwen2_5_VLVisionFullAttention(nn.Module):
         super().__init__()
         self._origin_model = model
         self.num_heads = model.num_heads
-        self.head_dim = model.head_dim
+        self.head_dim = getattr(model, "head_dim", model.proj.in_features // model.num_heads)
         self.qkv = model.qkv
         self.proj = model.proj
 
@@ -114,7 +114,7 @@ class Qwen2_5_VLVisionWindowAttention(nn.Module):
         super().__init__()
         self._origin_model = model
         self.num_heads = model.num_heads
-        self.head_dim = model.head_dim
+        self.head_dim = getattr(model, "head_dim", model.proj.in_features // model.num_heads)
         self.qkv = model.qkv
         self.proj = model.proj
         self.window_seq_len = window_seq_len
@@ -162,7 +162,8 @@ class Qwen2_5_VL_LanguageModelWrapper(DecoderOnlyWrapper):
         input_ids = None if self.use_inputs_embeds else args.pop(0)
         inputs_embeds = args.pop(0) if self.use_inputs_embeds else None
         cache_position = args.pop(0)
-        block_tables = args.pop(0)
+        global_block_tables = args.pop(0)
+        local_block_tables = None
         position_embeds = args.pop(0)
         query_position = args.pop(0) if self.phase == "prefill" else None
         position_ids = None
@@ -188,7 +189,8 @@ class Qwen2_5_VL_LanguageModelWrapper(DecoderOnlyWrapper):
             input_ids,
             inputs_embeds,
             cache_position,
-            block_tables,
+            global_block_tables,
+            local_block_tables,
             query_position,
             attention_mask,
             position_ids,

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

@@ -0,0 +1,23 @@
+# 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_resnet import RBLNResNetForImageClassificationConfig
+from .modeling_resnet import RBLNResNetForImageClassification
+
+
+__all__ = [
+    "RBLNResNetForImageClassificationConfig",
+    "RBLNResNetForImageClassification",
+]

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

@@ -0,0 +1,25 @@
+# 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_generic import RBLNModelForImageClassificationConfig
+
+
+class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConfig):
+    """
+    Configuration class for RBLNResNetForImageClassification.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized ResNet models for image classification tasks.
+    """

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

@@ -0,0 +1,26 @@
+# 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 ...modeling_generic import RBLNModelForImageClassification
+
+
+class RBLNResNetForImageClassification(RBLNModelForImageClassification):
+    """
+    RBLN optimized ResNet model for image classification tasks.
+
+    This class provides hardware-accelerated inference for ResNet models
+    on RBLN devices, supporting image classification with convolutional neural networks
+    designed for computer vision tasks.
+    """

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

@@ -0,0 +1,24 @@
+# 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_roberta import RBLNRobertaForMaskedLMConfig, RBLNRobertaForSequenceClassificationConfig
+from .modeling_roberta import RBLNRobertaForMaskedLM, RBLNRobertaForSequenceClassification
+
+
+__all__ = [
+    "RBLNRobertaForMaskedLMConfig",
+    "RBLNRobertaForSequenceClassificationConfig",
+    "RBLNRobertaForMaskedLM",
+    "RBLNRobertaForSequenceClassification",
+]

optimum/rbln/transformers/{configuration_alias.py → models/roberta/configuration_roberta.py}

@@ -12,38 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_generic import (
-    RBLNModelForAudioClassificationConfig,
-    RBLNModelForImageClassificationConfig,
-    RBLNModelForMaskedLMConfig,
-    RBLNModelForQuestionAnsweringConfig,
-    RBLNModelForSequenceClassificationConfig,
-)
+from ...configuration_generic import RBLNModelForMaskedLMConfig, RBLNModelForSequenceClassificationConfig
 
 
-class RBLNASTForAudioClassificationConfig(RBLNModelForAudioClassificationConfig):
-    pass
-
-
-class RBLNDistilBertForQuestionAnsweringConfig(RBLNModelForQuestionAnsweringConfig):
-    pass
-
-
-class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConfig):
-    pass
-
+class RBLNRobertaForMaskedLMConfig(RBLNModelForMaskedLMConfig):
+    """
+    Configuration class for RBLNRobertaForMaskedLM.
 
-class RBLNXLMRobertaForSequenceClassificationConfig(RBLNModelForSequenceClassificationConfig):
-    pass
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized RoBERTa models for masked language modeling tasks.
+    """
 
 
 class RBLNRobertaForSequenceClassificationConfig(RBLNModelForSequenceClassificationConfig):
-    pass
-
-
-class RBLNRobertaForMaskedLMConfig(RBLNModelForMaskedLMConfig):
-    pass
-
+    """
+    Configuration class for RBLNRobertaForSequenceClassification.
 
-class RBLNViTForImageClassificationConfig(RBLNModelForImageClassificationConfig):
-    pass
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized RoBERTa models for sequence classification tasks.
+    """

optimum/rbln/transformers/{modeling_alias.py → models/roberta/modeling_roberta.py}

@@ -12,42 +12,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..utils.logging import get_logger
-from .modeling_generic import (
-    RBLNModelForAudioClassification,
-    RBLNModelForImageClassification,
-    RBLNModelForMaskedLM,
-    RBLNModelForQuestionAnswering,
-    RBLNModelForSequenceClassification,
-)
+from ...modeling_generic import RBLNModelForMaskedLM, RBLNModelForSequenceClassification
 
 
-logger = get_logger()
-
-
-class RBLNASTForAudioClassification(RBLNModelForAudioClassification):
-    pass
-
-
-class RBLNDistilBertForQuestionAnswering(RBLNModelForQuestionAnswering):
-    rbln_model_input_names = ["input_ids", "attention_mask"]
-
-
-class RBLNResNetForImageClassification(RBLNModelForImageClassification):
-    pass
+class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
+    """
+    RBLN optimized RoBERTa model for masked language modeling tasks.
 
+    This class provides hardware-accelerated inference for RoBERTa models
+    on RBLN devices, supporting masked language modeling tasks such as
+    token prediction and text completion.
+    """
 
-class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
 class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
-    rbln_model_input_names = ["input_ids", "attention_mask"]
+    """
+    RBLN optimized RoBERTa model for sequence classification tasks.
 
+    This class provides hardware-accelerated inference for RoBERTa models
+    on RBLN devices, supporting text classification tasks such as sentiment analysis,
+    topic classification, and other sequence-level prediction tasks.
+    """
 
-class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
     rbln_model_input_names = ["input_ids", "attention_mask"]
-
-
-class RBLNViTForImageClassification(RBLNModelForImageClassification):
-    pass

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_seq2seq2 import RBLNModelForSeq2SeqLMConfig
+from .configuration_seq2seq import RBLNModelForSeq2SeqLMConfig
 from .modeling_seq2seq import RBLNModelForSeq2SeqLM

optimum/rbln/transformers/models/seq2seq/{configuration_seq2seq2.py → configuration_seq2seq.py}

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 import rebel
 
@@ -31,7 +31,7 @@ class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
         dec_max_seq_len: Optional[int] = None,
         use_attention_mask: Optional[bool] = None,
         pad_token_id: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

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

@@ -26,7 +26,7 @@ from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
 from ....utils.logging import get_logger
 from ....utils.runtime_utils import RBLNPytorchRuntime
-from .configuration_seq2seq2 import RBLNModelForSeq2SeqLMConfig
+from .configuration_seq2seq import RBLNModelForSeq2SeqLMConfig
 
 
 logger = get_logger(__name__)
@@ -161,16 +161,20 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             if "key_value_states" in name:
                 context.mark_static_address(tensor)
 
-        compiled_encoder = super().compile(
+        compiled_encoder = cls.compile(
             wrapped_model.encoder,
             enc_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
             example_inputs=enc_example_inputs,
             compile_context=context,
         )
 
-        compiled_decoder = super().compile(
+        compiled_decoder = cls.compile(
             wrapped_model.decoder,
             dec_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
             example_inputs=dec_example_inputs,
             compile_context=context,
         )

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

@@ -148,7 +148,8 @@ class Seq2SeqDecoderWrapper(nn.Module):
         new_layers = []
         for layer in model.get_decoder().layers:
             self_attn = Seq2SeqSelfAttention(layer.self_attn)
-            new_layers.append(Seq2SeqDecoderLayer(layer, self_attn))
+            cross_attn = Seq2SeqCrossAttention(layer.encoder_attn)
+            new_layers.append(Seq2SeqDecoderLayer(layer, self_attn, cross_attn))
 
         decoder_model = Seq2SeqDecoder(model.get_decoder(), new_layers)
         new_model = Seq2SeqForConditionalGeneration(model, decoder_model)
@@ -341,10 +342,11 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
         self_attn (Seq2SeqSelfAttention): Modified self-attention layer optimized for RBLN
     """
 
-    def __init__(self, decoder_layer, self_attn):
+    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__()
 
     def __post_init__(self, **kwargs):
@@ -402,7 +404,8 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
         # Cross-Attention Block
         residual = hidden_states
         hidden_states = self.pre_cross_attn_layer_norm(hidden_states)
-        cross_attn_output = self.encoder_attn(
+
+        cross_attn_output = self.cross_attn(
             hidden_states=hidden_states,
             past_key_value=cross_past_key_value,
             attention_mask=encoder_attention_mask,
@@ -487,3 +490,38 @@ class Seq2SeqSelfAttention(nn.Module):
         attn_output = self.out_proj(attn_output)
 
         return attn_output
+
+
+class Seq2SeqCrossAttention(nn.Module):
+    def __init__(self, attn, **kwargs):
+        super().__init__()
+        self._original_mod = attn
+        self.__post_init__(**kwargs)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: torch.Tensor = None,
+        past_key_value: Optional[object] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ):
+        bsz, tgt_len, _ = hidden_states.size()
+        query_states = self.q_proj(hidden_states).view(bsz, -1, self.num_heads, self.head_dim).transpose(1, 2)
+
+        is_cross_attention = key_value_states is not None
+        if is_cross_attention:
+            key_states = past_key_value[0]
+            value_states = past_key_value[1]
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(bsz, tgt_len, self.embed_dim)
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, None, past_key_value