chore: import upstream snapshot with attribution

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

@@ -0,0 +1,302 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from sglang.srt.lora.backend.base_backend import BaseLoRABackend
+from sglang.srt.lora.torch_ops import (
+    sgemm_lora_a_embedding_fwd,
+    sgemm_lora_a_fwd,
+    sgemm_lora_b_fwd,
+)
+from sglang.srt.lora.utils import LoRABatchInfo, generate_sequence_lengths
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+@dataclass
+class TorchNativeLoRABatchInfo(LoRABatchInfo):
+    # ranks of each lora adapter, in shape (lora_num,) placed on cpu device
+    lora_ranks_cpu: Optional[torch.Tensor] = None
+
+    # Indice pointers of each segment in shape (num_segments + 1, ) placed on cpu device
+    seg_indptr_cpu: Optional[torch.Tensor] = None
+
+    # Lengths of each segments in shape (num_segments,) placed on cpu device
+    seg_lens_cpu: Optional[torch.Tensor] = None
+
+    # 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"
+
+    def __init__(
+        self,
+        max_loras_per_batch: int,
+        device: torch.device,
+        **kwargs,
+    ):
+        super().__init__(max_loras_per_batch, device)
+
+    def run_lora_a_embedding(
+        self,
+        input_ids: torch.Tensor,
+        weights: torch.Tensor,
+        vocab_size: int,
+        extra_embeddings: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        assert (
+            extra_embeddings is None
+        ), "Extra embeddings for lora a is not supported yet in chunked backend"
+        output_tensor = sgemm_lora_a_embedding_fwd(
+            inputs=input_ids,
+            weights=weights,
+            batch_info=self.batch_info,
+            vocab_size=vocab_size,
+        )
+
+        return output_tensor
+
+    def run_lora_a_sgemm(
+        self,
+        x: torch.Tensor,
+        weights: torch.Tensor,
+        stack_num: int = 1,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        output_tensor = sgemm_lora_a_fwd(
+            inputs=x,
+            weights=weights,
+            batch_info=self.batch_info,
+            num_slices=stack_num,
+        )
+
+        return output_tensor
+
+    def run_lora_b_sgemm(
+        self,
+        x: torch.Tensor,
+        weights: torch.Tensor,
+        output_offset_cpu: torch.Tensor,
+        base_output: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        _, weight_out_dim, _ = weights.shape
+
+        output_tensor = sgemm_lora_b_fwd(
+            inputs=x,
+            weights=weights,
+            batch_info=self.batch_info,
+            slice_offsets=output_offset_cpu,
+            base_output=base_output,
+        )
+
+        return output_tensor
+
+    def run_qkv_lora(
+        self,
+        x: torch.Tensor,
+        qkv_lora_a: torch.Tensor,
+        qkv_lora_b: torch.Tensor,
+        output_offset: torch.Tensor,
+        output_offset_cpu: torch.Tensor,
+        max_qkv_out_dim: int,
+        base_output: torch.Tensor = None,
+        n_slices: int = 3,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        lora_a_output = sgemm_lora_a_fwd(
+            inputs=x,
+            weights=qkv_lora_a,
+            batch_info=self.batch_info,
+            num_slices=n_slices,
+        )
+
+        output_tensor = sgemm_lora_b_fwd(
+            inputs=lora_a_output,
+            weights=qkv_lora_b,
+            batch_info=self.batch_info,
+            slice_offsets=output_offset_cpu,
+            base_output=base_output,
+        )
+
+        return output_tensor
+
+    def run_gate_up_lora(
+        self,
+        x: torch.Tensor,
+        gate_up_lora_a: torch.Tensor,
+        gate_up_lora_b: torch.Tensor,
+        output_offset_cpu: torch.Tensor,
+        base_output: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        num_slices = len(output_offset_cpu) - 1
+        _, weight_out_dim, _ = gate_up_lora_b.shape
+
+        lora_a_output = sgemm_lora_a_fwd(
+            inputs=x,
+            weights=gate_up_lora_a,
+            batch_info=self.batch_info,
+            num_slices=num_slices,
+        )
+
+        output_tensor = sgemm_lora_b_fwd(
+            inputs=lora_a_output,
+            weights=gate_up_lora_b,
+            batch_info=self.batch_info,
+            slice_offsets=output_offset_cpu,
+            base_output=base_output,
+        )
+
+        return output_tensor
+
+    def init_cuda_graph_batch_info(
+        self,
+        max_bs_in_cuda_graph: int,
+        num_tokens_per_bs: int,
+    ):
+        with torch.device("cuda"):
+            self.cuda_graph_batch_info = TorchNativeLoRABatchInfo(
+                use_cuda_graph=True,
+                bs=max_bs_in_cuda_graph,
+                num_segments=self.max_loras_per_batch,
+                seg_lens=torch.full(
+                    (max_bs_in_cuda_graph,), num_tokens_per_bs, dtype=torch.int32
+                ),
+                seg_indptr=torch.zeros(max_bs_in_cuda_graph + 1, dtype=torch.int32),
+                weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
+                lora_ranks=torch.zeros(self.max_loras_per_batch, dtype=torch.int32),
+                scalings=torch.zeros(self.max_loras_per_batch, dtype=torch.float),
+                permutation=None,
+                max_len=num_tokens_per_bs,
+            )
+
+            # Initialize seg_indptr for CUDA graph as they remain constant
+            # across batches.
+            torch.cumsum(
+                self.cuda_graph_batch_info.seg_lens[:max_bs_in_cuda_graph],
+                dim=0,
+                out=self.cuda_graph_batch_info.seg_indptr[1 : max_bs_in_cuda_graph + 1],
+            )
+
+    def prepare_lora_batch(
+        self,
+        forward_batch: ForwardBatch,
+        weight_indices: list[int],
+        lora_ranks: list[int],
+        scalings: list[float],
+        use_cuda_graph: bool,
+    ):
+        # Do not use merge optimization for graph mode
+        # Use pinned memory to avoid synchronizations during host-to-device transfer
+        original_seq_lens_cpu = generate_sequence_lengths(forward_batch, device="cpu")
+        if not use_cuda_graph:
+            original_weight_indices_tensor = torch.tensor(
+                weight_indices, dtype=torch.int32, device="cpu"
+            )
+
+            unique_weight_indices_tensor, inverse_weight_indices_tensor = (
+                torch.unique_consecutive(
+                    original_weight_indices_tensor, return_inverse=True
+                )
+            )
+
+            seg_lens_cpu = (
+                torch.zeros_like(
+                    unique_weight_indices_tensor, dtype=torch.int32, device="cpu"
+                )
+                .scatter_add_(
+                    0,
+                    inverse_weight_indices_tensor,
+                    original_seq_lens_cpu,
+                )
+                .pin_memory()
+            )
+
+            weight_indices_tensor = unique_weight_indices_tensor.pin_memory()
+        else:
+            weight_indices_tensor = torch.repeat_interleave(
+                torch.tensor(weight_indices, dtype=torch.int32, device="cpu"),
+                original_seq_lens_cpu,
+            ).pin_memory()
+            seg_lens_cpu = torch.ones_like(weight_indices_tensor).pin_memory()
+
+        seg_indptr_cpu = torch.zeros(
+            (len(seg_lens_cpu) + 1,), dtype=torch.int32, pin_memory=True
+        )
+        seg_indptr_cpu[1:] = torch.cumsum(seg_lens_cpu, dim=0)
+        lora_ranks_tensor = torch.tensor(
+            lora_ranks, dtype=torch.int32, pin_memory=True, device="cpu"
+        )
+        scalings_tensor = torch.tensor(
+            scalings, dtype=torch.float, pin_memory=True, device="cpu"
+        )
+
+        bs = forward_batch.batch_size
+        num_segments = len(weight_indices_tensor)
+
+        if use_cuda_graph:
+            assert (
+                self.cuda_graph_batch_info is not None
+            ), "CUDA Graph batch info is not initialized."
+            batch_info = self.cuda_graph_batch_info
+            batch_info.bs = forward_batch.batch_size
+            batch_info.num_segments = num_segments
+        else:
+            max_len = max(seg_lens_cpu)
+
+            batch_info = TorchNativeLoRABatchInfo(
+                bs=forward_batch.batch_size,
+                num_segments=num_segments,
+                max_len=max_len,
+                use_cuda_graph=False,
+                seg_lens=torch.empty((bs,), dtype=torch.int32, device=self.device),
+                seg_indptr=torch.empty(
+                    (bs + 1,), dtype=torch.int32, device=self.device
+                ),
+                weight_indices=torch.empty(
+                    (bs,), dtype=torch.int32, device=self.device
+                ),
+                lora_ranks=torch.empty(
+                    (self.max_loras_per_batch,), dtype=torch.int32, device=self.device
+                ),
+                scalings=torch.empty(
+                    (self.max_loras_per_batch,), dtype=torch.float, device=self.device
+                ),
+                permutation=None,
+            )
+
+        # Copy to device asynchronously
+        batch_info.lora_ranks[: self.max_loras_per_batch].copy_(
+            lora_ranks_tensor, non_blocking=True
+        )
+        batch_info.scalings[: self.max_loras_per_batch].copy_(
+            scalings_tensor, non_blocking=True
+        )
+        batch_info.weight_indices[:num_segments].copy_(
+            weight_indices_tensor, non_blocking=True
+        )
+        batch_info.seg_indptr[: len(seg_indptr_cpu)].copy_(
+            seg_indptr_cpu, non_blocking=True
+        )
+        batch_info.seg_lens[: len(seg_lens_cpu)].copy_(seg_lens_cpu, non_blocking=True)
+
+        batch_info.lora_ranks_cpu = lora_ranks_tensor
+        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
+
+        batch_info = self._add_moe_lora_info(forward_batch, batch_info)
+        self.batch_info = batch_info