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

optimum_rbln-0.7.3a3/.github/version.yaml

@@ -0,0 +1 @@
+rebel_compiler_version: 0.7.3.dev150+g305d8758

optimum_rbln-0.7.3a3/.github/workflows/rbln_check_compiler.yaml

@@ -0,0 +1,61 @@
+name: Check Rebel Compiler Version
+
+on:
+  workflow_call:
+    inputs:
+      compiler_version:
+        description: "Rebel compiler version to check"
+        required: true
+        type: string
+    outputs:
+      is_deploy_version:
+        description: "Whether the compiler is a deploy version"
+        value: ${{ jobs.check-rebel-compiler-version.outputs.is_deploy_version }}
+      compiler_version:
+        description: "The compiler version used"
+        value: ${{ jobs.check-rebel-compiler-version.outputs.compiler_version }}
+
+jobs:
+  check-rebel-compiler-version:
+    runs-on: rebel-k8s-runner
+    outputs:
+      is_deploy_version: ${{ steps.check_prod.outputs.IS_PROD }}
+      compiler_version: ${{ steps.install_compiler.outputs.COMPILER_VERSION }}
+    steps:
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.9"
+
+      - name: Install rebel-compiler
+        id: install_compiler
+        env:
+          REBEL_PYPI_ENDPOINT: ${{ vars.REBEL_PYPI_INTERNAL_ENDPOINT }}
+          REBEL_PYPI_USERNAME: ${{ secrets.REBEL_PYPI_USERNAME }}
+          REBEL_PYPI_PASSWORD: ${{ secrets.REBEL_PYPI_PASSWORD }}
+        run: |
+          set -e
+          PYPI_URL=$(echo $REBEL_PYPI_ENDPOINT | sed "s/\/\//\0$REBEL_PYPI_USERNAME:$REBEL_PYPI_PASSWORD@/")
+          pip3 install --extra-index-url $PYPI_URL rebel-compiler==${{ inputs.compiler_version }}
+          echo "COMPILER_VERSION=${{ inputs.compiler_version }}" >> $GITHUB_OUTPUT
+
+      - name: Run script to check ENV
+        id: check_env
+        run: |
+          echo "Running check for rebel-compiler version"
+          ENV_VALUE=$(python3 -c '${{ secrets.CHECK_DEPLOY }}')
+          echo $ENV_VALUE
+          echo "ENV_VALUE=$ENV_VALUE" >> $GITHUB_ENV
+
+      - name: Verify ENV is PROD
+        id: check_prod
+        run: |
+          if [ "$ENV_VALUE" = "PROD" ]; then
+            echo "IS_PROD=true" >> $GITHUB_OUTPUT
+            echo "version check pass(✅)!!"
+          else
+            echo "IS_PROD=false" >> $GITHUB_OUTPUT
+            echo "version check fail(❌)!!" 
+            echo "rebel-compiler must be prod version"
+            exit 1
+          fi 

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

@@ -15,6 +15,11 @@ on:
         description: "rebel_compiler version to run"
         required: true
         type: string
+      test_level:
+        description: "Test level for OPTIMUM_RBLN_TEST_LEVEL (default, full, essential)"
+        required: false
+        type: string
+        default: "default"
 
 env:
   REBEL_PYPI_ENDPOINT: ${{ vars.REBEL_PYPI_INTERNAL_ENDPOINT }}
@@ -53,18 +58,18 @@ jobs:
 
       - name: Run pytest (transformers)
         env:
-          OPTIMUM_RBLN_TEST_LEVEL: default
+          OPTIMUM_RBLN_TEST_LEVEL: ${{ inputs.test_level }}
         run: |
           uv run --no-sync pytest tests/test_transformers.py
 
       - name: Run pytest (diffusers)
         env:
-          OPTIMUM_RBLN_TEST_LEVEL: default
+          OPTIMUM_RBLN_TEST_LEVEL: ${{ inputs.test_level }}
         run: |
           uv run --no-sync pytest tests/test_diffusers.py
 
       - name: Run pytest (llm)
         env:
-          OPTIMUM_RBLN_TEST_LEVEL: default
+          OPTIMUM_RBLN_TEST_LEVEL: ${{ inputs.test_level }}
         run: |
           uv run --no-sync pytest tests/test_llm.py      

optimum_rbln-0.7.3a3/.github/workflows/rbln_scheduled_test.yaml

@@ -0,0 +1,51 @@
+name: Optimum-rbln / Scheduled Test
+
+on:
+  schedule:
+    # Run every day at 2am (17:00 UTC, 2:00am KST)
+    - cron: '0 17 * * *'
+
+env:
+  HF_USER_ID: ${{ secrets.HF_USER_ID }}
+  HF_AUTH_TOKEN: ${{ secrets.HF_AUTH_TOKEN }}
+
+jobs:
+  load-version:
+    runs-on: rebel-k8s-runner
+    outputs:
+      compiler_version: ${{ steps.get_version.outputs.compiler_version }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+      
+      - name: Get compiler version
+        id: get_version
+        run: |
+          VERSION=$(grep rebel_compiler_version .github/version.yaml | cut -d ':' -f2 | tr -d ' ')
+          echo "compiler_version=$VERSION" >> $GITHUB_OUTPUT
+
+  check-compiler:
+    needs: load-version
+    uses: ./.github/workflows/rbln_check_compiler.yaml
+    with:
+      compiler_version: ${{ needs.load-version.outputs.compiler_version }}
+    secrets: inherit
+
+  optimum-rbln-pytest:
+    needs: [load-version, check-compiler]
+    if: ${{ needs.check-compiler.outputs.is_deploy_version == 'true' }}
+    uses: ./.github/workflows/rbln_optimum_pytest.yaml
+    with:
+      ref: main
+      rebel_compiler_version: ${{ needs.check-compiler.outputs.compiler_version }}
+      test_level: "full"
+    secrets: inherit
+
+  optimum-rbln-inference-test:
+    needs: check-compiler
+    if: ${{ needs.check-compiler.outputs.is_deploy_version == 'true' }}
+    uses: ./.github/workflows/rbln_optimum_inference_test.yaml
+    with:
+      ref: main
+      rebel_compiler_version: ${{ needs.check-compiler.outputs.compiler_version }}
+    secrets: inherit

optimum_rbln-0.7.3a3/.github/workflows/rbln_trigger_on_pr.yaml

@@ -0,0 +1,94 @@
+name: Optimum-rbln / PR
+
+on:
+  pull_request:
+    branches:
+      - main
+
+env:
+  REBEL_PYPI_ENDPOINT: ${{ vars.REBEL_PYPI_INTERNAL_ENDPOINT }}
+  REBEL_PYPI_USERNAME: ${{ secrets.REBEL_PYPI_USERNAME }}
+  REBEL_PYPI_PASSWORD: ${{ secrets.REBEL_PYPI_PASSWORD }}
+
+jobs:
+  check-skip-ci:
+    runs-on: rebel-k8s-runner
+    outputs:
+      should_skip: ${{ contains(github.event.pull_request.head.commit.message, '[skip ci]') }}
+    steps:
+      - name: Check if [skip ci] is in commit message
+        run: |
+          if ${{ contains(github.event.pull_request.head.commit.message, '[skip ci]') }}; then
+            echo "Found [skip ci] in commit message, skipping CI"
+          else
+            echo "No [skip ci] found, continuing with CI"
+          fi
+
+  load-version:
+    runs-on: rebel-k8s-runner
+    needs: check-skip-ci
+    if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' }}
+    outputs:
+      compiler_version: ${{ steps.get_version.outputs.compiler_version }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+      
+      - name: Get compiler version
+        id: get_version
+        run: |
+          VERSION=$(grep rebel_compiler_version .github/version.yaml | cut -d ':' -f2 | tr -d ' ')
+          echo "compiler_version=$VERSION" >> $GITHUB_OUTPUT
+
+  check-compiler:
+    needs: [check-skip-ci, load-version]
+    if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' }}
+    uses: ./.github/workflows/rbln_check_compiler.yaml
+    with:
+      compiler_version: ${{ needs.load-version.outputs.compiler_version }}
+    secrets: inherit
+
+  check-team-member:
+    runs-on: rebel-k8s-runner
+    needs: [check-skip-ci, check-compiler]
+    if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' && needs.check-compiler.outputs.is_deploy_version == 'true' }}
+    outputs:
+      is_team_member: ${{ steps.check_member.outputs.IS_TEAM_MEMBER }}
+    steps:
+      - name: Fetch team members
+        id: fetch_team
+        run: |
+          response=$(curl -s -H "Authorization: Bearer ${{ secrets.GIT_PAT }}" \
+            -H "Content-Type: application/json" \
+            -d '{"query":"query { organization(login: \"rebellions-sw\") { team(slug: \"rebel-sw-team\") { members(first: 100) { nodes { login } } } } }"}' \
+            https://api.github.com/graphql)
+          echo "$response" | jq -r '.data.organization.team.members.nodes[].login' > team_members.txt
+
+      - name: Check if PR author is a team member
+        id: check_member
+        run: |
+          pr_author=${{ github.event.pull_request.user.login }}
+          if grep -qx "$pr_author" team_members.txt; then
+            echo "IS_TEAM_MEMBER=true" >> $GITHUB_OUTPUT
+          else
+            echo "IS_TEAM_MEMBER=false" >> $GITHUB_OUTPUT
+          fi
+
+  optimum-rbln-pytest:
+    needs: [check-skip-ci, check-compiler, check-team-member]
+    if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' && needs.check-team-member.outputs.is_team_member == 'true' }}
+    uses: ./.github/workflows/rbln_optimum_pytest.yaml
+    with:
+      ref: ${{ github.event.pull_request.head.sha }}
+      rebel_compiler_version: ${{ needs.check-compiler.outputs.compiler_version }}
+      test_level: "default"
+    secrets: inherit
+
+  optimum-rbln-inference-test:
+    needs: [check-skip-ci, check-compiler, check-team-member]
+    if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' && needs.check-team-member.outputs.is_team_member == 'true' }}
+    uses: ./.github/workflows/rbln_optimum_inference_test.yaml
+    with:
+      ref: ${{ github.event.pull_request.head.sha }}
+      rebel_compiler_version: ${{ needs.check-compiler.outputs.compiler_version }}
+    secrets: inherit

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

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

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

{optimum_rbln-0.7.3a1 → optimum_rbln-0.7.3a3}/src/optimum/rbln/ops/__init__.py

@@ -13,9 +13,9 @@
 # limitations under the License.
 
 from .attn import (
-    register_rbln_custom_attention_add_softmax,
-    register_rbln_custom_causal_masked_attention,
-    register_rbln_custom_masked_attention,
+    register_rbln_custom_add_softmax_attention,
+    register_rbln_custom_paged_attention,
+    register_rbln_custom_paged_causal_attention,
 )
-from .flash_attn import register_rbln_custom_flash_causal_masked_attention, register_rbln_custom_flash_masked_attention
+from .flash_attn import register_rbln_custom_paged_flash_attention, register_rbln_custom_paged_flash_causal_attention
 from .kv_cache_update import register_rbln_custom_cache_update

{optimum_rbln-0.7.3a1 → optimum_rbln-0.7.3a3}/src/optimum/rbln/ops/attn.py

@@ -25,14 +25,14 @@ else:
 
 
 @lru_cache
-def register_rbln_custom_masked_attention():
+def register_rbln_custom_paged_attention():
     torch.library.define(
-        "rbln_custom_ops::masked_attn_decode",
-        "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d) -> Tensor[]",
+        "rbln_custom_ops::paged_attn_decode",
+        "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e, int f) -> Tensor[]",
     )
 
-    @torch.library.impl("rbln_custom_ops::masked_attn_decode", "cpu")
-    def attn_decode_cpu(q, k, v, mask, kcache, vcache, seq, scale):
+    @torch.library.impl("rbln_custom_ops::paged_attn_decode", "cpu")
+    def attn_decode_cpu(q, k, v, mask, kcache, vcache, seq, scale, block_table, block_size):
         """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
@@ -51,8 +51,10 @@ def register_rbln_custom_masked_attention():
         - 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
+        - seq: [1, 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:
             Tuple[Tensor, Tensor, Tensor]:
@@ -66,8 +68,8 @@ def register_rbln_custom_masked_attention():
             torch.empty(*vcache.shape, device=vcache.device),
         )
 
-    @register_fake("rbln_custom_ops::masked_attn_decode")
-    def attn_decode_abstract(q, k, v, m, kcache, vcache, seq, partition):
+    @register_fake("rbln_custom_ops::paged_attn_decode")
+    def attn_decode_abstract(q, k, v, m, kcache, vcache, seq, scale, block_table, block_size):
         return (
             q,
             torch.empty(*kcache.shape, device=kcache.device),
@@ -75,12 +77,12 @@ def register_rbln_custom_masked_attention():
         )
 
     torch.library.define(
-        "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[]",
+        "rbln_custom_ops::paged_attn_prefill",
+        "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e, int f) -> Tensor[]",
     )
 
-    @torch.library.impl("rbln_custom_ops::masked_attn_prefill", "cpu")
-    def attn_prefill_cpu(q, k, v, mask, kcache, vcache, batch, seq, scale):
+    @torch.library.impl("rbln_custom_ops::paged_attn_prefill", "cpu")
+    def attn_prefill_cpu(q, k, v, mask, kcache, vcache, seq, scale, block_table, block_size):
         """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
@@ -97,9 +99,10 @@ def register_rbln_custom_masked_attention():
         - 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
-        - batch: [1] - Batch index for cache access
-        - seq: [1] - Starting sequence position
+        - 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:
             Tuple[Tensor, Tensor, Tensor]:
@@ -109,20 +112,20 @@ def register_rbln_custom_masked_attention():
         """
         return q, kcache, vcache
 
-    @register_fake("rbln_custom_ops::masked_attn_prefill")
-    def attn_prefill_abstract(q, k, v, m, kcache, vcache, batch, seq, partition):
+    @register_fake("rbln_custom_ops::paged_attn_prefill")
+    def attn_prefill_abstract(q, k, v, m, kcache, vcache, seq, scale, block_table, block_size):
         return q, kcache, vcache
 
 
 @lru_cache
-def register_rbln_custom_causal_masked_attention():
+def register_rbln_custom_paged_causal_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[]",
+        "rbln_custom_ops::paged_causal_attn_decode",
+        "(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::causal_masked_attn_decode", "cpu")
-    def attn_decode_cpu(q, k, v, kcache, vcache, seq, scale):
+    @torch.library.impl("rbln_custom_ops::paged_causal_attn_decode", "cpu")
+    def attn_decode_cpu(q, k, v, kcache, vcache, seq, scale, block_table, block_size):
         """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
@@ -140,8 +143,10 @@ def register_rbln_custom_causal_masked_attention():
         - 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
+        - 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:
             Tuple[Tensor, Tensor, Tensor]:
@@ -155,8 +160,8 @@ def register_rbln_custom_causal_masked_attention():
             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):
+    @register_fake("rbln_custom_ops::paged_causal_attn_decode")
+    def attn_decode_abstract(q, k, v, kcache, vcache, seq, scale, block_table, block_size):
         return (
             q,
             torch.empty(*kcache.shape, device=kcache.device),
@@ -164,12 +169,12 @@ def register_rbln_custom_causal_masked_attention():
         )
 
     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[]",
+        "rbln_custom_ops::paged_causal_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::causal_masked_attn_prefill", "cpu")
-    def attn_prefill_cpu(q, k, v, kcache, vcache, batch, seq, scale):
+    @torch.library.impl("rbln_custom_ops::paged_causal_attn_prefill", "cpu")
+    def attn_prefill_cpu(q, k, v, kcache, vcache, seq, scale, block_table, block_size):
         """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
@@ -186,8 +191,10 @@ def register_rbln_custom_causal_masked_attention():
         - 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
+        - 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:
             Tuple[Tensor, Tensor, Tensor]:
@@ -197,20 +204,20 @@ def register_rbln_custom_causal_masked_attention():
         """
         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):
