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

optimum/rbln/transformers/models/qwen2_vl/qwen2_vl_architecture.py

@@ -0,0 +1,165 @@
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+
+from ..decoderonly.decoderonly_architecture import (
+    DecoderOnlyWrapper,
+    apply_rotary_pos_emb,
+)
+
+
+class Qwen2VisionTransformerWrapper(nn.Module):
+    def __init__(self, model: torch.nn.Module):
+        super().__init__()
+        self._original_mod = model
+        self.merger = model.merger
+        self.blocks = self.wrap_vision_blocks(model.blocks)
+
+    def wrap_vision_blocks(self, blocks: torch.nn.ModuleList):
+        wrapped_blocks = []
+        for i, block in enumerate(blocks):
+            wrapped_blocks.append(Qwen2VLVisionBlock(block))
+        return nn.ModuleList(wrapped_blocks)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        full_attn_masks: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ):
+        full_attn_masks = (1 - full_attn_masks) * torch.finfo(torch.float32).min
+
+        for block in self.blocks:
+            hidden_states = block(hidden_states, full_attn_masks, [cos, sin])
+
+        return self.merger(hidden_states)
+
+
+class Qwen2VLVisionBlock(torch.nn.Module):
+    def __init__(self, model: torch.nn.Module):
+        super().__init__()
+        self._origin_model = model
+        self.norm1 = model.norm1
+        self.norm2 = model.norm2
+
+        self.attn = VisionAttention(model.attn)
+        self.mlp = model.mlp
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_masks: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(
+            self.norm1(hidden_states),
+            attn_masks,
+            position_embeddings,
+        )
+        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+        return hidden_states
+
+
+class VisionAttention(nn.Module):
+    def __init__(self, model: nn.Module) -> None:
+        super().__init__()
+        self._origin_model = model
+        self.num_heads = model.num_heads
+        self.head_dim = getattr(model, "head_dim", model.proj.in_features // model.num_heads)
+        self.qkv = model.qkv
+        self.proj = model.proj
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_masks: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        seq_length = hidden_states.shape[0]
+        hidden_states = hidden_states.unsqueeze(0)
+        q, k, v = (
+            self.qkv(hidden_states).reshape(1, seq_length, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4).unbind(0)
+        )
+
+        cos, sin = position_embeddings
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(self.head_dim)
+        attn_weights = attn_weights + attn_masks
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)
+        attn_output = torch.matmul(attn_weights, v)
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(1, seq_length, -1)
+        attn_output = self.proj(attn_output).squeeze(0)
+
+        return attn_output
+
+
+class Qwen2VL_LanguageModelWrapper(DecoderOnlyWrapper):
+    def prepare_forward_args(self, *args):
+        args = list(args)
+        input_ids = None if self.rbln_config.use_inputs_embeds else args.pop(0)
+        inputs_embeds = args.pop(0) if self.rbln_config.use_inputs_embeds else None
+        cache_position = args.pop(0)
+        global_block_tables = args.pop(0)
+        local_block_tables = None
+        position_embeds = args.pop(0)
+        query_position = args.pop(0) if self.phase == "prefill" else None
+        position_ids = None
+        attention_mask = args.pop(0) if self.rbln_config.use_attention_mask else None
+        lora_int_id = args.pop(0) if self.rbln_config.lora_config else None
+        past_key_values = args
+
+        if len(past_key_values) != 2 * self.num_hidden_layers:
+            raise ValueError(
+                f"Different past_key_values to model's config. {len(past_key_values)} != {2 * self.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
+
+        return (
+            input_ids,
+            inputs_embeds,
+            cache_position,
+            global_block_tables,
+            local_block_tables,
+            query_position,
+            attention_mask,
+            position_ids,
+            lora_int_id,
+            past_key_values,
+            position_embeds,
+        )
+
+    def convert_to_rbln_class(self, model: PreTrainedModel, max_seq_len: int):
+        new_layers = []
+
+        for layer_idx, layer in enumerate(model.model.language_model.layers):
+            is_sliding = layer_idx in self.rbln_config.sliding_window_layers
+            new_self_attn = self.get_rbln_attn_class()(
+                self.get_attn_layer(layer), self.rbln_config, is_sliding=is_sliding
+            )
+            new_layer = self.get_rbln_layer_class()(layer, new_self_attn)
+            new_layers.append(new_layer)
+
+        new_model = self.get_rbln_model_class()(
+            model.model.language_model,
+            new_layers,
+            self.rbln_config,
+            use_learned_pos_emb=self.__class__._use_learned_pos_emb,
+        )
+
+        new_model = self.get_rbln_causal_lm_class()(model.model, new_model)
+        return new_model

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

@@ -0,0 +1,16 @@
+# 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_qwen3 import RBLNQwen3ForCausalLMConfig, RBLNQwen3ModelConfig
+from .modeling_qwen3 import RBLNQwen3ForCausalLM, RBLNQwen3Model

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

@@ -0,0 +1,71 @@
+# 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 ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
+
+
+class RBLNQwen3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    """
+    Configuration class for RBLN Qwen3 models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNQwen3ForCausalLM, RBLNQwen3ForCausalLMConfig
+
+    # Create a configuration object
+    config = RBLNQwen3ForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=40960,
+        tensor_parallel_size=4,
+        kvcache_partition_len=16384
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNQwen3ForCausalLM.from_pretrained(
+        "Qwen/Qwen3-4B",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """
+
+
+class RBLNQwen3ModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN Qwen3 models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNQwen3Model, RBLNQwen3ModelConfig
+
+    # Create a configuration object
+    config = RBLNQwen3ModelConfig(
+        batch_size=1,
+        max_seq_len=40960,
+        tensor_parallel_size=4,
+        kvcache_partition_len=16384
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNQwen3Model.from_pretrained(
+        "Qwen/Qwen3-Embedding-4B",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

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

@@ -0,0 +1,133 @@
+# 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 TYPE_CHECKING
+
+from transformers import PretrainedConfig
+
+from ....utils import logging
+from ...models.decoderonly import (
+    RBLNDecoderOnlyModel,
+    RBLNDecoderOnlyModelForCausalLM,
+    RBLNDecoderOnlyModelForCausalLMConfig,
+)
+from .qwen3_architecture import Qwen3Wrapper
+
+
+logger = logging.get_logger(__name__)
+
+if TYPE_CHECKING:
+    from transformers import PretrainedConfig
+
+
+class RBLNQwen3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The Qwen3 Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    A class to convert and run pre-trained transformers based Qwen3ForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers Qwen3ForCausalLM model into a RBLN transformer model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+    **Configuration:**
+    This model uses [`RBLNQwen3ForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNQwen3ForCausalLMConfig`] or a dictionary conforming to its structure.
+    See the [`RBLNQwen3ForCausalLMConfig`] class for all available configuration options.
+    Examples:
+        ```python
+        from optimum.rbln import RBLNQwen3ForCausalLM
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNQwen3ForCausalLM.from_pretrained(
+            "Qwen/Qwen3-4B",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 40_960,
+            "tensor_parallel_size": 4,
+            "kvcache_partition_len": 8192,
+        }
+        model = RBLNQwen3ForCausalLM.from_pretrained(
+            "Qwen/Qwen3-4B",
+            export=True,
+            rbln_config=rbln_config
+        )
+        # Using a RBLNQwen3ForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNQwen3ForCausalLMConfig
+        config = RBLNQwen3ForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=40_960,
+            tensor_parallel_size=4,
+            kvcache_partition_len=8192,
+        )
+        model = RBLNQwen3ForCausalLM.from_pretrained(
+            "Qwen/Qwen3-4B",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    _decoder_wrapper_cls = Qwen3Wrapper
+
+    @classmethod
+    def _update_sliding_window_config(
+        cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
+    ):
+        # https://github.com/huggingface/transformers/issues/35896
+        # There seems to be a bug in transformers(v4.52.4). Therefore, similar to when attn_implementation is eager,
+        # we set all layers to use sliding window in this version. This should be updated once the bug is fixed.
+
+        rbln_config.cache_impl = "sliding_window"
+        rbln_config.sliding_window = model_config.sliding_window
+        rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
+        return rbln_config
+
+    def forward(self, *args, **kwargs):
+        kwargs["return_dict"] = True
+        return super().forward(*args, **kwargs)
+
+
+class RBLNQwen3Model(RBLNDecoderOnlyModel):
+    """
+    The bare Qwen3 Model outputting raw hidden-states without any specific head on top.
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    A class to convert and run pre-trained transformers based Qwen3Model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers Qwen3Model into a RBLN transformer model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+    **Configuration:**
+    This model uses [`RBLNQwen3ModelConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNQwen3ModelConfig`] or a dictionary conforming to its structure.
+    See the [`RBLNQwen3ModelConfig`] class for all available configuration options.
+    Examples:
+        ```python
+        from optimum.rbln import RBLNQwen3Model
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNQwen3Model.from_pretrained(
+            "Qwen/Qwen3-Embedding-4B",
+            export=True,
+            rbln_batch_size=1,
+            rbln_max_seq_len=40_960,
+            rbln_tensor_parallel_size=4,
+            rbln_kvcache_partition_len=8192,
+        )
+    """
+
+    _decoder_wrapper_cls = Qwen3Wrapper
+    _use_rotary_emb = True

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

@@ -0,0 +1,31 @@
+# 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 ..decoderonly.decoderonly_architecture import DecoderOnlyAttention, DecoderOnlyWrapper
+
+
+class Qwen3Wrapper(DecoderOnlyWrapper):
+    def get_rbln_attn_class(self):
+        return Qwen3Attention
+
+
+class Qwen3Attention(DecoderOnlyAttention):
+    def __post_init__(self):
+        self.k_proj = self._original_mod.k_proj
+        self.v_proj = self._original_mod.v_proj
+        self.q_proj = self._original_mod.q_proj
+        self.o_proj = self._original_mod.o_proj
+        self.q_norm = self._original_mod.q_norm
+        self.k_norm = self._original_mod.k_norm

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."
+                )