optimum-rbln 0.8.2rc0__py3-none-any.whl → 0.8.3a1__py3-none-any.whl

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

@@ -34,6 +34,9 @@ class RBLNDecoderOnlyModelConfig(RBLNModelConfig):
     KV cache management, and batching for inference.
     """
 
+    _default_phases = ["prefill"]
+    _default_logits_to_keep = 0
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -47,9 +50,12 @@ class RBLNDecoderOnlyModelConfig(RBLNModelConfig):
         quantization: Optional[Union[Dict[str, Any], RBLNQuantizationConfig]] = None,
         prefill_chunk_size: Optional[int] = None,
         kvcache_num_blocks: Optional[int] = None,
+        decoder_batch_sizes: Optional[List[int]] = None,
         cache_impl: Optional[CacheImplType] = None,
         sliding_window: Optional[int] = None,
         sliding_window_layers: Optional[List[int]] = None,
+        phases: Optional[List[PhaseType]] = None,
+        logits_to_keep: Optional[int] = None,
         **kwargs,
     ):
         """
@@ -80,6 +86,12 @@ class RBLNDecoderOnlyModelConfig(RBLNModelConfig):
             kvcache_num_blocks (Optional[int]): The total number of blocks to allocate for the
                 PagedAttention KV cache. See the "KV Cache Number of Blocks (`kvcache_num_blocks`)"
                 section below for details.
+            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.
             cache_impl (Optional[CacheImplType]): Specifies the KV cache implementation strategy. Defaults to "static".
                 - "static": Uses a fixed-size global KV cache for all layers, suitable for standard attention patterns.
                 - "sliding_window": Implements a sliding window KV cache, where each layer maintains a local cache of recent tokens.
@@ -88,6 +100,10 @@ class RBLNDecoderOnlyModelConfig(RBLNModelConfig):
                 you must specify the `sliding_window` size and optionally `sliding_window_layers` for hybrid mode.
             sliding_window (Optional[int]): The size of the sliding window. Defaults to None.
             sliding_window_layers (Optional[List[int]]): The layers to use for the sliding window used in the hybrid model. Defaults to None.
+            phases (Optional[List[PhaseType]]): The phases to compile the model for. Defaults to ["prefill"] if DecoderOnlyModel is used,
+                ["prefill", "decode"] if DecoderOnlyModelForCausalLM is used.
+            logits_to_keep (Optional[int]): The number of logits to keep for the decoder.  If set to 0, the decoder will keep all logits.
+                Defaults to 0 if DecoderOnlyModel is used, 1 if DecoderOnlyModelForCausalLM is used.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
@@ -156,8 +172,6 @@ class RBLNDecoderOnlyModelConfig(RBLNModelConfig):
         self.batch_size = batch_size or 1
         if not isinstance(self.batch_size, int) or self.batch_size < 0:
             raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
-        if self.batch_size > 1:
-            raise NotImplementedError("Batch size > 1 is not supported for RBLNDecoderOnlyModel.")
 
         self.max_seq_len = max_seq_len
         self.use_inputs_embeds = use_inputs_embeds or False
@@ -183,178 +197,12 @@ class RBLNDecoderOnlyModelConfig(RBLNModelConfig):
         self.sliding_window = sliding_window
         self.sliding_window_layers = sliding_window_layers or []
 
