optimum-rbln 0.8.2a4__py3-none-any.whl → 0.9.3__py3-none-any.whl

optimum/rbln/transformers/models/qwen3/qwen3_architecture.py

@@ -12,16 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import torch
-import torch.nn as nn
-from transformers import PreTrainedModel
 
-from ..decoderonly.decoderonly_architecture import (
-    DecoderOnlyAttention,
-    DecoderOnlyLayer,
-    DecoderOnlyWrapper,
-    RotaryEmbedding,
-)
+from ..decoderonly.decoderonly_architecture import DecoderOnlyAttention, DecoderOnlyWrapper
 
 
 class Qwen3Wrapper(DecoderOnlyWrapper):
@@ -37,239 +29,3 @@ class Qwen3Attention(DecoderOnlyAttention):
         self.o_proj = self._original_mod.o_proj
         self.q_norm = self._original_mod.q_norm
         self.k_norm = self._original_mod.k_norm
-
-
-class Qwen3ModelWrapper(nn.Module):
-    def __init__(
-        self,
-        model,
-        attn_impl=None,
-        use_inputs_embeds=None,
-        use_attention_mask=None,
-        use_rotary_emb=None,
-        cache_impl=None,
-        kvcache_partition_len=None,
-        max_seq_len=None,
-        kvcache_block_size=None,
-        sliding_window=None,
-        sliding_window_layers=None,
-    ):
-        super().__init__()
-        self.config = model.config
-
-        if use_rotary_emb:
-            rotary_embs = self.get_rotary_emb(max_seq_len=max_seq_len)
-            if isinstance(rotary_embs, tuple):
-                self.rotary_emb_global, self.rotary_emb_local = rotary_embs
-            else:
-                self.rotary_emb = rotary_embs
-        else:
-            self.rotary_emb = None
-
-        self._original_mod = model
-        self.use_inputs_embeds = use_inputs_embeds
-        self.attn_impl = attn_impl
-        self.cache_impl = cache_impl
-        self.use_attention_mask = use_attention_mask
-        self.kvcache_partition_len = kvcache_partition_len
-        self.kvcache_block_size = kvcache_block_size
-        self.max_seq_len = max_seq_len
-        self.sliding_window = sliding_window
-        self.sliding_window_layers = sliding_window_layers
-        self.model = self.convert_to_rbln_model(model)
-
-    def get_rotary_emb(self, max_seq_len):
-        return RotaryEmbedding(config=self.config, max_seq_len_cached=max_seq_len)
-
-    def convert_to_rbln_model(self, base_model: PreTrainedModel):
-        for layer_idx, layer in enumerate(base_model.layers):
-            is_sliding = layer_idx in self.sliding_window_layers
-            new_self_attn = Qwen3Attention(
-                layer.self_attn,
-                self.use_attention_mask if not is_sliding else True,
-                use_position_ids=None,
-                kvcache_block_size=self.sliding_window
-                if layer_idx in self.sliding_window_layers
-                else self.kvcache_block_size,
-                is_sliding=is_sliding,
-                attn_impl=self.attn_impl if not is_sliding else "eager",
-                kvcache_partition_len=self.kvcache_partition_len,
-            )
-            base_model.layers[layer_idx] = DecoderOnlyLayer(layer, new_self_attn)
-
-        return base_model
-
-    @property
-    def hidden_multiplier(self):
-        return 1
-
-    def get_last_layernorm(self) -> nn.LayerNorm:
-        return self._original_mod.norm
-
-    def get_embedding(self) -> nn.Embedding:
-        return self._original_mod.embed_tokens
-
-    def get_pos_embedding(self) -> nn.Embedding:
-        raise NotImplementedError(
-            "The 'get_pos_embedding' method is not implemented. Please define this method in a subclass."
-        )
-
-    def convert_sequence_positions_for_flash_attn(self, seq_positions, max_seq_len):
-        if self.attn_impl not in ["flash_attn"]:
-            raise NotImplementedError(f"Unknown attn_impl ({self.attn_impl}).")
-        partition_len = self.kvcache_partition_len
-        num_partition = max_seq_len // partition_len
-
-        cs = seq_positions.repeat(num_partition, 1).transpose(0, 1)
-        pidx = torch.arange(num_partition)
-        cache_pos_for_partitions = torch.clamp(cs - pidx * partition_len, 0, partition_len)
-        return cache_pos_for_partitions
-
-    def get_local_cache_positions(self, position_ids, query_position):
-        max_cache_len = self.model.config.sliding_window
-        valid_input_len = 1 if query_position is None else query_position + 1
-        cache_seq_len = torch.clamp(position_ids, max=max_cache_len)[:, :1]  # past seen tokens
-        cache_offset = (
-            torch.clamp(position_ids, max=max_cache_len)[:, :1] + valid_input_len
-        )  # cache offset for next steps
-
-        return cache_seq_len, cache_offset
-
-    def prepare_forward_args(self, *args):
-        args = list(args)
-        input_ids = None if self.use_inputs_embeds else args.pop(0)
-        inputs_embeds = args.pop(0) if self.use_inputs_embeds else None
-        cache_position = args.pop(0)
-        global_block_tables = args.pop(0) if self.cache_impl in ["hybrid", "static"] else None
-        local_block_tables = args.pop(0) if self.cache_impl in ["hybrid", "sliding_window"] else None
-        query_position = args.pop(0) if self.sliding_window else None
-        attention_mask = args.pop(0) if self.use_attention_mask else None
-        position_ids = None
-        past_key_values = args
-
-        if len(past_key_values) != 2 * self.config.num_hidden_layers:
-            raise ValueError(
-                f"Different past_key_values to model's config. {len(past_key_values)} != {2 * self.config.num_hidden_layers}"
-            )
-
-        # [key, value] * n_layer -> ( (key, value) ) * n_layer
-        # cache shape : batch, n_heads, 1, max_seq_len, head_dim
-        _past_key_values = []
-        for i in range(self.config.num_hidden_layers):
-            key_states = past_key_values[i * 2]
-            value_states = past_key_values[i * 2 + 1]
-            past_key_value = [key_states, value_states]
-            _past_key_values.append(past_key_value)
-        past_key_values = _past_key_values
-
-        if hasattr(self, "rotary_emb_global") and hasattr(self, "rotary_emb_local"):
-            rotary_emb = (self.rotary_emb_global, self.rotary_emb_local)
-        else:
-            rotary_emb = self.rotary_emb
-
-        return (
-            input_ids,
-            inputs_embeds,
-            cache_position,
-            global_block_tables,
-            local_block_tables,
-            attention_mask,
-            position_ids,
-            query_position,
-            past_key_values,
-            rotary_emb,
-        )
-
-    def forward(self, *args):
-        (
-            input_ids,
-            inputs_embeds,
-            cache_position,
-            global_block_tables,
-            local_block_tables,
-            attention_mask,
-            position_ids,
-            query_position,
-            past_key_values,
-            rotary_emb,
-        ) = self.prepare_forward_args(*args)
-
-        # retrieve input_ids and inputs_embeds
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError(
-                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
-            )
-
-        # embed positions
-        if inputs_embeds is None:
-            inputs_embeds = self.get_embedding()(input_ids)
-
-        hidden_states = inputs_embeds * self.hidden_multiplier
-
-        # get cos,sin vector if needed
-        position_ids = position_ids if position_ids is not None else cache_position
-        if rotary_emb is not None:
-            if isinstance(rotary_emb, torch.Tensor):
-                cos = rotary_emb[0]
-                sin = rotary_emb[1]
-            else:
-                cos, sin = rotary_emb(hidden_states, self.max_seq_len)  # dtype carrier, max_seq_len
-                cos, sin = slice_and_unsqueeze_cos_sin(cos, sin, position_ids)
-        else:
-            batch_size = inputs_embeds.shape[0]
-            if position_ids.shape[0] > 1:
-                position_embeds = []
-                for b_idx in range(batch_size):
-                    position_embed = self.get_pos_embedding()(position_ids[b_idx])
-                    position_embeds.append(position_embed)
-
-                position_embeds = torch.cat(position_embeds, dim=0).unsqueeze(1)
-            else:
-                position_embeds = self.get_pos_embedding()(position_ids)
-            hidden_states = hidden_states + position_embeds
-            cos, sin = None, None
-
-        # Get sequence positions for flash attention
-        if self.attn_impl == "flash_attn":
-            seq_positions = cache_position[:, 0]
-            seq_positions = self.convert_sequence_positions_for_flash_attn(
-                seq_positions=seq_positions, max_seq_len=self.max_seq_len
-            )
-        else:
-            seq_positions = cache_position[:, :1]
-
-        # Get local cache positions for sliding window layers
-        if len(self.sliding_window_layers) > 0:
-            sliding_cache_pos = self.get_local_cache_positions(position_ids, query_position)
-
-        for layer_idx, layer in enumerate(self.model.layers):
-            is_sliding = True if layer_idx in self.sliding_window_layers else False
-            hidden_states = layer(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                seq_positions=sliding_cache_pos if is_sliding else seq_positions,
-                past_key_values=past_key_values,
-                cos=cos,
-                sin=sin,
-                block_tables=local_block_tables if is_sliding else global_block_tables,
-            )
-
-        hidden_states = self.get_last_layernorm()(hidden_states)
-        return hidden_states
-
-
-def slice_and_unsqueeze_cos_sin(cos, sin, cache_position, unsqueeze_dim=1):
-    """Slice cos[cache_position], sin[cache_position] vector for the query."""
-    if cache_position.shape[0] > 1:
-        cos_all = []
-        sin_all = []
-        for i in range(cache_position.shape[0]):
-            cos_all.append(cos[cache_position[i : i + 1]].unsqueeze(unsqueeze_dim))
-            sin_all.append(sin[cache_position[i : i + 1]].unsqueeze(unsqueeze_dim))
-        cos = torch.cat(cos_all, dim=0)
-        sin = torch.cat(sin_all, dim=0)
-    else:
-        cos = cos[cache_position].unsqueeze(unsqueeze_dim)
-        sin = sin[cache_position].unsqueeze(unsqueeze_dim)
-
-    return cos, sin

