optimum-rbln 0.8.1rc1__py3-none-any.whl → 0.8.2__py3-none-any.whl

optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import torch
 from transformers import CLIPTextConfig, CLIPTextModel, CLIPVisionConfig, CLIPVisionModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.models.clip.modeling_clip import CLIPTextModelOutput, CLIPVisionModelOutput
 
 from ....configuration_utils import RBLNCompileConfig
@@ -111,12 +112,27 @@ class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
 
 
 class _VisionEncoder(torch.nn.Module):
-    def __init__(self, enc: CLIPVisionModel):
+    def __init__(
+        self,
+        enc: CLIPVisionModel,
+        interpolate_pos_encoding: bool,
+        output_hidden_states: bool,
+        output_attentions: bool,
+    ):
         super().__init__()
         self.enc = enc
+        self.interpolate_pos_encoding = interpolate_pos_encoding
+        self.output_hidden_states = output_hidden_states
+        self.output_attentions = output_attentions
 
     def forward(self, inp):
-        enc_out = self.enc(inp, output_hidden_states=True, return_dict=False)
+        enc_out = self.enc(
+            inp,
+            output_hidden_states=self.output_hidden_states,
+            interpolate_pos_encoding=self.interpolate_pos_encoding,
+            output_attentions=self.output_attentions,
+            return_dict=False,
+        )
         return enc_out
 
 
