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

optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.8.2a2'
-__version_tuple__ = version_tuple = (0, 8, 2, 'a2')
+__version__ = version = '0.8.2a4'
+__version_tuple__ = version_tuple = (0, 8, 2, 'a4')

optimum/rbln/diffusers/modeling_diffusers.py

@@ -70,8 +70,6 @@ class RBLNDiffusionMixin:
     _submodules = []
     _optional_submodules = []
     _prefix = {}
-    _rbln_config_class = None
-    _hf_class = None
 
     @staticmethod
     def _maybe_apply_and_fuse_lora(
@@ -114,14 +112,14 @@ class RBLNDiffusionMixin:
     @classmethod
     def get_rbln_config_class(cls) -> Type[RBLNModelConfig]:
         # Lazily loads and caches the corresponding RBLN model config class.
-        if cls._rbln_config_class is None:
+        if "_rbln_config_class" not in cls.__dict__ or cls._rbln_config_class is None:
             rbln_config_class_name = cls.__name__ + "Config"
             cls._rbln_config_class = get_rbln_config_class(rbln_config_class_name)
         return cls._rbln_config_class
 
     @classmethod
     def get_hf_class(cls):
-        if cls._hf_class is None:
+        if "_hf_class" not in cls.__dict__ or cls._hf_class is None:
             hf_cls_name = cls.__name__[4:]
             library = importlib.import_module("diffusers")
             cls._hf_class = getattr(library, hf_cls_name, None)

optimum/rbln/modeling.py

@@ -35,8 +35,6 @@ logger = get_logger(__name__)
 
 
 class RBLNModel(RBLNBaseModel):
-    _output_class = None
-
     @classmethod
     def update_kwargs(cls, kwargs):
         # Update user-given kwargs to get proper pytorch model.
@@ -289,7 +287,7 @@ class RBLNModel(RBLNBaseModel):
     @classmethod
     def get_hf_output_class(cls):
         # Dynamically gets the output class from the corresponding HuggingFace model class.
-        if cls._output_class:
+        if "_output_class" in cls.__dict__ and cls._output_class is not None:
             return cls._output_class
 
         hf_class = cls.get_hf_class()

optimum/rbln/modeling_base.py

@@ -23,9 +23,10 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
 import rebel
 import torch
 from transformers import AutoConfig, AutoModel, GenerationConfig, PretrainedConfig
+from transformers.utils.hub import PushToHubMixin
 
 from .configuration_utils import RBLNAutoConfig, RBLNCompileConfig, RBLNModelConfig, get_rbln_config_class
-from .utils.hub import PushToHubMixin, pull_compiled_model_from_hub, validate_files
+from .utils.hub import pull_compiled_model_from_hub, validate_files
 from .utils.logging import get_logger
 from .utils.runtime_utils import UnavailableRuntime, tp_and_devices_are_ok
 from .utils.save_utils import maybe_load_preprocessors
@@ -50,11 +51,8 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     model_type = "rbln_model"
     auto_model_class = AutoModel
     config_class = AutoConfig
-
     config_name = "config.json"
     hf_library_name = "transformers"
-    _hf_class = None
-    _rbln_config_class = None
 
     def __init__(
         self,
@@ -115,7 +113,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     def _load_compiled_model_dir(
         cls,
         model_id: Union[str, Path],
-        use_auth_token: Optional[Union[bool, str]] = None,
+        token: Optional[Union[bool, str]] = None,
         revision: Optional[str] = None,
         force_download: bool = False,
         cache_dir: Optional[str] = None,
@@ -134,7 +132,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             model_path = pull_compiled_model_from_hub(
                 model_id=model_id,
                 subfolder=subfolder,
-                use_auth_token=use_auth_token,
+                token=token,
                 revision=revision,
                 cache_dir=cache_dir,
                 force_download=force_download,
@@ -172,7 +170,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         cls,
         model_id: Union[str, Path],
         config: Optional["PretrainedConfig"] = None,
-        use_auth_token: Optional[Union[bool, str]] = None,
+        token: Optional[Union[bool, str]] = None,
         revision: Optional[str] = None,
         force_download: bool = False,
         cache_dir: Optional[str] = None,
@@ -189,7 +187,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         if rbln_compiled_models is None:
             model_path_subfolder = cls._load_compiled_model_dir(
                 model_id=model_id,
-                use_auth_token=use_auth_token,
+                token=token,
                 revision=revision,
                 force_download=force_download,
                 cache_dir=cache_dir,
@@ -232,7 +230,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                         cache_dir=cache_dir,
                         force_download=force_download,
                         revision=revision,
-                        token=use_auth_token,
+                        token=token,
                         trust_remote_code=trust_remote_code,
                     )
                 elif cls.hf_library_name == "diffusers":
@@ -250,7 +248,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                         force_download=force_download,
                         local_files_only=local_files_only,
                         revision=revision,
-                        token=use_auth_token,
+                        token=token,
                         subfolder=subfolder,
                     )
                     config = PretrainedConfig(**config)
@@ -421,7 +419,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
         # Returns:
         #     type: The original HuggingFace model class
-        if cls._hf_class is None:
+        if "_hf_class" not in cls.__dict__ or cls._hf_class is None:
             hf_cls_name = cls.__name__[4:]
             library = importlib.import_module(cls.hf_library_name)
             cls._hf_class = getattr(library, hf_cls_name, None)
@@ -430,7 +428,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     @classmethod
     def get_rbln_config_class(cls) -> Type[RBLNModelConfig]:
         # Lazily loads and caches the corresponding RBLN model config class.
-        if cls._rbln_config_class is None:
+        if "_rbln_config_class" not in cls.__dict__ or cls._rbln_config_class is None:
             rbln_config_class_name = cls.__name__ + "Config"
             cls._rbln_config_class = get_rbln_config_class(rbln_config_class_name)
         return cls._rbln_config_class
@@ -507,6 +505,9 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                 f"Please ensure the model directory exists and you have the necessary permissions to access it."
             )
 
+        if isinstance(self.config, PretrainedConfig):
+            self.config.save_pretrained(real_save_dir)
+
         if save_directory_path == real_save_dir:
             raise FileExistsError(
                 f"Cannot save model to '{save_directory}'. This directory already exists and contains the model files."
@@ -534,7 +535,10 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             raise e  # Re-raise the exception after cleanup
 
         if push_to_hub:
-            return super().push_to_hub(str(save_directory_path), **kwargs)
+            repo_id = kwargs.pop("repo_id", None)
+            if repo_id is None:
+                raise ValueError("`repo_id` must be provided to push the model to the HuggingFace model hub.")
+            return super().push_to_hub(repo_id=repo_id, **kwargs)
 
     @staticmethod
     def _raise_missing_compiled_file_error(missing_files: List[str]):

optimum/rbln/transformers/models/bart/bart_architecture.py

@@ -56,10 +56,7 @@ class BartDecoderWrapper(Seq2SeqDecoderWrapper):
 
 
 class BartForConditionalGeneration(Seq2SeqForConditionalGeneration):
-    has_rescaling = False
-
-    def __post_init__(self):
-        self.scaling = self.config.d_model**-0.5
+    pass
 
 
 class BartDecoder(Seq2SeqDecoder):

optimum/rbln/transformers/models/clip/configuration_clip.py

@@ -43,7 +43,15 @@ class RBLNCLIPTextModelWithProjectionConfig(RBLNCLIPTextModelConfig):
 
 
 class RBLNCLIPVisionModelConfig(RBLNModelConfig):
-    def __init__(self, batch_size: Optional[int] = None, image_size: Optional[int] = None, **kwargs: Dict[str, Any]):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        image_size: Optional[int] = None,
+        interpolate_pos_encoding: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        **kwargs: Dict[str, Any],
+    ):
         """
         Args:
             batch_size (Optional[int]): The batch size for image processing. Defaults to 1.
@@ -60,6 +68,9 @@ class RBLNCLIPVisionModelConfig(RBLNModelConfig):
             raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
 
         self.image_size = image_size
+        self.interpolate_pos_encoding = interpolate_pos_encoding or False
+        self.output_hidden_states = output_hidden_states
+        self.output_attentions = output_attentions
 
     @property
     def image_width(self):

optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import torch
 from transformers import CLIPTextConfig, CLIPTextModel, CLIPVisionConfig, CLIPVisionModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.models.clip.modeling_clip import CLIPTextModelOutput, CLIPVisionModelOutput
 
 from ....configuration_utils import RBLNCompileConfig
@@ -111,12 +112,27 @@ class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
 
 
 class _VisionEncoder(torch.nn.Module):
-    def __init__(self, enc: CLIPVisionModel):
+    def __init__(
+        self,
+        enc: CLIPVisionModel,
+        interpolate_pos_encoding: bool,
+        output_hidden_states: bool,
+        output_attentions: bool,
+    ):
         super().__init__()
         self.enc = enc
+        self.interpolate_pos_encoding = interpolate_pos_encoding
+        self.output_hidden_states = output_hidden_states
+        self.output_attentions = output_attentions
 
     def forward(self, inp):
-        enc_out = self.enc(inp, output_hidden_states=True, return_dict=False)
+        enc_out = self.enc(
+            inp,
+            output_hidden_states=self.output_hidden_states,
+            interpolate_pos_encoding=self.interpolate_pos_encoding,
+            output_attentions=self.output_attentions,
+            return_dict=False,
+        )
         return enc_out
 
 
@@ -130,7 +146,12 @@ class RBLNCLIPVisionModel(RBLNModel):
 
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
-        return _VisionEncoder(model).eval()
+        wrapper_cfg = {
+            "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
+            "output_hidden_states": rbln_config.output_hidden_states,
+            "output_attentions": rbln_config.output_attentions,
+        }
+        return _VisionEncoder(model, **wrapper_cfg).eval()
 
     @classmethod
     def update_rbln_config_using_pipe(
@@ -155,6 +176,12 @@ class RBLNCLIPVisionModel(RBLNModel):
         if rbln_config.image_size is None:
             raise ValueError("`rbln_image_size` should be specified!")
 
+        if rbln_config.output_attentions is None:
+            rbln_config.output_attentions = getattr(model_config, "output_attentions", False)
+
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
+
         rbln_compile_config = RBLNCompileConfig(
             input_info=[
                 (
@@ -176,27 +203,76 @@ class RBLNCLIPVisionModel(RBLNModel):
     def forward(
         self,
         pixel_values: Optional[torch.FloatTensor] = None,
-        return_dict: bool = None,
+        return_dict: bool = True,
+        output_attentions: bool = None,
+        output_hidden_states: bool = None,
+        interpolate_pos_encoding: bool = False,
         **kwargs,
-    ) -> Union[Tuple, CLIPVisionModelOutput]:
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
         if len(kwargs) > 0 and any(value is not None for value in kwargs.values()):
             logger.warning(
                 f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__.__name__}."
             )
+
+        output_attentions = output_attentions if output_attentions is not None else self.rbln_config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        if output_attentions != self.rbln_config.output_attentions:
+            raise ValueError(
+                f"Variable output_attentions {output_attentions} is not equal to rbln_config.output_attentions {self.rbln_config.output_attentions} "
+                f"Please compile again with the correct argument."
+            )
+
+        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."
+            )
+
+        if interpolate_pos_encoding != self.rbln_config.interpolate_pos_encoding:
+            raise ValueError(
+                f"Variable interpolate_pos_encoding {interpolate_pos_encoding} is not equal to rbln_config.interpolate_pos_encoding {self.rbln_config.interpolate_pos_encoding} "
+                f"Please compile again with the correct argument."
+            )
+
         output = super().forward(pixel_values, return_dict=return_dict)
         return output
 
     def _prepare_output(self, output, return_dict):
         # Prepare model output based on return_dict flag.
         # This method can be overridden by subclasses to provide task-specific output handling.
+        last_hidden_state = output.pop(0)
+        pooler_output = output.pop(0)
+        vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+
+        if self.rbln_config.output_hidden_states:
+            hidden_states = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers + 1):
+                hidden_states += (output.pop(0),)
+        else:
+            hidden_states = None
+
+        if self.rbln_config.output_attentions:
+            attentions = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers):
+                attentions += (output.pop(0),)
+        else:
+            attentions = None
 
         if not return_dict:
-            return (output,) if not isinstance(output, (tuple, list)) else output
+            return tuple(
+                item for item in (last_hidden_state, pooler_output, hidden_states, attentions) if item is not None
+            )
         else:
-            return CLIPVisionModelOutput(
-                image_embeds=output[0],
-                last_hidden_state=output[1],
-                hidden_states=output[2:],
+            return BaseModelOutputWithPooling(
+                last_hidden_state=last_hidden_state,
+                pooler_output=pooler_output,
+                hidden_states=hidden_states,
+                attentions=attentions,
             )
 
 
@@ -208,21 +284,40 @@ class RBLNCLIPVisionModelWithProjection(RBLNCLIPVisionModel):
     multimodal embedding alignment tasks.
     """
 
-    def forward(
-        self,
-        pixel_values: Optional[torch.FloatTensor] = None,
-        **kwargs,
-    ) -> Union[Tuple, CLIPVisionModelOutput]:
-        if len(kwargs) > 0 and any(kwargs.values()):
-            logger.warning(f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__}.")
-
-        output = super().forward(pixel_values)
-        image_embeds = output[0]
-        last_hidden_state = output[1]
-        hidden_states = output[2:]
-
-        return CLIPVisionModelOutput(
-            image_embeds=image_embeds,
-            last_hidden_state=last_hidden_state,
-            hidden_states=hidden_states,
-        )
+    def _prepare_output(self, output, return_dict):
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
+
+        image_embeds = output.pop(0) if isinstance(output, (tuple, list)) else output
+        last_hidden_state = output.pop(0)
+
+        vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+
+        if self.rbln_config.output_hidden_states:
+            hidden_states = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers + 1):
+                hidden_states += (output.pop(0),)
+        else:
+            hidden_states = None
+
+        if self.rbln_config.output_attentions:
+            attentions = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers):
+                attentions += (output.pop(0),)
+        else:
+            attentions = None
+
+        if not return_dict:
+            return tuple(
+                item for item in (image_embeds, last_hidden_state, hidden_states, attentions) if item is not None
+            )
+
+        else:
+            return CLIPVisionModelOutput(
+                image_embeds=image_embeds,
+                last_hidden_state=last_hidden_state,
+                hidden_states=hidden_states,
+                attentions=attentions,
+            )

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