optimum/rbln/transformers/models/resnet/configuration_resnet.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 
+from typing import Optional
+
 from ...configuration_generic import RBLNModelForImageClassificationConfig
 
 
@@ -23,3 +25,18 @@ class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConf
     This configuration class stores the configuration parameters specific to
     RBLN-optimized ResNet models for image classification tasks.
     """
+
+    def __init__(self, output_hidden_states: Optional[bool] = None, **kwargs):
+        """
+        Args:
+            image_size (Optional[Union[int, Tuple[int, int]]]): The size of input images.
+                Can be an integer for square images or a tuple (height, width).
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            output_hidden_states (bool, optional) — Whether or not to return the hidden states of all layers.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.output_hidden_states = output_hidden_states

optimum/rbln/transformers/models/resnet/modeling_resnet.py

@@ -13,7 +13,17 @@
 # limitations under the License.
 
 
+from typing import TYPE_CHECKING, Optional, Tuple, Union
+
+import torch
+from transformers.modeling_outputs import ImageClassifierOutputWithNoAttention
+
 from ...modeling_generic import RBLNModelForImageClassification
+from .configuration_resnet import RBLNResNetForImageClassificationConfig
+
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig, PreTrainedModel
 
 
 class RBLNResNetForImageClassification(RBLNModelForImageClassification):