@@ -130,7 +146,12 @@ class RBLNCLIPVisionModel(RBLNModel):
 
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
-        return _VisionEncoder(model).eval()
+        wrapper_cfg = {
+            "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
+            "output_hidden_states": rbln_config.output_hidden_states,
+            "output_attentions": rbln_config.output_attentions,
+        }
+        return _VisionEncoder(model, **wrapper_cfg).eval()
 
     @classmethod
     def update_rbln_config_using_pipe(
@@ -155,6 +176,12 @@ class RBLNCLIPVisionModel(RBLNModel):
         if rbln_config.image_size is None:
             raise ValueError("`rbln_image_size` should be specified!")
 
+        if rbln_config.output_attentions is None:
+            rbln_config.output_attentions = getattr(model_config, "output_attentions", False)
+
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
+
         rbln_compile_config = RBLNCompileConfig(
             input_info=[
                 (
@@ -176,27 +203,76 @@ class RBLNCLIPVisionModel(RBLNModel):
     def forward(
         self,
         pixel_values: Optional[torch.FloatTensor] = None,
-        return_dict: bool = None,
+        return_dict: bool = True,
+        output_attentions: bool = None,
+        output_hidden_states: bool = None,
+        interpolate_pos_encoding: bool = False,
         **kwargs,
-    ) -> Union[Tuple, CLIPVisionModelOutput]:
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
         if len(kwargs) > 0 and any(value is not None for value in kwargs.values()):
             logger.warning(
                 f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__.__name__}."
             )
+
+        output_attentions = output_attentions if output_attentions is not None else self.rbln_config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        if output_attentions != self.rbln_config.output_attentions:
+            raise ValueError(
+                f"Variable output_attentions {output_attentions} is not equal to rbln_config.output_attentions {self.rbln_config.output_attentions} "
+                f"Please compile again with the correct argument."
+            )
+
+        if output_hidden_states != self.rbln_config.output_hidden_states:
+            raise ValueError(
+                f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
+                f"Please compile again with the correct argument."
+            )
+
+        if interpolate_pos_encoding != self.rbln_config.interpolate_pos_encoding:
+            raise ValueError(
+                f"Variable interpolate_pos_encoding {interpolate_pos_encoding} is not equal to rbln_config.interpolate_pos_encoding {self.rbln_config.interpolate_pos_encoding} "
+                f"Please compile again with the correct argument."
+            )
+
         output = super().forward(pixel_values, return_dict=return_dict)
         return output
 
     def _prepare_output(self, output, return_dict):
         # Prepare model output based on return_dict flag.
         # This method can be overridden by subclasses to provide task-specific output handling.
+        last_hidden_state = output.pop(0)
+        pooler_output = output.pop(0)
+        vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+
+        if self.rbln_config.output_hidden_states:
+            hidden_states = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers + 1):
+                hidden_states += (output.pop(0),)
+        else:
+            hidden_states = None
+
+        if self.rbln_config.output_attentions:
+            attentions = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers):
+                attentions += (output.pop(0),)
+        else:
+            attentions = None
 
         if not return_dict:
-            return (output,) if not isinstance(output, (tuple, list)) else output
+            return tuple(
+                item for item in (last_hidden_state, pooler_output, hidden_states, attentions) if item is not None
+            )
         else:
-            return CLIPVisionModelOutput(
-                image_embeds=output[0],
-                last_hidden_state=output[1],
-                hidden_states=output[2:],
+            return BaseModelOutputWithPooling(
+                last_hidden_state=last_hidden_state,
+                pooler_output=pooler_output,
+                hidden_states=hidden_states,
+                attentions=attentions,
             )
 
 
@@ -208,21 +284,40 @@ class RBLNCLIPVisionModelWithProjection(RBLNCLIPVisionModel):
     multimodal embedding alignment tasks.
     """
 
-    def forward(
-        self,
-        pixel_values: Optional[torch.FloatTensor] = None,
-        **kwargs,
-    ) -> Union[Tuple, CLIPVisionModelOutput]:
-        if len(kwargs) > 0 and any(kwargs.values()):
-            logger.warning(f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__}.")
-
-        output = super().forward(pixel_values)
-        image_embeds = output[0]
-        last_hidden_state = output[1]
-        hidden_states = output[2:]
-
-        return CLIPVisionModelOutput(
-            image_embeds=image_embeds,
-            last_hidden_state=last_hidden_state,
-            hidden_states=hidden_states,
-        )
+    def _prepare_output(self, output, return_dict):
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
+
+        image_embeds = output.pop(0) if isinstance(output, (tuple, list)) else output
+        last_hidden_state = output.pop(0)
+
+        vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+
+        if self.rbln_config.output_hidden_states:
+            hidden_states = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers + 1):
+                hidden_states += (output.pop(0),)
+        else:
+            hidden_states = None
+
+        if self.rbln_config.output_attentions:
+            attentions = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers):
+                attentions += (output.pop(0),)
+        else:
+            attentions = None
+
+        if not return_dict:
+            return tuple(
+                item for item in (image_embeds, last_hidden_state, hidden_states, attentions) if item is not None
+            )
+
+        else:
+            return CLIPVisionModelOutput(
+                image_embeds=image_embeds,
+                last_hidden_state=last_hidden_state,
+                hidden_states=hidden_states,
+                attentions=attentions,
+            )

optimum/rbln/transformers/models/colpali/colpali_architecture.py

@@ -4,10 +4,7 @@ import torch
 from torch import nn
 from transformers import GemmaForCausalLM, GemmaModel
 
-from ..decoderonly.decoderonly_architecture import (
-    RotaryEmbedding,
-    apply_rotary_pos_emb,
-)
+from ..decoderonly.decoderonly_architecture import RotaryEmbedding, apply_rotary_pos_emb
 
 
 def slice_and_unsqueeze_cos_sin(cos, sin, position_ids):

optimum/rbln/transformers/models/colpali/configuration_colpali.py

@@ -11,7 +11,7 @@
 # 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 List, Optional, Union
+from typing import Any, List, Optional, Union
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -50,7 +50,7 @@ class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
         max_seq_lens: Union[int, List[int]] = None,
         output_hidden_states: Optional[bool] = None,
         vision_tower: Optional[RBLNModelConfig] = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         """
         Args:

optimum/rbln/transformers/models/colpali/modeling_colpali.py

@@ -17,10 +17,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 import torch
-from transformers import (
-    PretrainedConfig,
-    PreTrainedModel,
-)
+from transformers import PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.modeling_utils import no_init_weights
 from transformers.models.colpali.modeling_colpali import ColPaliForRetrievalOutput
@@ -32,12 +29,7 @@ from .colpali_architecture import RBLNColPaliForRetrievalWrapper
 
 
 if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
-        PretrainedConfig,
-    )
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
 
 
 class LoopVisionTower:

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

@@ -22,5 +22,5 @@ from ....ops import (
     paged_flash_causal_attn_decode,
     paged_flash_causal_attn_prefill,
 )
-from .configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
-from .modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM
+from .configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
+from .modeling_decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM

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

@@ -12,9 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, List, Literal, Optional, Union
-
-import rebel
+from typing import Any, Dict, List, Literal, Optional, Union, get_args
 
 from ....configuration_utils import RBLNModelConfig
 from ....utils.logging import get_logger
@@ -24,11 +22,12 @@ from ...utils.rbln_quantization import RBLNQuantizationConfig
 logger = get_logger()
 
 CacheImplType = Literal["static", "sliding_window", "hybrid"]
+PhaseType = Literal["prefill", "image_prefill", "decode"]
 
 
-class RBLNDecoderOnlyModelForCausalLMConfig(RBLNModelConfig):
+class RBLNDecoderOnlyModelConfig(RBLNModelConfig):
     """
-    Configuration class for RBLN decoder-only models for Causal Language Modeling.
+    Configuration class for RBLN decoder-only models.
 
     This class extends RBLNModelConfig with parameters specific to decoder-only transformer
     architectures optimized for RBLN devices. It controls aspects like attention implementation,
@@ -48,7 +47,6 @@ class RBLNDecoderOnlyModelForCausalLMConfig(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,
@@ -76,20 +74,12 @@ class RBLNDecoderOnlyModelForCausalLMConfig(RBLNModelConfig):
             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.
@@ -166,58 +156,237 @@ class RBLNDecoderOnlyModelForCausalLMConfig(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
         self.use_position_ids = use_position_ids or False
-        self.use_attention_mask = use_attention_mask
-
-        npu = self.npu or rebel.get_npu_name()
-        if npu == "RBLN-CA02":
-            if self.use_attention_mask is False:
-                logger.warning("Attention mask should be used with RBLN-CA02. Setting use_attention_mask to True.")
-            self.use_attention_mask = True
-        else:
-            self.use_attention_mask = self.use_attention_mask or False
+        self.use_attention_mask = use_attention_mask or False
 
         if self.use_position_ids and not self.use_attention_mask:
             raise ValueError("Position IDs should be used with attention mask.")
 
-        self.attn_impl = attn_impl
-        self.kvcache_partition_len = kvcache_partition_len
-        self.kvcache_block_size = kvcache_block_size
         self.quantization = quantization or {}
         if self.quantization and isinstance(self.quantization, dict):
             self.quantization = RBLNQuantizationConfig(**self.quantization)
 
+        self.attn_impl = attn_impl
+        self.kvcache_partition_len = kvcache_partition_len
+        self.kvcache_block_size = kvcache_block_size
         self.prefill_chunk_size = prefill_chunk_size or 128
         if self.prefill_chunk_size % 64 != 0 or self.prefill_chunk_size <= 0:
             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)
-
         self.cache_impl = cache_impl or "static"
         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"]
+
+        if "decode" in self.phases:
+            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)
+
+    @staticmethod
+    def validate_phases_type(phases: List[PhaseType]):
+        if not isinstance(phases, list):
+            raise ValueError("`phases` must be a list.")
+        if not all(phase in get_args(PhaseType) for phase in phases):
+            raise ValueError(f"All elements in `phases` must be of type `PhaseType`({get_args(PhaseType)}).")
+
     @property
     def use_multiple_decoder(self):
         return isinstance(self.decoder_batch_sizes, list) and len(self.decoder_batch_sizes) > 1
+
+    @property
+    def can_generate(self):
+        return "decode" in self.phases