optimum-rbln 0.8.1rc0__py3-none-any.whl → 0.8.2__py3-none-any.whl

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

@@ -19,8 +19,6 @@ import torch.nn as nn
 from ....utils import logging
 from ...models.decoderonly.decoderonly_architecture import (
     DecoderOnlyAttention,
-    DecoderOnlyFlashAttention,
-    DecoderOnlyForCausalLM,
     DecoderOnlyLayer,
     DecoderOnlyModel,
     DecoderOnlyWrapper,
@@ -36,38 +34,23 @@ logger = logging.get_logger(__name__)
 class ExaoneForCausalLMWrapper(DecoderOnlyWrapper):
     """A wrapper class for the Exaone model with a language modeling head."""
 
-    def convert_to_rbln_causal_lm(self, causal_lm: "ExaoneForCausalLM", max_seq_len: int):
-        new_layers = []
-        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,
-                    use_position_ids=self.use_position_ids,
-                )
-            elif self.attn_impl == "flash_attn":
-                new_self_attn = ExaoneFlashAttention(
-                    layer.attn.attention,
-                    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}")
-
-            new_layer = ExaoneLayer(layer, new_self_attn)
-            new_layers.append(new_layer)
-        new_model = ExaoneModel(
-            causal_lm.transformer,
-            new_layers,
-            partition_len=self.kvcache_partition_len,
-            max_seq_len=max_seq_len,
-            sliding_window_layers=self.sliding_window_layers,
-        )
-        new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
-        return new_causal_lm
+    def get_decoder_layers(self, causal_lm: "ExaoneForCausalLM"):
+        return causal_lm.transformer.h
+
+    def get_attn_layer(self, layer: nn.Module):
+        return layer.attn.attention
+
+    def get_model_layer(self, causal_lm: "ExaoneForCausalLM"):
+        return causal_lm.transformer
+
+    def get_rbln_attn_class(self):
+        return ExaoneAttention
+
+    def get_rbln_layer_class(self):
+        return ExaoneLayer
+
+    def get_rbln_model_class(self):
+        return ExaoneModel
 
 
 class ExaoneModel(DecoderOnlyModel):
@@ -92,11 +75,3 @@ class ExaoneAttention(DecoderOnlyAttention):
         self.k_proj = self._original_mod.k_proj
         self.v_proj = self._original_mod.v_proj
         self.o_proj = self._original_mod.out_proj
-
-
-class ExaoneFlashAttention(DecoderOnlyFlashAttention):
-    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.out_proj

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

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

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

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

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

@@ -12,54 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING
 
-from ...models.decoderonly.decoderonly_architecture import (
-    DecoderOnlyAttention,
-    DecoderOnlyFlashAttention,
-    DecoderOnlyForCausalLM,
-    DecoderOnlyLayer,
-    DecoderOnlyModel,
-    DecoderOnlyWrapper,
-)
-
-
-if TYPE_CHECKING:
-    from transformers import GemmaForCausalLM
+from ...models.decoderonly.decoderonly_architecture import DecoderOnlyModel, DecoderOnlyWrapper
 
 
 class GemmaWrapper(DecoderOnlyWrapper):
-    def convert_to_rbln_causal_lm(self, causal_lm: "GemmaForCausalLM", max_seq_len: int):
-        new_layers = []
-        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,
-                    use_position_ids=self.use_position_ids,
-                )
-            elif self.attn_impl == "flash_attn":
-                new_self_attn = DecoderOnlyFlashAttention(
-                    layer.self_attn,
-                    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}")
-            new_layer = DecoderOnlyLayer(layer, new_self_attn)
-            new_layers.append(new_layer)
-        new_model = GemmaModel(
-            causal_lm.model,
-            new_layers,
-            partition_len=self.kvcache_partition_len,
-            max_seq_len=max_seq_len,
-            sliding_window_layers=self.sliding_window_layers,
-        )
-        new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
-        return new_causal_lm
+    def get_rbln_model_class(self):
+        return GemmaModel
 
 
 class GemmaModel(DecoderOnlyModel):

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

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

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

@@ -11,9 +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
-
-import rebel
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
 from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
@@ -23,10 +21,11 @@ from ..siglip.configuration_siglip import RBLNSiglipVisionModelConfig
 class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     def __init__(
         self,
-        prefill_chunk_size: Optional[int] = None,
         use_position_ids: Optional[bool] = None,
         use_attention_mask: Optional[bool] = None,
-        **kwargs: Dict[str, Any],
+        prefill_chunk_size: Optional[int] = None,
+        image_prefill_chunk_size: Optional[int] = None,
+        **kwargs: Any,
     ):
         # use_attention_mask and use_position_ids are always True for Gemma3
         use_attention_mask = use_attention_mask or True
@@ -39,10 +38,15 @@ class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
             use_position_ids=use_position_ids,
             **kwargs,
         )
+        self.image_prefill_chunk_size = image_prefill_chunk_size
 
-        npu = self.npu or rebel.get_npu_name()
-        if npu == "RBLN-CA02":
-            raise NotImplementedError("Gemma3 is currently not supported on RBLN-CA02")
+    @property
+    def use_image_prefill(self):
+        return self.image_prefill_chunk_size is not None
+
+    @property
+    def decoder_runtime_idx(self):
+        return 2 if self.use_image_prefill else 1
 
 
 class RBLNGemma3ForConditionalGenerationConfig(RBLNModelConfig):
@@ -53,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:
@@ -72,3 +76,11 @@ class RBLNGemma3ForConditionalGenerationConfig(RBLNModelConfig):
 
         self.vision_tower = self.init_submodule_config(RBLNSiglipVisionModelConfig, vision_tower)
         self.language_model = self.init_submodule_config(RBLNGemma3ForCausalLMConfig, 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

@@ -13,15 +13,13 @@
 # limitations under the License.
 
 import copy
-from typing import TYPE_CHECKING, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 
 import torch
 from transformers.models.gemma3.modeling_gemma3 import Gemma3RMSNorm
 
 from ..decoderonly.decoderonly_architecture import (
     DecoderOnlyAttention,
-    DecoderOnlyFlashAttention,
-    DecoderOnlyForCausalLM,
     DecoderOnlyLayer,
     DecoderOnlyModel,
     DecoderOnlyWrapper,
@@ -30,10 +28,6 @@ from ..decoderonly.decoderonly_architecture import (
 )
 
 
-if TYPE_CHECKING:
-    from transformers import Gemma3ForCausalLM
-
-
 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)
@@ -45,49 +39,14 @@ class Gemma3ForCausalLMWrapper(DecoderOnlyWrapper):
 
         return (rotary_emb_global, rotary_emb_local)
 
-    def convert_to_rbln_causal_lm(self, causal_lm: "Gemma3ForCausalLM", max_seq_len: int):
-        new_layers = []
-        for layer_idx, layer in enumerate(causal_lm.model.layers):
-            if layer_idx in self.sliding_window_layers:
-                new_self_attn = Gemma3Attention(
-                    layer.self_attn,
-                    use_attention_mask=None,  # FIXME: no use in SWA
-                    use_position_ids=self.use_position_ids,
-                    kvcache_block_size=self.config.sliding_window,
-                    is_sliding=True,
-                )
-            else:
-                if self.attn_impl == "eager":
-                    new_self_attn = Gemma3Attention(
-                        layer.self_attn,
-                        use_attention_mask=self.use_attention_mask,
-                        use_position_ids=self.use_position_ids,
-                        kvcache_block_size=self.kvcache_block_size,
-                        is_sliding=False,
-                    )
-                elif self.attn_impl == "flash_attn":
-                    new_self_attn = Gemma3FlashAttention(
-                        layer.self_attn,
-                        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}")
-
-            new_layer = Gemma3DecoderLayer(layer, new_self_attn)
-            new_layers.append(new_layer)
-
-        new_model = Gemma3TextModel(
-            causal_lm.model,
-            new_layers,
-            partition_len=self.kvcache_partition_len,
-            max_seq_len=max_seq_len,
-            sliding_window_layers=self.sliding_window_layers,
-        )
-        new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
-        return new_causal_lm
+    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):
@@ -199,16 +158,3 @@ class Gemma3Attention(DecoderOnlyAttention):
 
     def get_attn_scale(self):
         return self._original_mod.config.query_pre_attn_scalar**-0.5
-
-
-class Gemma3FlashAttention(DecoderOnlyFlashAttention):
-    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