@@ -24,3 +34,66 @@ class RBLNResNetForImageClassification(RBLNModelForImageClassification):
     on RBLN devices, supporting image classification with convolutional neural networks
     designed for computer vision tasks.
     """
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
+        model: Optional["PreTrainedModel"] = None,
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_config: Optional["RBLNResNetForImageClassificationConfig"] = None,
+    ) -> "RBLNResNetForImageClassificationConfig":
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
+
+        rbln_config = super()._update_rbln_config(
+            preprocessors=preprocessors,
+            model=model,
+            model_config=model_config,
+            rbln_config=rbln_config,
+        )
+
+        return rbln_config
+
+    @classmethod
+    def _wrap_model_if_needed(
+        cls, model: torch.nn.Module, rbln_config: "RBLNResNetForImageClassificationConfig"
+    ) -> torch.nn.Module:
+        class _ResNetForImageClassification(torch.nn.Module):
+            def __init__(self, model: torch.nn.Module, output_hidden_states: bool):
+                super().__init__()
+                self.model = model
+                self.output_hidden_states = output_hidden_states
+
+            def forward(self, *args, **kwargs):
+                output = self.model(*args, output_hidden_states=self.output_hidden_states, **kwargs)
+                return output
+
+        return _ResNetForImageClassification(model, rbln_config.output_hidden_states)
+
+    def forward(
+        self, pixel_values: torch.Tensor, output_hidden_states: bool = None, return_dict: bool = None, **kwargs
+    ) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:
+        """
+        Foward pass for the RBLN-optimized ResNet model for image classification.
+
+        Args:
+            pixel_values (torch.FloatTensor of shape (batch_size, channels, height, width)): The tensors corresponding to the input images.
+            output_hidden_states (bool, *optional*, defaults to False): Whether or not to return the hidden states of all layers.
+                See hidden_states under returned tensors for more details.
+            return_dict (bool, *optional*, defaults to True): Whether to return a dictionary of outputs.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a ImageClassifierOutputWithNoAttention object.
+        """
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        if output_hidden_states != self.rbln_config.output_hidden_states:
+            raise ValueError(
+                f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
+                f"Please compile again with the correct argument."
+            )
+
+        return super().forward(pixel_values=pixel_values, return_dict=return_dict, **kwargs)

optimum/rbln/transformers/models/roberta/modeling_roberta.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Tuple, Union
+
+import torch
+from transformers.modeling_outputs import MaskedLMOutput, SequenceClassifierOutput
+
 from ...modeling_generic import RBLNModelForMaskedLM, RBLNModelForSequenceClassification
 
 
@@ -26,6 +31,19 @@ class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
 
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> Union[Tuple, MaskedLMOutput]:
+        """
+        Forward pass for the RBLN-optimized RoBERTa model for masked language modeling tasks.
+
+        Args:
+            input_ids (torch.LongTensor of shape (batch_size, sequence_length), optional): Indices of input sequence tokens in the vocabulary.
+            attention_mask (torch.FloatTensor of shape (batch_size, sequence_length), optional): Mask to avoid performing attention on padding token indices.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a MaskedLMOutput object.
+        """
+        return super().forward(input_ids, attention_mask, **kwargs)
+
 
 class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     """
