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

optimum/rbln/__init__.py

@@ -73,6 +73,7 @@ _import_structure = {
         "RBLNRobertaForMaskedLM",
         "RBLNViTForImageClassification",
         "RBLNBertForMaskedLM",
+        "RBLNTimeSeriesTransformerForPrediction",
     ],
     "diffusers": [
         "RBLNAutoencoderKL",
@@ -184,6 +185,7 @@ if TYPE_CHECKING:
         RBLNRobertaForSequenceClassification,
         RBLNT5EncoderModel,
         RBLNT5ForConditionalGeneration,
+        RBLNTimeSeriesTransformerForPrediction,
         RBLNViTForImageClassification,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,

optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.3.post2'
-__version_tuple__ = version_tuple = (0, 7, 3)
+__version__ = version = '0.7.4a1'
+__version_tuple__ = version_tuple = (0, 7, 4)

optimum/rbln/ops/__init__.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 
 from .attn import (
-    register_rbln_custom_add_softmax_attention,
+    register_rbln_custom_paged_add_softmax_attention,
     register_rbln_custom_paged_attention,
     register_rbln_custom_paged_causal_attention,
 )
 from .flash_attn import register_rbln_custom_paged_flash_attention, register_rbln_custom_paged_flash_causal_attention
 from .kv_cache_update import register_rbln_custom_cache_update
+from .linear import linear

optimum/rbln/ops/attn.py

@@ -182,14 +182,14 @@ def register_rbln_custom_paged_causal_attention():
 
 
 @lru_cache
-def register_rbln_custom_add_softmax_attention():
+def register_rbln_custom_paged_add_softmax_attention():
     torch.library.define(
-        "rbln_custom_ops::add_softmax_attn_decode",
-        "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d) -> Tensor",
+        "rbln_custom_ops::paged_add_softmax_attn_decode",
+        "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e, int f) -> Tensor",
     )
 
-    @torch.library.impl("rbln_custom_ops::add_softmax_attn_decode", "cpu")
-    def add_softmax_attn_decode_cpu(q, k, v, mask, kcache, vcache, seq, scale):
+    @torch.library.impl("rbln_custom_ops::paged_add_softmax_attn_decode", "cpu")
+    def paged_add_softmax_attn_decode_cpu(q, k, v, mask, kcache, vcache, seq, scale, block_table, block_size):
         """Defines the computation pattern for fused attention with KV cache updates.
 
         IMPORTANT: This op serves as a pattern definition for the RBLN compiler to generate
@@ -210,12 +210,14 @@ def register_rbln_custom_add_softmax_attention():
         - vcache: [batch_size, n_heads, 1, max_seq_len, head_dim] - Value cache
         - seq: [1] - Current sequence position
         - scale: [] - Attention scale factor
+        - block_table: [batch_size, max_seq_len // block_size] - Block indices for KV cache management
+        - block_size: [] - Number of tokens per block
 
         Returns:
             Tensor: attn_output: [batch=1, n_heads, 1, 1, head_dim] - Attention output
         """
         return q
 
-    @register_fake("rbln_custom_ops::add_softmax_attn_decode")
-    def add_softmax_attn_decode_abstract(q, k, v, m, kcache, vcache, seq, partition):
+    @register_fake("rbln_custom_ops::paged_add_softmax_attn_decode")
+    def paged_add_softmax_attn_decode_abstract(q, k, v, m, kcache, vcache, seq, partition, block_table, block_size):
         return q

optimum/rbln/ops/linear.py

@@ -0,0 +1,25 @@
+# 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 Optional
+
+import torch
+from torch import Tensor
+
+
+@torch.library.custom_op("rbln_custom_ops::linear", mutates_args=())
+def linear(input: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
+    output_shape = list(input.shape[:-1])
+    output_shape += [weight.shape[0]]
+    return torch.empty(size=output_shape, dtype=input.dtype, device=input.device, requires_grad=input.requires_grad)

optimum/rbln/transformers/__init__.py

@@ -52,6 +52,7 @@ _import_structure = {
         "RBLNPhiForCausalLM",
         "RBLNT5EncoderModel",
         "RBLNT5ForConditionalGeneration",
+        "RBLNTimeSeriesTransformerForPrediction",
         "RBLNLlavaNextForConditionalGeneration",
         "RBLNMidmLMHeadModel",
         "RBLNXLMRobertaModel",
@@ -113,6 +114,7 @@ if TYPE_CHECKING:
         RBLNQwen2ForCausalLM,
         RBLNT5EncoderModel,
         RBLNT5ForConditionalGeneration,
+        RBLNTimeSeriesTransformerForPrediction,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
         RBLNXLMRobertaModel,

optimum/rbln/transformers/models/__init__.py

@@ -50,6 +50,7 @@ _import_structure = {
     "mistral": ["RBLNMistralForCausalLM"],
     "phi": ["RBLNPhiForCausalLM"],
     "qwen2": ["RBLNQwen2ForCausalLM"],
+    "time_series_transformers": ["RBLNTimeSeriesTransformerForPrediction"],
     "t5": ["RBLNT5EncoderModel", "RBLNT5ForConditionalGeneration"],
     "wav2vec2": ["RBLNWav2Vec2ForCTC"],
     "whisper": ["RBLNWhisperForConditionalGeneration"],
@@ -90,6 +91,7 @@ if TYPE_CHECKING:
     from .phi import RBLNPhiForCausalLM
     from .qwen2 import RBLNQwen2ForCausalLM
     from .t5 import RBLNT5EncoderModel, RBLNT5ForConditionalGeneration
+    from .time_series_transformers import RBLNTimeSeriesTransformerForPrediction
     from .wav2vec2 import RBLNWav2Vec2ForCTC
     from .whisper import RBLNWhisperForConditionalGeneration
     from .xlm_roberta import RBLNXLMRobertaModel

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

@@ -94,12 +94,11 @@ class RBLNBartModel(RBLNModel):
             for model_input_name in rbln_model_input_names
         ]
 
-        enc_compile_config = RBLNCompileConfig(input_info=input_info, compiled_model_name="encoder")
-        dec_compile_config = RBLNCompileConfig(input_info=input_info, compiled_model_name="decoder")
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
 
         rbln_config = RBLNConfig(
             rbln_cls=cls.__name__,
-            compile_cfgs=[enc_compile_config, dec_compile_config],
+            compile_cfgs=[rbln_compile_config],
             rbln_kwargs=rbln_kwargs,
         )
 
@@ -108,6 +107,8 @@ class RBLNBartModel(RBLNModel):
 
 
 class RBLNBartForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    support_causal_attn = True
+
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
         enc_max_seq_len = (

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

@@ -222,8 +222,6 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
 
             attention_mask = self.dec_attn_mask
 
-            attention_mask = self.dec_attn_mask
-
         logits = super().forward(
             inputs,
             cache_position,
@@ -547,22 +545,27 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
         @QuantizationManager.with_quantization_env
         def compile_model(*args, **kwargs):
-            wrapped_model.phase = "prefill"
-            compiled_prefill = RBLNModel.compile(
-                wrapped_model,
-                prefill_compile_config,
-                example_inputs=prefill_example_inputs,
-                compile_context=context,
-            )
+            try:
+                original_linear = torch.nn.functional.linear
+                torch.nn.functional.linear = torch.ops.rbln_custom_ops.linear
+                wrapped_model.phase = "prefill"
+                compiled_prefill = RBLNModel.compile(
+                    wrapped_model,
+                    prefill_compile_config,
+                    example_inputs=prefill_example_inputs,
+                    compile_context=context,
+                )
 
-            wrapped_model.phase = "decode"
-            compiled_decoder = RBLNModel.compile(
-                wrapped_model,
-                dec_compile_config,
-                example_inputs=dec_example_inputs,
-                compile_context=context,
-            )
-            return {"prefill": compiled_prefill, "decoder": compiled_decoder}
+                wrapped_model.phase = "decode"
+                compiled_decoder = RBLNModel.compile(
+                    wrapped_model,
+                    dec_compile_config,
+                    example_inputs=dec_example_inputs,
+                    compile_context=context,
+                )
+                return {"prefill": compiled_prefill, "decoder": compiled_decoder}
+            finally:
+                torch.nn.functional.linear = original_linear
 
         return compile_model(quantize_config=quantize_config)
 

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

@@ -38,8 +38,8 @@ 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]))
+        output = super().forward(*args, **kwargs)
+        return BaseModelOutput(last_hidden_state=output)
 
 
 class RBLNRuntimeDecoder(RBLNPytorchRuntime):
@@ -94,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)
@@ -115,6 +115,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
 
     main_input_name = "input_ids"
     auto_model_class = AutoModelForSeq2SeqLM
+    support_causal_attn = None
 
     def __post_init__(self, **kwargs):
         batch_size = self.rbln_config.model_cfg["batch_size"]
@@ -186,13 +187,16 @@ 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 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
+        if cls.support_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
 
         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")
@@ -265,11 +269,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
                 [rbln_batch_size, 1],
                 "int32",
             ),
-            (
-                "block_tables",
-                [rbln_batch_size, 1],
-                "int16",
-            ),
+            ("block_tables", [rbln_batch_size, 1], "int16"),
         ]
         dec_input_info.extend(
             [

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

@@ -13,9 +13,8 @@
 # limitations under the License.
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
 
-import rebel
 import torch
 from transformers import (
     AutoModelForTextEncoding,
@@ -23,7 +22,7 @@ from transformers import (
     T5EncoderModel,
     T5ForConditionalGeneration,
 )
-from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
+from transformers.modeling_outputs import BaseModelOutput
 
 from ....diffusers.modeling_diffusers import RBLNDiffusionMixin
 from ....modeling import RBLNModel
@@ -58,63 +57,6 @@ 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__()
@@ -247,20 +189,7 @@ 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,
-        )
+    support_causal_attn = False
 
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
@@ -279,139 +208,3 @@ 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_add_softmax_attention
+from ....ops import register_rbln_custom_paged_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_add_softmax_attention()
+        register_rbln_custom_paged_add_softmax_attention()
         self.num_layers = self.config.num_layers
         self.conditional_generation = self.convert_to_rbln_conditional_generation(model, dec_max_seq_len)
 
@@ -77,6 +77,7 @@ class T5DecoderWrapper(Seq2SeqDecoderWrapper):
         attention_mask,
         encoder_attention_mask,
         cache_position,
+        block_tables,
         cross_kv_cache,
         *self_kv_cache,
     ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor]]:
@@ -95,6 +96,7 @@ class T5DecoderWrapper(Seq2SeqDecoderWrapper):
             self_past_key_values=self_past_key_values,
             cross_past_key_values=cross_past_key_values,
             cache_position=cache_position,
+            block_tables=block_tables,
         )
 
         return lm_logits
@@ -162,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.add_softmax_attn_decode
+        self.attn_decode = torch.ops.rbln_custom_ops.paged_add_softmax_attn_decode
 
     def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         query_states = self.q_proj(hidden_states)
@@ -176,6 +178,7 @@ class T5LayerSelfAttention(Seq2SeqSelfAttention):
         past_key_value: Tuple[torch.Tensor],
         attention_mask: torch.Tensor,
         cache_position: torch.Tensor,
+        block_tables: torch.Tensor,
         **kwargs,
     ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
         bsz, tgt_len, _ = hidden_states.size()
@@ -185,6 +188,7 @@ class T5LayerSelfAttention(Seq2SeqSelfAttention):
         key_states = self._shape(key_states, -1, bsz)
         value_states = self._shape(value_states, -1, bsz)
 
+        block_size = past_key_value[0].shape[-2]
         attn_output = self.attn_decode(
             query_states,
             key_states,
@@ -196,6 +200,8 @@ class T5LayerSelfAttention(Seq2SeqSelfAttention):
             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,
         )
 
         attn_output = attn_output.view(bsz, self.num_heads, -1, self.head_dim).transpose(1, 2)

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

@@ -0,0 +1,24 @@
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from .modeling_time_series_transformers import RBLNTimeSeriesTransformerForPrediction