[LoRA] Support dual CUDA streams-Linear Layer

tests/lora/test_olmoe_tp.py

@@ -103,7 +103,7 @@ def test_olmoe_lora(olmoe_lora_files):
         max_model_len=1024,
         enable_lora=True,
         max_loras=4,
-        enforce_eager=True,
+        # enforce_eager=True,
         trust_remote_code=True,
         enable_chunked_prefill=True,
     )

vllm/lora/layers/base_linear.py

@@ -5,41 +5,72 @@
 import torch
 from transformers import PretrainedConfig
 
+from vllm.config import get_current_vllm_config
 from vllm.config.lora import LoRAConfig
 from vllm.distributed.utils import divide
+from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     LinearBase,
     ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import current_stream, direct_register_custom_op
 
 from .base import BaseLayerWithLoRA
 from .utils import _get_lora_device
 
+# aux_stream() is shared for:
+#   - LoRA dual-stream: base linear and LoRA compute on different CUDA streams
+_aux_stream: torch.cuda.Stream | None = None
+
+
+def aux_stream() -> torch.cuda.Stream | None:
+    """
+    Returns the auxiliary CUDA stream for overlapping compute.
+    Initialized only once on CUDA-alike platforms.
+    """
+    global _aux_stream
+
+    if _aux_stream is None and current_platform.is_cuda_alike():
+        _aux_stream = torch.cuda.Stream()
+
+    return _aux_stream
+
 
 class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
     def __init__(self, base_layer: LinearBase):
         super().__init__()
+
+        self._lora_stream = aux_stream()
         self.base_layer = base_layer
         self.input_size = self.base_layer.input_size
         # Ensure tp_size and tp_rank consistency with the base_layer.
         self.tp_size = self.base_layer.tp_size
         self.tp_rank = self.base_layer.tp_rank
         self.device = _get_lora_device(self.base_layer)
+        self._init_lora_stream_context()
         self.output_slices: tuple[int, ...]
         self.output_size: int
         self.n_slices: int
 
+    def _init_lora_stream_context(self) -> None:
+        vllm_config = get_current_vllm_config()
+        # lora_linear avoids prefix conflicts with the base layer
+        self.layer_name = self.base_layer.prefix + ".lora_linear"
+        compilation_config = vllm_config.compilation_config
+        if self.layer_name in compilation_config.static_forward_context:
+            raise ValueError("Duplicate layer name: {}".format(self.layer_name))
+        compilation_config.static_forward_context[self.layer_name] = self
+
     def create_lora_weights(
         self,
         max_loras: int,
         lora_config: LoRAConfig,
         model_config: PretrainedConfig | None = None,
     ) -> None:
         self.lora_config = lora_config
-        #
         if isinstance(self.base_layer, ReplicatedLinear):
             lora_a_out_size = lora_config.max_lora_rank
             lora_b_out_size = self.output_size
@@ -120,6 +151,18 @@ def set_lora(
         )
 
     def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
+        if self._lora_stream is None:
+            return self._apply_sync(x, bias)
+        else:
+            num_tokens = x.size(0)
+            output_size = sum(self.output_slices)
+            return torch.ops.vllm.lora_linear(
+                self.layer_name, num_tokens, output_size, x, bias
+            )
+
+    def _apply_sync(
+        self, x: torch.Tensor, bias: torch.Tensor | None = None
+    ) -> torch.Tensor:
         output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
 
         original_shape = output.shape if output.ndim == 3 else None
@@ -144,6 +187,59 @@ def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tens
 
         return output
 
