From 9fc04580867e30a9af04e00b5c3c55cadd1ff6ec Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Thu, 23 Jul 2026 04:46:45 -0700
Subject: [PATCH 01/14] Overlap LoRA computations with base model projections

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/attention.py | 29 ++++++++++----
 tensorrt_llm/_torch/modules/gated_mlp.py | 22 +++++++----
 tensorrt_llm/_torch/modules/linear.py    | 13 ++++--
 tensorrt_llm/_torch/modules/mlp.py       | 12 ++++--
 tensorrt_llm/_torch/peft/lora/layer.py   | 50 +++++++++++++++++++++++-
 5 files changed, 103 insertions(+), 23 deletions(-)

diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index b6fea37271ba..0131f6ec91d2 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -1052,16 +1052,29 @@ 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)
+            def qkv_lora():
+                return (
+                    self.splitted_qkv_lora(hidden_states, lora_params,
+                                           self.layer_idx),
+                    self.fused_qkv_lora(hidden_states, lora_params,
+                                        self.layer_idx),
+                )
+
+            qkv, qkv_lora_results = (
+                self.splitted_qkv_lora.execute_with_base(
+                    lambda: self.qkv_proj(hidden_states),
+                    qkv_lora,
+                    lora_params,
+                    self.layer_idx,
+                    (self.fused_qkv_lora, ),
+                ))
+            for qkv_lora_result in qkv_lora_results:
+                if qkv_lora_result is not None:
+                    qkv = qkv + qkv_lora_result
+        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).
diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py
index 8c49a116bf4c..b78098afe268 100644
--- a/tensorrt_llm/_torch/modules/gated_mlp.py
+++ b/tensorrt_llm/_torch/modules/gated_mlp.py
@@ -353,14 +353,22 @@ 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)
+        def gate_up_lora():
+            return (
+                self.splitted_gate_up_lora(x, lora_params, self.layer_idx),
+                self.fused_gate_up_lora(x, lora_params, self.layer_idx),
+            )
 
