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

optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py

@@ -26,13 +26,15 @@ from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig, Pre
 from transformers.modeling_utils import no_init_weights
 from transformers.utils import ModelOutput
 
+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 ...utils.rbln_quantization import QuantizationManager
+from .configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
 from .decoderonly_architecture import (
     DecoderOnlyWrapper,
+    set_default_values,
     validate_attention_method,
 )
 
@@ -356,17 +358,9 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     _use_rotary_emb = True
 
     def __post_init__(self, **kwargs):
-        self.batch_size = self.rbln_config.model_cfg["batch_size"]
-        self.max_seq_len = self.rbln_config.model_cfg["max_seq_len"]
-        self.prefill_chunk_size = self.rbln_config.model_cfg["prefill_chunk_size"]
-        self.kvcache_block_size = self.rbln_config.model_cfg["kvcache_block_size"]
-        # FIXME get kvcache_num_blocks from compiled results.
-        self.kvcache_num_blocks = self.rbln_config.model_cfg["kvcache_num_blocks"]
-        self.use_attention_mask = self.rbln_config.model_cfg["use_attention_mask"]
-        attn_impl = self.rbln_config.model_cfg["attn_impl"]
         main_input_name = self.main_input_name
 
-        if self.rbln_config.model_cfg["use_inputs_embeds"]:
+        if self.rbln_config.use_inputs_embeds:
             main_input_name = "inputs_embeds"
             artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
             with no_init_weights():
@@ -380,40 +374,44 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             self.embed_tokens = None
 
         # Initialize shared resources to be used across Runtime instances (prefill and decode phases)
-        dec_attn_mask = torch.zeros(self.batch_size, 1, 1, self.max_seq_len, dtype=torch.float32)
+        dec_attn_mask = torch.zeros(
+            self.rbln_config.batch_size, 1, 1, self.rbln_config.max_seq_len, dtype=torch.float32
+        )
         block_tables = torch.zeros(
-            self.batch_size, self.max_seq_len // self.kvcache_block_size, dtype=torch.int16
+            self.rbln_config.batch_size,
+            self.rbln_config.max_seq_len // self.rbln_config.kvcache_block_size,
+            dtype=torch.int16,
         ).fill_(-1)
-        free_block_pool = deque(x for x in range(self.kvcache_num_blocks))
+        free_block_pool = deque(x for x in range(self.rbln_config.kvcache_num_blocks))
 
         self.prefill_decoder = RBLNRuntimeModel(
             runtime=self.model[0],
             main_input_name=main_input_name,
             embed_tokens=self.embed_tokens,
             phase="prefill",
-            batch_size=self.batch_size,
+            batch_size=self.rbln_config.batch_size,
             dec_attn_mask=dec_attn_mask,
             block_tables=block_tables,
             free_block_pool=free_block_pool,
-            kvcache_block_size=self.kvcache_block_size,
+            kvcache_block_size=self.rbln_config.kvcache_block_size,
             vocab_size=self.config.vocab_size,
-            prefill_chunk_size=self.prefill_chunk_size,
-            max_seq_len=self.max_seq_len,
-            use_attention_mask=self.use_attention_mask,
-            attn_impl=attn_impl,
+            prefill_chunk_size=self.rbln_config.prefill_chunk_size,
+            max_seq_len=self.rbln_config.max_seq_len,
+            use_attention_mask=self.rbln_config.use_attention_mask,
+            attn_impl=self.rbln_config.attn_impl,
         )
         self.decoder = RBLNRuntimeModel(
             runtime=self.model[1],
             main_input_name=main_input_name,
             embed_tokens=self.embed_tokens,
             phase="decode",
-            batch_size=self.batch_size,
+            batch_size=self.rbln_config.batch_size,
             dec_attn_mask=dec_attn_mask,
             block_tables=block_tables,
             free_block_pool=free_block_pool,
-            kvcache_block_size=self.kvcache_block_size,
-            use_attention_mask=self.use_attention_mask,
-            attn_impl=attn_impl,
+            kvcache_block_size=self.rbln_config.kvcache_block_size,
+            use_attention_mask=self.rbln_config.use_attention_mask,
+            attn_impl=self.rbln_config.attn_impl,
         )
 
     @classmethod
@@ -422,13 +420,13 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         model: "PreTrainedModel",
         save_dir_path: Path,
         subfolder: str,
-        rbln_config: RBLNConfig,
+        rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
     ):
         """
         If you are unavoidably running on a CPU rather than an RBLN device,
         store the torch tensor, weight, etc. in this function.
         """
