optimum-rbln 0.7.3a1__py3-none-any.whl → 0.7.3a2__py3-none-any.whl

optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py

@@ -12,13 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 from torch import nn
 from transformers.utils import logging
 
-from ....ops import register_rbln_custom_cache_update, register_rbln_custom_masked_attention
+from ....ops import (
+    register_rbln_custom_cache_update,
+    register_rbln_custom_paged_attention,
+    register_rbln_custom_paged_causal_attention,
+)
 
 
 logger = logging.get_logger(__name__)
@@ -87,7 +91,7 @@ class Seq2SeqEncoderWrapper(nn.Module):
         input_ids: torch.Tensor,
         attention_mask: torch.Tensor,
         cross_key_values: torch.Tensor,
-        batch_position: torch.Tensor,
+        b_idx: torch.Tensor,
     ) -> Tuple[torch.Tensor]:
         # 1. get encoder last_hidden_states
         encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
@@ -111,7 +115,7 @@ class Seq2SeqEncoderWrapper(nn.Module):
         # 3. update the cross_attention's past_key_value direct to the device-dram for optimization.
         batch_axis = torch.tensor(1, dtype=torch.int16)
         cross_key_values = torch.ops.rbln_custom_ops.rbln_cache_update(
-            cross_key_values, cross_kv, batch_position, batch_axis
+            cross_key_values, cross_kv, b_idx[0], batch_axis
         )
 
         return cross_key_values
@@ -131,9 +135,10 @@ class Seq2SeqDecoderWrapper(nn.Module):
         **kwargs: Additional arguments for decoder configuration.
     """
 
-    def __init__(self, model: nn.Module, **kwargs):
+    def __init__(self, model: nn.Module, use_attention_mask: bool = True, **kwargs):
         super().__init__()
         self.config = model.config
+        self.use_attention_mask = use_attention_mask
         self.__post_init__(model, **kwargs)
 
     def __post_init__(self, model: nn.Module, **kwargs):
@@ -143,7 +148,11 @@ class Seq2SeqDecoderWrapper(nn.Module):
         It is inspired by the BART architecture, but it is designed to be flexible and can be overridden
         by subclasses to modify or add custom attributes as necessary.
         """
-        register_rbln_custom_masked_attention()
+        if self.use_attention_mask:
+            register_rbln_custom_paged_attention()
+        else:
+            register_rbln_custom_paged_causal_attention()
+
         self.num_layers = self.config.decoder_layers
         self.conditional_generation = self.convert_to_rbln_conditional_generation(model)
 
@@ -160,13 +169,23 @@ class Seq2SeqDecoderWrapper(nn.Module):
 
     def forward(
         self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-        encoder_attention_mask: torch.Tensor,
-        cache_position: torch.Tensor,
-        cross_kv_cache: torch.Tensor,
-        *self_kv_cache: torch.Tensor,
+        *args,
     ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor]]:
+        if self.use_attention_mask:
+            (
+                input_ids,
+                attention_mask,
+                encoder_attention_mask,
+                cache_position,
+                block_tables,
+                cross_kv_cache,
+                *self_kv_cache,
+            ) = args
+
+        else:
+            attention_mask = None
+            (input_ids, encoder_attention_mask, cache_position, block_tables, cross_kv_cache, *self_kv_cache) = args
+
         self_past_key_values = ()
         cross_past_key_values = ()
         for i in range(0, self.num_layers * 2, 2):
@@ -181,6 +200,7 @@ class Seq2SeqDecoderWrapper(nn.Module):
             self_past_key_values=self_past_key_values,
             cross_past_key_values=cross_past_key_values,
             cache_position=cache_position,
+            block_tables=block_tables,
         )
 
         outputs = (lm_logits,) + self_present_key_values
@@ -228,6 +248,7 @@ class Seq2SeqForConditionalGeneration(nn.Module):
         self_past_key_values,
         cross_past_key_values,
         cache_position,
+        block_tables: Optional[torch.Tensor] = None,
     ):
         hidden_states, self_present_key_values = self.decoder(
             input_ids=input_ids,
@@ -236,6 +257,7 @@ class Seq2SeqForConditionalGeneration(nn.Module):
             self_past_key_values=self_past_key_values,
             cross_past_key_values=cross_past_key_values,
             cache_position=cache_position,
+            block_tables=block_tables,
         )
 
         if self.has_rescaling and self.config.tie_word_embeddings:
@@ -292,6 +314,7 @@ class Seq2SeqDecoder(torch.nn.Module):
         self_past_key_values: torch.Tensor,
         cross_past_key_values: torch.Tensor,
         cache_position: torch.Tensor,
+        block_tables: Optional[torch.Tensor] = None,
     ):
         # embedding
         hidden_states = self.get_embedding()(input_ids)
@@ -314,6 +337,7 @@ class Seq2SeqDecoder(torch.nn.Module):
                 self_past_key_value=self_past_key_value,
                 cross_past_key_value=cross_past_key_value,
                 cache_position=cache_position,
+                block_tables=block_tables,
             )
             self_present_key_values += self_present_key_value
 
@@ -373,6 +397,7 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
         self_past_key_value: Tuple[torch.Tensor],
         cross_past_key_value: Tuple[torch.Tensor],
         cache_position: torch.Tensor,
+        block_tables: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
         dummy_encoder_hidden_states = torch.zeros(1, encoder_attention_mask.shape[-1])
 
@@ -384,6 +409,7 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
             past_key_value=self_past_key_value,
             attention_mask=attention_mask,
             cache_position=cache_position,
+            block_tables=block_tables,
         )
         hidden_states = residual + hidden_states
         hidden_states = self.post_self_attn_layer_norm(hidden_states)
@@ -407,10 +433,10 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
 
 
 class Seq2SeqSelfAttention(nn.Module):
-    def __init__(self, attn):
+    def __init__(self, attn, **kwargs):
         super().__init__()
         self._original_mod = attn
