chore: import upstream snapshot with attribution

python/sglang/kernels/ops/gemm/__init__.py

@@ -0,0 +1,138 @@
+"""GEMM and fused-GEMM kernels."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.selector import get_kernel
+from sglang.kernels.spec import (
+    CapabilityRequirement,
+    FormatSignature,
+    KernelBackend,
+    KernelSpec,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+_CUDA = CapabilityRequirement(requires_cuda=True)
+
+register_kernel(
+    KernelSpec(
+        op="gemm.fp8_scaled_mm",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel:fp8_scaled_mm",
+        format_signature=FormatSignature(
+            supported_dtypes=("float8_e4m3fn",),
+            description="C = (A_fp8 @ B_fp8) * scales_a * scales_b (+ bias)",
+        ),
+        description="FP8 scaled matmul (sgl_kernel wheel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="gemm.dsv3_fused_a_gemm",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel:dsv3_fused_a_gemm",
+        format_signature=FormatSignature(
+            supported_dtypes=("bfloat16",),
+            description="DeepSeek-V3 fused QKV-A GEMM",
+        ),
+        description="DeepSeek-V3 fused-A GEMM (sgl_kernel wheel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="gemm.dsv3_fused_a_gemm",
+        backend=KernelBackend.CUDA_JIT,
+        target="sglang.jit_kernel.dsv3_fused_a_gemm:dsv3_fused_a_gemm",
+        capability=_CUDA,
+        format_signature=FormatSignature(
+            supported_dtypes=("bfloat16",),
+            description="DeepSeek-V3 fused QKV-A GEMM (drop-in with AOT signature)",
+        ),
+        description="DeepSeek-V3 fused-A GEMM (sglang.jit_kernel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="gemm.dsv3_router_gemm",
+        backend=KernelBackend.CUDA_JIT,
+        target="sglang.jit_kernel.dsv3_router_gemm:dsv3_router_gemm",
+        capability=_CUDA,
+        format_signature=FormatSignature(
+            supported_dtypes=("bfloat16",),
+            description="DeepSeek-V3 router GEMM; num_tokens in [1, 16]",
+        ),
+        description="DeepSeek-V3 router GEMM (sglang.jit_kernel, JIT-only).",
+    )
+)
+
+
+def fp8_scaled_mm(
+    mat_a: torch.Tensor,
+    mat_b: torch.Tensor,
+    scales_a: torch.Tensor,
+    scales_b: torch.Tensor,
+    out_dtype: torch.dtype,
+    bias: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 scaled matmul: ``(mat_a @ mat_b) * scales_a * scales_b (+ bias)``."""
+    return get_kernel("gemm.fp8_scaled_mm", KernelBackend.CUDA_AOT)(
+        mat_a, mat_b, scales_a, scales_b, out_dtype, bias
+    )
+
+
+def dsv3_fused_a_gemm(
+    mat_a: torch.Tensor,
+    mat_b: torch.Tensor,
+    output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """DeepSeek-V3 fused QKV-A GEMM."""
+    return get_kernel("gemm.dsv3_fused_a_gemm", KernelBackend.CUDA_AOT)(
+        mat_a, mat_b, output
+    )
+
+
+def dsv3_router_gemm(
+    hidden_states: torch.Tensor,
+    router_weights: torch.Tensor,
+    out_dtype: Optional[torch.dtype] = None,
+    output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """DeepSeek-V3 router GEMM (JIT-backed). ``out_dtype`` defaults to bfloat16."""
+    impl = get_kernel("gemm.dsv3_router_gemm", KernelBackend.CUDA_JIT)
+    if out_dtype is None:
+        return impl(hidden_states, router_weights, output=output)
+    return impl(hidden_states, router_weights, out_dtype, output)
+
+
+__all__ = ["fp8_scaled_mm", "dsv3_fused_a_gemm", "dsv3_router_gemm"]
+
+
+# LoRA SGMV Triton kernels migrated into this group (from lora/triton_ops);
+# registered for inventory. Import them from their modules.
+_TRITON_KERNELS = [
+    ("chunked_embedding_lora_a", "chunked_embedding_lora_a_forward"),
+    ("chunked_sgmv_expand", "chunked_sgmv_lora_expand_forward"),
+    ("chunked_sgmv_shrink", "chunked_sgmv_lora_shrink_forward"),
+    ("embedding_lora_a", "embedding_lora_a_fwd"),
+    ("gate_up_lora_b", "gate_up_lora_b_fwd"),
+    ("qkv_lora_b", "qkv_lora_b_fwd"),
+    ("sgemm_lora_a", "sgemm_lora_a_fwd"),
+    ("sgemm_lora_b", "sgemm_lora_b_fwd"),
+    ("kv_b_lora_absorbed", "step_a_q_fwd"),
+    ("kv_b_lora_absorbed", "step_b_q_fwd"),
+    ("kv_b_lora_absorbed", "step_a_v_fwd"),
+    ("kv_b_lora_absorbed", "step_b_v_fwd"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"gemm.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.gemm.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn

python/sglang/kernels/ops/gemm/chunked_embedding_lora_a.py

@@ -0,0 +1,142 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _chunked_embedding_lora_a_kernel(
+    # Pointers to tensors
+    input_ids,
+    weights,
+    output,
+    # Dimensions
+    vocab_size,
+    rank,
+    num_loras,
+    # Strides
+    w_stride_0,  # stride for lora index
+    w_stride_1,  # stride for rank
+    w_stride_2,  # stride for vocab
+    output_stride_0,
+    output_stride_1,
+    # Chunk info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    num_segments,
+    permutation,
+    # Meta-parameters
+    BLOCK_RANK: tl.constexpr,
+):
+    """
+    Embedding lookup for LoRA A weights without support for extra tokens.
+
+    Each program handles one chunk of tokens across rank dimension
+    """
+    chunk_idx = tl.program_id(axis=0)
+    # If chunk id is larger than actual number of chunks, skip
+    if chunk_idx >= num_segments:
+        return
+    chunk_start = tl.load(seg_indptr + chunk_idx)
+    chunk_end = tl.load(seg_indptr + chunk_idx + 1)
+    if chunk_start == chunk_end:
+        return
+    # Load LoRA adapter index for this segment, then look up the rank
+    lora_index = tl.load(weight_indices + chunk_idx)
+    rank_val = tl.load(lora_ranks + lora_index)
+    # If rank is 0, skip
+    if rank_val == 0:
+        return
+    # for each token in chunk, load embedding across rank dimension
+    for c in range(chunk_start, chunk_end):
+        s_index = tl.load(permutation + c)
+        # Load the token ID
+        token_id = tl.load(input_ids + s_index)
+        # Process in chunks of BLOCK_RANK dimensions
+        num_blocks = tl.cdiv(rank_val, BLOCK_RANK)
+
+        for block_id in range(num_blocks):
+            rank_offset = tl.arange(0, BLOCK_RANK) + block_id * BLOCK_RANK
+            rank_mask = rank_offset < rank_val
+
+            # Use regular LoRA A weights
+            # weights shape: (num_loras, rank, vocab_size)
+            # We need to load weights[lora_index, rank_offset, token_id]
+            weight_ptr = (
+                weights
+                + lora_index * w_stride_0
+                + rank_offset * w_stride_1
+                + token_id * w_stride_2
+            )
+            emb_values = tl.load(weight_ptr, mask=rank_mask, other=0.0)
+
+            # Write to output
+            output_ptr = (
+                output + s_index * output_stride_0 + rank_offset * output_stride_1
+            )
+            tl.store(output_ptr, emb_values, mask=rank_mask)
+
+
+def chunked_embedding_lora_a_forward(
+    input_ids: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    vocab_size: int,
+) -> torch.Tensor:
+    """
+    Chunked Forward pass for LoRA A embedding lookup; each program handles one chunk of embedding lookup work
+    belonging to the same adapter
+
+    Args:
+        input_ids: (s,) token IDs
+        weights: (num_loras, rank, vocab_size) LoRA A embedding weights
+        batch_info: LoRABatchInfo containing batch information
+        vocab_size: base vocabulary size
+
+    Returns:
+        output: (s, rank) embedded features
+    """
+    assert input_ids.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(input_ids.shape) == 1
+    assert len(weights.shape) == 3
+
+    S = input_ids.shape[0]
+    num_loras = weights.shape[0]
+    rank = weights.shape[1]
+
+    # Block size for rank dimension
+    BLOCK_RANK = 128
+    num_segments = batch_info.num_segments
+    segment_grid = (
+        batch_info.weight_indices.shape[0]
+        if batch_info.use_cuda_graph
+        else num_segments
+    )
+    # 1D Grid: one program per chunk of embedding lookup work
+    grid = (segment_grid,)
+    output = torch.zeros((S, rank), device=input_ids.device, dtype=weights.dtype)
+
+    _chunked_embedding_lora_a_kernel[grid](
+        input_ids,
+        weights,
+        output,
+        vocab_size,
+        rank,
+        num_loras,
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        segment_grid,
+        batch_info.permutation,
+        BLOCK_RANK,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/chunked_sgmv_expand.py

@@ -0,0 +1,238 @@
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.lora_tuning_config import get_lora_expand_config
+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"], kwargs["OUTPUT_DIM"])
+)
+@triton.jit(do_not_specialize=["num_segs", "output_stride_0", "output_stride_1"])
+def _chunked_lora_expand_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Output strides may differ from OUTPUT_DIM when compact LoRA output is
+    # accumulated into a wider base projection.
+    output_stride_0,
+    output_stride_1,
+    # 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) or a wider base output.
+    """
+    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
+
+    pid_s = tl.program_id(axis=2)
+    if pid_s >= num_segs:
+        return
+
+    seg_start = tl.load(seg_indptr + pid_s)
+    seg_end = tl.load(seg_indptr + pid_s + 1)
+    if seg_start == seg_end:
+        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
+
+    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, other=0
+    )
+
+    # 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
+
+    # Block shapes — use auto-tuned config if available, else defaults
+    BLOCK_M = batch_info.max_len
+    config = get_lora_expand_config(
+        K=OUTPUT_DIM, R=MAX_RANK, num_slices=num_slices, chunk_size=BLOCK_M
+    )
+    BLOCK_K = config["BLOCK_K"]
+    BLOCK_N = config["BLOCK_N"]
+
+    num_segments = batch_info.num_segments
+    segment_grid = (
+        batch_info.weight_indices.shape[0]
+        if batch_info.use_cuda_graph
+        else num_segments
+    )
+
+    grid = (
+        triton.cdiv(max_slice_size, BLOCK_N),
+        num_slices,  # number of slices in the input/output
+        segment_grid,
+    )
+
+    if base_output is None:
+        output = torch.zeros((M, OUTPUT_DIM), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    # Optional launch params from tuned config
+    extra_kwargs = {}
+    if "num_warps" in config:
+        extra_kwargs["num_warps"] = config["num_warps"]
+    if "num_stages" in config:
+        extra_kwargs["num_stages"] = config["num_stages"]
+    if "maxnreg" in config:
+        extra_kwargs["maxnreg"] = config["maxnreg"]
+
+    _chunked_lora_expand_kernel[grid](
+        x=x,
+        weights=weights,
+        output=output,
+        output_stride_0=output.stride(0),
+        output_stride_1=output.stride(1),
+        seg_indptr=batch_info.seg_indptr,
+        weight_indices=batch_info.weight_indices,
+        lora_ranks=batch_info.lora_ranks,
+        permutation=batch_info.permutation,
+        num_segs=segment_grid,
+        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,
+        **extra_kwargs,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/chunked_sgmv_shrink.py

@@ -0,0 +1,196 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.lora_tuning_config import get_lora_shrink_config
+from sglang.srt.lora.utils import LoRABatchInfo
+from sglang.srt.utils import cached_triton_kernel
+
+
+@cached_triton_kernel(
+    lambda _, kwargs: (kwargs["K"], kwargs["NUM_SLICES"], kwargs["BLOCK_M"])
+)
+@triton.jit(do_not_specialize=["num_segs"])
+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)
+
+    seg_start = tl.load(seg_indptr + pid_s)
+    seg_end = tl.load(seg_indptr + pid_s + 1)
+    if seg_start == seg_end:
+        return
+
+    # 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
+
+    # 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, other=0
+    )
+
+    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 — use auto-tuned config if available, else defaults
+    BLOCK_M = batch_info.max_len
+    # weights shape is (num_lora, num_slices * rank, input_dim)
+    MAX_RANK = weights.shape[1] // num_slices
+    config = get_lora_shrink_config(
+        K=weights.shape[2], R=MAX_RANK, num_slices=num_slices, chunk_size=BLOCK_M
+    )
+    BLOCK_N = config["BLOCK_N"]
+    BLOCK_K = config["BLOCK_K"]
+
+    S = x.shape[0]
+    N = weights.shape[1]
+    K = weights.shape[2]
+    assert x.shape[-1] == K
+
+    num_segments = batch_info.num_segments
+    segment_grid = (
+        batch_info.weight_indices.shape[0]
+        if batch_info.use_cuda_graph
+        else num_segments
+    )
+    grid = (
+        triton.cdiv(N, BLOCK_N),
+        segment_grid,
+    )
+
+    # Optional launch params from tuned config
+    extra_kwargs = {}
+    if "num_warps" in config:
+        extra_kwargs["num_warps"] = config["num_warps"]
+    if "num_stages" in config:
+        extra_kwargs["num_stages"] = config["num_stages"]
+
+    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=segment_grid,
+        # constants
+        N=N,
+        K=K,
+        NUM_SLICES=num_slices,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+        **extra_kwargs,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=1024,R=64,S=1,device=NVIDIA_H200.json

@@ -0,0 +1,29 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 1,
+        "maxnreg": 128
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2,
+        "maxnreg": 160
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 16,
+        "num_warps": 8,
+        "num_stages": 2,
+        "maxnreg": 128
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=4096,R=64,S=3,device=NVIDIA_H200.json

@@ -0,0 +1,29 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 3,
+        "maxnreg": 160
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 1,
+        "maxnreg": 160
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 16,
+        "num_warps": 8,
+        "num_stages": 2,
+        "maxnreg": 128
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_expand,K=6144,R=64,S=2,device=NVIDIA_H200.json

@@ -0,0 +1,29 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2,
+        "maxnreg": 112
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 3,
+        "maxnreg": 160
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 16,
+        "num_warps": 8,
+        "num_stages": 3,
+        "maxnreg": 128
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=2,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 4
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=3,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 3
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=2048,R=64,S=1,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 128,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 4
+    }
+}

python/sglang/kernels/ops/gemm/csgmv_configs/triton_3_5_1/lora_shrink,K=3072,R=64,S=1,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 128,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 4
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 3
+    }
+}

python/sglang/kernels/ops/gemm/embedding_lora_a.py

@@ -0,0 +1,186 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _embedding_lora_a_kernel(
+    # Pointers to tensors
+    input_ids,
+    weights,
+    output,
+    extra_embeddings,
+    # Dimensions
+    vocab_size,
+    rank,
+    num_loras,
+    # Strides
+    w_stride_0,  # stride for lora index
+    w_stride_1,  # stride for rank
+    w_stride_2,  # stride for vocab
+    output_stride_0,
+    output_stride_1,
+    extra_emb_stride_0,  # stride for lora index
+    extra_emb_stride_1,  # stride for token
+    extra_emb_stride_2,  # stride for hidden dim (= rank for extra embeddings)
+    # Batch info
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    # Meta-parameters
+    BLOCK_RANK: tl.constexpr,
+    HAS_EXTRA_EMBEDDINGS: tl.constexpr,
+):
+    """
+    Embedding lookup for LoRA A weights with support for extra tokens.
+
+    Each program handles one token across a block of rank dimensions.
+    Grid: (cdiv(max_len, 1), bs) - one program per token in each batch
+    """
+    batch_id = tl.program_id(axis=1)
+    token_idx = tl.program_id(axis=0)
+
+    w_index = tl.load(weight_indices + batch_id)
+    rank_val = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, skip
+    if rank_val == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_len = tl.load(seg_lens + batch_id)
+
+    # Check if this token is within the segment
+    if token_idx >= seg_len:
+        return
+
+    # Load the token ID
+    token_id = tl.load(input_ids + seg_start + token_idx)
+
+    # Process in chunks of BLOCK_RANK dimensions
+    num_blocks = tl.cdiv(rank_val, BLOCK_RANK)
+
+    for block_id in range(num_blocks):
+        rank_offset = tl.arange(0, BLOCK_RANK) + block_id * BLOCK_RANK
+        rank_mask = rank_offset < rank_val
+
+        # Check if this is an extra token
+        is_extra_token = token_id >= vocab_size
+
+        if HAS_EXTRA_EMBEDDINGS and is_extra_token:
+            # Use extra embeddings
+            extra_token_id = token_id - vocab_size
+            extra_emb_ptr = (
+                extra_embeddings
+                + w_index * extra_emb_stride_0
+                + extra_token_id * extra_emb_stride_1
+                + rank_offset * extra_emb_stride_2
+            )
+            emb_values = tl.load(extra_emb_ptr, mask=rank_mask, other=0.0)
+        else:
+            # Use regular LoRA A weights
+            # weights shape: (num_loras, rank, vocab_size)
+            # We need to load weights[w_index, rank_offset, token_id]
+            token_id_clamped = tl.minimum(token_id, vocab_size - 1)
+            weight_ptr = (
+                weights
+                + w_index * w_stride_0
+                + rank_offset * w_stride_1
+                + token_id_clamped * w_stride_2
+            )
+            emb_values = tl.load(weight_ptr, mask=rank_mask, other=0.0)
+
+        # Write to output
+        output_ptr = (
+            output
+            + (seg_start + token_idx) * output_stride_0
+            + rank_offset * output_stride_1
+        )
+        tl.store(output_ptr, emb_values, mask=rank_mask)
+
+
+def embedding_lora_a_fwd(
+    input_ids: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    vocab_size: int,
+    extra_embeddings: torch.Tensor = None,
+) -> torch.Tensor:
+    """
+    Forward pass for LoRA A embedding lookup.
+
+    Args:
+        input_ids: (s,) token IDs
+        weights: (num_loras, rank, vocab_size) LoRA A embedding weights
+        batch_info: LoRABatchInfo containing batch information
+        vocab_size: base vocabulary size
+        extra_embeddings: (num_loras, num_extra_tokens, rank) extra token embeddings
+
+    Returns:
+        output: (s, rank) embedded features
+    """
+    assert input_ids.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(input_ids.shape) == 1
+    assert len(weights.shape) == 3
+
+    S = input_ids.shape[0]
+    num_loras = weights.shape[0]
+    rank = weights.shape[1]
+    vocab_size_weights = weights.shape[2]
+
+    # Block size for rank dimension
+    BLOCK_RANK = 128
+
+    has_extra_embeddings = extra_embeddings is not None
+
+    if has_extra_embeddings:
+        assert extra_embeddings.is_contiguous()
+        extra_emb_stride = (
+            extra_embeddings.stride(0),
+            extra_embeddings.stride(1),
+            extra_embeddings.stride(2),
+        )
+    else:
+        # Create dummy tensor to satisfy Triton
+        extra_embeddings = torch.empty(
+            (1, 1, 1), device=input_ids.device, dtype=weights.dtype
+        )
+        extra_emb_stride = (1, 1, 1)
+
+    # Grid: one program per token in each batch segment
+    grid = (
+        batch_info.max_len,
+        batch_info.bs,
+    )
+
+    output = torch.zeros((S, rank), device=input_ids.device, dtype=weights.dtype)
+
+    _embedding_lora_a_kernel[grid](
+        input_ids,
+        weights,
+        output,
+        extra_embeddings,
+        vocab_size,
+        rank,
+        num_loras,
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        extra_emb_stride[0],
+        extra_emb_stride[1],
+        extra_emb_stride[2],
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        BLOCK_RANK,
+        has_extra_embeddings,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/gate_up_lora_b.py

@@ -0,0 +1,204 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _gate_up_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Parameters of size
+    K,  # K = R
+    output_dim,
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths,ranks and weight id
+    seg_lens,
+    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,
+    # For fused output scaling
+    scalings,
+):
+    """
+    This kernel packs 2 sgemms (gate/up) into a single kernel. The multiplication
+    results are accumulated into the output tensor.
+
+    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, 2 * 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, 2 * output_dim, K).
+        output (Tensor): The output tensor where the result is stored.
+            Shape: (s, 2 * output_dim).
+    """
+    # output_dim >> K
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len.
+    # gate_up_id decides which of gate or up (0: gate, 1: up)
+    batch_id = tl.program_id(axis=2)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    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)
+
+    # Adjust K (rank) according to the specific LoRA adapter
+    K = tl.minimum(K, rank)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    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
+    # and accumulate
+    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)
+
+    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
+    )
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (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, :] < output_dim),
+            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
+        + 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)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def gate_up_lora_b_fwd(
+    x: torch.Tensor,
+    gate_up_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_dim: int,
+    base_output: torch.Tensor = None,
+) -> torch.Tensor:
+
+    # x: (s, 2 * r)
+    # gate_up_lora_b: (num_lora, 2 * output_dim, r)
+    # output: (s, 2 * output_dim)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # lora_output[:, :output_dim] = sgemm(x[:, :r], gate_up_lora_b[:, :output_dim, :])
+    # lora_output[:, output_dim:]
+    #      = sgemm(x[:, r:], gate_up_lora_b[:, output_dim:, :])
+
+    # Get dims
+    s = x.shape[0]
+    input_dim = x.shape[1]
+    r = gate_up_lora_b.shape[-1]
+    assert input_dim == 2 * r
+
+    BLOCK_S = 16
+    BLOCK_R = 16
+    BLOCK_OUT = 64
+
+    grid_b = (
+        triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(output_dim, BLOCK_OUT),
+        2,  # this dimension decides current block computes on gate or up proj
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((s, 2 * output_dim), device=x.device, dtype=x.dtype)
+    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,
+        output,
+        r,
+        output_dim,
+        x.stride(0),
+        x.stride(1),
+        gate_up_lora_b.stride(0),
+        gate_up_lora_b.stride(1),
+        gate_up_lora_b.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_OUT,
+        BLOCK_R,
+        batch_info.scalings,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/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/kernels/ops/gemm/kv_b_lora_absorbed.py

@@ -0,0 +1,853 @@
+"""Triton kernels for absorbed-MLA ``kv_b_proj`` LoRA correction.
+
+The absorbed-MLA path bypasses ``kv_b_proj.forward()`` and folds the K/V
+sides as plain BMMs ``q_nope @ w_kc`` and ``attn_output @ w_vc``.  When a
+LoRA adapter is active on ``kv_b_proj`` we add the LoRA delta to
+``q_nope_out`` / ``attn_bmm_output`` manually.
+
+Using the standard LoRA factored math we *never* materialize ``B @ A``:
+
+    q_correction = q_nope     @ B_kc @ A   * scaling          # K-side
+    v_correction = attn_output @ A.T @ B_vc.T * scaling       # V-side
+
+where ``A: (slot, rank, kv_lora_rank)`` is the LoRA-A of ``kv_b_proj``
+(shared across heads) and ``B: (slot, num_heads*(qk_nope+v_head_dim), rank)``
+is the LoRA-B; ``B_kc`` / ``B_vc`` are its K-half / V-half slices.
+
+Four kernels split the math along the factorization boundary, all using
+the SGMM idiom from ``sgemm_lora_a`` / ``qkv_lora_b`` and the segment-indptr
+routing used by ``chunked_sgmv_*``:
+
+  * ``step_a_q_fwd``: per-head per-slot SGMM, ``(S,H,qk_nope) -> (S,H,rank)``
+  * ``step_b_q_fwd``: shared-A per-slot SGMM, scaled+accumulated,
+    ``(S,H,rank) -> (S,H,kv_lora_rank)``
+  * ``step_a_v_fwd``: shared-A.T per-slot SGMM, ``(S,H,kv_lora_rank) -> (S,H,rank)``
+  * ``step_b_v_fwd``: per-head per-slot SGMM with V-half of B, transposed,
+    scaled+accumulated, ``(S,H,rank) -> (S,H,v_head_dim)``
+
+Grid axes for each kernel:
+  axis 0 : output tile in (S, N)         -- tile_id = pid_s * num_pid_n + pid_n
+  axis 1 : head_id                       -- per-head weight slice
+  axis 2 : batch_id (segment / request)  -- per-slot weight routing via weight_indices
+
+Per-segment routing: each program derives its segment length from
+``seg_indptr[segment_id + 1] - seg_indptr[segment_id]``, loads
+``weight_indices[segment_id]`` once, and uses that slot's slice of the LoRA
+weight stack.  When ``permutation`` is present, rows are routed through it,
+matching the csgmv backend's adapter-grouped chunks.  No Python loops over slots
+or heads.
+
+The math also stays in the input dtype (no fp32 round-trip) -- the
+contraction dim ``rank`` is small (typically 16-64), so bf16 accumulation
+over it is acceptable.  ``tl.dot`` itself uses fp32 accumulation internally.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
+from sglang.srt.lora.utils import LoRABatchInfo
+
+# ---------------------------------------------------------------------------
+# Block sizes -- chosen per-kernel from the natural shape of each step.
+#
+# The factored math gives the four kernels these contraction (K) and output
+# (N) ranges (for Kimi-K2.5: rank=16-32, qk_nope=v_head_dim=128, kv_lora_rank=512):
+#
+#                              K (contraction)      N (output)
+#   step_a_q                   qk_nope (~128)       rank (~16-32)
+#   step_b_q                   rank (~16-32)        kv_lora_rank (~512)
+#   step_a_v                   kv_lora_rank (~512)  rank (~16-32)
+#   step_b_v                   rank (~16-32)        v_head_dim (~128)
+#
+# So the "step_a_*" kernels want a large BLOCK_K (to keep loop iters small)
+# and a small BLOCK_N (matched to rank to avoid wasted tile lanes), while
+# the "step_b_*" kernels are the inverse.  Kernels aren't autotuned -- the
+# decode-shape workload is too small to benefit and the sweep surface is
+# wide.
+# ---------------------------------------------------------------------------
+
+_BLOCK_S = 16
+_STEP_A_BLOCK_K = 64  # contraction over qk_nope (~128) or kv_lora_rank (~512)
+_STEP_A_BLOCK_N = 16  # output is rank
+_STEP_B_BLOCK_K = 16  # contraction is rank
+_STEP_B_BLOCK_N = 64  # output is kv_lora_rank (~512) or v_head_dim (~128)
+
+
+def _num_segments(batch_info: LoRABatchInfo) -> int:
+    return batch_info.num_segments or batch_info.bs
+
+
+def _max_segment_len(batch_info: LoRABatchInfo) -> int:
+    if batch_info.max_len is not None:
+        return batch_info.max_len
+    if batch_info.seg_lens is not None:
+        return int(batch_info.seg_lens.max().item())
+    raise ValueError("LoRA batch_info must provide max_len or seg_lens.")
+
+
+def _segment_grid_size(batch_info: LoRABatchInfo, num_segments: int) -> int:
+    return (
+        batch_info.weight_indices.shape[0]
+        if batch_info.use_cuda_graph
+        else num_segments
+    )
+
+
+# ---------------------------------------------------------------------------
+# Kernel 1 -- Step A_q: per-head per-slot SGMM, reads K-half of B
+#
+#     q_lora_a[t, h, r] = sum_{i<qk_nope} q_nope[t, h, i] * B[slot, h*FULL_K + i, r]
+#
+# x      : (S, H, qk_nope)
+# w (B)  : (num_lora, H*FULL_K, rank)   -- FULL_K = qk_nope + v_head_dim
+# out    : (S, H, rank)                 -- fresh allocation, no accumulate
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_a_q_kernel(
+    x,
+    w,
+    out,
+    # dims
+    S,
+    H_FULL_K,  # H * (qk_nope + v_head_dim), the row-stride landmark
+    K,  # qk_nope (contraction)
+    N,  # rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,
+    w_stride_k,
+    out_stride_s,
+    out_stride_h,
+    out_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    num_segments,
+    # meta
+    FULL_K: tl.constexpr,  # per-head row stride in B (qk_nope + v_head_dim)
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+
+    # Truncate output N to this slot's rank (allows mixed-rank batches).
+    N_eff = tl.minimum(N, cur_rank)
+
+    num_pid_n = tl.cdiv(N_eff, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    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)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    # Clamp masked-lane indices into the valid range so pointer arithmetic
+    # stays in-bounds even before the load mask drops the values.
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    safe_n = tl.minimum(n_offset, N_eff - 1)
+
+    head_row_base = (
+        head_id * FULL_K
+    )  # row offset for this head's K-half (i in [0, qk_nope))
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K
+        safe_k = tl.minimum(cur_k, K - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # B[slot, h*FULL_K + i, r]: row dim of B carries i (= GEMM K),
+        # column dim carries r (= GEMM N).
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + (head_row_base + safe_k[:, None]) * w_stride_n
+            + safe_n[None, :] * w_stride_k,
+            mask=k_mask[:, None] & (n_offset[None, :] < N_eff),
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    out_offs = (
+        safe_row[:, None] * out_stride_s
+        + head_id * out_stride_h
+        + safe_n[None, :] * out_stride_n
+    )
+    out_mask = row_mask[:, None] & (n_offset[None, :] < N_eff)
+    tl.store(out + out_offs, partial_sum, mask=out_mask)
+
+
+def step_a_q_fwd(
+    q_nope: torch.Tensor,
+    B_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    full_K_per_head: int,
+) -> torch.Tensor:
+    """Step A of the q-side correction.
+
+    Args:
+        q_nope: ``(S, H, qk_nope)``, the absorbed-MLA q intermediate.
+        B_buf: ``(num_lora, H*full_K_per_head, rank)`` from the LoRA pool.
+        batch_info: standard ``LoRABatchInfo``.
+        full_K_per_head: ``qk_nope + v_head_dim``, the row stride per head in B.
+
+    Returns:
+        ``(S, H, rank)`` -- per-token, per-head low-rank intermediate, ready for step B_q.
+    """
+    S, H, qk_nope_dim = q_nope.shape
+    rank = B_buf.shape[-1]
+    out = torch.empty((S, H, rank), device=q_nope.device, dtype=q_nope.dtype)
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S) * triton.cdiv(rank, _STEP_A_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+
+    _step_a_q_kernel[grid](
+        q_nope,
+        B_buf,
+        out,
+        S,
+        H * full_K_per_head,
+        qk_nope_dim,
+        rank,
+        q_nope.stride(0),
+        q_nope.stride(1),
+        q_nope.stride(2),
+        B_buf.stride(0),
+        B_buf.stride(1),
+        B_buf.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        segment_grid,
+        FULL_K=full_K_per_head,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_A_BLOCK_N,
+        BLOCK_K=_STEP_A_BLOCK_K,
+    )
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Kernel 2 -- Step B_q: shared-A per-slot SGMM, scaled + accumulated
+#
+#     base[t, h, k] += sum_r x[t, h, r] * A[slot, r, k] * scaling
+#
+# x      : (S, H, rank)
+# w (A)  : (num_lora, rank, kv_lora_rank)
+# base   : (S, H, kv_lora_rank), updated in-place (accumulated)
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_b_q_kernel(
+    x,
+    w,
+    base,
+    # dims
+    S,
+    K,  # rank (contraction)
+    N,  # kv_lora_rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_k,
+    w_stride_n,
+    b_stride_s,
+    b_stride_h,
+    b_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    scalings,
+    num_segments,
+    # meta
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+    scaling = tl.load(scalings + w_index)
+
+    # Truncate contraction K to this slot's rank.
+    K_eff = tl.minimum(K, cur_rank)
+
+    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
+
+    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)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    n_mask = n_offset[None, :] < N
+    safe_n = tl.minimum(n_offset, N - 1)
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K_eff, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K_eff
+        safe_k = tl.minimum(cur_k, K_eff - 1)
+
+        # x[s, h, k]  (k iterates over rank)
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # A[slot, k, n]: read k along contraction, n along output.
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + safe_n[None, :] * w_stride_n,
+            mask=k_mask[:, None] & n_mask,
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+
+    # Accumulate into base[s, h, n].
+    base_offs = (
+        safe_row[:, None] * b_stride_s
+        + head_id * b_stride_h
+        + safe_n[None, :] * b_stride_n
+    )
+    out_mask = row_mask[:, None] & n_mask
+    partial_sum += tl.load(base + base_offs, mask=out_mask, other=0.0)
+    tl.store(base + base_offs, partial_sum, mask=out_mask)
+
+
+def step_b_q_fwd(
+    q_lora_a: torch.Tensor,
+    A_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor,
+) -> torch.Tensor:
+    """Step B of the q-side correction, accumulating into ``base_output``.
+
+    Args:
+        q_lora_a: ``(S, H, rank)`` from step A_q.
+        A_buf: ``(num_lora, rank, kv_lora_rank)`` from the LoRA pool.
+        batch_info: standard ``LoRABatchInfo``.
+        base_output: ``(S, H, kv_lora_rank)``, modified in-place
+            (the absorbed ``q_nope @ w_kc`` result).
+
+    Returns:
+        ``base_output`` (same object, mutated).
+    """
+    S, H, rank = q_lora_a.shape
+    kv_lora_rank = A_buf.shape[-1]
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S)
+        * triton.cdiv(kv_lora_rank, _STEP_B_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+
+    _step_b_q_kernel[grid](
+        q_lora_a,
+        A_buf,
+        base_output,
+        S,
+        rank,
+        kv_lora_rank,
+        q_lora_a.stride(0),
+        q_lora_a.stride(1),
+        q_lora_a.stride(2),
+        A_buf.stride(0),
+        A_buf.stride(1),
+        A_buf.stride(2),
+        base_output.stride(0),
+        base_output.stride(1),
+        base_output.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        batch_info.scalings,
+        segment_grid,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_B_BLOCK_N,
+        BLOCK_K=_STEP_B_BLOCK_K,
+    )
+    return base_output
+
+
+# ---------------------------------------------------------------------------
+# Kernel 3 -- Step A_v: shared-A.T per-slot SGMM (no scaling, fresh output)
+#
+#     attn_lora_a[t, h, r] = sum_k attn_output[t, h, k] * A[slot, r, k]
+#
+# x      : (S, H, kv_lora_rank)
+# w (A)  : (num_lora, rank, kv_lora_rank) -- accessed transposed vs step B_q
+# out    : (S, H, rank), fresh allocation
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_a_v_kernel(
+    x,
+    w,
+    out,
+    # dims
+    S,
+    K,  # kv_lora_rank (contraction)
+    N,  # rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,  # A's "rank" axis (= GEMM N)
+    w_stride_k,  # A's "kv_lora_rank" axis (= GEMM K)
+    out_stride_s,
+    out_stride_h,
+    out_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    num_segments,
+    # meta
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+
+    # Truncate output N to this slot's rank.
+    N_eff = tl.minimum(N, cur_rank)
+
+    num_pid_n = tl.cdiv(N_eff, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    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)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    safe_n = tl.minimum(n_offset, N_eff - 1)
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K
+        safe_k = tl.minimum(cur_k, K - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # A[slot, r, k] -- here we want each (k, r) so we read along k
+        # (inner / contraction) and produce r as output.  Stride access:
+        # the row dim is r (= GEMM N), column dim is k (= GEMM K).
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + safe_n[None, :] * w_stride_n,
+            mask=k_mask[:, None] & (n_offset[None, :] < N_eff),
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    out_offs = (
+        safe_row[:, None] * out_stride_s
+        + head_id * out_stride_h
+        + safe_n[None, :] * out_stride_n
+    )
+    out_mask = row_mask[:, None] & (n_offset[None, :] < N_eff)
+    tl.store(out + out_offs, partial_sum, mask=out_mask)
+
+
+def step_a_v_fwd(
+    attn_output: torch.Tensor,
+    A_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+) -> torch.Tensor:
+    """Step A of the v-side correction.
+
+    Args:
+        attn_output: ``(S, H, kv_lora_rank)``, the post-attention intermediate.
+        A_buf: ``(num_lora, rank, kv_lora_rank)``.
+        batch_info: standard ``LoRABatchInfo``.
+
+    Returns:
+        ``(S, H, rank)`` -- per-token, per-head low-rank intermediate for step B_v.
+    """
+    S, H, kv_lora_rank = attn_output.shape
+    rank = A_buf.shape[1]
+    out = torch.empty((S, H, rank), device=attn_output.device, dtype=attn_output.dtype)
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S) * triton.cdiv(rank, _STEP_A_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+
+    _step_a_v_kernel[grid](
+        attn_output,
+        A_buf,
+        out,
+        S,
+        kv_lora_rank,
+        rank,
+        attn_output.stride(0),
+        attn_output.stride(1),
+        attn_output.stride(2),
+        A_buf.stride(0),
+        A_buf.stride(1),
+        A_buf.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        segment_grid,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_A_BLOCK_N,
+        BLOCK_K=_STEP_A_BLOCK_K,
+    )
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Kernel 4 -- Step B_v: per-head per-slot SGMM with V-half of B (transposed),
+# scaled + accumulated
+#
+#     base[t, h, j] += sum_r x[t, h, r] * B[slot, h*FULL_K + qk_nope + j, r] * scaling
+#
+# x      : (S, H, rank)
+# w (B)  : (num_lora, H*FULL_K, rank), V-half slice via offset
+# base   : (S, H, v_head_dim), updated in-place (accumulated)
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_b_v_kernel(
+    x,
+    w,
+    base,
+    # dims
+    S,
+    K,  # rank (contraction)
+    N,  # v_head_dim (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,  # B's row dim (h*FULL_K + j) -- this is GEMM N
+    w_stride_k,  # B's rank dim -- this is GEMM K
+    b_stride_s,
+    b_stride_h,
+    b_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    scalings,
+    num_segments,
+    # meta
+    FULL_K: tl.constexpr,  # qk_nope + v_head_dim
+    QK_NOPE_OFFSET: tl.constexpr,  # offset of V-half within each head's row block
+    SORTED_BY_ADAPTER: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+    scaling = tl.load(scalings + w_index)
+
+    K_eff = tl.minimum(K, cur_rank)
+
+    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
+
+    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)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    n_mask = n_offset[None, :] < N
+    safe_n = tl.minimum(n_offset, N - 1)
+
+    # V-half row base for this head: h*FULL_K + qk_nope
+    head_row_base = head_id * FULL_K + QK_NOPE_OFFSET
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K_eff, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K_eff
+        safe_k = tl.minimum(cur_k, K_eff - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # B[slot, h*FULL_K + qk_nope + j, r] -- row dim is j (= GEMM N),
+        # column dim is r (= GEMM K).  Transposed access vs step A_q.
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + (head_row_base + safe_n[None, :]) * w_stride_n,
+            mask=k_mask[:, None] & n_mask,
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+
+    base_offs = (
+        safe_row[:, None] * b_stride_s
+        + head_id * b_stride_h
+        + safe_n[None, :] * b_stride_n
+    )
+    out_mask = row_mask[:, None] & n_mask
+    partial_sum += tl.load(base + base_offs, mask=out_mask, other=0.0)
+    tl.store(base + base_offs, partial_sum, mask=out_mask)
+
+
+def step_b_v_fwd(
+    attn_lora_a: torch.Tensor,
+    B_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor,
+    qk_nope_head_dim: int,
+    v_head_dim: int,
+) -> torch.Tensor:
+    """Step B of the v-side correction, accumulating into ``base_output``.
+
+    Args:
+        attn_lora_a: ``(S, H, rank)`` from step A_v.
+        B_buf: ``(num_lora, H*(qk_nope+v_head_dim), rank)``.
+        batch_info: standard ``LoRABatchInfo``.
+        base_output: ``(S, H, v_head_dim)``, modified in-place
+            (the absorbed ``attn_output @ w_vc`` result).
+        qk_nope_head_dim: offset of V-half within each head's row block of B.
+        v_head_dim: output feature dim per head.
+
+    Returns:
+        ``base_output`` (same object, mutated).
+    """
+    S, H, rank = attn_lora_a.shape
+    full_K_per_head = qk_nope_head_dim + v_head_dim
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S)
+        * triton.cdiv(v_head_dim, _STEP_B_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+
+    _step_b_v_kernel[grid](
+        attn_lora_a,
+        B_buf,
+        base_output,
+        S,
+        rank,
+        v_head_dim,
+        attn_lora_a.stride(0),
+        attn_lora_a.stride(1),
+        attn_lora_a.stride(2),
+        B_buf.stride(0),
+        B_buf.stride(1),
+        B_buf.stride(2),
+        base_output.stride(0),
+        base_output.stride(1),
+        base_output.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        batch_info.scalings,
+        segment_grid,
+        FULL_K=full_K_per_head,
+        QK_NOPE_OFFSET=qk_nope_head_dim,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_B_BLOCK_N,
+        BLOCK_K=_STEP_B_BLOCK_K,
+    )
+    return base_output

python/sglang/kernels/ops/gemm/lora_tuning_config.py

@@ -0,0 +1,201 @@
+"""
+Configuration loader for auto-tuned LoRA CSGMV kernel block sizes.
+
+Follows the same pattern as fused_moe_triton_config.py:
+- Offline tuning script writes JSON files keyed by chunk_size (BLOCK_M)
+- At server startup, the config loader reads the best block sizes for each kernel
+- Kernels use these instead of hardcoded defaults
+
+Config file naming: lora_{kernel},K={K},R={R},S={S},device={device}.json
+Where kernel is "shrink" or "expand", K is input_dim, R is max_rank, S is num_slices.
+
+Config file format (keyed by chunk_size):
+{
+    "16": {"BLOCK_N": 16, "BLOCK_K": 256, "num_warps": 4, "num_stages": 3},
+    "32": {"BLOCK_N": 32, "BLOCK_K": 128, "num_warps": 4, "num_stages": 4},
+    "128": {"BLOCK_N": 64, "BLOCK_K": 256, "num_warps": 8, "num_stages": 3}
+}
+
+Usage:
+    python3 benchmark/kernels/lora_csgmv/tune_lora_csgmv.py \
+        --model Qwen/Qwen3-Embedding-0.6B --max-lora-rank 64
+
+    # Configs saved to python/sglang/kernels/ops/gemm/configs/
+
+    # Server automatically picks them up:
+    python3 -m sglang.launch_server --model ... --enable-lora --lora-backend csgmv
+"""
+
+from __future__ import annotations
+
+import functools
+import json
+import logging
+import os
+from typing import Any, Dict, Optional
+
+import triton
+
+from sglang.srt.utils import get_device_name
+
+logger = logging.getLogger(__name__)
+
+
+def get_lora_config_file_name(
+    kernel: str,
+    K: int,
+    R: int,
+    S: int,
+) -> str:
+    """Generate config filename for a LoRA kernel configuration.
+
+    Args:
+        kernel: "shrink" or "expand"
+        K: The large dimension (input_dim for shrink, output_dim for expand)
+        R: The max LoRA rank
+        S: num_slices (qkv=3, gate_up=2, others=1)
+    """
+    device_name = get_device_name().replace(" ", "_")
+    return f"lora_{kernel},K={K},R={R},S={S},device={device_name}.json"
+
+
+@functools.lru_cache
+def get_lora_configs(
+    kernel: str,
+    K: int,
+    R: int,
+    S: int,
+) -> Optional[Dict[int, Dict[str, Any]]]:
+    """Load pre-tuned LoRA kernel configs from JSON files.
+
+    Returns a dict mapping chunk_size (BLOCK_M) to block size configs,
+    or None if no config file is found.
+    """
+    json_file_name = get_lora_config_file_name(kernel, K, R, S)
+
+    config_dir = os.environ.get(
+        "SGLANG_LORA_CONFIG_DIR", os.path.dirname(os.path.realpath(__file__))
+    )
+    configs_root = os.path.join(config_dir, "csgmv_configs")
+
+    triton_version = triton.__version__
+    version_dir = f"triton_{triton_version.replace('.', '_')}"
+
+    # Try exact triton version first
+    config_file_path = os.path.join(configs_root, version_dir, json_file_name)
+    if os.path.exists(config_file_path):
+        with open(config_file_path) as f:
+            logger.info(f"Using LoRA {kernel} config from {config_file_path}.")
+            return {int(key): val for key, val in json.load(f).items()}
+
+    # Scan existing version directories as fallback (newest first)
+    if os.path.isdir(configs_root):
+        version_dirs = sorted(
+            (d for d in os.listdir(configs_root) if d.startswith("triton_")),
+            reverse=True,
+        )
+        for vdir in version_dirs:
+            if vdir == version_dir:
+                continue
+            try_path = os.path.join(configs_root, vdir, json_file_name)
+            if os.path.exists(try_path):
+                with open(try_path) as f:
+                    logger.warning(
+                        f"LoRA {kernel} config not found for Triton {triton_version}. "
+                        f"Falling back to {try_path}."
+                    )
+                    return {int(key): val for key, val in json.load(f).items()}
+
+    return None
+
+
+# Default block sizes (current hardcoded values)
+DEFAULT_SHRINK_CONFIG = {"BLOCK_N": 16, "BLOCK_K": 256}
+DEFAULT_EXPAND_CONFIG = {"BLOCK_N": 64, "BLOCK_K": 16}
+
+# Track which configs have been logged to avoid spamming on every forward pass
+_logged_configs: set = set()
+
+
+def get_lora_shrink_config(
+    K: int,
+    R: int,
+    num_slices: int,
+    chunk_size: int,
+) -> Dict[str, int]:
+    """Get block sizes for the CSGMV shrink (lora_a) kernel.
+
+    Args:
+        K: input_dim
+        R: max_rank
+        num_slices: number of slices (qkv=3, gate_up=2, others=1)
+        chunk_size: BLOCK_M value (= batch_info.max_len)
+    """
+    log_key = ("shrink", K, R, num_slices, chunk_size)
+    configs = get_lora_configs("shrink", K, R, num_slices)
+    if configs is not None:
+        config = configs.get(chunk_size)
+        if config is None:
+            closest = min(configs.keys(), key=lambda x: abs(x - chunk_size))
+            config = configs[closest]
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA shrink (K={K}, R={R}): no config for chunk_size={chunk_size}, "
+                    f"using closest={closest}: {config}"
+                )
+        else:
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA shrink (K={K}, R={R}, chunk_size={chunk_size}): tuned config {config}"
+                )
+        return config
+    if log_key not in _logged_configs:
+        _logged_configs.add(log_key)
+        logger.info(
+            f"LoRA shrink (K={K}, R={R}): no tuned config, using defaults {DEFAULT_SHRINK_CONFIG}"
+        )
+    return dict(DEFAULT_SHRINK_CONFIG)
+
+
+def get_lora_expand_config(
+    K: int,
+    R: int,
+    num_slices: int,
+    chunk_size: int,
+) -> Dict[str, int]:
+    """Get block sizes for the CSGMV expand (lora_b) kernel.
+
+    Args:
+        K: output_dim
+        R: max_rank
+        num_slices: number of slices (qkv=3, gate_up=2, others=1)
+        chunk_size: BLOCK_M value (= batch_info.max_len)
+    """
+    log_key = ("expand", K, R, num_slices, chunk_size)
+    configs = get_lora_configs("expand", K, R, num_slices)
+    if configs is not None:
+        config = configs.get(chunk_size)
+        if config is None:
+            closest = min(configs.keys(), key=lambda x: abs(x - chunk_size))
+            config = configs[closest]
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA expand (K={K}, R={R}): no config for chunk_size={chunk_size}, "
+                    f"using closest={closest}: {config}"
+                )
+        else:
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA expand (K={K}, R={R}, chunk_size={chunk_size}): tuned config {config}"
+                )
+        return config
+    if log_key not in _logged_configs:
+        _logged_configs.add(log_key)
+        logger.info(
+            f"LoRA expand (K={K}, R={R}): no tuned config, using defaults {DEFAULT_EXPAND_CONFIG}"
+        )
+    return dict(DEFAULT_EXPAND_CONFIG)

python/sglang/kernels/ops/gemm/qkv_lora_b.py

@@ -0,0 +1,216 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _qkv_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Parameters of size
+    K,  # K = R
+    max_qkv_out_dim,  # max(output_q_dim, output_kv_dim)
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    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,
+    # For fused output scaling
+    scalings,
+):
+    """
+    This kernel packs 3 sgemms (q/k/v) into a single kernel. The multiplication
+    results are accumulated into the output tensor.
+
+    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, 3 * K), where s is the sum of all sequence lengths in the
+            batch and K is the maximum LoRA rank. The second dimension is partitioned
+            for Q, K, and V.
+        weights (Tensor): The LoRA B weights for all adapters.
+            Shape: (num_lora, N_Q + 2 * N_KV, K).
+        output (Tensor): The output tensor where the result is stored.
+            Shape: (s, N_Q + 2 * N_KV).
+    """
+
+    # 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)
+    batch_id = tl.program_id(axis=2)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    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
+    scaling = tl.load(scalings + w_index)
+    # Adjust K (rank) according to the specific LoRA adapter
+    K = tl.minimum(K, rank)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    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
+    # and accumulate
+    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)
+
+    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
+    )
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (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, :] < n_size),
+            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
+        + 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)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def qkv_lora_b_fwd(
+    x: torch.Tensor,
+    qkv_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_offset: torch.Tensor,
+    max_qkv_out_dim: int,
+    base_output: torch.Tensor = None,
+    n_slices: int = 3,
+) -> torch.Tensor:
+
+    # x: (s, n_slices * r)
+    # qkv_lora_b: (num_lora, output_dim_q + 2 * output_dim_kv, r)
+    # output_offset = [0, output_dim_q, output_dim_q + output_dim_kv,
+    #                     output_dim_q + 2 * output_dim_kv]  (length n_slices + 1)
+    # max_qkv_out_dim = max(output_dim_q, output_dim_kv)
+    # output: (s, output_dim_q + 2 * output_dim_kv)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # lora_output[:, :output_dim_q] = sgemm(x[:, :r], qkv_lora_b[:, :outptu_dim_q, :])
+    # lora_output[:, output_dim_q: output_dim_q + output_dim_kv]
+    #      = sgemm(x[:, r: 2 * r], qkv_lora_b[:, outptu_dim_q: output_dim_q + output_dim_kv, :])
+    # lora_output[:, output_dim_q + output_dim_kv: ]
+    #      = sgemm(x[:, 2 * r: , qkv_lora_b[:, output_dim_q + output_dim_kv: , :])
+
+    # Get dims
+    s = x.shape[0]
+    input_dim = x.shape[1]
+    r = qkv_lora_b.shape[-1]
+    output_dim = qkv_lora_b.shape[-2]
+    assert input_dim == n_slices * r
+    assert output_offset.shape[0] == n_slices + 1
+
+    BLOCK_S = 16
+    BLOCK_R = 16
+    BLOCK_OUT = 64
+
+    grid_b = (
+        triton.cdiv(batch_info.max_len, BLOCK_S)
+        * triton.cdiv(max_qkv_out_dim, BLOCK_OUT),
+        n_slices,
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((s, output_dim), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    sorted_by_adapter = batch_info.permutation is not None
+    _qkv_lora_b_kernel[grid_b](
+        x,
+        qkv_lora_b,
+        output,
+        r,
+        max_qkv_out_dim,
+        x.stride(0),
+        x.stride(1),
+        qkv_lora_b.stride(0),
+        qkv_lora_b.stride(1),
+        qkv_lora_b.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        output_offset,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_OUT,
+        BLOCK_R,
+        batch_info.scalings,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/sgemm_lora_a.py

@@ -0,0 +1,182 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _sgemm_lora_a_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Matrix dimensions
+    N,  # stack_num * r
+    K,  # input_dim
+    stack_num,
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths,ranks and weight id
+    seg_lens,
+    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,
+):
+    """
+    Computes a segmented batched matrix multiplication for the LoRA A matrix.
+
+    The kernel ensures that output[seg_start:seg_start + seg_len, :rank * stack_num]
+    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)`.
+        output (torch.Tensor): The output tensor of shape `(s, N)`.
+    """
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len
+    batch_id = tl.program_id(axis=1)
+    w_index = tl.load(weight_indices + batch_id)
+    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
+
+    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)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    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
+    # and accumulate
+    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)
+    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
+    )
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (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, :] < 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_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)
+
+
+def sgemm_lora_a_fwd(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    stack_num: int = 1,
+) -> torch.Tensor:
+    # x: (s, input_dim)
+    # weights: (num_lora, stack_num * r, input_dim)
+    # output: (s, stack_num * r)
+    # stack_num: run_qkv_lora: 3, run_gate_up_lora: 2
+    # when called by run_qkv_lora, the weights.shape[-2] will be 3 * 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
+
+    S = x.shape[0]
+    R = weights.shape[-2]
+    K = weights.shape[-1]
+    assert x.shape[-1] == K
+
+    # Block shapes
+    BLOCK_S = 16
+    BLOCK_K = 256
+    BLOCK_R = 16
+
+    grid = (
+        triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(R, BLOCK_R),
+        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,
+        weights,
+        output,
+        R,
+        K,
+        stack_num,
+        x.stride(0),
+        x.stride(1),
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_R,
+        BLOCK_K,
+    )
+    return output

