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

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

@@ -38,7 +38,10 @@ class GPT2Wrapper(DecoderOnlyWrapper):
         new_layers = []
         for layer in causal_lm.transformer.h:
             new_self_attn = GPT2Attention(
-                layer.attn, self.use_attention_mask, kvcache_block_size=self.kvcache_block_size
+                layer.attn,
+                self.use_attention_mask,
+                kvcache_block_size=self.kvcache_block_size,
+                use_position_ids=self.use_position_ids,
             )
             new_layer = GPT2Layer(layer, new_self_attn)
             new_layers.append(new_layer)

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

@@ -421,6 +421,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
         image_hidden_states: Optional[torch.FloatTensor] = None,
         cache_position: torch.Tensor = None,
         generate_idx: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
         **kwargs,
     ) -> Union[Tuple, Idefics3CausalLMOutputWithPast]:
         # Prefill
@@ -434,14 +435,14 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
 
             for b_idx in range(batch_size):
                 cache_position = torch.arange(0, generate_idx[b_idx].item(), dtype=torch.int32).unsqueeze(0)
-                logit = self.text_model.prefill_decoder(
+                output = self.text_model.prefill_decoder(
                     input_ids=inputs[b_idx : b_idx + 1] if inputs_embeds is None else None,
                     inputs_embeds=inputs[b_idx : b_idx + 1] if inputs_embeds is not None else None,
                     attention_mask=attention_mask[b_idx] if attention_mask is not None else None,
                     cache_position=cache_position,
                     batch_idx=b_idx,
                 )
-                logits.append(logit)
+                logits.append(output.logits)
 
             logits = torch.cat(logits, dim=0)
 
@@ -451,9 +452,12 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
                 input_ids=input_ids,
                 inputs_embeds=inputs_embeds,
                 cache_position=cache_position,
-            )
+            ).logits
 
-        return RBLNDecoderOnlyOutput(
-            logits=logits,
-            generate_idx=generate_idx,
-        )
+        if not return_dict:
+            return logits, generate_idx
+        else:
+            return RBLNDecoderOnlyOutput(
+                logits=logits,
+                generate_idx=generate_idx,
+            )

optimum/rbln/transformers/models/llava_next/modeling_llava_next.py

@@ -372,7 +372,7 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
             inputs_embeds = [inputs_embeds[i : i + 1, attention_mask[i].bool()] for i in range(batch_size)]
             for batch_idx in range(batch_size):
                 generate_idx[batch_idx] = inputs_embeds[batch_idx].shape[-2]
