optimum-rbln 0.8.4a8__py3-none-any.whl → 0.9.2__py3-none-any.whl

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

@@ -11,6 +11,7 @@
 # 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.
+import importlib
 import inspect
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -25,7 +26,9 @@ from transformers.models.gemma3.modeling_gemma3 import Gemma3TextScaledWordEmbed
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ...modeling_outputs import RBLNDecoderOnlyOutput
+from ...utils.rbln_runtime_wrapper import LoopProcessor
 from ..decoderonly.decoderonly_runtime_utils import RBLNPageTableManager
+from ..decoderonly.generation_decoderonly import RBLNDecoderOnlyGenerationMixin
 from ..decoderonly.modeling_decoderonly import (
     RBLNDecoderOnlyModelForCausalLM,
 )
@@ -38,61 +41,44 @@ if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, Gemma3ForConditionalGeneration
 
 
-class LoopVisionTower:
-    def __init__(self, vision_tower: RBLNModel) -> None:
-        self.vision_tower = vision_tower
+class LoopVisionTower(LoopProcessor):
+    def __init__(self, vision_tower: "RBLNModel"):
+        super().__init__(model=vision_tower)
 
-    def forward(self, *args, **kwargs):
-        # Loop instead of batch
-        # shape of pixel_values : [batch, num_channel, height, width]
-        pixel_values = args[0]
+    def _get_batch_size(self, pixel_values, **kwargs):
+        return pixel_values.shape[0]
 
-        batch_size = pixel_values.shape[0]
-        outputs = []
-        for i in range(batch_size):
-            outputs.append(self.vision_tower(pixel_values=pixel_values[i : i + 1], return_dict=True))
+    def _prepare_inputs_for_iteration(self, index, common_inputs, pixel_values, **kwargs):
+        pixel_values_item = pixel_values[index : index + 1]
+        out_buffer = [tensor[index : index + 1] for tensor in kwargs["out"]]
+        return ([pixel_values_item], {"out": out_buffer})
 
-        last_hidden_states = [output.last_hidden_state for output in outputs]
-
-        # FIXME:: This can be optimized using out= API of rbln runtime.
-        last_hidden_states = torch.cat(last_hidden_states, dim=0)
+    def _process_outputs(self, outputs: list, **kwargs) -> "BaseModelOutputWithPooling":
+        output = kwargs["out"]
 
         return BaseModelOutputWithPooling(
-            last_hidden_state=last_hidden_states,
+            last_hidden_state=output[0],
         )
 
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
-
-    def __repr__(self) -> str:
-        return repr(self.vision_tower)
-
 
-class LoopProjector:
-    def __init__(self, multi_modal_projector) -> None:
-        self.multi_modal_projector = multi_modal_projector
+class LoopProjector(LoopProcessor):
+    def __init__(self, multi_modal_projector: "RBLNModel"):
+        super().__init__(model=multi_modal_projector)
 
-    def forward(self, *args, **kwargs):
-        # Loop instead of batch
-        image_feature = args[0]
+    def _get_batch_size(self, image_feature, **kwargs):
+        return image_feature.shape[0]
 
-        batch_size = image_feature.shape[0]
-        outputs = []
-        for i in range(batch_size):
-            outputs.append(self.multi_modal_projector(image_feature[i : i + 1]))
+    def _prepare_inputs_for_iteration(self, index, common_inputs, image_feature, **kwargs):
+        image_feature_item = image_feature[index : index + 1]
+        out_buffer = [tensor[index : index + 1] for tensor in kwargs["out"]]
+        return ([image_feature_item], {"out": out_buffer})
 
-        # FIXME:: This can be optimized using out= API of rbln runtime.
-        outputs = torch.cat(outputs, dim=0)
-        return outputs
+    def _process_outputs(self, outputs: list, **kwargs):
+        output = kwargs["out"]
+        return output[0]
 
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
 