-        self.__post_init__()
+        self.__post_init__(**kwargs)
 
     def __post_init__(self, **kwargs):
         """
@@ -442,6 +468,7 @@ class Seq2SeqSelfAttention(nn.Module):
         past_key_value: Tuple[torch.Tensor],
         attention_mask: torch.Tensor,
         cache_position: torch.Tensor,
+        block_tables: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
         bsz, tgt_len, _ = hidden_states.size()
 
@@ -450,18 +477,22 @@ class Seq2SeqSelfAttention(nn.Module):
         key_states = self._shape(key_states, -1, bsz)
         value_states = self._shape(value_states, -1, bsz)
 
-        attn_output, key_states, value_states = self.attn_decode(
+        block_size = past_key_value[0].shape[-2]
+        args = [
             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,
             torch.tensor(1.0, dtype=torch.float32),  # scale
-        )
+            block_tables,
+            block_size,
+        ]
+        if attention_mask is not None:
+            args.insert(3, attention_mask.unsqueeze(2))
+
+        attn_output, key_states, value_states = self.attn_decode(*args)
 
         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)

optimum/rbln/transformers/models/t5/modeling_t5.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
 
+import rebel
 import torch
 from transformers import (
     AutoModelForTextEncoding,
@@ -22,7 +23,7 @@ from transformers import (
     T5EncoderModel,
     T5ForConditionalGeneration,
 )
-from transformers.modeling_outputs import BaseModelOutput
+from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
 
 from ....diffusers.modeling_diffusers import RBLNDiffusionMixin
 from ....modeling import RBLNModel
@@ -57,6 +58,63 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         )
 
 
+class RBLNRuntimeEncoder(RBLNPytorchRuntime):
+    mandatory_members = ["main_input_name"]
+
+    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
+        _ = super().forward(*args, **kwargs)
+        return BaseModelOutput(last_hidden_state=torch.tensor([1.0]))
+
+
+class RBLNRuntimeDecoder(RBLNPytorchRuntime):
+    mandatory_members = ["main_input_name"]
+
+    def __init__(
+        self,
+        runtime: rebel.Runtime,
+        batch_size: int,
+        dec_max_seq_len: int,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(runtime, **kwargs)
+        self.batch_size = batch_size
+        self.dec_max_seq_len = dec_max_seq_len
+
+    def forward(
+        self,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        decoder_attention_mask: Optional[torch.BoolTensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor]:
+        batch_size = decoder_input_ids.shape[0]
+        if batch_size != self.batch_size:
+            raise RuntimeError(
+                f"Batch size mismatch: got {batch_size}, expected {self.batch_size} (compiled batch size)."
+            )
+
+        if batch_size != cache_position.shape[0]:
+            raise RuntimeError(f"Cache position size mismatch: got {cache_position.shape[0]}, expected {batch_size}.")
+
+        for b_idx in range(self.batch_size):
+            decoding_step = cache_position[b_idx].item()
+            if not (0 <= decoding_step < self.dec_max_seq_len):
+                raise ValueError(
+                    f"Decoding step {decoding_step} out of bounds for attention mask with shape {self.dec_attn_mask.shape}."
+                )
+            decoder_attention_mask[b_idx, : decoding_step + 1] = 1
+
+        lm_logits = super().forward(
+            decoder_input_ids,
+            decoder_attention_mask,
+            attention_mask,
+            cache_position,
+        )
+
+        return Seq2SeqLMOutput(logits=lm_logits)
+
+
 class T5EncoderWrapper(torch.nn.Module):
     def __init__(self, model: "T5EncoderModel") -> None:
         super().__init__()
@@ -189,6 +247,21 @@ class RBLNT5EncoderModel(RBLNModel):
 
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    def __post_init__(self, **kwargs):
+        batch_size = self.rbln_config.model_cfg["batch_size"]
+        dec_max_seq_len = self.rbln_config.model_cfg["dec_max_seq_len"]
+
+        self.encoder = RBLNRuntimeEncoder(
+            runtime=self.model[0],
+            main_input_name="input_ids",
+        )
+        self.decoder = RBLNRuntimeDecoder(
+            runtime=self.model[1],
+            main_input_name="input_ids",
+            batch_size=batch_size,
+            dec_max_seq_len=dec_max_seq_len,
+        )
+
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
         enc_max_seq_len = rbln_config.model_cfg["enc_max_seq_len"]
@@ -206,3 +279,139 @@ class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
             return redirect(val)
 
         return val
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model_config: "PretrainedConfig",
+        rbln_kwargs: Dict[str, Any] = {},
+    ) -> RBLNConfig:
+        rbln_enc_max_seq_len = rbln_kwargs.get("enc_max_seq_len", None)
+        rbln_dec_max_seq_len = rbln_kwargs.get("dec_max_seq_len", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
+
+        n_layer = getattr(model_config, "decoder_layers", None) or getattr(model_config, "num_layers")
+        n_head = getattr(model_config, "decoder_attention_heads", None) or getattr(model_config, "num_heads")
+        d_kv = (
+            model_config.d_kv
+            if hasattr(model_config, "d_kv")
+            else model_config.d_model // model_config.encoder_attention_heads
+        )
+
+        max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
+            model_config, "max_position_embeddings", None
+        )
+
+        rbln_pad_token_id = getattr(model_config, "pad_token_id", None)
+        if rbln_pad_token_id is None:
+            rbln_pad_token_id = getattr(model_config, "bos_token_id", None)
+            if rbln_pad_token_id is None:
+                rbln_pad_token_id = getattr(model_config, "eos_token_id", None)
+                if rbln_pad_token_id is None:
+                    rbln_pad_token_id = -1
+
+        if rbln_enc_max_seq_len is None:
+            rbln_enc_max_seq_len = max_position_embeddings
+            if rbln_enc_max_seq_len is None:
+                for tokenizer in preprocessors:
+                    if hasattr(tokenizer, "model_max_length"):
+                        rbln_enc_max_seq_len = tokenizer.model_max_length
+                        break
+                if rbln_enc_max_seq_len is None:
+                    raise ValueError("`rbln_enc_max_seq_len` should be specified!")
+        if max_position_embeddings is not None and rbln_enc_max_seq_len > max_position_embeddings:
+            raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
+
+        if rbln_dec_max_seq_len is None:
+            rbln_dec_max_seq_len = max_position_embeddings
+            if rbln_dec_max_seq_len is None:
+                for tokenizer in preprocessors:
+                    if hasattr(tokenizer, "model_max_length"):
+                        rbln_dec_max_seq_len = tokenizer.model_max_length
+                        break
+                if rbln_dec_max_seq_len is None:
+                    raise ValueError("`rbln_dec_max_seq_len` should be specified!")
+
+        if max_position_embeddings is not None and rbln_dec_max_seq_len > max_position_embeddings:
+            raise ValueError("`rbln_dec_max_seq_len` should be less or equal than max_position_embeddings!")
+
+        # model input info
+        enc_input_info = [
+            ("input_ids", [1, rbln_enc_max_seq_len], "int64"),
+            ("attention_mask", [1, rbln_enc_max_seq_len], "float32"),
+            (
+                "cross_key_value_states",
+                [
+                    n_layer * 2,
+                    rbln_batch_size,
+                    n_head,
+                    rbln_enc_max_seq_len,
+                    d_kv,
+                ],
+                "float32",
+            ),
+            ("block_tables", [1], "int16"),
+        ]
+
+        dec_input_info = [
+            ("input_ids", [rbln_batch_size, 1], "int64"),
+            ("attention_mask", [rbln_batch_size, rbln_dec_max_seq_len], "float32"),
+            ("encoder_attention_mask", [rbln_batch_size, rbln_enc_max_seq_len], "float32"),
+            (
+                "cache_position",
+                [rbln_batch_size, 1],
+                "int32",
+            ),
+        ]
+        dec_input_info.extend(
+            [
+                (
+                    "cross_key_value_states",
+                    [
+                        n_layer * 2,
+                        rbln_batch_size,
+                        n_head,
+                        rbln_enc_max_seq_len,
+                        d_kv,
+                    ],
+                    "float32",
+                )
+            ]
+        )
+        dec_input_info.extend(
+            [
+                (
+                    f"self_key_value_states_{i}",
+                    [
+                        rbln_batch_size,
+                        n_head,
+                        rbln_dec_max_seq_len,
+                        d_kv,
+                    ],
+                    "float32",
+                )
+                for i in range(n_layer * 2)
+            ]
+        )
+
+        enc_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=enc_input_info)
+        dec_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
+
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[enc_compile_config, dec_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+
+        rbln_config.model_cfg.update(
+            {
+                "enc_max_seq_len": rbln_enc_max_seq_len,
+                "dec_max_seq_len": rbln_dec_max_seq_len,
+                "batch_size": rbln_batch_size,
+                "pad_token_id": rbln_pad_token_id,
+            }
+        )
+
+        return rbln_config

optimum/rbln/transformers/models/t5/t5_architecture.py

@@ -18,7 +18,7 @@ import torch
 from torch import nn
 from transformers.utils import logging
 
-from ....ops import register_rbln_custom_attention_add_softmax
+from ....ops import register_rbln_custom_add_softmax_attention
 from ..seq2seq.seq2seq_architecture import (
     Seq2SeqDecoder,
     Seq2SeqDecoderLayer,
@@ -55,7 +55,7 @@ class T5EncoderWrapper(Seq2SeqEncoderWrapper):
 
 class T5DecoderWrapper(Seq2SeqDecoderWrapper):
     def __post_init__(self, model, dec_max_seq_len: int = None):
-        register_rbln_custom_attention_add_softmax()
+        register_rbln_custom_add_softmax_attention()
         self.num_layers = self.config.num_layers
         self.conditional_generation = self.convert_to_rbln_conditional_generation(model, dec_max_seq_len)
 
@@ -71,6 +71,36 @@ class T5DecoderWrapper(Seq2SeqDecoderWrapper):
 
         return new_model
 
+    def forward(
+        self,
+        input_ids,
+        attention_mask,
+        encoder_attention_mask,
+        cache_position,
+        cross_kv_cache,
+        *self_kv_cache,
+    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor]]:
+        self_past_key_values = ()
+        cross_past_key_values = ()
+
+        for i in range(0, self.num_layers * 2, 2):
+            self_past_key_values = self_past_key_values + ((self_kv_cache[i], self_kv_cache[i + 1]),)
+            cross_past_key_values = cross_past_key_values + ((cross_kv_cache[i], cross_kv_cache[i + 1]),)
+
+        # decode
+        lm_logits, self_present_key_values = self.conditional_generation(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            encoder_attention_mask=encoder_attention_mask,
+            self_past_key_values=self_past_key_values,
+            cross_past_key_values=cross_past_key_values,
+            cache_position=cache_position,
+        )
+
+        outputs = (lm_logits,) + self_present_key_values
+
+        return outputs
+
 
 class T5ForConditionalGeneration(Seq2SeqForConditionalGeneration):
     has_rescaling = True
@@ -134,7 +164,7 @@ class T5LayerSelfAttention(Seq2SeqSelfAttention):
         self.out_proj = self._original_mod.o
         self.num_heads = self._original_mod.n_heads
         self.head_dim = self._original_mod.key_value_proj_dim
-        self.attn_decode = torch.ops.rbln_custom_ops.attn_decode_add_softmax
+        self.attn_decode = torch.ops.rbln_custom_ops.add_softmax_attn_decode
 
     def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         query_states = self.q_proj(hidden_states)
@@ -142,6 +172,42 @@ class T5LayerSelfAttention(Seq2SeqSelfAttention):
         value_states = self.v_proj(hidden_states)
         return query_states, key_states, value_states
 
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        past_key_value: Tuple[torch.Tensor],
+        attention_mask: torch.Tensor,
+        cache_position: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
+        bsz, tgt_len, _ = hidden_states.size()
+
+        query_states, key_states, value_states = self.projection(hidden_states=hidden_states)
+        query_states = self._shape(query_states, tgt_len, bsz)
+        key_states = self._shape(key_states, -1, bsz)
+        value_states = self._shape(value_states, -1, bsz)
+
+        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,
+            torch.tensor(1.0, dtype=torch.float32),  # scale
+        )
+
+        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)
+
+        return attn_output, present_key_value
+
 
 class T5CrossAttention(nn.Module):
     def __init__(self, attn):

{optimum_rbln-0.7.3a1.dist-info → optimum_rbln-0.7.3a2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.7.3a1
+Version: 0.7.3a2
 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

{optimum_rbln-0.7.3a1.dist-info → optimum_rbln-0.7.3a2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 optimum/rbln/__init__.py,sha256=eHi15YM3989AcX52jka9rUmgAtlp1PHqMNwBEdOfuu8,6554
-optimum/rbln/__version__.py,sha256=3XXLT-7KoOXBM5ecjGQ9vxdHcJ06x38tTkK1veoUkmQ,513
+optimum/rbln/__version__.py,sha256=bShBukYvw7AqWtLsut0yClygDEGsFRmxrXypqIeEXcQ,513
 optimum/rbln/modeling.py,sha256=3XE0IrCYbkjw9_Q9BFzZ_ri_Kyxw1g6iwfdohZB46-s,8289
 optimum/rbln/modeling_base.py,sha256=ELSPbjx7awBRM2SckkD-5gI3TIa01mfzz7gDRC1Pljk,21778
 optimum/rbln/modeling_config.py,sha256=7104bxmrvKW4Q6XTruQayiIGl8GHDFmPkJ3cknMIInE,11335
@@ -39,9 +39,9 @@ optimum/rbln/diffusers/pipelines/stable_diffusion_xl/__init__.py,sha256=9iIMZYvp
 optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py,sha256=OvB5bxX6HUiqJeIc3uukuEmUXYEx1pTqGNOtdG2l1m8,902
 optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py,sha256=3aB1Rw-OgKytQOHwOaShbEvq_XVHPOGvsGm8pstEmKU,930
 optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py,sha256=MzVP1wscaO1sUIiBIPJqG6zuGyez9VUbA42-JSIm-mk,930
-optimum/rbln/ops/__init__.py,sha256=-jcOGX3B8w5Znpr1z2eUsrK3TN-w9trrkSoqJRWgXdA,945
-optimum/rbln/ops/attn.py,sha256=WUsy4I25gm2j9Xdns9W2NNd3jCNcueqJuisDzp0jPaA,13899
-optimum/rbln/ops/flash_attn.py,sha256=aQRupKPvJsNFWKrHaeyXg-LemyUJWmCJaVrA__Mjabo,3869
+optimum/rbln/ops/__init__.py,sha256=TxOmsN0u3PmyK4Sb89qbiC4rePOlkvUT7Lm6wVoTnY0,941
+optimum/rbln/ops/attn.py,sha256=LbJAmFtNj05i6BURfKV3KybsPItFe8w-YdSe5SuWkEc,12365
+optimum/rbln/ops/flash_attn.py,sha256=4shKNY13skPoYnbEsGrXDzgNwBIhHZEFrnUnWx1ESZU,4076
 optimum/rbln/ops/kv_cache_update.py,sha256=9W4WCO1Dtfy0u5i978JJRa7uLbqrfR2lHuoPynb07fw,3143
 optimum/rbln/transformers/__init__.py,sha256=AGo3BqVIZrsOzYsQAnnQ25HCstTPBclrXbvvUxVMlqE,4255
 optimum/rbln/transformers/modeling_alias.py,sha256=yx7FnZQWAnrWzivaO5hI7T6i-fyLzt2tMIXG2oDNbPo,1657
@@ -52,25 +52,25 @@ optimum/rbln/transformers/models/auto/__init__.py,sha256=GvGbb3ZpMv-h6euXeZ42jSi
 optimum/rbln/transformers/models/auto/auto_factory.py,sha256=IK9jFrJ3EEzYQa9_aKpcp2TO68M5YGkA-HcfBVpA2QU,7027
 optimum/rbln/transformers/models/auto/modeling_auto.py,sha256=Un9qoqdy3dO8JBza_bTJF_6_fRVNM9QisihSgTRFI-o,3933
 optimum/rbln/transformers/models/bart/__init__.py,sha256=32HPe0_GIO0hp9U464Iv6Jd7M-1nop9g8hA1UZMHhyw,674
-optimum/rbln/transformers/models/bart/bart_architecture.py,sha256=ZV-6Y3xABJsGAw2wi3vGYZNXbeVp-DlI2uUsdsa-8M8,5486
-optimum/rbln/transformers/models/bart/modeling_bart.py,sha256=QZCTJA0K90YBzkCXxs3JR9Ol9lbmAn50RDeN2hcWtx8,5673
+optimum/rbln/transformers/models/bart/bart_architecture.py,sha256=Oo-Cdne7igKEex8wwP-gztKJHgs5GLHQjK1oc3IZIDE,5801
+optimum/rbln/transformers/models/bart/modeling_bart.py,sha256=iI3ubPOVvHmhLt0wEz_vkOfMyNTHVNjmnkLtbpOX760,5797
 optimum/rbln/transformers/models/bert/__init__.py,sha256=YVV7k_laU6yJBawZrgjIWjRmIF-Y4oQQHqyf8lsraQs,691
 optimum/rbln/transformers/models/bert/modeling_bert.py,sha256=p3utRqf3dv9_RkHwaMCa1EfXttNJkqCJUIZo3CeZ9YY,4674
 optimum/rbln/transformers/models/clip/__init__.py,sha256=H9vuBwrmFO0-CqZhXUrKF-uQL6igCqMlqrT1X_ELaAI,754
 optimum/rbln/transformers/models/clip/modeling_clip.py,sha256=NiSm7bHs4SReHDUr53BBWSX0Y8bkKOeUSpsBDrp8YDw,6628
 optimum/rbln/transformers/models/decoderonly/__init__.py,sha256=pDogsdpJKKB5rqnVFrRjwfhUvOSV-jZ3oARMsqSvOOQ,665
-optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py,sha256=6X87ZVvz4wHoTATdaxxSLy8wBfsEkUwWQISlo_mXPKM,40822
-optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py,sha256=BYENVqueqR121nPyh2LnV_sMdhl95GRgqHLnCcX2sz8,29067
+optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py,sha256=x8_xQ5aGXbadJyajpJQyfgxx4YPSj62VlmmGDMnC-1E,41819
+optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py,sha256=dyl8tDBjfe5VfU1XbKAoZS7g7F90JTYVmMuz0HTmCoE,35345
 optimum/rbln/transformers/models/dpt/__init__.py,sha256=gP1tkR3XMNlHq1GT87ugIVvb2o_1eAUg1JaniXjy1Lw,651
 optimum/rbln/transformers/models/dpt/modeling_dpt.py,sha256=ZsS2SOiqcA4azULB-WFEMQZbgIoOyVUKqVKqrw_tWzA,3430
 optimum/rbln/transformers/models/exaone/__init__.py,sha256=zYH_5tVa8-juEdsOIky7I33WSC3Zuhoq1upI0OHYeVw,859
-optimum/rbln/transformers/models/exaone/exaone_architecture.py,sha256=aPit1EOe3s3g3IVhztU1wydiTjYGA_j02btV9dl8W_I,3119
+optimum/rbln/transformers/models/exaone/exaone_architecture.py,sha256=ZM5vvz8KBipOiMVi8vqfvejkDSknW69xh4GrvJix-g0,3350
 optimum/rbln/transformers/models/exaone/modeling_exaone.py,sha256=WjyH8PmsMljSea7kJn_Cq1FJ96OXwXAoU7hv2Q8zUnI,1747
 optimum/rbln/transformers/models/gemma/__init__.py,sha256=7qUrekuBwCI9a6_Fq6j--FzCirRtUDz3ApY17mQS5Y4,648
-optimum/rbln/transformers/models/gemma/gemma_architecture.py,sha256=_GPIcSY5Q3PPuTehEseEf43mMBkW9Gl6pJlnHnjmkkM,2055
+optimum/rbln/transformers/models/gemma/gemma_architecture.py,sha256=bmCx405FVcffhgrQ53qMMZDbSlPxWOjucMHbvq19Gnw,2286
 optimum/rbln/transformers/models/gemma/modeling_gemma.py,sha256=-U3w3cEOv3ps1S8aL7uOq6Kq2siCPZz7Z8MXhDQgQqo,1530
 optimum/rbln/transformers/models/gpt2/__init__.py,sha256=UwwPPYVTB9ywDWy314L2bNL0i7wfkQFA71qjgXicEPg,646
-optimum/rbln/transformers/models/gpt2/gpt2_architecture.py,sha256=eN_UFcaaxWrtXvAGYck7J19Im2GZub2pwOejF3VWR6I,2934
+optimum/rbln/transformers/models/gpt2/gpt2_architecture.py,sha256=1IxqHmB-GlH2Dv2Yk4z0rMxL9CpxMGHhSu_x8_4cxvs,3008
 optimum/rbln/transformers/models/gpt2/modeling_gpt2.py,sha256=qBDanUk_O-HtOIVCA4IE3FYyCsnL9xIDK00vft-0caw,1490
 optimum/rbln/transformers/models/llama/__init__.py,sha256=jo_j_eIrHYGNEhR5lb6g3r5sO0ewe3fm2TWO8mLrT58,648
 optimum/rbln/transformers/models/llama/llama_architecture.py,sha256=S7MCPfyjG5eUqgaS-QNBB0ApUD6wnb5fR0RHq7k7-pA,728
@@ -78,23 +78,23 @@ optimum/rbln/transformers/models/llama/modeling_llama.py,sha256=Z3iony7icoFhRQ11
 optimum/rbln/transformers/models/llava_next/__init__.py,sha256=VLieyWm-UgvuNxw9B38wrL1Jsa09NBDX_ebABmdpTbs,670
 optimum/rbln/transformers/models/llava_next/modeling_llava_next.py,sha256=w_plsUOzxnhkQBhQeUqW9aJqGCvCvLtsx0XNKYjOprU,26203
 optimum/rbln/transformers/models/midm/__init__.py,sha256=UJSaErsF-z6dZERIS143WTaygffZyzEGqoQ2ZPDiM-c,855
-optimum/rbln/transformers/models/midm/midm_architecture.py,sha256=au6jHs7UQjthXDOrL7aqlOw7fkwM0-vkKkLGWeV1KKQ,5370
+optimum/rbln/transformers/models/midm/midm_architecture.py,sha256=357iviqQkzI0s_lU_teH1sVOChNRDUABe3GA0HuhZZY,5444
 optimum/rbln/transformers/models/midm/modeling_midm.py,sha256=GG25BozEZriAL-OPFGpzOjyDtSFB-NfeiLJTDAqxe20,1734
 optimum/rbln/transformers/models/mistral/__init__.py,sha256=jpGdNtRLoV7WmuYpRGVXR27BTC8RIi_nhmvYlxuhqRc,652
 optimum/rbln/transformers/models/mistral/mistral_architecture.py,sha256=_aU8TE_tdvfo0K7QpgTlz_d0qwk4O82dl9268lPL16E,733
 optimum/rbln/transformers/models/mistral/modeling_mistral.py,sha256=7nrddoBIHf8S12LZWBUpotnvG3gND11vMQda9yYXJ-s,1560
 optimum/rbln/transformers/models/phi/__init__.py,sha256=mZLt1M7BbYEvSon5UlkniMUPa15SfjZFdw0kMSAF3VA,644
 optimum/rbln/transformers/models/phi/modeling_phi.py,sha256=j-6Pqd5rR2JE8I1pnKFlCi4nW5Dv3wZjoPWxohissoo,1516
-optimum/rbln/transformers/models/phi/phi_architecture.py,sha256=QQBf5tlJDYuEHy8wLRpQW9vhYV3R6kr5OLTt4ZXrwl8,4039
+optimum/rbln/transformers/models/phi/phi_architecture.py,sha256=rBQjr6MOYBo1i5yLekMSR81TzYlHrHAA30kyKDdR7ww,4132
 optimum/rbln/transformers/models/qwen2/__init__.py,sha256=RAMWc21W_2I6DH9xBjeNxPECmAcTrbKhSIefq3Lass0,648
 optimum/rbln/transformers/models/qwen2/modeling_qwen2.py,sha256=9-aFDvjMzPNUyGOz0qo33RE18bUFGYZ3Wt_68zb5uJY,1530
 optimum/rbln/transformers/models/qwen2/qwen2_architecture.py,sha256=XlNAMYAcDLohnSAhIFGKOPuCB5XLgzYs5ABWdeQSaZs,720
 optimum/rbln/transformers/models/seq2seq/__init__.py,sha256=EmEMV4rOYqKyruX85d0fR73-b8N6BSD6CPcbpYdBuVk,651
-optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py,sha256=HG_-8ufRWIls67imU1547V0bk9FUWC0haOBL7eyRV6k,16365
-optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py,sha256=jPiRo2woijKd8VOHKb0qhBmy0vw4_WHQQh1JotlTx1w,18390
+optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py,sha256=NPfJf9Uk_bYOae7hXGHwteGiWH0va63Z-D93RmAMENg,17611
+optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py,sha256=QXIGWSu9PsKWE3WhkgmBj3VeszqXIo2MPOwcrb54Tbs,19348
 optimum/rbln/transformers/models/t5/__init__.py,sha256=1skR1RmnG62WTAP3-F5P1x-V_ReFhMyirH3u56vWwvc,675
-optimum/rbln/transformers/models/t5/modeling_t5.py,sha256=9AHRoGsr4eD_dIm1JA6ojafqIxd4w5Upzw3HmKOADkk,8049
-optimum/rbln/transformers/models/t5/t5_architecture.py,sha256=kkjErS42mW2jv5O_xL7BaKobvvqy7BGmYOowKyHakvI,7189
+optimum/rbln/transformers/models/t5/modeling_t5.py,sha256=nKRR3eH1EAu1YkKvhlqGyTrJXDRd-IWB5LOeG9jrcb4,16021
+optimum/rbln/transformers/models/t5/t5_architecture.py,sha256=oCdmF4eCTayAVjx3c-SVpmhrjnWE92jh79dMIYCwotY,9690
 optimum/rbln/transformers/models/wav2vec2/__init__.py,sha256=YpgA0K-vyg9veh0eL_jxauosbRpb_kpGKHvvQLBspKM,649
 optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py,sha256=JYJmV52j6cBwim4RanVJryfKnV80V96ol0A-oR6o7cg,3856
 optimum/rbln/transformers/models/whisper/__init__.py,sha256=ktnNe5ri3ycCWZ_W_voFB9y9-vgGgxS1X9s8LBRZmWc,665
@@ -114,7 +114,7 @@ optimum/rbln/utils/model_utils.py,sha256=DfD_Z2qvZHqcddXqnzTM1AN8khanj3-DXK2lJvV
 optimum/rbln/utils/runtime_utils.py,sha256=5-DYniyP59nx-mrrbi7AqA77L85b4Cm5oLpaxidSyss,3699
 optimum/rbln/utils/save_utils.py,sha256=hG5uOtYmecSXZuGTvCXsTM-SiyZpr5q3InUGCCq_jzQ,3619
 optimum/rbln/utils/submodule.py,sha256=oZoGrItB8WqY4i-K9WJPlLlcLohc1YGB9OHB8_XZw3A,4071
-optimum_rbln-0.7.3a1.dist-info/METADATA,sha256=59cCm0xXF4GfQ4oMuDLjOYHjkhUMqnZayBGMBstYd0Q,5300
-optimum_rbln-0.7.3a1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-optimum_rbln-0.7.3a1.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-optimum_rbln-0.7.3a1.dist-info/RECORD,,
+optimum_rbln-0.7.3a2.dist-info/METADATA,sha256=C-IWumO-veJFZPHpF8wcOTOE0TCDzKU1Xk_ylaqrvPM,5300
+optimum_rbln-0.7.3a2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+optimum_rbln-0.7.3a2.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+optimum_rbln-0.7.3a2.dist-info/RECORD,,