diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index de15d4e16..4758e640f 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/python/sglang/srt/lora/backend/triton_backend.py @@ -48,6 +48,12 @@ class TritonLoRABackend(BaseLoRABackend): extra_embeddings=extra_embeddings, ) + def _sgemm_info(self, pruned_batch_info=None): + """Return the sgemm batch_info (merged segments when available).""" + if pruned_batch_info is not None: + return pruned_batch_info + return getattr(self, "sgemm_batch_info", None) or self.batch_info + def run_lora_a_sgemm( self, x: torch.Tensor, @@ -57,10 +63,9 @@ class TritonLoRABackend(BaseLoRABackend): *args, **kwargs, ) -> torch.Tensor: - batch_info = ( - pruned_batch_info if pruned_batch_info is not None else self.batch_info + return sgemm_lora_a_fwd( + x, weights, self._sgemm_info(pruned_batch_info), stack_num=stack_num ) - return sgemm_lora_a_fwd(x, weights, batch_info, stack_num=stack_num) def run_lora_b_sgemm( self, @@ -71,10 +76,9 @@ class TritonLoRABackend(BaseLoRABackend): *args, **kwargs, ) -> torch.Tensor: - batch_info = ( - pruned_batch_info if pruned_batch_info is not None else self.batch_info + return sgemm_lora_b_fwd( + x, weights, self._sgemm_info(pruned_batch_info), base_output ) - return sgemm_lora_b_fwd(x, weights, batch_info, base_output) def run_qkv_lora( self, @@ -93,11 +97,12 @@ class TritonLoRABackend(BaseLoRABackend): # qkv_lora_b: (num_lora, output_dim_q + 2 * output_dim_kv, r) assert isinstance(qkv_lora_b, torch.Tensor) - lora_a_output = sgemm_lora_a_fwd(x, qkv_lora_a, self.batch_info, stack_num=3) + sgemm_info = self._sgemm_info() + lora_a_output = sgemm_lora_a_fwd(x, qkv_lora_a, sgemm_info, stack_num=3) lora_output = qkv_lora_b_fwd( lora_a_output, qkv_lora_b, - self.batch_info, + sgemm_info, output_offset, max_qkv_out_dim, base_output, @@ -120,14 +125,13 @@ class TritonLoRABackend(BaseLoRABackend): assert isinstance(gate_up_lora_b, torch.Tensor) output_dim = gate_up_lora_b.shape[-2] // 2 + sgemm_info = self._sgemm_info() # lora_a_output: (s, 2 * r) - lora_a_output = sgemm_lora_a_fwd( - x, gate_up_lora_a, self.batch_info, stack_num=2 - ) + lora_a_output = sgemm_lora_a_fwd(x, gate_up_lora_a, sgemm_info, stack_num=2) lora_output = gate_up_lora_b_fwd( lora_a_output, gate_up_lora_b, - self.batch_info, + sgemm_info, output_dim, base_output, ) @@ -138,6 +142,8 @@ class TritonLoRABackend(BaseLoRABackend): max_bs_in_cuda_graph: int, num_tokens_per_bs: int, ): + max_tokens = max_bs_in_cuda_graph * num_tokens_per_bs + mlpb = self.max_loras_per_batch with torch.device("cuda"): self.cuda_graph_batch_info = LoRABatchInfo( bs=max_bs_in_cuda_graph, @@ -149,19 +155,75 @@ class TritonLoRABackend(BaseLoRABackend): seg_indptr=torch.zeros(max_bs_in_cuda_graph + 1, dtype=torch.int32), max_len=num_tokens_per_bs, 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), + lora_ranks=torch.zeros(mlpb, dtype=torch.int32), + scalings=torch.zeros(mlpb, dtype=torch.float), permutation=None, ) - # 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], ) + # Sgemm batch_info with segments merged by adapter. + # Updated each batch by compute_sgemm_routing(). + self.cuda_graph_sgemm_batch_info = LoRABatchInfo( + bs=mlpb, + use_cuda_graph=True, + num_segments=mlpb, + seg_lens=torch.zeros(mlpb, dtype=torch.int32), + seg_indptr=torch.zeros(mlpb + 1, dtype=torch.int32), + max_len=max_tokens, + weight_indices=torch.arange(mlpb, dtype=torch.int32), + lora_ranks=torch.zeros(mlpb, dtype=torch.int32), + scalings=torch.zeros(mlpb, dtype=torch.float), + permutation=torch.zeros(max_tokens, dtype=torch.int32), + ) + + def compute_sgemm_routing(self, use_cuda_graph: bool): + """Sort tokens by adapter and build merged segments for sgemm LoRA.""" + bi = self.batch_info + bs = bi.bs + mlpb = self.max_loras_per_batch + wi = bi.weight_indices[:bs] + + perm = torch.argsort(wi, stable=True).to(torch.int32) + sorted_wi = wi[perm] + adapter_ids = torch.arange(mlpb, device=wi.device, dtype=torch.int32) + seg_starts = torch.searchsorted(sorted_wi, adapter_ids) + seg_ends = torch.searchsorted(sorted_wi, adapter_ids, right=True) + seg_lens = seg_ends - seg_starts + + if use_cuda_graph: + sgemm = getattr(self, "cuda_graph_sgemm_batch_info", None) + if sgemm is None: + return + sgemm.permutation[:bs] = perm + sgemm.seg_lens[:] = seg_lens + sgemm.seg_indptr[0] = 0 + torch.cumsum(sgemm.seg_lens, dim=0, out=sgemm.seg_indptr[1:]) + sgemm.max_len = bs + sgemm.lora_ranks[:mlpb] = bi.lora_ranks[:mlpb] + sgemm.scalings[:mlpb] = bi.scalings[:mlpb] + else: + seg_indptr = torch.zeros(mlpb + 1, dtype=torch.int32, device=wi.device) + seg_indptr[1:] = torch.cumsum(seg_lens, dim=0) + sgemm = LoRABatchInfo( + bs=mlpb, + use_cuda_graph=False, + num_segments=mlpb, + seg_lens=seg_lens, + seg_indptr=seg_indptr, + max_len=bs, + weight_indices=adapter_ids, + lora_ranks=bi.lora_ranks[:mlpb].clone(), + scalings=bi.scalings[:mlpb].clone(), + permutation=perm, + ) + + self.sgemm_batch_info = sgemm + def prepare_lora_batch( self, forward_batch: ForwardBatch, @@ -234,6 +296,14 @@ class TritonLoRABackend(BaseLoRABackend): batch_info.weight_indices[:bs].copy_(weight_indices_tensor, non_blocking=True) self.batch_info = batch_info + + # Biggest win is in decode. + is_decode = not forward_batch.forward_mode.is_extend() + if is_decode: + self.compute_sgemm_routing(use_cuda_graph) + else: + self.sgemm_batch_info = None + self.lm_head_batch_info, self.lm_head_pass_batch_infos = ( self._prepare_lm_head_batch_info(forward_batch, weight_indices, batch_info) ) diff --git a/python/sglang/srt/lora/triton_ops/gate_up_lora_b.py b/python/sglang/srt/lora/triton_ops/gate_up_lora_b.py index fc4574dd3..16ade8b44 100644 --- a/python/sglang/srt/lora/triton_ops/gate_up_lora_b.py +++ b/python/sglang/srt/lora/triton_ops/gate_up_lora_b.py @@ -2,6 +2,7 @@ import torch import triton import triton.language as tl +from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions from sglang.srt.lora.utils import LoRABatchInfo @@ -27,7 +28,9 @@ def _gate_up_lora_b_kernel( seg_indptr, weight_indices, lora_ranks, + sorted_token_ids, # Meta parameters + SORTED_BY_ADAPTER: tl.constexpr, BLOCK_S: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, @@ -67,6 +70,8 @@ def _gate_up_lora_b_kernel( gate_up_id = tl.program_id(axis=1) pid = tl.program_id(axis=0) seg_len = tl.load(seg_lens + batch_id) + if seg_len == 0: + return seg_start = tl.load(seg_indptr + batch_id) n_start = gate_up_id * output_dim # offset on output dim scaling = tl.load(scalings + w_index) @@ -78,6 +83,8 @@ def _gate_up_lora_b_kernel( num_pid_n = tl.cdiv(output_dim, BLOCK_N) pid_s = pid // num_pid_n pid_n = pid % num_pid_n + if pid_s * BLOCK_S >= seg_len: + return # Create pointers for the first block of x and weights # The pointers will be advanced as we move in the K direction @@ -86,8 +93,13 @@ def _gate_up_lora_b_kernel( n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N k_offset = tl.arange(0, BLOCK_K) - x_ptrs = (x + seg_start * x_stride_0 + (gate_up_id * K) * x_stride_1) + ( - s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + s_physical = _resolve_token_positions( + sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER + ) + x_ptrs = ( + x + + (gate_up_id * K) * x_stride_1 + + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1) ) w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + ( k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 @@ -115,8 +127,10 @@ def _gate_up_lora_b_kernel( # Store result to output matrix partial_sum *= scaling partial_sum = partial_sum.to(x.dtype.element_ty) - output_ptr = (output + seg_start * output_stride_0 + n_start * output_stride_1) + ( - s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + output_ptr = ( + output + + n_start * output_stride_1 + + (s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1) ) output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < output_dim) partial_sum += tl.load(output_ptr, mask=output_mask) @@ -161,6 +175,7 @@ def gate_up_lora_b_fwd( else: output = base_output + sorted_by_adapter = batch_info.permutation is not None _gate_up_lora_b_kernel[grid_b]( x, gate_up_lora_b, @@ -178,6 +193,8 @@ def gate_up_lora_b_fwd( batch_info.seg_indptr, batch_info.weight_indices, batch_info.lora_ranks, + batch_info.permutation, + sorted_by_adapter, BLOCK_S, BLOCK_OUT, BLOCK_R, diff --git a/python/sglang/srt/lora/triton_ops/kernel_utils.py b/python/sglang/srt/lora/triton_ops/kernel_utils.py new file mode 100644 index 000000000..788a7c305 --- /dev/null +++ b/python/sglang/srt/lora/triton_ops/kernel_utils.py @@ -0,0 +1,19 @@ +import triton +import triton.language as tl + + +@triton.jit +def _resolve_token_positions( + sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER: tl.constexpr +): + """Map logical segment offsets to physical token positions. + + When SORTED_BY_ADAPTER is True, segments are grouped by adapter and + sorted_token_ids provides the indirection to the original token rows. + When False, tokens are already contiguous starting at seg_start. + """ + if SORTED_BY_ADAPTER: + return tl.load( + sorted_token_ids + seg_start + s_offset, mask=s_offset < seg_len + ).to(tl.int64) + return (seg_start + s_offset).to(tl.int64) diff --git a/python/sglang/srt/lora/triton_ops/qkv_lora_b.py b/python/sglang/srt/lora/triton_ops/qkv_lora_b.py index 1d6663dbe..08ecc40ed 100644 --- a/python/sglang/srt/lora/triton_ops/qkv_lora_b.py +++ b/python/sglang/srt/lora/triton_ops/qkv_lora_b.py @@ -2,6 +2,7 @@ import torch import triton import triton.language as tl +from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions from sglang.srt.lora.utils import LoRABatchInfo @@ -29,7 +30,9 @@ def _qkv_lora_b_kernel( lora_ranks, # Offsets of q/k/v slice on output dimension n_offs, + sorted_token_ids, # Meta parameters + SORTED_BY_ADAPTER: tl.constexpr, BLOCK_S: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, @@ -69,6 +72,8 @@ def _qkv_lora_b_kernel( qkv_id = tl.program_id(axis=1) pid = tl.program_id(axis=0) seg_len = tl.load(seg_lens + batch_id) + if seg_len == 0: + return seg_start = tl.load(seg_indptr + batch_id) n_start = tl.load(n_offs + qkv_id) n_size = tl.load(n_offs + qkv_id + 1) - n_start @@ -80,6 +85,8 @@ def _qkv_lora_b_kernel( num_pid_n = tl.cdiv(max_qkv_out_dim, BLOCK_N) pid_s = pid // num_pid_n pid_n = pid % num_pid_n + if pid_s * BLOCK_S >= seg_len: + return # 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 @@ -88,8 +95,13 @@ def _qkv_lora_b_kernel( n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N k_offset = tl.arange(0, BLOCK_K) - x_ptrs = (x + seg_start * x_stride_0 + (qkv_id * K) * x_stride_1) + ( - s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + s_physical = _resolve_token_positions( + sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER + ) + x_ptrs = ( + x + + (qkv_id * K) * x_stride_1 + + (s_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1) ) w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + ( k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 @@ -116,8 +128,10 @@ def _qkv_lora_b_kernel( # Store result to output matrix partial_sum *= scaling partial_sum = partial_sum.to(x.dtype.element_ty) - output_ptr = (output + seg_start * output_stride_0 + n_start * output_stride_1) + ( - s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + output_ptr = ( + output + + n_start * output_stride_1 + + (s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1) ) output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < n_size) partial_sum += tl.load(output_ptr, mask=output_mask) @@ -171,6 +185,7 @@ def qkv_lora_b_fwd( else: output = base_output + sorted_by_adapter = batch_info.permutation is not None _qkv_lora_b_kernel[grid_b]( x, qkv_lora_b, @@ -189,6 +204,8 @@ def qkv_lora_b_fwd( batch_info.weight_indices, batch_info.lora_ranks, output_offset, + batch_info.permutation, + sorted_by_adapter, BLOCK_S, BLOCK_OUT, BLOCK_R, diff --git a/python/sglang/srt/lora/triton_ops/sgemm_lora_a.py b/python/sglang/srt/lora/triton_ops/sgemm_lora_a.py index dded64bcf..0dd3e5bbb 100644 --- a/python/sglang/srt/lora/triton_ops/sgemm_lora_a.py +++ b/python/sglang/srt/lora/triton_ops/sgemm_lora_a.py @@ -2,6 +2,7 @@ import torch import triton import triton.language as tl +from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions from sglang.srt.lora.utils import LoRABatchInfo @@ -28,7 +29,9 @@ def _sgemm_lora_a_kernel( seg_indptr, weight_indices, lora_ranks, + sorted_token_ids, # Meta parameters + SORTED_BY_ADAPTER: tl.constexpr, BLOCK_S: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, @@ -62,6 +65,8 @@ def _sgemm_lora_a_kernel( pid = tl.program_id(axis=0) seg_start = tl.load(seg_indptr + batch_id) seg_len = tl.load(seg_lens + batch_id) + if seg_len == 0: + return # Adjust N (stack_num * max_rank) according to the specific LoRA adapter N = tl.minimum(N, rank * stack_num) @@ -70,6 +75,8 @@ def _sgemm_lora_a_kernel( num_pid_n = tl.cdiv(N, BLOCK_N) pid_s = pid // num_pid_n pid_n = pid % num_pid_n + if pid_s * BLOCK_S >= seg_len: + return # Create pointers for the first block of x and weights[batch_id] # The pointers will be advanced as we move in the K direction @@ -77,9 +84,10 @@ def _sgemm_lora_a_kernel( s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N k_offset = tl.arange(0, BLOCK_K) - x_ptrs = (x + seg_start * x_stride_0) + ( - s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + s_physical = _resolve_token_positions( + sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER ) + x_ptrs = x + (s_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 ) @@ -104,10 +112,10 @@ def _sgemm_lora_a_kernel( # Store result to output matrix partial_sum = partial_sum.to(x.dtype.element_ty) - output_ptr = (output + seg_start * output_stride_0) + ( - s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 - ) output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < N) + output_ptr = output + ( + s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + ) tl.store(output_ptr, partial_sum, mask=output_mask) @@ -144,6 +152,8 @@ def sgemm_lora_a_fwd( batch_info.bs, ) + sorted_by_adapter = batch_info.permutation is not None + output = torch.empty((S, R), device=x.device, dtype=x.dtype) _sgemm_lora_a_kernel[grid]( x, @@ -163,6 +173,8 @@ def sgemm_lora_a_fwd( batch_info.seg_indptr, batch_info.weight_indices, batch_info.lora_ranks, + batch_info.permutation, + sorted_by_adapter, BLOCK_S, BLOCK_R, BLOCK_K, diff --git a/python/sglang/srt/lora/triton_ops/sgemm_lora_b.py b/python/sglang/srt/lora/triton_ops/sgemm_lora_b.py index b796cdd0e..fc7f844e2 100644 --- a/python/sglang/srt/lora/triton_ops/sgemm_lora_b.py +++ b/python/sglang/srt/lora/triton_ops/sgemm_lora_b.py @@ -2,6 +2,7 @@ import torch import triton import triton.language as tl +from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions from sglang.srt.lora.utils import LoRABatchInfo @@ -27,7 +28,9 @@ def _sgemm_lora_b_kernel( seg_indptr, weight_indices, lora_ranks, + sorted_token_ids, # Meta parameters + SORTED_BY_ADAPTER: tl.constexpr, BLOCK_S: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, @@ -63,6 +66,8 @@ def _sgemm_lora_b_kernel( pid = tl.program_id(axis=0) seg_len = tl.load(seg_lens + batch_id) + if seg_len == 0: + return seg_start = tl.load(seg_indptr + batch_id) scaling = tl.load(scalings + w_index) # Adjust K (rank) according to the specific LoRA adapter @@ -72,6 +77,8 @@ def _sgemm_lora_b_kernel( num_pid_n = tl.cdiv(N, BLOCK_N) pid_s = pid // num_pid_n pid_n = pid % num_pid_n + if pid_s * BLOCK_S >= seg_len: + return # Create pointers for the first block of x and weights[batch_id] # The pointers will be advanced as we move in the K direction @@ -79,9 +86,10 @@ def _sgemm_lora_b_kernel( s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N k_offset = tl.arange(0, BLOCK_K) - x_ptrs = (x + seg_start * x_stride_0) + ( - s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + s_physical = _resolve_token_positions( + sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER ) + x_ptrs = x + (s_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 ) @@ -108,8 +116,8 @@ def _sgemm_lora_b_kernel( # Store result to output matrix partial_sum *= scaling partial_sum = partial_sum.to(x.dtype.element_ty) - output_ptr = (output + seg_start * output_stride_0) + ( - s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + output_ptr = output + ( + s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 ) output_mask = (s_offset[:, None] < seg_len) & n_mask partial_sum += tl.load(output_ptr, mask=output_mask, other=0.0) @@ -152,6 +160,7 @@ def sgemm_lora_b_fwd( else: output = base_output + sorted_by_adapter = batch_info.permutation is not None _sgemm_lora_b_kernel[grid]( x, weights, @@ -169,6 +178,8 @@ def sgemm_lora_b_fwd( batch_info.seg_indptr, batch_info.weight_indices, batch_info.lora_ranks, + batch_info.permutation, + sorted_by_adapter, BLOCK_S, BLOCK_N, BLOCK_R, diff --git a/test/registered/lora/test_sgemm_sorted_by_adapter.py b/test/registered/lora/test_sgemm_sorted_by_adapter.py new file mode 100644 index 000000000..a8c787267 --- /dev/null +++ b/test/registered/lora/test_sgemm_sorted_by_adapter.py @@ -0,0 +1,236 @@ +"""Test that sgemm kernels produce identical results with and without SORTED_BY_ADAPTER.""" + +from typing import Any + +import pytest +import torch + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=15, suite="stage-b-test-1-gpu-large") + + +def _make_batch_info( + bs: int, + weight_indices: list[int], + lora_ranks: list[int], + scalings: list[float], + device: str = "cuda", +) -> Any: + """Build a per-sequence LoRABatchInfo (no permutation).""" + from sglang.srt.lora.utils import LoRABatchInfo + + seg_lens = torch.ones(bs, dtype=torch.int32, device=device) + seg_indptr = torch.zeros(bs + 1, dtype=torch.int32, device=device) + seg_indptr[1:] = torch.cumsum(seg_lens, dim=0) + return LoRABatchInfo( + bs=bs, + use_cuda_graph=False, + num_segments=bs, + seg_lens=seg_lens, + seg_indptr=seg_indptr, + max_len=1, + weight_indices=torch.tensor(weight_indices, dtype=torch.int32, device=device), + lora_ranks=torch.tensor(lora_ranks, dtype=torch.int32, device=device), + scalings=torch.tensor(scalings, dtype=torch.float, device=device), + permutation=None, + ) + + +def _make_sorted_batch_info( + weight_indices: list[int], + lora_ranks: list[int], + scalings: list[float], + max_loras: int, + device: str = "cuda", +) -> Any: + from sglang.srt.lora.utils import LoRABatchInfo + + """Build a merged-by-adapter LoRABatchInfo (with permutation).""" + wi = torch.tensor(weight_indices, dtype=torch.int32, device=device) + bs = wi.shape[0] + + perm = torch.argsort(wi, stable=True).to(torch.int32) + sorted_wi = wi[perm] + adapter_ids = torch.arange(max_loras, device=device, dtype=torch.int32) + seg_starts = torch.searchsorted(sorted_wi, adapter_ids) + seg_ends = torch.searchsorted(sorted_wi, adapter_ids, right=True) + seg_lens = seg_ends - seg_starts + + seg_indptr = torch.zeros(max_loras + 1, dtype=torch.int32, device=device) + seg_indptr[1:] = torch.cumsum(seg_lens, dim=0) + + return LoRABatchInfo( + bs=max_loras, + use_cuda_graph=False, + num_segments=max_loras, + seg_lens=seg_lens, + seg_indptr=seg_indptr, + max_len=bs, + weight_indices=adapter_ids, + lora_ranks=torch.tensor(lora_ranks, dtype=torch.int32, device=device), + scalings=torch.tensor(scalings, dtype=torch.float, device=device), + permutation=perm, + ) + + +def _check_close( + a: torch.Tensor, b: torch.Tensor, name: str, atol: float = 1e-4, rtol: float = 1e-3 +) -> None: + diff = (a - b).abs().max().item() + assert torch.allclose(a, b, atol=atol, rtol=rtol), f"{name}: max diff = {diff}" + + +def test_sgemm_lora_a(): + from sglang.srt.lora.triton_ops import sgemm_lora_a_fwd + + torch.manual_seed(42) + bs, input_dim, rank, num_loras = 8, 256, 16, 3 + x = torch.randn(bs, input_dim, device="cuda", dtype=torch.bfloat16) + weights = torch.randn( + num_loras, rank, input_dim, device="cuda", dtype=torch.bfloat16 + ) + wi = [i % num_loras for i in range(bs)] + lora_ranks = [rank] * num_loras + scalings = [1.0] * num_loras + + bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) + bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) + + out_plain = sgemm_lora_a_fwd(x, weights, bi_plain) + out_sorted = sgemm_lora_a_fwd(x, weights, bi_sorted) + _check_close(out_plain, out_sorted, "sgemm_lora_a") + + +def test_sgemm_lora_b(): + from sglang.srt.lora.triton_ops import sgemm_lora_b_fwd + + torch.manual_seed(42) + bs, output_dim, rank, num_loras = 8, 256, 16, 3 + x = torch.randn(bs, rank, device="cuda", dtype=torch.bfloat16) + weights = torch.randn( + num_loras, output_dim, rank, device="cuda", dtype=torch.bfloat16 + ) + wi = [i % num_loras for i in range(bs)] + lora_ranks = [rank] * num_loras + scalings = [0.5] * num_loras + + bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) + bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) + + base_plain = torch.randn(bs, output_dim, device="cuda", dtype=torch.bfloat16) + base_sorted = base_plain.clone() + + out_plain = sgemm_lora_b_fwd(x, weights, bi_plain, base_plain) + out_sorted = sgemm_lora_b_fwd(x, weights, bi_sorted, base_sorted) + _check_close(out_plain, out_sorted, "sgemm_lora_b") + + +def test_qkv_lora_b(): + from sglang.srt.lora.triton_ops import qkv_lora_b_fwd + + torch.manual_seed(42) + bs, rank, num_loras = 8, 16, 3 + n_slices = 3 + q_dim, kv_dim = 128, 64 + total_out = q_dim + 2 * kv_dim + x = torch.randn(bs, n_slices * rank, device="cuda", dtype=torch.bfloat16) + weights = torch.randn( + num_loras, total_out, rank, device="cuda", dtype=torch.bfloat16 + ) + output_offset = torch.tensor( + [0, q_dim, q_dim + kv_dim, total_out], device="cuda", dtype=torch.int32 + ) + wi = [i % num_loras for i in range(bs)] + lora_ranks = [rank] * num_loras + scalings = [1.0] * num_loras + + bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) + bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) + + base_plain = torch.randn(bs, total_out, device="cuda", dtype=torch.bfloat16) + base_sorted = base_plain.clone() + + max_qkv_out_dim = max(q_dim, kv_dim) + out_plain = qkv_lora_b_fwd( + x, weights, bi_plain, output_offset, max_qkv_out_dim, base_plain + ) + out_sorted = qkv_lora_b_fwd( + x, weights, bi_sorted, output_offset, max_qkv_out_dim, base_sorted + ) + _check_close(out_plain, out_sorted, "qkv_lora_b") + + +def test_gate_up_lora_b(): + from sglang.srt.lora.triton_ops import gate_up_lora_b_fwd + + torch.manual_seed(42) + bs, rank, num_loras = 8, 16, 3 + output_dim = 128 + x = torch.randn(bs, 2 * rank, device="cuda", dtype=torch.bfloat16) + weights = torch.randn( + num_loras, 2 * output_dim, rank, device="cuda", dtype=torch.bfloat16 + ) + wi = [i % num_loras for i in range(bs)] + lora_ranks = [rank] * num_loras + scalings = [1.0] * num_loras + + bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) + bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) + + base_plain = torch.randn(bs, 2 * output_dim, device="cuda", dtype=torch.bfloat16) + base_sorted = base_plain.clone() + + out_plain = gate_up_lora_b_fwd(x, weights, bi_plain, output_dim, base_plain) + out_sorted = gate_up_lora_b_fwd(x, weights, bi_sorted, output_dim, base_sorted) + _check_close(out_plain, out_sorted, "gate_up_lora_b") + + +def test_mixed_ranks(): + """Test with different LoRA ranks per adapter.""" + from sglang.srt.lora.triton_ops import sgemm_lora_a_fwd + + torch.manual_seed(42) + bs, input_dim, num_loras = 12, 256, 4 + max_rank = 32 + lora_ranks = [8, 16, 32, 16] + scalings = [0.25, 0.5, 1.0, 2.0] + # Use max_rank for weight shape, kernel handles per-adapter rank + weights = torch.randn( + num_loras, max_rank, input_dim, device="cuda", dtype=torch.bfloat16 + ) + x = torch.randn(bs, input_dim, device="cuda", dtype=torch.bfloat16) + wi = [i % num_loras for i in range(bs)] + + bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) + bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) + + out_plain = sgemm_lora_a_fwd(x, weights, bi_plain) + out_sorted = sgemm_lora_a_fwd(x, weights, bi_sorted) + _check_close(out_plain, out_sorted, "sgemm_lora_a_mixed_ranks") + + +def test_single_adapter(): + """All sequences use the same adapter.""" + from sglang.srt.lora.triton_ops import sgemm_lora_a_fwd + + torch.manual_seed(42) + bs, input_dim, rank, num_loras = 16, 256, 16, 2 + x = torch.randn(bs, input_dim, device="cuda", dtype=torch.bfloat16) + weights = torch.randn( + num_loras, rank, input_dim, device="cuda", dtype=torch.bfloat16 + ) + wi = [0] * bs # all adapter 0 + lora_ranks = [rank, rank] + scalings = [1.0, 1.0] + + bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) + bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) + + out_plain = sgemm_lora_a_fwd(x, weights, bi_plain) + out_sorted = sgemm_lora_a_fwd(x, weights, bi_sorted) + _check_close(out_plain, out_sorted, "sgemm_lora_a_single_adapter") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])