-    @property
-    def use_global_attention(self):
-        return self.cache_impl in ["static", "hybrid"]
-
-    @property
-    def use_local_attention(self):
-        return self.cache_impl in ["sliding_window", "hybrid"]
-
-
-class RBLNDecoderOnlyModelForCausalLMConfig(RBLNDecoderOnlyModelConfig):
-    """
-    Configuration class for RBLN decoder-only models for Causal Language Modeling.
-
-    This class extends RBLNModelConfig with parameters specific to decoder-only transformer
-    architectures optimized for RBLN devices. It controls aspects like attention implementation,
-    KV cache management, and batching for inference.
-    """
-
-    def __init__(
-        self,
-        batch_size: Optional[int] = None,
-        max_seq_len: Optional[int] = None,
-        use_inputs_embeds: Optional[bool] = None,
-        use_attention_mask: Optional[bool] = None,
-        use_position_ids: Optional[bool] = None,
-        attn_impl: Optional[str] = None,
-        kvcache_partition_len: Optional[int] = None,
-        kvcache_block_size: Optional[int] = None,
-        quantization: Optional[Union[Dict[str, Any], RBLNQuantizationConfig]] = None,
-        prefill_chunk_size: Optional[int] = None,
-        kvcache_num_blocks: Optional[int] = None,
-        decoder_batch_sizes: Optional[List[int]] = None,
-        cache_impl: Optional[CacheImplType] = None,
-        sliding_window: Optional[int] = None,
-        sliding_window_layers: Optional[List[int]] = None,
-        phases: Optional[List[PhaseType]] = None,
-        **kwargs,
-    ):
-        """
-        Args:
-            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
-            max_seq_len (Optional[int]): The maximum sequence length supported by the model.
-                If not provided, it attempts to infer from the model's configuration
-                (`max_position_embeddings` or `n_positions`). Must be specified if not available
-                in the model config.
-            use_inputs_embeds (Optional[bool]): Whether to use input embeddings (`inputs_embeds`)
-                directly instead of `input_ids`. Defaults to False. Requires the model to be
-                compiled with this option enabled.
-            use_attention_mask (Optional[bool]): Whether the model requires attention masks during
-                inference. This is typically determined based on the target device and model
-                architecture. Defaults are often set automatically based on the model and RBLN NPU.
-            use_position_ids (Optional[bool]): Whether to use position IDs. Defaults to False.
-            attn_impl (Optional[str]): Specifies the attention implementation to use.
-                See the "Attention Implementation (`attn_impl`)" section below for details.
-            kvcache_partition_len (Optional[int]): Defines the partition length for the KV cache
-                when using "flash_attn". See the "KV Cache Partition Length (`kvcache_partition_len`)"
-                section below for details.
-            kvcache_block_size (Optional[int]): Sets the size (in number of tokens) of each block
-                in the PagedAttention KV cache. See the "KV Cache Block Size (`kvcache_block_size`)"
-                section below for details.
-            quantization (Optional[Dict[str, Any]]): Configuration dictionary for applying model
-                quantization. Specifies format, etc.
-            prefill_chunk_size (Optional[int]): The chunk size used during the prefill phase for
-                processing input sequences. Defaults to 128. Must be a positive integer
-                divisible by 64. Affects prefill performance and memory usage.
-            kvcache_num_blocks (Optional[int]): The total number of blocks to allocate for the
-                PagedAttention KV cache. See the "KV Cache Number of Blocks (`kvcache_num_blocks`)"
-                section below for details.
-            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.
-            cache_impl (Optional[CacheImplType]): Specifies the KV cache implementation strategy. Defaults to "static".
-                - "static": Uses a fixed-size global KV cache for all layers, suitable for standard attention patterns.
-                - "sliding_window": Implements a sliding window KV cache, where each layer maintains a local cache of recent tokens.
-                - "hybrid": Combines both static and sliding window approaches, allowing different layers to use different cache strategies.
-                The choice affects memory usage and attention patterns. When using "sliding_window" or "hybrid",
-                you must specify the `sliding_window` size and optionally `sliding_window_layers` for hybrid mode.
-            sliding_window (Optional[int]): The size of the sliding window. Defaults to None.
-            sliding_window_layers (Optional[List[int]]): The layers to use for the sliding window used in the hybrid model. Defaults to None.
-            phases (Optional[List[PhaseType]]): The phases to compile the model for. Defaults to None.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
-
-        Raises:
-            ValueError: If `batch_size` is not a positive integer.
-            ValueError: If `prefill_chunk_size` is not a positive integer divisible by 64.
-            ValueError: If `max_seq_len` cannot be determined and is required.
-            ValueError: If attention parameter constraints are violated (e.g., `max_seq_len` vs
-                `kvcache_partition_len` for flash attention).
-
-
-        Attention Implementation:
-            `attn_impl` determines the underlying attention mechanism used by the model.
-
-            - **`"eager"`** (Default if `kvcache_partition_len` is not set): Uses the standard PyTorch
-                attention implementation. Suitable for sequences up to a certain limit (e.g., 32,768 tokens).
-            - **`"flash_attn"`**: Utilizes an optimized Flash Attention implementation, beneficial for
-                longer sequences and potentially faster execution. Requires `max_seq_len` to be at least
-                8,192. If `kvcache_partition_len` is specified, `attn_impl` automatically defaults
-                to `"flash_attn"`. When using `"flash_attn"`, `kvcache_block_size` must equal
-                `kvcache_partition_len`.
-
-            The choice impacts performance and memory usage, especially for long sequences.
-            Constraints related to `max_seq_len` and `kvcache_partition_len` apply when using
-            `"flash_attn"`.
-
-
-        KV Cache Partition Length:
-            `kvcache_partition_len` is relevant **only** when `attn_impl` is `"flash_attn"`.
-
-            - It defines the length (number of tokens) of each partition within the Key-Value (KV) cache.
-            - Must be between 4,096 and 32,768 (inclusive).
-            - When using `"flash_attn"`, `max_seq_len` must be a multiple of `kvcache_partition_len`
-                and at least twice its value (`max_seq_len >= 2 * kvcache_partition_len`).
-            - If `attn_impl` is `"flash_attn"` and `kvcache_partition_len` is `None`, it defaults to
-                16,384.
-
-
-        KV Cache Number of Blocks:
-            `kvcache_num_blocks` controls the total number of memory blocks allocated for the PagedAttention KV cache.
-            Each block holds `kvcache_block_size` tokens of Key and Value states.
-
-            - **Automatic Estimation (Default)**: If `kvcache_num_blocks` is `None`, the system estimates
-                the maximum number of blocks that can fit into the available RBLN device memory. This
-                calculation considers the model size (kernel memory), required buffer memory, the number
-                of layers and heads, `kvcache_block_size`, tensor parallelism, and available RBLN NPU DRAM.
-                This aims to maximize cache capacity for potentially better performance with long sequences
-                or larger batches without manual tuning.
-            - **Manual Setting**: You can explicitly set the number of blocks. This provides finer control
-                but requires careful consideration of memory limits. Setting it too high may lead to
-                compilation errors if it exceeds available memory. The system will issue warnings if your
-                setting exceeds the estimated maximum.
-            - **Performance Impact**: A larger number of blocks reduces the likelihood of cache eviction,
-                which is beneficial for tasks involving many long sequences or large batch sizes, enabling
-                higher throughput. However, allocating more blocks consumes more memory.
-            - **Minimum Requirement**: The system requires a minimum number of blocks to function,
-                calculated based on `max_seq_len`, `kvcache_block_size`, and `batch_size`. The number of
-                allocated blocks must be sufficient to hold at least one full sequence length per item
-                in the batch concurrently. The system will log warnings or raise errors if constraints
-                are violated (e.g., if `kvcache_num_blocks` is less than `batch_size` when using Flash Attention).
-
-            The optimal value depends on the specific model, task, hardware, and desired trade-off
-            between performance and memory usage. The automatic estimation provides a robust starting point.
-        """
-
-        super().__init__(
-            max_seq_len=max_seq_len,
-            use_inputs_embeds=use_inputs_embeds,
-            use_attention_mask=use_attention_mask,
-            use_position_ids=use_position_ids,
-            attn_impl=attn_impl,
-            kvcache_partition_len=kvcache_partition_len,
-            kvcache_block_size=kvcache_block_size,
-            quantization=quantization,
-            prefill_chunk_size=prefill_chunk_size,
-            kvcache_num_blocks=kvcache_num_blocks,
-            cache_impl=cache_impl,
-            sliding_window=sliding_window,
-            sliding_window_layers=sliding_window_layers,
-            **kwargs,
-        )
-
-        # override batch_size for causal lm
-        self.batch_size = batch_size or 1
-        if not isinstance(self.batch_size, int) or self.batch_size < 0:
-            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
-
         if phases is not None:
             self.validate_phases_type(phases)
