optimum-rbln 0.7.4a8__py3-none-any.whl → 0.7.5a0__py3-none-any.whl

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py

@@ -37,6 +37,7 @@ from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
 
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...configurations import RBLNStableDiffusionXLControlNetPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
@@ -47,6 +48,7 @@ logger = logging.get_logger(__name__)
 
 class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusionXLControlNetPipeline):
     original_class = StableDiffusionXLControlNetPipeline
+    _rbln_config_class = RBLNStableDiffusionXLControlNetPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]
 
     # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.py

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py

@@ -37,6 +37,7 @@ from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module
 
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...configurations import RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
@@ -47,6 +48,7 @@ logger = logging.get_logger(__name__)
 
 class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionXLControlNetImg2ImgPipeline):
     original_class = StableDiffusionXLControlNetImg2ImgPipeline
+    _rbln_config_class = RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]
 
     # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl_img2img.py

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py

@@ -14,11 +14,13 @@
 
 from diffusers import KandinskyV22Pipeline
 
+from ...configurations import RBLNKandinskyV22PipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22Pipeline(RBLNDiffusionMixin, KandinskyV22Pipeline):
     original_class = KandinskyV22Pipeline
+    _rbln_config_class = RBLNKandinskyV22PipelineConfig
     _submodules = ["unet", "movq"]
 
     def get_compiled_image_size(self):

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py

@@ -29,6 +29,7 @@ from transformers import (
     CLIPVisionModelWithProjection,
 )
 
+from ...configurations import RBLNKandinskyV22CombinedPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 from .pipeline_kandinsky2_2 import RBLNKandinskyV22Pipeline
 from .pipeline_kandinsky2_2_img2img import RBLNKandinskyV22Img2ImgPipeline
@@ -38,6 +39,7 @@ from .pipeline_kandinsky2_2_prior import RBLNKandinskyV22PriorPipeline
 
 class RBLNKandinskyV22CombinedPipeline(RBLNDiffusionMixin, KandinskyV22CombinedPipeline):
     original_class = KandinskyV22CombinedPipeline
+    _rbln_config_class = RBLNKandinskyV22CombinedPipelineConfig
     _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22Pipeline}
     _submodules = ["prior_image_encoder", "prior_text_encoder", "prior_prior", "unet", "movq"]
     _prefix = {"prior_pipe": "prior_"}

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py

@@ -14,11 +14,13 @@
 
 from diffusers import KandinskyV22Img2ImgPipeline
 
+from ...configurations import RBLNKandinskyV22Img2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22Img2ImgPipeline(RBLNDiffusionMixin, KandinskyV22Img2ImgPipeline):
     original_class = KandinskyV22Img2ImgPipeline
+    _rbln_config_class = RBLNKandinskyV22Img2ImgPipelineConfig
     _submodules = ["unet", "movq"]
 
     def get_compiled_image_size(self):

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpaint.py

@@ -14,11 +14,13 @@
 
 from diffusers import KandinskyV22InpaintPipeline
 
+from ...configurations import RBLNKandinskyV22InpaintPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22InpaintPipeline(RBLNDiffusionMixin, KandinskyV22InpaintPipeline):
     original_class = KandinskyV22InpaintPipeline
+    _rbln_config_class = RBLNKandinskyV22InpaintPipelineConfig
     _submodules = ["unet", "movq"]
 
     def get_compiled_image_size(self):

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py

@@ -14,9 +14,11 @@
 
 from diffusers import KandinskyV22PriorPipeline
 
+from ...configurations import RBLNKandinskyV22PriorPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22PriorPipeline(RBLNDiffusionMixin, KandinskyV22PriorPipeline):
     original_class = KandinskyV22PriorPipeline
+    _rbln_config_class = RBLNKandinskyV22PriorPipelineConfig
     _submodules = ["text_encoder", "image_encoder", "prior"]

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

@@ -15,9 +15,11 @@
 
 from diffusers import StableDiffusionPipeline
 
+from ...configurations import RBLNStableDiffusionPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionPipeline(RBLNDiffusionMixin, StableDiffusionPipeline):
     original_class = StableDiffusionPipeline
+    _rbln_config_class = RBLNStableDiffusionPipelineConfig
     _submodules = ["vae", "text_encoder", "unet"]

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusionImg2ImgPipeline
 
+from ...configurations import RBLNStableDiffusionImg2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionImg2ImgPipeline):
     original_class = StableDiffusionImg2ImgPipeline