+    def _apply_async_impl(
+        self, x: torch.Tensor, bias: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """
+        Forward pass with base linear and LoRA on separate CUDA streams for overlap.
+        Base layer runs on default stream; LoRA runs on aux stream.
+        """
+        assert isinstance(self._lora_stream, torch.cuda.Stream)  # Make mypy happy
+        self._lora_stream.wait_stream(current_stream())
+        output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
+
+        # Pre-allocate buffer for LoRA-only output (add_inputs=False writes here).
+        # Infer shape from x and output_slices; avoid using output tensor.
+        num_tokens = x.size(0) if x.ndim == 2 else x.size(1)
+        output_size = sum(self.output_slices)
+
+        lora_output = torch.empty(
+            (num_tokens, output_size),
+            device=self.device,
+            dtype=x.dtype,
+        )
+        with torch.cuda.stream(self._lora_stream):
+            # LoRA stream waits for base layer output before reading.
+            self._lora_stream.wait_stream(current_stream())
+            self.punica_wrapper.add_lora_linear(
+                lora_output,
+                x,
+                self.lora_a_stacked,
+                self.lora_b_stacked,
+                1.0,
+                self.output_slices,
+                add_inputs=False,
+            )
+        # Default stream waits for LoRA to complete before combining.
+        current_stream().wait_stream(self._lora_stream)
+        original_shape = output.shape if output.ndim == 3 else None
+
+        # In transformers backend, x and output have extra batch dimension like
+        # (1, seq_len, hidden_dim), while punica expects (seq_len, hidden_dim),
+        # therefore we need to flatten the batch dimensions.
+        if x.ndim == 3 and output.ndim == 3:
+            output = output.flatten(0, 1)
+            x = x.flatten(0, 1)
+
+        output = output + lora_output
+
+        # Reshape the flattened output back to its original shape,
+        # as some MM encoders cannot handle flattened inputs.
+        if original_shape is not None:
+            output = output.reshape(original_shape)
+
+        return output
+
     @property
     def weight(self) -> torch.Tensor:
         # unquantizedLinear
@@ -167,3 +263,36 @@ def bias(self) -> torch.Tensor | None:
             return self.base_layer.bias
         else:
             return None
+
+
+def lora_linear(
+    layer_name: str,
+    num_tokens: int,
+    output_size: int,
+    x: torch.Tensor,
+    bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    forward_context: ForwardContext = get_forward_context()
+    self = forward_context.no_compile_layers[layer_name]
+    return self._apply_async_impl(x, bias)
+
+
+def lora_linear_fake(
+    layer_name: str,
+    num_tokens: int,
+    output_size: int,
+    x: torch.Tensor,
+    bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    return torch.empty(
+        (num_tokens, output_size),
+        device=x.device,
+        dtype=x.dtype,
+    )
+
+
+direct_register_custom_op(
+    op_name="lora_linear",
+    op_func=lora_linear,
+    fake_impl=lora_linear_fake,
+)

vllm/lora/punica_wrapper/punica_gpu.py

@@ -144,7 +144,9 @@ def add_expand(
             x (torch.Tensor): Input tensors
             lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight
             output_slices (tuple[int, ...]): Every slice's size
-            add_inputs (bool): Defaults to True.
+            add_inputs (bool): If True, add LoRA output to y; if False, write
+                LoRA-only output to y (used for dual-stream when base and LoRA
+                run on different CUDA streams). Defaults to True.
         """
         y_org = y
         y = y.view(-1, y.shape[-1])
@@ -161,7 +163,7 @@ def add_expand(
                 num_tokens, self.lora_config.specialize_active_lora
             ),
             offset_start=offset_start,
-            add_inputs=True,
+            add_inputs=add_inputs,
         )
 
         y = y.view_as(y_org)
@@ -244,7 +246,7 @@ def add_lora_linear(
         buffer = torch.empty(
             (len(output_slices), x.size(0), r), dtype=torch.float32, device=x.device
         )
-
+        add_inputs = kwargs.pop("add_inputs", True)
         self.add_shrink(
             buffer,  # type: ignore
             x,
@@ -257,7 +259,7 @@ def add_lora_linear(
             buffer,  # type: ignore
             lora_b_stacked,
             output_slices,
-            add_inputs=True,
+            add_inputs=add_inputs,
             **kwargs,
         )