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

optimum/rbln/ops/flash_attn.py

@@ -12,71 +12,165 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from functools import lru_cache
-
 import torch
-from packaging import version
-
-
-if version.parse(torch.__version__) > version.parse("2.4.0"):
-    register_fake = torch.library.register_fake
-else:
-    register_fake = torch.library.impl_abstract
-
-
-@lru_cache
-def register_rbln_custom_paged_flash_attention():
-    torch.library.define(
-        "rbln_custom_ops::paged_flash_attn_decode",
-        "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e, int f, int g) -> Tensor",
-    )
-
-    @torch.library.impl("rbln_custom_ops::paged_flash_attn_decode", "cpu")
-    def flash_attn_decode_cpu(q, k, v, mask, kcache, vcache, seq, scale, block_table, block_size, partition):
-        return q
-
-    @register_fake("rbln_custom_ops::paged_flash_attn_decode")
-    def flash_attn_decode_abstract(q, k, v, m, kcache, vcache, seq, scale, block_table, block_size, partition):
-        return q
-
-    torch.library.define(
-        "rbln_custom_ops::paged_flash_attn_prefill",
-        "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e, int f, int g) -> Tensor",
-    )
-
-    @torch.library.impl("rbln_custom_ops::flash_attn_prefill", "cpu")
-    def flash_attn_prefill_cpu(q, k, v, mask, kcache, vcache, seq, scale, block_table, block_size, partition):
-        return q
-
-    @register_fake("rbln_custom_ops::paged_flash_attn_prefill")
-    def flash_attn_prefill_abstract(q, k, v, m, kcache, vcache, seq, scale, block_table, block_size, partition):
-        return q
-
-
-@lru_cache
-def register_rbln_custom_paged_flash_causal_attention():
-    torch.library.define(
-        "rbln_custom_ops::paged_flash_causal_attn_decode",
-        "(Tensor x, Tensor y, Tensor z, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e, int f, int g) -> Tensor",
-    )
-
-    @torch.library.impl("rbln_custom_ops::paged_flash_causal_attn_decode", "cpu")
-    def flash_attn_decode_cpu(q, k, v, kcache, vcache, seq, scale, block_table, block_size, partition):
-        return q
-
-    @register_fake("rbln_custom_ops::paged_flash_causal_attn_decode")
-    def flash_attn_decode_abstract(q, k, v, kcache, vcache, seq, scale, block_table, block_size, partition):
-        return q
-
-    torch.library.define(
-        "rbln_custom_ops::paged_flash_causal_attn_prefill",
-        "(Tensor x, Tensor y, Tensor z, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e, int f, int g) -> Tensor",
-    )
-
-    @torch.library.impl("rbln_custom_ops::paged_flash_causal_attn_prefill", "cpu")
-    def flash_attn_prefill_cpu(q, k, v, kcache, vcache, seq, scale, block_table, block_size, partition):
-        return q
-
-    @register_fake("rbln_custom_ops::paged_flash_causal_attn_prefill")
-    def flash_attn_prefill_abstract(q, k, v, kcache, vcache, seq, scale, block_table, block_size, partition):
-        return q
+from torch import Tensor
+
+
+@torch.library.custom_op(
+    "rbln_custom_ops::paged_flash_attn_decode",
+    mutates_args=(["kcache", "vcache"]),
+)
+def paged_flash_attn_decode(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+    partition: int,
+) -> Tensor:
+    """Defines the computation pattern for fused flash attention with KV cache for decoding.
+
+    Returns a tensor with the same shape as q.
+    """
+    return torch.empty_like(q)
+
+
+@paged_flash_attn_decode.register_fake
+def paged_flash_attn_decode_fake(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+    partition: int,
+) -> Tensor:
+    return torch.empty_like(q)
+
+
+@torch.library.custom_op(
+    "rbln_custom_ops::paged_flash_attn_prefill",
+    mutates_args=(["kcache", "vcache"]),
+)
+def paged_flash_attn_prefill(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+    partition: int,
+) -> Tensor:
+    """Defines the computation pattern for fused flash attention with KV cache for prefill.
+
+    Returns a tensor with the same shape as q.
+    """
+    return torch.empty_like(q)
+
+
+@paged_flash_attn_prefill.register_fake
+def paged_flash_attn_prefill_fake(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+    partition: int,
+) -> Tensor:
+    return torch.empty_like(q)
+
+
+@torch.library.custom_op(
+    "rbln_custom_ops::paged_flash_causal_attn_decode",
+    mutates_args=(["kcache", "vcache"]),
+)
+def paged_flash_causal_attn_decode(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+    partition: int,
+) -> Tensor:
+    """Defines the computation pattern for fused causal flash attention with KV cache for decoding.
+
+    Returns a tensor with the same shape as q.
+    """
+    return torch.empty_like(q)
+
+
+@paged_flash_causal_attn_decode.register_fake
+def paged_flash_causal_attn_decode_fake(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+    partition: int,
+) -> Tensor:
+    return torch.empty_like(q)
+
+
+@torch.library.custom_op(
+    "rbln_custom_ops::paged_flash_causal_attn_prefill",
+    mutates_args=(["kcache", "vcache"]),
+)
+def paged_flash_causal_attn_prefill(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+    partition: int,
+) -> Tensor:
+    """Defines the computation pattern for fused causal flash attention with KV cache for prefill.
+
+    Returns a tensor with the same shape as q.
+    """
+    return torch.empty_like(q)
+
+
+@paged_flash_causal_attn_prefill.register_fake
+def paged_flash_causal_attn_prefill_fake(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+    partition: int,
+) -> Tensor:
+    return torch.empty_like(q)