@@ -303,6 +303,8 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             position_embed = (
                 position_embed[:, :, :, attention_mask.bool(), :] if attention_mask is not None else position_embed
             )
+        if token_type_ids is not None:
+            token_type_ids = token_type_ids[:, attention_mask.bool()] if attention_mask is not None else token_type_ids
 
         query_length = inputs.shape[1]
         if query_length > self.rbln_config.max_seq_len:
@@ -352,8 +354,11 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             if position_embed is not None:
                 position_embed = torch.nn.functional.pad(position_embed, (0, 0, 0, padding_size))
 
+            if token_type_ids is not None:
+                token_type_ids = torch.nn.functional.pad(token_type_ids, (0, padding_size), value=-1)
+
         # Overwrite position_ids and padded_cache_lengths
-        position_ids = None
+        position_ids = cache_position.clone()
         padded_cache_lengths = 0
 
         return (
@@ -365,6 +370,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             position_embed,
             padded_cache_lengths,
             query_length,
+            token_type_ids,
         )
 
     def prefill_forward(
@@ -393,6 +399,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             position_embed,
             padded_cache_lengths,
             query_length,
+            token_type_ids,
         ) = self._prepare_prefill_inputs(
             inputs, cache_position, attention_mask, position_embed, token_type_ids=token_type_ids
         )

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

