optimum-rbln 0.8.2a4__py3-none-any.whl → 0.9.3__py3-none-any.whl

optimum/rbln/transformers/models/colqwen2/modeling_colqwen2.py

@@ -0,0 +1,446 @@
+# 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 TYPE_CHECKING, Optional, Union
+
+import torch
+from transformers import (
+    PretrainedConfig,
+    PreTrainedModel,
+)
+from transformers.modeling_utils import no_init_weights
+from transformers.models.colqwen2.modeling_colqwen2 import ColQwen2ForRetrievalOutput
+from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+    Qwen2_5_VLModel,
+    Qwen2_5_VLRotaryEmbedding,
+)
+from transformers.models.qwen2_vl.modeling_qwen2_vl import (
+    Qwen2VLModel,
+    Qwen2VLRotaryEmbedding,
+)
+
+from optimum.rbln.transformers.models.decoderonly.modeling_decoderonly import (
+    RBLNDecoderOnlyModel,
+)
+
+from .configuration_colqwen2 import (
+    RBLNColQwen2ForRetrievalConfig,
+)
+
+
+if TYPE_CHECKING:
+    from transformers import (
+        AutoFeatureExtractor,
+        AutoProcessor,
+        AutoTokenizer,
+        PretrainedConfig,
+    )
+
+from .colqwen2_architecture import ColQwen2LanguageModelWrapper
+
+
+class RBLNColQwen2ForRetrieval(RBLNDecoderOnlyModel):
+    """
+    The ColQwen Model transformer for document retrieval using vision-language models.
+    This model inherits from [`RBLNDecoderOnlyModel`]. 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 `ColQwen2ForRetrieval` model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers `ColQwen2ForRetrieval` 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 [`RBLNColQwen2ForRetrievalConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNColQwen2ForRetrievalConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNColQwen2ForRetrievalConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNColQwen2ForRetrieval
+
+        # Using a config dictionary
+        rbln_config = {
+            "visual": {
+                "max_seq_lens": 6400,
+            },
+            "max_seq_len": 32_768,
+            "tensor_parallel_size": 4,
+            "device": [0, 1, 2, 3],
+            "output_hidden_states": False,
+        }
+        model = RBLNColQwen2ForRetrieval.from_pretrained(
+            "vidore/colqwen2-v1.0-hf",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+        # Using a RBLNColQwen2ForRetrievalConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNColQwen2ForRetrievalConfig
+
+        config = RBLNColQwen2ForRetrievalConfig(
+            visual={
+                "max_seq_lens": 6400,
+                "device": 0,
+            },
+            max_seq_len=32_768,
+            tensor_parallel_size=4,
+            device=[0, 1, 2, 3],
+            output_hidden_states=False,
+        )
+        model = RBLNColQwen2ForRetrieval.from_pretrained(
+            "vidore/colqwen2-v1.0-hf",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    main_input_name = "inputs_embeds"
+    auto_model_class = None
+    _rbln_submodules = [
+        {"name": "visual"},
+    ]
+    _decoder_wrapper_cls = ColQwen2LanguageModelWrapper
+    _use_rotary_emb = False
+
+    def __post_init__(self, **kwargs):
+        self.config = self.config.vlm_config if hasattr(self.config, "vlm_config") else self.config
+
+        artifacts = torch.load(
+            self.model_save_dir / self.subfolder / "torch_artifacts.pth",
+            weights_only=False,
+        )
+        self.embed_tokens = self._create_embedding_layer()
+        self.embed_tokens.load_state_dict(artifacts["embed_tokens"])
+        self.visual = self.rbln_submodules[0]
+        self.prefill_runtime = self.model[0]
+        self.mrope_section = self.config.text_config.rope_scaling["mrope_section"]
+        self.is_colqwen2_5 = "qwen2_5_vl" in self.config.model_type
+
+        if self.is_colqwen2_5:
+            self.rotary_emb = Qwen2_5_VLRotaryEmbedding(self.config.text_config)
+        else:
+            self.rotary_emb = Qwen2VLRotaryEmbedding(self.config.text_config)
+        self.block_tables = torch.arange(self.rbln_config.kvcache_num_blocks, dtype=torch.int16)
+
+    @classmethod
+    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+        if hasattr(model, "vlm"):
+            model.visual = model.vlm.visual
+            model.language_model = model.vlm.language_model
+
+        # FIXME: temporary fix for ColQwen2ForRetrieval dtype issue
+        return model.to(torch.float32)
+
+    def _create_embedding_layer(self):
+        with no_init_weights():
+            embed_tokens = torch.nn.Embedding(
+                self.config.text_config.vocab_size,
+                self.config.text_config.hidden_size,
+                self.config.text_config.pad_token_id,
+            )
+        return embed_tokens
+
+    @classmethod
+    def get_input_info(
+        cls,
+        batch_size: int,
+        query_length: int,
+        rbln_config: RBLNColQwen2ForRetrievalConfig,
+        model_config: PretrainedConfig,
+    ):
+        text_config = model_config.text_config
+        input_info = super().get_input_info(
+            batch_size,
+            query_length,
+            rbln_config=rbln_config,
+            model_config=text_config,
+        )
+
+        pos_idx = 3
+        input_info.insert(
+            pos_idx,
+            (
+                "position_emb",
+                [
+                    2,
+                    batch_size,
+                    1,
+                    query_length,
+                    text_config.hidden_size // text_config.num_attention_heads,
+                ],
+                rbln_config.torch_dtype,
+            ),
+        )
+
+        # remove query postion from input_info
+        if "query_position" in input_info:
+            query_position = input_info.pop(4)
+            assert query_position[0] == "query_position", print(query_position[0], "is deleted.")
+        return input_info
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
+        model: Optional["PreTrainedModel"] = None,
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_config: Optional[RBLNColQwen2ForRetrievalConfig] = None,
+    ) -> RBLNColQwen2ForRetrievalConfig:
+        model_config = model_config.vlm_config if hasattr(model_config, "vlm_config") else model_config
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = getattr(model_config.text_config, "output_hidden_states", False)
+
+        return super()._update_rbln_config(
+            preprocessors=preprocessors, model=model, model_config=model_config, rbln_config=rbln_config
+        )
+
+    def _get_position_embeddings(self, hidden_states, position_ids):
+        cos, sin = self.rotary_emb(hidden_states, position_ids)
+        mrope_section = self.mrope_section * 2
+        cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(1)
+        sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(1)
+        return torch.stack([cos, sin])
+
+    def get_rope_index(self, *args, **kwargs):
+        if self.is_colqwen2_5:
+            return Qwen2_5_VLModel.get_rope_index(self, *args, **kwargs)
+        else:
+            return Qwen2VLModel.get_rope_index(self, *args, **kwargs)
+
+    def _preprocess_visual(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: torch.Tensor = None,
+        pixel_values: torch.Tensor = None,
+        pixel_values_videos: torch.FloatTensor = None,
+        image_grid_thw: torch.LongTensor = None,
+        video_grid_thw: torch.LongTensor = None,
+        second_per_grid_ts: torch.Tensor = None,
+    ):
+        batch_size = input_ids.shape[0]
+        inputs_embeds = self.embed_tokens(input_ids)
+
+        if pixel_values is not None:
+            image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
+            n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
+            n_image_features = image_embeds.shape[0]
+            if n_image_tokens != n_image_features:
+                raise ValueError(
+                    f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+                )
+
+            mask = input_ids == self.config.image_token_id
+            mask_unsqueezed = mask.unsqueeze(-1)
+            mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
+
+            image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+            inputs_embeds = inputs_embeds.masked_scatter(mask_expanded, image_embeds)
+
+        if pixel_values_videos is not None:
+            video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
+            n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
+            n_video_features = video_embeds.shape[0]
+            if n_video_tokens != n_video_features:
+                raise ValueError(
+                    f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
+                )
+
+            mask = input_ids == self.config.video_token_id
+            mask_unsqueezed = mask.unsqueeze(-1)
+            mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
+            inputs_embeds = inputs_embeds.masked_scatter(mask_expanded, video_embeds)
+
+        max_inputs_len = input_ids.shape[1]
+        head_dim = self.config.text_config.hidden_size // self.config.text_config.num_attention_heads
+        all_position_embeds = torch.zeros(2, batch_size, 1, max_inputs_len, head_dim)
+        all_rope_deltas = []
+
+        image_token_id = self.config.image_token_id
+        video_token_id = self.config.video_token_id
+        vision_start_token_id = self.config.vision_start_token_id
+        image_idx, video_idx = 0, 0
+
+        for b_idx in range(batch_size):
+            input_id = input_ids[b_idx : b_idx + 1][:, attention_mask[b_idx].bool()]
+            vision_start_indices = torch.argwhere(input_id == vision_start_token_id).squeeze(1)
+            vision_tokens = input_id[0][vision_start_indices + 1]
+            image_nums = (vision_tokens == image_token_id).sum()
+            video_nums = (vision_tokens == video_token_id).sum()
+            args = [
+                input_id,
+                image_grid_thw[image_idx : image_idx + image_nums] if image_grid_thw is not None else None,
+                video_grid_thw[video_idx : video_idx + video_nums] if video_grid_thw is not None else None,
+            ]
+            if self.config.model_type == "qwen2_5_vl":
+                args.append(
+                    second_per_grid_ts[video_idx : video_idx + video_nums] if second_per_grid_ts is not None else None
+                )
+            position_ids, rope_deltas = self.get_rope_index(*args)
+            image_idx += image_nums
+            video_idx += video_nums
+
+            position_embed = self._get_position_embeddings(inputs_embeds, position_ids)
+            mask_indices = torch.nonzero(attention_mask[b_idx], as_tuple=True)[0]
+            all_position_embeds[:, b_idx : b_idx + 1].index_copy_(dim=-2, index=mask_indices, source=position_embed)
+            all_rope_deltas.append(rope_deltas)
+
+        rope_deltas = torch.stack(all_rope_deltas)
+
+        return inputs_embeds, all_position_embeds, rope_deltas
+
+    def _preprocess_chunked_prefill(self, inputs_embeds, attention_mask, position_embed):
+        # valid sequence length of inputs_embeds
+        query_length = inputs_embeds.shape[1] if attention_mask is None else torch.sum(attention_mask.view(-1)).item()
+
+        # extract valid inputs
+        inputs_embeds = inputs_embeds[:, attention_mask.bool()] if attention_mask is not None else inputs_embeds
+        position_embed = (
+            position_embed[:, :, :, attention_mask.bool(), :] if attention_mask is not None else position_embed
+        )
+
+        # add padding for chunked prefill
+        padding_size = (
+            self.rbln_config.prefill_chunk_size - (query_length % self.rbln_config.prefill_chunk_size)
+        ) % self.rbln_config.prefill_chunk_size
+        padded_len = query_length + padding_size
+
+        inputs_embeds = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, padding_size))
+        position_embed = torch.nn.functional.pad(position_embed, (0, 0, 0, padding_size))
+        cache_position = torch.arange(padded_len, dtype=torch.int32).unsqueeze(0)
+
+        return inputs_embeds, position_embed, cache_position, query_length
+
+    def _chunked_prefill_forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_embed: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = False,
+    ):
+        padded_inputs_embeds, padded_position_embed, cache_position, query_length = self._preprocess_chunked_prefill(
+            inputs_embeds, attention_mask, position_embed
+        )
+
+        # Chunked prefill
+        projs = []
+        all_hidden_states = [] if output_hidden_states else None
+        for step in range(0, query_length, self.rbln_config.prefill_chunk_size):
+            # Extract the current chunk of inputs and cache positions
+            input_chunk = padded_inputs_embeds[:, step : step + self.rbln_config.prefill_chunk_size]
+            cache_pos_chunk = cache_position[:, step : step + self.rbln_config.prefill_chunk_size]
+            position_embed_chunk = padded_position_embed[:, :, :, step : step + self.rbln_config.prefill_chunk_size, :]
+
+            # Forward pass for the current chunk
+            proj = self.prefill_runtime(
+                inputs_embeds=input_chunk,
+                cache_position=cache_pos_chunk,
+                block_tables=self.block_tables,
+                position_emb=position_embed_chunk,
+            )
+
+            if output_hidden_states:
+                projs.append(proj[0])
+                all_hidden_states.append(proj[1:])
+            else:
+                projs.append(proj)
+
+        projs = torch.concat(projs, dim=-2)[:, :query_length]
+        if output_hidden_states:
+            # Concatenate chunks for each layer
+            concatenated_hidden_states = [
+                torch.concat(hs_chunks, dim=-2)[:, :query_length] for hs_chunks in list(zip(*all_hidden_states))
+            ]
+            all_hidden_states = tuple(concatenated_hidden_states)
+
+        return self._postprocess_chunked_prefill(projs, attention_mask), all_hidden_states
+
+    def _postprocess_chunked_prefill(self, projs, attention_mask):
+        # index copy for attention mask
+        if attention_mask is not None:
+            embedding = torch.full(
+                (1, attention_mask.shape[-1], projs.shape[-1]),
+                fill_value=1e-10,
+                dtype=projs.dtype,
+            )
+            mask_indices = torch.nonzero(attention_mask, as_tuple=True)[0]
+            embedding.index_copy_(dim=-2, index=mask_indices, source=projs)
+        else:
+            embedding = projs
+        return embedding
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        second_per_grid_ts: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        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."
+            )
+
+        # Handle the custom "pixel_values" input obtained with `ColQwen2Processor` through unpadding
+        if pixel_values is not None and image_grid_thw is not None:
+            offsets = image_grid_thw[:, 1] * image_grid_thw[:, 2]  # (batch_size,)
+            pixel_values = torch.cat(
+                [pixel_sequence[:offset] for pixel_sequence, offset in zip(pixel_values, offsets)],
+                dim=0,
+            )
+        # visual preprocessing
+        inputs_embeds, position_embed, _ = self._preprocess_visual(
+            input_ids,
+            attention_mask,
+            pixel_values,
+            pixel_values_videos,
+            image_grid_thw,
+            video_grid_thw,
+            second_per_grid_ts,
+        )
+        batch_size = inputs_embeds.shape[0]
+
+        projs = []
+        for b_idx in range(batch_size):
+            proj = self._chunked_prefill_forward(
+                inputs_embeds[b_idx : b_idx + 1],
+                attention_mask[b_idx] if attention_mask is not None else None,
+                position_embed[:, b_idx : b_idx + 1],
+                output_hidden_states=output_hidden_states,
+            )
+            projs.append(proj[0])
+            all_hidden_states = proj[1] if output_hidden_states else ()
+
+        # postprocess
+        projs = torch.cat(projs, dim=0)
+        projs = projs / projs.norm(dim=-1, keepdim=True)
+        projs = projs * attention_mask.unsqueeze(-1)
+
+        return ColQwen2ForRetrievalOutput(
+            embeddings=projs,
+            hidden_states=all_hidden_states,
+        )

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

@@ -22,5 +22,6 @@ 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 .configuration_lora import RBLNLoRAAdapterConfig, RBLNLoRAConfig
+from .modeling_decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM

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

@@ -12,29 +12,32 @@
 # 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
 from ...utils.rbln_quantization import RBLNQuantizationConfig
+from .configuration_lora import RBLNLoRAConfig
 
 
 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,
     KV cache management, and batching for inference.
     """
 
+    _default_phases = ["prefill"]
+    _default_logits_to_keep = 0
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -46,12 +49,15 @@ class RBLNDecoderOnlyModelForCausalLMConfig(RBLNModelConfig):
         kvcache_partition_len: Optional[int] = None,
         kvcache_block_size: Optional[int] = None,
         quantization: Optional[Union[Dict[str, Any], RBLNQuantizationConfig]] = None,
+        lora_config: Optional[Union[Dict[str, Any], RBLNLoRAConfig]] = 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,
     ):
         """
@@ -78,6 +84,10 @@ class RBLNDecoderOnlyModelForCausalLMConfig(RBLNModelConfig):
                 section below for details.
             quantization (Optional[Dict[str, Any]]): Configuration dictionary for applying model
                 quantization. Specifies format, etc.
+            lora_config (Optional[Union[Dict[str, Any], RBLNLoRAConfig]]): Configuration for LoRA
+                (Low-Rank Adaptation) settings when using (multi-)LoRA support. Can be provided as
+                a dictionary or an RBLNLoRAConfig instance. When provided, enables LoRA functionality
+                for the model compilation. Defaults to None (no LoRA).
             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.
@@ -98,7 +108,11 @@ class RBLNDecoderOnlyModelForCausalLMConfig(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.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            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:
             ValueError: If `batch_size` is not a positive integer.
@@ -170,54 +184,117 @@ class RBLNDecoderOnlyModelForCausalLMConfig(RBLNModelConfig):
         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.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.lora_config = lora_config
+        if self.lora_config and isinstance(self.lora_config, dict):
+            self.lora_config = RBLNLoRAConfig(**self.lora_config)
 
-        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]
+        # Validate LoRA adapters if LoRA is enabled
+        if self.lora_config is not None:
+            validation_results = self.lora_config.validate_adapter_weights()
+            failed_adapters = [adapter_id for adapter_id, is_valid in validation_results.items() if not is_valid]
 
-        if self.use_multiple_decoder:
-            if max(self.decoder_batch_sizes) > self.batch_size:
+            if failed_adapters:
                 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})."
+                    f"Some LoRA adapters failed validation and may not be accessible at compile time: {failed_adapters}. "
+                    "Please ensure all adapter weights are available and properly formatted."
                 )
-            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)
+            logger.info(
+                f"LoRA configuration initialized with {self.lora_config.num_adapters} adapters: "
+                f"{self.lora_config.adapter_ids}. Max rank: {self.lora_config.max_lora_rank}"
+            )
+
+        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.cache_impl = cache_impl or "static"
         self.sliding_window = sliding_window
         self.sliding_window_layers = sliding_window_layers or []
 
+        if phases is not None:
+            self.validate_phases_type(phases)
+        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.")
+
+        self.decoder_batch_sizes = None
+        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):
+    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 use_lora(self):
+        return self.lora_config is not None
+
+    @property
+    def can_generate(self) -> bool:
+        return "decode" in self.phases
+
+    @property
+    def nbits_per_param(self) -> int:
+        if self.quantization:
+            return self.quantization.nbits_per_param
+        return 16
+
+
+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