+    @register_fake("rbln_custom_ops::paged_causal_attn_prefill")
+    def attn_prefill_abstract(q, k, v, kcache, vcache, seq, scale, block_table, block_size):
         return q, kcache, vcache
 
 
 @lru_cache
-def register_rbln_custom_attention_add_softmax():
+def register_rbln_custom_add_softmax_attention():
     torch.library.define(
-        "rbln_custom_ops::attn_decode_add_softmax",
+        "rbln_custom_ops::add_softmax_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_add_softmax", "cpu")
-    def attn_decode_add_softmax_cpu(q, k, v, mask, kcache, vcache, seq, scale):
+    @torch.library.impl("rbln_custom_ops::add_softmax_attn_decode", "cpu")
+    def add_softmax_attn_decode_cpu(q, k, v, mask, 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
@@ -244,57 +251,10 @@ def register_rbln_custom_attention_add_softmax():
             torch.empty(*vcache.shape, device=vcache.device),
         )
 
-    @register_fake("rbln_custom_ops::attn_decode_add_softmax")
-    def attn_decode_add_softmax_abstract(q, k, v, m, kcache, vcache, seq, partition):
+    @register_fake("rbln_custom_ops::add_softmax_attn_decode")
+    def add_softmax_attn_decode_abstract(q, k, v, m, 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::attn_prefill_add_softmax",
-        "(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_add_softmax", "cpu")
-    def attn_prefill_add_softmax_cpu(q, k, v, mask, 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
-        - 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
-        - 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, seq_len, 1, 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,
-            torch.empty(1, *kcache.shape[1:], device=kcache.device),
-            torch.empty(1, *vcache.shape[1:], device=vcache.device),
-        )
-
-    @register_fake("rbln_custom_ops::attn_prefill_add_softmax")
-    def attn_prefill_add_softmax_abstract(q, k, v, m, kcache, vcache, batch, seq, partition):
-        return (
-            q,
-            torch.empty(1, *kcache.shape[1:], device=kcache.device),
-            torch.empty(1, *vcache.shape[1:], device=vcache.device),
-        )