+    _rbln_config_class = RBLNStableDiffusionImg2ImgPipelineConfig
     _submodules = ["text_encoder", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusionInpaintPipeline
 
+from ...configurations import RBLNStableDiffusionInpaintPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionInpaintPipeline(RBLNDiffusionMixin, StableDiffusionInpaintPipeline):
     original_class = StableDiffusionInpaintPipeline
+    _rbln_config_class = RBLNStableDiffusionInpaintPipelineConfig
     _submodules = ["text_encoder", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusion3Pipeline
 
+from ...configurations import RBLNStableDiffusion3PipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3Pipeline(RBLNDiffusionMixin, StableDiffusion3Pipeline):
     original_class = StableDiffusion3Pipeline
+    _rbln_config_class = RBLNStableDiffusion3PipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusion3Img2ImgPipeline
 
+from ...configurations import RBLNStableDiffusion3Img2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3Img2ImgPipeline(RBLNDiffusionMixin, StableDiffusion3Img2ImgPipeline):
     original_class = StableDiffusion3Img2ImgPipeline
+    _rbln_config_class = RBLNStableDiffusion3Img2ImgPipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusion3InpaintPipeline
 
+from ...configurations import RBLNStableDiffusion3InpaintPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3InpaintPipeline(RBLNDiffusionMixin, StableDiffusion3InpaintPipeline):
     original_class = StableDiffusion3InpaintPipeline
+    _rbln_config_class = RBLNStableDiffusion3InpaintPipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusionXLPipeline
 
+from ...configurations import RBLNStableDiffusionXLPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLPipeline(RBLNDiffusionMixin, StableDiffusionXLPipeline):
     original_class = StableDiffusionXLPipeline
+    _rbln_config_class = RBLNStableDiffusionXLPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusionXLImg2ImgPipeline
 
+from ...configurations import RBLNStableDiffusionXLImg2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionXLImg2ImgPipeline):
     original_class = StableDiffusionXLImg2ImgPipeline
+    _rbln_config_class = RBLNStableDiffusionXLImg2ImgPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusionXLInpaintPipeline
 
+from ...configurations import RBLNStableDiffusionXLInpaintPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLInpaintPipeline(RBLNDiffusionMixin, StableDiffusionXLInpaintPipeline):
     original_class = StableDiffusionXLInpaintPipeline
+    _rbln_config_class = RBLNStableDiffusionXLInpaintPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

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, Optional
+from typing import Any, Dict, List, Optional
 
 import rebel
 
@@ -37,6 +37,7 @@ class RBLNDecoderOnlyModelForCausalLMConfig(RBLNModelConfig):
         quantization: Optional[Dict[str, Any]] = None,
         prefill_chunk_size: Optional[int] = None,
         kvcache_num_blocks: Optional[int] = None,
+        decoder_batch_sizes: Optional[List[int]] = None,
         **kwargs,
     ):
         """
@@ -53,6 +54,13 @@ class RBLNDecoderOnlyModelForCausalLMConfig(RBLNModelConfig):
             prefill_chunk_size (Optional[int]): The chunk size for prefilling the KV cache. Defaults to 128,
                 and must be a positive integer divisible by 64.
             kvcache_num_blocks (Optional[int]): The number of blocks in the KV cache.
+            decoder_batch_sizes (Optional[List[int]]): A list of batch sizes for which separate decoder models will be compiled.
+                This allows the model to handle varying batch sizes efficiently during generation. If not specified,
+                defaults to a list containing only the model's main batch size. When specifying multiple batch sizes:
+                1) All values must be less than or equal to the main batch size.
+                2) The list will be sorted in descending order (larger batch sizes first).
+                3) If using multiple decoders, at least one batch size should match the main batch size.
+
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
@@ -88,3 +96,25 @@ class RBLNDecoderOnlyModelForCausalLMConfig(RBLNModelConfig):
             raise ValueError("`prefill_chunk_size` must be a positive integer divisible by 64.")
 
         self.kvcache_num_blocks = kvcache_num_blocks
+        self.decoder_batch_sizes = decoder_batch_sizes
+        if self.decoder_batch_sizes is None:
+            self.decoder_batch_sizes = [self.batch_size]
+
+        if self.use_multiple_decoder:
+            if max(self.decoder_batch_sizes) > self.batch_size:
+                raise ValueError(
+                    f"Decoder batch size ({max(self.decoder_batch_sizes)}) must be less than or equal to the runtime batch size ({self.batch_size})."
+                )
+            if max(self.decoder_batch_sizes) < self.batch_size:
+                logger.warning(
+                    f"Maximum decoder batch size ({max(self.decoder_batch_sizes)}) is less than the model's batch size ({self.batch_size}). "
+                    "Appending the model's batch size to the decoder batch size."
+                )
+                self.decoder_batch_sizes.append(self.batch_size)
+
+            # Larger batch size should be at the beginning of the list.
+            self.decoder_batch_sizes.sort(reverse=True)
+
+    @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/modeling_decoderonly.py

@@ -229,6 +229,12 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
 
             attention_mask = self.dec_attn_mask
 
+        if self.batch_size < block_tables.shape[0]:
+            block_tables = block_tables[: self.batch_size]
+
+        if self.batch_size < attention_mask.shape[0]:
+            attention_mask = attention_mask[: self.batch_size]
+
         logits = super().forward(
             inputs,
             cache_position,
@@ -417,19 +423,24 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             use_attention_mask=self.rbln_config.use_attention_mask,
             attn_impl=self.rbln_config.attn_impl,
         )
-        self.decoder = RBLNRuntimeModel(
-            runtime=self.model[1],
-            main_input_name=main_input_name,
-            embed_tokens=self.embed_tokens,
-            phase="decode",
-            batch_size=self.rbln_config.batch_size,
-            dec_attn_mask=dec_attn_mask,
-            block_tables=block_tables,
-            free_block_pool=free_block_pool,
-            kvcache_block_size=self.rbln_config.kvcache_block_size,
-            use_attention_mask=self.rbln_config.use_attention_mask,
-            attn_impl=self.rbln_config.attn_impl,
-        )
+        self.decoders = {}
+        for i, batch_size in enumerate(self.rbln_config.decoder_batch_sizes):
+            self.decoders[batch_size] = RBLNRuntimeModel(
+                runtime=self.model[i + 1],
+                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,
+                kvcache_block_size=self.rbln_config.kvcache_block_size,
+                use_attention_mask=self.rbln_config.use_attention_mask,
+                attn_impl=self.rbln_config.attn_impl,
+            )
+
+        # NOTE(eunji): Use a decoder whose batch size matches the model's main batch size for compatibility.
+        self.decoder = self.decoders[self.rbln_config.batch_size]
 
     @classmethod
     def save_torch_artifacts(
@@ -547,7 +558,6 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
         rbln_compile_configs = rbln_config.compile_cfgs
         prefill_compile_config = rbln_compile_configs[0]
-        dec_compile_config = rbln_compile_configs[1]
 
         context = CompileContext(use_weight_sharing=True)
 
@@ -562,33 +572,42 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
                 static_tensors[name] = tensor
                 context.mark_static_address(tensor)
 
-        dec_example_inputs = dec_compile_config.get_dummy_inputs(fill=0, static_tensors=static_tensors)
-
         @QuantizationManager.with_quantization_env
-        def compile_model(*args, **kwargs):
+        def compile_model(wrapped_model, compile_config, example_inputs, compile_context, **kwargs):
             try:
                 original_linear = torch.nn.functional.linear
                 torch.nn.functional.linear = torch.ops.rbln_custom_ops.linear
-                wrapped_model.phase = "prefill"
-                compiled_prefill = RBLNModel.compile(
+                compiled_model = RBLNModel.compile(
                     wrapped_model,
-                    prefill_compile_config,
-                    example_inputs=prefill_example_inputs,
-                    compile_context=context,
+                    compile_config,
+                    example_inputs=example_inputs,
+                    compile_context=compile_context,
                 )
-
-                wrapped_model.phase = "decode"
-                compiled_decoder = RBLNModel.compile(
-                    wrapped_model,
-                    dec_compile_config,
-                    example_inputs=dec_example_inputs,
-                    compile_context=context,
-                )
-                return {"prefill": compiled_prefill, "decoder": compiled_decoder}
+                return compiled_model
             finally:
                 torch.nn.functional.linear = original_linear
 
-        compiled_models = compile_model(quantize_config=rbln_config.quantization)
+        wrapped_model.phase = "prefill"
+        compiled_prefill = compile_model(
+            wrapped_model,
+            prefill_compile_config,
+            prefill_example_inputs,
+            context,
+            quantize_config=rbln_config.quantization,
+        )
+
+        wrapped_model.phase = "decode"
+        compiled_models = {"prefill": compiled_prefill}
+        for batch_size, dec_compile_config in zip(rbln_config.decoder_batch_sizes, rbln_compile_configs[1:]):
+            dec_example_inputs = dec_compile_config.get_dummy_inputs(fill=0, static_tensors=static_tensors)
+            compiled_decoder = compile_model(
+                wrapped_model,
+                dec_compile_config,
+                dec_example_inputs,
+                context,
+                quantize_config=rbln_config.quantization,
+            )
+            compiled_models[f"decoder_batch_{batch_size}"] = compiled_decoder
 
         # check if the memory is enough to have additional blocks
         required_num_blocks = (rbln_config.max_seq_len // rbln_config.kvcache_block_size) * rbln_config.batch_size
@@ -613,8 +632,11 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         alloc_memory_by_key: Dict[str, int] = {
             key: sum(memory_per_node) for key, memory_per_node in alloc_memory_per_node_by_key.items()
         }
-        for key, memory_per_node in compiled_models["decoder"].get_alloc_per_node_by_key().items():
-            alloc_memory_by_key[key] += sum(memory_per_node)
+        for batch_size in rbln_config.decoder_batch_sizes:
+            for key, memory_per_node in (
+                compiled_models[f"decoder_batch_{batch_size}"].get_alloc_per_node_by_key().items()
+            ):
+                alloc_memory_by_key[key] += sum(memory_per_node)
         alloc_memory_by_key.pop("PortRecur")  # kv-cache
         kernel_size = alloc_memory_by_key.pop("Kernel")  # model weight
 
@@ -650,6 +672,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         n_model_params: Optional[int] = None,
         kernel_size: Optional[int] = None,
         buffer: Optional[int] = None,
+        num_runtimes: int = 2,
     ) -> int:
         """
         We are finding max_n_blocks(x) that satisfies the following equation:
@@ -721,7 +744,8 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
         if buffer is None:
             # TODO: Accurate buffer estimation
-            buffer_per_core = 2**29  # 500MB per npu
+            buffer_per_runtime_per_core = 2**28  # 256MB per runtime
+            buffer_per_core = buffer_per_runtime_per_core * num_runtimes  # 1 for prefill, 1 for decoder
             buffer = buffer_per_core * tensor_parallel_size
         available_dram -= buffer
 
@@ -839,6 +863,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
                 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 + len(rbln_config.decoder_batch_sizes),
             )
 
             max_num_blocks = min(max_num_blocks, estimated_max_num_blocks)
@@ -881,24 +906,28 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             hidden_size=hidden_size,
             head_dim=head_dim,
         )
-        dec_input_info = cls.get_input_info(
-            batch_size=rbln_config.batch_size,
-            query_length=1,
-            use_inputs_embeds=rbln_config.use_inputs_embeds,
-            use_attention_mask=rbln_config.use_attention_mask,
-            max_seq_len=rbln_config.max_seq_len,
-            kvcache_block_size=rbln_config.kvcache_block_size,
-            kvcache_num_blocks=rbln_config.kvcache_num_blocks,
-            num_key_value_heads=num_key_value_heads,
-            num_hidden_layers=num_hidden_layers,
-            hidden_size=hidden_size,
-            head_dim=head_dim,
-        )
 
         prefill_compile_config = RBLNCompileConfig(compiled_model_name="prefill", input_info=prefill_input_info)
