optimum-rbln 0.7.3a6__tar.gz → 0.7.3.post1__tar.gz

optimum_rbln-0.7.3.post1/.github/version.yaml

@@ -0,0 +1 @@
+rebel_compiler_version: 0.7.3

{optimum_rbln-0.7.3a6 → optimum_rbln-0.7.3.post1}/.github/workflows/rbln_optimum_pytest.yaml

@@ -38,6 +38,13 @@ jobs:
           submodules: recursive
           fetch-depth: 0
 
+      - name: Get commit message if not provided
+        id: get_commit_message
+        if: ${{ inputs.commit_message == '' }}
+        run: |
+          COMMIT_MESSAGE=$(git log -1 --pretty=%B)
+          echo "message=$COMMIT_MESSAGE" >> $GITHUB_OUTPUT
+
       - name: Setup uv
         uses: astral-sh/setup-uv@v3
         with:
@@ -55,21 +62,34 @@ jobs:
         run: |
           PYPI_URL=$(echo ${{ env.REBEL_PYPI_ENDPOINT }} | sed "s/\/\//\0${{ env.REBEL_PYPI_USERNAME}}:${{ env.REBEL_PYPI_PASSWORD}}@/")
           uv pip install --extra-index-url $PYPI_URL rebel-compiler==${{ inputs.rebel_compiler_version }}
-
+      
       - name: Run pytest (transformers)
         env:
           OPTIMUM_RBLN_TEST_LEVEL: ${{ inputs.test_level }}
         run: |
-          uv run --no-sync pytest tests/test_transformers.py
+          echo 
+          if ${{ !contains( steps.get_commit_message.outputs.message , '[skip-transformers]') }}; then
+            uv run --no-sync pytest tests/test_transformers.py
+          else
+            echo "Found [skip-transformers] in commit message, skipping CI"
+          fi
 
       - name: Run pytest (diffusers)
         env:
           OPTIMUM_RBLN_TEST_LEVEL: ${{ inputs.test_level }}
         run: |
-          uv run --no-sync pytest tests/test_diffusers.py
+          if ${{ !contains( steps.get_commit_message.outputs.message , '[skip-diffusers]') }}; then
+            uv run --no-sync pytest tests/test_diffusers.py
+          else
+            echo "Found [skip-diffusers] in commit message, skipping CI"
+          fi
 
       - name: Run pytest (llm)
         env:
           OPTIMUM_RBLN_TEST_LEVEL: ${{ inputs.test_level }}
         run: |
-          uv run --no-sync pytest tests/test_llm.py      
+          if ${{ !contains( steps.get_commit_message.outputs.message , '[skip-llms]') }}; then
+            uv run --no-sync pytest tests/test_llm.py
+          else
+            echo "Found [skip-llms] in commit message, skipping CI"
+          fi

{optimum_rbln-0.7.3a6 → optimum_rbln-0.7.3.post1}/PKG-INFO

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

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

{optimum_rbln-0.7.3a6 → optimum_rbln-0.7.3.post1}/src/optimum/rbln/diffusers/modeling_diffusers.py

@@ -23,7 +23,6 @@ from ..modeling import RBLNModel
 from ..modeling_config import RUNTIME_KEYWORDS, ContextRblnConfig, use_rbln_config
 from ..utils.decorator_utils import remove_compile_time_kwargs
 from ..utils.logging import get_logger
-from . import pipelines
 
 
 logger = get_logger(__name__)
