optimum-rbln 0.7.2rc2__tar.gz → 0.7.3a0__tar.gz

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/.github/scripts/auto_code_review.py

@@ -97,50 +97,6 @@ def remove_file_from_diff(diff_content, file_to_remove):
     return "\n".join(result)
 
 
-def skip_bot(pr):
-    global force_review
-    """
-    Review if
-    1. last commit messages starts with "[autoreview]"
-    2. last comment contains "/autoreview"
-    """
-
-    # Check commit message
-    commits = list(pr.get_commits())
-    if len(commits) == 0:
-        return True
-
-    last_commit = commits[-1]
-    try:
-        commit_message = last_commit.raw_data["commit"]["message"]
-    except KeyError:
-        commit_message = ""
-
-    if commit_message.startswith("[autoreview]"):
-        return False
-
-    # Check the last comment
-    comments = list(pr.get_issue_comments())
-    if len(comments) == 0:
-        return True
-
-    last = comments[-1]
-    if last.user.login.find("github-actions") != -1:
-        return True
-
-    if last.body.find("/autoreview") == -1:
-        return True
-
-    if last.reactions["heart"] > 0:
-        return True
-
-    if last.body.find("force") != -1:
-        force_review = True
-
-    last.create_reaction("heart")
-    return False
-
-
 def main():
     github_token = os.getenv("GITHUB_TOKEN")
     pr_number = os.getenv("PR_NUMBER")
@@ -155,13 +111,6 @@ def main():
     repo = g.get_repo(os.getenv("GITHUB_REPOSITORY"))
     pr = repo.get_pull(int(pr_number))
 
-    if skip_bot(pr):
-        print(
-            "To invoke review, Write '/autoreview' and re-run github actions,"
-            " or start the commit message with '[autoreview]'. "
-        )
-        sys.exit(0)
-
     # Get PR diff
     diff = get_pr_diff()
     diff = remove_file_from_diff(diff, "uv.lock")

optimum_rbln-0.7.3a0/.github/workflows/auto_code_review.yml

@@ -0,0 +1,72 @@
+name: Auto Code Review
+
+on:
+  pull_request:
+  issue_comment:
+    types: [created]
+  push:
+    branches:
+      - '**'
+
+env:
+  GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
+  GOOGLE_MODEL_ID: ${{ vars.GOOGLE_MODEL_ID }}
+
+jobs:
+  auto-review:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 2
+      
+      - name: Check if review should run
+        id: check
+        run: |
+          PR_NUMBER=""
+          SHOULD_RUN="false"
+          
+          # For push events, check commit message
+          if [[ "${{ github.event_name }}" == "push" ]]; then
+            if [[ "${{ contains(github.event.head_commit.message, '[autoreview]') }}" == "true" ]]; then
+              SHOULD_RUN="true"
+              # Use GitHub CLI to find PR associated with this commit
+              PR_NUMBER=$(gh pr list --head ${{ github.ref_name }} --json number --jq '.[0].number')
+            fi
+          
+          # For PR events
+          elif [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            PR_NUMBER="${{ github.event.pull_request.number }}"
+          
+          # For comment events, check if it's "/autoreview"
+          elif [[ "${{ github.event_name }}" == "issue_comment" ]]; then
+            if [[ "${{ github.event.issue.pull_request != null }}" == "true" && "${{ contains(github.event.comment.body, '/autoreview') }}" == "true" ]]; then
+              SHOULD_RUN="true"
+              PR_NUMBER="${{ github.event.issue.number }}"
+            fi
+          fi
+          
+          echo "should_run=$SHOULD_RUN" >> $GITHUB_OUTPUT
+          echo "pr_number=$PR_NUMBER" >> $GITHUB_OUTPUT
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Set up Python
+        if: steps.check.outputs.should_run == 'true' && steps.check.outputs.pr_number != ''
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.x'
+
+      - name: Install dependencies
+        if: steps.check.outputs.should_run == 'true' && steps.check.outputs.pr_number != ''
+        run: |
+          python -m pip install --upgrade pip
+          pip install google-generativeai PyGithub
+
+      - name: Run Auto Code Review
+        if: steps.check.outputs.should_run == 'true' && steps.check.outputs.pr_number != ''
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          PR_NUMBER: ${{ steps.check.outputs.pr_number }}
+        run: python .github/scripts/auto_code_review.py

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/.github/workflows/rbln_trigger_on_pr.yaml

@@ -9,7 +9,7 @@ env:
   REBEL_PYPI_ENDPOINT: ${{ vars.REBEL_PYPI_INTERNAL_ENDPOINT }}
   REBEL_PYPI_USERNAME: ${{ secrets.REBEL_PYPI_USERNAME }}
   REBEL_PYPI_PASSWORD: ${{ secrets.REBEL_PYPI_PASSWORD }}
-  REBEL_COMPILER_VERSION: 0.7.2.dev213+g4f75bbe9
+  REBEL_COMPILER_VERSION: 0.7.3.dev100+g3fd6ed0a
 
 jobs:
   check-rebel-compiler-version:

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.7.2rc2
+Version: 0.7.3a0
 Summary: Optimum RBLN is the interface between the Hugging Face Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/src/optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.2rc2'
-__version_tuple__ = version_tuple = (0, 7, 2)
+__version__ = version = '0.7.3a0'
+__version_tuple__ = version_tuple = (0, 7, 3)

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/src/optimum/rbln/diffusers/modeling_diffusers.py

@@ -71,13 +71,11 @@ class RBLNDiffusionMixin:
     _prefix = {}
 
     @classmethod
-    @property
-    def img2img_pipeline(cls):
+    def is_img2img_pipeline(cls):
         return "Img2Img" in cls.__name__
 
     @classmethod
-    @property
-    def inpaint_pipeline(cls):
+    def is_inpaint_pipeline(cls):
         return "Inpaint" in cls.__name__
 
     @classmethod
@@ -100,8 +98,8 @@ class RBLNDiffusionMixin:
             submodule_config.update({k: v for k, v in pipe_global_config.items() if k not in submodule_config})
             submodule_config.update(
                 {
-                    "img2img_pipeline": cls.img2img_pipeline,
-                    "inpaint_pipeline": cls.inpaint_pipeline,
+                    "img2img_pipeline": cls.is_img2img_pipeline(),
+                    "inpaint_pipeline": cls.is_inpaint_pipeline(),
                 }
             )
             submodule_config = submodule_cls.update_rbln_config_using_pipe(model, submodule_config)

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/src/optimum/rbln/modeling.py

@@ -196,7 +196,7 @@ class RBLNModel(RBLNBaseModel):
         **kwargs,
     ) -> "PreTrainedModel":
         kwargs = cls.update_kwargs(kwargs)