-        if rbln_config.model_cfg["use_inputs_embeds"]:
+        if rbln_config.use_inputs_embeds:
             save_dict = {}
             save_dict["embed_tokens"] = model.get_input_embeddings().state_dict()
             torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
@@ -493,33 +491,35 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         return val
 
     @classmethod
-    def get_pytorch_model(cls, *args, **kwargs) -> "PreTrainedModel":
-        logger.debug("Loading the LLM model to the CPU.")  # TODO(jongho): Remove.
-
-        rbln_kwargs = kwargs.get("rbln_kwargs", {})
-        rbln_quantization = rbln_kwargs.get("quantization", None)
-        if rbln_quantization is not None and rbln_quantization["format"] == "rbln":
+    def get_pytorch_model(
+        cls, *args, rbln_config: Optional[RBLNDecoderOnlyModelForCausalLMConfig] = None, **kwargs
+    ) -> "PreTrainedModel":
+        if (
+            rbln_config is not None
+            and "format" in rbln_config.quantization
+            and rbln_config.quantization["format"] == "rbln"
+        ):
             model = cls.get_quantized_model(*args, **kwargs)
         else:
             model = super().get_pytorch_model(*args, **kwargs)
 
-        logger.debug("Loaded the LLM model to the CPU.")
         return model
 
     @classmethod
-    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
-        wrapper_cfg = {"max_seq_len": rbln_config.model_cfg["max_seq_len"]}
-        wrapper_cfg["attn_impl"] = rbln_config.model_cfg.get("attn_impl")
-        wrapper_cfg["kvcache_partition_len"] = rbln_config.model_cfg.get("kvcache_partition_len")
-        wrapper_cfg["kvcache_block_size"] = rbln_config.model_cfg.get("kvcache_block_size")
-        wrapper_cfg["use_rotary_emb"] = cls._use_rotary_emb
-        wrapper_cfg["use_attention_mask"] = rbln_config.model_cfg.get("use_attention_mask")
-
+    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: "RBLNDecoderOnlyModelForCausalLMConfig"):
+        wrapper_cfg = {
+            "max_seq_len": rbln_config.max_seq_len,
+            "attn_impl": rbln_config.attn_impl,
+            "kvcache_partition_len": rbln_config.kvcache_partition_len,
+            "kvcache_block_size": rbln_config.kvcache_block_size,
+            "use_rotary_emb": cls._use_rotary_emb,
+            "use_attention_mask": rbln_config.use_attention_mask,
+        }
         return cls._decoder_wrapper_cls(model, **wrapper_cfg).eval()
 
     @classmethod
     @torch.inference_mode()
-    def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNConfig):
+    def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
         wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
 
         rbln_compile_configs = rbln_config.compile_cfgs
@@ -541,8 +541,6 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
         dec_example_inputs = dec_compile_config.get_dummy_inputs(fill=0, static_tensors=static_tensors)
 
-        quantize_config = rbln_config.model_cfg.get("quantization", None)
-
         @QuantizationManager.with_quantization_env
         def compile_model(*args, **kwargs):
             try:
@@ -567,7 +565,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             finally:
                 torch.nn.functional.linear = original_linear
 
-        return compile_model(quantize_config=quantize_config)
+        return compile_model(quantize_config=rbln_config.quantization)
 
     @classmethod
     def get_maximum_num_blocks(
@@ -654,78 +652,57 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         return max_n_blocks
 
     @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_max_seq_len = rbln_kwargs.get("max_seq_len", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-        rbln_use_inputs_embeds = rbln_kwargs.get("use_inputs_embeds", None)
-        rbln_use_attention_mask = rbln_kwargs.get("use_attention_mask", None)
-        rbln_attn_impl = rbln_kwargs.get("attn_impl", None)
-        rbln_kvcache_partition_len = rbln_kwargs.get("kvcache_partition_len", None)
-        rbln_kvcache_block_size = rbln_kwargs.get("kvcache_block_size", None)
-        rbln_quantization = QuantizationManager.validate_quantization_config(rbln_kwargs.get("quantization", None))
-        rbln_prefill_chunk_size = rbln_kwargs.get("prefill_chunk_size", 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
-
-        if rbln_prefill_chunk_size is None:
-            rbln_prefill_chunk_size = 128
-        elif rbln_prefill_chunk_size % 64 != 0 or rbln_prefill_chunk_size == 0:
-            raise ValueError(
-                f"Invalid rbln_prefill_chunk_size: {rbln_prefill_chunk_size}. It must be a nonzero multiple of 64."
-            )
-
-        if rbln_max_seq_len is None:
-            rbln_max_seq_len = getattr(model_config, "max_position_embeddings", None) or getattr(
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
+        model: Optional["PreTrainedModel"] = None,
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_config: Optional[RBLNDecoderOnlyModelForCausalLMConfig] = None,
+    ) -> RBLNDecoderOnlyModelForCausalLMConfig:
+        if rbln_config.max_seq_len is None:
+            rbln_config.max_seq_len = getattr(model_config, "max_position_embeddings", None) or getattr(
                 model_config, "n_positions", None
             )
-        if rbln_max_seq_len is None:
-            raise ValueError("`rbln_max_seq_len` should be specified.")
-
-        rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
-        rbln_use_inputs_embeds = False if rbln_use_inputs_embeds is None else rbln_use_inputs_embeds
+        if rbln_config.max_seq_len is None:
+            raise ValueError("`max_seq_len` should be specified.")
+
+        rbln_config.attn_impl, rbln_config.kvcache_partition_len, rbln_config.kvcache_block_size = set_default_values(
+            attn_impl=rbln_config.attn_impl,
+            kvcache_partition_len=rbln_config.kvcache_partition_len,
+            kvcache_block_size=rbln_config.kvcache_block_size,
+            max_seq_len=rbln_config.max_seq_len,
+        )
 
-        rbln_attn_impl, rbln_kvcache_partition_len, rbln_kvcache_block_size = validate_attention_method(
-            rbln_attn_impl=rbln_attn_impl,
-            rbln_kvcache_partition_len=rbln_kvcache_partition_len,
-            rbln_kvcache_block_size=rbln_kvcache_block_size,
-            rbln_max_seq_len=rbln_max_seq_len,
+        validate_attention_method(
+            attn_impl=rbln_config.attn_impl,
+            kvcache_partition_len=rbln_config.kvcache_partition_len,
+            kvcache_block_size=rbln_config.kvcache_block_size,
+            max_seq_len=rbln_config.max_seq_len,
         )
 
-        if rbln_kvcache_block_size is None:
-            if rbln_attn_impl == "eager":
-                rbln_kvcache_block_size = rbln_max_seq_len
-            else:
-                rbln_kvcache_block_size = rbln_kvcache_partition_len
+        rbln_config.kvcache_num_blocks = (
+            rbln_config.max_seq_len // rbln_config.kvcache_block_size
+        ) * rbln_config.batch_size
 
-        rbln_kvcache_num_blocks = (rbln_max_seq_len // rbln_kvcache_block_size) * rbln_batch_size
-        if rbln_attn_impl == "flash_attn":
+        if rbln_config.attn_impl == "flash_attn":
             max_num_blocks = cls.get_maximum_num_blocks(
                 config=model_config,
-                tensor_parallel_size=rbln_kwargs.get("tensor_parallel_size", 1),
-                kvcache_block_size=rbln_kvcache_block_size,
-                nbits_per_param=16 if rbln_quantization is None else 4,  # TODO(jongho): FIX Ad-hoc
-                n_model_params=rbln_kwargs["n_model_params"],
+                tensor_parallel_size=rbln_config.tensor_parallel_size or 1,
+                kvcache_block_size=rbln_config.kvcache_block_size,
+                nbits_per_param=16 if not rbln_config.quantization else 4,  # TODO(jongho): FIX Ad-hoc
+                n_model_params=sum(p.numel() for p in model.parameters()),
             )
-            rbln_kvcache_num_blocks = min(rbln_kvcache_num_blocks, max_num_blocks)
+            rbln_config.kvcache_num_blocks = min(rbln_config.kvcache_num_blocks, max_num_blocks)
 
-            required_blocks = rbln_max_seq_len // rbln_kvcache_block_size + 1
-            if rbln_kvcache_num_blocks < required_blocks:
-                rbln_kvcache_num_blocks = required_blocks
+            required_blocks = rbln_config.max_seq_len // rbln_config.kvcache_block_size + 1
+            if rbln_config.kvcache_num_blocks < required_blocks:
+                rbln_config.kvcache_num_blocks = required_blocks
 
-            logger.info(f"[KVCache] Compiling with num_blocks: {rbln_kvcache_num_blocks}")
+            logger.info(f"[KVCache] Compiling with num_blocks: {rbln_config.kvcache_num_blocks}")
 
-            if rbln_kvcache_num_blocks < rbln_batch_size:
+            if rbln_config.kvcache_num_blocks < rbln_config.batch_size:
                 raise RuntimeError(
-                    f"Batch size ({rbln_batch_size}) exceeds available KV cache blocks ({rbln_kvcache_num_blocks}). "
+                    f"Batch size ({rbln_config.batch_size}) exceeds available KV cache blocks ({rbln_config.kvcache_num_blocks}). "
                     "Ensure the number of blocks is at least equal to the batch size."
                 )
 
@@ -740,6 +717,10 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             query_length,
             use_inputs_embeds,
             hidden_size,
+            use_attention_mask,
+            max_seq_len,
+            kvcache_block_size,
+            kvcache_num_blocks,
         ):
             if use_inputs_embeds:
                 main_input = ("inputs_embeds", [batch_size, query_length, hidden_size], "float32")
@@ -755,10 +736,10 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
                 ),
             ]
 
-            if rbln_use_attention_mask:
+            if use_attention_mask:
                 input_info.extend(
                     [
-                        ("attention_mask", [batch_size, 1, query_length, rbln_max_seq_len], "float32"),
+                        ("attention_mask", [batch_size, 1, query_length, max_seq_len], "float32"),
                     ]
                 )
 
@@ -769,7 +750,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
                     ]
                 )
 
-            max_block_cnt = rbln_max_seq_len // rbln_kvcache_block_size
+            max_block_cnt = max_seq_len // kvcache_block_size
 
             if query_length > 1:
                 input_info.extend([("block_tables", [max_block_cnt], "int16")])
@@ -781,9 +762,9 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
                     (
                         f"past_key_values_{i}",
                         [
-                            rbln_kvcache_num_blocks,
+                            kvcache_num_blocks,
                             num_key_value_heads,
-                            rbln_kvcache_block_size,
+                            kvcache_block_size,
                             head_dim,
                         ],
                         "float32",
@@ -796,42 +777,29 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
         prefill_input_info = get_input_info(
             batch_size=1,
-            query_length=rbln_prefill_chunk_size,
-            use_inputs_embeds=rbln_use_inputs_embeds,
+            query_length=rbln_config.prefill_chunk_size,
+            use_inputs_embeds=rbln_config.use_inputs_embeds,
             hidden_size=hidden_size,
+            use_attention_mask=rbln_config.use_attention_mask,
+            max_seq_len=rbln_config.max_seq_len,
+            kvcache_block_size=rbln_config.kvcache_block_size,
+            kvcache_num_blocks=rbln_config.kvcache_num_blocks,
         )
         dec_input_info = get_input_info(
-            batch_size=rbln_batch_size,
+            batch_size=rbln_config.batch_size,
             query_length=1,
-            use_inputs_embeds=rbln_use_inputs_embeds,
+            use_inputs_embeds=rbln_config.use_inputs_embeds,
             hidden_size=hidden_size,
+            use_attention_mask=rbln_config.use_attention_mask,
+            max_seq_len=rbln_config.max_seq_len,
+            kvcache_block_size=rbln_config.kvcache_block_size,
+            kvcache_num_blocks=rbln_config.kvcache_num_blocks,
         )
 
         prefill_compile_config = RBLNCompileConfig(compiled_model_name="prefill", input_info=prefill_input_info)
         dec_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
 
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[prefill_compile_config, dec_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-
-        rbln_config.model_cfg.update(
-            {
-                "max_seq_len": rbln_max_seq_len,
-                "batch_size": rbln_batch_size,
-                "prefill_chunk_size": rbln_prefill_chunk_size,
-                "use_attention_mask": rbln_use_attention_mask,
-                "use_inputs_embeds": rbln_use_inputs_embeds,
-                "kvcache_partition_len": rbln_kvcache_partition_len,
-                "kvcache_block_size": rbln_kvcache_block_size,
-                "attn_impl": rbln_attn_impl,
-                "kvcache_num_blocks": rbln_kvcache_num_blocks,
-            }
-        )
-
-        if rbln_quantization is not None:
-            rbln_config.model_cfg.update({"quantization": rbln_quantization})
+        rbln_config.set_compile_cfgs([prefill_compile_config, dec_compile_config])
 
         return rbln_config
 
@@ -839,18 +807,23 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     def _create_runtimes(
         cls,
         compiled_models: List[rebel.RBLNCompiledModel],
-        rbln_device_map: Dict[str, int],
-        activate_profiler: Optional[bool] = None,
+        rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
     ) -> List[rebel.Runtime]:
-        if any(model_name not in rbln_device_map for model_name in ["prefill", "decoder"]):
+        if any(model_name not in rbln_config.device_map for model_name in ["prefill", "decoder"]):
             cls._raise_missing_compiled_file_error(["prefill", "decoder"])
 
         return [
-            compiled_models[0].create_runtime(
-                tensor_type="pt", device=rbln_device_map["prefill"], activate_profiler=activate_profiler
+            rebel.Runtime(
+                compiled_models[0],
+                tensor_type="pt",
+                device=rbln_config.device_map["prefill"],
+                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,
             ),
         ]
 
@@ -887,7 +860,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             model_inputs.update({"input_ids": input_ids})
 
         if inputs_embeds is not None:
-            if self.rbln_config.model_cfg["use_inputs_embeds"]:
+            if self.rbln_config.use_inputs_embeds:
                 model_inputs.update({"inputs_embeds": inputs_embeds})
             else:
                 raise ValueError(

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

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .configuration_dpt import RBLNDPTForDepthEstimationConfig
 from .modeling_dpt import RBLNDPTForDepthEstimation

optimum/rbln/transformers/models/dpt/configuration_dpt.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 ...configuration_generic import RBLNModelForDepthEstimationConfig
+
+
+class RBLNDPTForDepthEstimationConfig(RBLNModelForDepthEstimationConfig):
+    pass

optimum/rbln/transformers/models/dpt/modeling_dpt.py

@@ -12,82 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Union
 
-from transformers import AutoModelForDepthEstimation
-from transformers.modeling_outputs import DepthEstimatorOutput
+from ...modeling_generic import RBLNModelForDepthEstimation
 
-from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
-from ....utils.logging import get_logger
 
-
-logger = get_logger(__name__)
-
-if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
-
-
-class RBLNDPTForDepthEstimation(RBLNModel):
-    auto_model_class = AutoModelForDepthEstimation
-    main_input_name = "pixel_values"
-
-    @classmethod
-    def _get_rbln_config(
-        cls,
-        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
-        model_config: Optional["PretrainedConfig"] = None,
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_image_size = rbln_kwargs.get("image_size", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        if rbln_image_size is None:
-            for processor in preprocessors:
-                image_size = getattr(processor, "size", None)
-
-                if image_size is not None:
-                    if isinstance(image_size, Iterable):
-                        if "shortest_edge" in image_size:
-                            rbln_image_size = image_size["shortest_edge"]
-                            break
-                        elif "height" in image_size and "width" in image_size:
-                            rbln_image_size = image_size["height"], image_size["width"]
-                            break
-                    else:
-                        rbln_image_size = image_size
-
-            if rbln_image_size is None:
-                rbln_image_size = getattr(model_config, "image_size", None)
-
-            if rbln_image_size is None:
-                raise ValueError("`rbln_image_size` should be specified!")
-
-        if isinstance(rbln_image_size, int):
-            rbln_image_size = rbln_image_size, rbln_image_size
-
-        input_info = [("pixel_values", [rbln_batch_size, 3, rbln_image_size[0], rbln_image_size[1]], "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.model_cfg.update(
-            {
-                "image_size": rbln_image_size,
-                "batch_size": rbln_batch_size,
-            }
-        )
-
-        return rbln_config
-
-    def forward(self, *args, **kwargs):
-        predicted_depth = super().forward(*args, **kwargs)
-        return DepthEstimatorOutput(predicted_depth=predicted_depth)
+class RBLNDPTForDepthEstimation(RBLNModelForDepthEstimation):
+    pass

optimum/rbln/transformers/models/exaone/__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_exaone import RBLNExaoneForCausalLMConfig
 from .modeling_exaone import RBLNExaoneForCausalLM

optimum/rbln/transformers/models/exaone/configuration_exaone.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 RBLNExaoneForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    pass

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

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .configuration_gemma import RBLNGemmaForCausalLMConfig
 from .modeling_gemma import RBLNGemmaForCausalLM

optimum/rbln/transformers/models/gemma/configuration_gemma.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 RBLNGemmaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    pass

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

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .configuration_gpt2 import RBLNGPT2LMHeadModelConfig
 from .modeling_gpt2 import RBLNGPT2LMHeadModel

optimum/rbln/transformers/models/gpt2/configuration_gpt2.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 RBLNGPT2LMHeadModelConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    pass

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

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .configuration_llama import RBLNLlamaForCausalLMConfig
 from .modeling_llama import RBLNLlamaForCausalLM

optimum/rbln/transformers/models/llama/configuration_llama.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 RBLNLlamaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    pass