{optimum_rbln-0.7.3a1 → optimum_rbln-0.7.3a3}/src/optimum/rbln/ops/flash_attn.py

@@ -25,22 +25,22 @@ else:
 
 
 @lru_cache
-def register_rbln_custom_flash_masked_attention():
+def register_rbln_custom_paged_flash_attention():
     torch.library.define(
-        "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[]",
+        "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::flash_masked_attn_decode", "cpu")
-    def flash_attn_decode_cpu(q, k, v, mask, kcache, vcache, seq, scale, partition):
+    @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,
             torch.empty(*kcache.shape, device=kcache.device),
             torch.empty(*vcache.shape, device=vcache.device),
         )
 
-    @register_fake("rbln_custom_ops::flash_masked_attn_decode")
-    def flash_attn_decode_abstract(q, k, v, m, kcache, vcache, seq, scale, partition):
+    @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.empty(*kcache.shape, device=kcache.device),
@@ -48,36 +48,36 @@ def register_rbln_custom_flash_masked_attention():
         )
 
     torch.library.define(
-        "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[]",
+        "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, batch, seq, scale, partition):
+    def flash_attn_prefill_cpu(q, k, v, mask, kcache, vcache, seq, scale, block_table, block_size, partition):
         return q, kcache, vcache
 
-    @register_fake("rbln_custom_ops::flash_masked_attn_prefill")
-    def flash_attn_prefill_abstract(q, k, v, m, kcache, vcache, batch, seq, scale, partition):
+    @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, kcache, vcache
 
 
 @lru_cache
-def register_rbln_custom_flash_causal_masked_attention():
+def register_rbln_custom_paged_flash_causal_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[]",
+        "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::flash_causal_masked_attn_decode", "cpu")
-    def flash_attn_decode_cpu(q, k, v, kcache, vcache, seq, scale, partition):
+    @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,
             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):
+    @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.empty(*kcache.shape, device=kcache.device),
@@ -85,14 +85,14 @@ def register_rbln_custom_flash_causal_masked_attention():
         )
 
     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[]",
+        "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::flash_causal_masked_attn_prefill", "cpu")
-    def flash_attn_prefill_cpu(q, k, v, kcache, vcache, batch, seq, scale, partition):
+    @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, 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):
+    @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, kcache, vcache