optimum-rbln 0.9.3.post1__py3-none-any.whl

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/__init__.py

@@ -0,0 +1,23 @@
+# 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_resnet import RBLNResNetForImageClassificationConfig
+from .modeling_resnet import RBLNResNetForImageClassification
+
+
+__all__ = [
+    "RBLNResNetForImageClassificationConfig",
+    "RBLNResNetForImageClassification",
+]

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

@@ -0,0 +1,42 @@
+# 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
+
+from ...configuration_generic import RBLNModelForImageClassificationConfig
+
+
+class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConfig):
+    """
+    Configuration class for RBLNResNetForImageClassification.
+
+    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

@@ -0,0 +1,99 @@
+# 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, 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):
+    """
+    RBLN optimized ResNet model for image classification tasks.
+
+    This class provides hardware-accelerated inference for ResNet models
+    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/__init__.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_roberta import RBLNRobertaForMaskedLMConfig, RBLNRobertaForSequenceClassificationConfig
+from .modeling_roberta import RBLNRobertaForMaskedLM, RBLNRobertaForSequenceClassification
+
+
+__all__ = [
+    "RBLNRobertaForMaskedLMConfig",
+    "RBLNRobertaForSequenceClassificationConfig",
+    "RBLNRobertaForMaskedLM",
+    "RBLNRobertaForSequenceClassification",
+]

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

@@ -0,0 +1,33 @@
+# 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 RBLNModelForMaskedLMConfig, RBLNModelForSequenceClassificationConfig
+
+
+class RBLNRobertaForMaskedLMConfig(RBLNModelForMaskedLMConfig):
+    """
+    Configuration class for RBLNRobertaForMaskedLM.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized RoBERTa models for masked language modeling tasks.
+    """
+
+
+class RBLNRobertaForSequenceClassificationConfig(RBLNModelForSequenceClassificationConfig):
+    """
+    Configuration class for RBLNRobertaForSequenceClassification.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized RoBERTa models for sequence classification tasks.
+    """

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

@@ -0,0 +1,72 @@
+# 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 Tuple, Union
+
+import torch
+from transformers.modeling_outputs import MaskedLMOutput, SequenceClassifierOutput
+
+from ...modeling_generic import RBLNModelForMaskedLM, RBLNModelForSequenceClassification
+
+
+class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
+    """
+    RBLN optimized RoBERTa model for masked language modeling tasks.
+
+    This class provides hardware-accelerated inference for RoBERTa models
+    on RBLN devices, supporting masked language modeling tasks such as
+    token prediction and text completion.
+    """
+
+    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):
+    """
+    RBLN optimized RoBERTa model for sequence classification tasks.
+
+    This class provides hardware-accelerated inference for RoBERTa models
+    on RBLN devices, supporting text classification tasks such as sentiment analysis,
+    topic classification, and other sequence-level prediction tasks.
+    """
+
+    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/__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_seq2seq import RBLNModelForSeq2SeqLMConfig
+from .modeling_seq2seq import RBLNModelForSeq2SeqLM

optimum/rbln/transformers/models/seq2seq/configuration_seq2seq.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 typing import Any, Optional
+
+from ....configuration_utils import RBLNModelConfig
+from ....utils.deprecation import deprecate_kwarg
+from ....utils.logging import get_logger
+
+
+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,
+        kvcache_num_blocks: Optional[int] = None,
+        kvcache_block_size: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        """
+        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.
+            use_attention_mask (Optional[bool]): Whether to use attention masks during inference.
+            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.
+        """
+        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.use_attention_mask = use_attention_mask
+
+        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."
+                )