optimum-rbln 0.2.1a1__tar.gz → 0.2.1a3__tar.gz

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/.github/workflows/rbln_trigger_on_pr.yaml

@@ -9,7 +9,7 @@ env:
   REBEL_PYPI_ENDPOINT: ${{ vars.REBEL_PYPI_INTERNAL_ENDPOINT }}
   REBEL_PYPI_USERNAME: ${{ secrets.REBEL_PYPI_USERNAME }}
   REBEL_PYPI_PASSWORD: ${{ secrets.REBEL_PYPI_PASSWORD }}
-  REBEL_COMPILER_VERSION: 0.7.1.dev63+ge6c4fb68
+  REBEL_COMPILER_VERSION: 0.7.2.dev151+g19f099fd
 
 jobs:
   check-rebel-compiler-version:

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.2.1a1
+Version: 0.2.1a3
 Summary: Optimum RBLN is the interface between the Hugging Face 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
@@ -28,7 +28,7 @@ Requires-Dist: packaging>=24.1
 Requires-Dist: torch<=2.5.1
 Requires-Dist: torchaudio<=2.5.1
 Requires-Dist: torchvision<=0.20.1
-Requires-Dist: transformers==4.45.2
+Requires-Dist: transformers==4.48.3
 Description-Content-Type: text/markdown
 
 

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/examples/speech-recognition/run_whisper.py

@@ -104,7 +104,7 @@ def main(
             **gen_kwargs,
         )
 
-        generated_ids = outputs
+        generated_ids = outputs.get("sequences") if isinstance(outputs, dict) else outputs
         transcriptions = processor.batch_decode(generated_ids, skip_special_tokens=True, decode_with_timestamps=True)
 
         print("---RBLN Longform Generate Result ---")

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
     "torchaudio<=2.5.1",
     "torchvision<=0.20.1",
     "accelerate>=1.0.1",
-    "transformers==4.45.2",
+    "transformers==4.48.3",
     "diffusers<=0.31.0",
     "packaging>=24.1",
 ]

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/__init__.py

@@ -71,6 +71,7 @@ _import_structure = {
         "RBLNRobertaForSequenceClassification",
         "RBLNRobertaForMaskedLM",
         "RBLNViTForImageClassification",
+        "RBLNBertForMaskedLM",
     ],
     "diffusers": [
         "RBLNStableDiffusionPipeline",
@@ -141,6 +142,7 @@ if TYPE_CHECKING:
         RBLNAutoModelForVision2Seq,
         RBLNBartForConditionalGeneration,
         RBLNBartModel,
+        RBLNBertForMaskedLM,
         RBLNBertForQuestionAnswering,
         RBLNBertModel,
         RBLNCLIPTextModel,

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/__version__.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1a1'
+__version__ = version = '0.2.1a3'
 __version_tuple__ = version_tuple = (0, 2, 1)

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/diffusers/models/autoencoders/autoencoder_kl.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import rebel
@@ -23,6 +22,7 @@ from transformers import PretrainedConfig
 
 from ....modeling import RBLNModel
 from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
+from ....utils.logging import get_logger
 from ...modeling_diffusers import RBLNDiffusionMixin
 from .vae import RBLNRuntimeVAEDecoder, RBLNRuntimeVAEEncoder, _VAEDecoder, _VAEEncoder
 
@@ -31,7 +31,7 @@ if TYPE_CHECKING:
     import torch
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class RBLNAutoencoderKL(RBLNModel):

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/diffusers/models/autoencoders/vae.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 from typing import TYPE_CHECKING
 
 import torch  # noqa: I001
@@ -20,13 +19,14 @@ from diffusers import AutoencoderKL
 from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
 from diffusers.models.modeling_outputs import AutoencoderKLOutput
 
+from ....utils.logging import get_logger
 from ....utils.runtime_utils import RBLNPytorchRuntime
 
 
 if TYPE_CHECKING:
     import torch
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class RBLNRuntimeVAEEncoder(RBLNPytorchRuntime):

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/diffusers/models/controlnet.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import importlib
-import logging
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import torch
@@ -22,6 +21,7 @@ from transformers import PretrainedConfig
 
 from ...modeling import RBLNModel
 from ...modeling_config import RBLNCompileConfig, RBLNConfig
+from ...utils.logging import get_logger
 from ..modeling_diffusers import RBLNDiffusionMixin
 
 
@@ -29,7 +29,7 @@ if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class _ControlNetModel(torch.nn.Module):

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/diffusers/models/transformers/transformer_sd3.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import torch
@@ -22,13 +21,14 @@ from transformers import PretrainedConfig
 
 from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
+from ....utils.logging import get_logger
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class SD3Transformer2DModelWrapper(torch.nn.Module):

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/diffusers/models/unets/unet_2d_condition.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
@@ -22,13 +21,14 @@ from transformers import PretrainedConfig
 
 from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
+from ....utils.logging import get_logger
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class _UNet_SD(torch.nn.Module):

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 import os
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
@@ -21,13 +20,14 @@ import torch
 from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
 
 from ....modeling import RBLNModel
+from ....utils.logging import get_logger
 from ...models.controlnet import RBLNControlNetModel
 
 
 if TYPE_CHECKING:
     pass
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class RBLNMultiControlNetModel(RBLNModel):

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py

@@ -34,16 +34,17 @@ from diffusers import StableDiffusionControlNetPipeline
 from diffusers.image_processor import PipelineImageInput
 from diffusers.pipelines.controlnet.pipeline_controlnet import retrieve_timesteps
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
-from diffusers.utils import deprecate, logging
+from diffusers.utils import deprecate
 from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
 
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ....utils.logging import get_logger
 from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
 
-logger = logging.get_logger(__name__)
+logger = get_logger(__name__)
 
 
 class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionControlNetPipeline):

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/modeling.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
@@ -24,13 +23,14 @@ from transformers import AutoConfig, PretrainedConfig
 
 from .modeling_base import RBLNBaseModel
 from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, use_rbln_config
