optimum-rbln 0.8.2a4__py3-none-any.whl → 0.9.3__py3-none-any.whl

optimum/rbln/transformers/models/pixtral/__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_pixtral import RBLNPixtralVisionModelConfig
+from .modeling_pixtral import RBLNPixtralVisionModel

optimum/rbln/transformers/models/pixtral/configuration_pixtral.py

@@ -0,0 +1,43 @@
+# 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, Optional, Tuple
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNPixtralVisionModelConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        max_image_size: Tuple = None,
+        batch_size: Optional[int] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            max_image_size (Tuple): The size of max input images. A tuple (max_height, max_width)
+            batch_size (Optional[int]): The batch size for image processing. Defaults to 1.
+            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.max_image_size = max_image_size
+        self.output_hidden_states = output_hidden_states

optimum/rbln/transformers/models/pixtral/modeling_pixtral.py

@@ -0,0 +1,322 @@
+# 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 pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
+
+import rebel
+import torch
+import torch.nn as nn
+from transformers import PixtralVisionConfig, PixtralVisionModel
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.modeling_utils import no_init_weights
+from transformers.models.pixtral.modeling_pixtral import PixtralRMSNorm, PixtralRotaryEmbedding
+
+from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
+from ....modeling import RBLNModel
+from ....utils.logging import get_logger
+from ....utils.runtime_utils import RBLNPytorchRuntime
+from .configuration_pixtral import RBLNPixtralVisionModelConfig
+from .pixtral_architecture import PixtralAttention
+
+
+logger = get_logger(__name__)
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+
+class RBLNRuntimePixtralVisionModel(RBLNPytorchRuntime):
+    mandatory_members = ["main_input_name"]
+
+    def __init__(
+        self,
+        runtime: rebel.Runtime,
+        config: PixtralVisionConfig,
+        rbln_config: RBLNPixtralVisionModelConfig,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(runtime, **kwargs)
+        self.patch_positional_embedding = PixtralRotaryEmbedding(config)
+        self.patch_size = config.patch_size
+        self.image_size = config.image_size
+        self.hidden_size = config.hidden_size
+        self.max_image_size = rbln_config.max_image_size
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        image_sizes: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ):
+        if pixel_values.shape[2] > self.max_image_size[0] or pixel_values.shape[3] > self.max_image_size[1]:
+            raise ValueError("The height() and width of pixel_values can't be larger than max_image_size.")
+
+        if pixel_values.shape[2] != self.max_image_size[0] or pixel_values.shape[3] != self.max_image_size[1]:
+            padded_pixel_values = [
+                torch.nn.functional.pad(
+                    image,
+                    pad=(
+                        0,
+                        self.max_image_size[1] - pixel_values.shape[3],
+                        0,
+                        self.max_image_size[0] - pixel_values.shape[2],
+                    ),
+                )
+                for image in pixel_values
+            ]
+            pixel_values = torch.stack(padded_pixel_values)
+
+        batch_size, _, H_max, W_max = pixel_values.shape
+        H_max_p = H_max // self.patch_size
+        W_max_p = W_max // self.patch_size
+
+        final_hidden_states = None
+
+        last_hidden_state_list = []
+        if output_hidden_states:
+            batch_hidden_states_list = []
+
+        for i in range(batch_size):
+            h_patched_original = image_sizes[i, 0] // self.patch_size
+            w_patched_original = image_sizes[i, 1] // self.patch_size
+
+            single_pixel_values = pixel_values[i : i + 1]
+            patch_embed = self.patch_conv(single_pixel_values)
+            patch_embed_seq = patch_embed[:, :, :h_patched_original, :w_patched_original].flatten(2).transpose(1, 2)
+            patch_embed_seq = self.ln_pre(patch_embed_seq)
+            patch_embed_seq = nn.functional.pad(
+                patch_embed_seq, (0, 0, 0, H_max_p * W_max_p - patch_embed_seq.shape[1]), "constant", value=0
+            )
+
+            max_w_from_config = self.image_size // self.patch_size
+            mesh = torch.meshgrid(torch.arange(h_patched_original), torch.arange(w_patched_original), indexing="ij")
+            h_grid, v_grid = torch.stack(mesh, dim=-1).reshape(-1, 2).chunk(2, -1)
+            ids = h_grid * max_w_from_config + v_grid
+            position_ids = ids[:, 0]
+
+            position_embeddings = self.patch_positional_embedding(patch_embed_seq, position_ids)
+            cos = nn.functional.pad(
+                position_embeddings[0],
+                (0, 0, 0, H_max_p * W_max_p - position_embeddings[0].shape[0]),
+                "constant",
+                value=0,
+            )
+            sin = nn.functional.pad(
+                position_embeddings[1],
+                (0, 0, 0, H_max_p * W_max_p - position_embeddings[1].shape[0]),
+                "constant",
+                value=0,
+            )
+
+            attention_mask = torch.full(
+                (1, patch_embed_seq.shape[-2]), fill_value=torch.finfo(patch_embed_seq.dtype).min
+            )
+            attention_mask[:, : h_patched_original * w_patched_original] = 0
+            if "out" in kwargs:
+                super().forward(patch_embed_seq, attention_mask, cos, sin, **kwargs)
+                transformer_output = kwargs["out"]
+            else:
+                transformer_output = super().forward(patch_embed_seq, attention_mask, cos, sin, **kwargs)
+
+            last_hidden_state_list.append(transformer_output[0][:, : h_patched_original * w_patched_original, :])
+            hidden_states = transformer_output[1:]
+
+            if output_hidden_states:
+                batch_hidden_states_list.append(
+                    [hidden_state[:, : h_patched_original * w_patched_original, :] for hidden_state in hidden_states]
+                )
+
+        final_last_hidden_state = torch.cat(last_hidden_state_list, dim=1)
+
+        if output_hidden_states:
+            hidden_states = [
+                torch.cat(
+                    [batch_hidden_states[layer_idx] for batch_hidden_states in batch_hidden_states_list],
+                    dim=1,
+                )
+                for layer_idx in range(len(batch_hidden_states_list[0]))
+            ]
+
+            final_hidden_states = tuple(hidden_states)
+
+        if not return_dict:
+            return tuple(v for v in (final_last_hidden_state, final_hidden_states) if v is not None)
+
+        # TODO: output_attentions
+        return BaseModelOutput(
+            last_hidden_state=final_last_hidden_state,
+            hidden_states=final_hidden_states,
+        )
+
+
+class _PixtralVisionModel(torch.nn.Module):
+    def __init__(self, model: PixtralVisionModel, output_hidden_states: bool):
+        super().__init__()
+        self.transformer = self.convert_to_rbln_pixtral_vision_model(model)
+        self.output_hidden_states = output_hidden_states
+
+    def convert_to_rbln_pixtral_vision_model(self, model: nn.Module):
+        for layer in model.transformer.layers:
+            layer.attention = PixtralAttention(layer.attention)
+        return model.transformer
+
+    def forward(self, patch_embeds, attention_mask, position_embeddings_1, position_embeddings_2):
+        output = self.transformer(
+            inputs_embeds=patch_embeds,
+            attention_mask=attention_mask,
+            position_embeddings=(position_embeddings_1, position_embeddings_2),
+            output_hidden_states=self.output_hidden_states,
+            return_dict=False,
+        )
+        return output
+
+
+class RBLNPixtralVisionModel(RBLNModel):
+    """
+    RBLN optimized Pixtral vision encoder model.
+
+    This class provides hardware-accelerated inference for Pixtral vision encoders
+    on RBLN devices, supporting image encoding for multimodal tasks.
+    """
+
+    def __post_init__(self, **kwargs):
+        artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
+        with no_init_weights():
+            self.patch_conv = nn.Conv2d(
+                in_channels=self.config.num_channels,
+                out_channels=self.config.hidden_size,
+                kernel_size=self.config.patch_size,
+                stride=self.config.patch_size,
+                bias=False,
+            )
+            self.ln_pre = PixtralRMSNorm(self.config.hidden_size, eps=1e-5)
+        self.patch_conv.load_state_dict(artifacts["patch_conv"])
+        self.ln_pre.load_state_dict(artifacts["ln_pre"])
+        self.model = RBLNRuntimePixtralVisionModel(
+            self.model[0],
+            main_input_name="pixel_values",
+            config=self.config,
+            rbln_config=self.rbln_config,
+            patch_conv=self.patch_conv,
+            ln_pre=self.ln_pre,
+        )
+
+    @classmethod
+    def save_torch_artifacts(
+        cls,
+        model: "PreTrainedModel",
+        save_dir_path: Path,
+        subfolder: str,
+        rbln_config: RBLNModelConfig,
+    ):
+        save_dict = {}
+        save_dict["patch_conv"] = model.get_input_embeddings().state_dict()
+        save_dict["ln_pre"] = model.ln_pre.state_dict()
+        torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
+
+    @classmethod
+    def _wrap_model_if_needed(
+        cls, model: torch.nn.Module, rbln_config: RBLNPixtralVisionModelConfig
+    ) -> torch.nn.Module:
+        wrapper_cfg = {
+            "output_hidden_states": rbln_config.output_hidden_states,
+        }
+        return _PixtralVisionModel(model, **wrapper_cfg).eval()
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: Optional["PreTrainedModel"] = None,
+        model_config: "PixtralVisionConfig" = None,
+        rbln_config: Optional[RBLNPixtralVisionModelConfig] = None,
+    ) -> RBLNPixtralVisionModelConfig:
+        if rbln_config.max_image_size is None:
+            rbln_config.max_image_size = (model_config.image_size, model_config.image_size)
+
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
+
+        num_total_patches = (rbln_config.max_image_size[0] // model_config.patch_size) * (
+            rbln_config.max_image_size[1] // model_config.patch_size
+        )
+
+        rbln_compile_config = RBLNCompileConfig(
+            input_info=[
+                (
+                    "patch_embeds",
+                    [1, num_total_patches, model_config.hidden_size],
+                    "float32",
+                ),
+                ("attention_mask", [1, num_total_patches], "float32"),
+                (
+                    "position_embeddings_1",
+                    [
+                        num_total_patches,
+                        model_config.head_dim,
+                    ],
+                    "float32",
+                ),
+                (
+                    "position_embeddings_2",
+                    [
+                        num_total_patches,
+                        model_config.head_dim,
+                    ],
+                    "float32",
+                ),
+            ]
+        )
+
+        rbln_config.set_compile_cfgs([rbln_compile_config])
+        return rbln_config
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        image_sizes: Optional[torch.FloatTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutput]:
+        """
+        Forward pass for the RBLN-optimized Pixtral vision model.
+
+        Args:
+            pixel_values (torch.Tensor of shape (batch_size, num_channels, image_size, image_size)) — The tensors corresponding to the input images. Pixel values can be obtained using PixtralImageProcessor. See PixtralImageProcessor.call() for details (PixtralProcessor uses PixtralImageProcessor for processing images).
+            image_sizes (torch.Tensor of shape (batch_size, 2), optional) — The sizes of the images in the batch, being (height, width) for each image.
+            output_hidden_states (bool, optional) — Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for more detail.
+            return_dict (bool, optional) — Whether or not to return a ModelOutput instead of a plain tuple.
+
+        Returns:
+            BaseModelOutput or tuple(torch.FloatTensor)
+        """
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        if output_hidden_states != self.rbln_config.output_hidden_states:
+            raise ValueError(
+                f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
+                f"Please compile again with the correct argument."
+            )
+
+        output = self.model(
+            pixel_values, image_sizes, output_hidden_states=output_hidden_states, return_dict=return_dict, **kwargs
+        )
+
+        return output

