optimum-rbln 0.7.3.post1__py3-none-any.whl → 0.7.3.post2__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.7.3.post1'
+__version__ = version = '0.7.3.post2'
 __version_tuple__ = version_tuple = (0, 7, 3)

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

@@ -108,8 +108,6 @@ class RBLNBartModel(RBLNModel):
 
 
 class RBLNBartForConditionalGeneration(RBLNModelForSeq2SeqLM):
-    support_paged_causal_attn = True
-
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
         enc_max_seq_len = (

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

@@ -50,7 +50,6 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         runtime: rebel.Runtime,
         batch_size: int,
         dec_max_seq_len: int,
-        support_paged_causal_attn: Optional[bool] = None,
         use_attention_mask: Optional[bool] = None,
         **kwargs: Any,
     ) -> None:
@@ -58,10 +57,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         self.batch_size = batch_size
         self.dec_max_seq_len = dec_max_seq_len
         self.use_attention_mask = use_attention_mask
-        if support_paged_causal_attn:
-            self.default_block_tables = torch.arange(0, self.batch_size, dtype=torch.int16).view(self.batch_size, 1)
-        else:
-            self.default_block_tables = None
+        self.default_block_tables = torch.arange(0, self.batch_size, dtype=torch.int16).view(self.batch_size, 1)
 
     def forward(
         self,
@@ -98,7 +94,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
             decoder_attention_mask if self.use_attention_mask else None,
             attention_mask,
             cache_position,
-            block_tables=block_tables,
+            block_tables,
         )
 
         return Seq2SeqLMOutput(logits=lm_logits)
@@ -119,7 +115,6 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
 
     main_input_name = "input_ids"
     auto_model_class = AutoModelForSeq2SeqLM
-    support_paged_causal_attn = None
 
     def __post_init__(self, **kwargs):
         batch_size = self.rbln_config.model_cfg["batch_size"]
@@ -135,7 +130,6 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             main_input_name="input_ids",
             batch_size=batch_size,
             dec_max_seq_len=dec_max_seq_len,
-            support_paged_causal_attn=self.support_paged_causal_attn,
             use_attention_mask=self.use_attention_mask,
         )
 
@@ -192,16 +186,13 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         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
+        rbln_use_attention_mask = rbln_kwargs.get("use_attention_mask", None)
 
-        if cls.support_paged_causal_attn:
-            rbln_use_attention_mask = rbln_kwargs.get("use_attention_mask", None)
-            if rbln_use_attention_mask is None:
-                rbln_use_attention_mask = False
-                rbln_npu = rbln_kwargs.get("npu", None) or rebel.get_npu_name()
-                if rbln_npu == "RBLN-CA02":
-                    rbln_use_attention_mask = True
-        else:
-            rbln_use_attention_mask = True
+        if rbln_use_attention_mask is None:
+            rbln_use_attention_mask = False
+            rbln_npu = rbln_kwargs.get("npu", None) or rebel.get_npu_name()
+            if rbln_npu == "RBLN-CA02":
+                rbln_use_attention_mask = True
 
         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")
@@ -274,6 +265,11 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
                 [rbln_batch_size, 1],
                 "int32",
             ),
+            (
+                "block_tables",
+                [rbln_batch_size, 1],
+                "int16",
+            ),
         ]
         dec_input_info.extend(
             [
@@ -306,8 +302,6 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             ]
         )
 
-        if cls.support_paged_causal_attn:
-            dec_input_info.insert(3, ("block_tables", [rbln_batch_size, 1], "int16"))
         if rbln_use_attention_mask:
             dec_input_info.insert(1, ("attention_mask", [rbln_batch_size, rbln_dec_max_seq_len], "float32"))
 

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,7 +247,20 @@ class RBLNT5EncoderModel(RBLNModel):
 
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
-    support_causal_paged_attn = False
+    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"):
@@ -208,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-0.7.3.post1.dist-info → optimum_rbln-0.7.3.post2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.7.3.post1
+Version: 0.7.3.post2
 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.3.post1.dist-info → optimum_rbln-0.7.3.post2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 optimum/rbln/__init__.py,sha256=ZDzXcl-oAcYJhKjJMpotjbTih9awo7HzUb6T3MUEP6Q,6894
-optimum/rbln/__version__.py,sha256=aegWGVZeZJ9bIegWWNAgPL2y9SAs5kPTsXCQi0EZ9go,517
+optimum/rbln/__version__.py,sha256=OJRzB6Y7xaNgH7EkerbumPEoY0Nlzs1_xYhBJvOXTzQ,517
 optimum/rbln/modeling.py,sha256=nJsAs5zs--VVOYGFjYNpqfxYIemJIK4Lr0WEzlDLdP0,8390
 optimum/rbln/modeling_base.py,sha256=dNCL-BhrWCpuOVkZaj8-MW567Tf4lLo3p3Z3ldjWJfU,21779
 optimum/rbln/modeling_config.py,sha256=7104bxmrvKW4Q6XTruQayiIGl8GHDFmPkJ3cknMIInE,11335
@@ -55,7 +55,7 @@ optimum/rbln/transformers/models/auto/auto_factory.py,sha256=IK9jFrJ3EEzYQa9_aKp
 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=Oo-Cdne7igKEex8wwP-gztKJHgs5GLHQjK1oc3IZIDE,5801
-optimum/rbln/transformers/models/bart/modeling_bart.py,sha256=6IpWXlBCd02v66KF77oEWfrv8-FnPBYjjjL_8KZL3Ow,5835
+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
@@ -92,10 +92,10 @@ optimum/rbln/transformers/models/qwen2/__init__.py,sha256=RAMWc21W_2I6DH9xBjeNxP
 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=9Pf9Y86ABDfhwIenlZqYfgqjbyFmtKBiPnbCD7zxw4M,18017
+optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py,sha256=NPfJf9Uk_bYOae7hXGHwteGiWH0va63Z-D93RmAMENg,17611
 optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py,sha256=tvzacIZam1sIr_1BvvZ_fDr8u5dXAiYiynFdX9tArtY,18877
 optimum/rbln/transformers/models/t5/__init__.py,sha256=1skR1RmnG62WTAP3-F5P1x-V_ReFhMyirH3u56vWwvc,675
-optimum/rbln/transformers/models/t5/modeling_t5.py,sha256=8PAhPlYT1dmpcWM7hUMmZV9lPd4d75CuMuFen1pzr3Q,8088
+optimum/rbln/transformers/models/t5/modeling_t5.py,sha256=nKRR3eH1EAu1YkKvhlqGyTrJXDRd-IWB5LOeG9jrcb4,16021
 optimum/rbln/transformers/models/t5/t5_architecture.py,sha256=AArCQhZRETVM583wlIRzMFOSYq7t2nzxaAeyhZxyxKk,9508
 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
@@ -116,7 +116,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.3.post1.dist-info/METADATA,sha256=dKER74SsqGQwVQgTXVM854y97xzhfRl5LKaGedd4IIw,5304
-optimum_rbln-0.7.3.post1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-optimum_rbln-0.7.3.post1.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-optimum_rbln-0.7.3.post1.dist-info/RECORD,,
+optimum_rbln-0.7.3.post2.dist-info/METADATA,sha256=YgOp5SEpJ_VfYEohAoBhSQ20TaX1usvkRAzV7s7mS5I,5304
+optimum_rbln-0.7.3.post2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+optimum_rbln-0.7.3.post2.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+optimum_rbln-0.7.3.post2.dist-info/RECORD,,