@@ -67,6 +66,7 @@ class RBLNDiffusionMixin:
           as keys in rbln_config
     """
 
+    _connected_classes = {}
     _submodules = []
     _prefix = {}
 
@@ -103,37 +103,6 @@ class RBLNDiffusionMixin:
                 }
             )
             submodule_config = submodule_cls.update_rbln_config_using_pipe(model, submodule_config)
-        elif hasattr(pipelines, submodule_class_name):
-            submodule_config = rbln_config.get(submodule_name, {})
-            submodule_config = copy.deepcopy(submodule_config)
-
-            submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), f"{submodule_class_name}")
-            prefix = cls._prefix.get(submodule_name, "")
-            connected_submodules = cls._connected_classes.get(submodule_name)._submodules
-            pipe_global_config = {k: v for k, v in submodule_config.items() if k not in connected_submodules}
-            submodule_config = {k: v for k, v in submodule_config.items() if k in connected_submodules}
-            for key in submodule_config.keys():
-                submodule_config[key].update(pipe_global_config)
-
-            for connected_submodule_name in connected_submodules:
-                connected_submodule_config = rbln_config.pop(prefix + connected_submodule_name, {})
-                if connected_submodule_name in submodule_config:
-                    submodule_config[connected_submodule_name].update(connected_submodule_config)
-                else:
-                    submodule_config[connected_submodule_name] = connected_submodule_config
-
-            pipe_global_config = {
-                k: v for k, v in rbln_config.items() if k != submodule_class_name and not isinstance(v, dict)
-            }
-
-            for connected_submodule_name in connected_submodules:
-                for k, v in pipe_global_config.items():
-                    if "guidance_scale" in k:
-                        if prefix + "guidance_scale" == k:
-                            submodule_config[connected_submodule_name]["guidance_scale"] = v
-                    else:
-                        submodule_config[connected_submodule_name][k] = v
-            rbln_config[submodule_name] = submodule_config
         else:
             raise ValueError(f"submodule {submodule_name} isn't supported")
         return submodule_config
@@ -199,25 +168,8 @@ class RBLNDiffusionMixin:
         else:
             # raise error if any of submodules are torch module.
             model_index_config = cls.load_config(pretrained_model_name_or_path=model_id)
-            if cls._load_connected_pipes:
-                submodules = []
-                for submodule in cls._submodules:
-                    submodule_config = rbln_config.pop(submodule, {})
-                    prefix = cls._prefix.get(submodule, "")
-                    connected_submodules = cls._connected_classes.get(submodule)._submodules
-                    for connected_submodule_name in connected_submodules:
-                        connected_submodule_config = submodule_config.pop(connected_submodule_name, {})
-                        if connected_submodule_config:
-                            rbln_config[prefix + connected_submodule_name] = connected_submodule_config
-                        submodules.append(prefix + connected_submodule_name)
-                pipe_global_config = {k: v for k, v in rbln_config.items() if k not in submodules}
-                for submodule in submodules:
-                    if submodule in rbln_config:
-                        rbln_config[submodule].update(pipe_global_config)
-            else:
-                submodules = cls._submodules
-
-            for submodule_name in submodules:
+            rbln_config = cls._flatten_rbln_config(rbln_config)
+            for submodule_name in cls._submodules:
                 if isinstance(kwargs.get(submodule_name), torch.nn.Module):
                     raise AssertionError(
                         f"{submodule_name} is not compiled torch module. If you want to compile, set `export=True`."
@@ -266,9 +218,89 @@ class RBLNDiffusionMixin:
             lora_scales=lora_scales,
         )
 
-        compiled_submodules = cls._compile_submodules(model, passed_submodules, model_save_dir, rbln_config)
+        if cls._load_connected_pipes:
+            compiled_submodules = cls._compile_pipelines(model, passed_submodules, model_save_dir, rbln_config)
+        else:
+            compiled_submodules = cls._compile_submodules(model, passed_submodules, model_save_dir, rbln_config)
         return cls._construct_pipe(model, compiled_submodules, model_save_dir, rbln_config)
 
+    @classmethod
+    def _prepare_rbln_config(
+        cls,
+        rbln_config,
+    ) -> Dict[str, Any]:
+        prepared_config = {}
+        for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
+            connected_pipe_config = rbln_config.pop(connected_pipe_name, {})
+            prefix = cls._prefix.get(connected_pipe_name, "")
+            guidance_scale = rbln_config.pop(f"{prefix}guidance_scale", None)
+            if "guidance_scale" not in connected_pipe_config and guidance_scale is not None:
+                connected_pipe_config["guidance_scale"] = guidance_scale
+            for submodule_name in connected_pipe_cls._submodules:
+                submodule_config = rbln_config.pop(prefix + submodule_name, {})
+                if submodule_name not in connected_pipe_config:
+                    connected_pipe_config[submodule_name] = {}
+                connected_pipe_config[submodule_name].update(
+                    {k: v for k, v in submodule_config.items() if k not in connected_pipe_config[submodule_name]}
+                )
+            prepared_config[connected_pipe_name] = connected_pipe_config
+        prepared_config.update(rbln_config)
+        return prepared_config
+
+    @classmethod
+    def _flatten_rbln_config(
+        cls,
+        rbln_config,
+    ) -> Dict[str, Any]:
+        prepared_config = cls._prepare_rbln_config(rbln_config)
+        flattened_config = {}
+        pipe_global_config = {k: v for k, v in prepared_config.items() if k not in cls._connected_classes.keys()}
+        for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
+            connected_pipe_config = prepared_config.pop(connected_pipe_name)
+            prefix = cls._prefix.get(connected_pipe_name, "")
+            connected_pipe_global_config = {
+                k: v for k, v in connected_pipe_config.items() if k not in connected_pipe_cls._submodules
+            }
+            for submodule_name in connected_pipe_cls._submodules:
+                flattened_config[prefix + submodule_name] = connected_pipe_config[submodule_name]
+                flattened_config[prefix + submodule_name].update(
+                    {
+                        k: v
+                        for k, v in connected_pipe_global_config.items()
+                        if k not in flattened_config[prefix + submodule_name]
+                    }
+                )
+        flattened_config.update(pipe_global_config)
+        return flattened_config
+
+    @classmethod
+    def _compile_pipelines(
+        cls,
+        model: torch.nn.Module,
+        passed_submodules: Dict[str, RBLNModel],
+        model_save_dir: Optional[PathLike],
+        rbln_config: Dict[str, Any],
+    ) -> Dict[str, RBLNModel]:
+        compiled_submodules = {}
+
+        rbln_config = cls._prepare_rbln_config(rbln_config)
+        pipe_global_config = {k: v for k, v in rbln_config.items() if k not in cls._connected_classes.keys()}
+        for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
+            connected_pipe_submodules = {}
+            prefix = cls._prefix.get(connected_pipe_name, "")
+            for submodule_name in connected_pipe_cls._submodules:
+                connected_pipe_submodules[submodule_name] = passed_submodules.get(prefix + submodule_name, None)
+            connected_pipe = getattr(model, connected_pipe_name)
+            connected_pipe_config = {}
+            connected_pipe_config.update(pipe_global_config)
+            connected_pipe_config.update(rbln_config[connected_pipe_name])
+            connected_pipe_compiled_submodules = connected_pipe_cls._compile_submodules(
+                connected_pipe, connected_pipe_submodules, model_save_dir, connected_pipe_config, prefix
+            )
+            for submodule_name, compiled_submodule in connected_pipe_compiled_submodules.items():
+                compiled_submodules[prefix + submodule_name] = compiled_submodule
+        return compiled_submodules
+
     @classmethod
     def _compile_submodules(
         cls,
@@ -307,41 +339,6 @@ class RBLNDiffusionMixin:
                     model_save_dir=model_save_dir,
                     rbln_config=submodule_rbln_config,
                 )
-            elif hasattr(pipelines, submodule.__class__.__name__):
-                connected_pipe = submodule
-                connected_pipe_model_save_dir = model_save_dir
-                connected_pipe_rbln_config = submodule_rbln_config
-                connected_pipe_cls: RBLNDiffusionMixin = getattr(
-                    importlib.import_module("optimum.rbln"), connected_pipe.__class__.__name__
-                )
-                submodule_dict = {}
-                for name in connected_pipe.config.keys():
-                    if hasattr(connected_pipe, name):
-                        submodule_dict[name] = getattr(connected_pipe, name)
-                connected_pipe = connected_pipe_cls(**submodule_dict)
-                connected_pipe_submodules = {}
-                prefix = cls._prefix.get(submodule_name, "")
-                for name in connected_pipe_cls._submodules:
-                    if prefix + name in passed_submodules:
-                        connected_pipe_submodules[name] = passed_submodules.get(prefix + name)
-
-                connected_pipe_compiled_submodules = connected_pipe_cls._compile_submodules(
-                    model=connected_pipe,
-                    passed_submodules=connected_pipe_submodules,
-                    model_save_dir=model_save_dir,
-                    rbln_config=connected_pipe_rbln_config,
-                    prefix=prefix,
-                )
-                connected_pipe = connected_pipe_cls._construct_pipe(
-                    connected_pipe,
-                    connected_pipe_compiled_submodules,
-                    connected_pipe_model_save_dir,
-                    connected_pipe_rbln_config,
-                )
-
-                for name in connected_pipe_cls._submodules:
-                    compiled_submodules[prefix + name] = getattr(connected_pipe, name)
-                submodule = connected_pipe
             else:
                 raise ValueError(f"Unknown class of submodule({submodule_name}) : {submodule.__class__.__name__} ")
 
@@ -374,23 +371,16 @@ class RBLNDiffusionMixin:
     @classmethod
     def _construct_pipe(cls, model, submodules, model_save_dir, rbln_config):
         # Construct finalize pipe setup with compiled submodules and configurations
-        submodule_names = []
-        for submodule_name in cls._submodules:
-            submodule = getattr(model, submodule_name)
-            if hasattr(pipelines, submodule.__class__.__name__):
-                prefix = cls._prefix.get(submodule_name, "")
-                connected_pipe_submodules = submodules[submodule_name].__class__._submodules
-                connected_pipe_submodules = [prefix + name for name in connected_pipe_submodules]
-                submodule_names += connected_pipe_submodules
-                setattr(model, submodule_name, submodules[submodule_name])
-            else:
-                submodule_names.append(submodule_name)
-
         if model_save_dir is not None:
             # To skip saving original pytorch modules
-            for submodule_name in submodule_names:
+            for submodule_name in cls._submodules:
                 delattr(model, submodule_name)
 
+            if cls._load_connected_pipes:
+                for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
+                    for submodule_name in connected_pipe_cls._submodules:
+                        delattr(getattr(model, connected_pipe_name), submodule_name)
+
             # Direct calling of `save_pretrained` causes config.unet = (None, None).
             # So config must be saved again, later.
             model.save_pretrained(model_save_dir)
@@ -398,10 +388,15 @@ class RBLNDiffusionMixin:
             # Causing warning messeages.
 
         update_dict = {}
-        for submodule_name in submodule_names:
+        for submodule_name in cls._submodules:
             # replace submodule
             setattr(model, submodule_name, submodules[submodule_name])
             update_dict[submodule_name] = ("optimum.rbln", submodules[submodule_name].__class__.__name__)
+        if cls._load_connected_pipes:
+            for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
+                prefix = cls._prefix.get(connected_pipe_name, "")
+                for submodule_name in connected_pipe_cls._submodules:
+                    setattr(getattr(model, connected_pipe_name), submodule_name, submodules[prefix + submodule_name])
 
         # Update config to be able to load from model directory.
         #
@@ -420,16 +415,9 @@ class RBLNDiffusionMixin:
         if rbln_config.get("optimize_host_memory") is False:
             # Keep compiled_model objs to further analysis. -> TODO: remove soon...
             model.compiled_models = []
-            if model._load_connected_pipes:
-                for name in cls._submodules:
-                    connected_pipe = getattr(model, name)
-                    for submodule_name in connected_pipe.__class__._submodules:
-                        submodule = getattr(connected_pipe, submodule_name)
-                        model.compiled_models.extend(submodule.compiled_models)
-            else:
-                for name in cls._submodules:
-                    submodule = getattr(model, name)
-                    model.compiled_models.extend(submodule.compiled_models)
+            for name in cls._submodules:
+                submodule = getattr(model, name)
+                model.compiled_models.extend(submodule.compiled_models)
 
         return model
 

{optimum_rbln-0.7.3a6 → optimum_rbln-0.7.3.post1}/src/optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py

@@ -39,7 +39,7 @@ from .pipeline_kandinsky2_2_prior import RBLNKandinskyV22PriorPipeline
 class RBLNKandinskyV22CombinedPipeline(RBLNDiffusionMixin, KandinskyV22CombinedPipeline):
     original_class = KandinskyV22CombinedPipeline
     _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22Pipeline}
-    _submodules = ["prior_pipe", "decoder_pipe"]
+    _submodules = ["prior_image_encoder", "prior_text_encoder", "prior_prior", "unet", "movq"]
     _prefix = {"prior_pipe": "prior_"}
 
     def __init__(
@@ -90,7 +90,7 @@ class RBLNKandinskyV22CombinedPipeline(RBLNDiffusionMixin, KandinskyV22CombinedP
 class RBLNKandinskyV22Img2ImgCombinedPipeline(RBLNDiffusionMixin, KandinskyV22Img2ImgCombinedPipeline):
     original_class = KandinskyV22Img2ImgCombinedPipeline
     _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22Img2ImgPipeline}
-    _submodules = ["prior_pipe", "decoder_pipe"]
+    _submodules = ["prior_image_encoder", "prior_text_encoder", "prior_prior", "unet", "movq"]
     _prefix = {"prior_pipe": "prior_"}
 
     def __init__(
@@ -141,7 +141,7 @@ class RBLNKandinskyV22Img2ImgCombinedPipeline(RBLNDiffusionMixin, KandinskyV22Im
 class RBLNKandinskyV22InpaintCombinedPipeline(RBLNDiffusionMixin, KandinskyV22InpaintCombinedPipeline):
     original_class = KandinskyV22InpaintCombinedPipeline
     _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22InpaintPipeline}
-    _submodules = ["prior_pipe", "decoder_pipe"]
+    _submodules = ["prior_image_encoder", "prior_text_encoder", "prior_prior", "unet", "movq"]
     _prefix = {"prior_pipe": "prior_"}
 
     def __init__(

{optimum_rbln-0.7.3a6 → optimum_rbln-0.7.3.post1}/src/optimum/rbln/transformers/models/bart/modeling_bart.py

@@ -108,6 +108,8 @@ class RBLNBartModel(RBLNModel):
 
 
 class RBLNBartForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    support_paged_causal_attn = True
+
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
         enc_max_seq_len = (

{optimum_rbln-0.7.3a6 → optimum_rbln-0.7.3.post1}/src/optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py

@@ -98,6 +98,18 @@ def validate_attention_method(
                 "this requirement, or consider switching `rbln_attn_impl` to 'eager' for shorter lengths."
             )
 
+    if rbln_kvcache_block_size is not None:
+        if rbln_attn_impl == "flash_attn" and rbln_kvcache_partition_len != rbln_kvcache_block_size:
+            raise ValueError(
+                f" When using 'flash attention', the `rbln_kvcache_block_size` ({rbln_kvcache_block_size})  "
+                f"must always be set equal to the `rbln_kvcache_partition_len` {rbln_kvcache_partition_len}."
+            )
+        elif rbln_attn_impl == "eager" and rbln_kvcache_block_size != rbln_max_seq_len:
+            raise ValueError(
+                f" When using 'eager attention', the `rbln_kvcache_block_size` ({rbln_kvcache_block_size})  "
+                f"must always be set equal to the `rbln_max_seq_len` {rbln_max_seq_len}."
+            )
+
     return rbln_attn_impl, rbln_kvcache_partition_len, rbln_kvcache_block_size
 
 

{optimum_rbln-0.7.3a6 → optimum_rbln-0.7.3.post1}/src/optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py

@@ -50,6 +50,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         runtime: rebel.Runtime,
         batch_size: int,
         dec_max_seq_len: int,
+        support_paged_causal_attn: Optional[bool] = None,
         use_attention_mask: Optional[bool] = None,
         **kwargs: Any,
     ) -> None:
@@ -57,7 +58,10 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         self.batch_size = batch_size
         self.dec_max_seq_len = dec_max_seq_len
         self.use_attention_mask = use_attention_mask
-        self.default_block_tables = torch.arange(0, self.batch_size, dtype=torch.int16).view(self.batch_size, 1)
+        if support_paged_causal_attn:
+            self.default_block_tables = torch.arange(0, self.batch_size, dtype=torch.int16).view(self.batch_size, 1)
+        else:
+            self.default_block_tables = None
 
     def forward(
         self,
@@ -94,7 +98,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
             decoder_attention_mask if self.use_attention_mask else None,
             attention_mask,
             cache_position,
-            block_tables,
+            block_tables=block_tables,
         )
 
         return Seq2SeqLMOutput(logits=lm_logits)
@@ -115,6 +119,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
 
     main_input_name = "input_ids"
     auto_model_class = AutoModelForSeq2SeqLM
+    support_paged_causal_attn = None
 
     def __post_init__(self, **kwargs):
         batch_size = self.rbln_config.model_cfg["batch_size"]
@@ -130,6 +135,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             main_input_name="input_ids",
             batch_size=batch_size,
             dec_max_seq_len=dec_max_seq_len,
+            support_paged_causal_attn=self.support_paged_causal_attn,
             use_attention_mask=self.use_attention_mask,
         )
 
@@ -186,13 +192,16 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         rbln_dec_max_seq_len = rbln_kwargs.get("dec_max_seq_len", None)
         rbln_batch_size = rbln_kwargs.get("batch_size", None)
         rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
-        rbln_use_attention_mask = rbln_kwargs.get("use_attention_mask", None)
 
-        if rbln_use_attention_mask is None:
-            rbln_use_attention_mask = False
-            rbln_npu = rbln_kwargs.get("npu", None) or rebel.get_npu_name()
-            if rbln_npu == "RBLN-CA02":
-                rbln_use_attention_mask = True
+        if cls.support_paged_causal_attn:
+            rbln_use_attention_mask = rbln_kwargs.get("use_attention_mask", None)
+            if rbln_use_attention_mask is None:
+                rbln_use_attention_mask = False
+                rbln_npu = rbln_kwargs.get("npu", None) or rebel.get_npu_name()
+                if rbln_npu == "RBLN-CA02":
+                    rbln_use_attention_mask = True
+        else:
+            rbln_use_attention_mask = True
 
         n_layer = getattr(model_config, "decoder_layers", None) or getattr(model_config, "num_layers")
         n_head = getattr(model_config, "decoder_attention_heads", None) or getattr(model_config, "num_heads")
@@ -265,11 +274,6 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
                 [rbln_batch_size, 1],
                 "int32",
             ),
-            (
-                "block_tables",
-                [rbln_batch_size, 1],
-                "int16",
-            ),
         ]
         dec_input_info.extend(
             [
@@ -302,6 +306,8 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             ]
         )
 
+        if cls.support_paged_causal_attn:
+            dec_input_info.insert(3, ("block_tables", [rbln_batch_size, 1], "int16"))
         if rbln_use_attention_mask:
             dec_input_info.insert(1, ("attention_mask", [rbln_batch_size, rbln_dec_max_seq_len], "float32"))