minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[lora] Speedup triton backend sgemm calls with better grid (#22386)

Browse Source
This commit is contained in:
Kurt Shuster
2026-04-15 13:47:07 -07:00
committed by GitHub
parent 113d654152
commit 32d9fe5a32
7 changed files with 415 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/lora/backend/triton_backend.py
+86 -16
View File
@@ -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)
)
python/sglang/srt/lora/triton_ops/gate_up_lora_b.py
+21 -4
View File
@@ -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,
python/sglang/srt/lora/triton_ops/kernel_utils.py
+19
View File
@@ -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)
python/sglang/srt/lora/triton_ops/qkv_lora_b.py
+21 -4
View File
@@ -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,
python/sglang/srt/lora/triton_ops/sgemm_lora_a.py
+17 -5
View File
@@ -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,
python/sglang/srt/lora/triton_ops/sgemm_lora_b.py
+15 -4
View File
@@ -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,
test/registered/lora/test_sgemm_sorted_by_adapter.py
+236
View File
@@ -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"])