@@ -23,22 +23,21 @@ from ..siglip.configuration_siglip import RBLNSiglipVisionModelConfig
 class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     def __init__(
         self,
-        prefill_chunk_size: Optional[int] = None,
         use_position_ids: Optional[bool] = None,
         use_attention_mask: Optional[bool] = None,
+        image_prefill_chunk_size: Optional[int] = None,
         **kwargs: Dict[str, Any],
     ):
         # use_attention_mask and use_position_ids are always True for Gemma3
         use_attention_mask = use_attention_mask or True
         use_position_ids = use_position_ids or True
-        prefill_chunk_size = prefill_chunk_size or 256
 
         super().__init__(
-            prefill_chunk_size=prefill_chunk_size,
             use_attention_mask=use_attention_mask,
             use_position_ids=use_position_ids,
             **kwargs,
         )
+        self.image_prefill_chunk_size = image_prefill_chunk_size
 
         npu = self.npu or rebel.get_npu_name()
         if npu == "RBLN-CA02":

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

@@ -31,15 +31,11 @@ from transformers.models.gemma3.modeling_gemma3 import Gemma3TextScaledWordEmbed
 
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
-from ....utils.logging import get_logger
 from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyOutput, RBLNRuntimeModel
 from .configuration_gemma3 import RBLNGemma3ForCausalLMConfig
 from .gemma3_architecture import Gemma3ForCausalLMWrapper
 
 
-logger = get_logger()
-
-
 if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, Gemma3ForConditionalGeneration
 
@@ -320,194 +316,28 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
         self.prefill = self.runtime if self.phase == "prefill" else None  # FIXME
         self.decode = self.runtime if self.phase == "decode" else None
 