optimum/rbln/transformers/models/pixtral/pixtral_architecture.py

@@ -0,0 +1,73 @@
+# 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, Tuple
+
+import torch
+import torch.nn as nn
+
+from ..decoderonly.decoderonly_architecture import apply_rotary_pos_emb
+
+
+class PixtralAttention(nn.Module):
+    def __init__(self, self_attention):
+        super().__init__()
+        self.original_model = self_attention
+        self.num_heads = getattr(self.original_model, "num_heads", None) or getattr(
+            self.original_model.config, "num_attention_heads"
+        )
+        self.head_dim = self.original_model.head_dim
+        self.scaling = self.head_dim**-0.5
+
+        self.__post_init__()
+
+    def __post_init__(self):
+        self.q_proj = self.original_model.q_proj
+        self.k_proj = self.original_model.k_proj
+        self.v_proj = self.original_model.v_proj
+        self.o_proj = self.original_model.o_proj
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+    ):
+        batch_size, patches, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # TODO: return output attention
+        query_states = query_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
+        key_states = key_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
+        value_states = value_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
+
+        cos, sin = position_embeddings
+        cos = cos[None, None, None, :, :]
+        sin = sin[None, None, None, :, :]
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(3, 4)) * self.scaling
+        attn_weights = attn_weights + attention_mask
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.transpose(1, 3)
+
+        attn_output = attn_output.reshape(batch_size, patches, -1)
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, _

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

