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

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

@@ -18,12 +18,6 @@ import torch
 from torch import nn
 from transformers.utils import logging
 
-from ....ops import (
-    register_rbln_custom_cache_update,
-    register_rbln_custom_paged_attention,
-    register_rbln_custom_paged_causal_attention,
-)
-
 
 logger = logging.get_logger(__name__)
 
@@ -59,7 +53,6 @@ class Seq2SeqEncoderWrapper(nn.Module):
 
     def __init__(self, model: nn.Module, enc_max_seq_len: int):
         super().__init__()
-        register_rbln_custom_cache_update()
         self.config = model.config
         self.encoder = model.get_encoder()
         self.encoder_max_length = enc_max_seq_len
@@ -90,8 +83,8 @@ class Seq2SeqEncoderWrapper(nn.Module):
         self,
         input_ids: torch.Tensor,
         attention_mask: torch.Tensor,
-        cross_key_values: torch.Tensor,
         b_idx: torch.Tensor,
+        *cross_key_values: Tuple[torch.Tensor],
     ) -> Tuple[torch.Tensor]:
         # 1. get encoder last_hidden_states
         encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
@@ -110,13 +103,15 @@ class Seq2SeqEncoderWrapper(nn.Module):
             cross_kv.append(past_k)
             cross_kv.append(past_v)
 
-        cross_kv = torch.stack(cross_kv, dim=0)
-
         # 3. update the cross_attention's past_key_value direct to the device-dram for optimization.
-        batch_axis = torch.tensor(1, dtype=torch.int16)
-        enc_out = torch.ops.rbln_custom_ops.rbln_cache_update(cross_key_values, cross_kv, b_idx[0], batch_axis)
+        batch_axis = torch.tensor(0, dtype=torch.int16)
+        cross_key_values = list(cross_key_values)
+        for i in range(self.n_layer * 2):
+            cross_key_values[i] = torch.ops.rbln_custom_ops.rbln_cache_update(
+                cross_key_values[i], cross_kv[i], b_idx[0], batch_axis
+            )
 
-        return enc_out
+        return cross_key_values
 
 
 class Seq2SeqDecoderWrapper(nn.Module):