-        dec_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
 
-        rbln_config.set_compile_cfgs([prefill_compile_config, dec_compile_config])
+        dec_compile_configs = []
+        for batch_size in rbln_config.decoder_batch_sizes:
+            dec_input_info = cls.get_input_info(
+                batch_size=batch_size,
+                query_length=1,
+                use_inputs_embeds=rbln_config.use_inputs_embeds,
+                use_attention_mask=rbln_config.use_attention_mask,
+                max_seq_len=rbln_config.max_seq_len,
+                kvcache_block_size=rbln_config.kvcache_block_size,
+                kvcache_num_blocks=rbln_config.kvcache_num_blocks,
+                num_key_value_heads=num_key_value_heads,
+                num_hidden_layers=num_hidden_layers,
+                hidden_size=hidden_size,
+                head_dim=head_dim,
+            )
+            dec_compile_configs.append(
+                RBLNCompileConfig(compiled_model_name=f"decoder_batch_{batch_size}", input_info=dec_input_info)
+            )
+        rbln_config.set_compile_cfgs([prefill_compile_config, *dec_compile_configs])
 
         return rbln_config
 
@@ -908,8 +937,12 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         compiled_models: List[rebel.RBLNCompiledModel],
         rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
     ) -> List[rebel.Runtime]:
-        if any(model_name not in rbln_config.device_map for model_name in ["prefill", "decoder"]):
-            cls._raise_missing_compiled_file_error(["prefill", "decoder"])
+        expected_model_names = [
+            "prefill",
+            *[f"decoder_batch_{batch_size}" for batch_size in rbln_config.decoder_batch_sizes],
+        ]
+        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 [
             rebel.Runtime(
@@ -918,12 +951,15 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
                 device=rbln_config.device_map["prefill"],
                 activate_profiler=rbln_config.activate_profiler,
             ),
-            rebel.Runtime(
-                compiled_models[1],
-                tensor_type="pt",
-                device=rbln_config.device_map["decoder"],
-                activate_profiler=rbln_config.activate_profiler,
-            ),
+            *[
+                rebel.Runtime(
+                    compiled_models[i + 1],
+                    tensor_type="pt",
+                    device=rbln_config.device_map[f"decoder_batch_{batch_size}"],
+                    activate_profiler=rbln_config.activate_profiler,
+                )
+                for i, batch_size in enumerate(rbln_config.decoder_batch_sizes)
+            ],
         ]
 
     def get_decoder(self):
@@ -1024,7 +1060,15 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             logits = torch.cat(logits, dim=0)
         # Decoder
         else:
-            logits = self.decoder(
+            inputs = inputs_embeds if inputs_embeds is not None else input_ids
+            batch_size = inputs.shape[0]
+            if batch_size not in self.decoders:
+                raise ValueError(
+                    f"No decoder runtime available for batch size {batch_size}. "
+                    f"Available batch sizes are: {list(self.decoders.keys())}. "
+                    f"Please run your model with one of these batch sizes or add support for batch size {batch_size}."
+                )
+            logits = self.decoders[batch_size](
                 input_ids=input_ids,
                 inputs_embeds=inputs_embeds,
                 cache_position=cache_position,

optimum/rbln/utils/import_utils.py

@@ -144,10 +144,27 @@ def check_version_compats() -> None:
         except importlib.metadata.PackageNotFoundError:
             warnings.warn(f"optimum-rbln requires {compat.package_name} to be installed.", ImportWarning)
             continue
+        # For versions 0.7.2 and above, don't show warning for rebel-compiler if base versions match
 
-        if not Version(compat.min_version) <= Version(dep_version) < Version(compat.max_version):
-            warnings.warn(
-                f"optimum-rbln v{my_version} is compatible to {compat.package_name} v{compat.min_version} to v{compat.max_version}. (you are currently using v{dep_version})\n"
-                "Please refer to our SDK release notes at https://docs.rbln.ai/about_atom/release_note.html",
-                ImportWarning,
-            )
+        if compat.package_name == "rebel-compiler":
+            # For optimum-rbln versions 0.7.2 and above, suppress the warning if the base versions of
+            # optimum-rbln and rebel-compiler match (e.g., 0.7.x with 0.7.y).
+            if (
+                Version(my_version) >= Version("0.7.2")
+                and Version(my_version).base_version == Version(dep_version).base_version
+            ):
+                continue
+            else:
+                warnings.warn(
+                    f"Version mismatch detected: optimum-rbln v{my_version} and {compat.package_name} v{dep_version} have different base versions. "
+                    f"For optimal performance and compatibility, please ensure both packages share the same major and minor version numbers. "
+                    "Please refer to our SDK release notes at https://docs.rbln.ai/about_atom/release_note.html",
+                    ImportWarning,
+                )
+        else:
+            if not Version(compat.min_version) <= Version(dep_version) < Version(compat.max_version):
+                warnings.warn(
+                    f"optimum-rbln v{my_version} is compatible to {compat.package_name} v{compat.min_version} to v{compat.max_version}. (you are currently using v{dep_version})\n"
+                    "Please refer to our SDK release notes at https://docs.rbln.ai/about_atom/release_note.html",
+                    ImportWarning,
+                )

{optimum_rbln-0.7.4a8.dist-info → optimum_rbln-0.7.5a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.7.4a8
+Version: 0.7.5a0
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai