optimum-rbln 0.7.5a1__py3-none-any.whl → 0.7.5rc0__py3-none-any.whl

optimum/rbln/__init__.py

@@ -38,6 +38,7 @@ _import_structure = {
         "RBLNAutoModelForCTC",
         "RBLNAutoModelForDepthEstimation",
         "RBLNAutoModelForImageClassification",
+        "RBLNAutoModelForImageTextToText",
         "RBLNAutoModelForMaskedLM",
         "RBLNAutoModelForQuestionAnswering",
         "RBLNAutoModelForSeq2SeqLM",
@@ -78,6 +79,10 @@ _import_structure = {
         "RBLNExaoneForCausalLMConfig",
         "RBLNGemmaForCausalLM",
         "RBLNGemmaForCausalLMConfig",
+        "RBLNGemma3ForCausalLM",
+        "RBLNGemma3ForCausalLMConfig",
+        "RBLNGemma3ForConditionalGeneration",
+        "RBLNGemma3ForConditionalGenerationConfig",
         "RBLNGPT2LMHeadModel",
         "RBLNGPT2LMHeadModelConfig",
         "RBLNIdefics3VisionTransformer",
@@ -259,6 +264,7 @@ if TYPE_CHECKING:
         RBLNAutoModelForCTC,
         RBLNAutoModelForDepthEstimation,
         RBLNAutoModelForImageClassification,
+        RBLNAutoModelForImageTextToText,
         RBLNAutoModelForMaskedLM,
         RBLNAutoModelForQuestionAnswering,
         RBLNAutoModelForSeq2SeqLM,
@@ -297,6 +303,10 @@ if TYPE_CHECKING:
         RBLNDPTForDepthEstimationConfig,
         RBLNExaoneForCausalLM,
         RBLNExaoneForCausalLMConfig,
+        RBLNGemma3ForCausalLM,
+        RBLNGemma3ForCausalLMConfig,
+        RBLNGemma3ForConditionalGeneration,
+        RBLNGemma3ForConditionalGenerationConfig,
         RBLNGemmaForCausalLM,
         RBLNGemmaForCausalLMConfig,
         RBLNGPT2LMHeadModel,

optimum/rbln/__version__.py

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

optimum/rbln/transformers/__init__.py

@@ -34,6 +34,7 @@ _import_structure = {
         "RBLNAutoModelForCTC",
         "RBLNAutoModelForDepthEstimation",
         "RBLNAutoModelForImageClassification",
+        "RBLNAutoModelForImageTextToText",
         "RBLNAutoModelForMaskedLM",
         "RBLNAutoModelForQuestionAnswering",
         "RBLNAutoModelForSeq2SeqLM",
@@ -72,6 +73,10 @@ _import_structure = {
         "RBLNExaoneForCausalLMConfig",
         "RBLNGemmaForCausalLM",
         "RBLNGemmaForCausalLMConfig",
+        "RBLNGemma3ForCausalLM",
+        "RBLNGemma3ForCausalLMConfig",
+        "RBLNGemma3ForConditionalGeneration",
+        "RBLNGemma3ForConditionalGenerationConfig",
         "RBLNGPT2LMHeadModel",
         "RBLNGPT2LMHeadModelConfig",
         "RBLNIdefics3VisionTransformer",
@@ -148,6 +153,7 @@ if TYPE_CHECKING:
         RBLNAutoModelForCTC,
         RBLNAutoModelForDepthEstimation,
         RBLNAutoModelForImageClassification,
+        RBLNAutoModelForImageTextToText,
         RBLNAutoModelForMaskedLM,
         RBLNAutoModelForQuestionAnswering,
         RBLNAutoModelForSeq2SeqLM,
@@ -184,6 +190,10 @@ if TYPE_CHECKING:
         RBLNDPTForDepthEstimationConfig,
         RBLNExaoneForCausalLM,
         RBLNExaoneForCausalLMConfig,
+        RBLNGemma3ForCausalLM,
+        RBLNGemma3ForCausalLMConfig,
+        RBLNGemma3ForConditionalGeneration,
+        RBLNGemma3ForConditionalGenerationConfig,
         RBLNGemmaForCausalLM,
         RBLNGemmaForCausalLMConfig,
         RBLNGPT2LMHeadModel,

optimum/rbln/transformers/models/__init__.py

@@ -31,6 +31,7 @@ _import_structure = {
         "RBLNAutoModelForSequenceClassification",
         "RBLNAutoModelForSpeechSeq2Seq",
         "RBLNAutoModelForVision2Seq",
+        "RBLNAutoModelForImageTextToText",
     ],
     "bart": [
         "RBLNBartForConditionalGeneration",
@@ -80,6 +81,12 @@ _import_structure = {
     ],
     "exaone": ["RBLNExaoneForCausalLM", "RBLNExaoneForCausalLMConfig"],
     "gemma": ["RBLNGemmaForCausalLM", "RBLNGemmaForCausalLMConfig"],
+    "gemma3": [
+        "RBLNGemma3ForCausalLM",
+        "RBLNGemma3ForCausalLMConfig",
+        "RBLNGemma3ForConditionalGeneration",
+        "RBLNGemma3ForConditionalGenerationConfig",
+    ],
     "gpt2": ["RBLNGPT2LMHeadModel", "RBLNGPT2LMHeadModelConfig"],
     "idefics3": [
         "RBLNIdefics3VisionTransformer",
@@ -121,6 +128,7 @@ if TYPE_CHECKING:
         RBLNAutoModelForCTC,
         RBLNAutoModelForDepthEstimation,
         RBLNAutoModelForImageClassification,
+        RBLNAutoModelForImageTextToText,
         RBLNAutoModelForMaskedLM,
         RBLNAutoModelForQuestionAnswering,
         RBLNAutoModelForSeq2SeqLM,
@@ -170,6 +178,12 @@ if TYPE_CHECKING:
     )
     from .exaone import RBLNExaoneForCausalLM, RBLNExaoneForCausalLMConfig
     from .gemma import RBLNGemmaForCausalLM, RBLNGemmaForCausalLMConfig
+    from .gemma3 import (
+        RBLNGemma3ForCausalLM,
+        RBLNGemma3ForCausalLMConfig,
+        RBLNGemma3ForConditionalGeneration,
+        RBLNGemma3ForConditionalGenerationConfig,
+    )
     from .gpt2 import RBLNGPT2LMHeadModel, RBLNGPT2LMHeadModelConfig
     from .idefics3 import (
         RBLNIdefics3ForConditionalGeneration,

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

@@ -19,6 +19,7 @@ from .modeling_auto import (
     RBLNAutoModelForCTC,
     RBLNAutoModelForDepthEstimation,
     RBLNAutoModelForImageClassification,
+    RBLNAutoModelForImageTextToText,
     RBLNAutoModelForMaskedLM,
     RBLNAutoModelForQuestionAnswering,
     RBLNAutoModelForSeq2SeqLM,

optimum/rbln/transformers/models/auto/modeling_auto.py

@@ -23,6 +23,8 @@ from transformers.models.auto.modeling_auto import (
     MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES,
     MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
     MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES,
+    MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING,
+    MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES,
     MODEL_FOR_MASKED_LM_MAPPING,
     MODEL_FOR_MASKED_LM_MAPPING_NAMES,
     MODEL_FOR_QUESTION_ANSWERING_MAPPING,
@@ -90,6 +92,11 @@ class RBLNAutoModelForVision2Seq(_BaseAutoModelClass):
     _model_mapping_names = MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES
 
 
+class RBLNAutoModelForImageTextToText(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING
+    _model_mapping_names = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES
+
+
 class RBLNAutoModelForMaskedLM(_BaseAutoModelClass):
     _model_mapping = MODEL_FOR_MASKED_LM_MAPPING
     _model_mapping_names = MODEL_FOR_MASKED_LM_MAPPING_NAMES

optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py

@@ -157,7 +157,11 @@ class DecoderOnlyWrapper(nn.Module):
         self.config = causal_lm.config
 
         if use_rotary_emb:
-            self.rotary_emb = self.get_rotary_emb(max_seq_len=max_seq_len)
+            rotary_embs = self.get_rotary_emb(max_seq_len=max_seq_len)
+            if isinstance(rotary_embs, tuple):
+                self.rotary_emb_global, self.rotary_emb_local = rotary_embs
+            else:
+                self.rotary_emb = rotary_embs
         else:
             self.rotary_emb = None
 
@@ -195,7 +199,10 @@ class DecoderOnlyWrapper(nn.Module):
         for layer in causal_lm.model.layers:
             if self.attn_impl == "eager":
                 new_self_attn = DecoderOnlyAttention(
-                    layer.self_attn, self.use_attention_mask, kvcache_block_size=self.kvcache_block_size
+                    layer.self_attn,
+                    self.use_attention_mask,
+                    self.use_position_ids,
+                    kvcache_block_size=self.kvcache_block_size,
                 )
             elif self.attn_impl == "flash_attn":
                 new_self_attn = DecoderOnlyFlashAttention(
@@ -203,6 +210,7 @@ class DecoderOnlyWrapper(nn.Module):
                     kvcache_partition_len=self.kvcache_partition_len,
                     kvcache_block_size=self.kvcache_block_size,
                     use_attention_mask=self.use_attention_mask,
+                    use_position_ids=self.use_position_ids,
                 )
             else:
                 raise NotImplementedError(f"Unknwon attn : {self.attn_impl}")
@@ -363,6 +371,13 @@ class DecoderOnlyForCausalLM(nn.Module):
             hidden_states = hidden_states[:, query_position.to(torch.int).unsqueeze(0)]
 
         logits = self.lm_head(hidden_states)
+
+        # Apply final logit softmaxing if configured, e.g. for Gemma2
+        if getattr(self.config, "final_logit_softcapping", None) is not None:
+            logits = logits / self.config.final_logit_softcapping
+            logits = torch.tanh(logits)
+            logits = logits * self.config.final_logit_softcapping
+
         return logits
 
 
@@ -610,7 +625,7 @@ class DecoderOnlyAttention(nn.Module):
         self_attn: Original attention module from the base model
     """
 
-    def __init__(self, self_attn, use_attention_mask, kvcache_block_size):
+    def __init__(self, self_attn, use_attention_mask, use_position_ids, kvcache_block_size):
         super().__init__()
         self._original_mod = self_attn
         self.layer_idx = self_attn.layer_idx
@@ -629,6 +644,7 @@ class DecoderOnlyAttention(nn.Module):
             self.num_key_value_heads = self.num_heads
 
         self.use_attention_mask = use_attention_mask
+        self.use_position_ids = use_position_ids
         self.attention = self.get_attention()
         self.kvcache_block_size = kvcache_block_size
         self.__post_init__()
@@ -643,7 +659,9 @@ class DecoderOnlyAttention(nn.Module):
         self.attention.phase = phase
 
     def get_attention(self):
-        return AttentionOp(self.num_heads, self.head_dim, self.num_key_value_heads, self.use_attention_mask)
+        return AttentionOp(
+            self.num_heads, self.head_dim, self.num_key_value_heads, self.use_attention_mask, self.use_position_ids
+        )
 
     def __post_init__(self):
         self.q_proj = self._original_mod.q_proj
@@ -716,13 +734,16 @@ class DecoderOnlyAttention(nn.Module):
 
 
 class AttentionOp(nn.Module):
-    def __init__(self, num_heads: int, head_dim: int, num_key_value_heads: int, use_attention_mask: bool):
+    def __init__(
+        self, num_heads: int, head_dim: int, num_key_value_heads: int, use_attention_mask: bool, use_position_ids: bool
+    ):
         super().__init__()
         self.num_heads = num_heads
         self.head_dim = head_dim
         self.num_key_value_heads = num_key_value_heads
         self.phase = "prefill"
         self.use_attention_mask = use_attention_mask
+        self.use_position_ids = use_position_ids
 
     def forward(
         self,
@@ -755,7 +776,8 @@ class AttentionOp(nn.Module):
         # reshape for removing repeat_kv (batch=1 , num_head, 1, q_len=1, head_dim)
         key_state = key_state.unsqueeze(2)  # 1, 32, 1, 128, 128
         value_state = value_state.unsqueeze(2)
-        if self.use_attention_mask:
+
+        if self.use_attention_mask and not self.use_position_ids:
             attn_mask = attn_mask.unsqueeze(2)
 
         if self.phase == "decode":
@@ -772,7 +794,7 @@ class AttentionOp(nn.Module):
         )
 
         if self.phase == "decode":
-            if self.use_attention_mask:
+            if self.use_attention_mask and not self.use_position_ids:
                 attn_output = torch.ops.rbln_custom_ops.paged_attn_decode(
                     q=query_state,
                     k=key_state,
@@ -796,11 +818,11 @@ class AttentionOp(nn.Module):
                     scale=scale,
                     block_table=block_tables,
                     block_size=block_size,
-                    mask=None,
+                    mask=attn_mask if self.use_position_ids else None,
                 )
 
         else:
-            if self.use_attention_mask:
+            if self.use_attention_mask and not self.use_position_ids:
                 attn_output = torch.ops.rbln_custom_ops.paged_attn_prefill(
                     q=query_state,
                     k=key_state,
@@ -824,8 +846,8 @@ class AttentionOp(nn.Module):
                     scale=scale,
                     block_table=block_tables,
                     block_size=block_size,
-                    is_bidirectional=False,
-                    mask=None,
+                    is_bidirectional=True if self.phase == "image_prefill" else False,  # FIXME, Hard-coded for Gemma3.
+                    mask=attn_mask if self.use_position_ids else None,
                 )
 
         attn_output = attn_output.view(batch_size, self.num_heads, -1, self.head_dim)
@@ -927,10 +949,13 @@ class RotaryEmbedding(nn.Module):
 
 
 class DecoderOnlyFlashAttention(DecoderOnlyAttention):
-    def __init__(self, self_attn, kvcache_partition_len, kvcache_block_size, use_attention_mask):
+    def __init__(self, self_attn, kvcache_partition_len, kvcache_block_size, use_attention_mask, use_position_ids):
         self.kvcache_partition_size = kvcache_partition_len
         super().__init__(
-            self_attn=self_attn, use_attention_mask=use_attention_mask, kvcache_block_size=kvcache_block_size
+            self_attn=self_attn,
+            use_attention_mask=use_attention_mask,
+            use_position_ids=use_position_ids,
+            kvcache_block_size=kvcache_block_size,
         )
 
     def get_attention(self):
@@ -940,6 +965,7 @@ class DecoderOnlyFlashAttention(DecoderOnlyAttention):
             self.num_key_value_heads,
             self.kvcache_partition_size,
             self.use_attention_mask,
+            self.use_position_ids,
         )
 
     def forward(
@@ -991,12 +1017,14 @@ class FlashAttentionOp(AttentionOp):
         num_key_value_heads: int,
         kvcache_partition_len: int,
         use_attention_mask: bool,
+        use_position_ids: bool,
     ):
         super().__init__(
             num_heads=num_heads,
             head_dim=head_dim,
             num_key_value_heads=num_key_value_heads,
             use_attention_mask=use_attention_mask,
+            use_position_ids=use_position_ids,
         )
         self.kvcache_partition_size = kvcache_partition_len
 
@@ -1016,7 +1044,7 @@ class FlashAttentionOp(AttentionOp):
         # reshape for removing repeat_kv (batch=1 , num_head, 1, q_len=1, head_dim)
         key_state = key_state.unsqueeze(2)
         value_state = value_state.unsqueeze(2)
-        if self.use_attention_mask:
+        if self.use_attention_mask and not self.use_position_ids:
             attn_mask = attn_mask.unsqueeze(2)
 
         if self.phase == "decode":
@@ -1033,7 +1061,7 @@ class FlashAttentionOp(AttentionOp):
         )
 
         if self.phase == "decode":
-            if self.use_attention_mask:
+            if self.use_attention_mask and not self.use_position_ids:
                 attn_output = torch.ops.rbln_custom_ops.paged_flash_attn_decode(
                     q=query_state,
                     k=key_state,
@@ -1059,10 +1087,10 @@ class FlashAttentionOp(AttentionOp):
                     block_table=block_tables,
                     block_size=kvcache_block_size,
                     partition=self.kvcache_partition_size,
-                    mask=None,
+                    mask=attn_mask if self.use_position_ids else None,
                 )
         else:
-            if self.use_attention_mask:
+            if self.use_attention_mask and not self.use_position_ids:
                 attn_output = torch.ops.rbln_custom_ops.paged_flash_attn_prefill(
                     q=query_state,
                     k=key_state,
@@ -1088,8 +1116,8 @@ class FlashAttentionOp(AttentionOp):
                     block_table=block_tables,
                     block_size=kvcache_block_size,
                     partition=self.kvcache_partition_size,
-                    is_bidirectional=False,
-                    mask=None,
+                    is_bidirectional=True if self.phase == "image_prefill" else False,
+                    mask=attn_mask if self.use_position_ids else None,
                 )
 
         # reshape for removing repeat_kv
@@ -1098,3 +1126,70 @@ class FlashAttentionOp(AttentionOp):
         attn_output = attn_output.reshape(batch_size, -1, self.num_heads * self.head_dim)
 
         return attn_output
+
+
+class SlidingWindowAttentionOp(AttentionOp):
+    def forward(
+        self,
+        query_state: torch.Tensor,
+        key_state: torch.Tensor,
+        value_state: torch.Tensor,
+        attn_mask: torch.Tensor,
+        past_key_state: torch.Tensor,
+        past_value_state: torch.Tensor,
+        seq_position: Tuple[torch.Tensor],
+        scale: torch.Tensor,
+        block_tables: torch.Tensor,
+        block_size: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # reshape for removing repeat_kv (batch=1 , num_head, 1, q_len=1, head_dim)
+        key_state = key_state.unsqueeze(2)
+        value_state = value_state.unsqueeze(2)
+
+        if self.phase == "decode":
+            batch_size = key_state.shape[0]
+        else:
+            batch_size = 1
+
+        query_state = query_state.view(
+            batch_size,
+            self.num_key_value_heads,
+            self.num_heads // self.num_key_value_heads,
+            -1,  # seq len
+            self.head_dim,
+        )
+
+        if self.phase == "decode":
+            attn_output = torch.ops.rbln_custom_ops.paged_sliding_window_attn_decode(
+                q=query_state,
+                k=key_state,
+                v=value_state,
+                kcache=past_key_state.unsqueeze(2),
+                vcache=past_value_state.unsqueeze(2),
+                cache_seq_len=seq_position[0],
+                cache_offset=seq_position[1],
+                scale=scale,
+                block_table=block_tables,
+                block_size=block_size,
+            )
+        else:
+            attn_output = torch.ops.rbln_custom_ops.paged_sliding_window_attn_prefill(
+                q=query_state,
+                k=key_state,
+                v=value_state,
+                kcache=past_key_state.unsqueeze(2),
+                vcache=past_value_state.unsqueeze(2),
+                cache_seq_len=seq_position[0],
+                cache_offset=seq_position[1],
+                scale=scale,
+                block_table=block_tables,
+                block_size=block_size,
+                is_bidirectional=True if self.phase == "image_prefill" else False,
+            )
+
+        # reshape for removing repeat_kv
+        attn_output = attn_output.view(batch_size, self.num_heads, -1, self.head_dim)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(batch_size, -1, self.num_heads * self.head_dim)
+
+        return attn_output

optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py

@@ -167,6 +167,8 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         block_tables: Optional[torch.Tensor] = None,
         position_embed: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        local_block_tables: Optional[torch.Tensor] = None,
     ):
         if input_ids is None and inputs_embeds is None:
             raise ValueError("Either `input_ids` or `inputs_embeds` must be provided.")
@@ -193,6 +195,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
                 attention_mask=attention_mask,
                 position_embed=position_embed,
                 position_ids=position_ids,
+                local_block_tables=local_block_tables,
             )
         else:
             return self.prefill_forward(
@@ -202,6 +205,8 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
                 batch_idx,
                 block_tables,
                 position_embed=position_embed,
+                token_type_ids=token_type_ids,
+                local_block_tables=local_block_tables,
             )
 
     def decode_forward(
@@ -213,6 +218,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         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:
@@ -262,6 +268,8 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         cache_position: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         position_embed: Optional[torch.Tensor] = None,
+        local_block_tables: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
     ):
         """
         Prepare inputs for prefill phase.
@@ -345,6 +353,8 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         block_tables: torch.Tensor = None,
         is_external_block_tables: bool = None,
         position_embed: Optional[torch.Tensor] = None,
+        local_block_tables: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
     ) -> torch.FloatTensor:
         """
         Performs chunked prefill for efficient KV-cache updates and memory optimization.
@@ -360,7 +370,9 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             position_embed,
             padded_cache_lengths,
             query_length,
-        ) = self._prepare_prefill_inputs(inputs, cache_position, attention_mask, position_embed)
+        ) = self._prepare_prefill_inputs(
+            inputs, cache_position, attention_mask, position_embed, token_type_ids=token_type_ids
+        )
 
         # Process input in chunks of size `prefill_chunk_size`
         for step in range(0, query_length, self.prefill_chunk_size):
@@ -373,7 +385,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             if position_embed is not None:
                 position_embed_chunk = position_embed[:, :, :, step : step + self.prefill_chunk_size, :]
 
-            if self.use_attention_mask:
+            if self.use_attention_mask and not self.use_position_ids:
                 # Update attention mask to ensure proper causal behavior
                 if step >= self.prefill_chunk_size:
                     chunked_attention_mask[:, :, :, step - self.prefill_chunk_size : step] = 1
@@ -387,10 +399,10 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
                 input_chunk,
                 cache_pos_chunk,
                 block_tables,
+                position_embed_chunk if position_embed is not None else None,
                 query_position,
                 chunked_attention_mask if self.use_attention_mask else None,
-                position_ids_chunk if position_ids is not None else None,
-                position_embed_chunk if position_embed is not None else None,
+                position_ids_chunk if self.use_position_ids else None,
                 out=out_buffers,
             )
 
@@ -440,12 +452,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         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)
-            with no_init_weights():
-                self.embed_tokens = torch.nn.Embedding(
-                    self.config.vocab_size,
-                    self.config.hidden_size,
-                    self.config.pad_token_id,
-                )
+            self.embed_tokens = self._create_embedding_layer()
             self.embed_tokens.load_state_dict(artifacts["embed_tokens"])
         else:
             self.embed_tokens = None
@@ -478,6 +485,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             attn_impl=self.rbln_config.attn_impl,
             use_position_ids=self.rbln_config.use_position_ids,
         )
+
         self.decoders = {}
         for i, batch_size in enumerate(self.rbln_config.decoder_batch_sizes):
             self.decoders[batch_size] = RBLNRuntimeModel(
@@ -515,6 +523,15 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             save_dict["embed_tokens"] = model.get_input_embeddings().state_dict()
             torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
+    def _create_embedding_layer(self):
+        with no_init_weights():
+            embed_tokens = torch.nn.Embedding(
+                self.config.vocab_size,
+                self.config.hidden_size,
+                self.config.pad_token_id,
+            )
+        return embed_tokens
+
     def get_input_embeddings(self):
         return self.embed_tokens
 
@@ -1101,6 +1118,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         generate_idx: Optional[torch.Tensor] = None,
         padded_cache_lengths: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
         return_dict: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> Tuple[torch.FloatTensor]:
@@ -1123,6 +1141,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
                     attention_mask=attention_mask[b_idx] if attention_mask is not None else None,
                     cache_position=cache_position,
                     batch_idx=b_idx,
+                    token_type_ids=token_type_ids[b_idx : b_idx + 1] if token_type_ids is not None else None,
                 )
                 padded_cache_lengths[b_idx] += output.padded_cache_lengths
                 logits.append(output.logits)

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

@@ -41,7 +41,10 @@ class ExaoneForCausalLMWrapper(DecoderOnlyWrapper):
         for layer in causal_lm.transformer.h:
             if self.attn_impl == "eager":
                 new_self_attn = ExaoneAttention(
-                    layer.attn.attention, self.use_attention_mask, kvcache_block_size=self.kvcache_block_size
+                    layer.attn.attention,
+                    self.use_attention_mask,
+                    kvcache_block_size=self.kvcache_block_size,
+                    use_position_ids=self.use_position_ids,
                 )
             elif self.attn_impl == "flash_attn":
                 new_self_attn = ExaoneFlashAttention(
@@ -49,6 +52,7 @@ class ExaoneForCausalLMWrapper(DecoderOnlyWrapper):
                     kvcache_partition_len=self.kvcache_partition_len,
                     use_attention_mask=self.use_attention_mask,
                     kvcache_block_size=self.kvcache_block_size,
+                    use_position_ids=self.use_position_ids,
                 )
             else:
                 raise NotImplementedError(f"Unknwon attn : {self.attn_impl}")

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

@@ -34,7 +34,10 @@ class GemmaWrapper(DecoderOnlyWrapper):
         for layer in causal_lm.model.layers:
             if self.attn_impl == "eager":
                 new_self_attn = DecoderOnlyAttention(
-                    layer.self_attn, self.use_attention_mask, kvcache_block_size=self.kvcache_block_size
+                    layer.self_attn,
+                    self.use_attention_mask,
+                    kvcache_block_size=self.kvcache_block_size,
+                    use_position_ids=self.use_position_ids,
                 )
             elif self.attn_impl == "flash_attn":
                 new_self_attn = DecoderOnlyFlashAttention(
@@ -42,6 +45,7 @@ class GemmaWrapper(DecoderOnlyWrapper):
                     kvcache_partition_len=self.kvcache_partition_len,
                     use_attention_mask=self.use_attention_mask,
                     kvcache_block_size=self.kvcache_block_size,
+                    use_position_ids=self.use_position_ids,
                 )
             else:
                 raise NotImplementedError(f"Unknwon attn : {self.attn_impl}")

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