optimum-rbln 0.8.2a5__py3-none-any.whl → 0.8.2a6__py3-none-any.whl

optimum/rbln/__init__.py

@@ -102,6 +102,8 @@ _import_structure = {
         "RBLNLlamaModelConfig",
         "RBLNOPTForCausalLM",
         "RBLNOPTForCausalLMConfig",
+        "RBLNLlavaForConditionalGeneration",
+        "RBLNLlavaForConditionalGenerationConfig",
         "RBLNLlavaNextForConditionalGeneration",
         "RBLNLlavaNextForConditionalGenerationConfig",
         "RBLNMidmLMHeadModel",
@@ -118,6 +120,8 @@ _import_structure = {
         "RBLNPegasusModelConfig",
         "RBLNPhiForCausalLM",
         "RBLNPhiForCausalLMConfig",
+        "RBLNPixtralVisionModel",
+        "RBLNPixtralVisionModelConfig",
         "RBLNPhiModel",
         "RBLNPhiModelConfig",
         "RBLNQwen2ForCausalLM",
@@ -369,6 +373,8 @@ if TYPE_CHECKING:
         RBLNLlamaForCausalLMConfig,
         RBLNLlamaModel,
         RBLNLlamaModelConfig,
+        RBLNLlavaForConditionalGeneration,
+        RBLNLlavaForConditionalGenerationConfig,
         RBLNLlavaNextForConditionalGeneration,
         RBLNLlavaNextForConditionalGenerationConfig,
         RBLNMidmLMHeadModel,
@@ -389,6 +395,8 @@ if TYPE_CHECKING:
         RBLNPhiForCausalLMConfig,
         RBLNPhiModel,
         RBLNPhiModelConfig,
+        RBLNPixtralVisionModel,
+        RBLNPixtralVisionModelConfig,
         RBLNQwen2_5_VisionTransformerPretrainedModel,
         RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
         RBLNQwen2_5_VLForConditionalGeneration,

optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.8.2a5'
-__version_tuple__ = version_tuple = (0, 8, 2, 'a5')
+__version__ = version = '0.8.2a6'
+__version_tuple__ = version_tuple = (0, 8, 2, 'a6')

optimum/rbln/transformers/__init__.py

@@ -86,6 +86,8 @@ _import_structure = {
         "RBLNIdefics3VisionTransformerConfig",
         "RBLNLlamaForCausalLM",
         "RBLNLlamaForCausalLMConfig",
+        "RBLNLlavaForConditionalGeneration",
+        "RBLNLlavaForConditionalGenerationConfig",
         "RBLNLlamaModel",
         "RBLNLlamaModelConfig",
         "RBLNOPTForCausalLM",
@@ -108,6 +110,8 @@ _import_structure = {
         "RBLNOPTModelConfig",
         "RBLNPhiForCausalLM",
         "RBLNPhiForCausalLMConfig",
+        "RBLNPixtralVisionModelConfig",
+        "RBLNPixtralVisionModel",
         "RBLNPhiModel",
         "RBLNPhiModelConfig",
         "RBLNQwen2_5_VisionTransformerPretrainedModel",
@@ -220,6 +224,8 @@ if TYPE_CHECKING:
         RBLNLlamaForCausalLMConfig,
         RBLNLlamaModel,
         RBLNLlamaModelConfig,
+        RBLNLlavaForConditionalGeneration,
+        RBLNLlavaForConditionalGenerationConfig,
         RBLNLlavaNextForConditionalGeneration,
         RBLNLlavaNextForConditionalGenerationConfig,
         RBLNMidmLMHeadModel,
@@ -240,6 +246,8 @@ if TYPE_CHECKING:
         RBLNPhiForCausalLMConfig,
         RBLNPhiModel,
         RBLNPhiModelConfig,
+        RBLNPixtralVisionModel,
+        RBLNPixtralVisionModelConfig,
         RBLNQwen2_5_VisionTransformerPretrainedModel,
         RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
         RBLNQwen2_5_VLForConditionalGeneration,

optimum/rbln/transformers/models/__init__.py

@@ -106,6 +106,7 @@ _import_structure = {
         "RBLNIdefics3ForConditionalGenerationConfig",
         "RBLNIdefics3VisionTransformerConfig",
     ],
+    "llava": ["RBLNLlavaForConditionalGeneration", "RBLNLlavaForConditionalGenerationConfig"],
     "llama": ["RBLNLlamaForCausalLM", "RBLNLlamaForCausalLMConfig", "RBLNLlamaModel", "RBLNLlamaModelConfig"],
     "opt": ["RBLNOPTForCausalLM", "RBLNOPTForCausalLMConfig", "RBLNOPTModel", "RBLNOPTModelConfig"],
     "pegasus": [
@@ -116,6 +117,7 @@ _import_structure = {
     ],
     "llava_next": ["RBLNLlavaNextForConditionalGeneration", "RBLNLlavaNextForConditionalGenerationConfig"],
     "midm": ["RBLNMidmLMHeadModel", "RBLNMidmLMHeadModelConfig"],
+    "pixtral": ["RBLNPixtralVisionModel", "RBLNPixtralVisionModelConfig"],
     "mistral": [
         "RBLNMistralForCausalLM",
         "RBLNMistralForCausalLMConfig",
@@ -241,6 +243,7 @@ if TYPE_CHECKING:
         RBLNIdefics3VisionTransformerConfig,
     )
     from .llama import RBLNLlamaForCausalLM, RBLNLlamaForCausalLMConfig, RBLNLlamaModel, RBLNLlamaModelConfig
+    from .llava import RBLNLlavaForConditionalGeneration, RBLNLlavaForConditionalGenerationConfig
     from .llava_next import RBLNLlavaNextForConditionalGeneration, RBLNLlavaNextForConditionalGenerationConfig
     from .midm import RBLNMidmLMHeadModel, RBLNMidmLMHeadModelConfig
     from .mistral import RBLNMistralForCausalLM, RBLNMistralForCausalLMConfig, RBLNMistralModel, RBLNMistralModelConfig
@@ -252,6 +255,7 @@ if TYPE_CHECKING:
         RBLNPegasusModelConfig,
     )
     from .phi import RBLNPhiForCausalLM, RBLNPhiForCausalLMConfig, RBLNPhiModel, RBLNPhiModelConfig
+    from .pixtral import RBLNPixtralVisionModel, RBLNPixtralVisionModelConfig
     from .qwen2 import RBLNQwen2ForCausalLM, RBLNQwen2ForCausalLMConfig, RBLNQwen2Model, RBLNQwen2ModelConfig
     from .qwen2_5_vl import (
         RBLNQwen2_5_VisionTransformerPretrainedModel,

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union, get_args
 
 from ....configuration_utils import RBLNModelConfig
 from ....utils.logging import get_logger
@@ -352,6 +352,8 @@ class RBLNDecoderOnlyModelForCausalLMConfig(RBLNDecoderOnlyModelConfig):
         if not isinstance(self.batch_size, int) or self.batch_size < 0:
             raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
 
+        if phases is not None:
+            self.validate_phases_type(phases)
         self.phases = phases or ["prefill", "decode"]
 
         if "decode" in self.phases:
@@ -374,6 +376,13 @@ class RBLNDecoderOnlyModelForCausalLMConfig(RBLNDecoderOnlyModelConfig):
                 # Larger batch size should be at the beginning of the list.
                 self.decoder_batch_sizes.sort(reverse=True)
 
+    @staticmethod
+    def validate_phases_type(phases: List[PhaseType]):
+        if not isinstance(phases, list):
+            raise ValueError("`phases` must be a list.")
+        if not all(phase in get_args(PhaseType) for phase in phases):
+            raise ValueError(f"All elements in `phases` must be of type `PhaseType`({get_args(PhaseType)}).")
+
     @property
     def use_multiple_decoder(self):
         return isinstance(self.decoder_batch_sizes, list) and len(self.decoder_batch_sizes) > 1

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

@@ -1024,8 +1024,9 @@ class SlidingWindowAttentionOp(AttentionOp):
             "block_size": block_size,
         }
 
-        if "prefill" in self.phase:
-            op_args["is_bidirectional"] = True
+        if self.phase == "prefill" or self.phase == "image_prefill":
+            if not self.use_attention_mask or self.use_position_ids:
+                op_args["is_bidirectional"] = self.phase == "image_prefill"  # FIXME, Hard-coded for Gemma3.
 
         attn_op_name = self.get_attn_op_name()
         attn_op = getattr(torch.ops.rbln_custom_ops, attn_op_name, None)

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

@@ -1213,7 +1213,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNDecoderOnlyModel):
                 kvcache_block_size=rbln_config.kvcache_block_size,
                 nbits_per_param=16 if not rbln_config.quantization else 4,  # TODO(jongho): FIX Ad-hoc
                 n_model_params=sum(p.numel() for p in model.parameters()),
-                num_runtimes=1 if rbln_config.can_generate else 1 + len(rbln_config.decoder_batch_sizes),
+                num_runtimes=1 if not rbln_config.can_generate else 1 + len(rbln_config.decoder_batch_sizes),
             )
 
             max_num_blocks = min(max_num_blocks, estimated_max_num_blocks)
@@ -1395,8 +1395,12 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNDecoderOnlyModel):
         # The decoder stage operates as usual, processing inputs in batch mode.
 
         # for only use forward
-        if not self.can_generate():
-            generate_idx = attention_mask.sum(dim=-1, keepdim=True).int()
+        if generate_idx is None:
+            generate_idx = (
+                attention_mask.sum(dim=-1, keepdim=True).int()
+                if attention_mask is not None
+                else torch.full((input_ids.shape[0], 1), input_ids.shape[1], dtype=torch.int32)
+            )
             padded_cache_lengths = torch.zeros_like(generate_idx)
 
         # Prefll

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

@@ -37,6 +37,14 @@ class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
         )
         self.image_prefill_chunk_size = image_prefill_chunk_size
 
+    @property
+    def use_image_prefill(self):
+        return self.image_prefill_chunk_size is not None
+
+    @property
+    def decoder_runtime_idx(self):
+        return 2 if self.use_image_prefill else 1
+
 
 class RBLNGemma3ForConditionalGenerationConfig(RBLNModelConfig):
     submodules = ["vision_tower", "language_model"]

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

@@ -337,11 +337,12 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
         chunked_attention_mask = torch.zeros(1, chunked_attention_mask.shape[-1], dtype=torch.float32)
 
         # as gemma3 has different prefill chunk size for image and text, we need to pad the inputs to the max of the two.
-        padding_size = max(self.rbln_config.prefill_chunk_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)
+        if self.rbln_config.use_image_prefill:
+            padding_size = max(self.rbln_config.prefill_chunk_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,
@@ -389,7 +390,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
         step = 0
         while step < query_length:
             # Check if the prefill chunk is an image prefill
-            is_image_prefill = torch.all(
+            is_image_prefill = self.rbln_config.use_image_prefill and torch.all(
                 token_type_ids[:, step : step + self.rbln_config.image_prefill_chunk_size] == 1
             )
             prefill_chunk_size = (
@@ -397,8 +398,10 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
             )
 
             # 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 + prefill_chunk_size] == 1
+            is_text_prefill_with_image_tokens = (
+                self.rbln_config.use_image_prefill
+                and not is_image_prefill
+                and torch.any(token_type_ids[:, step : step + prefill_chunk_size] == 1)
             )
 
             # Check if the prefill chunk crosses a block boundary, requiring padding to align with block boundaries
@@ -418,7 +421,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
             num_processed_tokens = prefill_chunk_size
             if is_text_prefill_with_image_tokens:
                 first_image_token_idx = torch.where(token_type_ids[:, step : step + prefill_chunk_size] == 1)[1][0]
-                num_processed_tokens = first_image_token_idx
+                num_processed_tokens = first_image_token_idx.item()
             if is_last_chunk:
                 num_processed_tokens = query_length - step
 
@@ -548,9 +551,10 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
             dtype=torch.int16,
         ).fill_(-1)
         free_block_pool = deque(x for x in range(self.rbln_config.kvcache_num_blocks))
+
         self.prefill_decoder = RBLNGemma3RuntimeModel(
             runtime=self.model[0],
-            image_prefill=self.model[1],
+            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",
@@ -565,7 +569,7 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         self.decoders = {}
         for i, batch_size in enumerate(self.rbln_config.decoder_batch_sizes):
             self.decoders[batch_size] = RBLNGemma3RuntimeModel(
-                runtime=self.model[i + 2],
+                runtime=self.model[i + self.rbln_config.decoder_runtime_idx],
                 main_input_name=main_input_name,
                 embed_tokens=self.embed_tokens,
                 phase="decode",
@@ -628,20 +632,21 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         if not (rbln_config.use_attention_mask and rbln_config.use_position_ids):
             raise ValueError("use_attention_mask and use_position_ids must be True for RBLNGemma3ForCausalLM")
 
-        # Update image prefill compile config
-        img_prefill_input_info = cls.get_input_info(
-            batch_size=1,
-            query_length=rbln_config.image_prefill_chunk_size,
-            rbln_config=rbln_config,
-            model_config=model_config,
-        )
-        image_prefill_compile_config = RBLNCompileConfig(
-            compiled_model_name="image_prefill", input_info=img_prefill_input_info
-        )
-        # Insert image_prefill compile config at index 1
-        compile_cfgs = rbln_config.compile_cfgs
-        compile_cfgs.insert(1, image_prefill_compile_config)
-        rbln_config.set_compile_cfgs(compile_cfgs)
+        if rbln_config.use_image_prefill:
+            # Update image prefill compile config
+            img_prefill_input_info = cls.get_input_info(
+                batch_size=1,
+                query_length=rbln_config.image_prefill_chunk_size,
+                rbln_config=rbln_config,
+                model_config=model_config,
+            )
+            image_prefill_compile_config = RBLNCompileConfig(
+                compiled_model_name="image_prefill", input_info=img_prefill_input_info
+            )
+            # Insert image_prefill compile config at index 1
+            compile_cfgs = rbln_config.compile_cfgs
+            compile_cfgs.insert(1, image_prefill_compile_config)
+            rbln_config.set_compile_cfgs(compile_cfgs)
 
         return rbln_config
 
@@ -694,23 +699,27 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
             context,
             rbln_config.quantization,
         )
+        compiled_models = {"prefill": compiled_prefill}
 
-        image_prefill_compile_config = rbln_compile_configs[1]
-        image_prefill_example_inputs = image_prefill_compile_config.get_dummy_inputs(
-            fill=0, static_tensors=static_tensors
-        )
-        wrapped_model.phase = "image_prefill"
-        compiled_image_prefill = compile_model(
-            wrapped_model,
-            image_prefill_compile_config,
-            image_prefill_example_inputs,
-            context,
-            rbln_config.quantization,
-        )
+        if rbln_config.use_image_prefill:
+            image_prefill_compile_config = rbln_compile_configs[1]
+            image_prefill_example_inputs = image_prefill_compile_config.get_dummy_inputs(
+                fill=0, static_tensors=static_tensors
+            )
+            wrapped_model.phase = "image_prefill"
+            compiled_image_prefill = compile_model(
+                wrapped_model,
+                image_prefill_compile_config,
+                image_prefill_example_inputs,
+                context,
+                rbln_config.quantization,
+            )
+            compiled_models["image_prefill"] = compiled_image_prefill
 
-        compiled_models = {"prefill": compiled_prefill, "image_prefill": compiled_image_prefill}
         wrapped_model.phase = "decode"
-        for batch_size, dec_compile_config in zip(rbln_config.decoder_batch_sizes, rbln_compile_configs[2:]):
+        for batch_size, dec_compile_config in zip(
+            rbln_config.decoder_batch_sizes, rbln_compile_configs[rbln_config.decoder_runtime_idx :]
+        ):
             dec_example_inputs = dec_compile_config.get_dummy_inputs(fill=0, static_tensors=static_tensors)
             compiled_decoder = compile_model(
                 wrapped_model,
@@ -731,35 +740,45 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     ) -> List[rebel.Runtime]:
         expected_model_names = [
             "prefill",
-            "image_prefill",
             *[f"decoder_batch_{batch_size}" for batch_size in rbln_config.decoder_batch_sizes],
         ]
+        if rbln_config.use_image_prefill:
+            expected_model_names.insert(1, "image_prefill")
+
         if any(model_name not in rbln_config.device_map for model_name in expected_model_names):
             cls._raise_missing_compiled_file_error(expected_model_names)
 
-        return [
+        ret_val = [
             rebel.Runtime(
                 compiled_models[0],
                 tensor_type="pt",
                 device=rbln_config.device_map["prefill"],
                 activate_profiler=rbln_config.activate_profiler,
                 timeout=rbln_config.timeout,
-            ),
-            rebel.Runtime(
-                compiled_models[1],
-                tensor_type="pt",
-                device=rbln_config.device_map["image_prefill"],
-                activate_profiler=rbln_config.activate_profiler,
-                timeout=rbln_config.timeout,
-            ),
-            *[
+            )
+        ]
+        if rbln_config.use_image_prefill:
+            ret_val.append(
+                rebel.Runtime(
+                    compiled_models[1],
+                    tensor_type="pt",
+                    device=rbln_config.device_map["image_prefill"],
+                    activate_profiler=rbln_config.activate_profiler,
+                    timeout=rbln_config.timeout,
+                ),
+            )
+
+        ret_val.extend(
+            [
                 rebel.Runtime(
-                    compiled_models[i + 2],
+                    compiled_models[i + rbln_config.decoder_runtime_idx],
                     tensor_type="pt",
                     device=rbln_config.device_map[f"decoder_batch_{batch_size}"],
                     activate_profiler=rbln_config.activate_profiler,
                     timeout=rbln_config.timeout,
                 )
                 for i, batch_size in enumerate(rbln_config.decoder_batch_sizes)
-            ],
-        ]
+            ]
+        )
+
+        return ret_val

optimum/rbln/transformers/models/llava/__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_llava import RBLNLlavaForConditionalGenerationConfig
+from .modeling_llava import RBLNLlavaForConditionalGeneration

optimum/rbln/transformers/models/llava/configuration_llava.py

@@ -0,0 +1,54 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, Optional
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNLlavaForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNLlavaForConditionalGenerationConfig.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized LLaVA models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
+    submodules = ["vision_tower", "language_model"]
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        language_model: Optional[RBLNModelConfig] = None,
+        **kwargs: Dict[str, Any],
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+            language_model (Optional[RBLNModelConfig]): Configuration for the language model component.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
+        self.vision_tower = vision_tower
+        self.language_model = language_model