Use torch.addmm instead of separate mm and add_ calls for LoRA torch.native (#20562)

Co-authored-by: Satyam Kumar <satyamk@linkedin.com>

python/sglang/srt/lora/backend/torch_backend.py

@@ -23,6 +23,9 @@ class TorchNativeLoRABatchInfo(LoRABatchInfo):
     # The index of lora adapter used by each segment, in shape (num_segments,) placed on cpu device
     weight_indices_cpu: Optional[torch.Tensor] = None
 
+    # Scaling factors for each lora adapter, in shape (lora_num,) placed on cpu device
+    scalings_cpu: Optional[torch.Tensor] = None
+
 
 class TorchNativeLoRABackend(BaseLoRABackend):
     name = "torch_native"
@@ -44,7 +47,7 @@ class TorchNativeLoRABackend(BaseLoRABackend):
             weight_indices=self.batch_info.weight_indices_cpu,
             seg_len_tensor=self.batch_info.seg_lens_cpu,
             lora_ranks=self.batch_info.lora_ranks_cpu,
-            scaling_tensor=self.batch_info.scalings,
+            scaling_tensor=self.batch_info.scalings_cpu,
             num_slices=1,
         )
 
@@ -93,7 +96,7 @@ class TorchNativeLoRABackend(BaseLoRABackend):
             weight_indices=self.batch_info.weight_indices_cpu,
             seg_len_tensor=self.batch_info.seg_lens_cpu,
             lora_ranks=self.batch_info.lora_ranks_cpu,
-            scaling_tensor=self.batch_info.scalings,
+            scaling_tensor=self.batch_info.scalings_cpu,
             num_slices=num_slices,
         )
 
@@ -131,7 +134,7 @@ class TorchNativeLoRABackend(BaseLoRABackend):
             weight_indices=self.batch_info.weight_indices_cpu,
             seg_len_tensor=self.batch_info.seg_lens_cpu,
             lora_ranks=self.batch_info.lora_ranks_cpu,
-            scaling_tensor=self.batch_info.scalings,
+            scaling_tensor=self.batch_info.scalings_cpu,
             num_slices=num_slices,
         )
 
@@ -274,5 +277,6 @@ class TorchNativeLoRABackend(BaseLoRABackend):
         batch_info.seg_indptr_cpu = seg_indptr_cpu
         batch_info.seg_lens_cpu = seg_lens_cpu
         batch_info.weight_indices_cpu = weight_indices_tensor
+        batch_info.scalings_cpu = scalings_tensor
 
         self.batch_info = batch_info

python/sglang/srt/lora/torch_ops/lora_ops.py

@@ -31,13 +31,19 @@ def sgemm_lora_a_fwd(
             continue
 
         if rank > 0:
-
             x_seq = inputs[token_offset : token_offset + seq_len, :]
             w_seq = weights[lora_idx, : num_slices * rank, :]
 
-            result = torch.mm(x_seq, w_seq.T)
+            out_slice = output[
-            output[token_offset : token_offset + seq_len, : num_slices * rank] = (
+                token_offset : token_offset + seq_len, : num_slices * rank
-                scaling_tensor[lora_idx] * result
+            ]
+            torch.addmm(
+                out_slice,
+                x_seq,
+                w_seq.T,
+                beta=0,
+                alpha=scaling_tensor[lora_idx].item(),
+                out=out_slice,
             )
 
         token_offset += seq_len
@@ -98,10 +104,11 @@ def sgemm_lora_b_fwd(
                 lora_idx, slice_start_output:slice_end_output, :rank
             ]  # (slice_dim, rank)
 
-            output[
+            out_slice = output[
                 token_offset : token_offset + seq_len,
                 slice_start_output:slice_end_output,
-            ].add_(torch.mm(x_slice, w_slice.T))
+            ]
+            torch.addmm(out_slice, x_slice, w_slice.T, beta=1, alpha=1, out=out_slice)
 
         token_offset += seq_len