-                logit = self.language_model.prefill_decoder(
+                output = self.language_model.prefill_decoder(
                     inputs_embeds=inputs_embeds[batch_idx],
                     batch_idx=batch_idx,
                     cache_position=torch.arange(
@@ -382,14 +382,14 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
                     ).unsqueeze(0),
                 )
 
-                logits.append(logit)
+                logits.append(output.logits)
             logits = torch.cat(logits, dim=0)
         else:
-            logits = self.language_model.decoder(
+            output = self.language_model.decoder(
                 inputs_embeds=inputs_embeds,
                 cache_position=cache_position,
             )
-
+            logits = output.logits
         return RBLNDecoderOnlyOutput(logits=logits, generate_idx=generate_idx)
 
     # Almost copied from : https://github.com/huggingface/transformers/blob/6b550462139655d488d4c663086a63e98713c6b9/src/transformers/models/llava_next/modeling_llava_next.py

optimum/rbln/transformers/models/midm/midm_architecture.py

@@ -61,7 +61,10 @@ class MidmLMHeadModelWrapper(DecoderOnlyWrapper):
         new_layers = []
         for layer in causal_lm.transformer.h:
             new_self_attn = MidmAttention(
-                layer.attn, self.use_attention_mask, kvcache_block_size=self.kvcache_block_size
+                layer.attn,
+                self.use_attention_mask,
+                kvcache_block_size=self.kvcache_block_size,
+                use_position_ids=self.use_position_ids,
             )
             new_layer = MidmLayer(layer, new_self_attn)
             new_layers.append(new_layer)

optimum/rbln/transformers/models/opt/__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_opt import RBLNOPTForCausalLMConfig
+from .modeling_opt import RBLNOPTForCausalLM

optimum/rbln/transformers/models/opt/configuration_opt.py

@@ -0,0 +1,19 @@
+# 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 ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+
+
+class RBLNOPTForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    pass

optimum/rbln/transformers/models/opt/modeling_opt.py

@@ -0,0 +1,80 @@
+# 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 torch.nn as nn
+from transformers import PreTrainedModel
+
+from ....utils import logging
+from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from .opt_architecture import OPTWrapper
+
+
+logger = logging.get_logger(__name__)
+
+
+class MLP(nn.Module):
+    def __init__(self, fc1, fc2, activation_fn):
+        super(MLP, self).__init__()
+        self.fc1 = fc1
+        self.fc2 = fc2
+        self.activation_fn = activation_fn
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.activation_fn(x)
+        x = self.fc2(x)
+        return x
+
+
+class RBLNOPTForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The OPT 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 OPTForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers OPTForCausalLM 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 = OPTWrapper
+    _use_rotary_emb = False
+
+    def modify_opt_decoder_layer(layer):
+        mlp = MLP(layer.fc1, layer.fc2, layer.activation_fn)
+        layer.mlp = mlp
+        del layer.fc1
+        del layer.fc2
+        del layer.activation_fn
+
+        return layer
+
+    @classmethod
+    def wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
+        wrapper_cfg = {
+            "max_seq_len": rbln_config.max_seq_len,
+            "attn_impl": rbln_config.attn_impl,
+            "kvcache_partition_len": rbln_config.kvcache_partition_len,
+            "kvcache_block_size": rbln_config.kvcache_block_size,
+            "use_rotary_emb": cls._use_rotary_emb,
+            "use_attention_mask": rbln_config.use_attention_mask,
+            "use_position_ids": rbln_config.use_position_ids,
+            "use_inputs_embeds": rbln_config.use_inputs_embeds,
+        }
+
+        for i in range(len(model.model.decoder.layers)):
+            model.model.decoder.layers[i] = cls.modify_opt_decoder_layer(model.model.decoder.layers[i])
+
+        return cls._decoder_wrapper_cls(model, **wrapper_cfg).eval()

optimum/rbln/transformers/models/opt/opt_architecture.py

@@ -0,0 +1,77 @@
+# 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 TYPE_CHECKING
+
+import torch.nn as nn
+
+from ...models.decoderonly.decoderonly_architecture import (
+    DecoderOnlyAttention,
+    DecoderOnlyForCausalLM,
+    DecoderOnlyLayer,
+    DecoderOnlyModel,
+    DecoderOnlyWrapper,
+)
+
+
+if TYPE_CHECKING:
+    from transformers import OPTForCausalLM
+
+
+class OPTWrapper(DecoderOnlyWrapper):
+    def convert_to_rbln_causal_lm(self, causal_lm: "OPTForCausalLM", max_seq_len: int):
+        if self.attn_impl != "eager":
+            raise NotImplementedError(f"flash attention ({self.attn_impl}) is not implemented for {self.__class__}")
+
+        new_layers = []
+
+        for layer in causal_lm.model.decoder.layers:
+            new_self_attn = OPTAttention(
+                layer.self_attn,
+                self.use_attention_mask,
+                kvcache_block_size=self.kvcache_block_size,
+                use_position_ids=self.use_position_ids,
+            )
+            new_layer = OPTDecoderLayer(layer, new_self_attn)
+            new_layers.append(new_layer)
+        new_model = OPTModel(causal_lm.model.decoder, new_layers, max_seq_len=max_seq_len, use_learned_pos_emb=True)
+        new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
+        return new_causal_lm
+
+
+class OPTAttention(DecoderOnlyAttention):
+    def __post_init__(self):
+        self.k_proj = self._original_mod.k_proj
+        self.v_proj = self._original_mod.v_proj
+        self.q_proj = self._original_mod.q_proj
+        self.o_proj = self._original_mod.out_proj
+
+
+class OPTModel(DecoderOnlyModel):
+    def get_embedding(self) -> nn.Embedding:
+        return self._original_mod.embed_tokens
+
+    def get_pos_embedding(self):
+        return self._original_mod.embed_positions
+
+    def get_last_layernorm(self) -> nn.LayerNorm:
+        return self._original_mod.final_layer_norm
+
+
+class OPTDecoderLayer(DecoderOnlyLayer):
+    def get_pre_attention_layernorm(self) -> nn.LayerNorm:
+        return self._original_mod.self_attn_layer_norm
+
+    def get_post_attention_layernorm(self) -> nn.LayerNorm:
+        return self._original_mod.final_layer_norm

optimum/rbln/transformers/models/phi/phi_architecture.py

@@ -37,7 +37,10 @@ class PhiWrapper(DecoderOnlyWrapper):
         for layer in causal_lm.model.layers:
             if self.attn_impl == "eager":
                 new_self_attn = PhiAttention(
-                    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":
                 raise NotImplementedError(f"flash attn for {self.__class__} is not implemented yet.")

optimum/rbln/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py

@@ -371,6 +371,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         query_length: int,
         use_inputs_embeds: bool,
         use_attention_mask: bool,
+        use_position_ids: bool,
         max_seq_len: int,
         kvcache_block_size: int,
         kvcache_num_blocks: int,
@@ -384,6 +385,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
             query_length,
             use_inputs_embeds,
             use_attention_mask,
+            use_position_ids,
             max_seq_len,
             kvcache_block_size,
             kvcache_num_blocks,
@@ -392,8 +394,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
             hidden_size,
             head_dim,
         )
-        pos_idx = 5 if query_length > 1 else 4
-        pos_idx = pos_idx if use_attention_mask else pos_idx - 1
+        pos_idx = 3
         input_info.insert(pos_idx, ("position_emb", [2, batch_size, 1, query_length, head_dim], "float32"))
 
         return input_info
@@ -562,7 +563,8 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         video_grid_thw: Optional[torch.LongTensor] = None,
         cache_position: Optional[torch.LongTensor] = None,
         second_per_grid_ts: Optional[torch.Tensor] = None,
-        generate_idx: torch.Tensor = None,
+        generate_idx: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
         **kwargs,
     ) -> RBLNDecoderOnlyOutput:
         # Prefill
@@ -584,25 +586,30 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
             for b_idx in range(batch_size):
                 cache_position = torch.arange(0, generate_idx[b_idx].item(), dtype=torch.int32).unsqueeze(0)
 
-                logit = self.prefill_decoder(
+                output = self.prefill_decoder(
                     inputs_embeds=inputs_embeds[b_idx : b_idx + 1],
                     attention_mask=attention_mask[b_idx] if attention_mask is not None else None,
                     cache_position=cache_position,
                     batch_idx=b_idx,
                     position_embed=position_embed[:, b_idx : b_idx + 1],
                 )
-                logits.append(logit)
+                logits.append(output.logits)
             logits = torch.cat(logits, dim=0)
-        # Decoder
+            # Decoder
         else:
+            print(input_ids[0], cache_position[0])
             inputs_embeds, position_embed = self._preprocess_decoder(input_ids, cache_position)
-            logits = self.decoder(
+            output = self.decoder(
                 inputs_embeds=inputs_embeds,
                 cache_position=cache_position,
                 position_embed=position_embed,
             )
+            logits = output.logits
 
-        return RBLNDecoderOnlyOutput(
-            logits=logits,
-            generate_idx=generate_idx,
-        )
+        if not return_dict:
+            return logits, generate_idx
+        else:
+            return RBLNDecoderOnlyOutput(
+                logits=logits,
+                generate_idx=generate_idx,
+            )

optimum/rbln/transformers/models/qwen2_5_vl/qwen2_5_vl_architecture.py

@@ -157,58 +157,41 @@ class Qwen2_5_VLVisionWindowAttention(nn.Module):
 
 
 class Qwen2_5_VL_LanguageModelWrapper(DecoderOnlyWrapper):
-    def forward(self, *args):
-        if self.phase == "decode":
-            if self.use_attention_mask:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    attention_mask,
-                    block_tables,
-                    position_emb,
-                    *past_key_values,
-                ) = args
-            else:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    block_tables,
-                    position_emb,
-                    *past_key_values,
-                ) = args
-                attention_mask = None
-            query_position = None
-        elif self.phase == "prefill":
-            if self.use_attention_mask:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    attention_mask,
-                    query_position,
-                    block_tables,
-                    position_emb,
-                    *past_key_values,
-                ) = args
-            else:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    query_position,
-                    block_tables,
-                    position_emb,
-                    *past_key_values,
-                ) = args
-                attention_mask = None
-
-        else:
-            raise ValueError(f"Unknown phase: {self.phase}")
-
-        return self.forward_common(
-            input_ids_or_inputs_embeds,
+    def prepare_forward_args(self, *args):
+        args = list(args)
+        input_ids = None if self.use_inputs_embeds else args.pop(0)
+        inputs_embeds = args.pop(0) if self.use_inputs_embeds else None
+        cache_position = args.pop(0)
+        block_tables = args.pop(0)
+        position_embeds = args.pop(0)
+        query_position = args.pop(0) if self.phase == "prefill" else None
+        position_ids = None
+        attention_mask = args.pop(0) if self.use_attention_mask else None
+        past_key_values = args
+
+        if len(past_key_values) != 2 * self.num_hidden_layers:
+            raise ValueError(
+                f"Different past_key_values to model's config. {len(past_key_values)} != {2 * self.num_hidden_layers}"
+            )
+
+        # [key, value] * n_layer -> ( (key, value) ) * n_layer
+        # cache shape : batch, n_heads, 1, max_seq_len, head_dim
+        _past_key_values = []
+        for i in range(self.config.num_hidden_layers):
+            key_states = past_key_values[i * 2]
+            value_states = past_key_values[i * 2 + 1]
+            past_key_value = [key_states, value_states]
+            _past_key_values.append(past_key_value)
+        past_key_values = _past_key_values
+
+        return (
+            input_ids,
+            inputs_embeds,
             cache_position,
-            attention_mask,
-            query_position,
             block_tables,
-            position_emb,
-            *past_key_values,
+            query_position,
+            attention_mask,
+            position_ids,
+            past_key_values,
+            position_embeds,
         )

optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py

@@ -476,6 +476,8 @@ class Seq2SeqSelfAttention(nn.Module):
         ]
         if attention_mask is not None:
             args.insert(3, attention_mask.unsqueeze(2))
