optimum-rbln 0.7.3.post2__py3-none-any.whl → 0.7.4__py3-none-any.whl

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

@@ -12,28 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
+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 import CLIPTextConfig, CLIPTextModel, CLIPVisionConfig, CLIPVisionModel
 from transformers.models.clip.modeling_clip import CLIPTextModelOutput, CLIPVisionModelOutput
 
-from ....diffusers.modeling_diffusers import RBLNDiffusionMixin
+from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.logging import get_logger
+from .configuration_clip import RBLNCLIPTextModelConfig, RBLNCLIPVisionModelConfig
 
 
 logger = get_logger(__name__)
 
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, CLIPTextModel
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, CLIPTextModel, PreTrainedModel
+
+    from ....diffusers.modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
 
 
 class _TextEncoder(torch.nn.Module):
@@ -48,53 +44,55 @@ class _TextEncoder(torch.nn.Module):
 
 class RBLNCLIPTextModel(RBLNModel):
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPTextModelConfig) -> torch.nn.Module:
         return _TextEncoder(model).eval()
 
     @classmethod
-    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_config: str
+    ) -> "RBLNDiffusionMixinConfig":
         return rbln_config
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
-        model_config: "CLIPTextConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-        rbln_batch_size: Optional[int] = None,
-    ) -> RBLNConfig:
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        model_config.return_dict = False
-
+        model: Optional["PreTrainedModel"] = None,
+        model_config: "CLIPTextConfig" = None,
+        rbln_config: Optional[RBLNCLIPTextModelConfig] = None,
+    ) -> RBLNCLIPTextModelConfig:
         input_info = [
             (
                 "input_ids",
                 [
-                    rbln_batch_size,
+                    rbln_config.batch_size,
                     model_config.max_position_embeddings,
                 ],
                 "int64",
             ),
         ]
 
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
+        rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
         return rbln_config
 
