optimum-rbln 0.9.3.post1__py3-none-any.whl

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

@@ -0,0 +1,104 @@
+# 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 ....utils import logging
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
+from .gemma_architecture import GemmaWrapper
+
+
+logger = logging.get_logger(__name__)
+
+
+class RBLNGemmaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The Gemma Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. 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 GemmaForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers GemmaForCausalLM 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 [`RBLNGemmaForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNGemmaForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNGemmaForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNGemmaForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNGemmaForCausalLM.from_pretrained(
+            "google/gemma-7b",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNGemmaForCausalLM.from_pretrained(
+            "google/gemma-7b",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+
+        # Using a RBLNGemmaForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNGemmaForCausalLMConfig
+
+        config = RBLNGemmaForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNGemmaForCausalLM.from_pretrained(
+            "google/gemma-7b",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    _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/__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_gemma3 import RBLNGemma3ForCausalLMConfig, RBLNGemma3ForConditionalGenerationConfig
+from .modeling_gemma3 import RBLNGemma3ForCausalLM, RBLNGemma3ForConditionalGeneration

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

@@ -0,0 +1,109 @@
+# 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 Any, Optional
+
+from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+
+
+logger = get_logger(__name__)
+
+
+class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    def __init__(
+        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: 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
+
+    @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):
+    submodules = ["vision_tower", "language_model"]
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        language_model: Optional[RBLNModelConfig] = None,
+        **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.
+
+        Raises:
+            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}")
+
+        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

@@ -0,0 +1,170 @@
+# 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 copy
+from typing import Optional, Tuple, Union
+
+import torch
+from transformers.models.gemma3.modeling_gemma3 import Gemma3RMSNorm
+
+from ..decoderonly.decoderonly_architecture import (
+    DecoderOnlyAttention,
+    DecoderOnlyLayer,
+    DecoderOnlyModel,
+    DecoderOnlyWrapper,
+    RotaryEmbedding,
+    slice_and_unsqueeze_cos_sin,
+)
+
+
+class Gemma3ForCausalLMWrapper(DecoderOnlyWrapper):
+    def get_rotary_emb(self, max_seq_len):
+        rotary_emb_global = RotaryEmbedding(config=self.config, max_seq_len_cached=max_seq_len)
+
+        config = copy.deepcopy(self.config)
+        config.rope_theta = config.rope_local_base_freq
+        config.rope_scaling = {"rope_type": "default"}
+        rotary_emb_local = RotaryEmbedding(config=config, max_seq_len_cached=max_seq_len)
+
+        return (rotary_emb_global, rotary_emb_local)
+
+    def get_rbln_attn_class(self):
+        return Gemma3Attention
+
+    def get_rbln_layer_class(self):
+        return Gemma3DecoderLayer
+
+    def get_rbln_model_class(self):
+        return Gemma3TextModel
+
+
+class Gemma3TextModel(DecoderOnlyModel):
+    # Different from DecoderOnlyModel, this model has global and local rotary embeddings.
+    def forward(
+        self,
+        input_ids: torch.Tensor = None,
+        inputs_embeds: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        cache_position: torch.Tensor = None,
+        position_ids: torch.Tensor = None,
+        query_position: torch.Tensor = None,
+        past_key_values: Tuple[Tuple[torch.Tensor]] = None,
+        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):
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+            )
+
+        # embed positions
+        if inputs_embeds is None:
+            inputs_embeds = self.get_embedding()(input_ids)
+
+        hidden_states = inputs_embeds
+
+        # Global Position Embeddings
+        cos_global, sin_global = rotary_emb[0](hidden_states, self.max_seq_len)
+        cos_global, sin_global = slice_and_unsqueeze_cos_sin(cos_global, sin_global, position_ids)
+
+        # Local Position Embeddings
+        cos_local, sin_local = rotary_emb[1](hidden_states, self.max_seq_len)
+        cos_local, sin_local = slice_and_unsqueeze_cos_sin(cos_local, sin_local, position_ids)
+
+        # (batch, seq_len) -> (batch,)
+        if self.attn_impl == "flash_attn":
+            seq_positions = cache_position[:, 0]
+            seq_positions = self.convert_sequence_positions_for_flash_attn(
+                seq_positions=seq_positions, max_seq_len=self.max_seq_len
+            )
+        else:
+            seq_positions = cache_position[:, :1]
+
+        sliding_cache_pos = self.get_local_cache_positions(position_ids, query_position)
+
+        for layer_idx, layer in enumerate(self.layers):
+            is_sliding = True if layer_idx in self.sliding_window_layers else False
+            hidden_states = layer(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                seq_positions=sliding_cache_pos if is_sliding else seq_positions,
+                past_key_values=past_key_values,
+                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)
+        return hidden_states
+
+
+class Gemma3DecoderLayer(DecoderOnlyLayer):
+    def get_pre_feedforward_layernorm(self) -> Gemma3RMSNorm:
+        return self._original_mod.pre_feedforward_layernorm
+
+    def get_post_feedforward_layernorm(self) -> Gemma3RMSNorm:
+        return self._original_mod.post_feedforward_layernorm
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        seq_positions: Union[torch.LongTensor, Tuple[torch.LongTensor]],
+        past_key_values: Tuple[Tuple[torch.Tensor]],
+        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=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
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.get_pre_feedforward_layernorm()(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
+
+        return hidden_states
+
+
+class Gemma3Attention(DecoderOnlyAttention):
+    def __post_init__(self):
+        self.q_proj = self._original_mod.q_proj
+        self.k_proj = self._original_mod.k_proj
+        self.v_proj = self._original_mod.v_proj
+        self.o_proj = self._original_mod.o_proj
+        self.q_norm = self._original_mod.q_norm
+        self.k_norm = self._original_mod.k_norm
+
+    def get_attn_scale(self):
+        return self._original_mod.config.query_pre_attn_scalar**-0.5

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)