optimum-rbln 0.7.4a3__py3-none-any.whl → 0.7.4a5__py3-none-any.whl

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

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .configuration_llava_next import RBLNLlavaNextForConditionalGenerationConfig
 from .modeling_llava_next import RBLNLlavaNextForConditionalGeneration

optimum/rbln/transformers/models/llava_next/configuration_llava_next.py

@@ -0,0 +1,46 @@
+# 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
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNLlavaNextForConditionalGenerationConfig(RBLNModelConfig):
+    submodules = ["vision_tower", "language_model"]
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        language_model: Optional[RBLNModelConfig] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+            language_model (Optional[RBLNModelConfig]): Configuration for the language model component.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
+        self.vision_tower = vision_tower
+        self.language_model = language_model

optimum/rbln/transformers/models/llava_next/modeling_llava_next.py

@@ -26,8 +26,8 @@ from transformers import (
 )
 from transformers.modeling_outputs import BaseModelOutputWithPooling
 
+from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.logging import get_logger
 from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyOutput
 
@@ -134,7 +134,7 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
         model: "LlavaNextForConditionalGeneration",
         save_dir_path: Path,
         subfolder: str,
-        rbln_config: RBLNConfig,
+        rbln_config: RBLNModelConfig,
     ):
         """
         If you are unavoidably running on a CPU rather than an RBLN device,
@@ -161,42 +161,31 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
         return self.language_model.get_input_embeddings()
 
     @classmethod
-    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNConfig):
+    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         return model.multi_modal_projector
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        model: Optional["PreTrainedModel"] = None,
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_kwargs={},
-    ) -> RBLNConfig:
-        vision_feature_select_strategy = rbln_kwargs.get("vision_feature_select_strategy", None)
-
-        # 1. Multi-modal projection layer
-        batch_size = rbln_kwargs.get("rbln_batch_size", None)
-        if batch_size is None:
-            batch_size = 1
-
+        rbln_config: Optional[RBLNModelConfig] = None,
+    ) -> RBLNModelConfig:
         feature_size = model_config.vision_config.hidden_size
 
-        # See forward function to see more details.
-        vision_feature_select_strategy = (
-            vision_feature_select_strategy
-            if vision_feature_select_strategy is not None
-            else model_config.vision_feature_select_strategy
-        )
-
         # Calculating `num_positions` : See CLIPVisionEmbeddings of transformers for more details.
         num_positions = (model_config.vision_config.image_size // model_config.vision_config.patch_size) ** 2 + 1
-        if vision_feature_select_strategy == "default":
+        if model_config.vision_feature_select_strategy == "default":
             selected_image_feature_dim = num_positions - 1
         else:
             selected_image_feature_dim = num_positions
 
-        input_info = [("image_features", [batch_size, selected_image_feature_dim, feature_size], "float32")]
+        input_info = [
+            ("image_features", [rbln_config.batch_size, selected_image_feature_dim, feature_size], "float32")
+        ]
         rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-        rbln_config = RBLNConfig(rbln_cls=cls.__name__, compile_cfgs=[rbln_compile_config], rbln_kwargs=rbln_kwargs)
+        rbln_config.set_compile_cfgs([rbln_compile_config])
         return rbln_config
 
     def prepare_inputs_for_generation(

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

@@ -20,4 +20,5 @@ this_path = os.path.abspath(__file__)
 local_dir = "/" + os.path.join(*this_path.split("/")[:-1]) + "/hf_hub_cached"
 environ["LOCAL_CACHE_ROOT_CUSTOM_CODE_MIDM"] = local_dir
 
+from .configuration_midm import RBLNMidmLMHeadModelConfig
 from .modeling_midm import RBLNMidmLMHeadModel

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

@@ -0,0 +1,19 @@
+# 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 RBLNMidmLMHeadModelConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    pass

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

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .configuration_mistral import RBLNMistralForCausalLMConfig
 from .modeling_mistral import RBLNMistralForCausalLM

optimum/rbln/transformers/models/mistral/configuration_mistral.py

@@ -0,0 +1,19 @@
+# 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 RBLNMistralForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    pass

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

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .configuration_phi import RBLNPhiForCausalLMConfig
 from .modeling_phi import RBLNPhiForCausalLM

optimum/rbln/transformers/models/phi/configuration_phi.py

@@ -0,0 +1,19 @@
+# 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 RBLNPhiForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    pass

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

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .configuration_qwen2 import RBLNQwen2ForCausalLMConfig
 from .modeling_qwen2 import RBLNQwen2ForCausalLM

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

@@ -0,0 +1,19 @@
+# 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 RBLNQwen2ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    pass

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

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

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

@@ -0,0 +1,66 @@
+# 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 rebel
+
+from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger()
+
+
+class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
+    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,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            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.
+                This is automatically set to True for RBLN-CA02 devices.
+            pad_token_id (Optional[int]): The ID of the padding token in the vocabulary.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
+        self.enc_max_seq_len = enc_max_seq_len
+        self.dec_max_seq_len = dec_max_seq_len
+
+        self.use_attention_mask = use_attention_mask
+        npu = self.npu or rebel.get_npu_name()
+        if npu == "RBLN-CA02":
+            if self.use_attention_mask is False:
+                logger.warning("Attention mask should be used with RBLN-CA02. Setting use_attention_mask to True.")
+            self.use_attention_mask = True
+        else:
+            self.use_attention_mask = self.use_attention_mask or False
+
+        self.pad_token_id = pad_token_id

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

@@ -22,10 +22,11 @@ from rebel.compile_context import CompileContext
 from transformers import AutoModelForSeq2SeqLM, PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
 
+from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.logging import get_logger
 from ....utils.runtime_utils import RBLNPytorchRuntime
+from .configuration_seq2seq2 import RBLNModelForSeq2SeqLMConfig
 
 
 logger = get_logger(__name__)
@@ -118,9 +119,9 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
     support_causal_attn = None
 
     def __post_init__(self, **kwargs):
-        batch_size = self.rbln_config.model_cfg["batch_size"]
-        dec_max_seq_len = self.rbln_config.model_cfg["dec_max_seq_len"]
-        self.use_attention_mask = self.rbln_config.model_cfg.get("use_attention_mask", None)
+        batch_size = self.rbln_config.batch_size
+        dec_max_seq_len = self.rbln_config.dec_max_seq_len
+        self.use_attention_mask = self.rbln_config.use_attention_mask
 
         self.encoder = RBLNRuntimeEncoder(
             runtime=self.model[0],
@@ -136,7 +137,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
 
     @classmethod
     @torch.inference_mode()
-    def get_compiled_model(cls, model: PreTrainedModel, rbln_config: RBLNConfig):
+    def get_compiled_model(cls, model: PreTrainedModel, rbln_config: RBLNModelForSeq2SeqLMConfig):
         wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
 
         enc_compile_config = rbln_config.compile_cfgs[0]
@@ -177,26 +178,15 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         return {"encoder": compiled_encoder, "decoder": compiled_decoder}
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
-        model_config: "PretrainedConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_enc_max_seq_len = rbln_kwargs.get("enc_max_seq_len", None)
-        rbln_dec_max_seq_len = rbln_kwargs.get("dec_max_seq_len", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-        rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
-
-        if cls.support_causal_attn:
-            rbln_use_attention_mask = rbln_kwargs.get("use_attention_mask", None)
-            if rbln_use_attention_mask is None:
-                rbln_use_attention_mask = False
-                rbln_npu = rbln_kwargs.get("npu", None) or rebel.get_npu_name()
-                if rbln_npu == "RBLN-CA02":
-                    rbln_use_attention_mask = True
-        else:
-            rbln_use_attention_mask = True
+        model: Optional["PreTrainedModel"] = None,
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_config: Optional[RBLNModelForSeq2SeqLMConfig] = None,
+    ) -> RBLNModelForSeq2SeqLMConfig:
+        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")
@@ -210,43 +200,44 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             model_config, "max_position_embeddings", None
         )
 
-        rbln_pad_token_id = getattr(model_config, "pad_token_id", None)
-        if rbln_pad_token_id is None:
-            rbln_pad_token_id = getattr(model_config, "bos_token_id", None)
-            if rbln_pad_token_id is None:
-                rbln_pad_token_id = getattr(model_config, "eos_token_id", None)
-                if rbln_pad_token_id is None:
-                    rbln_pad_token_id = -1
-
-        if rbln_enc_max_seq_len is None:
-            rbln_enc_max_seq_len = max_position_embeddings
-            if rbln_enc_max_seq_len is None:
-                for tokenizer in preprocessors:
-                    if hasattr(tokenizer, "model_max_length"):
-                        rbln_enc_max_seq_len = tokenizer.model_max_length
-                        break
-                if rbln_enc_max_seq_len is None:
-                    raise ValueError("`rbln_enc_max_seq_len` should be specified!")
-        if max_position_embeddings is not None and rbln_enc_max_seq_len > max_position_embeddings:
-            raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
-
-        if rbln_dec_max_seq_len is None:
-            rbln_dec_max_seq_len = max_position_embeddings
-            if rbln_dec_max_seq_len is None:
-                for tokenizer in preprocessors:
-                    if hasattr(tokenizer, "model_max_length"):
-                        rbln_dec_max_seq_len = tokenizer.model_max_length
-                        break
-                if rbln_dec_max_seq_len is None:
-                    raise ValueError("`rbln_dec_max_seq_len` should be specified!")
-
-        if max_position_embeddings is not None and rbln_dec_max_seq_len > max_position_embeddings:
-            raise ValueError("`rbln_dec_max_seq_len` should be less or equal than max_position_embeddings!")
+        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:
+                if hasattr(tokenizer, "model_max_length"):
+                    enc_max_seq_len = enc_max_seq_len or tokenizer.model_max_length
+                    break
+
+            if enc_max_seq_len is None:
+                raise ValueError("`enc_max_seq_len` should be specified!")
+            rbln_config.enc_max_seq_len = enc_max_seq_len
+
+        if max_position_embeddings is not None and rbln_config.enc_max_seq_len > max_position_embeddings:
+            raise ValueError("`enc_max_seq_len` should be less or equal than max_position_embeddings!")
+
+        if rbln_config.dec_max_seq_len is None:
+            dec_max_seq_len = max_position_embeddings
+            for tokenizer in preprocessors:
+                if hasattr(tokenizer, "model_max_length"):
+                    dec_max_seq_len = dec_max_seq_len or tokenizer.model_max_length
+                    break
+
+            if dec_max_seq_len is None:
+                raise ValueError("`dec_max_seq_len` should be specified!")
+            rbln_config.dec_max_seq_len = dec_max_seq_len
+
+        if max_position_embeddings is not None and rbln_config.dec_max_seq_len > max_position_embeddings:
+            raise ValueError("`dec_max_seq_len` should be less or equal than max_position_embeddings!")
 
         # model input info
         enc_input_info = [
-            ("input_ids", [1, rbln_enc_max_seq_len], "int64"),
-            ("attention_mask", [1, rbln_enc_max_seq_len], "float32"),
+            ("input_ids", [1, rbln_config.enc_max_seq_len], "int64"),
+            ("attention_mask", [1, rbln_config.enc_max_seq_len], "float32"),
             ("block_tables", [1], "int16"),
         ]
         enc_input_info.extend(
@@ -254,9 +245,9 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
                 (
                     f"cross_key_value_states_{i}",
                     [
-                        rbln_batch_size,
+                        rbln_config.batch_size,
                         n_head,
-                        rbln_enc_max_seq_len,
+                        rbln_config.enc_max_seq_len,
                         d_kv,
                     ],
                     "float32",
@@ -266,23 +257,23 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         )
 
         dec_input_info = [
-            ("input_ids", [rbln_batch_size, 1], "int64"),
-            ("encoder_attention_mask", [rbln_batch_size, rbln_enc_max_seq_len], "float32"),
+            ("input_ids", [rbln_config.batch_size, 1], "int64"),
+            ("encoder_attention_mask", [rbln_config.batch_size, rbln_config.enc_max_seq_len], "float32"),
             (
                 "cache_position",
-                [rbln_batch_size, 1],
+                [rbln_config.batch_size, 1],
                 "int32",
             ),
-            ("block_tables", [rbln_batch_size, 1], "int16"),
+            ("block_tables", [rbln_config.batch_size, 1], "int16"),
         ]
         dec_input_info.extend(
             [
                 (
                     f"cross_key_value_states_{i}",
                     [
-                        rbln_batch_size,
+                        rbln_config.batch_size,
                         n_head,
-                        rbln_enc_max_seq_len,
+                        rbln_config.enc_max_seq_len,
                         d_kv,
                     ],
                     "float32",
@@ -295,9 +286,9 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
                 (
                     f"self_key_value_states_{i}",
                     [
-                        rbln_batch_size,
+                        rbln_config.batch_size,
                         n_head,
-                        rbln_dec_max_seq_len,
+                        rbln_config.dec_max_seq_len,
                         d_kv,
                     ],
                     "float32",
@@ -306,46 +297,38 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             ]
         )
 
-        if rbln_use_attention_mask:
-            dec_input_info.insert(1, ("attention_mask", [rbln_batch_size, rbln_dec_max_seq_len], "float32"))
+        if rbln_config.use_attention_mask:
+            dec_input_info.insert(
+                1, ("attention_mask", [rbln_config.batch_size, rbln_config.dec_max_seq_len], "float32")
+            )
 
         enc_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=enc_input_info)
         dec_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
 
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[enc_compile_config, dec_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-
-        rbln_config.model_cfg.update(
-            {
-                "enc_max_seq_len": rbln_enc_max_seq_len,
-                "dec_max_seq_len": rbln_dec_max_seq_len,
-                "batch_size": rbln_batch_size,
-                "pad_token_id": rbln_pad_token_id,
-                "use_attention_mask": rbln_use_attention_mask,
-            }
-        )
-
+        rbln_config.set_compile_cfgs([enc_compile_config, dec_compile_config])
         return rbln_config
 
     @classmethod
     def _create_runtimes(
         cls,
         compiled_models: List[rebel.RBLNCompiledModel],
-        rbln_device_map: Dict[str, int],
-        activate_profiler: Optional[bool] = None,
+        rbln_config: RBLNModelForSeq2SeqLMConfig,
     ) -> List[rebel.Runtime]:
-        if any(model_name not in rbln_device_map for model_name in ["encoder", "decoder"]):
+        if any(model_name not in rbln_config.device_map for model_name in ["encoder", "decoder"]):
             cls._raise_missing_compiled_file_error(["encoder", "decoder"])
 
         return [
-            compiled_models[0].create_runtime(
-                tensor_type="pt", device=rbln_device_map["encoder"], activate_profiler=activate_profiler
+            rebel.Runtime(
+                compiled_models[0],
+                tensor_type="pt",
+                device=rbln_config.device_map["encoder"],
+                activate_profiler=rbln_config.activate_profiler,
             ),
-            compiled_models[1].create_runtime(
-                tensor_type="pt", device=rbln_device_map["decoder"], activate_profiler=activate_profiler
+            rebel.Runtime(
+                compiled_models[1],
+                tensor_type="pt",
+                device=rbln_config.device_map["decoder"],
+                activate_profiler=rbln_config.activate_profiler,
             ),
         ]
 
@@ -367,7 +350,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
     ):
         cur_seq_len = input_ids.shape[-1]
         cache_position = cur_seq_len - 1
-        max_seq_len = self.rbln_config.model_cfg["dec_max_seq_len"]
+        max_seq_len = self.rbln_config.dec_max_seq_len
         decoder_batch_size = input_ids.shape[0]
         input_ids = input_ids[:, cur_seq_len - 1 : cur_seq_len].contiguous()
         decoder_attention_mask = torch.zeros(decoder_batch_size, max_seq_len, dtype=torch.float32)
@@ -387,7 +370,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         **kwargs,
     ) -> Tuple[torch.FloatTensor]:
         # common decoder
-        cache_position = torch.full((self.rbln_config.model_cfg["batch_size"], 1), cache_position, dtype=torch.int32)
+        cache_position = torch.full((self.rbln_config.batch_size, 1), cache_position, dtype=torch.int32)
         logits = self.decoder(decoder_input_ids=decoder_input_ids, cache_position=cache_position, **kwargs).logits
 
         return Seq2SeqLMOutput(
@@ -421,11 +404,11 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         batch_size, input_len = inputs_tensor.shape
         inputs_tensor = torch.nn.functional.pad(
             inputs_tensor,
-            (0, self.rbln_config.model_cfg["enc_max_seq_len"] - input_len),
-            value=self.rbln_config.model_cfg["pad_token_id"],
+            (0, self.rbln_config.enc_max_seq_len - input_len),
+            value=self.rbln_config.pad_token_id,
         )
         model_kwargs["attention_mask"] = torch.nn.functional.pad(
-            model_kwargs["attention_mask"], (0, self.rbln_config.model_cfg["enc_max_seq_len"] - input_len)
+            model_kwargs["attention_mask"], (0, self.rbln_config.enc_max_seq_len - input_len)
         )
 
         # 3. make sure that encoder returns `ModelOutput`

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

@@ -13,4 +13,5 @@
 # limitations under the License.
 
 from ....ops import paged_add_softmax_attn_decode
+from .configuration_t5 import RBLNT5EncoderModelConfig, RBLNT5ForConditionalGenerationConfig
 from .modeling_t5 import RBLNT5EncoderModel, RBLNT5ForConditionalGeneration

optimum/rbln/transformers/models/t5/configuration_t5.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_generic import RBLNTransformerEncoderForFeatureExtractionConfig
+from ..seq2seq import RBLNModelForSeq2SeqLMConfig
+
+
+class RBLNT5EncoderModelConfig(RBLNTransformerEncoderForFeatureExtractionConfig):
+    pass
+
+
+class RBLNT5ForConditionalGenerationConfig(RBLNModelForSeq2SeqLMConfig):
+    pass