@@ -37,3 +55,18 @@ class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     """
 
     rbln_model_input_names = ["input_ids", "attention_mask"]
+
+    def forward(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs
+    ) -> Union[Tuple, SequenceClassifierOutput]:
+        """
+        Forward pass for the RBLN-optimized RoBERTa model for sequence classification tasks.
+
+        Args:
+            input_ids (torch.LongTensor of shape (batch_size, sequence_length), optional): Indices of input sequence tokens in the vocabulary.
+            attention_mask (torch.FloatTensor of shape (batch_size, sequence_length), optional): Mask to avoid performing attention on padding token indices.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a SequenceClassifierOutput object.
+        """
+        return super().forward(input_ids, attention_mask, **kwargs)

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

@@ -12,11 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional
-
-import rebel
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.deprecation import deprecate_kwarg
 from ....utils.logging import get_logger
 
 
@@ -24,14 +23,18 @@ logger = get_logger()
 
 
 class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
+    support_paged_attention = None
+
+    @deprecate_kwarg(old_name="pad_token_id", version="0.10.0")
     def __init__(
         self,
         batch_size: Optional[int] = None,
         enc_max_seq_len: Optional[int] = None,
         dec_max_seq_len: Optional[int] = None,
         use_attention_mask: Optional[bool] = None,
-        pad_token_id: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        kvcache_num_blocks: Optional[int] = None,
+        kvcache_block_size: Optional[int] = None,
+        **kwargs: Any,
     ):
         """
         Args:
@@ -39,9 +42,11 @@ class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
             enc_max_seq_len (Optional[int]): Maximum sequence length for the encoder.
             dec_max_seq_len (Optional[int]): Maximum sequence length for the decoder.
             use_attention_mask (Optional[bool]): Whether to use attention masks during inference.
-                This is automatically set to True for RBLN-CA02 devices.
-            pad_token_id (Optional[int]): The ID of the padding token in the vocabulary.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kvcache_num_blocks (Optional[int]): The total number of blocks to allocate for the
+                PagedAttention KV cache for the SelfAttention. Defaults to batch_size.
+            kvcache_block_size (Optional[int]): Sets the size (in number of tokens) of each block
+                in the PagedAttention KV cache for the SelfAttention. Defaults to dec_max_seq_len.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.
@@ -55,12 +60,12 @@ class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
         self.dec_max_seq_len = dec_max_seq_len
 
         self.use_attention_mask = use_attention_mask
