optimum-rbln 0.7.4a1__tar.gz → 0.7.4a3__tar.gz

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a3}/.github/scripts/auto_code_review.py

@@ -22,8 +22,7 @@ from github import Github
 
 model_name = os.environ["GOOGLE_MODEL_ID"]
 genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
-max_context_token = 100000
-force_review = False
+max_context_token = 500000
 
 
 def get_pr_diff():
@@ -58,19 +57,6 @@ Review the following code changes(GIT DIFF) along with the pull request (PR) det
     return system_prompt, prompt
 
 
-def translate_review(review):
-    model = genai.GenerativeModel(
-        model_name,
-        system_instruction="You are a professional translator specializing in technical and software-related content. Keep the technical words in English, but understand the whole sentence and rephrase it in Korean.",
-    )
-    prompt = f"""Translate the following English text into Korean, maintaining technical accuracy and clarity. Include ONLY the translation, NO OTHER EXPLANATIONS or RESPONSES as a chatbot. :
-
-{review}"""
-    response = model.generate_content(prompt)
-
-    return response.text
-
-
 def review_code(system_prompt, prompt):
     model = genai.GenerativeModel(model_name, system_instruction=system_prompt)
     response = model.generate_content(prompt)
@@ -125,7 +111,7 @@ def main():
     system_prompt, prompt = get_prompt(diff, pr)
     model = genai.GenerativeModel(model_name=model_name, system_instruction=system_prompt)
     num_tokens = model.count_tokens(prompt).total_tokens
-    if num_tokens > max_context_token and not force_review:
+    if num_tokens > max_context_token:
         msg = f"Diff ({len(diff)}) exceeds maximum allowed tokens ({max_context_token}) > ({num_tokens}). Skipping review."
         print(msg)
         pr.create_issue_comment(msg)
@@ -133,14 +119,10 @@ def main():
 
     # Get Auto review
     review = review_code(system_prompt, prompt)
-    translation = translate_review(review)
 
     # Post comment on PR
-    pr.create_issue_comment(f"""# Auto Code Review
-
-- [참고] Auto Code Review를 invoke하려면, commit message의 시작을 [autoreview]로 시작하거나, "/autoreview" 를 comment로 작성한 후,
-해당 commit의 github action에서 code review를 re-run 하시면 됩니다.
-\n\n{review}\n\n{translation}""")
+    pr.create_issue_comment(f"""# Auto Code Review by {model_name}
+\n\n{review}""")
 
 
 if __name__ == "__main__":

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a3}/.github/workflows/auto_code_review.yml

@@ -14,7 +14,7 @@ env:
 
 jobs:
   auto-review:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-latest-rbln
     steps:
       - name: Checkout repository
         uses: actions/checkout@v3

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a3}/.github/workflows/check_code_quality.yml

@@ -25,7 +25,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: [3.9]
-        os: [ubuntu-latest]
+        os: [ubuntu-latest-rbln]
 
     runs-on: ${{ matrix.os }}
     steps:

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a3}/.github/workflows/deploy-on-tag.yaml

@@ -28,7 +28,7 @@ jobs:
               workflow_id: 'rebel_dispatch_model_generation_for_vllm.yaml',
               ref: 'dev',
               inputs: {
-                optimum_rbln_version: ${{ github.ref_name }},
+                optimum_rbln_version: "${{ github.ref_name }}",
               }
             })
             console.log(result)

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a3}/.github/workflows/pr-title-check.yaml

@@ -20,7 +20,7 @@ permissions:
 jobs:
   title-checker:
     name: Check PR title
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-latest-rbln
     steps:
       - uses: amannn/action-semantic-pull-request@v5
         id: lint_pr_title

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a3}/.github/workflows/pr_checklist_validator.yml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   validate-pr-checklist:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-latest-rbln
     steps:
       - name: Checkout repository
         uses: actions/checkout@v2

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a3}/.github/workflows/rbln_optimum_pytest.yaml

@@ -43,7 +43,9 @@ jobs:
         if: ${{ inputs.commit_message == '' }}
         run: |
           COMMIT_MESSAGE=$(git log -1 --pretty=%B)
-          echo "message=$COMMIT_MESSAGE" >> $GITHUB_OUTPUT
+          echo "message<<EOF" >> $GITHUB_OUTPUT
+          echo "$COMMIT_MESSAGE" >> $GITHUB_OUTPUT
+          echo "EOF" >> $GITHUB_OUTPUT
 
       - name: Setup uv
         uses: astral-sh/setup-uv@v3

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.7.4a1
+Version: 0.7.4a3
 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.4a1 → optimum_rbln-0.7.4a3}/src/optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.4a1'
+__version__ = version = '0.7.4a3'
 __version_tuple__ = version_tuple = (0, 7, 4)

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a3}/src/optimum/rbln/modeling.py

@@ -123,8 +123,15 @@ class RBLNModel(RBLNBaseModel):
                 config = AutoConfig.from_pretrained(config._name_or_path, **kwargs)
 
         if hasattr(model, "can_generate") and model.can_generate():
+            import json
+
             generation_config = model.generation_config
-            generation_config.save_pretrained(save_dir_path / subfolder)
+            generation_config_path = save_dir_path / subfolder / "generation_config.json"
+
+            generation_config.save_pretrained(generation_config_path.parent)
+            local_config = json.loads(generation_config_path.read_text(encoding="utf-8"))
+            local_config["transformers_version"] = generation_config.transformers_version
+            generation_config_path.write_text(json.dumps(local_config, indent=2) + "\n", encoding="utf-8")
 
         if not isinstance(config, PretrainedConfig):  # diffusers config
             config = PretrainedConfig(**config)