-    def __repr__(self) -> str:
-        return repr(self.multi_modal_projector)
-
-
-class RBLNGemma3ForConditionalGeneration(RBLNModel):
+class RBLNGemma3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixin):
     auto_model_class = AutoModelForImageTextToText
     _rbln_submodules = [
         {"name": "vision_tower"},
@@ -112,6 +98,23 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
     def can_generate(self):
         return True
 
+    @classmethod
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
+
+        with no_init_weights():
+            model_cls_name = model.model.language_model.__class__.__name__
+            causal_model_cls_name = model_cls_name.replace("TextModel", "ForCausalLM")
+            causal_model_cls = getattr(importlib.import_module("transformers"), causal_model_cls_name)
+            new_language_model = causal_model_cls(model.model.language_model.config)
+
+        new_language_model.lm_head = model.lm_head
+        new_language_model.model = model.model.language_model
+        model.model.language_model = new_language_model
+        model.lm_head = None
+        del model.lm_head
+        return model
+
     def __post_init__(self, **kwargs):
         self.vision_tower = LoopVisionTower(self.rbln_submodules[0])
         self.language_model = self.rbln_submodules[1]
@@ -210,8 +213,21 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
         # Returns:
         #     Image feature tensor of shape `(num_images, image_length, embed_dim)`.
 
-        vision_outputs = self.vision_tower(pixel_values).last_hidden_state
-        image_features = self.multi_modal_projector(vision_outputs)
+        vision_out_buffer = []
+        vision_out_size = [
+            pixel_values.shape[0],
+            (self.config.vision_config.image_size // self.config.vision_config.patch_size) ** 2,
+            self.config.vision_config.hidden_size,
+        ]
+        projector_out_size = [
+            pixel_values.shape[0],
+            self.config.mm_tokens_per_image,
+            self.config.text_config.hidden_size,
+        ]
+        vision_out_buffer.append(torch.empty(size=vision_out_size, dtype=torch.float32, device="cpu"))
+        projector_out_buffer = [torch.empty(size=projector_out_size, dtype=torch.float32, device="cpu")]
+        vision_outputs = self.vision_tower(pixel_values, out=vision_out_buffer).last_hidden_state
+        image_features = self.multi_modal_projector(vision_outputs, out=projector_out_buffer)
         return image_features
 
     def _preprocess_prefill(
@@ -393,6 +409,13 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     def _update_sliding_window_config(cls, model_config: PretrainedConfig, rbln_config: RBLNGemma3ForCausalLMConfig):
         sliding_window = getattr(model_config, "sliding_window", None)
         sliding_window_pattern = getattr(model_config, "sliding_window_pattern", None)
+        if sliding_window_pattern is None:
+            if hasattr(model_config, "layer_types"):
+                first_full_attention_index = model_config.layer_types.index("full_attention")
+                sliding_window_pattern = first_full_attention_index + 1
+            else:
+                raise ValueError("Cannot determine sliding_window_pattern from model_config")
+
         if sliding_window_pattern <= model_config.num_hidden_layers:
             rbln_config.cache_impl = "hybrid"
             rbln_config.sliding_window = sliding_window

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

@@ -75,7 +75,10 @@ class GPT2Attention(DecoderOnlyAttention):
         self.o_proj = self._original_mod.c_proj
         self.split_size = self._original_mod.split_size
 
-    def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def projection(self, hidden_states, lora_int_id) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if lora_int_id is not None:
+            raise NotImplementedError("LoRA is not supported for GPT2Attention")
+
         query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
         return query_states, key_states, value_states
 

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

@@ -15,10 +15,29 @@
 from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
 
 
 class RBLNIdefics3VisionTransformerConfig(RBLNModelConfig):
-    pass
+    """
+    Configuration class for RBLNIdefics3VisionTransformer.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Idefics3 vision transformer.
+    """
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        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}")
 
 
 class RBLNIdefics3ForConditionalGenerationConfig(RBLNModelConfig):
@@ -61,5 +80,10 @@ class RBLNIdefics3ForConditionalGenerationConfig(RBLNModelConfig):
         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_model = vision_model
-        self.text_model = text_model
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for Idefics3 vision transformer. It will be set to 1.")
+
+        self.vision_model = self.initialize_submodule_config(
+            submodule_config=vision_model, batch_size=1, force_kwargs=True
+        )
+        self.text_model = self.initialize_submodule_config(submodule_config=text_model)

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

@@ -35,6 +35,7 @@ from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ....utils.runtime_utils import RBLNPytorchRuntime
 from ...modeling_outputs import RBLNDecoderOnlyOutput
+from ..decoderonly.generation_decoderonly import RBLNDecoderOnlyGenerationMixin
 
 
 if TYPE_CHECKING:
@@ -75,10 +76,12 @@ class RBLNRuntimeVisionModel(RBLNPytorchRuntime):
 
         hidden_states = self.embeddings(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask)
 
-        return super().forward(hidden_states.contiguous())
+        return super().forward(hidden_states.contiguous(), **kwargs)
 
 
 class RBLNIdefics3VisionTransformer(RBLNModel):
+    _tp_support = False
+
     def __post_init__(self, **kwargs):
         artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
         with no_init_weights():
@@ -118,9 +121,6 @@ class RBLNIdefics3VisionTransformer(RBLNModel):
                 encoder_outputs = self.encoder(
                     inputs_embeds=hidden_states,
                     attention_mask=patch_attention_mask,
-                    output_attentions=None,
-                    output_hidden_states=None,
-                    return_dict=False,
                 )
                 last_hidden_state = encoder_outputs[0]
                 last_hidden_state = self.post_layernorm(last_hidden_state)
@@ -140,8 +140,7 @@ class RBLNIdefics3VisionTransformer(RBLNModel):
             (
                 "hidden_states",
                 [
-                    # batch_size * num_patches (dependent on image size) -> compile with 1 and use for loop
-                    1,
+                    rbln_config.batch_size,
                     (model_config.image_size // model_config.patch_size) ** 2,
                     model_config.hidden_size,
                 ],
@@ -160,29 +159,31 @@ class RBLNIdefics3VisionTransformer(RBLNModel):
         return_dict: Optional[bool] = None,
         **kwargs,
     ) -> Union[Tuple, BaseModelOutput]:
-        batch_size = pixel_values.shape[0]
-        last_hidden_state = []
-        for i in range(batch_size):
+        last_hidden_state_size = [
+            pixel_values.shape[0],
+            (self.config.image_size // self.config.patch_size) ** 2,
+            self.config.hidden_size,
+        ]
+        last_hidden_state = torch.empty(size=last_hidden_state_size, dtype=torch.float32, device="cpu")
+        for i in range(pixel_values.shape[0]):
             if patch_attention_mask is not None:
                 batch_attention_mask = patch_attention_mask[i : i + 1,]
             else:
                 batch_attention_mask = None
 
-            batch_hidden_state = self.model(
+            self.model(
                 pixel_values[i : i + 1,],
                 batch_attention_mask,
+                out=last_hidden_state[i : i + 1,],
                 return_dict=False,
             )
-            last_hidden_state.append(batch_hidden_state)
-        last_hidden_state = torch.cat(last_hidden_state, dim=0)
-
         if not return_dict:
             return (last_hidden_state,)
         else:
             return BaseModelOutput(last_hidden_state=last_hidden_state)
 
 
-class RBLNIdefics3ForConditionalGeneration(RBLNModel):
+class RBLNIdefics3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixin):
     """
     RBLNIdefics3ForConditionalGeneration 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.
@@ -285,8 +286,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
             (
                 "image_hidden_states",
                 [
-                    # batch_size * num_patches (dependent on image size) -> compile with 1 and use for loop
-                    1,
+                    rbln_config.vision_model.batch_size,
                     (model_config.vision_config.image_size // model_config.vision_config.patch_size) ** 2,
                     model_config.vision_config.hidden_size,
                 ],
@@ -425,10 +425,15 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
                 pixel_values=pixel_values, patch_attention_mask=patch_attention_mask, return_dict=True
             ).last_hidden_state
 
-            connector_outputs = []
+            connector_output_size = [
+                image_hidden_states.shape[0],
+                image_hidden_states.shape[1] // self.config.scale_factor**2,
+                self.config.text_config.hidden_size,
+            ]
+            connector_outputs = torch.empty(size=connector_output_size, dtype=torch.float32, device="cpu")
             for i in range(image_hidden_states.shape[0]):
-                connector_outputs.append(self.connector(image_hidden_states[i : i + 1,]))
-            image_hidden_states = torch.cat(connector_outputs, dim=0)
+                self.connector(image_hidden_states[i : i + 1,], out=connector_outputs[i : i + 1,])
+            image_hidden_states = connector_outputs
 
         elif image_hidden_states is not None:
             image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=input_ids.device)

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

@@ -15,6 +15,10 @@
 from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
 
 
 class RBLNLlavaForConditionalGenerationConfig(RBLNModelConfig):
@@ -54,5 +58,15 @@ class RBLNLlavaForConditionalGenerationConfig(RBLNModelConfig):
         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
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for Llava vision tower. It will be set to 1.")
+
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower,
+            batch_size=1,  # vision_tower batch_size is always 1 in Llava
+            force_kwargs=True,
+        )
+
+        self.language_model = self.initialize_submodule_config(
+            submodule_config=language_model,
+        )

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

@@ -12,18 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 import inspect
 from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union
 
 import torch
 from transformers import AutoModelForImageTextToText, LlavaForConditionalGeneration, PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.modeling_utils import no_init_weights
 from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast
 
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ....utils.logging import get_logger
 from ...modeling_outputs import RBLNDecoderOnlyOutput
+from ...utils.rbln_runtime_wrapper import LoopProcessor
+from ..decoderonly.generation_decoderonly import RBLNDecoderOnlyGenerationMixin
 
 
 logger = get_logger(__name__)
@@ -32,20 +36,32 @@ if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
 
 
-class LoopVisionTower:
-    def __init__(self, vision_tower: RBLNModel) -> None:
-        self.vision_tower = vision_tower
+class LoopVisionTower(LoopProcessor):
+    def __init__(self, vision_tower):
+        # FIXME: need to know RBLNModel or RuntimeWrapper
+        if hasattr(vision_tower.model, "runtime"):
+            super().__init__(model=vision_tower)
+        else:
+            super().__init__(model=vision_tower.model[0])
 
-    def forward(self, pixel_values, image_sizes: Optional[torch.Tensor] = None, **kwargs):
-        outputs = []
-        for i in range(pixel_values.shape[0]):
-            outputs.append(
-                self.vision_tower(
-                    pixel_values[i : i + 1], image_sizes[i : i + 1] if image_sizes is not None else None, **kwargs
-                )
-            )
+        self.rbln_config = vision_tower.rbln_config
+
+    def _get_batch_size(self, pixel_values, **kwargs):
+        return pixel_values.shape[0]
+
+    def _prepare_inputs_for_iteration(self, index, common_inputs, pixel_values, **kwargs):
+        pixel_values_item = pixel_values[index : index + 1]
+        if "image_sizes" in kwargs and kwargs["image_sizes"] is not None:
+            ret_val = [pixel_values_item, kwargs["image_sizes"][index : index + 1]]
+        else:
+            ret_val = [pixel_values_item]
+
+        out_buffer = [tensor[index : index + 1] for tensor in kwargs["out"]] if "out" in kwargs else None
+        return (ret_val, {"out": out_buffer})
 
-        if hasattr(self.vision_tower.rbln_config, "max_image_size"):
+    def _process_outputs(self, outputs: list, **kwargs) -> "BaseModelOutputWithPooling":
+        # when use another Wrapper
+        if hasattr(self.rbln_config, "max_image_size"):
             last_hidden_states = [output.last_hidden_state for output in outputs]
             last_hidden_states = torch.cat(last_hidden_states, dim=1)
             hidden_states = tuple(
@@ -55,52 +71,40 @@ class LoopVisionTower:
                 )
                 for layer_idx in range(len(outputs[0].hidden_states))
             )
-
         else:
-            last_hidden_states = [output.last_hidden_state for output in outputs]
-            last_hidden_states = torch.cat(last_hidden_states, dim=0)
-            hidden_states = [output.hidden_states for output in outputs]
-            hidden_states = tuple(
-                torch.cat(tuple((hidden_states[n][i] for n in range(pixel_values.shape[0]))), dim=0)
-                for i in range(len(hidden_states[0]))
-            )
+            output = kwargs["out"]
+            last_hidden_states = output[0]
+
+            if not output[2:]:
+                hidden_states = None
+            else:
+                hidden_states = tuple(output[2:])
 
         return BaseModelOutputWithPooling(
             last_hidden_state=last_hidden_states,
+            pooler_output=None,
             hidden_states=hidden_states,
         )
 
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
 
-    def __repr__(self) -> str:
-        return repr(self.vision_tower)
+class LoopProjector(LoopProcessor):
+    def __init__(self, multi_modal_projector: "RBLNModel"):
+        super().__init__(model=multi_modal_projector)
 
+    def _get_batch_size(self, image_feature, **kwargs):
+        return image_feature.shape[0]
 
-class LoopProjector:
-    def __init__(self, multi_modal_projector) -> None:
-        self.multi_modal_projector = multi_modal_projector
+    def _prepare_inputs_for_iteration(self, index, common_inputs, image_feature, **kwargs):
+        image_feature_item = image_feature[index : index + 1]
+        out_buffer = [tensor[index : index + 1] for tensor in kwargs["out"]]
+        return ([image_feature_item], {"out": out_buffer})
 
-    def forward(self, *args, **kwargs):
-        # Loop instead of batch
-        image_feature = args[0]
+    def _process_outputs(self, outputs: list, **kwargs):
+        output = kwargs["out"]
+        return output[0]
 
-        outputs = []
-        for i in range(image_feature.shape[0]):
-            outputs.append(self.multi_modal_projector(image_feature[i : i + 1]))
 
-        # FIXME:: This can be optimized using out= API of rbln runtime.
-        outputs = torch.cat(outputs, dim=0)
-        return outputs
-
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
-
-    def __repr__(self) -> str:
-        return repr(self.multi_modal_projector)
-
-
-class RBLNLlavaForConditionalGeneration(RBLNModel):
+class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixin):
     """
     RBLNLlavaForConditionalGeneration is a multi-modal model that combines vision and language processing capabilities,
     optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
@@ -170,6 +174,23 @@ class RBLNLlavaForConditionalGeneration(RBLNModel):
     def can_generate(self):
         return True
 
+    @classmethod
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
+
+        with no_init_weights():
+            model_cls_name = model.model.language_model.__class__.__name__
+            causal_model_cls_name = model_cls_name.replace("Model", "ForCausalLM")
+            causal_model_cls = getattr(importlib.import_module("transformers"), causal_model_cls_name)
+            new_language_model = causal_model_cls(model.model.language_model.config)
+
+        new_language_model.lm_head = model.lm_head
+        new_language_model.model = model.model.language_model
+        model.model.language_model = new_language_model
+        model.lm_head = None
+        del model.lm_head
+        return model
+
     def __post_init__(self, **kwargs):
         self.vision_tower = LoopVisionTower(self.rbln_submodules[0])
         self.language_model = self.rbln_submodules[1]
@@ -201,7 +222,7 @@ class RBLNLlavaForConditionalGeneration(RBLNModel):
         # support for pixtral that needs padding
         if hasattr(rbln_config.vision_tower, "max_image_size"):
             num_positions = (
-                rbln_config.vision_tower.batch_size
+                rbln_config.batch_size
                 * (rbln_config.vision_tower.max_image_size[0] // model_config.vision_config.patch_size)
                 * (rbln_config.vision_tower.max_image_size[1] // model_config.vision_config.patch_size)
             )
@@ -217,7 +238,11 @@ class RBLNLlavaForConditionalGeneration(RBLNModel):
         input_info = [
             (
                 "image_features",
-                [rbln_config.batch_size, selected_image_feature_dim, model_config.vision_config.hidden_size],
+                [
+                    1,
+                    selected_image_feature_dim,
+                    model_config.vision_config.hidden_size,
+                ],
                 "float32",
             )
         ]
@@ -290,7 +315,31 @@ class RBLNLlavaForConditionalGeneration(RBLNModel):
             raise ValueError(f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}")
 
         kwargs = {k: v for k, v in kwargs.items() if v is not None}
-        image_outputs = self.vision_tower(pixel_values, output_hidden_states=True, **kwargs)
+
+        # prepare out buffer for pre-allocation
+        if hasattr(self.rbln_config.vision_tower, "max_image_size"):
+            vision_out_size = [
+                pixel_values.shape[0],
+                (self.rbln_config.vision_tower.max_image_size[0] // self.config.vision_config.patch_size)
+                * (self.rbln_config.vision_tower.max_image_size[1] // self.config.vision_config.patch_size),
+                self.config.vision_config.hidden_size,
+            ]
+            pooler_out_size = None
+        else:
+            vision_out_size = [
+                pixel_values.shape[0],
+                (self.config.vision_config.image_size // self.config.vision_config.patch_size) ** 2 + 1,
+                self.config.vision_config.hidden_size,
+            ]
+            pooler_out_size = [pixel_values.shape[0], self.config.vision_config.hidden_size]
+
+        vision_out_buffer = []
+        for i in range(self.config.vision_config.num_hidden_layers + 2):
+            vision_out_buffer.append(torch.empty(size=vision_out_size, dtype=torch.float32, device="cpu"))
+        if pooler_out_size is not None:
+            vision_out_buffer.insert(1, torch.empty(size=pooler_out_size, dtype=torch.float32, device="cpu"))
+
+        image_outputs = self.vision_tower(pixel_values, output_hidden_states=True, out=vision_out_buffer, **kwargs)
 
         if isinstance(vision_feature_layer, int):
             selected_image_feature = image_outputs.hidden_states[vision_feature_layer]
@@ -311,15 +360,24 @@ class RBLNLlavaForConditionalGeneration(RBLNModel):
             )
             num_padding_patches = max_patches - num_real_patches
 
+            projector_out_size = [1, max_patches, self.config.text_config.hidden_size]
+            projector_out_buffer = [torch.empty(size=projector_out_size, dtype=torch.float32, device="cpu")]
+
             padding_tensor = torch.zeros(
                 (selected_image_feature.shape[0], num_padding_patches, selected_image_feature.shape[2]),
                 dtype=selected_image_feature.dtype,
             )
             padded_feature = torch.cat([selected_image_feature, padding_tensor], dim=1)
-            padded_projected_feature = self.multi_modal_projector(padded_feature)
+            padded_projected_feature = self.multi_modal_projector(padded_feature, out=projector_out_buffer)
             image_features = padded_projected_feature[:, :num_real_patches, :]
         else:
-            image_features = self.multi_modal_projector(selected_image_feature)
+            projector_out_size = [
+                pixel_values.shape[0] * pixel_values.shape[1],
+                (self.config.vision_config.image_size // self.config.vision_config.patch_size) ** 2,
+                self.config.text_config.hidden_size,
+            ]
+            projector_out_buffer = [torch.empty(size=projector_out_size, dtype=torch.float32, device="cpu")]
+            image_features = self.multi_modal_projector(selected_image_feature, out=projector_out_buffer)
 
         return image_features
 

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

@@ -16,7 +16,6 @@ from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
 from ....utils.logging import get_logger
-from ...models.clip import RBLNCLIPVisionModelConfig
 
 
 logger = get_logger(__name__)
@@ -55,17 +54,16 @@ class RBLNLlavaNextForConditionalGenerationConfig(RBLNModelConfig):
         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 = self.init_submodule_config(
-            RBLNCLIPVisionModelConfig,
-            vision_tower,
-        )
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for LlavaNext vision tower. It will be set to 1.")
 
-        if self.vision_tower.output_hidden_states is False:
-            raise ValueError(
-                f"LlavaNext requires output_hidden_states to be True, but found output_hidden_states={self.vision_tower.output_hidden_states}. "
-                f"Please compile again with the correct argument."
-            )
-        else:
-            self.vision_tower.output_hidden_states = True
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower,
+            batch_size=1,  # vision_tower batch_size is always 1 in LlavaNext
+            output_hidden_states=True,  # LlavaNext requires output_hidden_states to be True
+            force_kwargs=True,
+        )
 
-        self.language_model = language_model
+        self.language_model = self.initialize_submodule_config(
+            submodule_config=language_model,
+        )