-        return cls.hf_class.from_pretrained(
+        return cls.get_hf_class().from_pretrained(
             model_id,
             subfolder=subfolder,
             revision=revision,

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/src/optimum/rbln/modeling_base.py

@@ -389,8 +389,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         return rbln_config
 
     @classmethod
-    @property
-    def hf_class(cls):
+    def get_hf_class(cls):
         """
         Lazily loads and caches the corresponding Hugging Face model class.
         Removes 'RBLN' prefix from the class name to get the original class name
@@ -416,7 +415,20 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         return self.forward(*args, **kwargs)
 
     def __repr__(self):
-        return repr(self.model) + repr(self.rbln_submodules)
+        has_submodules = len(self.rbln_submodules) > 0
+        repr_str: str = f"<{self.__class__.__name__}>\n"
+        repr_str += f"- Total {len(self.model)} Runtimes"
+        repr_str += f" and {len(self.rbln_submodules)} Submodules\n" if has_submodules else "\n"
+        repr_str += "[Runtimes]\n"
+        repr_str += "\n".join([repr(model) for model in self.model])
+        repr_str += "\n"
+
+        if has_submodules > 0:
+            for i, submodule in enumerate(self.rbln_submodules):
+                repr_str += f"[Submodules {i} : {self._rbln_submodules[i]['name']}]\n"
+                repr_str += repr(submodule) + "\n"
+
+        return repr_str
 
     def __post_init__(self, **kwargs):
         pass

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/src/optimum/rbln/ops/__init__.py

@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .attn import register_rbln_custom_attention, register_rbln_custom_attention_add_softmax
-from .flash_attn import register_rbln_custom_flash_attention
+from .attn import (
+    register_rbln_custom_attention_add_softmax,
+    register_rbln_custom_causal_masked_attention,
+    register_rbln_custom_masked_attention,
+)
+from .flash_attn import register_rbln_custom_flash_causal_masked_attention, register_rbln_custom_flash_masked_attention
 from .kv_cache_update import register_rbln_custom_cache_update

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/src/optimum/rbln/ops/attn.py

@@ -25,13 +25,13 @@ else:
 
 
 @lru_cache
-def register_rbln_custom_attention():
+def register_rbln_custom_masked_attention():
     torch.library.define(
-        "rbln_custom_ops::attn_decode",
+        "rbln_custom_ops::masked_attn_decode",
         "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d) -> Tensor[]",
     )
 
-    @torch.library.impl("rbln_custom_ops::attn_decode", "cpu")
+    @torch.library.impl("rbln_custom_ops::masked_attn_decode", "cpu")
     def attn_decode_cpu(q, k, v, mask, kcache, vcache, seq, scale):
         """Defines the computation pattern for fused attention with KV cache updates.
 
@@ -66,7 +66,7 @@ def register_rbln_custom_attention():
             torch.empty(*vcache.shape, device=vcache.device),
         )
 
-    @register_fake("rbln_custom_ops::attn_decode")
+    @register_fake("rbln_custom_ops::masked_attn_decode")
     def attn_decode_abstract(q, k, v, m, kcache, vcache, seq, partition):
         return (
             q,
@@ -75,11 +75,11 @@ def register_rbln_custom_attention():
         )
 
     torch.library.define(
-        "rbln_custom_ops::attn_prefill",
+        "rbln_custom_ops::masked_attn_prefill",
         "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e) -> Tensor[]",
     )
 
-    @torch.library.impl("rbln_custom_ops::attn_prefill", "cpu")
+    @torch.library.impl("rbln_custom_ops::masked_attn_prefill", "cpu")
     def attn_prefill_cpu(q, k, v, mask, kcache, vcache, batch, seq, scale):
         """Defines the computation pattern for prefill phase attention with KV cache updates.
 
@@ -109,11 +109,99 @@ def register_rbln_custom_attention():
         """
         return q, kcache, vcache
 
-    @register_fake("rbln_custom_ops::attn_prefill")
+    @register_fake("rbln_custom_ops::masked_attn_prefill")
     def attn_prefill_abstract(q, k, v, m, kcache, vcache, batch, seq, partition):
         return q, kcache, vcache
 
 
+@lru_cache
+def register_rbln_custom_causal_masked_attention():
+    torch.library.define(
+        "rbln_custom_ops::causal_masked_attn_decode",
+        "(Tensor x, Tensor y, Tensor z, Tensor a, Tensor b, Tensor c, Tensor d) -> Tensor[]",
+    )
+
+    @torch.library.impl("rbln_custom_ops::causal_masked_attn_decode", "cpu")
+    def attn_decode_cpu(q, k, v, kcache, vcache, seq, scale):
+        """Defines the computation pattern for fused attention with KV cache updates.
+
+        IMPORTANT: This op serves as a pattern definition for the RBLN compiler to generate
+        a single optimized NPU operation. It is NOT meant for CPU execution.
+
+        Pattern components that compiler fuses into a single op:
+        1. KV cache updates with new key/value states
+        2. Scaled dot-product attention computation
+        3. Causal masked softmax operation
+        4. Final attention output computation
+
+        Expected tensor shapes:
+        - q: [batch=1, n_heads, n_groups, 1, head_dim] - Query states for single token
+        - k: [batch=1, n_heads, 1, 1, head_dim] - Key states for current input
+        - v: [batch=1, n_heads, 1, 1, head_dim] - Value states for current input
+        - kcache: [batch_size, n_heads, 1, max_seq_len, head_dim] - Key cache
+        - vcache: [batch_size, n_heads, 1, max_seq_len, head_dim] - Value cache
+        - seq: [1] - Current sequence position
+        - scale: [] - Attention scale factor
+
+        Returns:
+            Tuple[Tensor, Tensor, Tensor]:
+            - attn_output: [batch=1, n_heads, n_groups, 1, head_dim] - Attention output
+            - kcache: Same shape as input kcache, batch=1 - Placeholder for compiler
+            - vcache: Same shape as input vcache, batch=1 - Placeholder for compiler
+        """
+        return (
+            q,
+            torch.empty(*kcache.shape, device=kcache.device),
+            torch.empty(*vcache.shape, device=vcache.device),
+        )
+
+    @register_fake("rbln_custom_ops::causal_masked_attn_decode")
+    def attn_decode_abstract(q, k, v, kcache, vcache, seq, partition):
+        return (
+            q,
+            torch.empty(*kcache.shape, device=kcache.device),
+            torch.empty(*vcache.shape, device=vcache.device),
+        )
+
+    torch.library.define(
+        "rbln_custom_ops::causal_masked_attn_prefill",
+        "(Tensor x, Tensor y, Tensor z, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e) -> Tensor[]",
+    )
+
+    @torch.library.impl("rbln_custom_ops::causal_masked_attn_prefill", "cpu")
+    def attn_prefill_cpu(q, k, v, kcache, vcache, batch, seq, scale):
+        """Defines the computation pattern for prefill phase attention with KV cache updates.
+
+        IMPORTANT: This op serves as a pattern definition for the RBLN compiler to generate
+        a single optimized NPU operation. It is NOT meant for CPU execution.
+
+        Key differences from decode pattern:
+        - Handles prefill phase with multiple input tokens
+        - Takes explicit batch index for continuous batching
+
+        Expected tensor shapes:
+        - q: [batch=1, n_heads, n_groups, seq_len, head_dim] - Query states for multiple tokens
+        - k: [batch=1, n_heads, 1, seq_len, head_dim] - Key states for current input
+        - v: [batch=1, n_heads, 1, seq_len, head_dim] - Value states for current input
+        - kcache: [batch_size, n_heads, 1, max_seq_len, head_dim] - Key cache
+        - vcache: [batch_size, n_heads, 1, max_seq_len, head_dim] - Value cache
+        - batch: [1] - Batch index for cache access
+        - seq: [1] - Starting sequence position
+        - scale: [] - Attention scale factor
+
+        Returns:
+            Tuple[Tensor, Tensor, Tensor]:
+            - attn_output: [batch=1, n_heads, n_groups, seq_len, head_dim] - Attention output
+            - empty_kcache: Same shape as input kcache - Placeholder for compiler
+            - empty_vcache: Same shape as input vcache - Placeholder for compiler
+        """
+        return q, kcache, vcache
+
+    @register_fake("rbln_custom_ops::causal_masked_attn_prefill")
+    def attn_prefill_abstract(q, k, v, kcache, vcache, batch, seq, partition):
+        return q, kcache, vcache
+
+
 @lru_cache
 def register_rbln_custom_attention_add_softmax():
     torch.library.define(

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/src/optimum/rbln/ops/flash_attn.py

@@ -25,13 +25,13 @@ else:
 
 
 @lru_cache
-def register_rbln_custom_flash_attention():
+def register_rbln_custom_flash_masked_attention():
     torch.library.define(
-        "rbln_custom_ops::flash_attn_decode",
+        "rbln_custom_ops::flash_masked_attn_decode",
         "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d, int e) -> Tensor[]",
     )
 
-    @torch.library.impl("rbln_custom_ops::flash_attn_decode", "cpu")
+    @torch.library.impl("rbln_custom_ops::flash_masked_attn_decode", "cpu")
     def flash_attn_decode_cpu(q, k, v, mask, kcache, vcache, seq, scale, partition):
         return (
             q,
@@ -39,7 +39,7 @@ def register_rbln_custom_flash_attention():
             torch.empty(*vcache.shape, device=vcache.device),
         )
 
-    @register_fake("rbln_custom_ops::flash_attn_decode")
+    @register_fake("rbln_custom_ops::flash_masked_attn_decode")
     def flash_attn_decode_abstract(q, k, v, m, kcache, vcache, seq, scale, partition):
         return (
             q,
@@ -48,7 +48,7 @@ def register_rbln_custom_flash_attention():
         )
 
     torch.library.define(
-        "rbln_custom_ops::flash_attn_prefill",
+        "rbln_custom_ops::flash_masked_attn_prefill",
         "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e, int f) -> Tensor[]",
     )
 
@@ -56,6 +56,43 @@ def register_rbln_custom_flash_attention():
     def flash_attn_prefill_cpu(q, k, v, mask, kcache, vcache, batch, seq, scale, partition):
         return q, kcache, vcache
 
-    @register_fake("rbln_custom_ops::flash_attn_prefill")
+    @register_fake("rbln_custom_ops::flash_masked_attn_prefill")
     def flash_attn_prefill_abstract(q, k, v, m, kcache, vcache, batch, seq, scale, partition):
         return q, kcache, vcache
+
+
+@lru_cache
+def register_rbln_custom_flash_causal_masked_attention():
+    torch.library.define(
+        "rbln_custom_ops::flash_causal_masked_attn_decode",
+        "(Tensor x, Tensor y, Tensor z, Tensor a, Tensor b, Tensor c, Tensor d, int e) -> Tensor[]",
+    )
+
+    @torch.library.impl("rbln_custom_ops::flash_causal_masked_attn_decode", "cpu")
+    def flash_attn_decode_cpu(q, k, v, kcache, vcache, seq, scale, partition):
+        return (
+            q,
+            torch.empty(*kcache.shape, device=kcache.device),
+            torch.empty(*vcache.shape, device=vcache.device),
+        )
+
+    @register_fake("rbln_custom_ops::flash_causal_masked_attn_decode")
+    def flash_attn_decode_abstract(q, k, v, kcache, vcache, seq, scale, partition):
+        return (
+            q,
+            torch.empty(*kcache.shape, device=kcache.device),
+            torch.empty(*vcache.shape, device=vcache.device),
+        )
+
+    torch.library.define(
+        "rbln_custom_ops::flash_causal_masked_attn_prefill",
+        "(Tensor x, Tensor y, Tensor z, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e, int f) -> Tensor[]",
+    )
+
+    @torch.library.impl("rbln_custom_ops::flash_causal_masked_attn_prefill", "cpu")
+    def flash_attn_prefill_cpu(q, k, v, kcache, vcache, batch, seq, scale, partition):
+        return q, kcache, vcache
+
+    @register_fake("rbln_custom_ops::flash_causal_masked_attn_prefill")
+    def flash_attn_prefill_abstract(q, k, v, kcache, vcache, batch, seq, scale, partition):
+        return q, kcache, vcache

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/src/optimum/rbln/transformers/modeling_generic.py

@@ -73,7 +73,7 @@ class RBLNModelForQuestionAnswering(RBLNModel):
         if rbln_batch_size is None:
             rbln_batch_size = 1
 
-        signature_params = inspect.signature(cls.hf_class.forward).parameters.keys()
+        signature_params = inspect.signature(cls.get_hf_class().forward).parameters.keys()
 
         if rbln_model_input_names is None:
             for tokenizer in preprocessors:
@@ -289,7 +289,7 @@ class RBLNModelForSequenceClassification(RBLNModel):
         if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
             raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
 
-        signature_params = inspect.signature(cls.hf_class.forward).parameters.keys()
+        signature_params = inspect.signature(cls.get_hf_class().forward).parameters.keys()
 
         if rbln_model_input_names is None:
             for tokenizer in preprocessors:
@@ -362,7 +362,7 @@ class RBLNModelForMaskedLM(RBLNModel):
         if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
             raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
 
-        signature_params = inspect.signature(cls.hf_class.forward).parameters.keys()
+        signature_params = inspect.signature(cls.get_hf_class().forward).parameters.keys()
 
         if rbln_model_input_names is None:
             for tokenizer in preprocessors:

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/src/optimum/rbln/transformers/models/bart/bart_architecture.py

@@ -142,7 +142,7 @@ class BartSelfAttention(Seq2SeqSelfAttention):
         self.num_heads = self._original_mod.num_heads
         self.head_dim = self._original_mod.embed_dim // self._original_mod.num_heads
         self.scaling = self.head_dim**-0.5
-        self.attn_decode = torch.ops.rbln_custom_ops.attn_decode
+        self.attn_decode = torch.ops.rbln_custom_ops.masked_attn_decode
 
     def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         query_states = self.q_proj(hidden_states) * self.scaling

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/src/optimum/rbln/transformers/models/bart/modeling_bart.py

@@ -58,7 +58,7 @@ class RBLNBartModel(RBLNModel):
         if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
             raise ValueError("`rbln_max_seq_len` should be less or equal than max_position_embeddings!")
 
-        signature_params = inspect.signature(cls.hf_class.forward).parameters.keys()
+        signature_params = inspect.signature(cls.get_hf_class().forward).parameters.keys()
 
         if rbln_model_input_names is None:
             for tokenizer in preprocessors:

{optimum_rbln-0.7.2rc2 → optimum_rbln-0.7.3a0}/src/optimum/rbln/transformers/models/bert/modeling_bert.py

@@ -56,7 +56,7 @@ class RBLNBertModel(RBLNModel):
         if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
             raise ValueError("`rbln_max_seq_len` should be less or equal than max_position_embeddings!")
 
-        signature_params = inspect.signature(cls.hf_class.forward).parameters.keys()
+        signature_params = inspect.signature(cls.get_hf_class().forward).parameters.keys()
 
         if rbln_model_input_names is None:
             for tokenizer in preprocessors: