optimum-rbln 0.9.2a3__py3-none-any.whl → 0.9.2a5__py3-none-any.whl

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

@@ -187,6 +187,8 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
                 torch.ones(1, 1, self.rbln_config.prefill_chunk_size, self.rbln_config.prefill_chunk_size), diagonal=1
             )
 
+        self.lora_int_ids = None
+
     def inputs_embeddings_if_needed(
         self, input_ids: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None
     ):
@@ -210,6 +212,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         position_ids: Optional[torch.Tensor] = None,
         token_type_ids: Optional[torch.Tensor] = None,
         local_block_tables: Optional[torch.Tensor] = None,
+        lora_int_ids: Optional[torch.Tensor] = None,
     ):
         inputs = self.inputs_embeddings_if_needed(input_ids, inputs_embeds)
         block_tables, local_block_tables, is_external_block_tables = (
@@ -233,6 +236,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
                 position_embed=position_embed,
                 position_ids=position_ids,
                 local_block_tables=local_block_tables,
+                lora_int_ids=lora_int_ids,
             )
         else:
             return self.prefill_forward(
@@ -245,6 +249,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
                 position_embed=position_embed,
                 token_type_ids=token_type_ids,
                 local_block_tables=local_block_tables,
+                lora_int_ids=lora_int_ids,
             )
 
     def decode_forward(
@@ -257,7 +262,20 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         position_embed: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
         local_block_tables: Optional[torch.Tensor] = None,
+        lora_int_ids: Optional[torch.Tensor] = None,
     ) -> torch.FloatTensor:
+        if self.rbln_config.use_lora and lora_int_ids is None:
+            if self.lora_int_ids is None:
+                raise ValueError(
+                    "lora_int_id is required when using LoRA. "
+                    "You should call set_lora_int_ids() before forward() or pass lora_int_id to forward()."
+                )
+
+            lora_int_ids = self.lora_int_ids
+
+        if lora_int_ids is not None and lora_int_ids.shape[0] != self.batch_size:
+            raise ValueError(f"lora_int_ids size mismatch: got {lora_int_ids.shape[0]}, expected {self.batch_size}.")
+
         if self.batch_size != cache_position.shape[0]:
             raise RuntimeError(
                 f"Cache position size mismatch: got {cache_position.shape[0]}, expected {self.batch_size}."
@@ -287,6 +305,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             position_embed,
             attention_mask if self.rbln_config.use_attention_mask else None,
             position_ids if self.rbln_config.use_position_ids else None,
+            lora_int_ids if self.rbln_config.use_lora else None,
         )
 
         return RBLNDecoderOnlyOutput(logits=logits)
@@ -369,12 +388,25 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         position_embed: Optional[torch.Tensor] = None,
         token_type_ids: Optional[torch.Tensor] = None,
         local_block_tables: Optional[torch.Tensor] = None,
+        lora_int_ids: Optional[torch.Tensor] = None,
     ) -> torch.FloatTensor:
         """
         Performs chunked prefill for efficient KV-cache updates and memory optimization.
         Instead of processing the entire sequence at once, the input is divided into chunks of size `prefill_chunk_size`,
         and each chunk is processed sequentially. This allows for better memory utilization and compatibility with continuous batching.
         """
+        if self.rbln_config.use_lora and lora_int_ids is None:
+            if self.lora_int_ids is None:
+                raise ValueError(
+                    "lora_int_id is required when using LoRA. "
+                    "You should call set_lora_int_ids() before forward() or pass lora_int_id to forward()."
+                )
+
+            if batch_idx is not None:
+                lora_int_ids = self.lora_int_ids[batch_idx : batch_idx + 1].clone()
+            else:
+                lora_int_ids = self.lora_int_ids.clone()
+
         (
             inputs,
             cache_position,
@@ -426,6 +458,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
                 query_position,
                 chunked_attention_mask if self.rbln_config.use_attention_mask else None,
                 position_ids_chunk,
+                lora_int_ids if self.rbln_config.use_lora else None,
                 out=self.out_buffers,
             )
             output_logits.append(output_logit)

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

@@ -0,0 +1,204 @@
+import math
+from pathlib import Path
+from typing import Optional
+
+import safetensors.torch
+import torch
+from torch import nn
+
+from ....utils import logging
+from .configuration_lora import RBLNLoRAConfig
+
+
+logger = logging.get_logger()
+
+
+class LoRALinear(nn.Module):
+    """
+    A linear layer that supports multiple LoRA adapters compiled at static time.
+
+    This class replaces the original linear layer and handles both base weights
+    and multiple LoRA adapters in a single forward pass using custom ops.
+    """
+
+    def __init__(
+        self,
+        original_linear: nn.Linear,
+        lora_config: RBLNLoRAConfig,
+        projection_name: str = "",
+        layer_idx: int = 0,
+    ):
+        """
+        Args:
+            original_linear: The original linear layer to be replaced
+            lora_config: LoRA configuration containing all adapters
+            projection_name: Name of the projection (e.g., "q_proj", "k_proj")
+            layer_idx: Layer index for loading the correct LoRA weights
+        """
+        super().__init__()
+
+        self.in_features = original_linear.in_features
+        self.out_features = original_linear.out_features
+        self.projection_name = projection_name
+        self.layer_idx = layer_idx
+        self.lora_config = lora_config
+
+        # Store original linear weights and bias directly without cloning
+        self.register_buffer("weight", original_linear.weight.data)
+        if original_linear.bias is not None:
+            self.register_buffer("bias", original_linear.bias.data)
+        else:
+            self.bias = None
+
+        # Initialize LoRA weights
+        self._init_lora_weights()
+
+    def _should_apply_lora(self) -> bool:
+        """Check if this projection should have LoRA applied."""
+        # Check if any adapter targets this projection
+        return any(self.projection_name in adapter.target_modules for adapter in self.lora_config.adapters)
+
+    def _load_adapter_weights(self, adapter_path: Path):
+        """
+        Load adapter weights from local directory.
+
+        Args:
+            adapter_path: Path to local directory containing adapter weights
+
+        Returns:
+            Dictionary containing adapter weights
+
+        Raises:
+            FileNotFoundError: If no adapter weights are found in the directory
+        """
+        if not adapter_path.is_dir():
+            raise ValueError(f"Adapter path must be a directory, got: {adapter_path}")
+
+        # Try to load weights in order of preference
+        weight_files = [
+            ("adapter_model.safetensors", lambda p: safetensors.torch.load_file(p)),
+            ("adapter_model.bin", lambda p: torch.load(p, map_location="cpu")),
+            ("pytorch_model.bin", lambda p: torch.load(p, map_location="cpu")),
+        ]
+
+        for filename, load_fn in weight_files:
+            weight_path = adapter_path / filename
+            if weight_path.exists():
+                return load_fn(weight_path)
+
+        raise FileNotFoundError(
+            f"No adapter weights found in {adapter_path}. "
+            f"Expected one of: {', '.join(filename for filename, _ in weight_files)}"
+        )
+
+    def _init_lora_weights(self):
+        """Initialize LoRA adapter weights by loading and stacking them."""
+
+        lora_a_weights = []
+        lora_b_weights = []
+
+        for adapter in self.lora_config.adapters:
+            if self.projection_name not in adapter.target_modules:
+                # Create zero weights for adapters that don't target this projection
+                lora_a_weights.append(torch.zeros(adapter.r, self.in_features))
+                lora_b_weights.append(torch.zeros(self.out_features, adapter.r))
+                continue
+
+            adapter_weights = self._load_adapter_weights(adapter.local_adapter_path)
+
+            # Determine module type from projection name
+            attn_projs = {"q_proj", "k_proj", "v_proj", "o_proj"}
+            mlp_projs = {"gate_proj", "up_proj", "down_proj"}
+            if self.projection_name in attn_projs:
+                module_type = "self_attn"
+            elif self.projection_name in mlp_projs:
+                module_type = "mlp"
+            else:
+                module_type = "self_attn"
+
+            layer_key = f"base_model.model.model.layers.{self.layer_idx}.{module_type}.{self.projection_name}"
+            lora_a_key = f"{layer_key}.lora_A.weight"
+            lora_b_key = f"{layer_key}.lora_B.weight"
+
+            if lora_a_key in adapter_weights and lora_b_key in adapter_weights:
+                # Calculate scaling factor and fold it into lora_b
+                scaling = adapter.lora_alpha / adapter.r
+                if adapter.use_rslora:
+                    scaling = scaling / math.sqrt(adapter.r)
+                scaling = scaling * adapter.scaling_factor
+
+                lora_a_weights.append(adapter_weights[lora_a_key])
+                # scaling is pre-applied to lora_b_weights
+                lora_b_weights.append(adapter_weights[lora_b_key] * scaling)
+            else:
+                logger.warning(f"No LoRA weights found for {lora_a_key} or {lora_b_key}")
+                lora_a_weights.append(torch.zeros(adapter.r, self.in_features))
+                lora_b_weights.append(torch.zeros(self.out_features, adapter.r))
+
+        # Stack weights along adapter dimension
+        max_rank = self.lora_config.max_lora_rank
+
+        # Pad smaller ranks to max_rank
+        padded_lora_a = []
+        padded_lora_b = []
+
+        for i, (lora_a, lora_b) in enumerate(zip(lora_a_weights, lora_b_weights)):
+            current_rank = lora_a.shape[0]
+            if current_rank < max_rank:
+                # Pad with zeros
+                padded_a = torch.zeros(max_rank, self.in_features)
+                padded_b = torch.zeros(self.out_features, max_rank)
+                padded_a[:current_rank] = lora_a
+                padded_b[:, :current_rank] = lora_b
+                padded_lora_a.append(padded_a)
+                padded_lora_b.append(padded_b)
+            else:
+                padded_lora_a.append(lora_a)
+                padded_lora_b.append(lora_b)
+
+        lora_a_transposed = [lora_a.transpose(0, 1) for lora_a in padded_lora_a]  # [in_features, rank]
+        lora_b_transposed = [lora_b.transpose(0, 1) for lora_b in padded_lora_b]  # [rank, out_features]
+
+        self.register_buffer(
+            "lora_a_weights", torch.stack(lora_a_transposed, dim=0).to(self.weight.dtype)
+        )  # [num_adapters, in_features, rank]
+        self.register_buffer(
+            "lora_b_weights", torch.stack(lora_b_transposed, dim=0).to(self.weight.dtype)
+        )  # [num_adapters, rank, out_features]
+
+    def forward(self, x: torch.Tensor, lora_int_id: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Forward pass that combines base linear transformation with LoRA.
+
+        Args:
+            x: Input tensor [batch_size, seq_len, in_features]
+            lora_int_id: Adapter ID tensor [batch_size] indicating which adapter to use
+
+        Returns:
+            Output tensor [batch_size, seq_len, out_features]
+        """
+        # Base linear transformation
+        output = torch.nn.functional.linear(x, self.weight, self.bias)
+
+        # Apply LoRA if enabled and adapter ID is provided
+        if self._should_apply_lora() and lora_int_id is not None:
+            # Gather LoRA weights for each batch item
+            # lora_int_id: [batch_size] -> use as indices to select weights
+            selected_lora_a = self.lora_a_weights[lora_int_id]  # [batch_size, in_features, rank]
+            selected_lora_b = self.lora_b_weights[lora_int_id]  # [batch_size, rank, out_features]
+
+            # Batched matrix multiplication for LoRA computation
+            # x: [batch_size, seq_len, in_features]
+            # selected_lora_a: [batch_size, in_features, rank] (already transposed)
+            # selected_lora_b: [batch_size, rank, out_features] (already transposed)
+
+            # First matmul: x @ lora_a -> [batch_size, seq_len, rank]
+            temp = torch.bmm(x, selected_lora_a)
+
+            # Second matmul: temp @ lora_b -> [batch_size, seq_len, out_features]
+            lora_delta = torch.bmm(temp, selected_lora_b)
+
+            # Add LoRA delta to base output
+            output = output + lora_delta
+
+        return output

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

@@ -19,7 +19,7 @@ from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union
 import rebel
 import torch
 from rebel.compile_context import CompileContext
-from transformers import AutoModel, AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutputWithPast
 from transformers.modeling_utils import no_init_weights
 
@@ -317,12 +317,27 @@ class RBLNDecoderOnlyModel(RBLNModel, RBLNDecoderOnlyFlashAttentionMixin):
 
     @classmethod
     def get_pytorch_model(
-        cls, *args, rbln_config: Optional[RBLNDecoderOnlyModelConfig] = None, **kwargs
+        cls,
+        model_id: str,
+        *args,
+        rbln_config: Optional[RBLNDecoderOnlyModelConfig] = None,
+        num_hidden_layers: Optional[int] = None,
+        **kwargs,
     ) -> PreTrainedModel:
         if rbln_config and rbln_config.quantization:
-            model = cls.get_quantized_model(*args, rbln_config=rbln_config, **kwargs)
+            model = cls.get_quantized_model(model_id, *args, rbln_config=rbln_config, **kwargs)
         else:
-            model = super().get_pytorch_model(*args, **kwargs)
+            if num_hidden_layers is not None:
+                trust_remote_code = kwargs.get("trust_remote_code", None)
+                config, kwargs = AutoConfig.from_pretrained(
+                    model_id, return_unused_kwargs=True, num_hidden_layers=num_hidden_layers, **kwargs
+                )
+                if hasattr(config, "layer_types"):
+                    config.layer_types = config.layer_types[:num_hidden_layers]
+                kwargs["config"] = config
+                kwargs["trust_remote_code"] = trust_remote_code
+
+            model = super().get_pytorch_model(model_id, *args, **kwargs)
 
         return model
 
@@ -375,6 +390,9 @@ class RBLNDecoderOnlyModel(RBLNModel, RBLNDecoderOnlyFlashAttentionMixin):
         if rbln_config.use_position_ids:
             input_info.append(("position_ids", [batch_size, query_length], "int32"))
 
+        if rbln_config.use_lora:
+            input_info.append(("lora_int_ids", [batch_size], "int32"))
+
         kvcache_dtype = rbln_config.torch_dtype
         if rbln_config.quantization and rbln_config.quantization.kv_caches == "fp8":
             kvcache_dtype = "float8_e4m3fn"
@@ -667,6 +685,53 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNDecoderOnlyModel, RBLNDecoderOnlyGener
     def use_query_position(cls, use_local_attention: bool, is_prefill: bool = True):
         return is_prefill
 
+    def set_lora_int_ids(self, lora_int_ids: Optional[torch.Tensor]):
+        if isinstance(lora_int_ids, int):
+            lora_int_ids = torch.tensor([lora_int_ids], dtype=torch.int32)
+        elif isinstance(lora_int_ids, list):
+            lora_int_ids = torch.tensor(lora_int_ids, dtype=torch.int32)
+
+        self.lora_int_ids = lora_int_ids
+
+        self.prefill_decoder.lora_int_ids = lora_int_ids
+        if self.rbln_config.can_generate:
+            for batch_size in self.rbln_config.decoder_batch_sizes:
+                self.decoders[batch_size].lora_int_ids = lora_int_ids
+
+    def set_adapter(self, adapter_name: Union[str, List[str]]) -> None:
+        """
+        Sets the active adapter(s) for the model using adapter name(s).
+
+        Args:
+            adapter_name (Union[str, List[str]]): The name(s) of the adapter(s) to be activated.
+                Can be a single adapter name or a list of adapter names.
+
+        Raises:
+            ValueError: If the model is not configured with LoRA or if the adapter name is not found.
+        """
+        if not hasattr(self.rbln_config, "lora_config") or self.rbln_config.lora_config is None:
+            raise ValueError("Model is not configured with LoRA. Cannot set adapter.")
+
+        # Convert single adapter name to list for uniform processing
+        if isinstance(adapter_name, str):
+            adapter_names = [adapter_name]
+        else:
+            adapter_names = adapter_name
+
+        # Validate that all adapter names exist
+        available_adapters = {
+            adapter.lora_name: adapter.lora_int_id for adapter in self.rbln_config.lora_config.adapters
+        }
+        missing_adapters = [name for name in adapter_names if name not in available_adapters]
+        if missing_adapters:
+            raise ValueError(
+                f"Adapter(s) {missing_adapters} not found. Available adapters: {list(available_adapters.keys())}"
+            )
+
+        # Get the adapter IDs and set them
+        lora_int_ids = [available_adapters[name] for name in adapter_names]
+        self.set_lora_int_ids(torch.tensor(lora_int_ids, dtype=torch.int32))
+
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,
@@ -677,6 +742,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNDecoderOnlyModel, RBLNDecoderOnlyGener
         padded_cache_lengths: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
         token_type_ids: Optional[torch.Tensor] = None,
+        lora_int_ids: Optional[torch.Tensor] = None,
         return_dict: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> Tuple[torch.FloatTensor]:
@@ -684,6 +750,13 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNDecoderOnlyModel, RBLNDecoderOnlyGener
         # For continuous batching, the prefill stage processes one batch at a time and updates the KV cache using batch_idx.
         # A for-loop ensures synchronization with the HuggingFace generate API.
         # The decoder stage operates as usual, processing inputs in batch mode.
+        if self.rbln_config.use_lora and lora_int_ids is None:
+            if self.lora_int_ids is None:
+                raise ValueError(
+                    "lora_int_id is required when using LoRA. "
+                    "You should call set_lora_int_ids() before forward() or pass lora_int_id to forward()."
+                )
+            lora_int_ids = self.lora_int_ids
 
         # for only use forward
         if generate_idx is None:
@@ -708,6 +781,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNDecoderOnlyModel, RBLNDecoderOnlyGener
                     cache_position=cache_position,
                     batch_idx=b_idx,
                     token_type_ids=token_type_ids[b_idx : b_idx + 1] if token_type_ids is not None else None,
+                    lora_int_ids=lora_int_ids[b_idx : b_idx + 1] if lora_int_ids is not None else None,
                 )
                 padded_cache_lengths[b_idx] += output.padded_cache_lengths
                 logits.append(output.logits)
@@ -727,6 +801,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNDecoderOnlyModel, RBLNDecoderOnlyGener
                 inputs_embeds=inputs_embeds,
                 cache_position=cache_position,
                 position_ids=position_ids if self.rbln_config.use_position_ids else None,
+                lora_int_ids=lora_int_ids,
             ).logits
 
         if not return_dict:

optimum/rbln/transformers/models/gemma3/gemma3_architecture.py

@@ -63,6 +63,7 @@ class Gemma3TextModel(DecoderOnlyModel):
         rotary_emb: torch.nn.Module = None,
         global_block_tables: Optional[torch.Tensor] = None,
         local_block_tables: Optional[torch.Tensor] = None,
+        lora_int_id: Optional[torch.Tensor] = None,
     ):
         # retrieve input_ids and inputs_embeds
         if (input_ids is None) ^ (inputs_embeds is not None):
@@ -105,6 +106,7 @@ class Gemma3TextModel(DecoderOnlyModel):
                 cos=cos_local if is_sliding else cos_global,
                 sin=sin_local if is_sliding else sin_global,
                 block_tables=local_block_tables if is_sliding else global_block_tables,
+                lora_int_id=lora_int_id,
             )
 
         hidden_states = self.get_last_layernorm()(hidden_states)
@@ -127,12 +129,20 @@ class Gemma3DecoderLayer(DecoderOnlyLayer):
         cos: Optional[torch.Tensor] = None,
         sin: Optional[torch.Tensor] = None,
         block_tables: Optional[torch.Tensor] = None,
+        lora_int_id: Optional[torch.Tensor] = None,
     ):
         residual = hidden_states
         hidden_states = self.get_pre_attention_layernorm()(hidden_states)
 
         hidden_states = self.self_attn(
-            hidden_states, attention_mask, seq_positions, past_key_values, cos, sin, block_tables
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            seq_positions=seq_positions,
+            past_key_values=past_key_values,
+            cos=cos,
+            sin=sin,
+            block_tables=block_tables,
+            lora_int_id=lora_int_id,
         )
         hidden_states = self.get_post_attention_layernorm()(hidden_states)
         hidden_states = residual + hidden_states
@@ -140,7 +150,7 @@ class Gemma3DecoderLayer(DecoderOnlyLayer):
         # Fully Connected
         residual = hidden_states
         hidden_states = self.get_pre_feedforward_layernorm()(hidden_states)
-        hidden_states = self._original_mod.mlp(hidden_states)
+        hidden_states = self.forward_mlp(hidden_states, lora_int_id)
         hidden_states = self.get_post_feedforward_layernorm()(hidden_states)
         hidden_states = residual + hidden_states
 

optimum/rbln/transformers/models/gemma3/gemma3_runtime_utils.py

@@ -17,15 +17,16 @@ import rebel
 import torch
 
 from ...modeling_outputs import RBLNDecoderOnlyOutput, RBLNGemma3ForCausalLMOutput
+from ..decoderonly.decoderonly_runtime_utils import RBLNPytorchRuntime
 from ..decoderonly.modeling_decoderonly import RBLNRuntimeModel
 
 
 class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
     def __init__(self, *args, image_prefill: Optional[rebel.Runtime] = None, **kwargs):
         super().__init__(*args, **kwargs)
-        self.image_prefill = image_prefill  # FIXME(taehoon)
-        self.prefill = self.runtime if self.phase == "prefill" else None  # FIXME
-        self.decode = self.runtime if self.phase == "decode" else None
+        self.image_prefill = RBLNPytorchRuntime(image_prefill)  # FIXME(taehoon)
+        self.prefill = RBLNPytorchRuntime(self.runtime) if self.phase == "prefill" else None  # FIXME
+        self.decode = RBLNPytorchRuntime(self.runtime) if self.phase == "decode" else None
 
     def _prepare_prefill_inputs(self, *args, **kwargs):
         (
@@ -73,12 +74,24 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
         position_embed: Optional[torch.Tensor] = None,
         token_type_ids: Optional[torch.Tensor] = None,
         local_block_tables: Optional[torch.Tensor] = None,
+        lora_int_ids: Optional[torch.Tensor] = None,
     ) -> torch.FloatTensor:
         """
         Performs chunked prefill for efficient KV-cache updates and memory optimization.
         Instead of processing the entire sequence at once, the input is divided into chunks of size `prefill_chunk_size`,
         and each chunk is processed sequentially. This allows for better memory utilization and compatibility with continuous batching.
         """
+        if self.rbln_config.use_lora and lora_int_ids is None:
+            if self.lora_int_ids is None:
+                raise ValueError(
+                    "lora_int_id is required when using LoRA. "
+                    "You should call set_lora_int_ids() before forward() or pass lora_int_id to forward()."
+                )
+            if batch_idx is not None:
+                lora_int_ids = self.lora_int_ids[batch_idx : batch_idx + 1].clone()
+            else:
+                lora_int_ids = self.lora_int_ids.clone()
+
         (
             inputs,
             cache_position,
@@ -141,6 +154,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
                     query_position,
                     chunked_attention_mask,
                     position_ids_chunk,
+                    lora_int_ids if self.rbln_config.use_lora else None,
                 )
             else:
                 logits = self.prefill(
@@ -151,6 +165,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
                     query_position,
                     chunked_attention_mask,
                     position_ids_chunk,
+                    lora_int_ids if self.rbln_config.use_lora else None,
                 )
 
             padded_cache_lengths += current_padded_cache_lengths
@@ -173,7 +188,20 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
         position_embed: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
         local_block_tables: Optional[torch.Tensor] = None,
+        lora_int_ids: Optional[torch.Tensor] = None,
     ) -> torch.FloatTensor:
+        if self.rbln_config.use_lora and lora_int_ids is None:
+            if self.lora_int_ids is None:
+                raise ValueError(
+                    "lora_int_id is required when using LoRA. "
+                    "You should call set_lora_int_ids() before forward() or pass lora_int_id to forward()."
+                )
+
+            lora_int_ids = self.lora_int_ids
+
+        if lora_int_ids is not None and lora_int_ids.shape[0] != self.batch_size:
+            raise ValueError(f"lora_int_ids size mismatch: got {lora_int_ids.shape[0]}, expected {self.batch_size}.")
+
         batch_size = inputs.shape[0]
         if batch_size != self.batch_size:
             raise RuntimeError(

optimum/rbln/transformers/models/gemma3/modeling_gemma3.py

@@ -28,6 +28,7 @@ from ....modeling import RBLNModel
 from ...modeling_outputs import RBLNDecoderOnlyOutput
 from ...utils.rbln_runtime_wrapper import LoopProcessor
 from ..decoderonly.decoderonly_runtime_utils import RBLNPageTableManager
+from ..decoderonly.generation_decoderonly import RBLNDecoderOnlyGenerationMixin
 from ..decoderonly.modeling_decoderonly import (
     RBLNDecoderOnlyModelForCausalLM,
 )
@@ -77,7 +78,7 @@ class LoopProjector(LoopProcessor):
         return output[0]
 
 
-class RBLNGemma3ForConditionalGeneration(RBLNModel):
+class RBLNGemma3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixin):
     auto_model_class = AutoModelForImageTextToText
     _rbln_submodules = [
         {"name": "vision_tower"},
@@ -408,6 +409,13 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     def _update_sliding_window_config(cls, model_config: PretrainedConfig, rbln_config: RBLNGemma3ForCausalLMConfig):
         sliding_window = getattr(model_config, "sliding_window", None)
         sliding_window_pattern = getattr(model_config, "sliding_window_pattern", None)
+        if sliding_window_pattern is None:
+            if hasattr(model_config, "layer_types"):
+                first_full_attention_index = model_config.layer_types.index("full_attention")
+                sliding_window_pattern = first_full_attention_index + 1
+            else:
+                raise ValueError("Cannot determine sliding_window_pattern from model_config")
+
         if sliding_window_pattern <= model_config.num_hidden_layers:
             rbln_config.cache_impl = "hybrid"
             rbln_config.sliding_window = sliding_window

optimum/rbln/transformers/models/gpt2/gpt2_architecture.py

@@ -75,7 +75,10 @@ class GPT2Attention(DecoderOnlyAttention):
         self.o_proj = self._original_mod.c_proj
         self.split_size = self._original_mod.split_size
 
-    def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def projection(self, hidden_states, lora_int_id) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if lora_int_id is not None:
+            raise NotImplementedError("LoRA is not supported for GPT2Attention")
+
         query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
         return query_states, key_states, value_states
 

optimum/rbln/transformers/models/idefics3/modeling_idefics3.py

@@ -35,6 +35,7 @@ from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ....utils.runtime_utils import RBLNPytorchRuntime
 from ...modeling_outputs import RBLNDecoderOnlyOutput
+from ..decoderonly.generation_decoderonly import RBLNDecoderOnlyGenerationMixin
 
 
 if TYPE_CHECKING:
@@ -120,9 +121,6 @@ class RBLNIdefics3VisionTransformer(RBLNModel):
                 encoder_outputs = self.encoder(
                     inputs_embeds=hidden_states,
                     attention_mask=patch_attention_mask,
-                    output_attentions=None,
-                    output_hidden_states=None,
-                    return_dict=False,
                 )
                 last_hidden_state = encoder_outputs[0]
                 last_hidden_state = self.post_layernorm(last_hidden_state)
@@ -185,7 +183,7 @@ class RBLNIdefics3VisionTransformer(RBLNModel):
             return BaseModelOutput(last_hidden_state=last_hidden_state)
 
 
-class RBLNIdefics3ForConditionalGeneration(RBLNModel):
+class RBLNIdefics3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixin):
     """
     RBLNIdefics3ForConditionalGeneration is a multi-modal model that integrates vision and language processing capabilities,
     optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.

optimum/rbln/transformers/models/llava/modeling_llava.py

@@ -27,6 +27,7 @@ from ....modeling import RBLNModel
 from ....utils.logging import get_logger
 from ...modeling_outputs import RBLNDecoderOnlyOutput
 from ...utils.rbln_runtime_wrapper import LoopProcessor
+from ..decoderonly.generation_decoderonly import RBLNDecoderOnlyGenerationMixin
 
 
 logger = get_logger(__name__)
@@ -103,7 +104,7 @@ class LoopProjector(LoopProcessor):
         return output[0]
 
 
-class RBLNLlavaForConditionalGeneration(RBLNModel):
+class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixin):
     """
     RBLNLlavaForConditionalGeneration is a multi-modal model that combines vision and language processing capabilities,
     optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.

optimum/rbln/transformers/models/llava_next/modeling_llava_next.py

@@ -32,6 +32,7 @@ from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ....utils.logging import get_logger
 from ...utils.rbln_runtime_wrapper import LoopProcessor
+from ..decoderonly.generation_decoderonly import RBLNDecoderOnlyGenerationMixin
 from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyOutput
 
 
@@ -87,7 +88,7 @@ class LoopProjector(LoopProcessor):
         return output[0]
 
 
-class RBLNLlavaNextForConditionalGeneration(RBLNModel):
+class RBLNLlavaNextForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixin):
     """
     RBLNLlavaNextForConditionalGeneration is a multi-modal model that combines vision and language processing capabilities,
     optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.