optimum/rbln/ops/kv_cache_update.py

@@ -12,49 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from functools import lru_cache
-
 import torch
-from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_4
-
+from torch import Tensor
 
-if is_torch_greater_or_equal_than_2_4:
-    register_fake = torch.library.register_fake
-else:
-    register_fake = torch.library.impl_abstract
 
-
-@lru_cache
-def register_rbln_custom_cache_update():
+@torch.library.custom_op("rbln_custom_ops::rbln_cache_update", mutates_args=(["cache"]))
+def rbln_cache_update(cache: Tensor, state: Tensor, position: Tensor, axis: Tensor) -> Tensor:
     # Define the RBLN custom operation "rbln_cache_update" which updates a cache tensor with a given state tensor.
     # This operation is designed to perform in-place updates directly on the device without needing to transfer the cache back to the host.
     # The `position` parameter specifies the start index for the update along the specified axis, allowing flexible updates to any part of the cache tensor.
-    torch.library.define("rbln_custom_ops::rbln_cache_update", "(Tensor x, Tensor y, Tensor z, Tensor w) -> Tensor")
-
-    # Implementation of the "rbln_cache_update" operation for the CPU.
-    @torch.library.impl("rbln_custom_ops::rbln_cache_update", "cpu")
-    def rbln_cache_update_cpu(cache, state, position, axis):
-        assert position.dim() == 0
-        assert axis.dim() == 0
-
-        # Calculate the start (s) and end (e) indices for the update based on the position and the shape of the state tensor along the specified axis.
-        s = position  # Start index for the update, specified by the position.
-        e = (
-            position + state.shape[axis]
-        )  # End index is determined by adding the size of the state along the given axis.
-
-        # Update the specified portion of the cache tensor with the state tensor, using `slice_scatter`.
-        # This operation modifies the cache tensor in-place directly on the device, avoiding any unnecessary transfers between host and device.
-        cache.slice_scatter(state, dim=axis, start=s, end=e)
-
-        # 'rbln_cache_update' is an in-place operation that isn't tracked in JIT trace, so a dummy output was added to the return value.
-        return torch.empty([256])
-
-    # Register a "fake" implementation of the "rbln_cache_update" operation.
-    # This serves as an abstract definition for the RBLN compiler to recognize the operation and generate an optimized implementation.
-    @register_fake("rbln_custom_ops::rbln_cache_update")
-    def rbln_cache_update_abstract(cache, state, position, axis):
-        # Return a tensor with the same shape as the input cache tensor.
-        # This is a placeholder for the abstract implementation and does not perform any actual computation.
-        # Like the actual implementation, the abstraction assumes in-place device-side updates.
-        return torch.empty([256])
+    return torch.empty_like(cache)

optimum/rbln/ops/linear.py