@@ -12,5 +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
+from .configuration_qwen2 import RBLNQwen2ForCausalLMConfig, RBLNQwen2ModelConfig
+from .modeling_qwen2 import RBLNQwen2ForCausalLM, RBLNQwen2Model

optimum/rbln/transformers/models/qwen2/configuration_qwen2.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 RBLNQwen2ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -40,3 +40,11 @@ class RBLNQwen2ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     )
     ```
     """
+
+
+class RBLNQwen2ModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN Qwen2 models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+    """

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

@@ -15,7 +15,11 @@
 from transformers import PretrainedConfig
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyModelForCausalLMConfig
+from ...models.decoderonly import (
+    RBLNDecoderOnlyModel,
+    RBLNDecoderOnlyModelForCausalLM,
+    RBLNDecoderOnlyModelForCausalLMConfig,
+)
 from .qwen2_architecture import QWEN2Wrapper
 
 
@@ -95,3 +99,25 @@ class RBLNQwen2ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         rbln_config.sliding_window = model_config.sliding_window
         rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
         return rbln_config
+
+
+class RBLNQwen2Model(RBLNDecoderOnlyModel):
+    """
+    The Qwen2 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.
+    """
+
+    _decoder_wrapper_cls = QWEN2Wrapper
+
+    @classmethod
+    def _update_sliding_window_config(
+        cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
+    ):
+        # https://github.com/huggingface/transformers/issues/35896
+        # There seems to be a bug in transformers(v4.52.4). Therefore, similar to when attn_implementation is eager,
+        # we set all layers to use sliding window in this version. This should be updated once the bug is fixed.
+
+        rbln_config.cache_impl = "sliding_window"
+        rbln_config.sliding_window = model_config.sliding_window
+        rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
+        return rbln_config

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, List, Optional, Union
 
 from ....configuration_utils import RBLNModelConfig
 from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
@@ -31,10 +31,22 @@ class RBLNQwen2_5_VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausal
 
     def __init__(
         self,
-        visual: Optional[RBLNModelConfig] = None,
         use_inputs_embeds: bool = True,
-        **kwargs: Dict[str, Any],
+        visual: Optional[RBLNModelConfig] = None,
+        **kwargs: Any,
     ):
+        """
+        Args:
+            use_inputs_embeds (bool): Whether or not to use `inputs_embeds` as input. Defaults to `True`.
+            visual (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+            kwargs: Additional arguments passed to the parent `RBLNDecoderOnlyModelForCausalLMConfig`.
+
+        Raises:
+            ValueError: If `use_inputs_embeds` is False.
+            ValueError: If the visual configuration is provided but contains invalid settings, such as an invalid max_seq_lens (e.g., not a positive integer, not a multiple of the window-based attention unit, or insufficient for the expected resolution).
+            ValueError: If visual is None and no default vision configuration can be inferred for the model architecture.
+            ValueError: If any inherited parameters violate constraints defined in the parent class, such as batch_size not being a positive integer, prefill_chunk_size not being divisible by 64, or max_seq_len not meeting requirements for Flash Attention.
+        """
         super().__init__(use_inputs_embeds=use_inputs_embeds, **kwargs)
         if not self.use_inputs_embeds:
             raise ValueError(
@@ -53,7 +65,7 @@ class RBLNQwen2_5_VisionTransformerPretrainedModelConfig(RBLNModelConfig):
     mechanisms for processing images and videos.
     """
 
-    def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs: Dict[str, Any]):
+    def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs: Any):
         """
         Args:
             max_seq_lens (Optional[Union[int, List[int]]]): Maximum sequence lengths for Vision
@@ -66,10 +78,13 @@ class RBLNQwen2_5_VisionTransformerPretrainedModelConfig(RBLNModelConfig):
                 making 256 (64 * 4) valid. RBLN optimization runs inference per image or video
                 frame, so set `max_seq_len` to match the maximum expected resolution to reduce
                 computation. If not provided, a `ValueError` is raised.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
-            ValueError: If batch_size is not a positive integer.
+            ValueError: If `max_seq_lens` is None or not provided.
+            ValueError: If `max_seq_lens` (or any value in the list) is not a positive integer.
+            ValueError: If `max_seq_lens` is not a multiple of (window_size / patch_size)^2 for window-based attention, or is insufficient for the expected image/video resolution.
+            ValueError: If `batch_size` (inherited from RBLNModelConfig) is not a positive integer.
 
         Max Seq Lens:
             Since `Qwen2_5_VLForConditionalGeneration` performs inference on a per-image or per-frame basis,

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

@@ -17,24 +17,21 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union
 
 import torch
-from transformers import (
-    AutoModelForVision2Seq,
-    PretrainedConfig,
-    PreTrainedModel,
-    Qwen2_5_VLForConditionalGeneration,
-)
+from transformers import AutoModelForVision2Seq, PretrainedConfig, PreTrainedModel, Qwen2_5_VLForConditionalGeneration
 from transformers.modeling_utils import no_init_weights
 from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
     Qwen2_5_VisionPatchEmbed,
     Qwen2_5_VisionRotaryEmbedding,
     Qwen2_5_VisionTransformerPretrainedModel,
+    Qwen2_5_VLModel,
     Qwen2_5_VLRotaryEmbedding,
 )
 
 from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
 from ....utils.logging import get_logger
-from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyOutput
+from ...modeling_outputs import RBLNDecoderOnlyOutput
+from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM
 from .configuration_qwen2_5_vl import (
     RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
     RBLNQwen2_5_VLForConditionalGenerationConfig,
@@ -45,12 +42,7 @@ from .qwen2_5_vl_architecture import Qwen2_5_VisionTransformerWrapper, Qwen2_5_V
 logger = get_logger(__name__)
 
 if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
-        PretrainedConfig,
-    )
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
 
 
 class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
@@ -96,7 +88,7 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     @classmethod
-    def wrap_model_if_needed(
+    def _wrap_model_if_needed(
         cls, model: "PreTrainedModel", rbln_config: RBLNQwen2_5_VisionTransformerPretrainedModelConfig
     ):
         return Qwen2_5_VisionTransformerWrapper(model).eval()
@@ -381,6 +373,8 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         ```
     """
 
+    _supports_non_fp32 = False
+
     auto_model_class = AutoModelForVision2Seq
     _rbln_submodules = [
         {"name": "visual"},
@@ -399,13 +393,11 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         return True
 
     @classmethod
-    def update_kwargs(cls, kwargs):
-        kwargs.update(
-            {
-                "_attn_implementation": "eager",
-            }
-        )
-        return super().update_kwargs(kwargs)
+    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+        model.model.lm_head = model.lm_head
+        model.lm_head = None
+        del model.lm_head
+        return model
 
     @classmethod
     def get_input_info(
@@ -539,7 +531,8 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
             vision_tokens = input_id[0][vision_start_indices + 1]
             image_nums = (vision_tokens == image_token_id).sum()
             video_nums = (vision_tokens == video_token_id).sum()
-            position_ids, rope_deltas = self.get_rope_index(
+            position_ids, rope_deltas = Qwen2_5_VLModel.get_rope_index(
+                self,
                 input_id,
                 image_grid_thw[image_idx : image_idx + image_nums] if image_grid_thw is not None else None,
                 video_grid_thw[video_idx : video_idx + video_nums] if video_grid_thw is not None else None,