optimum-rbln 0.8.2rc0__py3-none-any.whl → 0.8.3a0__py3-none-any.whl

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

@@ -90,6 +90,15 @@ class RBLNGemmaModel(RBLNDecoderOnlyModel):
 
     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,7 +11,7 @@
 # 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
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
 from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
@@ -25,7 +25,7 @@ class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
         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,
     ):
         # use_attention_mask and use_position_ids are always True for Gemma3
         use_attention_mask = use_attention_mask or True
@@ -57,7 +57,7 @@ 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:

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

@@ -0,0 +1,217 @@
+# 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.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 = image_prefill  # FIXME(taehoon)
+        self.prefill = self.runtime if self.phase == "prefill" else None  # FIXME
+        self.decode = 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,
+    ) -> 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.
+        """
+        (
+            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,
+                )
+            else:
+                logits = self.prefill(
+                    input_chunk,
+                    cache_pos_chunk,
+                    block_tables,
+                    local_block_tables,
+                    query_position,
+                    chunked_attention_mask,
+                    position_ids_chunk,
+                )
+
+            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,
+    ) -> torch.FloatTensor:
+        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)

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

@@ -12,43 +12,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import inspect
-from collections import deque
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
 
 import rebel
 import torch
 from rebel.compile_context import CompileContext
-from transformers import (
-    AutoModelForImageTextToText,
-    Gemma3ForConditionalGeneration,
-    PretrainedConfig,
-    PreTrainedModel,
-)
+from transformers import AutoModelForImageTextToText, Gemma3ForConditionalGeneration, PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.modeling_utils import no_init_weights
 from transformers.models.gemma3.modeling_gemma3 import Gemma3TextScaledWordEmbedding
 
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
+from ...modeling_outputs import RBLNDecoderOnlyOutput
+from ..decoderonly.decoderonly_runtime_utils import RBLNPageTableManager
 from ..decoderonly.modeling_decoderonly import (
-    RBLNDecoderOnlyForCausalLMOutput,
     RBLNDecoderOnlyModelForCausalLM,
-    RBLNRuntimeModel,
 )
 from .configuration_gemma3 import RBLNGemma3ForCausalLMConfig
 from .gemma3_architecture import Gemma3ForCausalLMWrapper
+from .gemma3_runtime_utils import RBLNGemma3RuntimeModel
 
 
 if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, Gemma3ForConditionalGeneration
 
 
-@dataclass
-class RBLNGemma3ForCausalLMOutput(RBLNDecoderOnlyForCausalLMOutput):
-    attention_mask: Optional[torch.Tensor] = None
-
-
 class LoopVisionTower:
     def __init__(self, vision_tower: RBLNModel) -> None:
         self.vision_tower = vision_tower
@@ -201,7 +190,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
 
     def _update_model_kwargs_for_generation(
         self,
-        outputs: RBLNDecoderOnlyForCausalLMOutput,
+        outputs: RBLNDecoderOnlyOutput,
         model_kwargs: Dict[str, Any],
         **kwargs,
     ) -> Dict[str, Any]:
@@ -298,7 +287,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
         padded_cache_lengths: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
         **lm_kwargs: Dict[str, Any],
-    ) -> Union[Tuple, RBLNDecoderOnlyForCausalLMOutput]:
+    ) -> Union[Tuple, RBLNDecoderOnlyOutput]:
         # prefill
         if cache_position is None:
             logits = []
@@ -339,213 +328,11 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
                 position_ids=position_ids if self.rbln_config.language_model.use_position_ids else None,
             ).logits
 
-        return RBLNDecoderOnlyForCausalLMOutput(
+        return RBLNDecoderOnlyOutput(
             logits=logits, generate_idx=generate_idx, padded_cache_lengths=padded_cache_lengths
         )
 
 
-class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
-    def __init__(self, *args, image_prefill: Optional[rebel.Runtime] = None, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.image_prefill = image_prefill  # FIXME(taehoon)
-        self.prefill = self.runtime if self.phase == "prefill" else None  # FIXME
-        self.decode = self.runtime if self.phase == "decode" else None
-
-    def _prepare_prefill_inputs(self, *args, **kwargs):
-        (
-            inputs,
-            cache_position,
-            chunked_attention_mask,
-            out_buffers,
-            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,
-            out_buffers,
-            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,
-    ) -> 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.
-        """
-        (
-            inputs,
-            cache_position,
-            chunked_attention_mask,
-            out_buffers,
-            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,
-                    out=out_buffers,
-                )
-            else:
-                logits = self.prefill(
-                    input_chunk,
-                    cache_pos_chunk,
-                    block_tables,
-                    local_block_tables,
-                    query_position,
-                    chunked_attention_mask,
-                    position_ids_chunk,
-                    out=out_buffers,
-                )
-
-            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,
-    ) -> torch.FloatTensor:
-        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 RBLNDecoderOnlyForCausalLMOutput(logits=logits)
-
-
 class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
     The Gemma3 Model transformer with a language modeling head (linear layer) on top.
@@ -559,52 +346,34 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     _decoder_wrapper_cls = Gemma3ForCausalLMWrapper
 
-    def __post_init__(self, **kwargs):
-        main_input_name = self.main_input_name
-
-        if self.rbln_config.use_inputs_embeds:
-            main_input_name = "inputs_embeds"
-            artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
-            self.embed_tokens = self._create_embedding_layer()
-            self.embed_tokens.load_state_dict(artifacts["embed_tokens"])
-        else:
-            self.embed_tokens = None
-
+    def setup_runtime(self):
         # Initialize shared resources to be used across Runtime instances (prefill and decode phases)
         dec_attn_mask = torch.zeros(self.rbln_config.batch_size, self.rbln_config.max_seq_len, dtype=torch.float32)
-        block_tables = torch.zeros(
-            self.rbln_config.batch_size,
-            self.rbln_config.max_seq_len // self.rbln_config.kvcache_block_size,
-            dtype=torch.int16,
-        ).fill_(-1)
-        free_block_pool = deque(x for x in range(self.rbln_config.kvcache_num_blocks))
+        page_table_manager = RBLNPageTableManager(self.rbln_config)
+
+        common_kwargs = {
+            "main_input_name": "inputs_embeds" if self.rbln_config.use_inputs_embeds else "input_ids",
+            "embed_tokens": self.embed_tokens,
+            "dec_attn_mask": dec_attn_mask,
+            "page_table_manager": page_table_manager,
+            "rbln_config": self.rbln_config,
+        }
 
         self.prefill_decoder = RBLNGemma3RuntimeModel(
             runtime=self.model[0],
             image_prefill=self.model[1] if self.rbln_config.use_image_prefill else None,
-            main_input_name=main_input_name,
-            embed_tokens=self.embed_tokens,
             phase="prefill",
             batch_size=self.rbln_config.batch_size,
-            dec_attn_mask=dec_attn_mask,
-            block_tables=block_tables,
-            vocab_size=self.config.vocab_size,
-            free_block_pool=free_block_pool,
-            rbln_config=self.rbln_config,
+            **common_kwargs,
         )
 
         self.decoders = {}
         for i, batch_size in enumerate(self.rbln_config.decoder_batch_sizes):
             self.decoders[batch_size] = RBLNGemma3RuntimeModel(
                 runtime=self.model[i + self.rbln_config.decoder_runtime_idx],
-                main_input_name=main_input_name,
-                embed_tokens=self.embed_tokens,
                 phase="decode",
                 batch_size=batch_size,
-                dec_attn_mask=dec_attn_mask,
-                block_tables=block_tables,
-                free_block_pool=free_block_pool,
-                rbln_config=self.rbln_config,
+                **common_kwargs,
             )
 
         # NOTE(eunji): Use a decoder whose batch size matches the model's main batch size for compatibility.

optimum/rbln/transformers/models/gpt2/modeling_gpt2.py

@@ -47,6 +47,8 @@ class RBLNGPT2Model(RBLNDecoderOnlyModel):
 
     A class to convert and run pre-trained transformers based GPT2Model model on RBLN devices.
     It implements the methods to convert a pre-trained transformers GPT2Model 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.
     """
 
     _decoder_wrapper_cls = GPT2Wrapper

optimum/rbln/transformers/models/idefics3/configuration_idefics3.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -39,7 +39,7 @@ class RBLNIdefics3ForConditionalGenerationConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         vision_model: Optional[RBLNModelConfig] = None,
         text_model: Optional[RBLNModelConfig] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:

optimum/rbln/transformers/models/idefics3/modeling_idefics3.py

@@ -34,17 +34,11 @@ from transformers.models.idefics3.modeling_idefics3 import Idefics3CausalLMOutpu
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ....utils.runtime_utils import RBLNPytorchRuntime
-from ..decoderonly.modeling_decoderonly import (
-    RBLNDecoderOnlyForCausalLMOutput,
-)
+from ...modeling_outputs import RBLNDecoderOnlyOutput
 
 
 if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
-    )
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
 
 
 class RBLNRuntimeVisionModel(RBLNPytorchRuntime):
@@ -494,7 +488,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
         if not return_dict:
             return logits, generate_idx
         else:
-            return RBLNDecoderOnlyForCausalLMOutput(
+            return RBLNDecoderOnlyOutput(
                 logits=logits,
                 generate_idx=generate_idx,
             )