-        self.phases = phases or ["prefill", "decode"]
+        self.phases = phases or self._default_phases
+        self.logits_to_keep = logits_to_keep or self._default_logits_to_keep
+        if self.logits_to_keep is not None and self.logits_to_keep > 1:
+            raise NotImplementedError("`logits_to_keep` > 1 is currently not supported for RBLN models.")
 
         if "decode" in self.phases:
             self.decoder_batch_sizes = decoder_batch_sizes
@@ -384,9 +232,30 @@ class RBLNDecoderOnlyModelForCausalLMConfig(RBLNDecoderOnlyModelConfig):
             raise ValueError(f"All elements in `phases` must be of type `PhaseType`({get_args(PhaseType)}).")
 
     @property
-    def use_multiple_decoder(self):
+    def use_global_attention(self) -> bool:
+        return self.cache_impl in ["static", "hybrid"]
+
+    @property
+    def use_local_attention(self) -> bool:
+        return self.cache_impl in ["sliding_window", "hybrid"]
+
+    @property
+    def use_multiple_decoder(self) -> bool:
         return isinstance(self.decoder_batch_sizes, list) and len(self.decoder_batch_sizes) > 1
 
     @property
-    def can_generate(self):
+    def can_generate(self) -> bool:
         return "decode" in self.phases
+
+
+class RBLNDecoderOnlyModelForCausalLMConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN decoder-only models for Causal Language Modeling.
+
+    This class extends RBLNModelConfig with parameters specific to decoder-only transformer
+    architectures optimized for RBLN devices. It controls aspects like attention implementation,
+    KV cache management, and batching for inference.
+    """
+
+    _default_phases = ["prefill", "decode"]
+    _default_logits_to_keep = 1