optimum-rbln 0.7.4a0__py3-none-any.whl → 0.7.4a2__py3-none-any.whl

optimum/rbln/__init__.py

@@ -73,6 +73,7 @@ _import_structure = {
         "RBLNRobertaForMaskedLM",
         "RBLNViTForImageClassification",
         "RBLNBertForMaskedLM",
+        "RBLNTimeSeriesTransformerForPrediction",
     ],
     "diffusers": [
         "RBLNAutoencoderKL",
@@ -184,6 +185,7 @@ if TYPE_CHECKING:
         RBLNRobertaForSequenceClassification,
         RBLNT5EncoderModel,
         RBLNT5ForConditionalGeneration,
+        RBLNTimeSeriesTransformerForPrediction,
         RBLNViTForImageClassification,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,

optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.4a0'
+__version__ = version = '0.7.4a2'
 __version_tuple__ = version_tuple = (0, 7, 4)

optimum/rbln/ops/__init__.py

@@ -19,3 +19,4 @@ from .attn import (
 )
 from .flash_attn import register_rbln_custom_paged_flash_attention, register_rbln_custom_paged_flash_causal_attention
 from .kv_cache_update import register_rbln_custom_cache_update
+from .linear import linear

optimum/rbln/ops/linear.py

@@ -0,0 +1,25 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+@torch.library.custom_op("rbln_custom_ops::linear", mutates_args=())
+def linear(input: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
+    output_shape = list(input.shape[:-1])
+    output_shape += [weight.shape[0]]
+    return torch.empty(size=output_shape, dtype=input.dtype, device=input.device, requires_grad=input.requires_grad)

optimum/rbln/transformers/__init__.py

@@ -52,6 +52,7 @@ _import_structure = {
         "RBLNPhiForCausalLM",
         "RBLNT5EncoderModel",
         "RBLNT5ForConditionalGeneration",
+        "RBLNTimeSeriesTransformerForPrediction",
         "RBLNLlavaNextForConditionalGeneration",
         "RBLNMidmLMHeadModel",
         "RBLNXLMRobertaModel",
@@ -113,6 +114,7 @@ if TYPE_CHECKING:
         RBLNQwen2ForCausalLM,
         RBLNT5EncoderModel,
         RBLNT5ForConditionalGeneration,
+        RBLNTimeSeriesTransformerForPrediction,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
         RBLNXLMRobertaModel,

optimum/rbln/transformers/models/__init__.py

@@ -50,6 +50,7 @@ _import_structure = {
     "mistral": ["RBLNMistralForCausalLM"],
     "phi": ["RBLNPhiForCausalLM"],
     "qwen2": ["RBLNQwen2ForCausalLM"],
+    "time_series_transformers": ["RBLNTimeSeriesTransformerForPrediction"],
     "t5": ["RBLNT5EncoderModel", "RBLNT5ForConditionalGeneration"],
     "wav2vec2": ["RBLNWav2Vec2ForCTC"],
     "whisper": ["RBLNWhisperForConditionalGeneration"],
@@ -90,6 +91,7 @@ if TYPE_CHECKING:
     from .phi import RBLNPhiForCausalLM
     from .qwen2 import RBLNQwen2ForCausalLM
     from .t5 import RBLNT5EncoderModel, RBLNT5ForConditionalGeneration
+    from .time_series_transformers import RBLNTimeSeriesTransformerForPrediction
     from .wav2vec2 import RBLNWav2Vec2ForCTC
     from .whisper import RBLNWhisperForConditionalGeneration
     from .xlm_roberta import RBLNXLMRobertaModel

optimum/rbln/transformers/models/bart/modeling_bart.py

@@ -94,12 +94,11 @@ class RBLNBartModel(RBLNModel):
             for model_input_name in rbln_model_input_names
         ]
 
-        enc_compile_config = RBLNCompileConfig(input_info=input_info, compiled_model_name="encoder")
-        dec_compile_config = RBLNCompileConfig(input_info=input_info, compiled_model_name="decoder")
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
 
         rbln_config = RBLNConfig(
             rbln_cls=cls.__name__,
-            compile_cfgs=[enc_compile_config, dec_compile_config],
+            compile_cfgs=[rbln_compile_config],
             rbln_kwargs=rbln_kwargs,
         )
 

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

@@ -222,8 +222,6 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
 
             attention_mask = self.dec_attn_mask
 
-            attention_mask = self.dec_attn_mask
-
         logits = super().forward(
             inputs,
             cache_position,
@@ -547,22 +545,27 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
         @QuantizationManager.with_quantization_env
         def compile_model(*args, **kwargs):
-            wrapped_model.phase = "prefill"
-            compiled_prefill = RBLNModel.compile(
-                wrapped_model,
-                prefill_compile_config,
-                example_inputs=prefill_example_inputs,
-                compile_context=context,
-            )
+            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(
+                    wrapped_model,
+                    prefill_compile_config,
+                    example_inputs=prefill_example_inputs,
+                    compile_context=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}
+                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}
+            finally:
+                torch.nn.functional.linear = original_linear
 
         return compile_model(quantize_config=quantize_config)
 
@@ -575,11 +578,41 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         nbits_per_param: int,
         n_model_params: int,
     ) -> int:
+        """
+        We are finding max_n_blocks(x) that satisfies the following equation:
+
+        available_dram - kernel_size - buffer
+            - num_layers * 2 * tensor_parallel_size
+            * align_2MB(
+                x
+                * block_size
+                * align_64(head_dim)
+                * math.ceil(num_key_value_heads / tensor_parallel_size)
+                * 2
+            ) > 0
+
+        This inequality can be rewritten as follows:
+
+        a - c * align_2MB(b * x) > 0
+        where
+           a = available_dram - kernel_size - buffer
+           b = block_size * align_64(head_dim) * math.ceil(num_key_value_heads / tensor_parallel_size) * 2
+           c = num_layers * 2 * tensor_parallel_size
+
+        We can rewrite the inequality as follows:
+        k > align_2MB(b*x)
+        where
+           k = a / c
+
+        After that, we can derive the following equation:
+        x = floor(2**21 / b * floor((k - 1) / 2**21))
+        """
+
         def align(x: int, nbytes: int) -> int:
             return int(math.ceil(x / nbytes) * nbytes)
 
         def align_2MB(x: int) -> int:
-            return align(x, 2 * 1024 * 1024)
+            return align(x, 2**21)
 
         num_attention_heads = getattr(config, "n_head", None) or getattr(config, "num_attention_heads")
         num_layers = getattr(config, "n_layer", None) or getattr(config, "num_hidden_layers")
@@ -609,27 +642,16 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         available_dram -= kernel_size
 
         # TODO: Accurate buffer estimation
-        buffer = 2**30  # 1GB Buffer
-        if tensor_parallel_size <= 4:
-            buffer /= 4
-
+        buffer_per_core = 2**29  # 500MB per npu
+        buffer = buffer_per_core * tensor_parallel_size
         available_dram -= buffer
 
-        # Estimate nbytes per a single kvcache block
-        nbytes_per_block = (
-            align_2MB(
-                kvcache_block_size
-                * head_dim
-                * math.ceil(num_key_value_heads / tensor_parallel_size)  # Shard
-                * 2  # (fp16)
-            )
-            * num_layers
-            * 2  # (k, v)
-            * tensor_parallel_size
-        )
-        n_blocks = available_dram // nbytes_per_block
+        b = kvcache_block_size * align(head_dim, 64) * math.ceil(num_key_value_heads / tensor_parallel_size) * 2
+        c = num_layers * 2 * tensor_parallel_size
+        k = available_dram / c
+        max_n_blocks = math.floor(2**21 / b * math.floor((k - 1) / 2**21))
 
-        return n_blocks, nbytes_per_block
+        return max_n_blocks
 
     @classmethod
     def _get_rbln_config(
@@ -686,7 +708,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
         rbln_kvcache_num_blocks = (rbln_max_seq_len // rbln_kvcache_block_size) * rbln_batch_size
         if rbln_attn_impl == "flash_attn":
-            max_num_blocks, _ = cls.get_maximum_num_blocks(
+            max_num_blocks = cls.get_maximum_num_blocks(
                 config=model_config,
                 tensor_parallel_size=rbln_kwargs.get("tensor_parallel_size", 1),
                 kvcache_block_size=rbln_kvcache_block_size,

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

@@ -38,8 +38,8 @@ class RBLNRuntimeEncoder(RBLNPytorchRuntime):
     mandatory_members = ["main_input_name"]
 
     def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
-        _ = super().forward(*args, **kwargs)
-        return BaseModelOutput(last_hidden_state=torch.tensor([1.0]))
+        output = super().forward(*args, **kwargs)
+        return BaseModelOutput(last_hidden_state=output)
 
 
 class RBLNRuntimeDecoder(RBLNPytorchRuntime):

optimum/rbln/transformers/models/t5/modeling_t5.py

@@ -36,19 +36,50 @@ from .t5_architecture import T5Wrapper
 logger = get_logger()
 
 if TYPE_CHECKING:
+    from rebel import Runtime
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
 
 
 class RBLNRuntimeModel(RBLNPytorchRuntime):
+    def __init__(
+        self,
+        runtime: "Runtime",
+        max_seq_len: int,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(runtime, **kwargs)
+        self.max_seq_len = max_seq_len
+
+    def _prepare_inputs(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: torch.LongTensor,
+    ):
+        input_len = input_ids.shape[-1]
+        pad_len = None
+        if input_len > self.max_seq_len:
+            raise ValueError(f"Error input_len({input_len}) exceed max_seq_len({self.max_seq_len}).")
+        elif input_len < self.max_seq_len and input_len > 0:
+            pad_len = self.max_seq_len - input_len
+            logger.warning(
+                f"Warning: The input was padded with {pad_len} tokens to meet the compiled model's requirements. "
+                "For optimal performance, consider recompiling with a shorter 'rbln_max_seq_len'."
+            )
+            input_ids = torch.nn.functional.pad(input_ids, (0, pad_len))
+            attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_len), value=0)
+
+        return input_ids, attention_mask, pad_len
+
     def forward(
         self,
         input_ids: torch.LongTensor,
-        attention_mask: torch.FloatTensor,
+        attention_mask: torch.LongTensor,
         head_mask: torch.FloatTensor,
         inputs_embeds: torch.FloatTensor,
         **kwargs,
     ):
-        return super().forward(
+        input_ids, attention_mask, pad_len = self._prepare_inputs(input_ids, attention_mask)
+        logits = super().forward(
             input_ids,
             attention_mask,
             head_mask,
@@ -56,6 +87,8 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             **kwargs,
         )
 
+        return logits[:, :-pad_len, :] if pad_len is not None else logits
+
 
 class T5EncoderWrapper(torch.nn.Module):
     def __init__(self, model: "T5EncoderModel") -> None:
@@ -72,7 +105,8 @@ class RBLNT5EncoderModel(RBLNModel):
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
     def __post_init__(self, **kwargs):
-        self.model = RBLNRuntimeModel(runtime=self.model[0])
+        max_seq_len = self.rbln_config.model_cfg["max_seq_len"]
+        self.model = RBLNRuntimeModel(runtime=self.model[0], max_seq_len=max_seq_len)
 
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):

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

@@ -0,0 +1,24 @@
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from .modeling_time_series_transformers import RBLNTimeSeriesTransformerForPrediction