optimum-rbln 0.9.1__py3-none-any.whl → 0.9.2__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,
@@ -388,6 +420,16 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             inputs, cache_position, attention_mask, position_embed, token_type_ids=token_type_ids
         )
 
+        # Assumed that prefix caching was performed externally if cache_position doesn't start from 0.
+        prefix_cached_len = cache_position[0][0].item()
+        if prefix_cached_len > 0:
+            if prefix_cached_len % self.rbln_config.prefill_chunk_size != 0:
+                raise NotImplementedError(
+                    "Prefix Caching is not supported yet for non-multiple of prefill_chunk_size."
+                )
+            if self.rbln_config.use_attention_mask:
+                chunked_attention_mask[:, :, :, :prefix_cached_len] = 1
+
         # Process input in chunks of size `prefill_chunk_size`
         output_logits = []
         for step in range(0, query_length, self.rbln_config.prefill_chunk_size):
@@ -402,9 +444,14 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             if self.rbln_config.use_attention_mask and not self.rbln_config.use_position_ids:
                 if step > 0:  # update previous chunk
                     chunked_attention_mask[
-                        :, :, :, s - self.rbln_config.prefill_chunk_size : e - self.rbln_config.prefill_chunk_size
+                        :,
+                        :,
+                        :,
+                        s - self.rbln_config.prefill_chunk_size + prefix_cached_len : e
+                        - self.rbln_config.prefill_chunk_size
+                        + prefix_cached_len,
                     ] = 1
-                chunked_attention_mask[:, :, :, s:e] = self.causal_mask
+                chunked_attention_mask[:, :, :, s + prefix_cached_len : e + prefix_cached_len] = self.causal_mask
 
             # Calculate query position if needed
             if self.rbln_config.use_local_attention or self.rbln_config.logits_to_keep > 0:
@@ -426,6 +473,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

@@ -57,7 +57,6 @@ class RBLNDecoderOnlyModel(RBLNModel, RBLNDecoderOnlyFlashAttentionMixin):
     1. Converting pre-trained transformer models to RBLN-optimized format
     2. Handling the compilation process for RBLN devices
     3. Managing inference operations for decoder-only architectures
-
     This class inherits from RBLNModel and implements specific methods required for
     decoder-only architectures.
 
@@ -68,6 +67,8 @@ class RBLNDecoderOnlyModel(RBLNModel, RBLNDecoderOnlyFlashAttentionMixin):
         - The class handles RBLN-specific optimizations automatically during compilation
     """
 
+    _tp_support = True
+
     main_input_name = "input_ids"
     auto_model_class = AutoModel
     _decoder_wrapper_cls = DecoderOnlyWrapper
@@ -259,10 +260,12 @@ class RBLNDecoderOnlyModel(RBLNModel, RBLNDecoderOnlyFlashAttentionMixin):
 
         # Mark static tensors (self kv states)
         static_tensors = {}
+        idx = 0
         for (name, _, _), tensor in zip(compile_config.input_info, example_inputs):
             if "past_key_values" in name:
                 static_tensors[name] = tensor
-                context.mark_static_address(tensor)
+                context.mark_static_address(tensor, f"kv_cache_{idx}")
+                idx += 1
 
         return context, static_tensors
 
@@ -374,6 +377,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"
@@ -642,7 +648,6 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNDecoderOnlyModel, RBLNDecoderOnlyGener
     1. Converting pre-trained transformer models to RBLN-optimized format
     2. Handling the compilation process for RBLN devices
     3. Managing inference operations for causal language modeling
-
     This class inherits from RBLNModel and implements specific methods required for
     decoder-only architectures and causal language modeling tasks.
 
@@ -667,6 +672,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 +729,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 +737,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 +768,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 +788,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/configuration_gemma3.py

@@ -14,8 +14,11 @@
 from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
 from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
-from ..siglip.configuration_siglip import RBLNSiglipVisionModelConfig
+
+
+logger = get_logger(__name__)
 
 
 class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -89,8 +92,13 @@ class RBLNGemma3ForConditionalGenerationConfig(RBLNModelConfig):
         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}")
 
-        self.vision_tower = self.init_submodule_config(RBLNSiglipVisionModelConfig, vision_tower)
-        self.language_model = self.init_submodule_config(RBLNGemma3ForCausalLMConfig, language_model)
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for Gemma3 vision tower. It will be set to 1.")
+
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower, batch_size=1, force_kwargs=True
+        )
+        self.language_model = self.initialize_submodule_config(submodule_config=language_model)
 
     @property
     def image_prefill_chunk_size(self):

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(