-    def pad_for_chunked_images(
-        self,
-        inputs: torch.Tensor,
-        attention_mask: torch.Tensor,
-        position_ids: torch.Tensor,
-        token_type_ids: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, torch.Tensor]:
-        """
-        Pads inputs, attention_mask, and position_ids so image token groups (256 tokens with token_type_ids == 1)
-        start at multiples of prefill_chunk_size (256). Returns padded tensors and total padded length.
-
-        Args:
-            inputs: (1, seq_len, hidden_size) tensor.
-            attention_mask: (1, seq_len) tensor, 1 for valid, 0 for masked.
-            position_ids: (1, seq_len) tensor for RoPE.
-            token_type_ids: (1, seq_len) tensor, 0 for text, 1 for image.
-
-        Returns:
-            (inputs_padded, attention_mask_padded, position_ids_padded, padded_len, token_type_ids_padded).
-        """
-
-        if token_type_ids is None:
-            return inputs, attention_mask, position_ids, 0, torch.zeros(inputs.shape[:2], dtype=torch.long)
-
-        seq_len = inputs.shape[1]
-
-        # Find image start positions
-        image_starts = [
-            s
-            for s in range(seq_len - self.rbln_config.prefill_chunk_size + 1)
-            if torch.all(token_type_ids[:, s : s + self.rbln_config.prefill_chunk_size] == 1)
-        ]
-
-        # Initialize padded tensors
-        padded_input_len = seq_len
-        for image_start in image_starts:
-            pad_needed = (
-                self.rbln_config.prefill_chunk_size
-                - (image_start + padded_input_len - seq_len) % self.rbln_config.prefill_chunk_size
-            ) % self.rbln_config.prefill_chunk_size
-            padded_input_len += pad_needed
-        total_padding = padded_input_len - seq_len
-
-        if inputs.dim() == 3:
-            inputs_padded = torch.zeros(1, padded_input_len, inputs.shape[2], dtype=inputs.dtype)
-        else:
-            inputs_padded = torch.zeros(1, padded_input_len, dtype=inputs.dtype)
-        attention_mask_padded = torch.zeros(1, padded_input_len, dtype=attention_mask.dtype)
-        position_ids_padded = torch.zeros(1, padded_input_len, dtype=position_ids.dtype)
-        token_type_ids_padded = torch.zeros(1, padded_input_len, dtype=token_type_ids.dtype)
-
-        # Fill padded tensors
-        dest_pos = 0
-        src_pos = 0
-        last_pos_id = -1
-        for image_start in image_starts + [seq_len]:
-            # Text segment
-            if src_pos < image_start:
-                length = image_start - src_pos
-                inputs_padded[:, dest_pos : dest_pos + length] = inputs[:, src_pos:image_start]
-                attention_mask_padded[:, dest_pos : dest_pos + length] = attention_mask[:, src_pos:image_start]
-                position_ids_padded[:, dest_pos : dest_pos + length] = position_ids[:, src_pos:image_start]
-                token_type_ids_padded[:, dest_pos : dest_pos + length] = token_type_ids[:, src_pos:image_start]
-                dest_pos += length
-                last_pos_id = position_ids[0, image_start - 1].item()
-                src_pos = image_start
-
-            # Padding
-            pad_needed = (
-                self.rbln_config.prefill_chunk_size - dest_pos % self.rbln_config.prefill_chunk_size
-            ) % self.rbln_config.prefill_chunk_size
-            if pad_needed and dest_pos < padded_input_len:
-                position_ids_padded[:, dest_pos : dest_pos + pad_needed] = torch.arange(
-                    last_pos_id + 1, last_pos_id + pad_needed + 1, dtype=position_ids.dtype
-                ).unsqueeze(0)
-                dest_pos += pad_needed
-
-            # Image segment
-            if src_pos < seq_len and src_pos == image_start:
-                inputs_padded[:, dest_pos : dest_pos + self.rbln_config.prefill_chunk_size] = inputs[
-                    :, src_pos : src_pos + self.rbln_config.prefill_chunk_size
-                ]
-                attention_mask_padded[:, dest_pos : dest_pos + self.rbln_config.prefill_chunk_size] = attention_mask[
-                    :, src_pos : src_pos + self.rbln_config.prefill_chunk_size
-                ]
-                position_ids_padded[:, dest_pos : dest_pos + self.rbln_config.prefill_chunk_size] = position_ids[
-                    :, src_pos : src_pos + self.rbln_config.prefill_chunk_size
-                ]
-                token_type_ids_padded[:, dest_pos : dest_pos + self.rbln_config.prefill_chunk_size] = token_type_ids[
-                    :, src_pos : src_pos + self.rbln_config.prefill_chunk_size
-                ]
-                dest_pos += self.rbln_config.prefill_chunk_size
-                src_pos += self.rbln_config.prefill_chunk_size
-                last_pos_id = position_ids[0, image_start + self.rbln_config.prefill_chunk_size - 1].item()
-
-        return inputs_padded, attention_mask_padded, position_ids_padded, total_padding, token_type_ids_padded
-
-    def _prepare_prefill_inputs(
-        self,
-        inputs: torch.Tensor,
-        cache_position: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_embed: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-    ):
-        """
-        Prepare inputs for prefill phase.
-        """
-        # Handle continuous batching in a compiled graph by extracting valid inputs
-        # If an attention mask is provided, select only the valid (non-masked) inputs
-        inputs = inputs[:, attention_mask.bool()] if attention_mask is not None else inputs
-        token_type_ids = (
-            token_type_ids[:, attention_mask.bool()]
-            if attention_mask is not None and token_type_ids is not None
-            else token_type_ids
-        )
-
-        if position_embed is not None:
-            position_embed = (
-                position_embed[:, :, :, attention_mask.bool(), :] if attention_mask is not None else position_embed
-            )
-
-        seq_len = inputs.shape[1]
-        # Initialize attention mask for chunked processing
-        if self.rbln_config.use_attention_mask:
-            chunked_attention_mask = (
-                torch.ones(1, seq_len, dtype=torch.float32)
-                if self.rbln_config.use_position_ids
-                else torch.zeros(
-                    1, 1, self.rbln_config.prefill_chunk_size, self.rbln_config.max_seq_len, dtype=torch.float32
-                )
-            )
-        else:
-            chunked_attention_mask = None
-
-        # Buffer for storing output logits
-        out_buffers = [
-            torch.empty(
-                size=self.output_size,
-                dtype=torch.float32,
-                device="cpu",
-            )
-        ]
-
-        inputs, chunked_attention_mask, position_ids, padded_cache_lengths, token_type_ids_padded = (
-            self.pad_for_chunked_images(inputs, chunked_attention_mask, cache_position, token_type_ids)
-        )
-
-        query_length = inputs.shape[1]
-        if query_length > self.rbln_config.max_seq_len:
-            raise ValueError(
-                f"Input length ({query_length}) exceeds the maximum allowed sequence length ({self.rbln_config.max_seq_len})."
-            )
-
-        # Align attention_mask to compiled shape
-        if self.rbln_config.use_position_ids:
-            chunked_attention_mask = torch.nn.functional.pad(
-                chunked_attention_mask, (0, self.rbln_config.max_seq_len - query_length)
-            )
-
-        # Pad input and cache_position if the last chunk is smaller than `prefill_chunk_size`
-        padding_size = 0
-        if query_length % self.rbln_config.prefill_chunk_size != 0:
-            padding_size = (self.rbln_config.prefill_chunk_size - query_length) % self.rbln_config.prefill_chunk_size
-            # inputs_embeds
-            if inputs.dim() == 3:
-                inputs = torch.nn.functional.pad(inputs, (0, 0, 0, padding_size))
-            # inputs_ids
-            else:
-                inputs = torch.nn.functional.pad(inputs, (0, padding_size))
-
-            position_ids = torch.cat(
-                [
-                    position_ids,
-                    torch.arange(
-                        query_length,
-                        query_length + padding_size,
-                        dtype=torch.int32,
-                    ).unsqueeze(0),
-                ],
-                dim=-1,
-            )
-            token_type_ids_padded = torch.nn.functional.pad(token_type_ids_padded, (0, padding_size))
+    def _prepare_prefill_inputs(self, *args, **kwargs):
+        (
+            inputs,
+            cache_position,
+            chunked_attention_mask,
+            out_buffers,
+            position_ids,
+            position_embed,
+            padded_cache_lengths,
+            query_length,
+            token_type_ids,
+        ) = super()._prepare_prefill_inputs(*args, **kwargs)
 
-            if position_embed is not None:
-                position_embed = torch.nn.functional.pad(position_embed, (0, 0, 0, padding_size))
+        # chunked_attention_mask shape
+        chunked_attention_mask = torch.zeros(1, chunked_attention_mask.shape[-1], dtype=torch.float32)
 
-        cache_position = torch.arange(0, query_length + padding_size, dtype=torch.int32).unsqueeze(0)
+        # 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)
 
         return (
             inputs,
@@ -518,7 +348,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
             position_embed,
             padded_cache_lengths,
             query_length,
-            token_type_ids_padded,
+            token_type_ids,
         )
 
     def prefill_forward(
@@ -541,65 +371,73 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
         (
             inputs,
             cache_position,
-            padded_attention_mask,
+            chunked_attention_mask,
             out_buffers,
             position_ids,
             position_embed,
             padded_cache_lengths,
             query_length,
-            token_type_ids_padded,
+            token_type_ids,
         ) = self._prepare_prefill_inputs(
             inputs, cache_position, attention_mask, position_embed, token_type_ids=token_type_ids
         )
-        if not is_external_block_tables:
-            local_block_tables = torch.tensor([batch_idx], dtype=torch.int16)
-            self.dec_attn_mask[batch_idx : batch_idx + 1] = padded_attention_mask[:1]
-
-        if self.rbln_config.use_attention_mask and self.rbln_config.use_position_ids:
-            chunked_attention_mask = torch.zeros(1, self.rbln_config.max_seq_len, dtype=torch.float32)
-
-        # Process input in chunks of size `prefill_chunk_size`
-        for step in range(0, query_length, self.rbln_config.prefill_chunk_size):
-            # Extract the current chunk of inputs and cache positions
-            input_chunk = inputs[:, step : step + self.rbln_config.prefill_chunk_size]
-            cache_pos_chunk = cache_position[:, step : step + self.rbln_config.prefill_chunk_size]
-            position_ids_chunk = (
-                position_ids[:, step : step + self.rbln_config.prefill_chunk_size]
-                if position_ids is not None
-                else None
+
+        step = 0
+        while step < query_length:
+            # Check if the prefill chunk is an image prefill
+            is_image_prefill = torch.all(
+                token_type_ids[:, step : step + self.rbln_config.image_prefill_chunk_size] == 1
+            )
+            prefill_chunk_size = (
+                self.rbln_config.image_prefill_chunk_size if is_image_prefill else self.rbln_config.prefill_chunk_size
             )
 
-            if self.rbln_config.use_attention_mask:
-                if self.rbln_config.use_position_ids:
-                    chunked_attention_mask[0, step : step + self.rbln_config.prefill_chunk_size] = (
-                        padded_attention_mask[0, step : step + self.rbln_config.prefill_chunk_size]
-                    )
-
-            # Define query position
-            query_position = (
-                torch.sum(
-                    chunked_attention_mask[0][step : step + self.rbln_config.prefill_chunk_size],
-                    dim=-1,
-                    dtype=torch.int16,
-                ).squeeze(0)
-                - 1
+            # 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
             )
-            if token_type_ids_padded[:, step] == 1:
-                if torch.any(token_type_ids_padded[:, step : step + self.rbln_config.prefill_chunk_size] == 0):
-                    raise ValueError("All tokens of image_prefill should be the same image.")
-                else:
-                    logits = self.image_prefill(
-                        input_chunk,
-                        cache_pos_chunk,
-                        block_tables,
-                        local_block_tables,
-                        query_position,
-                        chunked_attention_mask,
-                        position_ids_chunk,
-                        out=out_buffers,
-                    )
+
+            # Check if the prefill chunk crosses a block boundary, requiring padding to align with block boundaries
+            is_cross_block_boundary = (
+                step // self.rbln_config.kvcache_block_size
+                != (step + prefill_chunk_size) // self.rbln_config.kvcache_block_size
+            )
+
+            # Check if the prefill chunk is the last chunk
+            is_last_chunk = step + prefill_chunk_size >= query_length
+
+            if is_cross_block_boundary:
+                padding_size = prefill_chunk_size - (step + prefill_chunk_size) % self.rbln_config.kvcache_block_size
+                padded_cache_lengths += padding_size
+
+            # if text_prefill end with image_tokens, we only treat the text part.
+            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
+            if is_last_chunk:
+                num_processed_tokens = query_length - step
+
+            input_chunk = inputs[:, step : step + prefill_chunk_size]
+            cache_pos_chunk = cache_position[:, step : step + prefill_chunk_size].clone() + padded_cache_lengths
+            position_ids_chunk = position_ids[:, step : step + prefill_chunk_size].clone()
+            chunked_attention_mask[
+                :, step + padded_cache_lengths : step + num_processed_tokens + padded_cache_lengths
+            ] = 1
+            query_position = torch.tensor(num_processed_tokens - 1, dtype=torch.int16)
+
+            if is_image_prefill:
+                logits = self.image_prefill(
+                    input_chunk,
+                    cache_pos_chunk,
+                    block_tables,
+                    local_block_tables,
+                    query_position,
+                    chunked_attention_mask,
+                    position_ids_chunk,
+                    out=out_buffers,
+                )
             else:
-                # Forward pass for the current chunk
                 logits = self.prefill(
                     input_chunk,
                     cache_pos_chunk,
@@ -611,6 +449,11 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
                     out=out_buffers,
                 )
 
+            step += num_processed_tokens
+
+        if not is_external_block_tables:
+            self.dec_attn_mask[batch_idx : batch_idx + 1] = chunked_attention_mask
+
         return RBLNGemma3ForCausalLMOutput(
             logits=logits, padded_cache_lengths=padded_cache_lengths, attention_mask=chunked_attention_mask
         )
@@ -757,13 +600,14 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     @classmethod
     def _update_submodule_config(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
-        if rbln_config.prefill_chunk_size is None:
-            rbln_config.prefill_chunk_size = model.config.mm_tokens_per_image
+        if rbln_config.image_prefill_chunk_size is None:
+            rbln_config.image_prefill_chunk_size = model.config.mm_tokens_per_image
 
-        if rbln_config.prefill_chunk_size != model.config.mm_tokens_per_image:
-            logger.warning(
-                f"Prefill chunk size is different from mm_tokens_per_image: {rbln_config.prefill_chunk_size} != {model.config.mm_tokens_per_image}"
+        if rbln_config.image_prefill_chunk_size != model.config.mm_tokens_per_image:
+            raise ValueError(
+                f"Image prefill chunk size is different from mm_tokens_per_image: {rbln_config.image_prefill_chunk_size} != {model.config.mm_tokens_per_image}"
             )
+
         return rbln_config
 
     @classmethod
@@ -777,14 +621,22 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         # Update rbln_config with super class
         rbln_config = super()._update_rbln_config(preprocessors, model, model_config, rbln_config)
 
-        # Assume that prefill compile config is at index 0
-        compile_cfgs = rbln_config.compile_cfgs
+        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=compile_cfgs[0].input_info
+            compiled_model_name="image_prefill", input_info=img_prefill_input_info
         )
         # Insert image_prefill compile config at index 1
-        image_idx = 1
-        compile_cfgs.insert(image_idx, image_prefill_compile_config)
+        compile_cfgs = rbln_config.compile_cfgs
+        compile_cfgs.insert(1, image_prefill_compile_config)
         rbln_config.set_compile_cfgs(compile_cfgs)
 
         return rbln_config
@@ -840,11 +692,14 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         )
 
         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,
-            prefill_example_inputs,
+            image_prefill_example_inputs,
             context,
             rbln_config.quantization,
         )

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

@@ -15,6 +15,11 @@
 from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+from ...models.clip import RBLNCLIPVisionModelConfig
+
+
+logger = get_logger(__name__)
 
 
 class RBLNLlavaNextForConditionalGenerationConfig(RBLNModelConfig):
@@ -50,5 +55,17 @@ 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 = vision_tower
+        self.vision_tower = self.init_submodule_config(
+            RBLNCLIPVisionModelConfig,
+            vision_tower,
+        )
+
+        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.language_model = language_model

optimum/rbln/utils/hub.py

@@ -12,59 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 from pathlib import Path
 from typing import List, Optional, Union
 
-from huggingface_hub import HfApi, HfFolder, hf_hub_download
-
-
-class PushToHubMixin:
-    def push_to_hub(
-        self,
-        save_directory: str,
-        repository_id: str,
-        private: Optional[bool] = None,
-        use_auth_token: Union[bool, str] = True,
-    ) -> str:
-        huggingface_token = _get_huggingface_token(use_auth_token)
-        api = HfApi()
-
-        api.create_repo(
-            token=huggingface_token,
-            repo_id=repository_id,
-            exist_ok=True,
-            private=private,
-        )
-        for path, subdirs, files in os.walk(save_directory):
-            for name in files:
-                local_file_path = os.path.join(path, name)
-                _, hub_file_path = os.path.split(local_file_path)
-                # FIXME: when huggingface_hub fixes the return of upload_file
-                try:
-                    api.upload_file(
-                        token=huggingface_token,
-                        repo_id=f"{repository_id}",
-                        path_or_fileobj=os.path.join(os.getcwd(), local_file_path),
-                        path_in_repo=hub_file_path,
-                    )
-                except KeyError:
-                    pass
-                except NameError:
-                    pass
+from huggingface_hub import HfApi, get_token, hf_hub_download
 
 
 def pull_compiled_model_from_hub(
     model_id: Union[str, Path],
     subfolder: str,
-    use_auth_token: Optional[Union[bool, str]],
+    token: Union[bool, str],
     revision: Optional[str],
     cache_dir: Optional[str],
     force_download: bool,
     local_files_only: bool,
 ) -> Path:
     """Pull model files from the HuggingFace Hub."""
-    huggingface_token = _get_huggingface_token(use_auth_token)
+    huggingface_token = _get_huggingface_token(token)
     repo_files = list(
         map(
             Path,
@@ -87,7 +51,7 @@ def pull_compiled_model_from_hub(
             repo_id=model_id,
             filename=filename,
             subfolder=subfolder,
-            use_auth_token=use_auth_token,
+            token=token,
             revision=revision,
             cache_dir=cache_dir,
             force_download=force_download,
@@ -113,10 +77,7 @@ def validate_files(
         raise FileExistsError(f"Multiple rbln_config.json files found in {location}. This is not expected.")
 
 
-def _get_huggingface_token(use_auth_token: Union[bool, str]) -> str:
-    if isinstance(use_auth_token, str):
-        return use_auth_token
-    elif use_auth_token:
-        return HfFolder.get_token()
-    else:
-        raise ValueError("`use_auth_token` must be provided to interact with the HuggingFace Hub.")
+def _get_huggingface_token(token: Union[bool, str]) -> str:
+    if isinstance(token, str):
+        return token
+    return get_token()

{optimum_rbln-0.8.2a2.dist-info → optimum_rbln-0.8.2a4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.2a2
+Version: 0.8.2a4
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai

{optimum_rbln-0.8.2a2.dist-info → optimum_rbln-0.8.2a4.dist-info}/RECORD

@@ -1,10 +1,10 @@
 optimum/rbln/__init__.py,sha256=Pl_On5dBWYbmJk6776cc0eU49oc0pebdCvgb2l4hQRA,15564
-optimum/rbln/__version__.py,sha256=0WA5HC48CQUocr8qaZZFi2jxgajnt6fwHhbsacDdi4g,519
+optimum/rbln/__version__.py,sha256=al82X96cfOX7yNTvVBaPgx2coI4eq0-fqG5LYKf53zE,519
 optimum/rbln/configuration_utils.py,sha256=tnR9XZfzZ9oHc7nU_kX33oo3qFFyicICSVQSujYPiOM,32911
-optimum/rbln/modeling.py,sha256=wZ_Cw83tgPDDJAB84_iLEB_H7ycL-rrlTsCzrHk-GWo,14293
-optimum/rbln/modeling_base.py,sha256=QpNkU_Do__JKmnHjaPzv47OhQwgGfVohisip1jqXa7A,23871
+optimum/rbln/modeling.py,sha256=gww-H-q16_mGw2qGnFwOjEj3J9yMjBKnRTKlnCkVlx8,14315
+optimum/rbln/modeling_base.py,sha256=AShxAt3KIOcCqfyF4U83dIrKwoj4p2Kxtc1ns_9-ltU,24154
 optimum/rbln/diffusers/__init__.py,sha256=_3FoA0uxCdFd6YK9PMaptFmR9XvdB_PcvYR1MkbGpN8,6957
-optimum/rbln/diffusers/modeling_diffusers.py,sha256=0Pwqw2PbLSQ6FXsXzbx44Jfz3Vm9hfSrfVRH5zcN0yk,19885
+optimum/rbln/diffusers/modeling_diffusers.py,sha256=gnCsDZ6WD6N7yOfcZU8pK5rtwyh4zwSCtmR6Wa4zRG8,19913
 optimum/rbln/diffusers/configurations/__init__.py,sha256=vMRnPY4s-Uju43xP038D2EA18X_mhy2YfsZVpSU-VoA,1322
 optimum/rbln/diffusers/configurations/models/__init__.py,sha256=7q95gtgDzCeIBogGw8SLQoHT4Wch7vpLJVF2UQovuoo,567
 optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl.py,sha256=ySetuNq6koleFIZ542zZLTzEEyl_CTul9l12ufWlQ_Y,3218
@@ -83,7 +83,7 @@ optimum/rbln/transformers/models/auto/__init__.py,sha256=34Xghf1ogG4u-jhBMlj134n
 optimum/rbln/transformers/models/auto/auto_factory.py,sha256=1CA52xV2dS1Uzumcgqe4zobdpoi-Xt2oNjP3uLFtm08,8020
 optimum/rbln/transformers/models/auto/modeling_auto.py,sha256=Ex2ARRRIt3LtKhazr0UWy67R6WFAwSKRoi5n6B8TvoI,4213
 optimum/rbln/transformers/models/bart/__init__.py,sha256=fVo-gZEmJ0yxkIxEX6ciuRAGgXNyuvaXE2s88bhbjAE,830
-optimum/rbln/transformers/models/bart/bart_architecture.py,sha256=oxb3l0lw0qrnbeQxT-R3c0OH4IEp7UzEp2Iz_jhZe_Q,6373
+optimum/rbln/transformers/models/bart/bart_architecture.py,sha256=BjVXyk7Ron4ga_Ac_cvQiRvrdJ9LULEHJcKK0LfCqas,6277
 optimum/rbln/transformers/models/bart/configuration_bart.py,sha256=L466cRy4egEbu_lTb_v1sK6t3Lbi4Bra6Jjkj3FjdQw,1315
 optimum/rbln/transformers/models/bart/modeling_bart.py,sha256=H4MmQZbofb9kJq5WKqoFVjmj3HVtgns3t2F3QdSU-QQ,2337
 optimum/rbln/transformers/models/bert/__init__.py,sha256=86FuGRBLw315_Roa9D5OUx6Ku2PM0DqSPZ-YSqbF-io,806
@@ -93,8 +93,8 @@ optimum/rbln/transformers/models/blip_2/__init__.py,sha256=L01gPXcUCa8Vg-bcng20v
 optimum/rbln/transformers/models/blip_2/configuration_blip_2.py,sha256=ke75GqPU139dNOY1nm6QE661LepbD_0V9Bx1QbtHhKA,3210
 optimum/rbln/transformers/models/blip_2/modeling_blip_2.py,sha256=2sIVGrIn1B2nUZ8hw1sgW3VbJ2vxrlBRN37GgDiw0GU,16191
 optimum/rbln/transformers/models/clip/__init__.py,sha256=TLeXDqcFK6M6v9x7Xr64kBbqGu3hFHM7p754dQ8UVQc,938
-optimum/rbln/transformers/models/clip/configuration_clip.py,sha256=D7CIWpbMhXUrGv-CnhxRtSS3vAYb427-w7zSkfuJHEU,3455
-optimum/rbln/transformers/models/clip/modeling_clip.py,sha256=QbYrt7pUWNal-p93fxmuKrHa2CPlCaq8F16qOfMAst0,8090
+optimum/rbln/transformers/models/clip/configuration_clip.py,sha256=HeKqLtTKrpQbx1Gq5GRSSLvjjgXHYD0rWIUhQOZDQCY,3836
+optimum/rbln/transformers/models/clip/modeling_clip.py,sha256=knK7gINAluSHcWvg3zaByb3XRLNmSEGw2NcsOGHnIow,12364
 optimum/rbln/transformers/models/colpali/__init__.py,sha256=n3rueXT_oC0N8myoZiic0YkVK24CW5hZBPa-0L8so6Y,119
 optimum/rbln/transformers/models/colpali/colpali_architecture.py,sha256=bWG7TehWRZkTh2y6mGkpd85_onWAyiyKdaQC9TFsy3E,8065
 optimum/rbln/transformers/models/colpali/configuration_colpali.py,sha256=ieY-tuyDPObFUIJ5sfpcfuCsJ_HTAizN7ZGqirqeFRU,2636
@@ -102,7 +102,7 @@ optimum/rbln/transformers/models/colpali/modeling_colpali.py,sha256=jzvJCBrrCXSp
 optimum/rbln/transformers/models/decoderonly/__init__.py,sha256=vQYZDDdoddwA7yKc5zzrq2Zs9sax-0p8rNF_aYfF4bk,1006
 optimum/rbln/transformers/models/decoderonly/configuration_decoderonly.py,sha256=cakn8RGo8gS3nmXdEqOfC2xUBOMGInROgLEbCOoLFR0,13398
 optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py,sha256=HrI12t9X9wV_-AZuTBSs-W7c5yVUkvd0secWlI72x2A,46325
-optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py,sha256=qe0icOCkPY1k1Db9gP5fXu-kHIjZV3M_tOjrs-EeULQ,53701
+optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py,sha256=lmQtd50LS6xGfLvSCSv1u7gwPk0ggjb3wSdzqIHRvLM,54080
 optimum/rbln/transformers/models/distilbert/__init__.py,sha256=zXL78SOEORTnUN_wrdoaDaYpntG8lcFHvPobM6jC0CI,841
 optimum/rbln/transformers/models/distilbert/configuration_distilbert.py,sha256=O3BW9JjyYk9PLyiofvOKEgTdMZ_jpIuPfot281pSsyg,984
 optimum/rbln/transformers/models/distilbert/modeling_distilbert.py,sha256=LUh6zYGa8AR3Yxaj3gtyJRc-czBN3qnHTc-JTAhuqY0,1099
@@ -118,9 +118,9 @@ optimum/rbln/transformers/models/gemma/configuration_gemma.py,sha256=3hAxl7LL9vF
 optimum/rbln/transformers/models/gemma/gemma_architecture.py,sha256=2Ivay8NTSHmQAqXFh9JvG6Ja5rMThcRAjYPzyipcRI8,956
 optimum/rbln/transformers/models/gemma/modeling_gemma.py,sha256=Ojvum34EhDHWfMB4D6S1BrwoTNwuBSZuBzwdnAgvq38,3095
 optimum/rbln/transformers/models/gemma3/__init__.py,sha256=6rugk3615SEt4lh7gduo_J9VyGiSReuEIvL0Uno0eaI,790
-optimum/rbln/transformers/models/gemma3/configuration_gemma3.py,sha256=eupMGTHJGJNNrAZ3GE6M6GQBAQzBb7KFJvalyDmbM-A,3063
+optimum/rbln/transformers/models/gemma3/configuration_gemma3.py,sha256=LwzlThcIXlpK1PdvgaIobp3uHQgaDCiKngi1XR9hDvU,3028
 optimum/rbln/transformers/models/gemma3/gemma3_architecture.py,sha256=fpLDAXCe5paWVsfc0tL59JkRQMRF-WNgIzOIb_QpSLU,6191
-optimum/rbln/transformers/models/gemma3/modeling_gemma3.py,sha256=iVzH7dm4ZO1raw6mznhmTkCY4UYbTzI5lm-8X7I1vN0,38528
+optimum/rbln/transformers/models/gemma3/modeling_gemma3.py,sha256=Hbh_Cfz0j_so8zCLP3ExO2lrCkyJqVnHI2Vqfpglfn4,31895
 optimum/rbln/transformers/models/gpt2/__init__.py,sha256=socBMIBZSiLbrVN12rQ4nL9gFeT0axMgz6SWaCaD4Ac,704
 optimum/rbln/transformers/models/gpt2/configuration_gpt2.py,sha256=9sS6-EGapmow3rG9ViejK9qwrqy_X86VBxQ7u9x0Yqk,923
 optimum/rbln/transformers/models/gpt2/gpt2_architecture.py,sha256=O7hBiaFJrpLSswGwW83cX9S9Q2wKRBDrpAqOgOS7zQg,2733
@@ -133,7 +133,7 @@ optimum/rbln/transformers/models/llama/configuration_llama.py,sha256=X6SXtRXGBC8
 optimum/rbln/transformers/models/llama/llama_architecture.py,sha256=S7MCPfyjG5eUqgaS-QNBB0ApUD6wnb5fR0RHq7k7-pA,728
 optimum/rbln/transformers/models/llama/modeling_llama.py,sha256=G91Yd8BfhiD8GT5bF46ZIHuc-ajGC-PO-mOQN3BhE1A,3122
 optimum/rbln/transformers/models/llava_next/__init__.py,sha256=kDXKr7wMkp1XqE__DER2B8kQF_NYMxhzsQS5ytGg56I,752
-optimum/rbln/transformers/models/llava_next/configuration_llava_next.py,sha256=b3roSXoAApUXk8dI4X__f3mWUnHcZDI1GAY34ckyYc4,2170
+optimum/rbln/transformers/models/llava_next/configuration_llava_next.py,sha256=U6_DQoaXugN2Bc4ntUb7WkelbNmw1L4VbgqsMRVuuE4,2776
 optimum/rbln/transformers/models/llava_next/modeling_llava_next.py,sha256=paYtCk58--FSZp8xjVrfZAxkJxO02X-jxaVPqL-l7ZU,27421
 optimum/rbln/transformers/models/midm/__init__.py,sha256=IC3FETwgYinbp3wDj7tp4zIHJhbqM-c6GfTRdYcMNj8,913
 optimum/rbln/transformers/models/midm/configuration_midm.py,sha256=DxhcSJlApxfi00XxYmSkKZ6bY9vfLXT0zh-oMKkZot0,1365
@@ -202,14 +202,14 @@ optimum/rbln/transformers/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm
 optimum/rbln/transformers/utils/rbln_quantization.py,sha256=um0N2ZruU_lNibo6rlzgwHAI2_8QOrYE7W7pA1qfXKM,9396
 optimum/rbln/utils/__init__.py,sha256=ieDBT2VFTt2E0M4v_POLBpuGW9LxSydpb_DuPd6PQqc,712
 optimum/rbln/utils/decorator_utils.py,sha256=xu-TrsNi33SRC2a7DBsyoo6-pEQxWKZPZSmM9QlDe2Y,3745
-optimum/rbln/utils/hub.py,sha256=Z_R9Ic9VAew8bUmlaAlxZf5JGMDBivHvvFRI557pILY,4196
+optimum/rbln/utils/hub.py,sha256=FPBGslHJAMeyfBID3viLmh51xJzcR29xWtYtMN8y2CI,2765
 optimum/rbln/utils/import_utils.py,sha256=fpOERIIxXm-cDYGn1NN6c7aWDPQYVitPQW2MiyZ9NEY,5471
 optimum/rbln/utils/logging.py,sha256=VKKBmlQSdg6iZCGmAXaWYiW67K84jyp1QJhLQSSjPPE,3453
 optimum/rbln/utils/model_utils.py,sha256=4k5879Kh75m3x_vS4-qOGfqsOiAvc2kdNFFfvsFvz3k,1748
 optimum/rbln/utils/runtime_utils.py,sha256=nIJioiN16nAyAzoArKjsy5ocLUsrr0UEy4f3LNT82SA,7961
 optimum/rbln/utils/save_utils.py,sha256=hG5uOtYmecSXZuGTvCXsTM-SiyZpr5q3InUGCCq_jzQ,3619
 optimum/rbln/utils/submodule.py,sha256=w5mgPgncI740gVKMu3S-69DGNdUSI0bTZxegQGcZ98Y,5011
-optimum_rbln-0.8.2a2.dist-info/METADATA,sha256=ncxtO1-Q-UrBAF4eatV3ZKNIpdpj5nwFNmHdxvJ0c3U,5299
-optimum_rbln-0.8.2a2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-optimum_rbln-0.8.2a2.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-optimum_rbln-0.8.2a2.dist-info/RECORD,,
+optimum_rbln-0.8.2a4.dist-info/METADATA,sha256=4R2Q_7kuyLzvC65FTrT3xcRXzdP-xIAcaNyM9JFMxQg,5299
+optimum_rbln-0.8.2a4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+optimum_rbln-0.8.2a4.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+optimum_rbln-0.8.2a4.dist-info/RECORD,,