refactor lora tp1

lmdeploy/pytorch/backends/base.py

@@ -20,7 +20,7 @@ class OpType(Enum):
     GeluAndMul = auto()
     RMSNorm = auto()
     LayerNorm = auto()
-    SLoRA = auto()
+    LoRA = auto()
     LinearW8A8 = auto()
     RMSNormW8A8 = auto()
     MultinomialSampling = auto()

lmdeploy/pytorch/backends/cuda/lora.py

@@ -0,0 +1,80 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from dataclasses import dataclass
+
+import torch
+
+from lmdeploy.pytorch.kernels.cuda.fused_lora import fused_lora
+from lmdeploy.pytorch.model_inputs import StepContextManager
+
+from ..lora import AdapterInfo, LoRABuilder, LoRAImpl
+
+
+@dataclass
+class PackedLoRAInput:
+    """packed lora input."""
+    x: torch.Tensor
+    q_start_loc: torch.Tensor
+    q_seqlens: torch.Tensor
+    adapter_ids: torch.Tensor
+    max_seq_len: int
+    is_decoding: bool
+
+
+class TritonLoRAImpl(LoRAImpl):
+    """triton lora implementation."""
+
+    @staticmethod
+    def _make_packed_lora_input(x, ctx_mgr):
+        """make PackedLoRAInput."""
+        context = ctx_mgr.current_context()
+
+        # adapter cache
+        max_q_seq_length = x.numel() // x.size(-1)
+
+        return PackedLoRAInput(x=x.flatten(0, -2).contiguous(),
+                               q_start_loc=context.q_start_loc,
+                               q_seqlens=context.q_seqlens,
+                               adapter_ids=context.local_adapter_ids,
+                               max_seq_len=max_q_seq_length,
+                               is_decoding=context.is_decoding)
+
+    def forward(self,
+                x: torch.Tensor,
+                lora_A: torch.Tensor,
+                lora_B: torch.Tensor,
+                base_output: torch.Tensor,
+                adapter_info: AdapterInfo,
+                ctx_mgr: StepContextManager,
+                colwise: bool,
+                is_tp: bool = True):
+        """forward."""
+        lora_input = self._make_packed_lora_input(x, ctx_mgr)
+
+        lora_out = fused_lora(lora_input.x,
+        lora_A,
+        lora_B,
+        scaling=adapter_info.scalings,
+        rank_start=adapter_info.rank_offsets,
+        ranks=adapter_info.ranks,
+        seq_start=lora_input.q_start_loc,
+        seq_lens=lora_input.q_seqlens,
+        adapter_ids=lora_input.adapter_ids,
+        max_rank=adapter_info.max_rank,
+        max_seqlen=lora_input.max_seq_len,
+        )
+
+        base_slice = adapter_info.base_slice
+        sliced_base = base_output[..., base_slice]
+        lora_out = lora_out.reshape(sliced_base.shape)
+        sliced_base.add_(lora_out)
+        output = base_output
+        return output
+
+
+class TritonLoRABuilder(LoRABuilder):
+    """triton lora layer builder."""
+
+    @staticmethod
+    def build():
+        """build."""
+        return TritonLoRAImpl()

lmdeploy/pytorch/backends/cuda/op_backend.py

@@ -32,9 +32,9 @@ def get_layer_impl_builder(cls, layer_type: OpType):
         elif layer_type == OpType.RMSNorm:
             from .norm import TritonRMSNormBuilder
             return TritonRMSNormBuilder
-        elif layer_type == OpType.SLoRA:
-            from .slora import TritonSLoRABuilder
-            return TritonSLoRABuilder
+        elif layer_type == OpType.LoRA:
+            from .lora import TritonLoRABuilder
+            return TritonLoRABuilder
         elif layer_type == OpType.LinearW8A8:
             from .qmodules import TritonLinearW8A8Builder
             return TritonLinearW8A8Builder

lmdeploy/pytorch/backends/slora.py → lmdeploy/pytorch/backends/lora.py

@@ -1,6 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 import torch
 
@@ -14,20 +14,28 @@ class AdapterInfo:
     out_features: int
     ranks: torch.Tensor
     scalings: torch.Tensor
-    rank_offsets: torch.Tensor
-    a_cache: torch.Tensor
-    b_cache: torch.Tensor
     base_slice: slice
-    max_rank: int
+    rank_offsets: torch.Tensor = field(init=False)
+    max_rank: int = field(init=False)
 
+    def __post_init__(self):
+        """post init."""
+        ranks = self.ranks
+        rank_offsets = ranks.cumsum(0) - ranks
+        max_rank = ranks.max().item()
+        self.rank_offsets = rank_offsets
+        self.max_rank = max_rank
 
-class SLoRAImpl(ABC):
-    """slora implementation api."""
+
+class LoRAImpl(ABC):
+    """lora implementation."""
 
     @abstractmethod
     def forward(self,
                 x: torch.Tensor,
                 base_output: torch.Tensor,
+                lora_A: torch.Tensor,
+                lora_B: torch.Tensor,
                 adapter_info: AdapterInfo,
                 ctx_mgr: StepContextManager,
                 colwise: bool,
@@ -36,8 +44,8 @@ def forward(self,
         raise NotImplementedError
 
 
-class SLoRABuilder(ABC):
-    """slora implementation builder."""
+class LoRABuilder(ABC):
+    """lora implementation builder."""
 
     @staticmethod
     @abstractmethod