@@ -0,0 +1,25 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+@torch.library.custom_op("rbln_custom_ops::linear", mutates_args=())
+def linear(input: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
+    output_shape = list(input.shape[:-1])
+    output_shape += [weight.shape[0]]
+    return torch.empty(size=output_shape, dtype=input.dtype, device=input.device, requires_grad=input.requires_grad)

optimum/rbln/transformers/__init__.py

@@ -18,7 +18,15 @@ from transformers.utils import _LazyModule
 
 
 _import_structure = {
-    "cache_utils": ["RebelDynamicCache"],
+    "configuration_alias": [
+        "RBLNASTForAudioClassificationConfig",
+        "RBLNDistilBertForQuestionAnsweringConfig",
+        "RBLNResNetForImageClassificationConfig",
+        "RBLNXLMRobertaForSequenceClassificationConfig",
+        "RBLNRobertaForSequenceClassificationConfig",
+        "RBLNRobertaForMaskedLMConfig",
+        "RBLNViTForImageClassificationConfig",
+    ],
     "models": [
         "RBLNAutoModel",
         "RBLNAutoModelForAudioClassification",
@@ -33,29 +41,66 @@ _import_structure = {
         "RBLNAutoModelForSpeechSeq2Seq",
         "RBLNAutoModelForVision2Seq",
         "RBLNBartForConditionalGeneration",
+        "RBLNBartForConditionalGenerationConfig",
         "RBLNBartModel",
-        "RBLNBertModel",
+        "RBLNBartModelConfig",
         "RBLNBertForMaskedLM",
+        "RBLNBertForMaskedLMConfig",
         "RBLNBertForQuestionAnswering",
+        "RBLNBertForQuestionAnsweringConfig",
+        "RBLNBertModel",
+        "RBLNBertModelConfig",
         "RBLNCLIPTextModel",
+        "RBLNCLIPTextModelConfig",
         "RBLNCLIPTextModelWithProjection",
+        "RBLNCLIPTextModelWithProjectionConfig",
         "RBLNCLIPVisionModel",
+        "RBLNCLIPVisionModelConfig",
         "RBLNCLIPVisionModelWithProjection",
+        "RBLNCLIPVisionModelWithProjectionConfig",
+        "RBLNDecoderOnlyModelForCausalLM",
+        "RBLNDecoderOnlyModelForCausalLMConfig",
         "RBLNDPTForDepthEstimation",
+        "RBLNDPTForDepthEstimationConfig",
         "RBLNExaoneForCausalLM",
+        "RBLNExaoneForCausalLMConfig",
         "RBLNGemmaForCausalLM",
+        "RBLNGemmaForCausalLMConfig",
         "RBLNGPT2LMHeadModel",
-        "RBLNQwen2ForCausalLM",
-        "RBLNWav2Vec2ForCTC",
-        "RBLNWhisperForConditionalGeneration",
+        "RBLNGPT2LMHeadModelConfig",
+        "RBLNIdefics3VisionTransformer",
+        "RBLNIdefics3ForConditionalGeneration",
+        "RBLNIdefics3ForConditionalGenerationConfig",
+        "RBLNIdefics3VisionTransformerConfig",
         "RBLNLlamaForCausalLM",
+        "RBLNLlamaForCausalLMConfig",
+        "RBLNLlavaNextForConditionalGeneration",
+        "RBLNLlavaNextForConditionalGenerationConfig",
+        "RBLNMidmLMHeadModel",
+        "RBLNMidmLMHeadModelConfig",
+        "RBLNMistralForCausalLM",
+        "RBLNMistralForCausalLMConfig",
         "RBLNPhiForCausalLM",
+        "RBLNPhiForCausalLMConfig",
+        "RBLNQwen2ForCausalLM",
+        "RBLNQwen2ForCausalLMConfig",
+        "RBLNQwen2_5_VisionTransformerPretrainedModel",
+        "RBLNQwen2_5_VisionTransformerPretrainedModelConfig",
+        "RBLNQwen2_5_VLForConditionalGeneration",
+        "RBLNQwen2_5_VLForConditionalGenerationConfig",
         "RBLNT5EncoderModel",
+        "RBLNT5EncoderModelConfig",
         "RBLNT5ForConditionalGeneration",
+        "RBLNT5ForConditionalGenerationConfig",
+        "RBLNWav2Vec2ForCTC",
+        "RBLNWav2Vec2ForCTCConfig",
+        "RBLNWhisperForConditionalGeneration",
+        "RBLNWhisperForConditionalGenerationConfig",
+        "RBLNTimeSeriesTransformerForPrediction",
+        "RBLNTimeSeriesTransformerForPredictionConfig",
         "RBLNLlavaNextForConditionalGeneration",
-        "RBLNMidmLMHeadModel",
         "RBLNXLMRobertaModel",
-        "RBLNMistralForCausalLM",
+        "RBLNXLMRobertaModelConfig",
     ],
     "modeling_alias": [
         "RBLNASTForAudioClassification",
@@ -69,7 +114,15 @@ _import_structure = {
 }
 
 if TYPE_CHECKING:
-    from .cache_utils import RebelDynamicCache
+    from .configuration_alias import (
+        RBLNASTForAudioClassificationConfig,
+        RBLNDistilBertForQuestionAnsweringConfig,
+        RBLNResNetForImageClassificationConfig,
+        RBLNRobertaForMaskedLMConfig,
+        RBLNRobertaForSequenceClassificationConfig,
+        RBLNViTForImageClassificationConfig,
+        RBLNXLMRobertaForSequenceClassificationConfig,
+    )
     from .modeling_alias import (
         RBLNASTForAudioClassification,
         RBLNDistilBertForQuestionAnswering,
@@ -93,29 +146,65 @@ if TYPE_CHECKING:
         RBLNAutoModelForSpeechSeq2Seq,
         RBLNAutoModelForVision2Seq,
         RBLNBartForConditionalGeneration,
+        RBLNBartForConditionalGenerationConfig,
         RBLNBartModel,
+        RBLNBartModelConfig,
         RBLNBertForMaskedLM,
+        RBLNBertForMaskedLMConfig,
         RBLNBertForQuestionAnswering,
+        RBLNBertForQuestionAnsweringConfig,
         RBLNBertModel,
+        RBLNBertModelConfig,
         RBLNCLIPTextModel,
+        RBLNCLIPTextModelConfig,
         RBLNCLIPTextModelWithProjection,
+        RBLNCLIPTextModelWithProjectionConfig,
         RBLNCLIPVisionModel,
+        RBLNCLIPVisionModelConfig,
         RBLNCLIPVisionModelWithProjection,
+        RBLNCLIPVisionModelWithProjectionConfig,
+        RBLNDecoderOnlyModelForCausalLM,
+        RBLNDecoderOnlyModelForCausalLMConfig,
         RBLNDPTForDepthEstimation,
+        RBLNDPTForDepthEstimationConfig,
         RBLNExaoneForCausalLM,
+        RBLNExaoneForCausalLMConfig,
         RBLNGemmaForCausalLM,
+        RBLNGemmaForCausalLMConfig,
         RBLNGPT2LMHeadModel,
+        RBLNGPT2LMHeadModelConfig,
+        RBLNIdefics3ForConditionalGeneration,
+        RBLNIdefics3ForConditionalGenerationConfig,
+        RBLNIdefics3VisionTransformer,
+        RBLNIdefics3VisionTransformerConfig,
         RBLNLlamaForCausalLM,
+        RBLNLlamaForCausalLMConfig,
         RBLNLlavaNextForConditionalGeneration,
+        RBLNLlavaNextForConditionalGenerationConfig,
         RBLNMidmLMHeadModel,
+        RBLNMidmLMHeadModelConfig,
         RBLNMistralForCausalLM,
+        RBLNMistralForCausalLMConfig,
         RBLNPhiForCausalLM,
+        RBLNPhiForCausalLMConfig,
+        RBLNQwen2_5_VisionTransformerPretrainedModel,
+        RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
+        RBLNQwen2_5_VLForConditionalGeneration,
+        RBLNQwen2_5_VLForConditionalGenerationConfig,
         RBLNQwen2ForCausalLM,
+        RBLNQwen2ForCausalLMConfig,
         RBLNT5EncoderModel,
+        RBLNT5EncoderModelConfig,
         RBLNT5ForConditionalGeneration,
+        RBLNT5ForConditionalGenerationConfig,
+        RBLNTimeSeriesTransformerForPrediction,
+        RBLNTimeSeriesTransformerForPredictionConfig,
         RBLNWav2Vec2ForCTC,
+        RBLNWav2Vec2ForCTCConfig,
         RBLNWhisperForConditionalGeneration,
+        RBLNWhisperForConditionalGenerationConfig,
         RBLNXLMRobertaModel,
+        RBLNXLMRobertaModelConfig,
     )
 else:
     import sys

optimum/rbln/transformers/configuration_alias.py

@@ -0,0 +1,49 @@
+# 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 (
+    RBLNModelForAudioClassificationConfig,
+    RBLNModelForImageClassificationConfig,
+    RBLNModelForMaskedLMConfig,
+    RBLNModelForQuestionAnsweringConfig,
+    RBLNModelForSequenceClassificationConfig,
+)
+
+
+class RBLNASTForAudioClassificationConfig(RBLNModelForAudioClassificationConfig):
+    pass
+
+
+class RBLNDistilBertForQuestionAnsweringConfig(RBLNModelForQuestionAnsweringConfig):
+    pass
+
+
+class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConfig):
+    pass
+
+
+class RBLNXLMRobertaForSequenceClassificationConfig(RBLNModelForSequenceClassificationConfig):
+    pass
+
+
+class RBLNRobertaForSequenceClassificationConfig(RBLNModelForSequenceClassificationConfig):
+    pass
+
+
+class RBLNRobertaForMaskedLMConfig(RBLNModelForMaskedLMConfig):
+    pass
+
+
+class RBLNViTForImageClassificationConfig(RBLNModelForImageClassificationConfig):
+    pass

optimum/rbln/transformers/configuration_generic.py

@@ -0,0 +1,142 @@
+# 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 List, Optional, Tuple, Union
+
+from ..configuration_utils import RBLNModelConfig
+
+
+class _RBLNTransformerEncoderConfig(RBLNModelConfig):
+    rbln_model_input_names: Optional[List[str]] = None
+
+    def __init__(
+        self,
+        max_seq_len: Optional[int] = None,
+        batch_size: Optional[int] = None,
+        model_input_names: Optional[List[str]] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            max_seq_len (Optional[int]): Maximum sequence length supported by the model.
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            model_input_names (Optional[List[str]]): Names of the input tensors for the model.
+                Defaults to class-specific rbln_model_input_names if not provided.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.max_seq_len = max_seq_len
+        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.model_input_names = model_input_names or self.rbln_model_input_names
+
+
+class _RBLNImageModelConfig(RBLNModelConfig):
+    def __init__(
+        self, image_size: Optional[Union[int, Tuple[int, int]]] = None, batch_size: Optional[int] = 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.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.image_size = image_size
+        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}")
+
+    @property
+    def image_width(self):
+        if isinstance(self.image_size, int):
+            return self.image_size
+        elif isinstance(self.image_size, (list, tuple)):
+            return self.image_size[1]
+        else:
+            return self.image_size["width"]
+
+    @property
+    def image_height(self):
+        if isinstance(self.image_size, int):
+            return self.image_size
+        elif isinstance(self.image_size, (list, tuple)):
+            return self.image_size[0]
+        else:
+            return self.image_size["height"]
+
+
+class RBLNModelForQuestionAnsweringConfig(_RBLNTransformerEncoderConfig):
+    pass
+
+
+class RBLNModelForSequenceClassificationConfig(_RBLNTransformerEncoderConfig):
+    pass
+
+
+class RBLNModelForMaskedLMConfig(_RBLNTransformerEncoderConfig):
+    pass
+
+
+class RBLNModelForTextEncodingConfig(_RBLNTransformerEncoderConfig):
+    pass
+
+
+# FIXME : Appropriate name ?
+class RBLNTransformerEncoderForFeatureExtractionConfig(_RBLNTransformerEncoderConfig):
+    pass
+
+
+class RBLNModelForImageClassificationConfig(_RBLNImageModelConfig):
+    pass
+
+
+class RBLNModelForDepthEstimationConfig(_RBLNImageModelConfig):
+    pass
+
+
+class RBLNModelForAudioClassificationConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        num_mel_bins: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            max_length (Optional[int]): Maximum length of the audio input in time dimension.
+            num_mel_bins (Optional[int]): Number of Mel frequency bins for audio processing.
+            **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.max_length = max_length
+        self.num_mel_bins = num_mel_bins