optimum-rbln 0.9.5a4__py3-none-any.whl → 0.10.0.post1__py3-none-any.whl

optimum/rbln/transformers/models/gpt_oss/configuration_gpt_oss.py

@@ -17,7 +17,7 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 class RBLNGptOssForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     """
-    Configuration class for RBLN GPT-OSS models.
+    Configuration class for RBLN GptOss models.
 
     This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
 
@@ -28,14 +28,15 @@ class RBLNGptOssForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     # Create a configuration object
     config = RBLNGptOssForCausalLMConfig(
         batch_size=1,
-        tensor_parallel_size=4
+        tensor_parallel_size=8,
+        kvcache_partition_len=8192,
     )
 
     # Use the configuration with from_pretrained
     model = RBLNGptOssForCausalLM.from_pretrained(
         "openai/gpt-oss-20b",
         export=True,
-        rbln_config=config
+        rbln_config=config,
     )
     ```
     """

optimum/rbln/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -41,8 +41,8 @@ class RBLNGptOssForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     The GPT-OSS 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 GPT-OSSForCausalLM model on RBLN devices.
-    It implements the methods to convert a pre-trained transformers GPT-OSSForCausalLM model into a RBLN transformer model by:
+    A class to convert and run pre-trained transformers based GptOssForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers GptOssForCausalLM 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.
 
@@ -62,19 +62,21 @@ class RBLNGptOssForCausalLM(RBLNDecoderOnlyModelForCausalLM):
             "openai/gpt-oss-20b",
             export=True,
             rbln_batch_size=1,
-            rbln_tensor_parallel_size=4,
+            rbln_tensor_parallel_size=8,
+            rbln_kvcache_partition_len=8192,
         )
 
 
         # Using a config dictionary
         rbln_config = {
             "batch_size": 1,
-            "tensor_parallel_size": 4,
+            "tensor_parallel_size": 8,
+            "kvcache_partition_len": 8192,
         }
         model = RBLNGptOssForCausalLM.from_pretrained(
             "openai/gpt-oss-20b",
             export=True,
-            rbln_config=rbln_config
+            rbln_config=rbln_config,
         )
 
 
@@ -83,12 +85,13 @@ class RBLNGptOssForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
         config = RBLNGptOssForCausalLMConfig(
             batch_size=1,
-            tensor_parallel_size=4
+            tensor_parallel_size=8,
+            kvcache_partition_len=8192,
         )
         model = RBLNGptOssForCausalLM.from_pretrained(
             "openai/gpt-oss-20b",
             export=True,
-            rbln_config=config
+            rbln_config=config,
         )
         ```
     """

optimum/rbln/transformers/models/mixtral/__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_mixtral import RBLNMixtralForCausalLMConfig
+from .modeling_mixtral import RBLNMixtralForCausalLM

optimum/rbln/transformers/models/mixtral/configuration_mixtral.py

@@ -0,0 +1,38 @@
+# Copyright 2026 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 RBLNMixtralForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    """
+    Configuration class for RBLN Mixtral models.
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNMixtralForCausalLM, RBLNMixtralForCausalLMConfig
+    # Create a configuration object
+    config = RBLNMixtralForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=32768,
+        tensor_parallel_size=4
+    )
+    # Use the configuration with from_pretrained
+    model = RBLNMixtralForCausalLM.from_pretrained(
+        "mistralai/Mixtral-8x7B-Instruct-v0.1",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

optimum/rbln/transformers/models/mixtral/mixtral_architecture.py

@@ -0,0 +1,76 @@
+# Copyright 2026 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
+
+import torch
+from torch import nn
+
+from ..decoderonly.configuration_decoderonly import RBLNLoRAConfig
+from ..decoderonly.decoderonly_architecture import DecoderOnlyAttention, DecoderOnlyLayer, DecoderOnlyWrapper
+
+
+class MixtralWrapper(DecoderOnlyWrapper):
+    def get_rbln_layer_class(self):
+        return MixtralLayer
+
+
+class MixtralLayer(DecoderOnlyLayer):
+    _MLP_ATTR = ("block_sparse_moe",)
+
+    def __init__(self, layer, self_attn: DecoderOnlyAttention, lora_config: Optional[RBLNLoRAConfig] = None):
+        super().__init__(layer, self_attn, lora_config)
+        self.block_sparse_moe = MixtralSparseMoeBlock(self.mlp)
+
+    def get_mlp(self) -> nn.Module:
+        return self.block_sparse_moe
+
+
+class MixtralSparseMoeBlock(nn.Module):
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        # self.num_experts = model.num_experts
+        self.top_k = model.top_k
+        self.gate = model.gate
+        self.experts = MixtralBlockSparseTop2MLP(model.experts, self.top_k)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        # router_logits: (batch * sequence_length, n_experts)
+        router_logits = self.gate(hidden_states)
+        final_hidden_states = self.experts(hidden_states, router_logits)
+        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+        return final_hidden_states
+
+
+class MixtralBlockSparseTop2MLP(nn.Module):
+    def __init__(self, expert_list, top_k):
+        super().__init__()
+        self.hidden_dim = expert_list[0].hidden_dim
+        self.ffn_dim = expert_list[0].ffn_dim
+        self.top_k = top_k
+
+        self.num_experts = len(expert_list)
+        self.w1 = nn.Linear(self.hidden_dim, self.num_experts * self.ffn_dim, bias=False)
+        self.w2 = nn.Linear(self.num_experts * self.ffn_dim, self.hidden_dim, bias=False)
+        self.w3 = nn.Linear(self.hidden_dim, self.num_experts * self.ffn_dim, bias=False)
+        self.w1.weight.data = torch.stack([expert.w1.weight.data for expert in expert_list], dim=0)
+        self.w2.weight.data = torch.stack([expert.w2.weight.data for expert in expert_list], dim=0)
+        self.w3.weight.data = torch.stack([expert.w3.weight.data for expert in expert_list], dim=0)
+
+    def forward(self, x, router_logits):
+        return torch.ops.rbln_custom_ops.custom_moe_glu(
+            x, self.w1.weight, self.w3.weight, self.w2.weight, router_logits, self.top_k, True
+        )

optimum/rbln/transformers/models/mixtral/modeling_mixtral.py

@@ -0,0 +1,68 @@
+# Copyright 2026 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 ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from .mixtral_architecture import MixtralWrapper
+
+
+class RBLNMixtralForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The Mixtral is a Mixture-of-Experts (MoE) variant of Mixtral, available as a base model and an aligned chat model.
+    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 MixtralForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers MixtralForCausalLM 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 [`RBLNMixtralForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNMixtralForCausalLMConfig`] or a dictionary conforming to its structure.
+    See the [`RBLNMixtralForCausalLMConfig`] class for all available configuration options.
+    Examples:
+        ```python
+        from optimum.rbln import RBLNMixtralForCausalLM
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNMixtralForCausalLM.from_pretrained(
+            "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 32768,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNMixtralForCausalLM.from_pretrained(
+            "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            export=True,
+            rbln_config=rbln_config
+        )
+        # Using a RBLNMixtralForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNMixtralForCausalLMConfig
+        config = RBLNMixtralForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=32768,
+            tensor_parallel_size=4
+        )
+        model = RBLNMixtralForCausalLM.from_pretrained(
+            "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    _decoder_wrapper_cls = MixtralWrapper

optimum/rbln/transformers/models/paligemma/modeling_paligemma.py

@@ -525,13 +525,13 @@ class RBLNPaliGemmaModel(RBLNModel):
         Forward pass for the RBLN-optimized PaliGemmaModel model.
 
         Args:
-            input_ids (torch.LongTensor of shape (batch_size, sequence_length)) — Indices of input sequence tokens in the vocabulary.
-            pixel_values (torch.Tensor of shape (batch_size, num_channels, image_size, image_size)) — The tensors corresponding to the input images.
-            attention_mask (torch.Tensor of shape (batch_size, sequence_length)) — Mask to avoid performing attention on padding token indices.
-            position_ids (torch.LongTensor of shape (batch_size, sequence_length)) — Indices of positions of each input sequence tokens in the position embeddings.
-            token_type_ids (torch.LongTensor of shape (batch_size, sequence_length)) — Segment token indices to indicate first and second portions of the inputs.
-            output_hidden_states (bool, optional) — Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for more detail.
-            return_dict (bool, optional) — Whether or not to return a ModelOutput instead of a plain tuple.
+            input_ids (torch.LongTensor of shape (batch_size, sequence_length)): Indices of input sequence tokens in the vocabulary.
+            pixel_values (torch.Tensor of shape (batch_size, num_channels, image_size, image_size)): The tensors corresponding to the input images.
+            attention_mask (torch.Tensor of shape (batch_size, sequence_length)): Mask to avoid performing attention on padding token indices.
+            position_ids (torch.LongTensor of shape (batch_size, sequence_length)): Indices of positions of each input sequence tokens in the position embeddings.
+            token_type_ids (torch.LongTensor of shape (batch_size, sequence_length)): Segment token indices to indicate first and second portions of the inputs.
+            output_hidden_states (bool, optional): Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for more detail.
+            return_dict (bool, optional): Whether or not to return a ModelOutput instead of a plain tuple.
 
         Returns:
             PaligemmaModelOutputWithPast or tuple(torch.FloatTensor)

optimum/rbln/transformers/models/pixtral/modeling_pixtral.py

@@ -297,13 +297,17 @@ class RBLNPixtralVisionModel(RBLNModel):
         Forward pass for the RBLN-optimized Pixtral vision model.
 
         Args:
-            pixel_values (torch.Tensor of shape (batch_size, num_channels, image_size, image_size)) — The tensors corresponding to the input images. Pixel values can be obtained using PixtralImageProcessor. See PixtralImageProcessor.call() for details (PixtralProcessor uses PixtralImageProcessor for processing images).
-            image_sizes (torch.Tensor of shape (batch_size, 2), optional) — The sizes of the images in the batch, being (height, width) for each image.
-            output_hidden_states (bool, optional) — Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for more detail.
-            return_dict (bool, optional) — Whether or not to return a ModelOutput instead of a plain tuple.
+            pixel_values: Input images as a tensor of shape (batch_size, num_channels, image_size, image_size).
+                Pixel values can be obtained using PixtralImageProcessor. See PixtralImageProcessor.__call__()
+                for details (PixtralProcessor uses PixtralImageProcessor for processing images).
+            image_sizes: The sizes of the images in the batch as a tensor of shape (batch_size, 2),
+                being (height, width) for each image. Optional.
+            output_hidden_states: Whether or not to return the hidden states of all layers. Optional.
+                See hidden_states under returned tensors for more detail.
+            return_dict: Whether or not to return a ModelOutput instead of a plain tuple. Optional.
 
         Returns:
-            BaseModelOutput or tuple(torch.FloatTensor)
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a BaseModelOutput object.
         """
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states

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

@@ -61,6 +61,8 @@ class RBLNQwen2_5_VLModelConfig(RBLNDecoderOnlyModelConfig):
     Configuration class for RBLNQwen2_5_VLModel.
     """
 
+    submodules = ["visual"]
+
     def __init__(self, visual: Optional[RBLNModelConfig] = None, **kwargs: Any):
         super().__init__(**kwargs)
         self.visual = self.initialize_submodule_config(submodule_config=visual)

optimum/rbln/transformers/models/qwen2_vl/configuration_qwen2_vl.py

@@ -53,6 +53,8 @@ class RBLNQwen2VLModelConfig(RBLNDecoderOnlyModelConfig):
     Configuration class for RBLNQwen2VLModel.
     """
 
+    submodules = ["visual"]
+
     def __init__(self, visual: Optional[RBLNModelConfig] = None, **kwargs: Dict[str, Any]):
         super().__init__(**kwargs)
         self.visual = self.initialize_submodule_config(submodule_config=visual)

optimum/rbln/transformers/models/resnet/configuration_resnet.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from typing import Optional
+from typing import Any, Optional, Tuple, Union
 
 from ...configuration_generic import RBLNModelForImageClassificationConfig
 
@@ -26,17 +26,23 @@ class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConf
     RBLN-optimized ResNet models for image classification tasks.
     """
 
-    def __init__(self, output_hidden_states: Optional[bool] = None, **kwargs):
+    def __init__(
+        self,
+        image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        batch_size: Optional[int] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs: Any,
+    ):
         """
         Args:
             image_size (Optional[Union[int, Tuple[int, int]]]): The size of input images.
                 Can be an integer for square images or a tuple (height, width).
             batch_size (Optional[int]): The batch size for inference. Defaults to 1.
-            output_hidden_states (bool, optional) — Whether or not to return the hidden states of all layers.
+            output_hidden_states (bool, optional): Whether or not to return the hidden states of all layers.
             kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.
         """
-        super().__init__(**kwargs)
+        super().__init__(image_size=image_size, batch_size=batch_size, **kwargs)
         self.output_hidden_states = output_hidden_states

optimum/rbln/transformers/models/whisper/generation_whisper.py

@@ -51,22 +51,22 @@ class RBLNWhisperGenerationMixin(WhisperGenerationMixin, GenerationMixin):
         return_segments: Optional[bool] = None,
         return_timestamps: Optional[bool] = None,
         return_token_timestamps: Optional[bool] = None,
-        **kwargs,
+        **kwargs: Optional[Dict[str, Any]],
     ) -> Union[ModelOutput, Dict[str, Any], torch.LongTensor]:
         """
         The generate function is utilized in its standard form as in the HuggingFace transformers library. User can use this function to generate text from the model.
         Check the [HuggingFace transformers documentation](https://huggingface.co/docs/transformers/v4.57.1/en/model_doc/whisper#transformers.WhisperForConditionalGeneration.generate) for more details.
 
         Args:
-            input_features(torch.Tensor, optional): The input features to the model.
-            attention_mask(torch.Tensor, optional): Attention mask needs to be passed when doing long-form transcription using a batch size > 1.
-            generation_config(GenerationConfig, optional): The generation configuration to be used as base parametrization for the generation call. **kwargs passed to generate matching the attributes of generation_config will override them.
+            input_features (torch.Tensor, optional): The input features to the model.
+            attention_mask (torch.Tensor, optional): Attention mask needs to be passed when doing long-form transcription using a batch size > 1.
+            generation_config (GenerationConfig, optional): The generation configuration to be used as base parametrization for the generation call. **kwargs passed to generate matching the attributes of generation_config will override them.
                 If generation_config is not provided, the default will be used, which had the following loading priority: 1) from the generation_config.json model file, if it exists; 2) from the model configuration.
                 Please note that unspecified parameters will inherit [GenerationConfig](https://huggingface.co/docs/transformers/v4.57.1/en/main_classes/text_generation#transformers.GenerationConfig)’s default values.
-            return_segments(bool, optional): Whether to return segments.
-            return_timestamps(bool, optional): Whether to return the timestamps with the text. For audios longer than 30 seconds, it is necessary to set return_timestamps=True.
-            return_token_timestamps(bool, optional): Whether to return token timestamps.
-            kwargs(dict[str, Any], optional): Additional arguments passed to the generate function.
+            return_segments (bool, optional): Whether to return segments.
+            return_timestamps (bool, optional): Whether to return the timestamps with the text. For audios longer than 30 seconds, it is necessary to set return_timestamps=True.
+            return_token_timestamps (bool, optional): Whether to return token timestamps.
+            kwargs (dict[str, Any], optional): Additional arguments passed to the generate function.
 
         Returns:
             Transcribes or translates log-mel input features to a sequence of auto-regressively generated token ids.

optimum/rbln/utils/deprecation.py

@@ -194,7 +194,7 @@ def deprecate_kwarg(
                 message = f"{message} {additional_message}"
 
             # update minimum_action if argument is ALREADY deprecated (current version >= deprecated version)
-            if is_greater_or_equal_version:
+            if is_greater_or_equal_version and message is not None:
                 # change to NOTIFY -> RAISE  in case we want to raise error for already deprecated arguments
                 if raise_if_greater_or_equal_version:
                     minimum_action = Action.RAISE
@@ -211,3 +211,80 @@ def deprecate_kwarg(
         return wrapped_func
 
     return wrapper
+
+
+def deprecate_method(
+    version: str,
+    new_method: Optional[str] = None,
+    raise_if_greater_or_equal_version: bool = True,
+    additional_message: Optional[str] = None,
+):
+    """
+    Decorator to mark a method as deprecated, optionally pointing to a replacement method.
+    This decorator allows you to:
+    - Notify users when a method is deprecated.
+    - Optionally specify a new method name that should be used instead.
+    - Raise an error if the deprecated method is called after the specified version.
+    Parameters:
+        version (`str`):
+            The version in which the method was (or will be) deprecated.
+        new_method (`Optional[str]`, *optional*):
+            The name of the new method to use instead. If specified, users will be directed to use this method.
+        raise_if_greater_or_equal_version (`bool`, *optional*, defaults to `True`):
+            Whether to raise `ValueError` if current `optimum.rbln` version is greater than or equal to the deprecated version.
+        additional_message (`Optional[str]`, *optional*):
+            An additional message to append to the default deprecation message.
+    Returns:
+        Callable:
+            A wrapped function that handles the deprecation warning or error.
+    Examples:
+        >>> class MyClass:
+        ...     @deprecate_method(version="0.12.0", new_method="from_pretrained")
+        ...     def load(self, path):
+        ...         return self.from_pretrained(path)
+    """
+
+    deprecated_version = packaging.version.parse(version)
+    current_version = packaging.version.parse(__version__)
+    is_greater_or_equal_version = current_version >= deprecated_version
+
+    if is_greater_or_equal_version:
+        version_message = f"and removed starting from version {version}"
+    else:
+        version_message = f"and will be removed in version {version}"
+
+    def wrapper(func):
+        sig = inspect.signature(func)
+        function_named_args = set(sig.parameters.keys())
+        is_instance_method = "self" in function_named_args
+        is_class_method = "cls" in function_named_args
+
+        @wraps(func)
+        def wrapped_func(*args, **kwargs):
+            # Get class + method name for better warning message
+            method_name = func.__name__
+            if is_instance_method:
+                method_name = f"{args[0].__class__.__name__}.{method_name}"
+            elif is_class_method:
+                method_name = f"{args[0].__name__}.{method_name}"
+
+            # Build deprecation message
+            if new_method is not None:
+                message = f"`{method_name}` is deprecated {version_message}. Use `{new_method}` instead."
+            else:
+                message = f"`{method_name}` is deprecated {version_message}."
+
+            if additional_message is not None:
+                message = f"{message} {additional_message}"
+
+            # Determine action based on version
+            if is_greater_or_equal_version and raise_if_greater_or_equal_version:
+                raise ValueError(message)
+            else:
+                logger.warning(message, stacklevel=2)
+
+            return func(*args, **kwargs)
+
+        return wrapped_func
+
+    return wrapper

optimum/rbln/utils/hub.py

@@ -16,7 +16,8 @@ import json
 from pathlib import Path
 from typing import List, Optional, Union
 
-from huggingface_hub import HfApi, get_token, hf_hub_download
+from huggingface_hub import HfApi, get_token, hf_hub_download, try_to_load_from_cache
+from huggingface_hub.errors import LocalEntryNotFoundError
 
 
 def pull_compiled_model_from_hub(
@@ -29,6 +30,97 @@ def pull_compiled_model_from_hub(
     local_files_only: bool,
 ) -> Path:
     """Pull model files from the HuggingFace Hub."""
+    config_filename = "rbln_config.json" if subfolder == "" else f"{subfolder}/rbln_config.json"
+
+    # Try to find config file in cache first.
+    config_cache_path = try_to_load_from_cache(
+        repo_id=str(model_id),
+        filename=config_filename,
+        revision=revision,
+        cache_dir=cache_dir,
+    )
+
+    # If config is cached and we're not forcing download, try to use cached files
+    if config_cache_path and isinstance(config_cache_path, str) and not force_download:
+        config_path = Path(config_cache_path)
+        if config_path.exists():
+            cache_dir_path = config_path.parent
+
+            # Look for .rbln files in the same directory
+            pattern_rbln = "*.rbln"
+            rbln_files = list(cache_dir_path.glob(pattern_rbln))
+
+            # Validate files found in cache
+            rbln_config_filenames = [config_path] if config_path.exists() else []
+            validate_files(rbln_files, rbln_config_filenames, f"cached repository {model_id}")
+
+            # If local_files_only is True, return cached directory without API call
+            if local_files_only:
+                return cache_dir_path
+
+            # If local_files_only is False, ensure all files are downloaded
+            # Download config file (will use cache if available, download if missing)
+            rbln_config_cache_path = hf_hub_download(
+                repo_id=model_id,
+                filename=config_filename,
+                token=token,
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                local_files_only=False,
+            )
+            cache_dir_path = Path(rbln_config_cache_path).parent
+
+            # Download all .rbln files found in cache (hf_hub_download will use cache if available)
+            for rbln_file in rbln_files:
+                filename = rbln_file.name if subfolder == "" else f"{subfolder}/{rbln_file.name}"
+                try:
+                    hf_hub_download(
+                        repo_id=model_id,
+                        filename=filename,
+                        token=token,
+                        revision=revision,
+                        cache_dir=cache_dir,
+                        force_download=force_download,
+                        local_files_only=False,
+                    )
+                except LocalEntryNotFoundError:
+                    # File might not exist in repo, skip it
+                    pass
+
+            # Note: We skip the API call here since we're using cached files
+            # If there are additional files in the repo that aren't cached,
+            # they won't be downloaded.
+            # If the user needs all files, they should use force_download=True
+            return cache_dir_path
+
+    # If local_files_only is True and config not found in cache, try to download with local_files_only
+    if local_files_only:
+        try:
+            rbln_config_cache_path = hf_hub_download(
+                repo_id=model_id,
+                filename=config_filename,
+                token=token,
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                local_files_only=True,
+            )
+            cache_dir_path = Path(rbln_config_cache_path).parent
+            rbln_files = list(cache_dir_path.glob("*.rbln"))
+            rbln_config_filenames = [Path(rbln_config_cache_path)] if Path(rbln_config_cache_path).exists() else []
+            validate_files(rbln_files, rbln_config_filenames, f"cached repository {model_id}")
+            return cache_dir_path
+        except LocalEntryNotFoundError as err:
+            raise FileNotFoundError(
+                f"Could not find compiled model files for {model_id} in local cache. "
+                f"Set local_files_only=False to download from HuggingFace Hub."
+            ) from err
+
+    # List files from repository. This only happens when:
+    # 1. Config is not cached, OR
+    # 2. force_download=True, OR
+    # 3. local_files_only=False and we need to discover all files in the repo
     huggingface_token = _get_huggingface_token(token)
     repo_files = list(
         map(
@@ -51,7 +143,6 @@ def pull_compiled_model_from_hub(
         rbln_config_cache_path = hf_hub_download(
             repo_id=model_id,
             filename=filename,
-            subfolder=subfolder,
             token=token,
             revision=revision,
             cache_dir=cache_dir,

optimum/rbln/utils/runtime_utils.py

@@ -98,7 +98,7 @@ def tp_and_devices_are_ok(
             return None
         if rebel.get_npu_name(device_id) is None:
             return (
-                f"Device {device_id} is not a valid NPU device. Please check your NPU status with 'rbln-stat' command."
+                f"Device {device_id} is not a valid NPU device. Please check your NPU status with 'rbln-smi' command."
             )
 
     if rebel.device_count() < tensor_parallel_size:
@@ -185,7 +185,7 @@ class UnavailableRuntime:
             "This model was loaded with create_runtimes=False. To use this model for inference:\n"
             "1. Load the model with runtime creation enabled:\n"
             "   model = RBLNModel.from_pretrained(..., rbln_create_runtimes=True)\n"
-            "2. Ensure your NPU hardware is properly configured (check with 'rbln-stat' command)\n"
+            "2. Ensure your NPU hardware is properly configured (check with 'rbln-smi' command)\n"
             "3. If you're on a machine without NPU hardware, you need to transfer the model files\n"
             "   to a compatible system with NPU support."
         )

{optimum_rbln-0.9.5a4.dist-info → optimum_rbln-0.10.0.post1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.9.5a4
+Version: 0.10.0.post1
 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