optimum-rbln 0.8.2a4__py3-none-any.whl → 0.8.2a6__py3-none-any.whl

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

@@ -12,5 +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
+from .configuration_gemma import RBLNGemmaForCausalLMConfig, RBLNGemmaModelConfig
+from .modeling_gemma import RBLNGemmaForCausalLM, RBLNGemmaModel

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNGemmaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -40,3 +40,11 @@ class RBLNGemmaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     )
     ```
     """
+
+
+class RBLNGemmaModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN Gemma models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+    """

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

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .gemma_architecture import GemmaWrapper
 
 
@@ -81,3 +81,15 @@ class RBLNGemmaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
 
     _decoder_wrapper_cls = GemmaWrapper
+
+
+class RBLNGemmaModel(RBLNDecoderOnlyModel):
+    """
+    The Gemma Model transformer without a language modeling head.
+    This model inherits from [`RBLNDecoderOnlyModel`]. 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 GemmaModel model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers GemmaModel model into a RBLN transformer model by:
+    """
+
+    _decoder_wrapper_cls = GemmaWrapper

optimum/rbln/transformers/models/gemma3/configuration_gemma3.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 from typing import Any, Dict, Optional
 
-import rebel
-
 from ....configuration_utils import RBLNModelConfig
 from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
 from ..siglip.configuration_siglip import RBLNSiglipVisionModelConfig
@@ -39,9 +37,13 @@ class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
         )
         self.image_prefill_chunk_size = image_prefill_chunk_size
 
-        npu = self.npu or rebel.get_npu_name()
-        if npu == "RBLN-CA02":
-            raise NotImplementedError("Gemma3 is currently not supported on RBLN-CA02")
+    @property
+    def use_image_prefill(self):
+        return self.image_prefill_chunk_size is not None
+
+    @property
+    def decoder_runtime_idx(self):
+        return 2 if self.use_image_prefill else 1
 
 
 class RBLNGemma3ForConditionalGenerationConfig(RBLNModelConfig):

optimum/rbln/transformers/models/gemma3/modeling_gemma3.py

@@ -31,7 +31,11 @@ from transformers.models.gemma3.modeling_gemma3 import Gemma3TextScaledWordEmbed
 
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
-from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyOutput, RBLNRuntimeModel
+from ..decoderonly.modeling_decoderonly import (
+    RBLNDecoderOnlyForCausalLMOutput,
+    RBLNDecoderOnlyModelForCausalLM,
+    RBLNRuntimeModel,
+)
 from .configuration_gemma3 import RBLNGemma3ForCausalLMConfig
 from .gemma3_architecture import Gemma3ForCausalLMWrapper
 
@@ -41,7 +45,7 @@ if TYPE_CHECKING:
 
 
 @dataclass
-class RBLNGemma3ForCausalLMOutput(RBLNDecoderOnlyOutput):
+class RBLNGemma3ForCausalLMOutput(RBLNDecoderOnlyForCausalLMOutput):
     attention_mask: Optional[torch.Tensor] = None
 
 
@@ -197,7 +201,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
 
     def _update_model_kwargs_for_generation(
         self,
-        outputs: RBLNDecoderOnlyOutput,
+        outputs: RBLNDecoderOnlyForCausalLMOutput,
         model_kwargs: Dict[str, Any],
         **kwargs,
     ) -> Dict[str, Any]:
@@ -266,7 +270,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
         position_ids: Optional[torch.Tensor] = None,
         token_type_ids: Optional[torch.Tensor] = None,
         **lm_kwargs: Dict[str, Any],
-    ) -> Union[Tuple, RBLNDecoderOnlyOutput]:
+    ) -> Union[Tuple, RBLNDecoderOnlyForCausalLMOutput]:
         # prefill
         if cache_position is None:
             logits = []
@@ -304,7 +308,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
                 position_ids=position_ids if self.rbln_config.language_model.use_position_ids else None,
             ).logits
 
-        return RBLNDecoderOnlyOutput(
+        return RBLNDecoderOnlyForCausalLMOutput(
             logits=logits, generate_idx=generate_idx, padded_cache_lengths=padded_cache_lengths
         )
 
@@ -333,11 +337,12 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
         chunked_attention_mask = torch.zeros(1, chunked_attention_mask.shape[-1], dtype=torch.float32)
 
         # as gemma3 has different prefill chunk size for image and text, we need to pad the inputs to the max of the two.
-        padding_size = max(self.rbln_config.prefill_chunk_size, self.rbln_config.image_prefill_chunk_size)
-        inputs = torch.nn.functional.pad(inputs, (0, 0, 0, padding_size))
-        cache_position = torch.nn.functional.pad(cache_position, (0, padding_size))
-        position_ids = torch.nn.functional.pad(position_ids, (0, padding_size))
-        token_type_ids = torch.nn.functional.pad(token_type_ids, (0, padding_size), value=-1)
+        if self.rbln_config.use_image_prefill:
+            padding_size = max(self.rbln_config.prefill_chunk_size, self.rbln_config.image_prefill_chunk_size)
+            inputs = torch.nn.functional.pad(inputs, (0, 0, 0, padding_size))
+            cache_position = torch.nn.functional.pad(cache_position, (0, padding_size))
+            position_ids = torch.nn.functional.pad(position_ids, (0, padding_size))
+            token_type_ids = torch.nn.functional.pad(token_type_ids, (0, padding_size), value=-1)
 
         return (
             inputs,
@@ -385,7 +390,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
         step = 0
         while step < query_length:
             # Check if the prefill chunk is an image prefill
-            is_image_prefill = torch.all(
+            is_image_prefill = self.rbln_config.use_image_prefill and torch.all(
                 token_type_ids[:, step : step + self.rbln_config.image_prefill_chunk_size] == 1
             )
             prefill_chunk_size = (
@@ -393,8 +398,10 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
             )
 
             # Check if the prefill chunk is a text prefill which have image_tokens in it.
-            is_text_prefill_with_image_tokens = not is_image_prefill and torch.any(
-                token_type_ids[:, step : step + prefill_chunk_size] == 1
+            is_text_prefill_with_image_tokens = (
+                self.rbln_config.use_image_prefill
+                and not is_image_prefill
+                and torch.any(token_type_ids[:, step : step + prefill_chunk_size] == 1)
             )
 
             # Check if the prefill chunk crosses a block boundary, requiring padding to align with block boundaries
@@ -414,7 +421,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
             num_processed_tokens = prefill_chunk_size
             if is_text_prefill_with_image_tokens:
                 first_image_token_idx = torch.where(token_type_ids[:, step : step + prefill_chunk_size] == 1)[1][0]
-                num_processed_tokens = first_image_token_idx
+                num_processed_tokens = first_image_token_idx.item()
             if is_last_chunk:
                 num_processed_tokens = query_length - step
 
@@ -509,7 +516,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
 
         logits = self.decode(inputs, cache_position, block_tables, local_block_tables, attention_mask, position_ids)
 
-        return RBLNDecoderOnlyOutput(logits=logits)
+        return RBLNDecoderOnlyForCausalLMOutput(logits=logits)
 
 
 class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
@@ -544,9 +551,10 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
             dtype=torch.int16,
         ).fill_(-1)
         free_block_pool = deque(x for x in range(self.rbln_config.kvcache_num_blocks))
+
         self.prefill_decoder = RBLNGemma3RuntimeModel(
             runtime=self.model[0],
-            image_prefill=self.model[1],
+            image_prefill=self.model[1] if self.rbln_config.use_image_prefill else None,
             main_input_name=main_input_name,
             embed_tokens=self.embed_tokens,
             phase="prefill",
@@ -561,7 +569,7 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         self.decoders = {}
         for i, batch_size in enumerate(self.rbln_config.decoder_batch_sizes):
             self.decoders[batch_size] = RBLNGemma3RuntimeModel(
-                runtime=self.model[i + 2],
+                runtime=self.model[i + self.rbln_config.decoder_runtime_idx],
                 main_input_name=main_input_name,
                 embed_tokens=self.embed_tokens,
                 phase="decode",
@@ -624,20 +632,21 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         if not (rbln_config.use_attention_mask and rbln_config.use_position_ids):
             raise ValueError("use_attention_mask and use_position_ids must be True for RBLNGemma3ForCausalLM")
 
-        # Update image prefill compile config
-        img_prefill_input_info = cls.get_input_info(
-            batch_size=1,
-            query_length=rbln_config.image_prefill_chunk_size,
-            rbln_config=rbln_config,
-            model_config=model_config,
-        )
-        image_prefill_compile_config = RBLNCompileConfig(
-            compiled_model_name="image_prefill", input_info=img_prefill_input_info
-        )
-        # Insert image_prefill compile config at index 1
-        compile_cfgs = rbln_config.compile_cfgs
-        compile_cfgs.insert(1, image_prefill_compile_config)
-        rbln_config.set_compile_cfgs(compile_cfgs)
+        if rbln_config.use_image_prefill:
+            # Update image prefill compile config
+            img_prefill_input_info = cls.get_input_info(
+                batch_size=1,
+                query_length=rbln_config.image_prefill_chunk_size,
+                rbln_config=rbln_config,
+                model_config=model_config,
+            )
+            image_prefill_compile_config = RBLNCompileConfig(
+                compiled_model_name="image_prefill", input_info=img_prefill_input_info
+            )
+            # Insert image_prefill compile config at index 1
+            compile_cfgs = rbln_config.compile_cfgs
+            compile_cfgs.insert(1, image_prefill_compile_config)
+            rbln_config.set_compile_cfgs(compile_cfgs)
 
         return rbln_config
 
@@ -690,23 +699,27 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
             context,
             rbln_config.quantization,
         )
+        compiled_models = {"prefill": compiled_prefill}
 
-        image_prefill_compile_config = rbln_compile_configs[1]
-        image_prefill_example_inputs = image_prefill_compile_config.get_dummy_inputs(
-            fill=0, static_tensors=static_tensors
-        )
-        wrapped_model.phase = "image_prefill"
-        compiled_image_prefill = compile_model(
-            wrapped_model,
-            image_prefill_compile_config,
-            image_prefill_example_inputs,
-            context,
-            rbln_config.quantization,
-        )
+        if rbln_config.use_image_prefill:
+            image_prefill_compile_config = rbln_compile_configs[1]
+            image_prefill_example_inputs = image_prefill_compile_config.get_dummy_inputs(
+                fill=0, static_tensors=static_tensors
+            )
+            wrapped_model.phase = "image_prefill"
+            compiled_image_prefill = compile_model(
+                wrapped_model,
+                image_prefill_compile_config,
+                image_prefill_example_inputs,
+                context,
+                rbln_config.quantization,
+            )
+            compiled_models["image_prefill"] = compiled_image_prefill
 
-        compiled_models = {"prefill": compiled_prefill, "image_prefill": compiled_image_prefill}
         wrapped_model.phase = "decode"
-        for batch_size, dec_compile_config in zip(rbln_config.decoder_batch_sizes, rbln_compile_configs[2:]):
+        for batch_size, dec_compile_config in zip(
+            rbln_config.decoder_batch_sizes, rbln_compile_configs[rbln_config.decoder_runtime_idx :]
+        ):
             dec_example_inputs = dec_compile_config.get_dummy_inputs(fill=0, static_tensors=static_tensors)
             compiled_decoder = compile_model(
                 wrapped_model,
@@ -727,35 +740,45 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     ) -> List[rebel.Runtime]:
         expected_model_names = [
             "prefill",
-            "image_prefill",
             *[f"decoder_batch_{batch_size}" for batch_size in rbln_config.decoder_batch_sizes],
         ]
+        if rbln_config.use_image_prefill:
+            expected_model_names.insert(1, "image_prefill")
+
         if any(model_name not in rbln_config.device_map for model_name in expected_model_names):
             cls._raise_missing_compiled_file_error(expected_model_names)
 
-        return [
+        ret_val = [
             rebel.Runtime(
                 compiled_models[0],
                 tensor_type="pt",
                 device=rbln_config.device_map["prefill"],
                 activate_profiler=rbln_config.activate_profiler,
                 timeout=rbln_config.timeout,
-            ),
-            rebel.Runtime(
-                compiled_models[1],
-                tensor_type="pt",
-                device=rbln_config.device_map["image_prefill"],
-                activate_profiler=rbln_config.activate_profiler,
-                timeout=rbln_config.timeout,
-            ),
-            *[
+            )
+        ]
+        if rbln_config.use_image_prefill:
+            ret_val.append(
+                rebel.Runtime(
+                    compiled_models[1],
+                    tensor_type="pt",
+                    device=rbln_config.device_map["image_prefill"],
+                    activate_profiler=rbln_config.activate_profiler,
+                    timeout=rbln_config.timeout,
+                ),
+            )
+
+        ret_val.extend(
+            [
                 rebel.Runtime(
-                    compiled_models[i + 2],
+                    compiled_models[i + rbln_config.decoder_runtime_idx],
                     tensor_type="pt",
                     device=rbln_config.device_map[f"decoder_batch_{batch_size}"],
                     activate_profiler=rbln_config.activate_profiler,
                     timeout=rbln_config.timeout,
                 )
                 for i, batch_size in enumerate(rbln_config.decoder_batch_sizes)
-            ],
-        ]
+            ]
+        )
+
+        return ret_val

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

@@ -12,5 +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
+from .configuration_gpt2 import RBLNGPT2LMHeadModelConfig, RBLNGPT2ModelConfig
+from .modeling_gpt2 import RBLNGPT2LMHeadModel, RBLNGPT2Model

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

@@ -12,11 +12,39 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNGPT2LMHeadModelConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     """
-    Configuration class for GPT-2 causal language model.
-    Inherits from RBLNDecoderOnlyModelForCausalLMConfig with no additional parameters.
+    Configuration class for RBLN GPT2 models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+    """
+
+
+class RBLNGPT2ModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN GPT2 models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNGPT2Model, RBLNGPT2ModelConfig
+
+    # Create a configuration object
+    config = RBLNGPT2ModelConfig(
+        batch_size=1,
+        max_seq_len=1024,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNGPT2Model.from_pretrained(
+        "openai/gpt2",
+        export=True,
+        rbln_config=config
+    )
+    ```
     """

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

@@ -13,11 +13,10 @@
 # limitations under the License.
 
 import math
-from typing import TYPE_CHECKING, Tuple
+from typing import TYPE_CHECKING, Tuple, Union
 
 import torch
 import torch.nn as nn
-from transformers import PreTrainedModel
 
 from ..decoderonly.decoderonly_architecture import (
     DecoderOnlyAttention,
@@ -28,7 +27,7 @@ from ..decoderonly.decoderonly_architecture import (
 
 
 if TYPE_CHECKING:
-    from transformers import GPT2LMHeadModel
+    from transformers import GPT2LMHeadModel, GPT2Model
 
 
 class GPT2Wrapper(DecoderOnlyWrapper):
@@ -44,11 +43,11 @@ class GPT2Wrapper(DecoderOnlyWrapper):
     def get_attn_layer(self, layer: nn.Module):
         return layer.attn
 
-    def get_model_layer(self, causal_lm: "GPT2LMHeadModel"):
-        return causal_lm.transformer
+    def get_model_layer(self, model: Union["GPT2LMHeadModel", "GPT2Model"]):
+        return model.transformer if self.is_causal_lm else model
 
-    def get_decoder_layers(self, causal_lm: PreTrainedModel):
-        return causal_lm.transformer.h
+    def get_decoder_layers(self, model: Union["GPT2LMHeadModel", "GPT2Model"]):
+        return model.transformer.h if self.is_causal_lm else model.h
 
 
 class GPT2Model(DecoderOnlyModel):

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

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .gpt2_architecture import GPT2Wrapper
 
 
@@ -36,3 +36,18 @@ class RBLNGPT2LMHeadModel(RBLNDecoderOnlyModelForCausalLM):
 
     _decoder_wrapper_cls = GPT2Wrapper
     _use_rotary_emb = False
+
+
+class RBLNGPT2Model(RBLNDecoderOnlyModel):
+    """
+    The GPT2 Model transformer without a language modeling head.
+
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model.
+
+    A class to convert and run pre-trained transformers based GPT2Model model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers GPT2Model model into a RBLN transformer model by:
+    """
+
+    _decoder_wrapper_cls = GPT2Wrapper
+    _use_rotary_emb = False

optimum/rbln/transformers/models/idefics3/modeling_idefics3.py

@@ -35,7 +35,7 @@ from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ....utils.runtime_utils import RBLNPytorchRuntime
 from ..decoderonly.modeling_decoderonly import (
-    RBLNDecoderOnlyOutput,
+    RBLNDecoderOnlyForCausalLMOutput,
 )
 
 
@@ -494,7 +494,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
         if not return_dict:
             return logits, generate_idx
         else:
-            return RBLNDecoderOnlyOutput(
+            return RBLNDecoderOnlyForCausalLMOutput(
                 logits=logits,
                 generate_idx=generate_idx,
             )

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

@@ -12,5 +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
+from .configuration_llama import RBLNLlamaForCausalLMConfig, RBLNLlamaModelConfig
+from .modeling_llama import RBLNLlamaForCausalLM, RBLNLlamaModel

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNLlamaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -40,3 +40,11 @@ class RBLNLlamaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     )
     ```
     """
+
+
+class RBLNLlamaModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN Llama models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+    """

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

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .llama_architecture import LlamaWrapper
 
 
@@ -81,3 +81,15 @@ class RBLNLlamaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
 
     _decoder_wrapper_cls = LlamaWrapper
+
+
+class RBLNLlamaModel(RBLNDecoderOnlyModel):
+    """
+    The Llama Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModel`]. 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 LlamaModel model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers LlamaModel model into a RBLN transformer model by:
+    """
+
+    _decoder_wrapper_cls = LlamaWrapper

optimum/rbln/transformers/models/llava/__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_llava import RBLNLlavaForConditionalGenerationConfig
+from .modeling_llava import RBLNLlavaForConditionalGeneration

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

@@ -0,0 +1,54 @@
+# 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 Any, Dict, Optional
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNLlavaForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNLlavaForConditionalGenerationConfig.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized LLaVA models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
+    submodules = ["vision_tower", "language_model"]
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        language_model: Optional[RBLNModelConfig] = None,
+        **kwargs: Dict[str, Any],
+    ):
+        """
+        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