sglang 0.5.2rc2__py3-none-any.whl → 0.5.3rc2__py3-none-any.whl

sglang/srt/lora/lora_manager.py

@@ -21,7 +21,6 @@ from typing import Dict, Iterable, List, Optional, Set, Tuple
 import torch
 
 from sglang.srt.configs.load_config import LoadConfig
-from sglang.srt.hf_transformers_utils import AutoConfig
 from sglang.srt.lora.backend.base_backend import BaseLoRABackend, get_backend_from_name
 from sglang.srt.lora.layers import BaseLayerWithLoRA, get_lora_layer
 from sglang.srt.lora.lora import LoRAAdapter
@@ -35,9 +34,11 @@ from sglang.srt.lora.utils import (
     get_normalized_target_modules,
     get_target_module_name,
 )
-from sglang.srt.managers.io_struct import LoRAUpdateResult
+from sglang.srt.managers.io_struct import LoRAUpdateOutput
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import replace_submodule
+from sglang.srt.utils.hf_transformers_utils import AutoConfig
 
 logger = logging.getLogger(__name__)
 
@@ -56,6 +57,7 @@ class LoRAManager:
         max_lora_rank: Optional[int] = None,
         target_modules: Optional[Iterable[str]] = None,
         lora_paths: Optional[List[LoRARef]] = None,
+        server_args: Optional[ServerArgs] = None,
     ):
         self.base_model: torch.nn.Module = base_model
         self.base_hf_config: AutoConfig = base_hf_config
@@ -69,7 +71,11 @@ class LoRAManager:
         # LoRA backend for running sgemm kernels
         logger.info(f"Using {lora_backend} as backend of LoRA kernels.")
         backend_type = get_backend_from_name(lora_backend)
-        self.lora_backend: BaseLoRABackend = backend_type(lora_backend)
+        self.lora_backend: BaseLoRABackend = backend_type(
+            max_loras_per_batch=max_loras_per_batch,
+            device=self.device,
+            server_args=server_args,
+        )
 
         # Initialize mutable internal state of the LoRAManager.
         self.init_state(
@@ -82,34 +88,27 @@ class LoRAManager:
         self.max_bs_in_cuda_graph = max_bs_in_cuda_graph
         with torch.device("cuda"):
             self.cuda_graph_batch_info = LoRABatchInfo(
-                bs=self.max_bs_in_cuda_graph,
-                seg_lens=torch.zeros(self.max_bs_in_cuda_graph, dtype=torch.int32),
-                seg_indptr=torch.zeros(
-                    self.max_bs_in_cuda_graph + 1, dtype=torch.int32
-                ),
+                bs=max_bs_in_cuda_graph,
+                use_cuda_graph=True,
+                num_segments=None,
+                seg_lens=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
+                seg_indptr=torch.zeros(max_bs_in_cuda_graph + 1, dtype=torch.int32),
                 max_len=1,
-                weight_indices=torch.zeros(
-                    self.max_bs_in_cuda_graph, dtype=torch.int32
-                ),
+                weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
+                permutation=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),
             )
 
-            # Initialize seg_lens and seg_indptr for CUDA graph as they remain constant
-            # across batches.
-            self.cuda_graph_batch_info.seg_lens[: self.max_bs_in_cuda_graph].fill_(1)
-            torch.cumsum(
-                self.cuda_graph_batch_info.seg_lens[: self.max_bs_in_cuda_graph],
-                dim=0,
-                out=self.cuda_graph_batch_info.seg_indptr[
-                    1 : self.max_bs_in_cuda_graph + 1
-                ],
-            )
+        self.lora_backend.init_cuda_graph_batch_info(
+            cuda_graph_batch_info=self.cuda_graph_batch_info,
+            max_bs_in_cuda_graph=max_bs_in_cuda_graph,
+        )
 
     def create_lora_update_result(
         self, success: bool, error_message: str = ""
-    ) -> LoRAUpdateResult:
-        return LoRAUpdateResult(
+    ) -> LoRAUpdateOutput:
+        return LoRAUpdateOutput(
             success=success,
             error_message=error_message,
             loaded_adapters={
@@ -118,7 +117,7 @@ class LoRAManager:
             },
         )
 
-    def load_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateResult:
+    def load_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateOutput:
         """
         Load a single LoRA adapter from the specified path.
 
@@ -175,7 +174,7 @@ class LoRAManager:
                 "`--max-loras-per-batch` or load it as unpinned LoRA adapters."
             )
 
-    def unload_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateResult:
+    def unload_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateOutput:
         """
         Unload LoRA adapters by their names. This will remove the adapters from the memory pool and
         delete the corresponding LoRA modules.
@@ -232,7 +231,6 @@ class LoRAManager:
         return required_slots <= mem_pool_vacancy
 
     def prepare_lora_batch(self, forward_batch: ForwardBatch):
-
         # Load active loras into lora memory pool
         cur_uids = set(forward_batch.lora_ids)
 
@@ -247,102 +245,30 @@ class LoRAManager:
         # set up batch info shared by all lora modules
         bs = forward_batch.batch_size
 
-        def transfer_adapter_info(
-            weight_indices_out: torch.Tensor,
-            lora_ranks_out: torch.Tensor,
-            scalings_out: torch.Tensor,
-        ):
-            """
-            Transfer adapter metadata (weight indices, LoRA rank, scalings) from host
-            to device (CUDA) asynchronously.
-            """
-            weight_indices = [0] * len(forward_batch.lora_ids)
-            lora_ranks = [0] * self.max_loras_per_batch
-            scalings = [0] * self.max_loras_per_batch
-            for i, uid in enumerate(forward_batch.lora_ids):
-                weight_indices[i] = self.memory_pool.get_buffer_id(uid)
-                if uid is not None:
-                    lora = self.loras[uid]
-                    lora_ranks[weight_indices[i]] = lora.config.r
-                    scalings[weight_indices[i]] = lora.scaling
-
-            # Use pinned memory to avoid synchronizations during host-to-device transfer
-            weight_indices_tensor = torch.tensor(
-                weight_indices, dtype=torch.int32, pin_memory=True, device="cpu"
-            )
-            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"
-            )
-
-            # Copy to device tensors asynchronously
-            weight_indices_out[:bs].copy_(weight_indices_tensor, non_blocking=True)
-            lora_ranks_out[: self.max_loras_per_batch].copy_(
-                lora_ranks_tensor, non_blocking=True
-            )
-            scalings_out[: self.max_loras_per_batch].copy_(
-                scalings_tensor, non_blocking=True
-            )
-
-        if (
+        use_cuda_graph = (
             hasattr(self, "max_bs_in_cuda_graph")
             and bs <= self.max_bs_in_cuda_graph
             and forward_batch.forward_mode.is_cuda_graph()
-        ):
-            # Do in-place updates when CUDA graph is enabled and the batch forward mode
-            # could use CUDA graph.
-
-            transfer_adapter_info(
-                self.cuda_graph_batch_info.weight_indices,
-                self.cuda_graph_batch_info.lora_ranks,
-                self.cuda_graph_batch_info.scalings,
-            )
-
-            self.cuda_graph_batch_info.bs = bs
-            self.cuda_graph_batch_info.max_len = 1
-            batch_info = self.cuda_graph_batch_info
-        else:
-            weight_indices = torch.empty((bs,), dtype=torch.int32, device=self.device)
-            lora_ranks = torch.zeros(
-                (self.max_loras_per_batch,), dtype=torch.int64, device=self.device
-            )
-            scalings = torch.zeros(
-                (self.max_loras_per_batch,), dtype=torch.float, device=self.device
-            )
-            transfer_adapter_info(
-                weight_indices,
-                lora_ranks,
-                scalings,
-            )
-
-            seg_lens = (
-                forward_batch.extend_seq_lens
-                if forward_batch.forward_mode.is_extend()
-                else torch.ones(bs, device=self.device)
-            )
-
-            max_len = (
-                # Calculate max_len from the CPU copy to avoid D2H transfer.
-                max(forward_batch.extend_seq_lens_cpu)
-                if forward_batch.forward_mode.is_extend()
-                else 1
-            )
+        )
 