optimum_rbln-0.7.4a3/src/optimum/rbln/ops/__init__.py

@@ -0,0 +1,18 @@
+# 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 .attn import *
+from .flash_attn import *
+from .kv_cache_update import *
+from .linear import linear

optimum_rbln-0.7.4a3/src/optimum/rbln/ops/attn.py

@@ -0,0 +1,287 @@
+# 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.
+
+
+import torch
+from torch import Tensor
+
+
+@torch.library.custom_op(
+    "rbln_custom_ops::paged_attn_decode",
+    mutates_args=(["kcache", "vcache"]),
+)
+def paged_attn_decode(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+) -> Tensor:
+    return torch.empty_like(q)
+
+
+@paged_attn_decode.register_fake
+def paged_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,
+) -> Tensor:
+    return torch.empty_like(q)
+
+
+@torch.library.custom_op(
+    "rbln_custom_ops::paged_attn_prefill",
+    mutates_args=(["kcache", "vcache"]),
+)
+def paged_attn_prefill(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+) -> Tensor:
+    """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
+    - mask: [batch=1, 1, 1, seq_len, max_seq_len] - Attention mask
+    - 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, 1] - Starting sequence position
+    - scale: [] - Attention scale factor
+    - block_table: [batch_size, max_seq_len // block_size] - Block indices for KV cache management
+    - block_size: [] - Number of tokens per block
+
+    Returns:
+        Tensor: attn_output: [batch=1, n_heads, n_groups, seq_len, head_dim] - Attention output
+    """
+    return torch.empty_like(q)
+
+
+@paged_attn_prefill.register_fake
+def paged_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,
+) -> Tensor:
+    return torch.empty_like(q)
+
+
+@torch.library.custom_op(
+    "rbln_custom_ops::paged_causal_attn_decode",
+    mutates_args=(["kcache", "vcache"]),
+)
+def paged_causal_attn_decode(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+) -> Tensor:
+    """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, 1] - Starting sequence position
+    - scale: [] - Attention scale factor
+    - block_table: [batch_size, max_seq_len // block_size] - Block indices for KV cache management
+    - block_size: [] - Number of tokens per block
+
+    Returns:
+        Tensor: attn_output: [batch=1, n_heads, n_groups, 1, head_dim] - Attention output
+    """
+    return torch.empty_like(q)
+
+
+@paged_causal_attn_decode.register_fake
+def paged_causal_attn_decode_fake(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+) -> Tensor:
+    return torch.empty_like(q)
+
+
+@torch.library.custom_op(
+    "rbln_custom_ops::paged_causal_attn_prefill",
+    mutates_args=(["kcache", "vcache"]),
+)
+def paged_causal_attn_prefill(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+) -> Tensor:
+    """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, 1] - Starting sequence position
+    - scale: [] - Attention scale factor
+    - block_table: [batch_size, max_seq_len // block_size] - Block indices for KV cache management
+    - block_size: [] - Number of tokens per block
+
+    Returns:
+        Tensor: attn_output: [batch=1, n_heads, n_groups, seq_len, head_dim] - Attention output
+    """
+    return torch.empty_like(q)
+
+
+@paged_causal_attn_prefill.register_fake
+def paged_causal_attn_prefill_fake(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+) -> Tensor:
+    return torch.empty_like(q)
+
+
+@torch.library.custom_op(
+    "rbln_custom_ops::paged_add_softmax_attn_decode",
+    mutates_args=(["kcache", "vcache"]),
+)
+def paged_add_softmax_attn_decode(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor,
+    kcache: Tensor,
+    vcache: Tensor,
+    seq: Tensor,
+    scale: Tensor,
+    block_table: Tensor,
+    block_size: int,
+) -> Tensor:
+    """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. add-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
+    - mask: [batch=1, n_heads, 1, 1, max_seq_len] - Attention mask
+    - 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
+    - block_table: [batch_size, max_seq_len // block_size] - Block indices for KV cache management
+    - block_size: [] - Number of tokens per block
+
+    Returns:
+        Tensor: attn_output: [batch=1, n_heads, 1, 1, head_dim] - Attention output
+    """
+    return torch.empty_like(q)
+
+
+@paged_add_softmax_attn_decode.register_fake
+def paged_add_softmax_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,
+) -> Tensor:
+    return torch.empty_like(q)

optimum_rbln-0.7.4a3/src/optimum/rbln/ops/flash_attn.py

@@ -0,0 +1,176 @@
+# 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.
+
+import torch
+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-0.7.4a3/src/optimum/rbln/ops/kv_cache_update.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.
+
+import torch
+from torch import Tensor
+
+
+@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.
+    return torch.empty_like(cache)

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a3}/src/optimum/rbln/transformers/models/bart/__init__.py

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from ....ops import paged_attn_decode, paged_causal_attn_decode
 from .modeling_bart import RBLNBartForConditionalGeneration, RBLNBartModel

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a3}/src/optimum/rbln/transformers/models/decoderonly/__init__.py

@@ -12,4 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from ....ops import (
+    paged_attn_decode,
+    paged_attn_prefill,
+    paged_causal_attn_decode,
+    paged_causal_attn_prefill,
+    paged_flash_attn_decode,
+    paged_flash_attn_prefill,
+    paged_flash_causal_attn_decode,
+    paged_flash_causal_attn_prefill,
+)
 from .modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM