[TRTLLM-14692][feat] Overlap LoRA and base model computations

tensorrt_llm/_torch/modules/attention.py

@@ -22,7 +22,7 @@
 from ..distributed import (AllReduceParams, HelixAllToAllNative, alltoall_helix,
                            cp_allgather, reducescatter)
 from ..model_config import ModelConfig
-from ..peft.lora.layer import LoraLayer, LoraModuleType, add_lora_result
+from ..peft.lora.layer import LoraLayer, LoraModuleType
 from ..pyexecutor.breakable_cuda_graph import (eager_on_graph,
                                                is_in_breakable_cuda_graph)
 from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs,
@@ -1052,16 +1052,16 @@ def forward(
         hidden_states = _helix_cp_allgather_input(hidden_states, attn_metadata,
                                                   self.mapping, self.layer_idx)
 
-        qkv = self.qkv_proj(hidden_states)
-
         if bool(lora_params):
-            qkv_lora = self.splitted_qkv_lora(hidden_states, lora_params,
-                                              self.layer_idx)
-            qkv = add_lora_result(qkv, qkv_lora)
-
-            qkv_lora = self.fused_qkv_lora(hidden_states, lora_params,
-                                           self.layer_idx)
-            qkv = add_lora_result(qkv, qkv_lora)
+            qkv = LoraLayer.forward_with_base(
+                lambda: self.qkv_proj(hidden_states),
+                (self.splitted_qkv_lora, self.fused_qkv_lora),
+                hidden_states,
+                lora_params,
+                self.layer_idx,
+            )
+        else:
+            qkv = self.qkv_proj(hidden_states)
 
         # For dynamic tree spec decoding with Python RoPE, adjust position_ids
         # to use tree offsets (same as C++ kernel: past_seq_len + offset).

tensorrt_llm/_torch/modules/gated_mlp.py

@@ -10,7 +10,7 @@
 
 from ..distributed import AllReduceParams
 from ..model_config import ModelConfig
-from ..peft.lora.layer import LoraLayer, LoraModuleType, add_lora_result
+from ..peft.lora.layer import LoraLayer, LoraModuleType
 from ..utils import Fp4QuantizedTensor
 from .linear import (Linear, TensorParallelMode, WeightMode,
                      WeightsLoadingConfig, is_static_nvfp4_input_eligible)
@@ -353,14 +353,13 @@ def forward_lora(
                 "LoRA is not supported with uneven TP for GatedMLP "
                 "(intermediate_size not divisible by tp_size).")
 
-        h1 = self.gate_up_proj(x)
-
-        h1_lora = self.splitted_gate_up_lora(x, lora_params, self.layer_idx)
-
-        h1 = add_lora_result(h1, h1_lora)
-
-        h1_lora = self.fused_gate_up_lora(x, lora_params, self.layer_idx)
-        h1 = add_lora_result(h1, h1_lora)
+        h1 = LoraLayer.forward_with_base(
+            lambda: self.gate_up_proj(x),
+            (self.splitted_gate_up_lora, self.fused_gate_up_lora),
+            x,
+            lora_params,
+            self.layer_idx,
+        )
 
         h2 = self._apply_activation(h1, has_lora=True)
         output = self.down_proj(h2,

tensorrt_llm/_torch/modules/linear.py

@@ -14,7 +14,7 @@
 
 import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils
 from tensorrt_llm._torch.custom_ops.torch_custom_ops import BufferKind
-from tensorrt_llm._torch.peft.lora.layer import LoraLayer, add_lora_result
+from tensorrt_llm._torch.peft.lora.layer import LoraLayer
 from tensorrt_llm._utils import is_device_integrated, mpi_disabled
 from tensorrt_llm.bindings import ipc_nvls_supported
 from tensorrt_llm.functional import (AllReduceFusionOp, AllReduceParams,
@@ -3823,10 +3823,16 @@ def apply_linear(self,
                      bias,
                      lora_params: Optional[dict] | None = None,
                      layer_idx: Optional[int] | None = None):
-        output = self.quant_method.apply(self, input, bias)
         if self.lora is not None and bool(lora_params):
-            lora_result = self.lora(input, lora_params, layer_idx)
-            output = add_lora_result(output, lora_result)
+            output = LoraLayer.forward_with_base(
+                lambda: self.quant_method.apply(self, input, bias),
+                (self.lora, ),
+                input,
+                lora_params,
+                layer_idx,
+            )
+        else:
+            output = self.quant_method.apply(self, input, bias)
         return output
 
     def apply_linear_allreduce(self,

tensorrt_llm/_torch/modules/mlp.py

@@ -8,7 +8,7 @@
 from tensorrt_llm.mapping import Mapping
 
 from ..model_config import ModelConfig
-from ..peft.lora.layer import LoraLayer, LoraModuleType, add_lora_result
+from ..peft.lora.layer import LoraLayer, LoraModuleType
 from ..utils import Fp4QuantizedTensor, gelu_tanh, relu2
 from .linear import (Linear, TensorParallelMode, WeightMode,
                      WeightsLoadingConfig, is_static_nvfp4_input_eligible)
@@ -268,11 +268,14 @@ def forward_lora(
     ) -> torch.Tensor:
         assert lora_params is not None
 
-        x_up = self.up_proj(x)
-
         assert self.layer_idx is not None, "layer_idx is required for lora"
-        x_up_lora = self.up_lora(x, lora_params, self.layer_idx)
-        x_up = add_lora_result(x_up, x_up_lora)
+        x_up = LoraLayer.forward_with_base(
+            lambda: self.up_proj(x),
+            (self.up_lora, ),
+            x,
+            lora_params,
+            self.layer_idx,
+        )
 
         x_act = self.activation(x_up)
         x_down = self.down_proj(x_act,

tensorrt_llm/_torch/peft/lora/config.py

@@ -111,6 +111,12 @@ class LoraConfig(StrictBaseModel):
         "Whether to swap gate/up projection order in fused gate_up_proj LoRA B weights. "
         "Set to False for models like Phi-4-MM that use a different weight order."
     )
+    overlap_lora_and_base: bool = Field(
+        default=False,
+        description=
+        "Whether to place LoRA operations on a secondary CUDA stream and overlap them "
+        "with base model computations. Improves latency in memory-bound regimes."
+    )
 
     @property
     def missing_qkv_modules(self) -> List[str]:

tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py

@@ -44,6 +44,7 @@ def __init__(
         max_lora_rank: int,
         model: torch.nn.Module,
         lora_model_config: Optional[LoraModelConfig],
+        overlap_lora_and_base: bool = False,
         device: str = "cuda",
         max_tokens_per_seq: int = 1,
     ):
@@ -56,6 +57,7 @@ def __init__(
             max_lora_rank: Maximum LoRA rank across all layers
             model: Model to get layerwise LoRA info
             lora_model_config: LoRA model configuration
+            overlap_lora_and_base: Whether to overlap LoRA and base model computations.
             device: Device to allocate tensors on
             max_tokens_per_seq: Maximum number of tokens per sequence (>1 for spec decode)
         """
@@ -67,6 +69,7 @@ def __init__(
         self.max_tokens_per_seq = max_tokens_per_seq
         self.adapter_slot_manager = AdapterSlotManager(max_lora_size)
         self.lora_model_config = lora_model_config
+        self.lora_aux_stream = torch.cuda.Stream(device=device) if overlap_lora_and_base else None
         lora_target_modules = lora_model_config.lora_target_modules
         self.target_modules_ids: Optional[tuple[int, ...]] = (
             tuple(map(LoraManager.LORA_MODULE_IDS.__getitem__, lora_target_modules))
@@ -225,6 +228,7 @@ def prepare_cuda_graph_lora_params(
             "num_seqs": attn_metadata.num_seqs,
             "use_cuda_graph_mode": True,  # Flag to indicate new mode
             "data_type": peft_cache_manager.data_type,
+            "lora_aux_stream": self.lora_aux_stream,
         }
 
         return lora_params

tensorrt_llm/_torch/peft/lora/layer.py

@@ -13,12 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections.abc import Callable
 from dataclasses import dataclass
 from enum import IntEnum
 from typing import Dict, List, Optional
 
 import torch
 
+from ...modules.multi_stream_utils import (do_multi_stream,
+                                           maybe_execute_in_parallel)
 from .cuda_graph_lora_params import CudaGraphLoraParams
 
 _FP8_LORA_TMA_ALIGNMENT = 16
@@ -213,9 +216,9 @@ def is_mamba(self) -> bool:
 
 def add_lora_result(output: torch.Tensor,
                     lora_result: Optional[torch.Tensor]) -> torch.Tensor:
-    if lora_result is None:
-        return output
-    return output + lora_result.to(output.dtype)
+    if lora_result is not None:
+        output.add_(lora_result.to(output.dtype))
+    return output
 
 
 class LoraLayer(torch.nn.Module):
@@ -228,6 +231,79 @@ def __init__(self, lora_module_types: List[LoraModuleType],
         self.output_hidden_sizes = output_hidden_sizes
         assert len(lora_module_types) == len(output_hidden_sizes)
 
+        self._par_events: List[torch.cuda.Event] | None = None
+
+    @staticmethod
+    def forward_with_base(
+        base_forward: Callable[[], torch.Tensor],
+        lora_layers: tuple["LoraLayer", ...],
+        x: torch.Tensor,
+        lora_params: dict,
+        layer_idx: int | None,
+    ) -> torch.Tensor:
+        """
+        Run the base and LoRA branches and merge their outputs.
+
+        Args:
+            base_forward: Forward call for base model projection
+            lora_layers: Tuple of LoRA layers to be called
+            x: Input tensor
+            lora_params: CUDA Graph compatible LoRA parameters
+            layer_idx: Current layer index
+
+        Returns:
+            LoRA + base model output tensor
+
+        Note that lora_layers needs to be a tuple in order to
+        handle fused/unfused modules (e.g., QKV), where both
+        variants are invoked but only one runs through.
+        """
+        cuda_graph_params = lora_params.get('cuda_graph_params')
+        has_lora_layer = bool(cuda_graph_params) and any(
+            CudaGraphLoraParams.LoraLayerKey(
+                layer_idx=layer_idx,
+                module_ids=tuple(layer.lora_module_types),
+            ) in cuda_graph_params.layer_info for layer in lora_layers)
+
+        lora_aux_stream = lora_params.get("lora_aux_stream")
+        execute_in_parallel = (has_lora_layer and lora_aux_stream is not None
+                               and do_multi_stream()
+                               and not torch.compiler.is_compiling())
+
+        # Pack all LoRA forwards (e.g., fused/unfused) in a single tuple
+        def lora_forward() -> tuple[torch.Tensor | None, ...]:
+            return tuple(
+                lora_layer(x, lora_params, layer_idx)
+                for lora_layer in lora_layers)
+
+        if execute_in_parallel:
+            assert lora_aux_stream is not None
+            # Lazy allocation of parallel events
+            if lora_layers[0]._par_events is None:
+                lora_layers[0]._par_events = [
+                    torch.cuda.Event(), torch.cuda.Event()
+                ]
+
+            base_output, lora_outputs = maybe_execute_in_parallel(
+                base_forward,
+                lora_forward,
+                lora_layers[0]._par_events[0],
+                lora_layers[0]._par_events[1],
+                lora_aux_stream,
+                disable_on_compile=True,
+            )
+        else:
+            base_output, lora_outputs = base_forward(), lora_forward()
+
+        for lora_output in lora_outputs:
+            if not isinstance(lora_output, torch.Tensor):
+                continue
+            if execute_in_parallel:
+                lora_output.record_stream(torch.cuda.current_stream())
+            base_output = add_lora_result(base_output, lora_output)
+
+        return base_output
+
     def forward(
         self,
         x,
@@ -567,7 +643,8 @@ def _forward_cuda_graph_mode(
             output_buffer = output_buffer.to(torch.bfloat16)
 
         # TODO: move to kernel
-        restored_output = torch.zeros_like(output_buffer)
+        # sorted_ids is a permutation, so index_copy_ initializes every row.
+        restored_output = torch.empty_like(output_buffer)
         restored_output.index_copy_(0,
                                     cuda_graph_params.sorted_ids[:batch_size],
                                     output_buffer)

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -1135,6 +1135,7 @@ def _init_cuda_graph_lora_manager(self, lora_config: LoraConfig):
                 max_lora_rank=lora_config.max_lora_rank,
                 model=self.model,
                 lora_model_config=self.lora_model_config,
+                overlap_lora_and_base=lora_config.overlap_lora_and_base,
                 device='cuda',
                 max_tokens_per_seq=max_tokens_per_seq)
 

tensorrt_llm/usage/llm_args_golden_manifest.json

@@ -996,6 +996,13 @@
       "kind": "value",
       "path": "lora_config.max_loras"
     },
+    {
+      "allowed_values": [],
+      "annotation": "<class 'bool'>",
+      "converter": "",
+      "kind": "value",
+      "path": "lora_config.overlap_lora_and_base"
+    },
     {
       "allowed_values": [],
       "annotation": "<class 'bool'>",