-            seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
-            seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
-
-            batch_info = LoRABatchInfo(
-                bs=bs,
-                seg_lens=seg_lens,
-                seg_indptr=seg_indptr,
-                max_len=max_len,
-                weight_indices=weight_indices,
-                lora_ranks=lora_ranks,
-                scalings=scalings,
-            )
-        self.lora_backend.set_batch_info(batch_info)
+        weight_indices = [0] * len(forward_batch.lora_ids)
+        lora_ranks = [0] * self.max_loras_per_batch
+        scalings = [0] * self.max_loras_per_batch
+        for i, uid in enumerate(forward_batch.lora_ids):
+            weight_indices[i] = self.memory_pool.get_buffer_id(uid)
+            if uid is not None:
+                lora = self.loras[uid]
+                lora_ranks[weight_indices[i]] = lora.config.r
+                scalings[weight_indices[i]] = lora.scaling
+        # Do in-place updates when CUDA graph is enabled and the batch forward mode
+        # could use CUDA graph.
+        self.lora_backend.prepare_lora_batch(
+            forward_batch=forward_batch,
+            weight_indices=weight_indices,
+            lora_ranks=lora_ranks,
+            scalings=scalings,
+            batch_info=self.cuda_graph_batch_info if use_cuda_graph else None,
+        )
 
     def update_lora_info(self):
         """
@@ -492,6 +418,10 @@ class LoRAManager:
         replace_submodule(self.base_model, module_name, lora_module)
         return lora_module
 
+    def should_skip_lora_for_vision_model(self, module_name):
+        # TODO: support different vision models
+        return module_name.find("vision_model.model") != -1
+
     def init_lora_modules(self):
         # Look-up table that essentially maps (layer_index, module_name) to the corresponding LoRA module.
         self.lora_modules: List[Dict[str, BaseLayerWithLoRA]] = [
@@ -509,6 +439,10 @@ class LoRAManager:
             ) and not self.base_model.should_apply_lora(module_name):
                 continue
 
+            # Skip vision model
+            if self.should_skip_lora_for_vision_model(module_name):
+                continue
+
             # The module should be converted if it is included in target_names
             if module_name.split(".")[-1] in self.target_modules:
                 layer_id = get_layer_id(module_name)

sglang/srt/lora/mem_pool.py

@@ -4,7 +4,6 @@ from typing import Callable, Dict, Iterable, List, Optional, Set, Tuple, Union
 import torch
 
 from sglang.srt.distributed import divide
-from sglang.srt.hf_transformers_utils import AutoConfig
 from sglang.srt.lora.layers import BaseLayerWithLoRA
 from sglang.srt.lora.lora import LoRAAdapter
 from sglang.srt.lora.lora_config import LoRAConfig
@@ -17,6 +16,7 @@ from sglang.srt.lora.utils import (
     get_stacked_multiply,
     get_target_module_name,
 )
+from sglang.srt.utils.hf_transformers_utils import AutoConfig
 
 logger = logging.getLogger(__name__)
 
@@ -104,12 +104,18 @@ class LoRAMemoryPool:
             return all(_can_support(x) for x in config)
 
     def get_lora_A_shape(
-        self, module_name: str, base_model: torch.nn.Module, max_lora_dim: int
+        self,
+        module_name: str,
+        base_model: torch.nn.Module,
+        max_lora_dim: int,
+        layer_idx: int,
     ) -> Tuple[int]:
         """
         Given a module_name (might be a stacked name), return the hidden dims of modules' input and output.
         """
-        input_dim, _ = get_hidden_dim(module_name, self.base_hf_config, base_model)
+        input_dim, _ = get_hidden_dim(
+            module_name, self.base_hf_config, base_model, layer_idx
+        )
         c = get_stacked_multiply(module_name)
         if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
             input_dim = divide(input_dim, self.tp_size)
@@ -120,12 +126,18 @@ class LoRAMemoryPool:
         )
 
     def get_lora_B_shape(
-        self, module_name: str, base_model: torch.nn.Module, max_lora_dim: int
+        self,
+        module_name: str,
+        base_model: torch.nn.Module,
+        max_lora_dim: int,
+        layer_idx: int,
     ) -> Tuple[int]:
         """
         Given a module_name (might be a stacked name), return the hidden dims of modules' input and output.
         """
-        _, output_dim = get_hidden_dim(module_name, self.base_hf_config, base_model)
+        _, output_dim = get_hidden_dim(
+            module_name, self.base_hf_config, base_model, layer_idx
+        )
         if self.tp_size > 1 and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES:
             output_dim = divide(output_dim, self.tp_size)
         return (
@@ -140,19 +152,21 @@ class LoRAMemoryPool:
         def init_buffer(
             buffer: Dict[str, List[torch.Tensor]],
             target_modules: Set[str],
-            get_lora_shape_fn: Callable[[str, torch.nn.Module, int], Tuple[int]],
+            get_lora_shape_fn: Callable[[str, torch.nn.Module, int, int], Tuple[int]],
         ):
             for module_name in target_modules:
-                lora_shape = get_lora_shape_fn(
-                    module_name, base_model, self.max_lora_rank
-                )
                 buffer[module_name] = [
                     torch.empty(
-                        lora_shape,
+                        get_lora_shape_fn(
+                            module_name,
+                            base_model,
+                            self.max_lora_rank,
+                            idx,
+                        ),
                         dtype=self.dtype,
                         device=device,
                     )
-                    for _ in range(self.num_layer)
+                    for idx in range(self.num_layer)
                 ]
 
         init_buffer(

sglang/srt/lora/triton_ops/__init__.py

@@ -1,3 +1,5 @@
+from .chunked_sgmv_expand import chunked_sgmv_lora_expand_forward
+from .chunked_sgmv_shrink import chunked_sgmv_lora_shrink_forward
 from .gate_up_lora_b import gate_up_lora_b_fwd
 from .qkv_lora_b import qkv_lora_b_fwd
 from .sgemm_lora_a import sgemm_lora_a_fwd
@@ -8,4 +10,6 @@ __all__ = [
     "qkv_lora_b_fwd",
     "sgemm_lora_a_fwd",
     "sgemm_lora_b_fwd",
+    "chunked_sgmv_lora_shrink_forward",
+    "chunked_sgmv_lora_expand_forward",
 ]

sglang/srt/lora/triton_ops/chunked_sgmv_expand.py

@@ -0,0 +1,214 @@
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.lora.utils import LoRABatchInfo
+from sglang.srt.utils import cached_triton_kernel
+
+
+@cached_triton_kernel(lambda _, kwargs: (kwargs["NUM_SLICES"], kwargs["BLOCK_M"]))
+@triton.jit
+def _chunked_lora_expand_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Information on sequence lengths and weight id
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    permutation,
+    num_segs,
+    # For fused output scaling
+    scalings,
+    # Offsets of q/k/v slice on output dimension
+    slice_offsets,
+    # Meta parameters
+    NUM_SLICES: tl.constexpr,
+    OUTPUT_DIM: tl.constexpr,
+    MAX_RANK: tl.constexpr,  # K = R
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    """
+    Computes a chunked SGMV for LoRA expand operations.
+
+    When a sequence's rank is 0, the kernel is essentially a no-op, following
+    the convention in pytorch where the product of two matrices of shape (m, 0)
+    and (0, n) is an all-zero matrix of shape (m, n).
+
+    Args:
+        x (Tensor): The input tensor, which is the result of the LoRA A projection.
+            Shape: (s, num_slices * K), where s is the sum of all sequence lengths in the
+            batch and K is the maximum LoRA rank.
+        weights (Tensor): The LoRA B weights for all adapters.
+            Shape: (num_lora, output_dim, K).
+        output (Tensor): The output tensor where the result is stored.
+            Shape: (s, output_dim).
+    """
+    tl.static_assert(NUM_SLICES <= 3)
+
+    x_stride_0: tl.constexpr = NUM_SLICES * MAX_RANK
+    x_stride_1: tl.constexpr = 1
+
+    w_stride_0: tl.constexpr = OUTPUT_DIM * MAX_RANK
+    w_stride_1: tl.constexpr = MAX_RANK
+    w_stride_2: tl.constexpr = 1
+
+    output_stride_0: tl.constexpr = OUTPUT_DIM
+    output_stride_1: tl.constexpr = 1
+
+    pid_s = tl.program_id(axis=2)
+    if pid_s >= num_segs:
+        return
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len.
+    # qkv_id decides which of q,k,v to compute (0: q, 1: k, 2: v)
+    w_index = tl.load(weight_indices + pid_s)
+    cur_rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + pid_s)
+    seg_end = tl.load(seg_indptr + pid_s + 1)
+
+    slice_id = tl.program_id(axis=1)
+    slice_start = tl.load(slice_offsets + slice_id)
+    slice_end = tl.load(slice_offsets + slice_id + 1)
+
+    scaling = tl.load(scalings + w_index)
+    # Adjust K (rank) according to the specific LoRA adapter
+    cur_rank = tl.minimum(MAX_RANK, cur_rank)
+
+    # Map logical sequence index to physical index
+    s_offset_logical = tl.arange(0, BLOCK_M) + seg_start
+    s_offset_physical = tl.load(
+        permutation + s_offset_logical, mask=s_offset_logical < seg_end
+    )
+
+    # Create pointers for the first block of x and weights[batch_id][n_start: n_end][:]
+    # The pointers will be advanced as we move in the K direction
+    # and accumulate
+    pid_n = tl.program_id(axis=0)
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + slice_start
+    k_offset = tl.arange(0, BLOCK_K)
+
+    x_ptrs = (
+        x
+        + slice_id * cur_rank * x_stride_1
+        + (s_offset_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1)
+    )
+    w_ptrs = (weights + w_index * w_stride_0) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(cur_rank, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset_logical[:, None] < seg_end)
+            & (k_offset[None, :] < cur_rank - k * BLOCK_K),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < cur_rank - k * BLOCK_K)
+            & (n_offset[None, :] < slice_end),
+            other=0.0,
+        )
+        partial_sum += tl.dot(x_tile, w_tile)
+
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+
+    # Store result to output matrix
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    output_ptr = output + (
+        s_offset_physical[:, None] * output_stride_0
+        + n_offset[None, :] * output_stride_1
+    )
+    output_mask = (s_offset_logical[:, None] < seg_end) & (
+        n_offset[None, :] < slice_end
+    )
+    partial_sum += tl.load(output_ptr, mask=output_mask, other=0.0)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def chunked_sgmv_lora_expand_forward(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    slice_offsets: torch.Tensor,
+    max_slice_size: int,
+    base_output: Optional[torch.Tensor],
+) -> torch.Tensor:
+
+    # x: (s, slice_num * r)
+    # weights: (num_lora, output_dim, r)
+    # slice_offsets: boundaries for different slices in the output dimension
+    # output: (s, output_dim)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # For each slice i, accumulates:
+    # lora_output[:, slice_offsets[i]:slice_offsets[i+1]] += scaling * sgemm(x[:, i*cur_rank:(i+1)*cur_rank], weights[:, slice_offsets[i]:slice_offsets[i+1], :])
+
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(x.shape) == 2
+    assert len(weights.shape) == 3
+
+    # Get dims
+    M = x.shape[0]
+    input_dim = x.shape[1]
+    OUTPUT_DIM = weights.shape[1]
+    MAX_RANK = weights.shape[2]
+    num_slices = len(slice_offsets) - 1
+    assert input_dim == num_slices * MAX_RANK
+
+    # TODO (lifuhuang): fine-tune per operation
+    BLOCK_M = batch_info.max_len
+    BLOCK_K = 16
+    BLOCK_N = 64
+
+    num_segments = batch_info.num_segments
+
+    grid = (
+        triton.cdiv(max_slice_size, BLOCK_N),
+        num_slices,  # number of slices in the input/output
+        batch_info.bs if batch_info.use_cuda_graph else num_segments,
+    )
+
+    if base_output is None:
+        output = torch.zeros((M, OUTPUT_DIM), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    _chunked_lora_expand_kernel[grid](
+        x=x,
+        weights=weights,
+        output=output,
+        seg_indptr=batch_info.seg_indptr,
+        weight_indices=batch_info.weight_indices,
+        lora_ranks=batch_info.lora_ranks,
+        permutation=batch_info.permutation,
+        num_segs=num_segments,
+        scalings=batch_info.scalings,
+        slice_offsets=slice_offsets,
+        # constants
+        NUM_SLICES=num_slices,
+        OUTPUT_DIM=OUTPUT_DIM,
+        MAX_RANK=MAX_RANK,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+    )
+
+    return output