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

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 = (
+        return sgemm_lora_a_fwd(
-            pruned_batch_info if pruned_batch_info is not None else self.batch_info
+            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 = (
+        return sgemm_lora_b_fwd(
-            pruned_batch_info if pruned_batch_info is not None else self.batch_info
+            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(
+        lora_a_output = sgemm_lora_a_fwd(x, gate_up_lora_a, sgemm_info, stack_num=2)
-            x, gate_up_lora_a, self.batch_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),
+                lora_ranks=torch.zeros(mlpb, dtype=torch.int32),
-                scalings=torch.zeros(self.max_loras_per_batch, dtype=torch.float),
+                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

@@ -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_physical = _resolve_token_positions(
-        s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1
+        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) + (
+    output_ptr = (
-        s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
+        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

@@ -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

@@ -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_physical = _resolve_token_positions(
-        s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1
+        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) + (
+    output_ptr = (
-        s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
+        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

@@ -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_physical = _resolve_token_positions(
-        s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1
+        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

@@ -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_physical = _resolve_token_positions(
-        s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1
+        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) + (
+    output_ptr = output + (
-        s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
+        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

@@ -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"])