-        npu = self.npu or rebel.get_npu_name()
-        if npu == "RBLN-CA02":
-            if self.use_attention_mask is False:
-                logger.warning("Attention mask should be used with RBLN-CA02. Setting use_attention_mask to True.")
-            self.use_attention_mask = True
-        else:
-            self.use_attention_mask = self.use_attention_mask or False
 
-        self.pad_token_id = pad_token_id
+        if self.support_paged_attention:
+            self.kvcache_num_blocks = kvcache_num_blocks
+            self.kvcache_block_size = kvcache_block_size
+        else:
+            if kvcache_num_blocks is not None or kvcache_block_size is not None:
+                raise ValueError(
+                    "You cannot set kvcache_num_blocks or kvcache_block_size as paged attention is not supported for the model."
+                )

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

@@ -20,7 +20,9 @@ import rebel
 import torch
 from rebel.compile_context import CompileContext
 from transformers import AutoModelForSeq2SeqLM, PretrainedConfig, PreTrainedModel
-from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutput, ModelOutput, Seq2SeqLMOutput
 
 from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
@@ -32,13 +34,13 @@ from .configuration_seq2seq import RBLNModelForSeq2SeqLMConfig
 logger = get_logger(__name__)
 
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, GenerationConfig, PretrainedConfig
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
 
 
 class RBLNRuntimeEncoder(RBLNPytorchRuntime):
     mandatory_members = ["main_input_name"]
 
-    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
+    def forward(self, *args: List[torch.Tensor], **kwargs: torch.Tensor):
         output = super().forward(*args, **kwargs)
         return BaseModelOutput(last_hidden_state=output)
 
@@ -83,7 +85,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
                 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}."
+                        f"Decoding step {decoding_step} out of bounds for decoder_max_seq_len ({self.dec_max_seq_len})."
                     )
                 decoder_attention_mask[b_idx, : decoding_step + 1] = 1
 