@@ -146,11 +141,6 @@ 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.
         """
-        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)
 
@@ -176,16 +166,17 @@ class Seq2SeqDecoderWrapper(nn.Module):
                 encoder_attention_mask,
                 cache_position,
                 block_tables,
-                cross_kv_cache,
-                *self_kv_cache,
+                *kv_cache,
             ) = args
 
         else:
             attention_mask = None
-            (input_ids, encoder_attention_mask, cache_position, block_tables, cross_kv_cache, *self_kv_cache) = args
+            (input_ids, encoder_attention_mask, cache_position, block_tables, *kv_cache) = args
 
         self_past_key_values = ()
         cross_past_key_values = ()
+        self_kv_cache = kv_cache[self.num_layers * 2 :]
+        cross_kv_cache = kv_cache[: self.num_layers * 2]
         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]),)

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

@@ -12,4 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from ....ops import paged_add_softmax_attn_decode
+from .configuration_t5 import RBLNT5EncoderModelConfig, RBLNT5ForConditionalGenerationConfig
 from .modeling_t5 import RBLNT5EncoderModel, RBLNT5ForConditionalGeneration

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

@@ -0,0 +1,24 @@
+# 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 ...configuration_generic import RBLNTransformerEncoderForFeatureExtractionConfig
+from ..seq2seq import RBLNModelForSeq2SeqLMConfig
+
+
+class RBLNT5EncoderModelConfig(RBLNTransformerEncoderForFeatureExtractionConfig):
+    pass
+
+
+class RBLNT5ForConditionalGenerationConfig(RBLNModelForSeq2SeqLMConfig):
+    pass

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

@@ -13,106 +13,21 @@
 # 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
 
-import rebel
 import torch
-from transformers import (
-    AutoModelForTextEncoding,
-    PretrainedConfig,
-    T5EncoderModel,
-    T5ForConditionalGeneration,
-)
-from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
+from transformers import AutoModelForTextEncoding, T5EncoderModel, T5ForConditionalGeneration
 
-from ....diffusers.modeling_diffusers import RBLNDiffusionMixin
-from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
-from ....utils.logging import get_logger
-from ....utils.runtime_utils import RBLNPytorchRuntime
+from ...modeling_generic import RBLNTransformerEncoderForFeatureExtraction
 from ...models.seq2seq import RBLNModelForSeq2SeqLM
+from .configuration_t5 import RBLNT5EncoderModelConfig, RBLNT5ForConditionalGenerationConfig
 from .t5_architecture import T5Wrapper
 
 
-logger = get_logger()
-
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
-
-
-class RBLNRuntimeModel(RBLNPytorchRuntime):
-    def forward(
-        self,
-        input_ids: torch.LongTensor,
-        attention_mask: torch.FloatTensor,
-        head_mask: torch.FloatTensor,
-        inputs_embeds: torch.FloatTensor,
-        **kwargs,
-    ):
-        return super().forward(
-            input_ids,
-            attention_mask,
-            head_mask,
-            inputs_embeds,
-            **kwargs,
-        )
-
-
-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]))
-
+    from transformers import PreTrainedModel
 
-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)
+    from ....diffusers.modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
 
 
 class T5EncoderWrapper(torch.nn.Module):
@@ -125,149 +40,35 @@ class T5EncoderWrapper(torch.nn.Module):
         return self.model(*args, **kwargs, return_dict=False)
 
 
-class RBLNT5EncoderModel(RBLNModel):
+class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
     auto_model_class = AutoModelForTextEncoding
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
-    def __post_init__(self, **kwargs):
-        self.model = RBLNRuntimeModel(runtime=self.model[0])
-
     @classmethod
-    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNT5EncoderModelConfig):
         return T5EncoderWrapper(model)
 
     @classmethod
-    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        batch_size = rbln_config.get("batch_size", 1)
-        max_sequence_length = rbln_config.get("max_sequence_length", 256)
-        model_input_names = ["input_ids"]
-
-        rbln_config.update(
-            {
-                "batch_size": batch_size,
-                "max_seq_len": max_sequence_length,
-                "model_input_names": model_input_names,
-            }
-        )
-
-        return rbln_config
-
-    @classmethod
-    def _get_rbln_config(
+    def update_rbln_config_using_pipe(
         cls,
-        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
-        model_config: Optional["PretrainedConfig"] = None,
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
-        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-
-        max_position_embeddings = getattr(model_config, "n_positions", None)
-
-        if rbln_max_seq_len is None:
-            rbln_max_seq_len = max_position_embeddings
-            if rbln_max_seq_len is None:
-                for tokenizer in preprocessors:
-                    if hasattr(tokenizer, "model_max_length"):
-                        rbln_max_seq_len = tokenizer.model_max_length
-                        break
-                if rbln_max_seq_len is None:
-                    raise ValueError("`rbln_max_seq_len` should be specified!")
-
-        if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
-            raise ValueError("`rbln_max_seq_len` should be less or equal than max_position_embeddings!")
-
-        signature_params = inspect.signature(cls.get_hf_class().forward).parameters.keys()
-
-        if rbln_model_input_names is None:
-            for tokenizer in preprocessors:
-                if hasattr(tokenizer, "model_input_names"):
-                    rbln_model_input_names = [name for name in signature_params if name in tokenizer.model_input_names]
-
-                    invalid_params = set(rbln_model_input_names) - set(signature_params)
-                    if invalid_params:
-                        raise ValueError(f"Invalid model input names: {invalid_params}")
-                    break
-            if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
-                rbln_model_input_names = cls.rbln_model_input_names
-            elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
-                raise ValueError(
-                    "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
-                    f"and be sure to make the order of the inputs same as T5EncoderModel forward() arguments like ({list(signature_params)})"
-                )
-        else:
-            invalid_params = set(rbln_model_input_names) - set(signature_params)
-            if invalid_params:
-                raise ValueError(f"Invalid model input names: {invalid_params}")
-            rbln_model_input_names = [name for name in signature_params if name in rbln_model_input_names]
-
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        input_info = [
-            (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
-            for model_input_name in rbln_model_input_names
-        ]
-
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-
-        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        pipe: "RBLNDiffusionMixin",
+        rbln_config: "RBLNDiffusionMixinConfig",
+        submodule_name: str,
+    ) -> "RBLNDiffusionMixinConfig":
+        submodule_config = getattr(rbln_config, submodule_name)
+        submodule_config.max_seq_len = rbln_config.max_seq_len or 256
+        submodule_config.model_input_names = ["input_ids"]
         return rbln_config
 
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple[torch.FloatTensor], BaseModelOutput]:
-        encoder_outputs = self.model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            inputs_embeds=inputs_embeds,
-            head_mask=head_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-        if not return_dict:
-            return (encoder_outputs,)
-        else:
-            return BaseModelOutput(last_hidden_state=encoder_outputs)
-
 
 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"):
-        enc_max_seq_len = rbln_config.model_cfg["enc_max_seq_len"]
-        dec_max_seq_len = rbln_config.model_cfg["dec_max_seq_len"]
-
-        return T5Wrapper(model, enc_max_seq_len=enc_max_seq_len, dec_max_seq_len=dec_max_seq_len)
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNT5ForConditionalGenerationConfig):
+        return T5Wrapper(
+            model, enc_max_seq_len=rbln_config.enc_max_seq_len, dec_max_seq_len=rbln_config.dec_max_seq_len
+        )
 
     def __getattr__(self, __name: str) -> Any:
         def redirect(func):
@@ -279,139 +80,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,6 @@ import torch
 from torch import nn
 from transformers.utils import logging
 
-from ....ops import register_rbln_custom_add_softmax_attention
 from ..seq2seq.seq2seq_architecture import (
     Seq2SeqDecoder,
     Seq2SeqDecoderLayer,
@@ -55,7 +54,6 @@ class T5EncoderWrapper(Seq2SeqEncoderWrapper):
 
 class T5DecoderWrapper(Seq2SeqDecoderWrapper):
     def __post_init__(self, model, dec_max_seq_len: int = None):
-        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)
 
@@ -77,11 +75,13 @@ class T5DecoderWrapper(Seq2SeqDecoderWrapper):
         attention_mask,
         encoder_attention_mask,
         cache_position,
-        cross_kv_cache,
-        *self_kv_cache,
+        block_tables,
+        *kv_cache,
     ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor]]:
         self_past_key_values = ()
         cross_past_key_values = ()
+        self_kv_cache = kv_cache[self.num_layers * 2 :]
+        cross_kv_cache = kv_cache[: self.num_layers * 2]
 
         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]),)
@@ -95,6 +95,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 +163,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 +177,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 +187,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 +199,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,26 @@
+# 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 ....ops import paged_add_softmax_attn_decode, rbln_cache_update
+from .configuration_time_series_transformer import RBLNTimeSeriesTransformerForPredictionConfig
+from .modeling_time_series_transformers import RBLNTimeSeriesTransformerForPrediction

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

@@ -0,0 +1,34 @@
+from typing import Optional
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNTimeSeriesTransformerForPredictionConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        enc_max_seq_len: Optional[int] = None,
+        dec_max_seq_len: Optional[int] = None,
+        num_parallel_samples: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            enc_max_seq_len (Optional[int]): Maximum sequence length for the encoder.
+            dec_max_seq_len (Optional[int]): Maximum sequence length for the decoder.
+            num_parallel_samples (Optional[int]): Number of samples to generate in parallel during prediction.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size <= 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
+        self.enc_max_seq_len = enc_max_seq_len
+        self.dec_max_seq_len = dec_max_seq_len
+        self.num_parallel_samples = num_parallel_samples