optimum-rbln 0.9.4a2__py3-none-any.whl → 0.10.0.post1__py3-none-any.whl

optimum/rbln/transformers/__init__.py

@@ -68,6 +68,8 @@ _import_structure = {
         "RBLNDecoderOnlyModelForCausalLMConfig",
         "RBLNDecoderOnlyModelConfig",
         "RBLNDecoderOnlyModel",
+        "RBLNDetrForObjectDetection",
+        "RBLNDetrForObjectDetectionConfig",
         "RBLNDistilBertForQuestionAnswering",
         "RBLNDistilBertForQuestionAnsweringConfig",
         "RBLNDPTForDepthEstimation",
@@ -78,6 +80,10 @@ _import_structure = {
         "RBLNExaoneForCausalLMConfig",
         "RBLNGemmaModel",
         "RBLNGemmaModelConfig",
+        "RBLNGemma2ForCausalLM",
+        "RBLNGemma2ForCausalLMConfig",
+        "RBLNGemma2Model",
+        "RBLNGemma2ModelConfig",
         "RBLNGemma3ForCausalLM",
         "RBLNGemma3ForCausalLMConfig",
         "RBLNGemma3ForConditionalGeneration",
@@ -88,6 +94,8 @@ _import_structure = {
         "RBLNGPT2LMHeadModelConfig",
         "RBLNGPT2Model",
         "RBLNGPT2ModelConfig",
+        "RBLNGptOssForCausalLM",
+        "RBLNGptOssForCausalLMConfig",
         "RBLNGroundingDinoDecoder",
         "RBLNGroundingDinoDecoderConfig",
         "RBLNGroundingDinoForObjectDetection",
@@ -110,6 +118,10 @@ _import_structure = {
         "RBLNPegasusForConditionalGenerationConfig",
         "RBLNPegasusModel",
         "RBLNPegasusModelConfig",
+        "RBLNPaliGemmaForConditionalGeneration",
+        "RBLNPaliGemmaForConditionalGenerationConfig",
+        "RBLNPaliGemmaModel",
+        "RBLNPaliGemmaModelConfig",
         "RBLNLlavaNextForConditionalGeneration",
         "RBLNLlavaNextForConditionalGenerationConfig",
         "RBLNLoRAAdapterConfig",
@@ -120,6 +132,8 @@ _import_structure = {
         "RBLNMistralForCausalLMConfig",
         "RBLNMistralModel",
         "RBLNMistralModelConfig",
+        "RBLNMixtralForCausalLM",
+        "RBLNMixtralForCausalLMConfig",
         "RBLNOPTForCausalLM",
         "RBLNOPTForCausalLMConfig",
         "RBLNOPTModel",
@@ -134,14 +148,22 @@ _import_structure = {
         "RBLNQwen2_5_VisionTransformerPretrainedModelConfig",
         "RBLNQwen2_5_VLForConditionalGeneration",
         "RBLNQwen2_5_VLForConditionalGenerationConfig",
+        "RBLNQwen2_5_VLModel",
+        "RBLNQwen2_5_VLModelConfig",
         "RBLNQwen2VisionTransformerPretrainedModel",
         "RBLNQwen2VisionTransformerPretrainedModelConfig",
         "RBLNQwen2VLForConditionalGeneration",
         "RBLNQwen2VLForConditionalGenerationConfig",
+        "RBLNQwen2VLModel",
+        "RBLNQwen2VLModelConfig",
         "RBLNQwen2Model",
         "RBLNQwen2ModelConfig",
         "RBLNQwen2ForCausalLM",
         "RBLNQwen2ForCausalLMConfig",
+        "RBLNQwen2MoeForCausalLM",
+        "RBLNQwen2MoeForCausalLMConfig",
+        "RBLNQwen3MoeForCausalLM",
+        "RBLNQwen3MoeForCausalLMConfig",
         "RBLNQwen3ForCausalLM",
         "RBLNQwen3ForCausalLMConfig",
         "RBLNQwen3Model",
@@ -228,12 +250,18 @@ if TYPE_CHECKING:
         RBLNDecoderOnlyModelForCausalLMConfig,
         RBLNDepthAnythingForDepthEstimation,
         RBLNDepthAnythingForDepthEstimationConfig,
+        RBLNDetrForObjectDetection,
+        RBLNDetrForObjectDetectionConfig,
         RBLNDistilBertForQuestionAnswering,
         RBLNDistilBertForQuestionAnsweringConfig,
         RBLNDPTForDepthEstimation,
         RBLNDPTForDepthEstimationConfig,
         RBLNExaoneForCausalLM,
         RBLNExaoneForCausalLMConfig,
+        RBLNGemma2ForCausalLM,
+        RBLNGemma2ForCausalLMConfig,
+        RBLNGemma2Model,
+        RBLNGemma2ModelConfig,
         RBLNGemma3ForCausalLM,
         RBLNGemma3ForCausalLMConfig,
         RBLNGemma3ForConditionalGeneration,
@@ -246,6 +274,8 @@ if TYPE_CHECKING:
         RBLNGPT2LMHeadModelConfig,
         RBLNGPT2Model,
         RBLNGPT2ModelConfig,
+        RBLNGptOssForCausalLM,
+        RBLNGptOssForCausalLMConfig,
         RBLNGroundingDinoDecoder,
         RBLNGroundingDinoDecoderConfig,
         RBLNGroundingDinoEncoder,
@@ -272,10 +302,16 @@ if TYPE_CHECKING:
         RBLNMistralForCausalLMConfig,
         RBLNMistralModel,
         RBLNMistralModelConfig,
+        RBLNMixtralForCausalLM,
+        RBLNMixtralForCausalLMConfig,
         RBLNOPTForCausalLM,
         RBLNOPTForCausalLMConfig,
         RBLNOPTModel,
         RBLNOPTModelConfig,
+        RBLNPaliGemmaForConditionalGeneration,
+        RBLNPaliGemmaForConditionalGenerationConfig,
+        RBLNPaliGemmaModel,
+        RBLNPaliGemmaModelConfig,
         RBLNPegasusForConditionalGeneration,
         RBLNPegasusForConditionalGenerationConfig,
         RBLNPegasusModel,
@@ -290,18 +326,26 @@ if TYPE_CHECKING:
         RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
         RBLNQwen2_5_VLForConditionalGeneration,
         RBLNQwen2_5_VLForConditionalGenerationConfig,
+        RBLNQwen2_5_VLModel,
+        RBLNQwen2_5_VLModelConfig,
         RBLNQwen2ForCausalLM,
         RBLNQwen2ForCausalLMConfig,
         RBLNQwen2Model,
         RBLNQwen2ModelConfig,
+        RBLNQwen2MoeForCausalLM,
+        RBLNQwen2MoeForCausalLMConfig,
         RBLNQwen2VisionTransformerPretrainedModel,
         RBLNQwen2VisionTransformerPretrainedModelConfig,
         RBLNQwen2VLForConditionalGeneration,
         RBLNQwen2VLForConditionalGenerationConfig,
+        RBLNQwen2VLModel,
+        RBLNQwen2VLModelConfig,
         RBLNQwen3ForCausalLM,
         RBLNQwen3ForCausalLMConfig,
         RBLNQwen3Model,
         RBLNQwen3ModelConfig,
+        RBLNQwen3MoeForCausalLM,
+        RBLNQwen3MoeForCausalLMConfig,
         RBLNResNetForImageClassification,
         RBLNResNetForImageClassificationConfig,
         RBLNRobertaForMaskedLM,

optimum/rbln/transformers/modeling_attention_utils.py

@@ -1,19 +1,16 @@
 import math
-from collections import Counter, defaultdict
-from typing import TYPE_CHECKING, Dict, Optional, Tuple
+from collections import defaultdict
+from typing import Optional, Tuple
 
 import rebel
 
 from ..utils.logging import get_logger
-from ..utils.runtime_utils import get_available_dram
+from ..utils.runtime_utils import get_available_dram, is_compiler_supports_buffer_resize
 from .models.decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
 
 
 logger = get_logger()
 
-if TYPE_CHECKING:
-    from transformers import PretrainedConfig, PreTrainedModel
-
 
 DEFAULT_FLASH_ATTN_PARTITION_LENGTH = 16_384
 DEFAULT_MAX_EAGER_ATTN_SEQUENCE_LENGTH = 32_768
@@ -123,7 +120,7 @@ def align_2MB(x: int) -> int:
     return align(x, 2**21)
 
 
-def get_alloc_memory_by_key(compiled_models: Dict[str, "rebel.RBLNCompiledModel"]) -> Dict[str, int]:
+def get_alloc_memory_by_key(compiled_models: dict[str, rebel.RBLNCompiledModel]) -> dict[str, int]:
     alloc_memory_by_key = defaultdict(int)
     # Get the actual memory allocation of each node by key
     for compiled_model in compiled_models.values():
@@ -147,239 +144,144 @@ def format_byte_size(nbytes: int) -> str:
 
 class RBLNDecoderOnlyFlashAttentionMixin:
     @classmethod
-    def get_maximum_num_blocks_by_model(
-        cls,
-        model: "PreTrainedModel",
-        model_config: "PretrainedConfig",
-        rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
-    ) -> int:
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-        available_dram = get_available_dram(rbln_config.npu) * tensor_parallel_size
-
-        kernel_memory = cls._get_kernel_memory(model, model_config=model_config, rbln_config=rbln_config)
-        buffer = cls._get_buffer(rbln_config)
-
-        remaining_dram = available_dram - kernel_memory - buffer
-        if remaining_dram <= 0:
+    def set_kvcache_num_blocks_after_compilation(
+        cls, compiled_models: dict[str, rebel.RBLNCompiledModel], rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
+    ):
+        rbln_config.kvcache_num_blocks = cls.estimate_num_kvcache_blocks(
+            compiled_models=compiled_models, rbln_config=rbln_config
+        )
+        if rbln_config.kvcache_num_blocks < rbln_config.num_min_blocks:
             raise ValueError(
-                "Insufficient available DRAM after accounting for kernel memory and buffer. "
-                "Cannot allocate any KV cache blocks."
-                f" (Available DRAM: {format_byte_size(available_dram)}, "
-                f"Kernel Memory: {format_byte_size(kernel_memory)}, "
-                f"Buffer: {format_byte_size(buffer)})"
+                "Memory is not enought for full sequence length. "
+                "Please consider decreasing `max_seq_len` to reduce the number of blocks."
             )
-        estimated_num_blocks = cls._estimate_num_blocks(
-            remaining_dram, model_config=model_config, rbln_config=rbln_config
-        )
-
-        return estimated_num_blocks
-
-    @classmethod
-    def _get_kernel_memory(
-        cls,
-        model: "PreTrainedModel",
-        model_config: "PretrainedConfig",
-        rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
-    ) -> int:
-        if model.get_output_embeddings() is None:
-            lm_head_nbytes = 0
-        else:
-            lm_head_nbytes = cls._get_lm_head_memory(model_config, rbln_config)
-
-        layer_nbytes = cls._get_layer_memory(model, model_config, rbln_config)
-        return lm_head_nbytes + layer_nbytes
-
-    @classmethod
-    def _get_lm_head_memory(
-        cls, model_config: "PretrainedConfig", rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
-    ) -> int:
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-        vocab_size = model_config.vocab_size
-        hidden_size = getattr(model_config, "n_embd", None) or model_config.hidden_size
-        lm_head_params = align(vocab_size, 64) * hidden_size
-
-        nbytes_per_param = 2  # Assuming lm_head is always not quantized
-        lm_head_memory_in_bytes = (
-            align_2MB(lm_head_params * nbytes_per_param / tensor_parallel_size) * tensor_parallel_size
+        cls.multiply_kv_cache_num_blocks(
+            compiled_models=compiled_models, rbln_config=rbln_config, multiplier=rbln_config.kvcache_num_blocks
         )
 
-        return lm_head_memory_in_bytes
-
     @classmethod
-    def _get_layer_memory(
+    def estimate_num_kvcache_blocks(
         cls,
-        model: "PreTrainedModel",
-        model_config: "PretrainedConfig",
+        compiled_models: dict[str, rebel.RBLNCompiledModel],
         rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
+        available_dram: Optional[int] = None,
     ) -> int:
-        # This is an *APPROXIMATE* calculation based on the number of parameters
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-        num_hidden_layers = getattr(model_config, "n_layer", None) or model_config.num_hidden_layers
+        if available_dram is None:
+            available_dram = get_available_dram(rbln_config.npu)
 
-        n_model_params = sum(p.numel() for p in model.parameters())
-        embed_token_params = sum(p.numel() for p in model.get_input_embeddings().parameters())
-
-        # Check : `embed_token` is same as `lm_head`
-        if model.get_output_embeddings() is not None:
-            params = n_model_params - 2 * embed_token_params
-        else:
-            params = n_model_params - embed_token_params
-
-        # Assuming all layers have the same number of parameters
-        # and all linear layers are quantized if quantization is enabled (This is not always true)
-        # TODO(jongho): More accurate calculation
-        nbits_per_param = rbln_config.nbits_per_param
-        layer_nbytes = (
-            (align_2MB(params // num_hidden_layers * nbits_per_param // 8 / tensor_parallel_size))
-            * num_hidden_layers
-            * tensor_parallel_size
-        )
-
-        return layer_nbytes
-
-    @classmethod
-    def _get_buffer(cls, rbln_config) -> int:
-        # TODO(jongho): Accurate buffer estimation
-        buffer_per_runtime_per_core = 2**28  # 256MB per runtime
-        num_runtimes = 1 if not rbln_config.can_generate else 1 + len(rbln_config.decoder_batch_sizes)
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-
-        buffer_per_core = buffer_per_runtime_per_core * num_runtimes
-        buffer = buffer_per_core * tensor_parallel_size
-        return buffer
-
-    @classmethod
-    def get_maximum_num_blocks_by_compiled_model(
-        cls,
-        compiled_models: Dict[str, "rebel.RBLNCompiledModel"],
-        model_config: "PretrainedConfig",
-        rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
-    ) -> int:
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-        available_dram = get_available_dram(rbln_config.npu) * tensor_parallel_size
-
-        alloc_memory_by_key = get_alloc_memory_by_key(compiled_models)
-        alloc_memory_by_key.pop("PortRecur", None)  # Old compiler's kv-cache Key
-        alloc_memory_by_key.pop("DramTensor", None)  # kv-cache
-        used_memory = sum(alloc_memory_by_key.values())
-
-        remaining_dram = available_dram - used_memory
-
-        if remaining_dram <= 0:
+        if "prefill" not in rbln_config.phases:
             logger.warning(
-                "Insufficient available DRAM after accounting for kernel memory and buffer. "
-                "Model cannot allocate any KV cache blocks."
+                "Not estimating number of KV cache blocks since `prefill` phase is not in the `phases` list."
             )
-
-        estimated_num_blocks = cls._estimate_num_blocks(
-            remaining_dram, model_config=model_config, rbln_config=rbln_config
-        )
-
-        return estimated_num_blocks
-
-    @classmethod
-    def _estimate_num_blocks(
-        cls, available_dram: int, model_config: "PretrainedConfig", rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
-    ) -> int:
-        """
-        Estimate the maximum number of KV cache blocks that can be allocated.
-
-        if all of the layers are full attention, the dram_per_block can be calculated simply as follows:
-            num_blocks = available_dram // dram_per_block
-
-        However, if the model contains a mix of full attention and sliding window attention layers,
-        we need to consider the memory occupied by the sliding window attention layers first,
-        since their memory usage is constant regardless of the number of blocks.
-            num_blocks = (available_dram - swa_kv_nbytes) // dram_per_block
-
-        """
-
-        def get_dram_per_block(seq_len: int, num_key_value_heads: int, tensor_parallel_size: int) -> int:
-            nbytes_per_param = 2  # Assuming kv-cache is always not quantized
-            dram_per_block = (
-                seq_len
-                * align(head_dim, 64)
-                * math.ceil(num_key_value_heads / tensor_parallel_size)
-                * nbytes_per_param
-                * tensor_parallel_size
-                * 2
-            )  # *2 for key and value
-
-            return dram_per_block
-
-        num_attention_heads = getattr(model_config, "n_head", None) or model_config.num_attention_heads
-        head_dim = getattr(model_config, "head_dim", None) or model_config.hidden_size // num_attention_heads
-        num_hidden_layers = getattr(model_config, "n_layer", None) or model_config.num_hidden_layers
-        num_key_value_heads = getattr(model_config, "num_key_value_heads", None) or num_attention_heads
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-
-        # Consider layer types if available
-        # If layer types are not found, assume all layers are full attention
-        layer_types = getattr(model_config, "layer_types", None)
-        if layer_types:
-            layer_types_dict = Counter(layer_types)
-            num_full_attention = layer_types_dict.pop("full_attention", 0)
-            num_sliding_window_attention = layer_types_dict.pop("sliding_attention", 0)
-            if len(layer_types_dict) > 0:
-                raise ValueError(f"Unknown layer types found in the config: {layer_types_dict.keys()}")
-
-        else:
-            num_full_attention = num_hidden_layers
-            num_sliding_window_attention = 0
-
-        # Reduce available DRAM by sliding window attention kv-cache
-        # Since memory occupation of swa layer is constant regardless of num_blocks
-        swa_kv_nbytes = 0
-        if num_sliding_window_attention > 0:
-            sliding_window = getattr(model_config, "sliding_window", None)
-            if sliding_window is None:
-                logger.warning(
-                    "`sliding_window` is not found in the config while `sliding_attention` layers are present. "
-                    "Assuming maximum sliding window size for estimation."
-                )
-                sliding_window = rbln_config.kvcache_block_size
-
-            swa_kv_nbytes = num_sliding_window_attention * get_dram_per_block(
-                seq_len=sliding_window,
-                num_key_value_heads=num_key_value_heads,
-                tensor_parallel_size=tensor_parallel_size,
+            return 1
+
+        num_node = rbln_config.tensor_parallel_size or 1
+        alloc_per_node_without_dram = [0] * num_node
+
+        for compiled_model in compiled_models.values():
+            for key, alloc_per_node in compiled_model.get_alloc_per_node_by_key().items():
+                if key == "DramTensor":
+                    continue
+
+                if len(alloc_per_node) != num_node:
+                    alloc_per_node += [0] * (num_node - len(alloc_per_node))
+
+                alloc_per_node_without_dram = [a + b for a, b in zip(alloc_per_node_without_dram, alloc_per_node)]
+
+        remaining_dram_at_node: list[int] = [
+            available_dram - without_dramtensor for without_dramtensor in alloc_per_node_without_dram
+        ]
+
+        # kvcache_tensor_sizes[key][node_id][chiplet_id] = alloc_size
+        kvcache_tensor_sizes: dict[str, list[list[int]]] = compiled_models["prefill"].exp_get_dram_tensor_sizes()
+        kvcache_meta_can_resize: dict[str, bool] = {
+            kvcache_meta.name: kvcache_meta.can_resize for kvcache_meta in rbln_config.kvcache_metas
+        }
+
+        def get_updated_kvcache_tensor_sizes(
+            kvcache_tensor_sizes: dict[str, list[list[int]]], multiplier: int
+        ) -> dict[str, list[list[int]]]:
+            # Get the updated KV cache tensor sizes by multiplying the multiplier
+            # with considering attention type (full or sliding), and memory alignment.
+            ret: dict[str, list[list[int]]] = {}
+            for key, sizes_at_node in kvcache_tensor_sizes.items():
+                m = multiplier if kvcache_meta_can_resize[key] else 1
+                ret[key] = [
+                    [align_2MB(size_at_chiplet * m) for size_at_chiplet in sizes_at_node_at_chiplet]
+                    for sizes_at_node_at_chiplet in sizes_at_node
+                ]
+            return ret
+
+        def check_memory_fits(multiplier: int) -> tuple[bool, list[int]]:
+            # Check if the given multiplier fits in memory
+            # Returns (fits: bool, kvcache_tensor_sizes_at_node: list[int])
+            updated_kvcache_tensor_sizes = get_updated_kvcache_tensor_sizes(kvcache_tensor_sizes, multiplier)
+
+            kvcache_tensor_sizes_at_node: list[int] = [0] * num_node
+            for tensor_sizes_at_node in updated_kvcache_tensor_sizes.values():
+                tensor_sizes_at_node: list[list[int]]
+                for node_id, sizes_at_chiplet in enumerate(tensor_sizes_at_node):
+                    sizes_at_chiplet: list[int]
+                    kvcache_tensor_sizes_at_node[node_id] += sum(sizes_at_chiplet)
+
+            fits = all(
+                remaining_dram_at_node[node_id] >= kvcache_tensor_sizes_at_node[node_id] for node_id in range(num_node)
             )
-
-            available_dram -= swa_kv_nbytes
-
-        dram_per_block = num_full_attention * get_dram_per_block(
-            seq_len=rbln_config.kvcache_block_size,
-            num_key_value_heads=num_key_value_heads,
-            tensor_parallel_size=tensor_parallel_size,
+            return fits, kvcache_tensor_sizes_at_node
+
+        # Fast path: try maximum blocks first (most common case)
+        fits, _ = check_memory_fits(rbln_config.num_full_blocks)
+        if fits:
+            # Best case: maximum blocks fit in memory
+            return rbln_config.num_full_blocks
+
+        # Slow path: binary search for optimal multiplier
+        logger.debug(
+            f"[KVCache] Not enough memory for {rbln_config.num_full_blocks} blocks. "
+            f"Searching for optimal multiplier..."
         )
 
-        if dram_per_block == 0:
-            raise ValueError("DRAM per block is calculated as zero, cannot estimate maximum number of blocks.")
+        left, right = 1, rbln_config.num_full_blocks - 1
+        multiplier = 1  # Default to minimum if no valid multiplier found
+
+        while left <= right:
+            mid = (left + right) // 2
+            fits, kvcache_tensor_sizes_at_node = check_memory_fits(mid)
+
+            if fits:
+                # Memory is sufficient, try larger multiplier
+                multiplier = mid
+                left = mid + 1
+            else:
+                # Memory is insufficient, try smaller multiplier
+                logger.debug(
+                    f"[KVCache] Not enough memory for {mid} blocks. Remaining DRAM: "
+                    f"{[format_byte_size(remaining_dram) for remaining_dram in remaining_dram_at_node]}, "
+                    f"KV cache tensor sizes: {[format_byte_size(size) for size in kvcache_tensor_sizes_at_node]}"
+                )
+                right = mid - 1
 
-        max_n_blocks = available_dram // dram_per_block
-        return max_n_blocks
+        return multiplier
 
     @classmethod
-    def maybe_suggest_kvcache_num_blocks(
+    def multiply_kv_cache_num_blocks(
         cls,
-        compiled_models: Dict[str, "rebel.RBLNCompiledModel"],
-        model_config: "PretrainedConfig",
+        compiled_models: dict[str, rebel.RBLNCompiledModel],
         rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
-    ) -> None:
-        max_num_blocks = cls.get_maximum_num_blocks_by_compiled_model(
-            compiled_models=compiled_models,
-            model_config=model_config,
-            rbln_config=rbln_config,
-        )
+        multiplier: int,
+    ):
+        if not is_compiler_supports_buffer_resize():
+            raise RuntimeError(
+                "The installed version of rebel-compiler does not support automatic kv cache size determination. "
+                "Please upgrade rebel-compiler to a version that supports this feature, "
+                "or explicitly set 'kvcache_num_blocks' in rbln_config to manually specify the cache size."
+            )
 
-        # Since our estimation logic is not always accurate,
-        # users can set `kvcache_num_blocks` to `max_num_blocks`.
-        # If the memory is not enough, the model will fail to compile.
-        if rbln_config.kvcache_num_blocks < max_num_blocks:
-            logger.warning(
-                f"Current `kvcache_num_blocks` setting is {rbln_config.kvcache_num_blocks}. "
-                "Our analysis indicates that additional memory is available for more blocks. "
-                f"Consider increasing `kvcache_num_blocks` to {max_num_blocks} for potentially improved performance. "
-                "Please be advised that our memory estimation algorithm has limitations, "
-                "and increasing this value may not guarantee successful model compilation."
+        for compiled_model in compiled_models.values():
+            compiled_model.exp_multiply_buffer_size(
+                {
+                    kvcache_meta.name: multiplier
+                    for kvcache_meta in rbln_config.kvcache_metas
+                    if kvcache_meta.can_resize
+                }
             )

optimum/rbln/transformers/modeling_outputs.py

@@ -18,6 +18,8 @@ from typing import Optional, Tuple
 import torch
 from transformers.modeling_outputs import ModelOutput
 
+from ..configuration_utils import RBLNModelConfig
+
 
 @dataclass
 class RBLNDecoderOnlyOutput(ModelOutput):
@@ -36,3 +38,26 @@ class RBLNGemma3ForCausalLMOutput(RBLNDecoderOnlyOutput):
 class RBLNSeq2SeqTSDecoderOutput(ModelOutput):
     last_hidden_states: torch.FloatTensor = None
     params: Tuple[torch.FloatTensor] = None
+
+
+def _validate_output_hidden_states(output_hidden_states: Optional[bool], rbln_config: RBLNModelConfig):
+    output_hidden_states = (
+        output_hidden_states if output_hidden_states is not None else rbln_config.output_hidden_states
+    )
+    if output_hidden_states != rbln_config.output_hidden_states:
+        raise ValueError(
+            f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {rbln_config.output_hidden_states} "
+            f"Please compile again with the correct argument."
+        )
+
+    return output_hidden_states
+
+
+def _validate_output_attentions(output_attentions: Optional[bool], rbln_config: RBLNModelConfig):
+    output_attentions = output_attentions if output_attentions is not None else rbln_config.output_attentions
+    if output_attentions != rbln_config.output_attentions:
+        raise ValueError(
+            f"Variable output_attentions {output_attentions} is not equal to rbln_config.output_attentions {rbln_config.output_attentions} "
+            f"Please compile again with the correct argument."
+        )
+    return output_attentions