@@ -101,7 +103,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         return Seq2SeqLMOutput(logits=lm_logits)
 
 
-class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
+class RBLNModelForSeq2SeqLM(RBLNModel, GenerationMixin, ABC):
     """
     This is a generic model class that will be instantiated as one of the model classes of the library (with a sequence-to-sequence language modeling head) when created with the from_pretrained() class method.
     This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
@@ -117,6 +119,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
     main_input_name = "input_ids"
     auto_model_class = AutoModelForSeq2SeqLM
     support_causal_attn = None
+    _is_stateful = False
 
     def __post_init__(self, **kwargs):
         batch_size = self.rbln_config.batch_size
@@ -138,7 +141,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
     @classmethod
     @torch.inference_mode()
     def get_compiled_model(cls, model: PreTrainedModel, rbln_config: RBLNModelForSeq2SeqLMConfig):
-        wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
+        wrapped_model = cls._wrap_model_if_needed(model, rbln_config)
 
         enc_compile_config = rbln_config.compile_cfgs[0]
         dec_compile_config = rbln_config.compile_cfgs[1]
@@ -181,6 +184,21 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
 
         return {"encoder": compiled_encoder, "decoder": compiled_decoder}
 
+    @classmethod
+    def _update_paged_attention_config(cls, model_config: PretrainedConfig, rbln_config: RBLNModelForSeq2SeqLMConfig):
+        rbln_config.kvcache_num_blocks = rbln_config.kvcache_num_blocks or rbln_config.batch_size
+        rbln_config.kvcache_block_size = rbln_config.kvcache_block_size or rbln_config.dec_max_seq_len
+
+        if rbln_config.kvcache_num_blocks != rbln_config.batch_size:
+            raise NotImplementedError(
+                f"kvcache_num_blocks ({rbln_config.kvcache_num_blocks}) must be equal to batch_size ({rbln_config.batch_size}) as flash attention is not supported yet."
+            )
+
+        if rbln_config.kvcache_block_size != rbln_config.dec_max_seq_len:
+            raise NotImplementedError(
+                f"kvcache_block_size ({rbln_config.kvcache_block_size}) must be equal to dec_max_seq_len ({rbln_config.dec_max_seq_len}) as flash attention is not supported yet."
+            )
+
     @classmethod
     def _update_rbln_config(
         cls,
@@ -204,12 +222,6 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             model_config, "max_position_embeddings", None
         )
 
-        pad_token_id = getattr(model_config, "pad_token_id", None)
-        pad_token_id = pad_token_id or getattr(model_config, "bos_token_id", None)
-        pad_token_id = pad_token_id or getattr(model_config, "eos_token_id", None)
-        pad_token_id = pad_token_id or -1
-        rbln_config.pad_token_id = pad_token_id
-
         if rbln_config.enc_max_seq_len is None:
             enc_max_seq_len = max_position_embeddings
             for tokenizer in preprocessors:
@@ -238,6 +250,9 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         if max_position_embeddings is not None and rbln_config.dec_max_seq_len > max_position_embeddings:
             raise ValueError("`dec_max_seq_len` should be less or equal than max_position_embeddings!")
 
+        if rbln_config.support_paged_attention:
+            cls._update_paged_attention_config(model_config, rbln_config)
+
         # model input info
         enc_input_info = [
             ("input_ids", [1, rbln_config.enc_max_seq_len], "int64"),
@@ -310,6 +325,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         dec_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
 
         rbln_config.set_compile_cfgs([enc_compile_config, dec_compile_config])
+
         return rbln_config
 
     @classmethod
@@ -411,7 +427,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         inputs_tensor = torch.nn.functional.pad(
             inputs_tensor,
             (0, self.rbln_config.enc_max_seq_len - input_len),
-            value=self.rbln_config.pad_token_id,
+            value=self.config.pad_token_id,
         )
         model_kwargs["attention_mask"] = torch.nn.functional.pad(
             model_kwargs["attention_mask"], (0, self.rbln_config.enc_max_seq_len - input_len)
@@ -430,3 +446,32 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             model_kwargs["encoder_outputs"] = encoder(**encoder_kwargs, block_tables=block_tables)
 
         return model_kwargs
+
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        **kwargs,
+    ) -> Union[ModelOutput, torch.LongTensor]:
+        """
+        The generate function is utilized in its standard form as in the HuggingFace transformers library. User can use this function to generate text from the model.
+        Check the [HuggingFace transformers documentation](https://huggingface.co/docs/transformers/v4.57.1/en/main_classes/text_generation#transformers.GenerationMixin.generate) for more details.
+
+        Args:
+            input_ids (torch.LongTensor): The input ids to the model.
+            attention_mask (torch.LongTensor, optional): The attention mask to the model.
+            generation_config (GenerationConfig, optional): The generation configuration to be used as base parametrization for the generation call. **kwargs passed to generate matching the attributes of generation_config will override them.
+                If generation_config is not provided, the default will be used, which had the following loading priority: 1) from the generation_config.json model file, if it exists; 2) from the model configuration.
+                Please note that unspecified parameters will inherit [GenerationConfig](https://huggingface.co/docs/transformers/v4.57.1/en/main_classes/text_generation#transformers.GenerationConfig)’s default values.
+            kwargs (dict[str, Any], optional): Additional arguments passed to the generate function. See the HuggingFace transformers documentation for more details.
+
+        Returns:
+            Generates sequences of token ids for models with a language modeling head.
+        """
+        if generation_config is not None:
+            kwargs["generation_config"] = generation_config
+        if attention_mask is not None:
+            kwargs["attention_mask"] = attention_mask
+
+        return super().generate(input_ids, **kwargs)

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

@@ -31,7 +31,7 @@ class Seq2SeqWrapper:
     Args:
         model (nn.Module): The Seq2Seq model to wrap.
         enc_max_seq_len (int): Maximum sequence length for the encoder's position embeddings and cache sizes.
-        **kwargs: Additional arguments to pass to the decoder wrapper.
+        kwargs: Additional arguments to pass to the decoder wrapper.
     """
 
     def __init__(self, model: nn.Module, enc_max_seq_len: int, **kwargs):
@@ -125,7 +125,7 @@ class Seq2SeqDecoderWrapper(nn.Module):
 
     Args:
         model (nn.Module): The Seq2Seq model containing the decoder.
-        **kwargs: Additional arguments for decoder configuration.
+        kwargs: Additional arguments for decoder configuration.
     """
 
     def __init__(self, model: nn.Module, use_attention_mask: bool = True, **kwargs):

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

@@ -12,9 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_siglip import (
-    RBLNSiglipVisionModelConfig,
-)
-from .modeling_siglip import (
-    RBLNSiglipVisionModel,
-)
+from .configuration_siglip import RBLNSiglipVisionModelConfig
+from .modeling_siglip import RBLNSiglipVisionModel