python/sglang/kernels/ops/gemm/sgemm_lora_b.py

@@ -0,0 +1,188 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _sgemm_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Matrix dimensions
+    N,  # output_dim
+    K,  # r
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths and weight id
+    seg_lens,
+    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,
+    # For fused output scaling
+    scalings,
+):
+    """
+    Computes a segmented batched matrix multiplication for the LoRA B matrix
+    and adds the result to the output in-place.
+
+    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 (torch.Tensor): The intermediate tensor from the LoRA 'A' multiplication,
+            of shape `(s, K)`, where `s` is the total number of tokens.
+        weights (torch.Tensor): The LoRA 'B' weights for all available adapters,
+            with shape `(num_lora, N, K)`.
+        output (torch.Tensor): The output tensor of shape `(s, N)`. This can be
+            the base model's output for a fused add operation.
+    """
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len
+    batch_id = tl.program_id(axis=1)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    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
+    K = tl.minimum(K, rank)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    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
+    # and accumulate
+    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)
+    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
+    )
+
+    # Iterate to compute the block in output matrix
+    n_mask = n_offset[None, :] < N
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (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_mask,
+            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_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)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def sgemm_lora_b_fwd(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor = None,
+) -> torch.Tensor:
+    # x: (s, max_r)
+    # weights: (num_lora, output_dim, max_r)
+    # output: (s, output_dim)
+    # output_dim is much larger than max_r
+
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(x.shape) == 2
+    assert len(weights.shape) == 3
+
+    S = x.shape[0]
+    N = weights.shape[-2]
+    R = weights.shape[-1]
+    assert x.shape[-1] == R
+
+    # Block shapes
+    BLOCK_S = 16
+    BLOCK_R = 16
+    BLOCK_N = 256
+
+    grid = (
+        triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(N, BLOCK_N),
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((S, N), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    sorted_by_adapter = batch_info.permutation is not None
+    _sgemm_lora_b_kernel[grid](
+        x,
+        weights,
+        output,
+        N,
+        R,
+        x.stride(0),
+        x.stride(1),
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_N,
+        BLOCK_R,
+        batch_info.scalings,
+    )
+    return output

python/sglang/kernels/ops/gemm/trtllm_lora_temp/__init__.py

@@ -0,0 +1,3 @@
+"""Experimental TRT-LLM LoRA kernel variants (gated by ``SGLANG_EXPERIMENTAL_LORA_OPTI`` / ``lora_envs``).
+
+Migrated from ``sglang.srt.lora.trtllm_lora_temp.triton_ops`` (RFC #29630)."""

python/sglang/kernels/ops/gemm/trtllm_lora_temp/gate_up_lora_b.py

@@ -0,0 +1,266 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
+    _resolve_token_positions,
+    get_pdl_launch_metadata,
+)
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
+from sglang.srt.lora.utils import LoRABatchInfo
+
+# Minimum total_tokens * rank for the single-adapter cuBLAS path; below this
+# the Triton kernel is faster (crossover measured at output_dim=1536/GPU:
+# cuBLAS wins rank64 from S>=256 and rank16 only from S>=2048).
+_CUBLAS_MIN_S_RANK = 16384
+
+
+@triton.jit
+def _gate_up_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Parameters of size
+    K,  # K = R
+    output_dim,
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths,ranks and weight id
+    seg_lens,
+    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,
+    # For fused output scaling
+    scalings,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    """
+    This kernel packs 2 sgemms (gate/up) into a single kernel. The multiplication
+    results are accumulated into the output tensor.
+
+    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, 2 * 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, 2 * output_dim, K).
+        output (Tensor): The output tensor where the result is stored.
+            Shape: (s, 2 * output_dim).
+    """
+    # output_dim >> K
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len.
+    # gate_up_id decides which of gate or up (0: gate, 1: up)
+    batch_id = tl.program_id(axis=2)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    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)
+
+    # Adjust K (rank) according to the specific LoRA adapter
+    K = tl.minimum(K, rank)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    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
+    # and accumulate
+    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)
+
+    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
+    )
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (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, :] < output_dim),
+            other=0.0,
+        )
+        partial_sum += tl.dot(
+            x_tile.to(w_tile.dtype), w_tile
+        )  # cast fused: split-K returns fp32, plain path bf16 (no-op)
+
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # Store result to output matrix
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    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)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def _gate_up_lora_b_cublas(
+    x: torch.Tensor,
+    gate_up_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_dim: int,
+    base_output: torch.Tensor,
+) -> torch.Tensor:
+    """Single-adapter dense path: one cuBLAS addmm_ per gate/up slice.
+
+    The LoRA-A output is rank-packed (slice i at columns [i*rank, (i+1)*rank)),
+    matching the Triton kernel's K = min(K, rank) slice stride. Slices are
+    disjoint output regions, so in-place addmm_ writes never collide.
+    """
+    r = gate_up_lora_b.shape[-1]
+    if base_output is None:
+        base_output = torch.zeros(
+            (x.shape[0], 2 * output_dim), device=x.device, dtype=x.dtype
+        )
+    w = gate_up_lora_b[0]
+    x_scaled = x[:, : 2 * r] * batch_info.scalings[0]
+    for i in range(2):
+        lo, hi = i * output_dim, (i + 1) * output_dim
+        base_output[:, lo:hi].addmm_(x_scaled[:, i * r : (i + 1) * r], w[lo:hi, :r].t())
+    return base_output
+
+
+def gate_up_lora_b_fwd(
+    x: torch.Tensor,
+    gate_up_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_dim: int,
+    base_output: torch.Tensor = None,
+) -> torch.Tensor:
+
+    # x: (s, 2 * r)
+    # gate_up_lora_b: (num_lora, 2 * output_dim, r)
+    # output: (s, 2 * output_dim)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # lora_output[:, :output_dim] = sgemm(x[:, :r], gate_up_lora_b[:, :output_dim, :])
+    # lora_output[:, output_dim:]
+    #      = sgemm(x[:, r:], gate_up_lora_b[:, output_dim:, :])
+
+    # Get dims
+    s = x.shape[0]
+    input_dim = x.shape[1]
+    r = gate_up_lora_b.shape[-1]
+    assert input_dim == 2 * r
+
+    if (
+        (
+            lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+            or lora_envs.SGLANG_OPT_LORA_CUBLAS_GATE_UP.get()
+        )
+        and s * r >= _CUBLAS_MIN_S_RANK
+        and gate_up_lora_b.shape[0] == 1
+    ):  # single-adapter fast path: only valid with one resident slot
+        return _gate_up_lora_b_cublas(
+            x, gate_up_lora_b, batch_info, output_dim, base_output
+        )
+
+    BLOCK_S = 16
+    BLOCK_R = 16
+    BLOCK_OUT = 64
+
+    grid_b = (
+        triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(output_dim, BLOCK_OUT),
+        2,  # this dimension decides current block computes on gate or up proj
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((s, 2 * output_dim), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+    _gate_up_lora_b_kernel[grid_b](
+        x,
+        gate_up_lora_b,
+        output,
+        r,
+        output_dim,
+        x.stride(0),
+        x.stride(1),
+        gate_up_lora_b.stride(0),
+        gate_up_lora_b.stride(1),
+        gate_up_lora_b.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_OUT,
+        BLOCK_R,
+        batch_info.scalings,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/trtllm_lora_temp/kernel_utils.py

@@ -0,0 +1,31 @@
+import triton
+import triton.language as tl
+
+from sglang.jit_kernel.utils import is_arch_support_pdl
+
+
+def get_pdl_launch_metadata() -> tuple[bool, dict]:
+    """Return (ENABLE_PDL constexpr value, extra launch kwargs) for LoRA kernels.
+
+    ``launch_pdl`` is NVIDIA-only Triton launch metadata; the HIP backend
+    rejects unknown kwargs, so it is only included when PDL is supported.
+    """
+    enable_pdl = is_arch_support_pdl()
+    return enable_pdl, ({"launch_pdl": True} if enable_pdl else {})
+
+
+@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/kernels/ops/gemm/trtllm_lora_temp/kv_b_lora_absorbed.py

@@ -0,0 +1,965 @@
+"""Triton kernels for absorbed-MLA ``kv_b_proj`` LoRA correction.
+
+The absorbed-MLA path bypasses ``kv_b_proj.forward()`` and folds the K/V
+sides as plain BMMs ``q_nope @ w_kc`` and ``attn_output @ w_vc``.  When a
+LoRA adapter is active on ``kv_b_proj`` we add the LoRA delta to
+``q_nope_out`` / ``attn_bmm_output`` manually.
+
+Using the standard LoRA factored math we *never* materialize ``B @ A``:
+
+    q_correction = q_nope     @ B_kc @ A   * scaling          # K-side
+    v_correction = attn_output @ A.T @ B_vc.T * scaling       # V-side
+
+where ``A: (slot, rank, kv_lora_rank)`` is the LoRA-A of ``kv_b_proj``
+(shared across heads) and ``B: (slot, num_heads*(qk_nope+v_head_dim), rank)``
+is the LoRA-B; ``B_kc`` / ``B_vc`` are its K-half / V-half slices.
+
+Four kernels split the math along the factorization boundary, all using
+the SGMM idiom from ``sgemm_lora_a`` / ``qkv_lora_b`` and the segment-indptr
+routing used by ``chunked_sgmv_*``:
+
+  * ``step_a_q_fwd``: per-head per-slot SGMM, ``(S,H,qk_nope) -> (S,H,rank)``
+  * ``step_b_q_fwd``: shared-A per-slot SGMM, scaled+accumulated,
+    ``(S,H,rank) -> (S,H,kv_lora_rank)``
+  * ``step_a_v_fwd``: shared-A.T per-slot SGMM, ``(S,H,kv_lora_rank) -> (S,H,rank)``
+  * ``step_b_v_fwd``: per-head per-slot SGMM with V-half of B, transposed,
+    scaled+accumulated, ``(S,H,rank) -> (S,H,v_head_dim)``
+
+Grid axes for each kernel:
+  axis 0 : output tile in (S, N)         -- tile_id = pid_s * num_pid_n + pid_n
+  axis 1 : head_id                       -- per-head weight slice
+  axis 2 : batch_id (segment / request)  -- per-slot weight routing via weight_indices
+
+Per-segment routing: each program derives its segment length from
+``seg_indptr[segment_id + 1] - seg_indptr[segment_id]``, loads
+``weight_indices[segment_id]`` once, and uses that slot's slice of the LoRA
+weight stack.  When ``permutation`` is present, rows are routed through it,
+matching the csgmv backend's adapter-grouped chunks.  No Python loops over slots
+or heads.
+
+The math also stays in the input dtype (no fp32 round-trip) -- the
+contraction dim ``rank`` is small (typically 16-64), so bf16 accumulation
+over it is acceptable.  ``tl.dot`` itself uses fp32 accumulation internally.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
+    _resolve_token_positions,
+    get_pdl_launch_metadata,
+)
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
+from sglang.srt.lora.utils import LoRABatchInfo
+
+# ---------------------------------------------------------------------------
+# Block sizes -- chosen per-kernel from the natural shape of each step.
+#
+# The factored math gives the four kernels these contraction (K) and output
+# (N) ranges (for Kimi-K2.5: rank=16-32, qk_nope=v_head_dim=128, kv_lora_rank=512):
+#
+#                              K (contraction)      N (output)
+#   step_a_q                   qk_nope (~128)       rank (~16-32)
+#   step_b_q                   rank (~16-32)        kv_lora_rank (~512)
+#   step_a_v                   kv_lora_rank (~512)  rank (~16-32)
+#   step_b_v                   rank (~16-32)        v_head_dim (~128)
+#
+# So the "step_a_*" kernels want a large BLOCK_K (to keep loop iters small)
+# and a small BLOCK_N (matched to rank to avoid wasted tile lanes), while
+# the "step_b_*" kernels are the inverse.  Kernels aren't autotuned -- the
+# decode-shape workload is too small to benefit and the sweep surface is
+# wide.
+# ---------------------------------------------------------------------------
+
+_BLOCK_S = 16
+
+
+def _num_segments(batch_info: LoRABatchInfo) -> int:
+    return batch_info.num_segments or batch_info.bs
+
+
+def _max_segment_len(batch_info: LoRABatchInfo) -> int:
+    if batch_info.max_len is not None:
+        return batch_info.max_len
+    if batch_info.seg_lens is not None:
+        return int(batch_info.seg_lens.max().item())
+    raise ValueError("LoRA batch_info must provide max_len or seg_lens.")
+
+
+def _segment_grid_size(batch_info: LoRABatchInfo, num_segments: int) -> int:
+    return (
+        batch_info.weight_indices.shape[0]
+        if batch_info.use_cuda_graph
+        else num_segments
+    )
+
+
+# ---------------------------------------------------------------------------
+# Kernel 1 -- Step A_q: per-head per-slot SGMM, reads K-half of B
+#
+#     q_lora_a[t, h, r] = sum_{i<qk_nope} q_nope[t, h, i] * B[slot, h*FULL_K + i, r]
+#
+# x      : (S, H, qk_nope)
+# w (B)  : (num_lora, H*FULL_K, rank)   -- FULL_K = qk_nope + v_head_dim
+# out    : (S, H, rank)                 -- fresh allocation, no accumulate
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_a_q_kernel(
+    x,
+    w,
+    out,
+    # dims
+    S,
+    H_FULL_K,  # H * (qk_nope + v_head_dim), the row-stride landmark
+    K,  # qk_nope (contraction)
+    N,  # rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,
+    w_stride_k,
+    out_stride_s,
+    out_stride_h,
+    out_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    num_segments,
+    # meta
+    FULL_K: tl.constexpr,  # per-head row stride in B (qk_nope + v_head_dim)
+    SORTED_BY_ADAPTER: tl.constexpr,
+    K_DIV: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+
+    # Truncate output N to this slot's rank (allows mixed-rank batches).
+    N_eff = tl.minimum(N, cur_rank)
+
+    num_pid_n = tl.cdiv(N_eff, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    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)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    # Clamp masked-lane indices into the valid range so pointer arithmetic
+    # stays in-bounds even before the load mask drops the values.
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    safe_n = tl.minimum(n_offset, N_eff - 1)
+
+    head_row_base = (
+        head_id * FULL_K
+    )  # row offset for this head's K-half (i in [0, qk_nope))
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K
+        safe_k = cur_k if K_DIV else tl.minimum(cur_k, K - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # B[slot, h*FULL_K + i, r]: row dim of B carries i (= GEMM K),
+        # column dim carries r (= GEMM N).
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + (head_row_base + safe_k[:, None]) * w_stride_n
+            + safe_n[None, :] * w_stride_k,
+            mask=k_mask[:, None] & (n_offset[None, :] < N_eff),
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    out_offs = (
+        safe_row[:, None] * out_stride_s
+        + head_id * out_stride_h
+        + safe_n[None, :] * out_stride_n
+    )
+    out_mask = row_mask[:, None] & (n_offset[None, :] < N_eff)
+    tl.store(out + out_offs, partial_sum, mask=out_mask)
+
+
+def step_a_q_fwd(
+    q_nope: torch.Tensor,
+    B_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    full_K_per_head: int,
+) -> torch.Tensor:
+    """Step A of the q-side correction.
+
+    Args:
+        q_nope: ``(S, H, qk_nope)``, the absorbed-MLA q intermediate.
+        B_buf: ``(num_lora, H*full_K_per_head, rank)`` from the LoRA pool.
+        batch_info: standard ``LoRABatchInfo``.
+        full_K_per_head: ``qk_nope + v_head_dim``, the row stride per head in B.
+
+    Returns:
+        ``(S, H, rank)`` -- per-token, per-head low-rank intermediate, ready for step B_q.
+    """
+    S, H, qk_nope_dim = q_nope.shape
+    rank = B_buf.shape[-1]
+
+    if (
+        lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+        or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
+    ) and B_buf.shape[
+        0
+    ] == 1:  # single-adapter fast path: only valid with one resident slot
+        # (S,H,r) view of a (H,S,r)-contiguous bmm result; step_b_q's dense
+        # path flattens in (h,s) order, so the chain needs no copies.
+        w_kc = B_buf[0].view(H, full_K_per_head, -1)[:, :qk_nope_dim, :]
+        return torch.bmm(q_nope.transpose(0, 1), w_kc).transpose(0, 1)
+
+    block_n = triton.next_power_of_2(rank)  # output N == rank -> one tile
+    out = torch.empty((S, H, rank), device=q_nope.device, dtype=q_nope.dtype)
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+    _STEP_A_Q_BLOCK_K = 128
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S) * triton.cdiv(rank, block_n),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+
+    _step_a_q_kernel[grid](
+        q_nope,
+        B_buf,
+        out,
+        S,
+        H * full_K_per_head,
+        qk_nope_dim,
+        rank,
+        q_nope.stride(0),
+        q_nope.stride(1),
+        q_nope.stride(2),
+        B_buf.stride(0),
+        B_buf.stride(1),
+        B_buf.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        segment_grid,
+        FULL_K=full_K_per_head,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        K_DIV=(qk_nope_dim % _STEP_A_Q_BLOCK_K == 0),
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=block_n,
+        BLOCK_K=_STEP_A_Q_BLOCK_K,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Kernel 2 -- Step B_q: shared-A per-slot SGMM, scaled + accumulated
+#
+#     base[t, h, k] += sum_r x[t, h, r] * A[slot, r, k] * scaling
+#
+# x      : (S, H, rank)
+# w (A)  : (num_lora, rank, kv_lora_rank)
+# base   : (S, H, kv_lora_rank), updated in-place (accumulated)
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_b_q_kernel(
+    x,
+    w,
+    base,
+    # dims
+    S,
+    K,  # rank (contraction)
+    N,  # kv_lora_rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_k,
+    w_stride_n,
+    b_stride_s,
+    b_stride_h,
+    b_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    scalings,
+    num_segments,
+    # meta
+    SORTED_BY_ADAPTER: tl.constexpr,
+    N_DIV: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+    scaling = tl.load(scalings + w_index)
+
+    # Truncate contraction K to this slot's rank.
+    K_eff = tl.minimum(K, cur_rank)
+
+    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
+
+    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)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    n_mask = n_offset[None, :] < N
+    safe_n = n_offset if N_DIV else tl.minimum(n_offset, N - 1)
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K_eff, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K_eff
+        safe_k = tl.minimum(cur_k, K_eff - 1)
+
+        # x[s, h, k]  (k iterates over rank)
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # A[slot, k, n]: read k along contraction, n along output.
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + safe_n[None, :] * w_stride_n,
+            mask=k_mask[:, None] & n_mask,
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+
+    # Accumulate into base[s, h, n].
+    base_offs = (
+        safe_row[:, None] * b_stride_s
+        + head_id * b_stride_h
+        + safe_n[None, :] * b_stride_n
+    )
+    out_mask = row_mask[:, None] & n_mask
+    tl.atomic_add(base + base_offs, partial_sum, mask=out_mask, sem="relaxed")
+
+
+def step_b_q_fwd(
+    q_lora_a: torch.Tensor,
+    A_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor,
+) -> torch.Tensor:
+    """Step B of the q-side correction, accumulating into ``base_output``.
+
+    Args:
+        q_lora_a: ``(S, H, rank)`` from step A_q.
+        A_buf: ``(num_lora, rank, kv_lora_rank)`` from the LoRA pool.
+        batch_info: standard ``LoRABatchInfo``.
+        base_output: ``(S, H, kv_lora_rank)``, modified in-place
+            (the absorbed ``q_nope @ w_kc`` result).
+
+    Returns:
+        ``base_output`` (same object, mutated).
+    """
+    S, H, rank = q_lora_a.shape
+    kv_lora_rank = A_buf.shape[-1]
+
+    if (
+        lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+        or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
+    ) and A_buf.shape[
+        0
+    ] == 1:  # single-adapter fast path: only valid with one resident slot
+        # Flatten (S,H) in whichever order base_output's storage allows
+        # without a copy (the absorbed q path passes a transpose view of a
+        # (H,S,kv)-contiguous bmm result). x is small; reshape may copy it.
+        base2d = x2d = None
+        if base_output.is_contiguous():
+            base2d = base_output.view(-1, kv_lora_rank)
+            x2d = q_lora_a[..., :rank].reshape(-1, rank)
+        elif base_output.transpose(0, 1).is_contiguous():
+            base2d = base_output.transpose(0, 1).view(-1, kv_lora_rank)
+            x2d = q_lora_a[..., :rank].transpose(0, 1).reshape(-1, rank)
+        if base2d is not None:
+            base2d.addmm_(x2d, A_buf[0, :, :], alpha=batch_info.scalings[0])
+            return base_output
+
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+    _STEP_B_Q_BLOCK_N = 128
+    _STEP_B_BLOCK_K = 16
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S)
+        * triton.cdiv(kv_lora_rank, _STEP_B_Q_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+
+    _step_b_q_kernel[grid](
+        q_lora_a,
+        A_buf,
+        base_output,
+        S,
+        rank,
+        kv_lora_rank,
+        q_lora_a.stride(0),
+        q_lora_a.stride(1),
+        q_lora_a.stride(2),
+        A_buf.stride(0),
+        A_buf.stride(1),
+        A_buf.stride(2),
+        base_output.stride(0),
+        base_output.stride(1),
+        base_output.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        batch_info.scalings,
+        segment_grid,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        N_DIV=(kv_lora_rank % _STEP_B_Q_BLOCK_N == 0),
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_B_Q_BLOCK_N,
+        BLOCK_K=_STEP_B_BLOCK_K,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+    return base_output
+
+
+# ---------------------------------------------------------------------------
+# Kernel 3 -- Step A_v: shared-A.T per-slot SGMM (no scaling, fresh output)
+#
+#     attn_lora_a[t, h, r] = sum_k attn_output[t, h, k] * A[slot, r, k]
+#
+# x      : (S, H, kv_lora_rank)
+# w (A)  : (num_lora, rank, kv_lora_rank) -- accessed transposed vs step B_q
+# out    : (S, H, rank), fresh allocation
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_a_v_kernel(
+    x,
+    w,
+    out,
+    # dims
+    S,
+    K,  # kv_lora_rank (contraction)
+    N,  # rank (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,  # A's "rank" axis (= GEMM N)
+    w_stride_k,  # A's "kv_lora_rank" axis (= GEMM K)
+    out_stride_s,
+    out_stride_h,
+    out_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    num_segments,
+    # meta
+    SORTED_BY_ADAPTER: tl.constexpr,
+    K_DIV: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+
+    # Truncate output N to this slot's rank.
+    N_eff = tl.minimum(N, cur_rank)
+
+    num_pid_n = tl.cdiv(N_eff, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    if pid_s * BLOCK_S >= seg_len:
+        return
+
+    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)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    safe_n = tl.minimum(n_offset, N_eff - 1)
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K
+        safe_k = cur_k if K_DIV else tl.minimum(cur_k, K - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # A[slot, r, k] -- here we want each (k, r) so we read along k
+        # (inner / contraction) and produce r as output.  Stride access:
+        # the row dim is r (= GEMM N), column dim is k (= GEMM K).
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + safe_n[None, :] * w_stride_n,
+            mask=k_mask[:, None] & (n_offset[None, :] < N_eff),
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    out_offs = (
+        safe_row[:, None] * out_stride_s
+        + head_id * out_stride_h
+        + safe_n[None, :] * out_stride_n
+    )
+    out_mask = row_mask[:, None] & (n_offset[None, :] < N_eff)
+    tl.store(out + out_offs, partial_sum, mask=out_mask)
+
+
+def step_a_v_fwd(
+    attn_output: torch.Tensor,
+    A_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+) -> torch.Tensor:
+    """Step A of the v-side correction.
+
+    Args:
+        attn_output: ``(S, H, kv_lora_rank)``, the post-attention intermediate.
+        A_buf: ``(num_lora, rank, kv_lora_rank)``.
+        batch_info: standard ``LoRABatchInfo``.
+
+    Returns:
+        ``(S, H, rank)`` -- per-token, per-head low-rank intermediate for step B_v.
+    """
+    S, H, kv_lora_rank = attn_output.shape
+    rank = A_buf.shape[1]
+
+    if (
+        (
+            lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+            or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
+        )
+        and attn_output.is_contiguous()
+        and A_buf.shape[0] == 1
+    ):  # single-adapter fast path: only valid with one resident slot
+        return torch.mm(
+            attn_output.view(-1, kv_lora_rank), A_buf[0, :rank, :].t()
+        ).view(S, H, rank)
+
+    block_n = triton.next_power_of_2(rank)
+    out = torch.empty((S, H, rank), device=attn_output.device, dtype=attn_output.dtype)
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+    _STEP_A_V_BLOCK_K = 256
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S) * triton.cdiv(rank, block_n),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+
+    _step_a_v_kernel[grid](
+        attn_output,
+        A_buf,
+        out,
+        S,
+        kv_lora_rank,
+        rank,
+        attn_output.stride(0),
+        attn_output.stride(1),
+        attn_output.stride(2),
+        A_buf.stride(0),
+        A_buf.stride(1),
+        A_buf.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        segment_grid,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        K_DIV=(kv_lora_rank % _STEP_A_V_BLOCK_K == 0),
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=block_n,
+        BLOCK_K=_STEP_A_V_BLOCK_K,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Kernel 4 -- Step B_v: per-head per-slot SGMM with V-half of B (transposed),
+# scaled + accumulated
+#
+#     base[t, h, j] += sum_r x[t, h, r] * B[slot, h*FULL_K + qk_nope + j, r] * scaling
+#
+# x      : (S, H, rank)
+# w (B)  : (num_lora, H*FULL_K, rank), V-half slice via offset
+# base   : (S, H, v_head_dim), updated in-place (accumulated)
+# ---------------------------------------------------------------------------
+
+
+@triton.jit(do_not_specialize=["num_segments"])
+def _step_b_v_kernel(
+    x,
+    w,
+    base,
+    # dims
+    S,
+    K,  # rank (contraction)
+    N,  # v_head_dim (output)
+    # strides
+    x_stride_s,
+    x_stride_h,
+    x_stride_k,
+    w_stride_l,
+    w_stride_n,  # B's row dim (h*FULL_K + j) -- this is GEMM N
+    w_stride_k,  # B's rank dim -- this is GEMM K
+    b_stride_s,
+    b_stride_h,
+    b_stride_n,
+    # batch info
+    seg_indptr,
+    weight_indices,
+    lora_ranks,
+    sorted_token_ids,
+    scalings,
+    num_segments,
+    # meta
+    FULL_K: tl.constexpr,  # qk_nope + v_head_dim
+    QK_NOPE_OFFSET: tl.constexpr,  # offset of V-half within each head's row block
+    SORTED_BY_ADAPTER: tl.constexpr,
+    N_DIV: tl.constexpr,  # N % BLOCK_N == 0 -> drop safe_n (keep the store coalesced)
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    batch_id = tl.program_id(axis=2)
+    head_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+
+    if batch_id >= num_segments:
+        return
+
+    w_index = tl.load(weight_indices + batch_id)
+    cur_rank = tl.load(lora_ranks + w_index)
+    if cur_rank == 0:
+        return
+
+    seg_start = tl.load(seg_indptr + batch_id)
+    seg_end = tl.load(seg_indptr + batch_id + 1)
+    seg_len = seg_end - seg_start
+    if seg_len == 0:
+        return
+    scaling = tl.load(scalings + w_index)
+
+    K_eff = tl.minimum(K, cur_rank)
+
+    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
+
+    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)
+
+    s_physical = _resolve_token_positions(
+        sorted_token_ids, seg_start, s_offset, seg_len, SORTED_BY_ADAPTER
+    )
+
+    row_mask = s_offset < seg_len
+    safe_row = tl.minimum(s_physical, S - 1)
+    n_mask = n_offset[None, :] < N
+    safe_n = n_offset if N_DIV else tl.minimum(n_offset, N - 1)
+
+    # V-half row base for this head: h*FULL_K + qk_nope
+    head_row_base = head_id * FULL_K + QK_NOPE_OFFSET
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k_block in range(0, tl.cdiv(K_eff, BLOCK_K)):
+        cur_k = k_block * BLOCK_K + k_offset
+        k_mask = cur_k < K_eff
+        safe_k = tl.minimum(cur_k, K_eff - 1)
+
+        # x[s, h, k]
+        x_tile = tl.load(
+            x
+            + safe_row[:, None] * x_stride_s
+            + head_id * x_stride_h
+            + safe_k[None, :] * x_stride_k,
+            mask=row_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # B[slot, h*FULL_K + qk_nope + j, r] -- row dim is j (= GEMM N),
+        # column dim is r (= GEMM K).  Transposed access vs step A_q.
+        w_tile = tl.load(
+            w
+            + w_index * w_stride_l
+            + safe_k[:, None] * w_stride_k
+            + (head_row_base + safe_n[None, :]) * w_stride_n,
+            mask=k_mask[:, None] & n_mask,
+            other=0.0,
+        )
+
+        partial_sum += tl.dot(x_tile, w_tile)
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+
+    base_offs = (
+        safe_row[:, None] * b_stride_s
+        + head_id * b_stride_h
+        + safe_n[None, :] * b_stride_n
+    )
+    out_mask = row_mask[:, None] & n_mask
+    tl.atomic_add(base + base_offs, partial_sum, mask=out_mask, sem="relaxed")
+
+
+def step_b_v_fwd(
+    attn_lora_a: torch.Tensor,
+    B_buf: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor,
+    qk_nope_head_dim: int,
+    v_head_dim: int,
+) -> torch.Tensor:
+    """Step B of the v-side correction, accumulating into ``base_output``.
+
+    Args:
+        attn_lora_a: ``(S, H, rank)`` from step A_v.
+        B_buf: ``(num_lora, H*(qk_nope+v_head_dim), rank)``.
+        batch_info: standard ``LoRABatchInfo``.
+        base_output: ``(S, H, v_head_dim)``, modified in-place
+            (the absorbed ``attn_output @ w_vc`` result).
+        qk_nope_head_dim: offset of V-half within each head's row block of B.
+        v_head_dim: output feature dim per head.
+
+    Returns:
+        ``base_output`` (same object, mutated).
+    """
+    S, H, rank = attn_lora_a.shape
+    full_K_per_head = qk_nope_head_dim + v_head_dim
+    num_segments = _num_segments(batch_info)
+    max_segment_len = _max_segment_len(batch_info)
+    segment_grid = _segment_grid_size(batch_info, num_segments)
+    _STEP_B_V_BLOCK_N = 64
+    _STEP_B_BLOCK_K = 16
+
+    grid = (
+        triton.cdiv(max_segment_len, _BLOCK_S)
+        * triton.cdiv(v_head_dim, _STEP_B_V_BLOCK_N),
+        H,
+        segment_grid,
+    )
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+
+    _step_b_v_kernel[grid](
+        attn_lora_a,
+        B_buf,
+        base_output,
+        S,
+        rank,
+        v_head_dim,
+        attn_lora_a.stride(0),
+        attn_lora_a.stride(1),
+        attn_lora_a.stride(2),
+        B_buf.stride(0),
+        B_buf.stride(1),
+        B_buf.stride(2),
+        base_output.stride(0),
+        base_output.stride(1),
+        base_output.stride(2),
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        batch_info.scalings,
+        segment_grid,
+        FULL_K=full_K_per_head,
+        QK_NOPE_OFFSET=qk_nope_head_dim,
+        SORTED_BY_ADAPTER=sorted_by_adapter,
+        N_DIV=(v_head_dim % _STEP_B_V_BLOCK_N == 0),
+        BLOCK_S=_BLOCK_S,
+        BLOCK_N=_STEP_B_V_BLOCK_N,
+        BLOCK_K=_STEP_B_BLOCK_K,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+    return base_output

python/sglang/kernels/ops/gemm/trtllm_lora_temp/qkv_lora_b.py

@@ -0,0 +1,297 @@
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
+    _resolve_token_positions,
+    get_pdl_launch_metadata,
+)
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
+from sglang.srt.lora.utils import LoRABatchInfo
+
+# Minimum max_len (longest segment) for the single-adapter cuBLAS path; below
+# this the Triton kernel is faster (measured crossover at the smallest
+# realistic N=768, where cuBLAS is weakest).
+_CUBLAS_MIN_MAX_LEN = 8
+
+
+@triton.jit
+def _qkv_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Parameters of size
+    K,  # K = R
+    max_qkv_out_dim,  # max(output_q_dim, output_kv_dim)
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    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,
+    # For fused output scaling
+    scalings,
+    ENABLE_PDL: tl.constexpr = False,
+    STORE_WRITEBACK: tl.constexpr = False,
+):
+    """
+    This kernel packs 3 sgemms (q/k/v) into a single kernel. The multiplication
+    results are accumulated into the output tensor.
+
+    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, 3 * K), where s is the sum of all sequence lengths in the
+            batch and K is the maximum LoRA rank. The second dimension is partitioned
+            for Q, K, and V.
+        weights (Tensor): The LoRA B weights for all adapters.
+            Shape: (num_lora, N_Q + 2 * N_KV, K).
+        output (Tensor): The output tensor where the result is stored.
+            Shape: (s, N_Q + 2 * N_KV).
+    """
+
+    # 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)
+    batch_id = tl.program_id(axis=2)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    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
+    scaling = tl.load(scalings + w_index)
+    # Adjust K (rank) according to the specific LoRA adapter.
+    K = tl.minimum(K, rank)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    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
+    # and accumulate
+    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)
+
+    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
+    )
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    n_mask = n_offset[None, :] < n_size
+    x_tile = tl.load(
+        x_ptrs,
+        mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K),
+        other=0.0,
+    )
+    w_tile = tl.load(
+        w_ptrs,
+        mask=(k_offset[:, None] < K) & n_mask,
+        other=0.0,
+    )
+    # cast fused: the split-K shrink returns fp32, plain path bf16 (no-op)
+    partial_sum = tl.dot(x_tile.to(w_tile.dtype), w_tile)
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # Store result to output matrix (cast to the OUTPUT dtype: x may be the fp32
+    # split-K shrink accumulator while base_output is bf16)
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(output.dtype.element_ty)
+    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)
+    if STORE_WRITEBACK:
+        # The expand-add output tiles are disjoint across all programs in this launch
+        # (distinct s-rows / n-cols / slice / segment), so each element is RMW'd by
+        # exactly one program -- a plain read-add-write is correct and avoids the bf16
+        # narrow-tile atomic (the dominant decode cost). base_output is a same-stream
+        # data dependency (base GEMM before apply_lora), not a concurrent writer.
+        partial_sum += tl.load(output_ptr, mask=output_mask, other=0.0)
+        tl.store(output_ptr, partial_sum, mask=output_mask)
+    else:
+        tl.atomic_add(output_ptr, partial_sum, mask=output_mask, sem="relaxed")
+
+
+def _qkv_lora_b_cublas(
+    x: torch.Tensor,
+    qkv_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_offset_cpu: torch.Tensor,
+    base_output: Optional[torch.Tensor],
+    n_slices: int,
+) -> torch.Tensor:
+    """Single-adapter dense path: one cuBLAS addmm_ per q/k/v slice.
+
+    The LoRA-A output is rank-packed (slice i at columns [i*rank, (i+1)*rank)),
+    matching the Triton kernel's K = min(K, rank) slice stride. Slice offsets
+    come from the pinned CPU copy (no GPU sync); slices are disjoint output
+    regions, so in-place addmm_ writes never collide.
+    """
+    r = qkv_lora_b.shape[-1]
+    if base_output is None:
+        base_output = torch.zeros(
+            (x.shape[0], qkv_lora_b.shape[-2]), device=x.device, dtype=x.dtype
+        )
+    w = qkv_lora_b[0]
+    x_scaled = x[:, : n_slices * r] * batch_info.scalings[0]
+    offsets = output_offset_cpu.tolist()
+    for i in range(n_slices):
+        lo, hi = offsets[i], offsets[i + 1]
+        base_output[:, lo:hi].addmm_(x_scaled[:, i * r : (i + 1) * r], w[lo:hi, :r].t())
+    return base_output
+
+
+def qkv_lora_b_fwd(
+    x: torch.Tensor,
+    qkv_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_offset: torch.Tensor,
+    max_qkv_out_dim: int,
+    base_output: torch.Tensor = None,
+    n_slices: int = 3,
+    output_offset_cpu: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+
+    # x: (s, n_slices * r)
+    # qkv_lora_b: (num_lora, output_dim_q + 2 * output_dim_kv, r)
+    # output_offset = [0, output_dim_q, output_dim_q + output_dim_kv,
+    #                     output_dim_q + 2 * output_dim_kv]  (length n_slices + 1)
+    # max_qkv_out_dim = max(output_dim_q, output_dim_kv)
+    # output: (s, output_dim_q + 2 * output_dim_kv)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # lora_output[:, :output_dim_q] = sgemm(x[:, :r], qkv_lora_b[:, :outptu_dim_q, :])
+    # lora_output[:, output_dim_q: output_dim_q + output_dim_kv]
+    #      = sgemm(x[:, r: 2 * r], qkv_lora_b[:, outptu_dim_q: output_dim_q + output_dim_kv, :])
+    # lora_output[:, output_dim_q + output_dim_kv: ]
+    #      = sgemm(x[:, 2 * r: , qkv_lora_b[:, output_dim_q + output_dim_kv: , :])
+
+    # Get dims
+    s = x.shape[0]
+    input_dim = x.shape[1]
+    r = qkv_lora_b.shape[-1]
+    output_dim = qkv_lora_b.shape[-2]
+    assert input_dim == n_slices * r
+    assert output_offset.shape[0] == n_slices + 1
+
+    if (
+        output_offset_cpu is not None
+        and (
+            lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+            or lora_envs.SGLANG_OPT_LORA_CUBLAS_QKV.get()
+        )
+        and batch_info.max_len >= _CUBLAS_MIN_MAX_LEN
+        and qkv_lora_b.shape[0]
+        == 1  # single-adapter fast path: only valid with one resident slot
+    ):
+        return _qkv_lora_b_cublas(
+            x, qkv_lora_b, batch_info, output_offset_cpu, base_output, n_slices
+        )
+
+    BLOCK_S = 16
+    BLOCK_R = triton.next_power_of_2(r)
+    # BLOCK_OUT stays 64: with the 1-adapter cuBLAS dispatch the Triton path
+    # only runs for decode-sized batches, where 128 halves the grid (96->48
+    # programs on Kimi r16 bs64) and slows the kernel ~60% (11.4->18.5us, B200).
+    # Re-swept for the store path on GB200: 32 vs 64 is within noise (one preset
+    # marginally each way), so the single value is kept for both writebacks.
+    BLOCK_OUT = 64
+
+    grid_b = (
+        triton.cdiv(batch_info.max_len, BLOCK_S)
+        * triton.cdiv(max_qkv_out_dim, BLOCK_OUT),
+        n_slices,
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((s, output_dim), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    sorted_by_adapter = batch_info.permutation is not None
+    store_writeback = lora_envs.SGLANG_OPT_LORA_QKV_B_STORE.get()
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+    _qkv_lora_b_kernel[grid_b](
+        x,
+        qkv_lora_b,
+        output,
+        r,
+        max_qkv_out_dim,
+        x.stride(0),
+        x.stride(1),
+        qkv_lora_b.stride(0),
+        qkv_lora_b.stride(1),
+        qkv_lora_b.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        output_offset,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_OUT,
+        BLOCK_R,
+        batch_info.scalings,
+        ENABLE_PDL=enable_pdl,
+        STORE_WRITEBACK=store_writeback,
+        **pdl_kwargs,
+    )
+
+    return output

python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_a.py

@@ -0,0 +1,269 @@
+import functools
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
+    _resolve_token_positions,
+    get_pdl_launch_metadata,
+)
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _sgemm_lora_a_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Matrix dimensions
+    N,  # stack_num * r
+    K,  # input_dim
+    stack_num,
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths,ranks and weight id
+    seg_lens,
+    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,
+    SPLIT_K: tl.constexpr = 1,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    """
+    Computes a segmented batched matrix multiplication for the LoRA A matrix.
+
+    The kernel ensures that output[seg_start:seg_start + seg_len, :rank * stack_num]
+    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)`.
+        output (torch.Tensor): The output tensor of shape `(s, N)`.
+    """
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len
+    batch_id = tl.program_id(axis=1)
+    w_index = tl.load(weight_indices + batch_id)
+    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
+
+    pid = tl.program_id(axis=0)
+    # Fold the split-K factor out of axis-0 (SPLIT_K == 1 -> pid_sk == 0, pid_tile == pid).
+    pid_sk = pid % SPLIT_K
+    pid_tile = pid // SPLIT_K
+
+    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) to this adapter's actual rank.
+    N = tl.minimum(N, rank * stack_num)
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    pid_s = pid_tile // num_pid_n
+    pid_n = pid_tile % 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
+    # and accumulate.
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K)
+    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
+    )
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    # Iterate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):
+        k_remaining = K - k * (BLOCK_K * SPLIT_K)
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < k_remaining),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < k_remaining) & (n_offset[None, :] < N),
+            other=0.0,
+        )
+        partial_sum += tl.dot(x_tile, w_tile)
+
+        x_ptrs += BLOCK_K * SPLIT_K * x_stride_1
+        w_ptrs += BLOCK_K * SPLIT_K * w_stride_2
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # Store result to output matrix
+    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
+    )
+    if SPLIT_K == 1:
+        tl.store(output_ptr, partial_sum.to(output.dtype.element_ty), mask=output_mask)
+    else:
+        tl.atomic_add(
+            output_ptr,
+            partial_sum.to(output.dtype.element_ty),
+            mask=output_mask,
+            sem="relaxed",
+        )
+
+
+@functools.lru_cache(maxsize=None)
+def _num_sms(device_index: int) -> int:
+    return torch.cuda.get_device_properties(device_index).multi_processor_count
+
+
+def sgemm_lora_a_fwd(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    stack_num: int = 1,
+    out_alloc_stream=None,
+) -> torch.Tensor:
+    # x: (s, input_dim)
+    # weights: (num_lora, stack_num * r, input_dim)
+    # output: (s, stack_num * r)
+    # stack_num: run_qkv_lora: 3, run_gate_up_lora: 2
+    # when called by run_qkv_lora, the weights.shape[-2] will be 3 * 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
+
+    S = x.shape[0]
+    R = weights.shape[-2]
+    K = weights.shape[-1]
+    assert x.shape[-1] == K
+
+    if (
+        lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+        or lora_envs.SGLANG_OPT_LORA_CUBLAS_A.get()
+    ) and weights.shape[
+        0
+    ] == 1:  # single-adapter fast path: only valid with one resident slot
+        # Honor out_alloc_stream like the Triton path below: under SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC
+        # the shrink output must be allocated on the MAIN (consumer) stream so the caching allocator
+        # frees/reuses it on the consumer's schedule (cuda-graph WAR). F.linear has no out=, so
+        # allocate explicitly and matmul into it.
+        if out_alloc_stream is not None:
+            with torch.cuda.stream(out_alloc_stream):
+                output = torch.empty((S, R), device=x.device, dtype=x.dtype)
+        else:
+            output = torch.empty((S, R), device=x.device, dtype=x.dtype)
+        return torch.matmul(x, weights[0].transpose(-2, -1), out=output)
+
+    # Block shapes
+    BLOCK_S = 16
+    BLOCK_K = 256
+    BLOCK_R = triton.next_power_of_2(R)
+
+    sorted_by_adapter = batch_info.permutation is not None
+
+    num_s_tiles = triton.cdiv(batch_info.max_len, BLOCK_S)
+    split_k = 1
+    if lora_envs.SGLANG_ENABLE_LORA_SHRINK_SPLIT_K.get() and x.is_cuda:
+        num_k_tiles = triton.cdiv(K, BLOCK_K)
+        base_grid = batch_info.bs * num_s_tiles
+        num_sms = _num_sms(x.device.index)
+        if base_grid < num_sms and num_k_tiles >= 16:
+            split_k = max(1, min(2 * num_sms // base_grid, num_k_tiles, 16))
+
+    launch_kwargs = {}
+    if split_k > 1:
+        # out_alloc_stream (SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC) is intentionally NOT honored here:
+        # torch.zeros launches its memset on the alloc stream, which would race the side-stream
+        # shrink without extra ordering. No current config exercises split-K together with the
+        # two-stream main-alloc overlap (qwen3.5 leaves split-K off; kimi is single-stream-coherent).
+        output = torch.zeros((S, R), device=x.device, dtype=torch.float32)
+        launch_kwargs = {
+            "num_warps": 2 if split_k <= 4 else 4,
+            "num_stages": 3,
+        }
+    elif out_alloc_stream is not None:
+        # Allocate the output on the MAIN (consumer) stream when requested, so the caching allocator
+        # frees/reuses it on the consumer's schedule, not the side stream's (cuda-graph WAR — see
+        # lora_overlap_alloc_stream / SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC).
+        with torch.cuda.stream(out_alloc_stream):
+            output = torch.empty((S, R), device=x.device, dtype=x.dtype)
+    else:
+        output = torch.empty((S, R), device=x.device, dtype=x.dtype)
+
+    grid = (
+        num_s_tiles * split_k,
+        batch_info.bs,
+    )
+
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+    _sgemm_lora_a_kernel[grid](
+        x,
+        weights,
+        output,
+        R,
+        K,
+        stack_num,
+        x.stride(0),
+        x.stride(1),
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_R,
+        BLOCK_K,
+        split_k,
+        ENABLE_PDL=enable_pdl,
+        **launch_kwargs,
+        **pdl_kwargs,
+    )
+    # split_k>1 returns the fp32 accumulator directly; the LoRA-B expand casts x to the weight dtype
+    # on-load (fused), dropping the standalone fp32->bf16 copy kernel. split_k==1 already returns x.dtype.
+    return output

python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_b.py

@@ -0,0 +1,223 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.gemm.trtllm_lora_temp.gate_up_lora_b import (
+    _CUBLAS_MIN_S_RANK,
+)
+from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
+    _resolve_token_positions,
+    get_pdl_launch_metadata,
+)
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+def _sgemm_lora_b_cublas(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor,
+) -> torch.Tensor:
+    """Single-adapter dense path: one cuBLAS addmm_ over the full output.
+
+    Mirrors the Triton kernel exactly: single slice, K = max_r (the kernel
+    reads the full K; the loader zero-pads the weight tail beyond the
+    adapter's rank), scaling fused via a pre-scaled x.
+    """
+    if base_output is None:
+        base_output = torch.zeros(
+            (x.shape[0], weights.shape[-2]), device=x.device, dtype=x.dtype
+        )
+    w = weights[batch_info.weight_indices[0]]
+    x_scaled = x * batch_info.scalings[0]
+    base_output.addmm_(x_scaled, w.t())
+    return base_output
+
+
+@triton.jit
+def _sgemm_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Matrix dimensions
+    N,  # output_dim
+    K,  # r
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths and weight id
+    seg_lens,
+    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,
+    # For fused output scaling
+    scalings,
+    ENABLE_PDL: tl.constexpr = False,
+):
+    """
+    Computes a segmented batched matrix multiplication for the LoRA B matrix
+    and adds the result to the output in-place.
+
+    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 (torch.Tensor): The intermediate tensor from the LoRA 'A' multiplication,
+            of shape `(s, K)`, where `s` is the total number of tokens.
+        weights (torch.Tensor): The LoRA 'B' weights for all available adapters,
+            with shape `(num_lora, N, K)`.
+        output (torch.Tensor): The output tensor of shape `(s, N)`. This can be
+            the base model's output for a fused add operation.
+    """
+
+    pid_s = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    batch_id = tl.program_id(axis=2)
+    w_index = tl.load(weight_indices + batch_id)
+    rank = tl.load(lora_ranks + w_index)
+
+    # If rank is 0, this kernel is a no-op.
+    if rank == 0:
+        return
+
+    seg_len = tl.load(seg_lens + batch_id)
+    if pid_s * BLOCK_S >= seg_len:  # also covers seg_len == 0
+        return
+    seg_start = tl.load(seg_indptr + batch_id)
+    scaling = tl.load(scalings + w_index)
+
+    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)
+    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
+    )
+
+    # GDC wait: ensure the prior kernel (producer of x) has fully completed
+    # before consuming its output.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    n_mask = n_offset[None, :] < N
+    output_ptr = output + (
+        s_physical[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
+    )
+    output_mask = (s_offset[:, None] < seg_len) & n_mask
+
+    x_tile = tl.load(
+        x_ptrs,
+        mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K),
+        other=0.0,
+    )
+    w_tile = tl.load(
+        w_ptrs,
+        mask=(k_offset[:, None] < K) & n_mask,
+        other=0.0,
+    )
+
+    # cast fused: the split-K shrink returns fp32, plain path bf16 (no-op)
+    partial_sum = tl.dot(x_tile.to(w_tile.dtype), w_tile) * scaling
+
+    # All input reads are done; hint the runtime to launch the dependent kernel.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # Store result to output matrix (cast to the OUTPUT dtype: x may be the fp32
+    # split-K shrink accumulator while base_output is bf16)
+    partial_sum = partial_sum.to(output.dtype.element_ty)
+    tl.atomic_add(output_ptr, partial_sum, mask=output_mask, sem="relaxed")
+
+
+def sgemm_lora_b_fwd(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor = None,
+) -> torch.Tensor:
+    # x: (s, max_r)
+    # weights: (num_lora, output_dim, max_r)
+    # output: (s, output_dim)
+    # output_dim is much larger than max_r
+
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert len(x.shape) == 2
+    assert len(weights.shape) == 3
+
+    S = x.shape[0]
+    N = weights.shape[-2]
+    R = weights.shape[-1]
+    assert x.shape[-1] == R
+
+    if (
+        (
+            lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
+            or lora_envs.SGLANG_OPT_LORA_CUBLAS_B.get()
+        )
+        and S * R >= _CUBLAS_MIN_S_RANK
+        and weights.shape[0] == 1
+    ):  # single-adapter fast path: only valid with one resident slot
+        return _sgemm_lora_b_cublas(x, weights, batch_info, base_output)
+    # Block shapes
+    BLOCK_S = 16
+    BLOCK_R = triton.next_power_of_2(R)
+    BLOCK_N = 256
+
+    grid = (
+        triton.cdiv(batch_info.max_len, BLOCK_S),
+        triton.cdiv(N, BLOCK_N),
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.zeros((S, N), device=x.device, dtype=x.dtype)
+    else:
+        output = base_output
+
+    sorted_by_adapter = batch_info.permutation is not None
+    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
+    _sgemm_lora_b_kernel[grid](
+        x,
+        weights,
+        output,
+        N,
+        R,
+        x.stride(0),
+        x.stride(1),
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        sorted_by_adapter,
+        BLOCK_S,
+        BLOCK_N,
+        BLOCK_R,
+        batch_info.scalings,
+        ENABLE_PDL=enable_pdl,
+        **pdl_kwargs,
+    )
+    return output