+from .utils.logging import get_logger
 
 
 if TYPE_CHECKING:
     from transformers import PreTrainedModel
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class RBLNModel(RBLNBaseModel):

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/modeling_base.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import importlib
-import logging
 import os
 import shutil
 from abc import ABC, abstractmethod
@@ -32,6 +31,7 @@ from transformers import (
 
 from .modeling_config import RBLNCompileConfig, RBLNConfig, use_rbln_config
 from .utils.hub import PushToHubMixin, pull_compiled_model_from_hub, validate_files
+from .utils.logging import get_logger
 from .utils.runtime_utils import UnavailableRuntime
 from .utils.save_utils import maybe_load_preprocessors
 from .utils.submodule import SubModulesMixin
@@ -40,7 +40,7 @@ from .utils.submodule import SubModulesMixin
 if TYPE_CHECKING:
     from transformers import PreTrainedModel
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class PreTrainedModel(ABC):  # noqa: F811
@@ -442,27 +442,47 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
             return
 
-        os.makedirs(save_directory, exist_ok=True)
-
         real_save_dir = self.model_save_dir / self.subfolder
         save_directory_path = Path(save_directory)
-        if os.path.exists(real_save_dir) and os.path.isdir(real_save_dir):
-            if save_directory_path.absolute() == real_save_dir.absolute():
-                raise FileExistsError(
-                    f"Cannot save model to '{save_directory}'. "
-                    f"This directory already exists and contains the model files."
-                )
-            shutil.copytree(real_save_dir, save_directory, dirs_exist_ok=True)
-            self.config.save_pretrained(save_directory)
-            if self.generation_config is not None:
-                self.generation_config.save_pretrained(save_directory)
-        else:
+
+        if not os.path.exists(real_save_dir) or not os.path.isdir(real_save_dir):
             raise FileNotFoundError(
                 f"Unable to save the model. The model directory '{real_save_dir}' does not exist or is not accessible. "
                 f"Cannot save to the specified destination '{save_directory}'. "
                 f"Please ensure the model directory exists and you have the necessary permissions to access it."
             )
 
+        if save_directory_path.absolute() == real_save_dir.absolute():
+            raise FileExistsError(
+                f"Cannot save model to '{save_directory}'. This directory already exists and contains the model files."
+            )
+
+        # Create a temporary directory next to the target directory
+        tmp_dir = save_directory + ".tmp"
+        try:
+            # Remove temporary directory if it exists from a previous failed attempt
+            if os.path.exists(tmp_dir):
+                shutil.rmtree(tmp_dir)
+
+            # First copy everything to a temporary directory
+            shutil.copytree(real_save_dir, tmp_dir)
+
+            # Save configs to the temporary directory
+            self.config.save_pretrained(tmp_dir)
+            if self.generation_config is not None:
+                self.generation_config.save_pretrained(tmp_dir)
+
+            # If everything succeeded, atomically replace the target directory
+            if os.path.exists(save_directory):
+                shutil.rmtree(save_directory)
+            os.rename(tmp_dir, save_directory)
+
+        except Exception as e:
+            # Clean up the temporary directory if anything fails
+            if os.path.exists(tmp_dir):
+                shutil.rmtree(tmp_dir)
+            raise e  # Re-raise the exception after cleanup
+
         if push_to_hub:
             return super().push_to_hub(save_directory, **kwargs)
 

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/ops/attn.py

@@ -152,16 +152,16 @@ def register_rbln_custom_attention_add_softmax():
         """
         return (
             q,
-            torch.empty(1, *kcache.shape[1:], device=kcache.device),
-            torch.empty(1, *vcache.shape[1:], device=vcache.device),
+            torch.empty(*kcache.shape, device=kcache.device),
+            torch.empty(*vcache.shape, device=vcache.device),
         )
 
     @register_fake("rbln_custom_ops::attn_decode_add_softmax")
     def attn_decode_add_softmax_abstract(q, k, v, m, kcache, vcache, seq, partition):
         return (
             q,
-            torch.empty(1, *kcache.shape[1:], device=kcache.device),
-            torch.empty(1, *vcache.shape[1:], device=vcache.device),
+            torch.empty(*kcache.shape, device=kcache.device),
+            torch.empty(*vcache.shape, device=vcache.device),
         )
 
     torch.library.define(

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/transformers/__init__.py

@@ -35,6 +35,7 @@ _import_structure = {
         "RBLNBartForConditionalGeneration",
         "RBLNBartModel",
         "RBLNBertModel",
+        "RBLNBertForMaskedLM",
         "RBLNBertForQuestionAnswering",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
@@ -92,6 +93,7 @@ if TYPE_CHECKING:
         RBLNAutoModelForVision2Seq,
         RBLNBartForConditionalGeneration,
         RBLNBartModel,
+        RBLNBertForMaskedLM,
         RBLNBertForQuestionAnswering,
         RBLNBertModel,
         RBLNCLIPTextModel,

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/transformers/models/__init__.py

@@ -33,7 +33,7 @@ _import_structure = {
         "RBLNAutoModelForVision2Seq",
     ],
     "bart": ["RBLNBartForConditionalGeneration", "RBLNBartModel"],
-    "bert": ["RBLNBertModel", "RBLNBertForQuestionAnswering"],
+    "bert": ["RBLNBertModel", "RBLNBertForQuestionAnswering", "RBLNBertForMaskedLM"],
     "clip": ["RBLNCLIPTextModel", "RBLNCLIPTextModelWithProjection", "RBLNCLIPVisionModel"],
     "dpt": ["RBLNDPTForDepthEstimation"],
     "exaone": ["RBLNExaoneForCausalLM"],
@@ -67,7 +67,7 @@ if TYPE_CHECKING:
         RBLNAutoModelForVision2Seq,
     )
     from .bart import RBLNBartForConditionalGeneration, RBLNBartModel
-    from .bert import RBLNBertForQuestionAnswering, RBLNBertModel
+    from .bert import RBLNBertForMaskedLM, RBLNBertForQuestionAnswering, RBLNBertModel
     from .clip import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection, RBLNCLIPVisionModel
     from .dpt import RBLNDPTForDepthEstimation
     from .exaone import RBLNExaoneForCausalLM

optimum_rbln-0.2.1a3/src/optimum/rbln/transformers/models/bert/__init__.py

@@ -0,0 +1,15 @@
+# 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 .modeling_bert import RBLNBertForMaskedLM, RBLNBertForQuestionAnswering, RBLNBertModel

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/transformers/models/bert/modeling_bert.py

@@ -13,17 +13,17 @@
 # limitations under the License.
 
 import inspect
-import logging
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 from transformers import PretrainedConfig
 
 from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
-from ...modeling_generic import RBLNModelForQuestionAnswering
+from ....utils.logging import get_logger
+from ...modeling_generic import RBLNModelForMaskedLM, RBLNModelForQuestionAnswering
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
@@ -100,5 +100,9 @@ class RBLNBertModel(RBLNModel):
         return rbln_config
 
 
+class RBLNBertForMaskedLM(RBLNModelForMaskedLM):
+    rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
+
+
 class RBLNBertForQuestionAnswering(RBLNModelForQuestionAnswering):
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 import torch
@@ -28,9 +27,10 @@ from transformers.models.clip.modeling_clip import CLIPTextModelOutput
 from ....diffusers.modeling_diffusers import RBLNDiffusionMixin
 from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
+from ....utils.logging import get_logger
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, CLIPTextModel

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py

@@ -544,15 +544,19 @@ class DecoderOnlyAttention(nn.Module):
         super().__init__()
         self._original_mod = self_attn
         self.layer_idx = self_attn.layer_idx
-        self.num_heads = self._original_mod.num_heads
+        self.num_heads = getattr(self._original_mod, "num_heads", None) or getattr(
+            self._original_mod.config, "num_attention_heads"
+        )
         self.head_dim = self._original_mod.head_dim
         self._phase = "prefill"
         self.scale = torch.tensor(self.get_attn_scale())
 
         if hasattr(self._original_mod, "num_key_value_heads"):
             self.num_key_value_heads = self._original_mod.num_key_value_heads
+        elif hasattr(self._original_mod, "config") and hasattr(self._original_mod.config, "num_key_value_heads"):
+            self.num_key_value_heads = self._original_mod.config.num_key_value_heads
         else:
-            self.num_key_value_heads = self._original_mod.num_heads
+            self.num_key_value_heads = self.num_heads
 
         self.attention = self.get_attention()
         self.__post_init__()

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/transformers/models/dpt/modeling_dpt.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Union
 
 from transformers import AutoModelForDepthEstimation
@@ -20,9 +19,10 @@ from transformers.modeling_outputs import DepthEstimatorOutput
 
 from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
+from ....utils.logging import get_logger
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/transformers/models/llava_next/modeling_llava_next.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import inspect
-import logging
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
 
@@ -30,10 +29,11 @@ from transformers.models.llava_next.modeling_llava_next import LlavaNextCausalLM
 
 from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
+from ....utils.logging import get_logger
 from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyOutput
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 if TYPE_CHECKING:
     from transformers import (

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import inspect
-import logging
 from abc import ABC
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
@@ -25,10 +24,11 @@ from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
 
 from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
+from ....utils.logging import get_logger
 from ....utils.runtime_utils import RBLNPytorchRuntime
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 if TYPE_CHECKING:
     from transformers import (

{optimum_rbln-0.2.1a1 → optimum_rbln-0.2.1a3}/src/optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py

@@ -420,7 +420,7 @@ class Seq2SeqSelfAttention(nn.Module):
         pass
 
     def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int) -> torch.Tensor:
-        return tensor.view(bsz, 1, seq_len, 1, self.num_heads, self.head_dim).transpose(2, 4)
+        return tensor.view(bsz, seq_len, 1, self.num_heads, self.head_dim).transpose(1, 3)
 
     def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Projects input hidden states into query, key, and value representations.
@@ -450,38 +450,21 @@ class Seq2SeqSelfAttention(nn.Module):
         key_states = self._shape(key_states, -1, bsz)
         value_states = self._shape(value_states, -1, bsz)
 
-        all_key_states = []
-        all_value_states = []
-        all_attn_output = []
-        for b_idx in range(bsz):
-            query_state = query_states[b_idx]
-            key_state = key_states[b_idx]
-            value_state = value_states[b_idx]
-            attn_mask = attention_mask[b_idx].unsqueeze(0).unsqueeze(2)
-            past_key_state = past_key_value[0].view(bsz, self.num_heads, 1, -1, self.head_dim)
-            past_value_state = past_key_value[1].view(bsz, self.num_heads, 1, -1, self.head_dim)
-
-            attn_output, key_state, value_state = self.attn_decode(
-                query_state,
-                key_state,
-                value_state,
-                attn_mask,
-                past_key_state,
-                past_value_state,
-                cache_position[b_idx][0],
-                torch.tensor(1.0, dtype=torch.float32),  # scale
-            )
-
-            attn_output = attn_output.view(1, self.num_heads, -1, self.head_dim).transpose(1, 2)
-            attn_output = attn_output.reshape(1, -1, self.num_heads * self.head_dim)
-
-            all_key_states.append(key_state.squeeze(2))
-            all_value_states.append(value_state.squeeze(2))
-            all_attn_output.append(attn_output)
+        attn_output, key_states, value_states = self.attn_decode(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask.unsqueeze(
+                2
+            ),  # Unsqueeze group axis since CustomKernel expects it for group query attention
+            past_key_value[0].view(bsz, self.num_heads, 1, -1, self.head_dim),
+            past_key_value[1].view(bsz, self.num_heads, 1, -1, self.head_dim),
+            cache_position.squeeze(1),
+            torch.tensor(1.0, dtype=torch.float32),  # scale
+        )
 
-        key_states = torch.cat(all_key_states, dim=0)
-        value_states = torch.cat(all_value_states, dim=0)
-        attn_output = torch.cat(all_attn_output, dim=0)
+        attn_output = attn_output.view(bsz, self.num_heads, -1, self.head_dim).transpose(1, 2)
+        attn_output = attn_output.reshape(bsz, -1, self.num_heads * self.head_dim)
 
         attn_output = self.out_proj(attn_output)
         present_key_value = (key_states, value_states)