-        h1_lora = self.fused_gate_up_lora(x, lora_params, self.layer_idx)
-        h1 = add_lora_result(h1, h1_lora)
+        h1, h1_lora_results = self.splitted_gate_up_lora.execute_with_base(
+            lambda: self.gate_up_proj(x),
+            gate_up_lora,
+            lora_params,
+            self.layer_idx,
+            (self.fused_gate_up_lora, ),
+        )
+        for h1_lora in h1_lora_results:
+            if h1_lora is not None:
+                h1 = h1 + h1_lora
 
         h2 = self._apply_activation(h1, has_lora=True)
         output = self.down_proj(h2,
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index 97334baf1e27..55803d03f6e2 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -3823,10 +3823,17 @@ 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, lora_result = self.lora.execute_with_base(
+                lambda: self.quant_method.apply(self, input, bias),
+                lambda: self.lora(input, lora_params, layer_idx),
+                lora_params,
+                layer_idx,
+            )
+            if lora_result is not None:
+                output = output + lora_result
+        else:
+            output = self.quant_method.apply(self, input, bias)
         return output
 
     def apply_linear_allreduce(self,
diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py
index bad1078027b6..00c7485d6e43 100644
--- a/tensorrt_llm/_torch/modules/mlp.py
+++ b/tensorrt_llm/_torch/modules/mlp.py
@@ -268,11 +268,15 @@ 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, x_up_lora = self.up_lora.execute_with_base(
+            lambda: self.up_proj(x),
+            lambda: self.up_lora(x, lora_params, self.layer_idx),
+            lora_params,
+            self.layer_idx,
+        )
+        if x_up_lora is not None:
+            x_up = x_up + x_up_lora
 
         x_act = self.activation(x_up)
         x_down = self.down_proj(x_act,
diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py
index cf580bd8d956..ffc070fb7fff 100644
--- a/tensorrt_llm/_torch/peft/lora/layer.py
+++ b/tensorrt_llm/_torch/peft/lora/layer.py
@@ -15,10 +15,12 @@
 
 from dataclasses import dataclass
 from enum import IntEnum
-from typing import Dict, List, Optional
+from typing import Any, Callable, 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
@@ -220,6 +222,8 @@ def add_lora_result(output: torch.Tensor,
 
 class LoraLayer(torch.nn.Module):
 
+    _aux_streams: Dict[int, torch.cuda.Stream] = {}
+
     def __init__(self, lora_module_types: List[LoraModuleType],
                  output_hidden_sizes: List[int]):
         super().__init__()
@@ -228,6 +232,50 @@ 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._aux_stream: Optional[torch.cuda.Stream] = None
+        self._parallel_events: Optional[List[torch.cuda.Event]] = None
+
+    def execute_with_base(
+        self,
+        base_forward: Callable[[], Any],
+        lora_forward: Callable[[], Any],
+        lora_params: Dict,
+        layer_idx: Optional[int],
+        additional_lora_layers: tuple["LoraLayer", ...] = (),
+    ) -> tuple[Any, Any]:
+        """Run the base and LoRA branches concurrently during graph capture."""
+        use_cuda_graph_mode = bool(lora_params) and lora_params.get(
+            'use_cuda_graph_mode', False)
+        cuda_graph_params = lora_params.get('cuda_graph_params')
+        lora_layers = (self, ) + additional_lora_layers
+        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)
+        execute_in_parallel = (use_cuda_graph_mode and has_lora_layer
+                               and do_multi_stream()
+                               and not torch.compiler.is_compiling())
+        if execute_in_parallel and self._aux_stream is None:
+            device = torch.cuda.current_device()
+            if device not in self._aux_streams:
+                self._aux_streams[device] = torch.cuda.Stream(device=device)
+            self._aux_stream = self._aux_streams[device]
+            self._parallel_events = [torch.cuda.Event(), torch.cuda.Event()]
+
+        if execute_in_parallel:
+            assert self._parallel_events is not None
+            return maybe_execute_in_parallel(
+                base_forward,
+                lora_forward,
+                self._parallel_events[0],
+                self._parallel_events[1],
+                self._aux_stream,
+                disable_on_compile=True,
+            )
+
+        return base_forward(), lora_forward()
+
     def forward(
         self,
         x,

From f24c4152078c3465ba2bb6185dd2e9c7c5430227 Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Fri, 24 Jul 2026 01:44:34 -0700
Subject: [PATCH 02/14] Refactor changes for concurrent LoRA execution

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/attention.py | 23 ++++-------
 tensorrt_llm/_torch/modules/gated_mlp.py | 13 ++----
 tensorrt_llm/_torch/modules/linear.py    |  6 ++-
 tensorrt_llm/_torch/modules/mlp.py       |  6 ++-
 tensorrt_llm/_torch/peft/lora/layer.py   | 51 +++++++++++++-----------
 5 files changed, 47 insertions(+), 52 deletions(-)

diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index 0131f6ec91d2..b7db954c0c0d 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -1053,23 +1053,14 @@ def forward(
                                                   self.mapping, self.layer_idx)
 
         if bool(lora_params):
+            qkv, qkv_lora_results = 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,
+            )
 
-            def qkv_lora():
-                return (
-                    self.splitted_qkv_lora(hidden_states, lora_params,
-                                           self.layer_idx),
-                    self.fused_qkv_lora(hidden_states, lora_params,
-                                        self.layer_idx),
-                )
-
-            qkv, qkv_lora_results = (
-                self.splitted_qkv_lora.execute_with_base(
-                    lambda: self.qkv_proj(hidden_states),
-                    qkv_lora,
-                    lora_params,
-                    self.layer_idx,
-                    (self.fused_qkv_lora, ),
-                ))
             for qkv_lora_result in qkv_lora_results:
                 if qkv_lora_result is not None:
                     qkv = qkv + qkv_lora_result
diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py
index b78098afe268..94d7d62492a6 100644
--- a/tensorrt_llm/_torch/modules/gated_mlp.py
+++ b/tensorrt_llm/_torch/modules/gated_mlp.py
@@ -353,19 +353,14 @@ def forward_lora(
                 "LoRA is not supported with uneven TP for GatedMLP "
                 "(intermediate_size not divisible by tp_size).")
 
-        def gate_up_lora():
-            return (
-                self.splitted_gate_up_lora(x, lora_params, self.layer_idx),
-                self.fused_gate_up_lora(x, lora_params, self.layer_idx),
-            )
-
-        h1, h1_lora_results = self.splitted_gate_up_lora.execute_with_base(
+        h1, h1_lora_results = LoraLayer.forward_with_base(
             lambda: self.gate_up_proj(x),
-            gate_up_lora,
+            (self.splitted_gate_up_lora, self.fused_gate_up_lora),
+            x,
             lora_params,
             self.layer_idx,
-            (self.fused_gate_up_lora, ),
         )
+
         for h1_lora in h1_lora_results:
             if h1_lora is not None:
                 h1 = h1 + h1_lora
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index 55803d03f6e2..db83b4d9f6ae 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -3824,12 +3824,14 @@ def apply_linear(self,
                      lora_params: Optional[dict] | None = None,
                      layer_idx: Optional[int] | None = None):
         if self.lora is not None and bool(lora_params):
-            output, lora_result = self.lora.execute_with_base(
+            output, (lora_result,) = LoraLayer.forward_with_base(
                 lambda: self.quant_method.apply(self, input, bias),
-                lambda: self.lora(input, lora_params, layer_idx),
+                (self.lora,),
+                input,
                 lora_params,
                 layer_idx,
             )
+
             if lora_result is not None:
                 output = output + lora_result
         else:
diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py
index 00c7485d6e43..dba04821d7b4 100644
--- a/tensorrt_llm/_torch/modules/mlp.py
+++ b/tensorrt_llm/_torch/modules/mlp.py
@@ -269,12 +269,14 @@ def forward_lora(
         assert lora_params is not None
 
         assert self.layer_idx is not None, "layer_idx is required for lora"
-        x_up, x_up_lora = self.up_lora.execute_with_base(
+        x_up, (x_up_lora,) = LoraLayer.forward_with_base(
             lambda: self.up_proj(x),
-            lambda: self.up_lora(x, lora_params, self.layer_idx),
+            (self.up_lora,),
+            x,
             lora_params,
             self.layer_idx,
         )
+
         if x_up_lora is not None:
             x_up = x_up + x_up_lora
 
diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py
index ffc070fb7fff..9465d5a8ba23 100644
--- a/tensorrt_llm/_torch/peft/lora/layer.py
+++ b/tensorrt_llm/_torch/peft/lora/layer.py
@@ -222,7 +222,8 @@ def add_lora_result(output: torch.Tensor,
 
 class LoraLayer(torch.nn.Module):
 
-    _aux_streams: Dict[int, torch.cuda.Stream] = {}
+    # Static aux stream for LoRA computations
+    _aux_stream: Optional[torch.cuda.Stream] = None
 
     def __init__(self, lora_module_types: List[LoraModuleType],
                  output_hidden_sizes: List[int]):
@@ -232,49 +233,53 @@ 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._aux_stream: Optional[torch.cuda.Stream] = None
-        self._parallel_events: Optional[List[torch.cuda.Event]] = None
+        self._par_events: Optional[List[torch.cuda.Event]] = None
 
-    def execute_with_base(
-        self,
+    @staticmethod
+    def forward_with_base(
         base_forward: Callable[[], Any],
-        lora_forward: Callable[[], Any],
+        lora_layers: tuple["LoraLayer", ...],
+        x: torch.Tensor,
         lora_params: Dict,
         layer_idx: Optional[int],
-        additional_lora_layers: tuple["LoraLayer", ...] = (),
     ) -> tuple[Any, Any]:
         """Run the base and LoRA branches concurrently during graph capture."""
-        use_cuda_graph_mode = bool(lora_params) and lora_params.get(
-            'use_cuda_graph_mode', False)
         cuda_graph_params = lora_params.get('cuda_graph_params')
-        lora_layers = (self, ) + additional_lora_layers
         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)
-        execute_in_parallel = (use_cuda_graph_mode and has_lora_layer
+
+        execute_in_parallel = (has_lora_layer
                                and do_multi_stream()
                                and not torch.compiler.is_compiling())
-        if execute_in_parallel and self._aux_stream is None:
-            device = torch.cuda.current_device()
-            if device not in self._aux_streams:
-                self._aux_streams[device] = torch.cuda.Stream(device=device)
-            self._aux_stream = self._aux_streams[device]
-            self._parallel_events = [torch.cuda.Event(), torch.cuda.Event()]
+
+        # Pack all LoRA forwards (e.g., fused/unfused) in a single tuple
+        def lora_forward():
+            return tuple(
+                lora_layer(x, lora_params, layer_idx)
+                for lora_layer in lora_layers
+            )
 
         if execute_in_parallel:
-            assert self._parallel_events is not None
+            # Lazy allocation of aux stream and parallel events
+            if LoraLayer._aux_stream is None:
+                LoraLayer._aux_stream = torch.cuda.Stream()
+            if lora_layers[0]._par_events is None:
+                lora_layers[0]._par_events = [
+                    torch.cuda.Event(), torch.cuda.Event()]
+
             return maybe_execute_in_parallel(
                 base_forward,
                 lora_forward,
-                self._parallel_events[0],
-                self._parallel_events[1],
-                self._aux_stream,
+                lora_layers[0]._par_events[0],
+                lora_layers[0]._par_events[1],
+                LoraLayer._aux_stream,
                 disable_on_compile=True,
             )
-
-        return base_forward(), lora_forward()
+        else:
+            return base_forward(), lora_forward()
 
     def forward(
         self,

From cea9c3c895d8dc3874d564590a5d1e5083dfd2e2 Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Fri, 24 Jul 2026 02:37:35 -0700
Subject: [PATCH 03/14] Encapsulate LoRA output handling inside
 forward_with_base(), switch to empty_like() for init

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/attention.py |  6 +-----
 tensorrt_llm/_torch/modules/gated_mlp.py |  6 +-----
 tensorrt_llm/_torch/modules/linear.py    |  5 +----
 tensorrt_llm/_torch/modules/mlp.py       |  5 +----
 tensorrt_llm/_torch/peft/lora/layer.py   | 23 +++++++++++++++++------
 5 files changed, 21 insertions(+), 24 deletions(-)

diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index b7db954c0c0d..a81197746311 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -1053,17 +1053,13 @@ def forward(
                                                   self.mapping, self.layer_idx)
 
         if bool(lora_params):
-            qkv, qkv_lora_results = LoraLayer.forward_with_base(
+            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,
             )
-
-            for qkv_lora_result in qkv_lora_results:
-                if qkv_lora_result is not None:
-                    qkv = qkv + qkv_lora_result
         else:
             qkv = self.qkv_proj(hidden_states)
 
diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py
index 94d7d62492a6..0f1d2ce9f0d4 100644
--- a/tensorrt_llm/_torch/modules/gated_mlp.py
+++ b/tensorrt_llm/_torch/modules/gated_mlp.py
@@ -353,7 +353,7 @@ def forward_lora(
                 "LoRA is not supported with uneven TP for GatedMLP "
                 "(intermediate_size not divisible by tp_size).")
 
-        h1, h1_lora_results = LoraLayer.forward_with_base(
+        h1 = LoraLayer.forward_with_base(
             lambda: self.gate_up_proj(x),
             (self.splitted_gate_up_lora, self.fused_gate_up_lora),
             x,
@@ -361,10 +361,6 @@ def forward_lora(
             self.layer_idx,
         )
 
-        for h1_lora in h1_lora_results:
-            if h1_lora is not None:
-                h1 = h1 + h1_lora
-
         h2 = self._apply_activation(h1, has_lora=True)
         output = self.down_proj(h2,
                                 all_reduce_params=final_all_reduce_params,
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index db83b4d9f6ae..2d191df9323a 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -3824,16 +3824,13 @@ def apply_linear(self,
                      lora_params: Optional[dict] | None = None,
                      layer_idx: Optional[int] | None = None):
         if self.lora is not None and bool(lora_params):
-            output, (lora_result,) = LoraLayer.forward_with_base(
+            output = LoraLayer.forward_with_base(
                 lambda: self.quant_method.apply(self, input, bias),
                 (self.lora,),
                 input,
                 lora_params,
                 layer_idx,
             )
-
-            if lora_result is not None:
-                output = output + lora_result
         else:
             output = self.quant_method.apply(self, input, bias)
         return output
diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py
index dba04821d7b4..ae9a7171aa6c 100644
--- a/tensorrt_llm/_torch/modules/mlp.py
+++ b/tensorrt_llm/_torch/modules/mlp.py
@@ -269,7 +269,7 @@ def forward_lora(
         assert lora_params is not None
 
         assert self.layer_idx is not None, "layer_idx is required for lora"
-        x_up, (x_up_lora,) = LoraLayer.forward_with_base(
+        x_up = LoraLayer.forward_with_base(
             lambda: self.up_proj(x),
             (self.up_lora,),
             x,
@@ -277,9 +277,6 @@ def forward_lora(
             self.layer_idx,
         )
 
-        if x_up_lora is not None:
-            x_up = x_up + x_up_lora
-
         x_act = self.activation(x_up)
         x_down = self.down_proj(x_act,
                                 lora_params=lora_params,
diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py
index 9465d5a8ba23..2b66402a774b 100644
--- a/tensorrt_llm/_torch/peft/lora/layer.py
+++ b/tensorrt_llm/_torch/peft/lora/layer.py
@@ -242,8 +242,8 @@ def forward_with_base(
         x: torch.Tensor,
         lora_params: Dict,
         layer_idx: Optional[int],
-    ) -> tuple[Any, Any]:
-        """Run the base and LoRA branches concurrently during graph capture."""
+    ) -> Any:
+        """Run the base and LoRA branches and merge their outputs."""
         cuda_graph_params = lora_params.get('cuda_graph_params')
         has_lora_layer = bool(cuda_graph_params) and any(
             CudaGraphLoraParams.LoraLayerKey(
@@ -251,7 +251,7 @@ def forward_with_base(
                 module_ids=tuple(layer.lora_module_types),
             ) in cuda_graph_params.layer_info for layer in lora_layers)
 
-        execute_in_parallel = (has_lora_layer
+        execute_in_parallel = (has_lora_layer 
                                and do_multi_stream()
                                and not torch.compiler.is_compiling())
 
@@ -270,7 +270,7 @@ def lora_forward():
                 lora_layers[0]._par_events = [
                     torch.cuda.Event(), torch.cuda.Event()]
 
-            return maybe_execute_in_parallel(
+            base_output, lora_outputs = maybe_execute_in_parallel(
                 base_forward,
                 lora_forward,
                 lora_layers[0]._par_events[0],
@@ -279,7 +279,17 @@ def lora_forward():
                 disable_on_compile=True,
             )
         else:
-            return base_forward(), lora_forward()
+            base_output, lora_outputs = base_forward(), lora_forward()
+
+        for lora_output in lora_outputs:
+            if lora_output is None:
+                continue
+            if cuda_graph_params and not torch.compiler.is_compiling():
+                base_output.add_(lora_output)
+            else:
+                base_output = base_output + lora_output
+
+        return base_output
 
     def forward(
         self,
@@ -620,7 +630,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)

From 88a1fb330d1fc82291bebbe8b14bfb5757cd80a5 Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Tue, 28 Jul 2026 00:58:02 -0700
Subject: [PATCH 04/14] Always use in-place add_() to combine base and LoRA
 outputs

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/peft/lora/layer.py | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py
index 2b66402a774b..e738a3a3e2aa 100644
--- a/tensorrt_llm/_torch/peft/lora/layer.py
+++ b/tensorrt_llm/_torch/peft/lora/layer.py
@@ -251,16 +251,14 @@ def forward_with_base(
                 module_ids=tuple(layer.lora_module_types),
             ) in cuda_graph_params.layer_info for layer in lora_layers)
 
-        execute_in_parallel = (has_lora_layer 
-                               and do_multi_stream()
+        execute_in_parallel = (has_lora_layer 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():
             return tuple(
                 lora_layer(x, lora_params, layer_idx)
-                for lora_layer in lora_layers
-            )
+                for lora_layer in lora_layers)
 
         if execute_in_parallel:
             # Lazy allocation of aux stream and parallel events
@@ -268,7 +266,8 @@ def lora_forward():
                 LoraLayer._aux_stream = torch.cuda.Stream()
             if lora_layers[0]._par_events is None:
                 lora_layers[0]._par_events = [
-                    torch.cuda.Event(), torch.cuda.Event()]
+                    torch.cuda.Event(), torch.cuda.Event()
+                ]
 
             base_output, lora_outputs = maybe_execute_in_parallel(
                 base_forward,
@@ -284,10 +283,7 @@ def lora_forward():
         for lora_output in lora_outputs:
             if lora_output is None:
                 continue
-            if cuda_graph_params and not torch.compiler.is_compiling():
-                base_output.add_(lora_output)
-            else:
-                base_output = base_output + lora_output
+            base_output.add_(lora_output)
 
         return base_output
 

From a6deba811c019b28370a3cb6cd5303adb09c202d Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Tue, 28 Jul 2026 01:32:37 -0700
Subject: [PATCH 05/14] Add overlap_lora_and_base flag to toggle aux CUDA
 stream

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/peft/lora/config.py                  | 6 ++++++
 tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py | 4 ++++
 tensorrt_llm/_torch/peft/lora/layer.py                   | 3 ++-
 tensorrt_llm/_torch/pyexecutor/model_engine.py           | 1 +
 tensorrt_llm/usage/llm_args_golden_manifest.json         | 7 +++++++
 5 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/peft/lora/config.py b/tensorrt_llm/_torch/peft/lora/config.py
index a501febb1b9b..f36a321a4bc1 100644
--- a/tensorrt_llm/_torch/peft/lora/config.py
+++ b/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=True,
+        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]:
diff --git a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py b/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py
index 6cda1e7da2f1..3fdf5f0cb231 100644
--- a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py
+++ b/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 = True,
         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.overlap_lora_and_base = overlap_lora_and_base
         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,
+            "overlap_lora_and_base": self.overlap_lora_and_base,
         }
 
         return lora_params
diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py
index e738a3a3e2aa..c127a66d40bd 100644
--- a/tensorrt_llm/_torch/peft/lora/layer.py
+++ b/tensorrt_llm/_torch/peft/lora/layer.py
@@ -251,7 +251,8 @@ def forward_with_base(
                 module_ids=tuple(layer.lora_module_types),
             ) in cuda_graph_params.layer_info for layer in lora_layers)
 
-        execute_in_parallel = (has_lora_layer and do_multi_stream()
+        execute_in_parallel = (has_lora_layer and lora_params.get(
+            "overlap_lora_and_base", True) and do_multi_stream()
                                and not torch.compiler.is_compiling())
 
         # Pack all LoRA forwards (e.g., fused/unfused) in a single tuple
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 3da7a73a7e9e..077296cb9143 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/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)
 
diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json
index 1dfc185d4a6e..f538f595195c 100644
--- a/tensorrt_llm/usage/llm_args_golden_manifest.json
+++ b/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'>",

From 01664808b171c6b39e13ab432377db3738a7ab08 Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Tue, 28 Jul 2026 05:37:13 -0700
Subject: [PATCH 06/14] CodeRabbit nitpicks and linting fixes

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/linear.py  |  2 +-
 tensorrt_llm/_torch/modules/mlp.py     |  2 +-
 tensorrt_llm/_torch/peft/lora/layer.py | 35 +++++++++++++++++++-------
 3 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index 2d191df9323a..d3631806e809 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -3826,7 +3826,7 @@ def apply_linear(self,
         if self.lora is not None and bool(lora_params):
             output = LoraLayer.forward_with_base(
                 lambda: self.quant_method.apply(self, input, bias),
-                (self.lora,),
+                (self.lora, ),
                 input,
                 lora_params,
                 layer_idx,
diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py
index ae9a7171aa6c..df654f627d88 100644
--- a/tensorrt_llm/_torch/modules/mlp.py
+++ b/tensorrt_llm/_torch/modules/mlp.py
@@ -271,7 +271,7 @@ def forward_lora(
         assert self.layer_idx is not None, "layer_idx is required for lora"
         x_up = LoraLayer.forward_with_base(
             lambda: self.up_proj(x),
-            (self.up_lora,),
+            (self.up_lora, ),
             x,
             lora_params,
             self.layer_idx,
diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py
index c127a66d40bd..4ed0b00b81f2 100644
--- a/tensorrt_llm/_torch/peft/lora/layer.py
+++ b/tensorrt_llm/_torch/peft/lora/layer.py
@@ -13,9 +13,10 @@
 # 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 Any, Callable, Dict, List, Optional
+from typing import Dict, List, Optional
 
 import torch
 
@@ -223,7 +224,7 @@ def add_lora_result(output: torch.Tensor,
 class LoraLayer(torch.nn.Module):
 
     # Static aux stream for LoRA computations
-    _aux_stream: Optional[torch.cuda.Stream] = None
+    _aux_stream: torch.cuda.Stream | None = None
 
     def __init__(self, lora_module_types: List[LoraModuleType],
                  output_hidden_sizes: List[int]):
@@ -233,17 +234,33 @@ 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: Optional[List[torch.cuda.Event]] = None
+        self._par_events: List[torch.cuda.Event] | None = None
 
     @staticmethod
     def forward_with_base(
-        base_forward: Callable[[], Any],
+        base_forward: Callable[[], torch.Tensor],
         lora_layers: tuple["LoraLayer", ...],
         x: torch.Tensor,
-        lora_params: Dict,
-        layer_idx: Optional[int],
-    ) -> Any:
-        """Run the base and LoRA branches and merge their outputs."""
+        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(
@@ -256,7 +273,7 @@ def forward_with_base(
                                and not torch.compiler.is_compiling())
 
         # Pack all LoRA forwards (e.g., fused/unfused) in a single tuple
-        def lora_forward():
+        def lora_forward() -> tuple[torch.Tensor | None, ...]:
             return tuple(
                 lora_layer(x, lora_params, layer_idx)
                 for lora_layer in lora_layers)

From 61014200f3a3ce218b16e0e2dcf52aeffe253b49 Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Fri, 31 Jul 2026 07:49:07 -0700
Subject: [PATCH 07/14] Disable LoRA overlap by default, add focused tests

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/peft/lora/config.py       |  2 +-
 .../tests_lora_modules/test_qwen3_sanity.py   | 20 ++++++++++++++++++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/peft/lora/config.py b/tensorrt_llm/_torch/peft/lora/config.py
index f36a321a4bc1..a58b82116ba5 100644
--- a/tensorrt_llm/_torch/peft/lora/config.py
+++ b/tensorrt_llm/_torch/peft/lora/config.py
@@ -112,7 +112,7 @@ class LoraConfig(StrictBaseModel):
         "Set to False for models like Phi-4-MM that use a different weight order."
     )
     overlap_lora_and_base: bool = Field(
-        default=True,
+        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."
diff --git a/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py b/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py
index 2f7b7fa6b512..f9bd6119c2fc 100644
--- a/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py
+++ b/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py
@@ -149,7 +149,7 @@ def _assert_lora_changes_output(out_lora, out_base):
     assert any_differ, "LoRA outputs identical to base model (same tokens AND same logprobs)"
 
 
-def _run_lora_test(model_path, target_modules, trtllm_modules, dtype=torch.bfloat16):
+def _run_lora_test(model_path, target_modules, trtllm_modules, dtype=torch.bfloat16, overlap=False):
     """End-to-end helper: create adapter, run inference, assert output differs."""
     with tempfile.TemporaryDirectory() as tmpdir:
         lora_dir = _create_lora_adapter(
@@ -164,6 +164,7 @@ def _run_lora_test(model_path, target_modules, trtllm_modules, dtype=torch.bfloa
             lora_target_modules=trtllm_modules,
             max_lora_rank=16,
             max_loras=2,
+            overlap_lora_and_base=overlap,
         )
         out_lora, out_base = _run_with_and_without_lora(
             model_path,
@@ -196,6 +197,23 @@ def test_qwen3_fp8_lora(self):
             dtype=torch.float8_e4m3fn,
         )
 
+    def test_qwen3_bf16_lora_overlap(self):
+        _run_lora_test(
+            self.model_path,
+            {**_ATTN_LORA_MODULES, **_MLP_LORA_MODULES},
+            _ATTN_TRTLLM_MODULES + _MLP_TRTLLM_MODULES,
+            overlap=True,
+        )
+
+    def test_qwen3_fp8_lora_overlap(self):
+        _run_lora_test(
+            self.model_path,
+            {**_ATTN_LORA_MODULES, **_MLP_LORA_MODULES},
+            _ATTN_TRTLLM_MODULES + _MLP_TRTLLM_MODULES,
+            dtype=torch.float8_e4m3fn,
+            overlap=True,
+        )
+
 
 @skip_gpu_memory_less_than_80gb
 class TestQwen3MoELoRA:

From 6380bb9944eb74c51fec5c3eba01c9d86b562ba0 Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Fri, 31 Jul 2026 08:50:53 -0700
Subject: [PATCH 08/14] Move ownership of aux LoRA CUDA stream to
 CudaGraphLoraManager

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 .../_torch/peft/lora/cuda_graph_lora_manager.py   |  4 ++--
 tensorrt_llm/_torch/peft/lora/layer.py            | 15 ++++++---------
 2 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py b/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py
index 3fdf5f0cb231..eef3596d14e2 100644
--- a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py
+++ b/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py
@@ -69,7 +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.overlap_lora_and_base = overlap_lora_and_base
+        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))
@@ -228,7 +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,
-            "overlap_lora_and_base": self.overlap_lora_and_base,
+            "lora_aux_stream": self.lora_aux_stream,
         }
 
         return lora_params
diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py
index 4ed0b00b81f2..dace2094fd7f 100644
--- a/tensorrt_llm/_torch/peft/lora/layer.py
+++ b/tensorrt_llm/_torch/peft/lora/layer.py
@@ -223,9 +223,6 @@ def add_lora_result(output: torch.Tensor,
 
 class LoraLayer(torch.nn.Module):
 
-    # Static aux stream for LoRA computations
-    _aux_stream: torch.cuda.Stream | None = None
-
     def __init__(self, lora_module_types: List[LoraModuleType],
                  output_hidden_sizes: List[int]):
         super().__init__()
@@ -268,8 +265,9 @@ def forward_with_base(
                 module_ids=tuple(layer.lora_module_types),
             ) in cuda_graph_params.layer_info for layer in lora_layers)
 
-        execute_in_parallel = (has_lora_layer and lora_params.get(
-            "overlap_lora_and_base", True) and do_multi_stream()
+        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
@@ -279,9 +277,8 @@ def lora_forward() -> tuple[torch.Tensor | None, ...]:
                 for lora_layer in lora_layers)
 
         if execute_in_parallel:
-            # Lazy allocation of aux stream and parallel events
-            if LoraLayer._aux_stream is None:
-                LoraLayer._aux_stream = torch.cuda.Stream()
+            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()
@@ -292,7 +289,7 @@ def lora_forward() -> tuple[torch.Tensor | None, ...]:
                 lora_forward,
                 lora_layers[0]._par_events[0],
                 lora_layers[0]._par_events[1],
-                LoraLayer._aux_stream,
+                lora_aux_stream,
                 disable_on_compile=True,
             )
         else:

From 040505498d80f5c46d1918cd70213206bb1c020f Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Wed, 5 Aug 2026 02:01:12 -0700
Subject: [PATCH 09/14] Call record_stream before lora_output add

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/peft/lora/layer.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py
index dace2094fd7f..138926147b5d 100644
--- a/tensorrt_llm/_torch/peft/lora/layer.py
+++ b/tensorrt_llm/_torch/peft/lora/layer.py
@@ -298,6 +298,8 @@ def lora_forward() -> tuple[torch.Tensor | None, ...]:
         for lora_output in lora_outputs:
             if lora_output is None:
                 continue
+            if execute_in_parallel:
+                lora_output.record_stream(torch.cuda.current_stream())
             base_output.add_(lora_output)
 
         return base_output

From 32d3f7f466fd662710cc85ec8410ea94f1c21b28 Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Wed, 5 Aug 2026 23:37:06 -0700
Subject: [PATCH 10/14] Change default for overlap_lora_and_base in
 CudaGraphLoraManager to False

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py b/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py
index eef3596d14e2..ad8ece4dff6c 100644
--- a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py
+++ b/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py
@@ -44,7 +44,7 @@ def __init__(
         max_lora_rank: int,
         model: torch.nn.Module,
         lora_model_config: Optional[LoraModelConfig],
-        overlap_lora_and_base: bool = True,
+        overlap_lora_and_base: bool = False,
         device: str = "cuda",
         max_tokens_per_seq: int = 1,
     ):

From 1d0572f6f3ad8b4d0b6c43b3b03080c98177616f Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Thu, 6 Aug 2026 00:13:16 -0700
Subject: [PATCH 11/14] Add focused test for LoraLayer.forward_with_base()

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 .../tests_lora_modules/test_qwen3_sanity.py   | 36 +++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py b/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py
index f9bd6119c2fc..cee182fa0ba4 100644
--- a/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py
+++ b/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py
@@ -17,6 +17,7 @@
 import json
 import os
 import tempfile
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -25,7 +26,10 @@
 from utils.util import skip_gpu_memory_less_than_80gb
 
 from tensorrt_llm import LLM, SamplingParams
+from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
+from tensorrt_llm._torch.peft.lora import layer as lora_layer_module
 from tensorrt_llm._torch.peft.lora.config import LoraConfig
+from tensorrt_llm._torch.peft.lora.layer import LoraLayer, LoraModuleType
 from tensorrt_llm.executor.request import LoRARequest
 
 # HF module name -> block path relative to layers.{idx}.
@@ -175,6 +179,38 @@ def _run_lora_test(model_path, target_modules, trtllm_modules, dtype=torch.bfloa
         _assert_lora_changes_output(out_lora, out_base)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="LoRA overlap requires CUDA streams.")
+def test_lora_forward_with_base_executes_overlap(monkeypatch):
+    """Verify that the LoRA overlap path uses its auxiliary CUDA stream."""
+    lora_layer = LoraLayer([LoraModuleType.ATTENTION_Q], [2])
+    layer_key = lora_layer_module.CudaGraphLoraParams.LoraLayerKey(
+        layer_idx=0, module_ids=tuple(lora_layer.lora_module_types)
+    )
+    aux_stream = torch.cuda.Stream()
+    parallel_streams = []
+    original_parallel_executor = lora_layer_module.maybe_execute_in_parallel
+
+    def record_parallel_executor(*args, **kwargs):
+        parallel_streams.append(args[4])
+        return original_parallel_executor(*args, **kwargs)
+
+    monkeypatch.setattr(lora_layer_module, "maybe_execute_in_parallel", record_parallel_executor)
+    monkeypatch.setattr(LoraLayer, "forward", lambda self, x, *_: torch.ones_like(x))
+
+    x = torch.ones((2, 2), device="cuda")
+    lora_params = {
+        "cuda_graph_params": SimpleNamespace(layer_info={layer_key: object()}),
+        "lora_aux_stream": aux_stream,
+    }
+    with with_multi_stream(True):
+        output = LoraLayer.forward_with_base(lambda: x.clone(), (lora_layer,), x, lora_params, 0)
+
+    assert lora_layer._par_events is not None
+    assert len(parallel_streams) == 1
+    assert parallel_streams[0] is aux_stream
+    torch.testing.assert_close(output, 2 * x)
+
+
 class TestQwen3LoRA:
     @pytest.fixture(autouse=True)
     def setup(self):

From b239a4e9d4a7f4c2a053522c9007c616ec8f1a2a Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Sun, 16 Aug 2026 23:53:31 -0700
Subject: [PATCH 12/14] Use add_lora_result() in forward_with_base()

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/attention.py | 2 +-
 tensorrt_llm/_torch/modules/gated_mlp.py | 2 +-
 tensorrt_llm/_torch/modules/linear.py    | 2 +-
 tensorrt_llm/_torch/modules/mlp.py       | 2 +-
 tensorrt_llm/_torch/peft/lora/layer.py   | 2 +-
 5 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index a81197746311..73af81f9cf26 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/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,
diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py
index 0f1d2ce9f0d4..afd8796c3e10 100644
--- a/tensorrt_llm/_torch/modules/gated_mlp.py
+++ b/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)
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index d3631806e809..cda84c8aa1c1 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/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,
diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py
index df654f627d88..c1a7bcab10cd 100644
--- a/tensorrt_llm/_torch/modules/mlp.py
+++ b/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)
diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py
index 138926147b5d..672ebf478d1b 100644
--- a/tensorrt_llm/_torch/peft/lora/layer.py
+++ b/tensorrt_llm/_torch/peft/lora/layer.py
@@ -300,7 +300,7 @@ def lora_forward() -> tuple[torch.Tensor | None, ...]:
                 continue
             if execute_in_parallel:
                 lora_output.record_stream(torch.cuda.current_stream())
-            base_output.add_(lora_output)
+            base_output = add_lora_result(base_output, lora_output)
 
         return base_output
 

From b1822cb477696b27a6930cc857ac91ac977e8778 Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Sun, 16 Aug 2026 23:59:35 -0700
Subject: [PATCH 13/14] Use in-place add_() in add_lora_result()

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/peft/lora/layer.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py
index 672ebf478d1b..bab159841de5 100644
--- a/tensorrt_llm/_torch/peft/lora/layer.py
+++ b/tensorrt_llm/_torch/peft/lora/layer.py
@@ -216,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):

From 2a9945130efeccf565c54e26987ff30907134546 Mon Sep 17 00:00:00 2001
From: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Date: Mon, 17 Aug 2026 00:04:05 -0700
Subject: [PATCH 14/14] More robust check for lora output type

Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
---
 tensorrt_llm/_torch/peft/lora/layer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py
index bab159841de5..f26ea410b535 100644
--- a/tensorrt_llm/_torch/peft/lora/layer.py
+++ b/tensorrt_llm/_torch/peft/lora/layer.py
@@ -296,7 +296,7 @@ def lora_forward() -> tuple[torch.Tensor | None, ...]:
             base_output, lora_outputs = base_forward(), lora_forward()
 
         for lora_output in lora_outputs:
-            if lora_output is None:
+            if not isinstance(lora_output, torch.Tensor):
                 continue
             if execute_in_parallel:
                 lora_output.record_stream(torch.cuda.current_stream())