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

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

sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py

@@ -0,0 +1,174 @@
+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_shrink_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Information on sequence lengths,ranks and weight id
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    permutation,
+    num_segs,
+    # Meta parameters
+    N: tl.constexpr,  # num_slices * r
+    K: tl.constexpr,  # input_dim
+    NUM_SLICES: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    """
+    Computes a chunked SGMV for LoRA shrink operations.
+
+    The kernel ensures that output[seg_start:seg_start + seg_len, :rank * num_slices]
+    stores the product of the input `x` and the LoRA weights for the corresponding
+    sequence. This implies that when rank is 0, the kernel is essentially a no-op,
+    as output[seg_start:seg_start + seg_len, :0] is trivially correct (empty).
+
+    Args:
+        x (torch.Tensor): The input activations tensor of shape `(s, K)`, where `s`
+            is the sum of all sequence lengths in the batch.
+        weights (torch.Tensor): The LoRA A weights for all available adapters,
+            with shape `(num_lora, N, K)` where N = num_slices * r.
+        output (torch.Tensor): The output tensor of shape `(s, N)`.
+    """
+    x_stride_1: tl.constexpr = 1
+    x_stride_0: tl.constexpr = K
+
+    w_stride_0: tl.constexpr = N * K
+    w_stride_1: tl.constexpr = K
+    w_stride_2: tl.constexpr = 1
+
+    output_stride_0: tl.constexpr = N
+    output_stride_1: tl.constexpr = 1
+
+    pid_s = tl.program_id(1)
+    if pid_s >= num_segs:
+        return
+
+    pid_n = tl.program_id(0)
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len
+    w_index = tl.load(weight_indices + pid_s)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel becomes a no-op as the output is always trivially correct.
+    if rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + pid_s)
+    seg_end = tl.load(seg_indptr + pid_s + 1)
+
+    # Adjust N dim according to the specific LoRA adapter
+    cur_n = tl.minimum(N, rank * NUM_SLICES)
+
+    # 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
+    )
+
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+    x_ptrs = x + (
+        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(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset_logical[:, None] < seg_end)
+            & (k_offset[None, :] < K - k * BLOCK_K),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < K - k * BLOCK_K) & (n_offset[None, :] < cur_n),
+            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 = 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, :] < cur_n)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def chunked_sgmv_lora_shrink_forward(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    num_slices: int,
+) -> torch.Tensor:
+    # x: (s, input_dim)
+    # weights: (num_lora, num_slices * r, input_dim)
+    # output: (s, num_slices * r)
+    # num_slices: qkv=3, gate_up=2, others=1
+    # when called with multiple slices, the weights.shape[-2] will be num_slices * r
+    # input_dim is much larger than r
+
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(x.shape) == 2
+    assert len(weights.shape) == 3
+
+    # Block shapes
+    # TODO (lifuhuang): experiment with split-k
+    BLOCK_M = batch_info.max_len
+    BLOCK_N = 16
+    BLOCK_K = 256
+
+    S = x.shape[0]
+    N = weights.shape[1]
+    K = weights.shape[2]
+    assert x.shape[-1] == K
+
+    num_segments = batch_info.num_segments
+    grid = (
+        triton.cdiv(N, BLOCK_N),
+        batch_info.bs if batch_info.use_cuda_graph else num_segments,
+    )
+
+    output = torch.empty((S, N), device=x.device, dtype=x.dtype)
+    _chunked_lora_shrink_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,
+        # constants
+        N=N,
+        K=K,
+        NUM_SLICES=num_slices,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+    )
+
+    return output

sglang/srt/lora/utils.py

@@ -5,7 +5,7 @@ from typing import Iterable, Optional, Set, Tuple
 
 import torch
 
-from sglang.srt.hf_transformers_utils import AutoConfig
+from sglang.srt.utils.hf_transformers_utils import AutoConfig
 
 
 @dataclass
@@ -19,6 +19,9 @@ class LoRABatchInfo:
     # Number of segments. For triton backend, it is equal to batch size.
     num_segments: int
 
+    # Maximum segment length of current batch
+    max_len: int
+
     # Indice pointers of each segment in shape (num_segments + 1, )
     seg_indptr: torch.Tensor
 
@@ -34,9 +37,6 @@ class LoRABatchInfo:
     # Lengths of each segments in shape (num_segments,)
     seg_lens: Optional[torch.Tensor]
 
-    # Maximum segment length of current batch
-    max_len: Optional[int]
-
     # The logical (re)ordering of input rows (tokens), in shape (num_tokens,)
     permutation: Optional[torch.Tensor]
 
@@ -98,6 +98,7 @@ def get_normalized_target_modules(
 ) -> set[str]:
     """
     Mapping a list of target module name to names of the normalized LoRA weights.
+    Handles both base module names (e.g., "gate_proj") and prefixed module names (e.g., "feed_forward.gate_proj").
     """
     params_mapping = {
         "q_proj": "qkv_proj",
@@ -109,7 +110,8 @@ def get_normalized_target_modules(
 
     result = set()
     for name in target_modules:
-        normalized_name = params_mapping.get(name, name)
+        base_name = name.split(".")[-1]
+        normalized_name = params_mapping.get(base_name, base_name)
         result.add(normalized_name)
     return result