+        else:
+            args.append(None)
 
         attn_output = self.attn_decode(*args)
 

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

@@ -0,0 +1,20 @@
+# 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_siglip import (
+    RBLNSiglipVisionModelConfig,
+)
+from .modeling_siglip import (
+    RBLNSiglipVisionModel,
+)

optimum/rbln/transformers/models/siglip/configuration_siglip.py

@@ -0,0 +1,66 @@
+# 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
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNSiglipVisionModelConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        image_size: Optional[int] = None,
+        interpolate_pos_encoding: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for image processing. Defaults to 1.
+            image_size (Optional[int]): The size of input images. Can be an integer for square images,
+                a tuple/list (height, width), or a dictionary with 'height' and 'width' keys.
+            interpolate_pos_encoding (Optional[bool]): Whether to interpolate the position encoding.
+            output_hidden_states: (Optional[bool]): Whether to return hidden states.
+            **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}")
+
+        self.image_size = image_size
+        self.interpolate_pos_encoding = interpolate_pos_encoding or False
+        self.output_hidden_states = output_hidden_states
+
+    @property
+    def image_width(self):
+        if isinstance(self.image_size, int):
+            return self.image_size
+        elif isinstance(self.image_size, (list, tuple)):
+            return self.image_size[1]
+        else:
+            return self.image_size["width"]
+
+    @property
+    def image_height(self):
+        if isinstance(self.image_size, int):
+            return self.image_size
+        elif isinstance(self.image_size, (list, tuple)):
+            return self.image_size[0]
+        else:
+            return self.image_size["height"]