optimum-rbln 0.8.2a4__py3-none-any.whl → 0.9.3rc0__py3-none-any.whl

optimum/rbln/transformers/models/depth_anything/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .configuration_depth_anything import RBLNDepthAnythingForDepthEstimationConfig
+from .modeling_depth_anything import RBLNDepthAnythingForDepthEstimation

optimum/rbln/transformers/models/depth_anything/configuration_depth_anything.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.
+
+from ...configuration_generic import RBLNModelForDepthEstimationConfig
+
+
+class RBLNDepthAnythingForDepthEstimationConfig(RBLNModelForDepthEstimationConfig):
+    """
+    Configuration class for DepthAnythingForDepthEstimation.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Depth Anything V2 Small models for depth estimation tasks.
+    """

optimum/rbln/transformers/models/depth_anything/modeling_depth_anything.py

@@ -0,0 +1,25 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...modeling_generic import RBLNModelForDepthEstimation
+
+
+class RBLNDepthAnythingForDepthEstimation(RBLNModelForDepthEstimation):
+    """
+    RBLN optimized DepthAnythingForDepthEstimation model for depth estimation tasks.
+
+    This class provides hardware-accelerated inference for Depth Anything V2
+    models on RBLN devices, providing the most capable monocular depth estimation (MDE) model.
+    """

optimum/rbln/transformers/models/exaone/modeling_exaone.py

@@ -14,11 +14,13 @@
 
 
 import inspect
-from typing import Any, Callable
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Union
 
 from transformers import AutoModelForCausalLM
 from transformers.generation.utils import GenerationMixin
 
+from ....configuration_utils import RBLNModelConfig
 from ....utils import logging
 from ..decoderonly import RBLNDecoderOnlyModelForCausalLM
 from .exaone_architecture import ExaoneForCausalLMWrapper
@@ -92,9 +94,45 @@ class RBLNExaoneForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     _supports_cache_class = True
 
     @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        kwargs.setdefault("trust_remote_code", True)
-        return super().from_pretrained(*args, **kwargs)
+    def from_pretrained(
+        cls,
+        model_id: Union[str, Path],
+        *,
+        export: Optional[bool] = None,
+        rbln_config: Optional[Union[Dict, RBLNModelConfig]] = None,
+        trust_remote_code: Optional[bool] = None,
+        **kwargs: Any,
+    ):
+        """
+        The `from_pretrained()` function is utilized in its standard form as in the HuggingFace transformers library.
+        User can use this function to load a pre-trained model from the HuggingFace library and convert it to a RBLN model to be run on RBLN NPUs.
+
+        Args:
+            model_id (Union[str, Path]): The model id of the pre-trained model to be loaded.
+                It can be downloaded from the HuggingFace model hub or a local path, or a model id of a compiled model using the RBLN Compiler.
+            export (Optional[bool]): A boolean flag to indicate whether the model should be compiled.
+                If None, it will be determined based on the existence of the compiled model files in the model_id.
+            rbln_config (Optional[Union[Dict, RBLNModelConfig]]): Configuration for RBLN model compilation and runtime.
+                This can be provided as a dictionary or an instance of the model's configuration class (e.g., `RBLNExaoneForCausalLMConfig` for EXAONE models).
+                For detailed configuration options, see the specific model's configuration class documentation.
+            trust_remote_code (bool): Whether or not to trust the remote code when loading a model from the Hub.
+            kwargs: Additional keyword arguments. Arguments with the prefix `rbln_` are passed to rbln_config, while the remaining arguments are passed to the HuggingFace library.
+
+        Returns:
+            (RBLNModel): A RBLN model instance ready for inference on RBLN NPU devices.
+        """
+
+        if trust_remote_code is not None:
+            kwargs["trust_remote_code"] = trust_remote_code
+        elif "trust_remote_code" not in kwargs:
+            kwargs["trust_remote_code"] = True
+
+        return super().from_pretrained(
+            model_id=model_id,
+            export=export,
+            rbln_config=rbln_config,
+            **kwargs,
+        )
 
     def __getattr__(self, __name: str) -> Any:
         def redirect(func):

optimum/rbln/transformers/models/gemma/__init__.py

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_gemma import RBLNGemmaForCausalLMConfig
-from .modeling_gemma import RBLNGemmaForCausalLM
+from .configuration_gemma import RBLNGemmaForCausalLMConfig, RBLNGemmaModelConfig
+from .modeling_gemma import RBLNGemmaForCausalLM, RBLNGemmaModel

optimum/rbln/transformers/models/gemma/configuration_gemma.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNGemmaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -40,3 +40,11 @@ class RBLNGemmaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     )
     ```
     """
+
+
+class RBLNGemmaModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN Gemma models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+    """

optimum/rbln/transformers/models/gemma/gemma_architecture.py

@@ -13,10 +13,7 @@
 # limitations under the License.
 
 
-from ...models.decoderonly.decoderonly_architecture import (
-    DecoderOnlyModel,
-    DecoderOnlyWrapper,
-)
+from ...models.decoderonly.decoderonly_architecture import DecoderOnlyModel, DecoderOnlyWrapper
 
 
 class GemmaWrapper(DecoderOnlyWrapper):

optimum/rbln/transformers/models/gemma/modeling_gemma.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .gemma_architecture import GemmaWrapper
 
 
@@ -81,3 +81,24 @@ class RBLNGemmaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
 
     _decoder_wrapper_cls = GemmaWrapper
+
+
+class RBLNGemmaModel(RBLNDecoderOnlyModel):
+    """
+    The Gemma Model transformer without a language modeling head.
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based GemmaModel model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers GemmaModel model into a RBLN transformer model by:
+
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    **Configuration:**
+    This model uses [`RBLNGemmaModelConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNGemmaModelConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNGemmaModelConfig`] class for all available configuration options.
+    """
+
+    _decoder_wrapper_cls = GemmaWrapper

optimum/rbln/transformers/models/gemma3/configuration_gemma3.py

@@ -11,13 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, Optional
-
-import rebel
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
 from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
-from ..siglip.configuration_siglip import RBLNSiglipVisionModelConfig
+
+
+logger = get_logger(__name__)
 
 
 class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -25,23 +26,45 @@ class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
         self,
         use_position_ids: Optional[bool] = None,
         use_attention_mask: Optional[bool] = None,
+        prefill_chunk_size: Optional[int] = None,
         image_prefill_chunk_size: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
+        """
+        Args:
+            use_position_ids (Optional[bool]): Whether or not to use `position_ids`, which is indices of positions of each input sequence tokens in the position embeddings.
+            use_attention_mask (Optional[bool]): Whether or not to use `attention_mask` to to avoid performing attention on padding token indices.
+            prefill_chunk_size (Optional[int]): The chunk size used during the prefill phase for
+                processing input sequences. Defaults to 256. Must be a positive integer
+                divisible by 64. Affects prefill performance and memory usage.
+            image_prefill_chunk_size (Optional[int]): The chunk size used during the prefill phase for
+                processing images. This config is used when `use_image_prefill` is True.
+                Currently, the `prefill_chunk_size` and `image_prefill_chunk_size` should be the same value.
+            kwargs: Additional arguments passed to the parent `RBLNDecoderOnlyModelForCausalLMConfig`.
+
+        Raises:
+            ValueError: If `use_attention_mask` or `use_position_ids` are False.
+        """
         # use_attention_mask and use_position_ids are always True for Gemma3
         use_attention_mask = use_attention_mask or True
         use_position_ids = use_position_ids or True
+        prefill_chunk_size = prefill_chunk_size or 256
 
         super().__init__(
+            prefill_chunk_size=prefill_chunk_size,
             use_attention_mask=use_attention_mask,
             use_position_ids=use_position_ids,
             **kwargs,
         )
         self.image_prefill_chunk_size = image_prefill_chunk_size
 
-        npu = self.npu or rebel.get_npu_name()
-        if npu == "RBLN-CA02":
-            raise NotImplementedError("Gemma3 is currently not supported on RBLN-CA02")
+    @property
+    def use_image_prefill(self):
+        return self.image_prefill_chunk_size is not None
+
+    @property
+    def decoder_runtime_idx(self):
+        return 2 if self.use_image_prefill else 1
 
 
 class RBLNGemma3ForConditionalGenerationConfig(RBLNModelConfig):
@@ -52,22 +75,35 @@ class RBLNGemma3ForConditionalGenerationConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         vision_tower: Optional[RBLNModelConfig] = None,
         language_model: Optional[RBLNModelConfig] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
             batch_size (Optional[int]): The batch size for inference. Defaults to 1.
             vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
             language_model (Optional[RBLNModelConfig]): Configuration for the language model component.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
-            ValueError: If batch_size is not a positive integer.
+            ValueError: If `batch_size` is not a positive integer.
         """
         super().__init__(**kwargs)
         self.batch_size = batch_size or 1
         if not isinstance(self.batch_size, int) or self.batch_size < 0:
             raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
 
-        self.vision_tower = self.init_submodule_config(RBLNSiglipVisionModelConfig, vision_tower)
-        self.language_model = self.init_submodule_config(RBLNGemma3ForCausalLMConfig, language_model)
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for Gemma3 vision tower. It will be set to 1.")
+
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower, batch_size=1, force_kwargs=True
+        )
+        self.language_model = self.initialize_submodule_config(submodule_config=language_model)
+
+    @property
+    def image_prefill_chunk_size(self):
+        return self.language_model.image_prefill_chunk_size
+
+    @property
+    def prefill_chunk_size(self):
+        return self.language_model.prefill_chunk_size

optimum/rbln/transformers/models/gemma3/gemma3_architecture.py

@@ -63,6 +63,7 @@ class Gemma3TextModel(DecoderOnlyModel):
         rotary_emb: torch.nn.Module = None,
         global_block_tables: Optional[torch.Tensor] = None,
         local_block_tables: Optional[torch.Tensor] = None,
+        lora_int_id: Optional[torch.Tensor] = None,
     ):
         # retrieve input_ids and inputs_embeds
         if (input_ids is None) ^ (inputs_embeds is not None):
@@ -105,6 +106,7 @@ class Gemma3TextModel(DecoderOnlyModel):
                 cos=cos_local if is_sliding else cos_global,
                 sin=sin_local if is_sliding else sin_global,
                 block_tables=local_block_tables if is_sliding else global_block_tables,
+                lora_int_id=lora_int_id,
             )
 
         hidden_states = self.get_last_layernorm()(hidden_states)
@@ -127,12 +129,20 @@ class Gemma3DecoderLayer(DecoderOnlyLayer):
         cos: Optional[torch.Tensor] = None,
         sin: Optional[torch.Tensor] = None,
         block_tables: Optional[torch.Tensor] = None,
+        lora_int_id: Optional[torch.Tensor] = None,
     ):
         residual = hidden_states
         hidden_states = self.get_pre_attention_layernorm()(hidden_states)
 
         hidden_states = self.self_attn(
-            hidden_states, attention_mask, seq_positions, past_key_values, cos, sin, block_tables
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            seq_positions=seq_positions,
+            past_key_values=past_key_values,
+            cos=cos,
+            sin=sin,
+            block_tables=block_tables,
+            lora_int_id=lora_int_id,
         )
         hidden_states = self.get_post_attention_layernorm()(hidden_states)
         hidden_states = residual + hidden_states
@@ -140,7 +150,7 @@ class Gemma3DecoderLayer(DecoderOnlyLayer):
         # Fully Connected
         residual = hidden_states
         hidden_states = self.get_pre_feedforward_layernorm()(hidden_states)
-        hidden_states = self._original_mod.mlp(hidden_states)
+        hidden_states = self.forward_mlp(hidden_states, lora_int_id)
         hidden_states = self.get_post_feedforward_layernorm()(hidden_states)
         hidden_states = residual + hidden_states
 

optimum/rbln/transformers/models/gemma3/gemma3_runtime_utils.py

@@ -0,0 +1,245 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Optional
+
+import rebel
+import torch
+
+from ...modeling_outputs import RBLNDecoderOnlyOutput, RBLNGemma3ForCausalLMOutput
+from ..decoderonly.decoderonly_runtime_utils import RBLNPytorchRuntime
+from ..decoderonly.modeling_decoderonly import RBLNRuntimeModel
+
+
+class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
+    def __init__(self, *args, image_prefill: Optional[rebel.Runtime] = None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.image_prefill = RBLNPytorchRuntime(image_prefill)  # FIXME(taehoon)
+        self.prefill = RBLNPytorchRuntime(self.runtime) if self.phase == "prefill" else None  # FIXME
+        self.decode = RBLNPytorchRuntime(self.runtime) if self.phase == "decode" else None
+
+    def _prepare_prefill_inputs(self, *args, **kwargs):
+        (
+            inputs,
+            cache_position,
+            chunked_attention_mask,
+            position_ids,
+            position_embed,
+            padded_cache_lengths,
+            query_length,
+            token_type_ids,
+        ) = super()._prepare_prefill_inputs(*args, **kwargs)
+
+        # chunked_attention_mask shape
+        chunked_attention_mask = torch.zeros(1, chunked_attention_mask.shape[-1], dtype=torch.float32)
+
+        # In case of Gemma3ForConditionalGeneration, the loop counter may not be a prefill_chunk_size,
+        # so we cannot guarantee that the last chunk starts at a position that is a multiple of prefill_chunk_size.
+        if self.rbln_config.use_image_prefill:
+            padding_size = self.rbln_config.image_prefill_chunk_size
+            inputs = torch.nn.functional.pad(inputs, (0, 0, 0, padding_size))
+            cache_position = torch.nn.functional.pad(cache_position, (0, padding_size))
+            position_ids = torch.nn.functional.pad(position_ids, (0, padding_size))
+            token_type_ids = torch.nn.functional.pad(token_type_ids, (0, padding_size), value=-1)
+
+        return (
+            inputs,
+            cache_position,
+            chunked_attention_mask,
+            position_ids,
+            position_embed,
+            padded_cache_lengths,
+            query_length,
+            token_type_ids,
+        )
+
+    def prefill_forward(
+        self,
+        inputs: torch.Tensor,
+        cache_position: torch.Tensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        batch_idx: int = None,
+        block_tables: torch.Tensor = None,
+        is_external_block_tables: bool = None,
+        position_embed: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        local_block_tables: Optional[torch.Tensor] = None,
+        lora_int_ids: Optional[torch.Tensor] = None,
+    ) -> torch.FloatTensor:
+        """
+        Performs chunked prefill for efficient KV-cache updates and memory optimization.
+        Instead of processing the entire sequence at once, the input is divided into chunks of size `prefill_chunk_size`,
+        and each chunk is processed sequentially. This allows for better memory utilization and compatibility with continuous batching.
+        """
+        if self.rbln_config.use_lora and lora_int_ids is None:
+            if self.lora_int_ids is None:
+                raise ValueError(
+                    "lora_int_id is required when using LoRA. "
+                    "You should call set_lora_int_ids() before forward() or pass lora_int_id to forward()."
+                )
+            if batch_idx is not None:
+                lora_int_ids = self.lora_int_ids[batch_idx : batch_idx + 1].clone()
+            else:
+                lora_int_ids = self.lora_int_ids.clone()
+
+        (
+            inputs,
+            cache_position,
+            chunked_attention_mask,
+            position_ids,
+            position_embed,
+            padded_cache_lengths,
+            query_length,
+            token_type_ids,
+        ) = self._prepare_prefill_inputs(
+            inputs, cache_position, attention_mask, position_embed, token_type_ids=token_type_ids
+        )
+
+        step = 0
+        while step < query_length:
+            if self.rbln_config.use_image_prefill:
+                # Check if the prefill chunk is an image prefill
+                is_image_prefill = torch.all(
+                    token_type_ids[:, step : step + self.rbln_config.image_prefill_chunk_size] == 1
+                )
+                # Check if the prefill chunk is a text prefill which have image_tokens in it.
+                is_text_prefill_with_image_tokens = not is_image_prefill and torch.any(
+                    token_type_ids[:, step : step + self.rbln_config.prefill_chunk_size] == 1
+                )
+            else:
+                is_image_prefill, is_text_prefill_with_image_tokens = False, False
+
+            # Check if the prefill chunk is the last chunk
+            is_last_chunk = step + self.rbln_config.prefill_chunk_size >= query_length
+
+            input_chunk = inputs[:, step : step + self.rbln_config.prefill_chunk_size]
+            cache_pos_chunk = (
+                cache_position[:, step : step + self.rbln_config.prefill_chunk_size] + padded_cache_lengths
+            )
+            position_ids_chunk = position_ids[:, step : step + self.rbln_config.prefill_chunk_size]
+
+            # if text_prefill end with image_tokens, we only treat the text part.
+            num_processed_tokens = self.rbln_config.prefill_chunk_size
+            current_padded_cache_lengths = 0
+            if is_text_prefill_with_image_tokens:
+                first_image_token_idx = torch.where(
+                    token_type_ids[:, step : step + self.rbln_config.prefill_chunk_size] == 1
+                )[1][0]
+                num_processed_tokens = first_image_token_idx.item()
+                current_padded_cache_lengths = self.rbln_config.prefill_chunk_size - num_processed_tokens
+            if is_last_chunk:
+                num_processed_tokens = query_length - step
+
+            chunked_attention_mask[
+                :, step + padded_cache_lengths : step + num_processed_tokens + padded_cache_lengths
+            ] = 1
+            query_position = torch.tensor(num_processed_tokens - 1, dtype=torch.int16)
+
+            if is_image_prefill:
+                logits = self.image_prefill(
+                    input_chunk,
+                    cache_pos_chunk,
+                    block_tables,
+                    local_block_tables,
+                    query_position,
+                    chunked_attention_mask,
+                    position_ids_chunk,
+                    lora_int_ids if self.rbln_config.use_lora else None,
+                )
+            else:
+                logits = self.prefill(
+                    input_chunk,
+                    cache_pos_chunk,
+                    block_tables,
+                    local_block_tables,
+                    query_position,
+                    chunked_attention_mask,
+                    position_ids_chunk,
+                    lora_int_ids if self.rbln_config.use_lora else None,
+                )
+
+            padded_cache_lengths += current_padded_cache_lengths
+            step += num_processed_tokens
+
+        if not is_external_block_tables:
+            self.dec_attn_mask[batch_idx : batch_idx + 1] = chunked_attention_mask
+
+        return RBLNGemma3ForCausalLMOutput(
+            logits=logits, padded_cache_lengths=padded_cache_lengths, attention_mask=chunked_attention_mask
+        )
+
+    def decode_forward(
+        self,
+        inputs: torch.Tensor,
+        cache_position: torch.Tensor = None,
+        block_tables: torch.Tensor = None,
+        is_external_block_tables: bool = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_embed: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        local_block_tables: Optional[torch.Tensor] = None,
+        lora_int_ids: Optional[torch.Tensor] = None,
+    ) -> torch.FloatTensor:
+        if self.rbln_config.use_lora and lora_int_ids is None:
+            if self.lora_int_ids is None:
+                raise ValueError(
+                    "lora_int_id is required when using LoRA. "
+                    "You should call set_lora_int_ids() before forward() or pass lora_int_id to forward()."
+                )
+
+            lora_int_ids = self.lora_int_ids
+
+        if lora_int_ids is not None and lora_int_ids.shape[0] != self.batch_size:
+            raise ValueError(f"lora_int_ids size mismatch: got {lora_int_ids.shape[0]}, expected {self.batch_size}.")
+
+        batch_size = inputs.shape[0]
+        if batch_size != self.batch_size:
+            raise RuntimeError(
+                f"Batch size mismatch: got {batch_size}, expected {self.batch_size} (compiled batch size)."
+            )
+
+        if batch_size != cache_position.shape[0]:
+            raise RuntimeError(f"Cache position size mismatch: got {cache_position.shape[0]}, expected {batch_size}.")
+
+        # FIXME(taehoon): how to handle pos_attn_mask with external block tables
+        if is_external_block_tables:
+            if attention_mask is None:
+                raise ValueError("attention_mask should be provided with external block tables.")
+            if local_block_tables is None:
+                raise ValueError("local_block_tables should be provided with external block tables.")
+        else:
+            local_block_tables = (
+                local_block_tables
+                if local_block_tables is not None
+                else torch.arange(0, self.batch_size, dtype=torch.int16).view(self.batch_size, -1)
+            )
+            if self.rbln_config.use_attention_mask and attention_mask is None:
+                for b_idx in range(batch_size):
+                    decoding_step = cache_position[b_idx].item()
+                    if not (0 <= decoding_step < self.dec_attn_mask.shape[-1]):
+                        raise ValueError(
+                            f"Decoding step {decoding_step} out of bounds for attention mask with shape {self.dec_attn_mask.shape}."
+                        )
+                    self.dec_attn_mask[b_idx, decoding_step] = 1
+
+                attention_mask = self.dec_attn_mask
+
+        if self.batch_size < block_tables.shape[0]:
+            block_tables = block_tables[: self.batch_size]
+
+        if attention_mask is not None and self.batch_size < attention_mask.shape[0]:
+            attention_mask = attention_mask[: self.batch_size]
+
+        logits = self.decode(inputs, cache_position, block_tables, local_block_tables, attention_mask, position_ids)
+
+        return RBLNDecoderOnlyOutput(logits=logits)