-    def forward(self, input_ids: "torch.Tensor", **kwargs):
-        text_output = super().forward(input_ids)
-        return CLIPTextModelOutput(
-            text_embeds=text_output[0],
-            last_hidden_state=text_output[1],
-            hidden_states=text_output[2:],
-        )
+    def forward(self, input_ids: torch.LongTensor, return_dict: bool = None, **kwargs) -> torch.FloatTensor:
+        # To ignore using attention_mask, we override forward method.
+        output = super().forward(input_ids, 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.
+        """
+        if not return_dict:
+            return (output,) if not isinstance(output, (tuple, list)) else output
+        else:
+            return CLIPTextModelOutput(
+                text_embeds=output[0],
+                last_hidden_state=output[1],
+                hidden_states=output[2:],
+            )
 
 
 class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
@@ -113,30 +111,30 @@ class _VisionEncoder(torch.nn.Module):
 
 class RBLNCLIPVisionModel(RBLNModel):
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
         return _VisionEncoder(model).eval()
 
     @classmethod
-    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
         return rbln_config
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
-        model_config: "CLIPVisionConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_batch_size = rbln_kwargs.get("batch_size", 1)
-        rbln_image_size = rbln_kwargs.get("image_size", None)
-
-        if rbln_image_size is None:
-            rbln_image_size = getattr(model_config, "image_size", None)
+        model: Optional["PreTrainedModel"] = None,
+        model_config: "CLIPVisionConfig" = None,
+        rbln_config: Optional[RBLNCLIPVisionModelConfig] = None,
+    ) -> RBLNCLIPVisionModelConfig:
+        if rbln_config.image_size is None:
+            rbln_config.image_size = getattr(model_config, "image_size", None)
 
-        if isinstance(rbln_image_size, int):
-            rbln_image_size = (rbln_image_size, rbln_image_size)
+        if isinstance(rbln_config.image_size, int):
+            rbln_config.image_size = (rbln_config.image_size, rbln_config.image_size)
 
-        if rbln_image_size is None:
+        if rbln_config.image_size is None:
             raise ValueError("`rbln_image_size` should be specified!")
 
         rbln_compile_config = RBLNCompileConfig(
@@ -144,45 +142,44 @@ class RBLNCLIPVisionModel(RBLNModel):
                 (
                     "pixel_values",
                     [
-                        rbln_batch_size,
+                        rbln_config.batch_size,
                         3,
-                        rbln_image_size[0],
-                        rbln_image_size[1],
+                        rbln_config.image_height,
+                        rbln_config.image_width,
                     ],
                     "float32",
                 )
             ]
         )
 
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-
-        rbln_config.model_cfg.update(
-            {
-                "batch_size": rbln_batch_size,
-                "image_size": rbln_image_size,
-            }
-        )
-
+        rbln_config.set_compile_cfgs([rbln_compile_config])
         return rbln_config
 
     def forward(
         self,
         pixel_values: Optional[torch.FloatTensor] = None,
+        return_dict: bool = None,
         **kwargs,
-    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+    ) -> 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)
-        return BaseModelOutputWithPooling(
-            last_hidden_state=output[0],
-            pooler_output=output[1],
-            hidden_states=output[2:],
-        )
+        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.
+        """
+        if not return_dict:
+            return (output,) if not isinstance(output, (tuple, list)) else output
+        else:
+            return CLIPVisionModelOutput(
+                image_embeds=output[0],
+                last_hidden_state=output[1],
+                hidden_states=output[2:],
+            )
 
 
 class RBLNCLIPVisionModelWithProjection(RBLNCLIPVisionModel):

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

@@ -12,4 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from ....ops import (
+    paged_attn_decode,
+    paged_attn_prefill,
+    paged_causal_attn_decode,
+    paged_causal_attn_prefill,
+    paged_flash_attn_decode,
+    paged_flash_attn_prefill,
+    paged_flash_causal_attn_decode,
+    paged_flash_causal_attn_prefill,
+)
+from .configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
 from .modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM

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

@@ -0,0 +1,90 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, Optional
+
+import rebel
+
+from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+from ...utils.rbln_quantization import QuantizationManager
+
+
+logger = get_logger()
+
+
+class RBLNDecoderOnlyModelForCausalLMConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        max_seq_len: Optional[int] = None,
+        use_inputs_embeds: Optional[bool] = None,
+        use_attention_mask: Optional[bool] = None,
+        attn_impl: Optional[str] = None,
+        kvcache_partition_len: Optional[int] = None,
+        kvcache_block_size: Optional[int] = None,
+        quantization: Optional[Dict[str, Any]] = None,
+        prefill_chunk_size: Optional[int] = None,
+        kvcache_num_blocks: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            max_seq_len (Optional[int]): The maximum sequence length supported by the model.
+            use_inputs_embeds (Optional[bool]): Whether to use input embeddings directly. Defaults to False.
+            use_attention_mask (Optional[bool]): Whether to use attention masks. This is automatically set to True
+                for RBLN-CA02 devices.
+            attn_impl (Optional[str]): The attention implementation to use.
+            kvcache_partition_len (Optional[int]): The length of each KV cache partition.
+            kvcache_block_size (Optional[int]): The block size for KV cache.
+            quantization (Optional[Dict[str, Any]]): Configuration for model quantization.
+            prefill_chunk_size (Optional[int]): The chunk size for prefilling the KV cache. Defaults to 128,
+                and must be a positive integer divisible by 64.
+            kvcache_num_blocks (Optional[int]): The number of blocks in the KV cache.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer or if prefill_chunk_size is not
+                a positive integer divisible by 64.
+        """
+        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_seq_len = max_seq_len
+        self.use_inputs_embeds = use_inputs_embeds or False
+
+        self.use_attention_mask = use_attention_mask
+        npu = self.npu or rebel.get_npu_name()
+        if npu == "RBLN-CA02":
+            if self.use_attention_mask is False:
+                logger.warning("Attention mask should be used with RBLN-CA02. Setting use_attention_mask to True.")
+            self.use_attention_mask = True
+        else:
+            self.use_attention_mask = self.use_attention_mask or False
+
+        self.attn_impl = attn_impl
+        self.kvcache_partition_len = kvcache_partition_len
+        self.kvcache_block_size = kvcache_block_size
+        self.quantization = quantization or {}
+        if self.quantization:
+            QuantizationManager.validate_quantization_config(self.quantization)
+
+        self.prefill_chunk_size = prefill_chunk_size or 128
+        if self.prefill_chunk_size % 64 != 0 or self.prefill_chunk_size <= 0:
+            raise ValueError("`prefill_chunk_size` must be a positive integer divisible by 64.")
+
+        self.kvcache_num_blocks = kvcache_num_blocks