chore: import upstream snapshot with attribution

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

@@ -0,0 +1,119 @@
+"""Mixture-of-Experts routing / bookkeeping 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="moe.moe_align_block_size",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel:moe_align_block_size",
+        format_signature=FormatSignature(
+            in_place=True,
+            description="align/sort expert token ids into block-padded buffers",
+        ),
+        description="MoE align-block-size (sgl_kernel wheel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="moe.moe_align_block_size",
+        backend=KernelBackend.CUDA_JIT,
+        target="sglang.jit_kernel.moe_align:moe_align_block_size",
+        capability=_CUDA,
+        format_signature=FormatSignature(
+            in_place=True,
+            description="MoE align-block-size (JIT variant, AOT signature)",
+        ),
+        description="MoE align-block-size (sglang.jit_kernel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="moe.topk_softmax",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel:topk_softmax",
+        format_signature=FormatSignature(
+            in_place=True,
+            description="top-k softmax routing weights/ids",
+        ),
+        description="MoE top-k softmax (sgl_kernel wheel).",
+    )
+)
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    experts_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+    cumsum_buffer: torch.Tensor,
+    pad_sorted_token_ids: bool = False,
+) -> None:
+    """Align and sort expert token ids into block-padded output buffers."""
+    return get_kernel("moe.moe_align_block_size", KernelBackend.CUDA_AOT)(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_token_ids,
+        experts_ids,
+        num_tokens_post_pad,
+        cumsum_buffer,
+        pad_sorted_token_ids,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool = False,
+    moe_softcapping: float = 0.0,
+    correction_bias: Optional[torch.Tensor] = None,
+) -> None:
+    """Compute top-k softmax routing weights/ids for MoE."""
+    return get_kernel("moe.topk_softmax", KernelBackend.CUDA_AOT)(
+        topk_weights,
+        topk_ids,
+        gating_output,
+        renormalize,
+        moe_softcapping,
+        correction_bias,
+    )
+
+
+__all__ = ["moe_align_block_size", "topk_softmax"]
+
+
+# Fused MoE-LoRA Triton kernels migrated into this group (from lora/triton_ops);
+# registered for inventory. Import them from their modules.
+_TRITON_KERNELS = [
+    ("fused_moe_lora_kernel", "fused_moe_lora"),
+    ("virtual_experts", "merged_experts_fused_moe_lora_add"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"moe.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.moe.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn

python/sglang/kernels/ops/moe/fused_moe_lora_kernel.py

@@ -0,0 +1,701 @@
+# Temporarily adapted from https://github.com/vllm-project/vllm/blob/main/vllm/lora/ops/triton_ops/fused_moe_lora_op.py, will optimize in future refactor
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.distributed import (
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_all_reduce,
+)
+from sglang.srt.utils.common import is_blackwell_supported, is_sm90_supported
+
+# Import SGLang's standard PDL support detection
+
+
+_LORA_PTR_DICT: dict[tuple[int, ...], torch.Tensor] = {}
+
+
+def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):
+    """
+    `_LORA_PTR_DICT` collects the required information during `profile_run`,
+    After this, it remains constant and subsequent usage is through LUT.
+    Refer to:
+    https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py
+    """
+    key = tuple(lora_weight.data_ptr() for lora_weight in lora_weights)
+
+    if (ptr_tensor := _LORA_PTR_DICT.get(key)) is not None:
+        return ptr_tensor
+
+    tensor_ptrs = []
+    for lora_weight in lora_weights:
+        tensor_ptrs.append(lora_weight.data_ptr())
+    ptr_tensor = torch.tensor(tensor_ptrs, device=device, dtype=torch.uint64)
+
+    _LORA_PTR_DICT[key] = ptr_tensor
+    return _LORA_PTR_DICT.get(key)
+
+
+@triton.jit(
+    do_not_specialize=[
+        "num_valid_tokens",
+        "EM",
+        "stride_tl",
+        "stride_el",
+        "slice_a_size",
+        "slice_c_size",
+    ]
+)
+def _fused_moe_lora_kernel(
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    # Matrix dimensions
+    N,
+    K,
+    EM,
+    num_valid_tokens,
+    num_experts,
+    lora_ids,
+    adapter_enabled,
+    # The stride variables represent how much to increase the ptr by when
+    # moving by 1 element in a particular dimension. E.g. `stride_am` is
+    # how much to increase `a_ptr` by to get the element one row down
+    # (A has M rows).
+    stride_am,
+    stride_ak,
+    stride_bl,
+    stride_be,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_tl,
+    stride_el,
+    slice_a_size,
+    slice_c_size,
+    # Meta-parameters
+    num_slice_a: tl.constexpr,
+    num_slice_c: tl.constexpr,
+    top_k: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+    USE_GDC: tl.constexpr,
+    launch_pdl: tl.constexpr,
+    IS_PRIMARY: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    slice_id = tl.program_id(axis=1)
+    lora_idx = tl.program_id(axis=2)
+    lora_id = tl.load(lora_ids + lora_idx)
+
+    if lora_id == -1:
+        # Early exit for the no-lora case.
+        return
+    moe_enabled = tl.load(adapter_enabled + lora_id)
+    if moe_enabled == 0:
+        # Early exit for the no moe lora case.
+        return
+    max_loras = tl.num_programs(axis=2)
+    grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)
+
+    # calculate pid_m,pid_n
+    pid_sk = pid % SPLIT_K
+    pid_m_n = pid // SPLIT_K
+    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid_m_n // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m)
+    pid_n = (pid_m_n % num_pid_in_group) // group_size_m
+
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr + lora_id)
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+    # get the expert_id to process curr shard
+    ind = lora_id * stride_el + pid_m
+    expert_id = tl.load(expert_ids_ptr + ind, ind < max_loras * stride_el, -1)
+    if expert_id == -1:
+        return
+
+    # get a_ptr,b_ptr,c_ptr
+    cur_a_ptr = a_ptr + (slice_id % num_slice_a) * slice_a_size
+    cur_b_ptr = tl.load(b_ptr + slice_id).to(tl.pointer_type(c_ptr.dtype.element_ty))
+    cur_c_ptr = c_ptr + (slice_id % num_slice_c) * slice_c_size
+
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
+    offs_k = pid_sk * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
+    # ================================================================= secure
+
+    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+    token_ind = stride_tl * lora_id + offs_token_id
+    offs_token = tl.load(
+        sorted_token_ids_ptr + token_ind, token_ind < max_loras * stride_tl, 0
+    )
+    token_mask = offs_token < num_valid_tokens
+
+    # ================================================================= secure
+
+    # get a_ptrs,b_ptrs
+    a_ptrs = cur_a_ptr + (
+        offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
+    )
+
+    b_ptrs = (
+        cur_b_ptr
+        + lora_id * stride_bl
+        + expert_id * stride_be
+        + offs_k[:, None] * stride_bk
+        + offs_bn[None, :] * stride_bn
+    )
+
+    if USE_GDC and IS_PRIMARY:
+        # GDC launch dependents hints the runtime system to launch dependent kernels.
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # ================================================================= secure
+
+    # accumulator
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    # ================================================================= secure
+
+    # GDC wait waits for ALL programs in the prior kernel to complete
+    # before continuing.
+    if USE_GDC and not IS_PRIMARY:
+        tl.extra.cuda.gdc_wait()
+
+    for k in range(0, grid_k):
+        k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K)
+        # pre-fetch lora weight
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
+        a = tl.load(
+            a_ptrs,
+            mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),
+            other=0.0,
+        )
+        accumulator += tl.dot(a, b.to(a.dtype))
+        # Advance the ptrs to the next K block.
+        a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk
+
+    if MUL_ROUTED_WEIGHT:
+        moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
+        accumulator = accumulator * moe_weight[:, None]
+    accumulator = accumulator.to(c_ptr.dtype.element_ty)
+    # Write back the block of the output
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = cur_c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
+
+    if SPLIT_K == 1:
+        tl.store(c_ptrs, accumulator, mask=c_mask)
+    else:
+        tl.atomic_add(c_ptrs, accumulator, mask=c_mask, sem="relaxed")
+
+
+@torch.inference_mode()
+def _fused_moe_lora_shrink(
+    a_intermediate_cache1: torch.Tensor,
+    # (num_slices, num_tokens, top_k_num, max_lora_rank)
+    qcurr_hidden_states: torch.Tensor,  # (num_tokens, K,)
+    lora_a_stacked: list[
+        torch.Tensor
+    ],  # [(max_loras, num_experts, max_lora_rank, K,),...]
+    topk_weights: torch.Tensor,  # (num_tokens, top_k_num)
+    sorted_token_ids: torch.Tensor,  # (max_loras, _)
+    expert_ids: torch.Tensor,  # (max_loras, _ ,)
+    num_tokens_post_padded: torch.Tensor,  # (max_loras, )
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    ## adding for kernel
+    device: torch.device,
+    N: int,
+    M: int,
+    EM: int,
+    K: int,
+    num_tokens: int,
+    num_experts: int,
+    num_slices: int,
+    block_size_m: int,
+    block_size_n: int,
+    block_size_k: int,
+    group_size_m: int,
+    num_warps: int,
+    num_stages: int,
+    split_k: int,
+    top_k_divisor: int = None,
+    mul_routed_weight: bool = False,
+) -> None:
+    w1_lora_a_stacked = lora_a_stacked[0]
+
+    use_gdc = is_sm90_supported() or is_blackwell_supported()
+    shrink_config = {
+        "BLOCK_SIZE_M": block_size_m,
+        "BLOCK_SIZE_N": block_size_n,
+        "BLOCK_SIZE_K": block_size_k,
+        "GROUP_SIZE_M": group_size_m,
+        "num_warps": num_warps,
+        "num_stages": num_stages,
+        "SPLIT_K": split_k,
+        "USE_GDC": use_gdc,
+        "launch_pdl": use_gdc,  # triton kernel metadata
+    }
+
+    b_ptr = _get_ptr(lora_a_stacked, device)
+
+    grid = lambda META: (
+        split_k
+        * triton.cdiv(EM, META["BLOCK_SIZE_M"])
+        * triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        len(lora_a_stacked),
+        lora_a_stacked[0].shape[0],
+    )
+    _fused_moe_lora_kernel[grid](
+        qcurr_hidden_states,
+        b_ptr,
+        a_intermediate_cache1,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        EM,
+        num_tokens,
+        num_experts,
+        lora_ids,
+        adapter_enabled,
+        qcurr_hidden_states.stride(0),
+        qcurr_hidden_states.stride(1),
+        w1_lora_a_stacked.stride(0),
+        w1_lora_a_stacked.stride(1),
+        w1_lora_a_stacked.stride(3),
+        w1_lora_a_stacked.stride(2),
+        a_intermediate_cache1.stride(2),
+        a_intermediate_cache1.stride(3),
+        sorted_token_ids.stride(0),
+        expert_ids.stride(0),
+        slice_a_size=qcurr_hidden_states.numel(),
+        slice_c_size=a_intermediate_cache1.numel() // num_slices,
+        num_slice_a=1,
+        num_slice_c=num_slices,
+        top_k=(
+            top_k_divisor
+            if top_k_divisor is not None
+            else (1 if mul_routed_weight else top_k_num)
+        ),
+        MUL_ROUTED_WEIGHT=False,
+        IS_PRIMARY=True,
+        **shrink_config,
+    )
+
+
+@torch.inference_mode()
+def _fused_moe_lora_expand(
+    output: torch.Tensor,  # (num_tokens, top_k_num, N*len(lora_a_stacked),)
+    a_intermediate_cache1: torch.Tensor,  # (num_slices, M, top_k_num, max_lora_rank)
+    b_intermediate_cache1: torch.Tensor,  # (num_slices, M, top_k_num, output_dim_size)
+    lora_b_stacked: list[
+        torch.Tensor
+    ],  # [(max_loras, num_experts, max_lora_rank, K,),...]
+    topk_weights: torch.Tensor,  # (num_tokens, top_k_num)
+    sorted_token_ids: torch.Tensor,  # (max_loras, _)
+    expert_ids: torch.Tensor,  # (max_loras, _ ,)
+    num_tokens_post_padded: torch.Tensor,  # (max_loras, )
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    ## adding for kernel
+    device: torch.device,
+    N: int,
+    M: int,
+    EM: int,
+    K: int,
+    num_tokens: int,
+    num_experts: int,
+    num_slices: int,
+    max_lora_rank: int,
+    w1_output_dim_size: int,
+    block_size_m: int,
+    block_size_n: int,
+    block_size_k: int,
+    group_size_m: int,
+    num_warps: int,
+    num_stages: int,
+    split_k: int,
+    mul_routed_weight: bool = False,
+    offset: int = 0,
+) -> None:
+
+    b_ptr = _get_ptr(lora_b_stacked, device)
+    K = max_lora_rank
+    N = w1_output_dim_size
+
+    w1_lora_b_stacked = lora_b_stacked[0]
+
+    a_intermediate_cache1 = a_intermediate_cache1.view(
+        -1, a_intermediate_cache1.shape[3]
+    )
+
+    use_gdc = is_sm90_supported() or is_blackwell_supported()
+    expand_config = {
+        "BLOCK_SIZE_M": block_size_m,
+        "BLOCK_SIZE_N": block_size_n,
+        "BLOCK_SIZE_K": block_size_k,
+        "GROUP_SIZE_M": group_size_m,
+        "num_warps": num_warps,
+        "num_stages": num_stages,
+        "SPLIT_K": split_k,  # Set split_k = 1 for expand calls
+        "USE_GDC": use_gdc,
+        "launch_pdl": use_gdc,  # triton kernel metadata
+    }
+
+    grid = lambda META: (
+        triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        len(lora_b_stacked),
+        lora_b_stacked[0].shape[0],
+    )
+    _fused_moe_lora_kernel[grid](
+        a_intermediate_cache1,
+        b_ptr,
+        b_intermediate_cache1,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        EM,
+        num_tokens,
+        num_experts,
+        lora_ids,
+        adapter_enabled,
+        a_intermediate_cache1.stride(0),
+        a_intermediate_cache1.stride(1),
+        w1_lora_b_stacked.stride(0),
+        w1_lora_b_stacked.stride(1),
+        w1_lora_b_stacked.stride(3),
+        w1_lora_b_stacked.stride(2),
+        b_intermediate_cache1.stride(2),
+        b_intermediate_cache1.stride(3),
+        sorted_token_ids.stride(0),
+        expert_ids.stride(0),
+        slice_a_size=a_intermediate_cache1.numel() // num_slices,
+        slice_c_size=b_intermediate_cache1.numel() // num_slices,
+        num_slice_a=num_slices,
+        num_slice_c=num_slices,
+        top_k=1,
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        IS_PRIMARY=False,
+        **expand_config,
+    )
+    for i in range(num_slices):
+        output[:, :, i * N + offset : (i + 1) * N + offset] += b_intermediate_cache1[i]
+
+
+@torch.inference_mode()
+def _fused_moe_lora(
+    output: torch.Tensor,  # (num_tokens, top_k_num, N*len(lora_a_stacked),)
+    qcurr_hidden_states: torch.Tensor,  # (num_tokens, K,)
+    lora_a_stacked: list[
+        torch.Tensor
+    ],  # [(max_loras, num_experts, max_lora_rank, K,),...]
+    lora_b_stacked: list[
+        torch.Tensor
+    ],  # [(max_loras, num_experts, N, max_lora_rank,),...]
+    topk_weights: torch.Tensor,  # (num_tokens, top_k_num)
+    sorted_token_ids: torch.Tensor,  # (max_loras, _)
+    expert_ids: torch.Tensor,  # (max_loras, _ ,)
+    num_tokens_post_padded: torch.Tensor,  # (max_loras, )
+    max_lora_rank: int,
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    shrink_block_size_m: int,
+    shrink_block_size_n: int,
+    shrink_block_size_k: int,
+    shrink_group_size_m: int,
+    shrink_num_warps: int,
+    shrink_num_stages: int,
+    shrink_split_k: int,
+    expand_block_size_m: int,
+    expand_block_size_n: int,
+    expand_block_size_k: int,
+    expand_group_size_m: int,
+    expand_num_warps: int,
+    expand_num_stages: int,
+    expand_split_k: int,
+    mul_routed_weight: bool = False,
+    fully_sharded: bool = False,
+    offset: int = 0,
+) -> None:
+    assert len(lora_a_stacked) == len(lora_b_stacked) > 0
+    assert (
+        sorted_token_ids.dim()
+        == expert_ids.dim()
+        == topk_weights.dim()
+        == qcurr_hidden_states.dim()
+        == 2
+    )
+    assert (
+        sorted_token_ids.shape[0]
+        == expert_ids.shape[0]
+        == num_tokens_post_padded.shape[0]
+    )
+    assert output.shape[0] == topk_weights.shape[0]
+    assert top_k_num == topk_weights.shape[1]
+    device = qcurr_hidden_states.device
+    num_slices = len(lora_a_stacked)
+    w1_lora_b_stacked = lora_b_stacked[0]
+    num_experts = lora_a_stacked[0].shape[1]
+    N = max_lora_rank
+    M = topk_weights.shape[0]
+    EM = sorted_token_ids.shape[1]
+    K = qcurr_hidden_states.shape[1]
+    num_tokens = M * top_k_num
+    w1_output_dim_size = w1_lora_b_stacked.shape[2]
+
+    # Detect whether input is already expanded (down path: [M*top_k, dim])
+    # or not (gate_up path: [M, dim]). Down path needs divisor=1.
+    input_is_expanded = qcurr_hidden_states.shape[0] == M * top_k_num
+    shrink_top_k_divisor = 1 if input_is_expanded else top_k_num
+
+    a_intermediate_cache1 = torch.zeros(
+        (num_slices, M, top_k_num, max_lora_rank),
+        dtype=output.dtype,
+        device=device,
+    )
+
+    b_intermediate_cache1 = torch.zeros(
+        (num_slices, M, top_k_num, w1_output_dim_size),
+        dtype=output.dtype,
+        device=device,
+    )
+
+    _fused_moe_lora_shrink(
+        a_intermediate_cache1,
+        qcurr_hidden_states,
+        lora_a_stacked,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        top_k_num,
+        lora_ids,
+        adapter_enabled,
+        ## adding for kernel
+        device,
+        N,
+        M,
+        EM,
+        K,
+        num_tokens,
+        num_experts,
+        num_slices,
+        shrink_block_size_m,
+        shrink_block_size_n,
+        shrink_block_size_k,
+        shrink_group_size_m,
+        shrink_num_warps,
+        shrink_num_stages,
+        shrink_split_k,
+        top_k_divisor=shrink_top_k_divisor,
+        mul_routed_weight=False,
+    )
+
+    if fully_sharded:
+        if max_lora_rank == w1_lora_b_stacked.shape[-1]:
+            a_intermediate_cache1 = tensor_model_parallel_all_reduce(
+                a_intermediate_cache1
+            )
+        else:
+            a_intermediate_cache1 = tensor_model_parallel_all_gather(
+                a_intermediate_cache1
+            )
+
+            # reset max_lora_rank to the full rank after allgather
+            max_lora_rank = a_intermediate_cache1.shape[-1]
+
+    _fused_moe_lora_expand(
+        output,
+        a_intermediate_cache1,
+        b_intermediate_cache1,
+        lora_b_stacked,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        top_k_num,
+        lora_ids,
+        adapter_enabled,
+        ## adding for kernel
+        device,
+        N,
+        M,
+        EM,
+        K,
+        num_tokens,
+        num_experts,
+        num_slices,
+        max_lora_rank,
+        w1_output_dim_size,
+        expand_block_size_m,
+        expand_block_size_n,
+        expand_block_size_k,
+        expand_group_size_m,
+        expand_num_warps,
+        expand_num_stages,
+        expand_split_k,
+        mul_routed_weight,
+        offset,
+    )
+
+
+def _fused_moe_lora_fake(
+    output: torch.Tensor,
+    qcurr_hidden_states: torch.Tensor,
+    lora_a_stacked: list[torch.Tensor],
+    lora_b_stacked: list[torch.Tensor],
+    topk_weights: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    max_lora_rank: int,
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    shrink_block_size_m: int,
+    shrink_block_size_n: int,
+    shrink_block_size_k: int,
+    shrink_group_size_m: int,
+    shrink_num_warps: int,
+    shrink_num_stages: int,
+    shrink_split_k: int,
+    expand_block_size_m: int,
+    expand_block_size_n: int,
+    expand_block_size_k: int,
+    expand_group_size_m: int,
+    expand_num_warps: int,
+    expand_num_stages: int,
+    expand_split_k: int,
+    mul_routed_weight: bool = False,
+    fully_sharded: bool = False,
+    offset: int = 0,
+) -> None:
+    return
+
+
+def _fused_moe_lora_shrink_fake(
+    a_intermediate_cache1: torch.Tensor,
+    qcurr_hidden_states: torch.Tensor,
+    lora_a_stacked: list[torch.Tensor],
+    topk_weights: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    device: torch.device,
+    N: int,
+    M: int,
+    EM: int,
+    K: int,
+    num_tokens: int,
+    num_experts: int,
+    num_slices: int,
+    block_size_m: int,
+    block_size_n: int,
+    block_size_k: int,
+    group_size_m: int,
+    num_warps: int,
+    num_stages: int,
+    split_k: int,
+    mul_routed_weight: bool = False,
+) -> None:
+    return
+
+
+def _fused_moe_lora_expand_fake(
+    output: torch.Tensor,
+    a_intermediate_cache1: torch.Tensor,
+    b_intermediate_cache1: torch.Tensor,
+    lora_b_stacked: list[torch.Tensor],
+    topk_weights: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    device: torch.device,
+    N: int,
+    M: int,
+    EM: int,
+    K: int,
+    num_tokens: int,
+    num_experts: int,
+    num_slices: int,
+    max_lora_rank: int,
+    w1_output_dim_size: int,
+    block_size_m: int,
+    block_size_n: int,
+    block_size_k: int,
+    group_size_m: int,
+    num_warps: int,
+    num_stages: int,
+    split_k: int,
+    mul_routed_weight: bool = False,
+    offset: int = 0,
+) -> None:
+    return
+
+
+# Register as SGLang custom ops following the same pattern as other ops
+try:
+    from sglang.srt.utils.common import direct_register_custom_op
+
+    direct_register_custom_op(
+        op_name="fused_moe_lora",
+        op_func=_fused_moe_lora,
+        mutates_args=["output"],
+        fake_impl=_fused_moe_lora_fake,
+    )
+
+    direct_register_custom_op(
+        op_name="fused_moe_lora_shrink",
+        op_func=_fused_moe_lora_shrink,
+        mutates_args=["a_intermediate_cache1"],
+        fake_impl=_fused_moe_lora_shrink_fake,
+    )
+
+    direct_register_custom_op(
+        op_name="fused_moe_lora_expand",
+        op_func=_fused_moe_lora_expand,
+        mutates_args=["output", "b_intermediate_cache1"],
+        fake_impl=_fused_moe_lora_expand_fake,
+    )
+
+    # Export through torch.ops.sglang namespace
+    fused_moe_lora = torch.ops.sglang.fused_moe_lora
+    fused_moe_lora_shrink = torch.ops.sglang.fused_moe_lora_shrink
+    fused_moe_lora_expand = torch.ops.sglang.fused_moe_lora_expand
+
+except AttributeError:
+    fused_moe_lora = _fused_moe_lora
+    fused_moe_lora_shrink = _fused_moe_lora_shrink
+    fused_moe_lora_expand = _fused_moe_lora_expand

python/sglang/kernels/ops/moe/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/moe/virtual_experts.py

@@ -0,0 +1,786 @@
+"""
+LoRA Virtual Experts Triton Ops.
+"""
+
+import functools
+from typing import Any
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.jit_kernel.moe_align import moe_align_block_size as jit_moe_align_block_size
+
+
+@triton.jit
+def _fused_virtual_topk_ids_kernel(
+    topk_ids_ptr,
+    token_lora_mapping_ptr,
+    virtual_topk_ids_ptr,
+    token_lora_mask_ptr,
+    num_experts_for_weight: tl.constexpr,
+    M,
+    top_k: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Fuses _get_virtual_topk_ids: comparison + clamp + arithmetic into one kernel.
+
+    For each (m, k):
+        lora_id = token_lora_mapping[m]
+        mask[m] = (lora_id >= 0)
+        safe_lora = max(lora_id, 0)
+        if shared_outer:  (handled by num_experts_for_weight == 0 sentinel)
+            virtual_topk_ids[m, k] = safe_lora * 1  (= safe_lora)
+        else:
+            virtual_topk_ids[m, k] = topk_ids[m, k] + safe_lora * num_experts_for_weight
+    """
+    pid = tl.program_id(0)
+    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    total = M * top_k
+    valid = offs < total
+
+    m = offs // top_k
+    # k = offs % top_k  # not needed directly
+
+    lora_id = tl.load(token_lora_mapping_ptr + m, mask=valid, other=0)
+    mask_val = lora_id >= 0
+    safe_lora = tl.maximum(lora_id, 0)
+
+    base = tl.load(topk_ids_ptr + offs, mask=valid, other=0)
+    # Preserve negative sentinel topk_ids (e.g. -1 for non-local experts after
+    # EP dispatch). Without this, `-1 + safe_lora * num_experts` would land on
+    # a real virtual-expert slot belonging to another adapter and trigger OOB
+    # loads in downstream LoRA kernels.
+    shifted = base + safe_lora * num_experts_for_weight
+    result = tl.where(base < 0, base, shifted)
+    tl.store(virtual_topk_ids_ptr + offs, result, mask=valid)
+
+    # Write mask once per row (at first k position)
+    k = offs % top_k
+    is_first_k = k == 0
+    tl.store(token_lora_mask_ptr + m, mask_val, mask=valid & is_first_k)
+
+
+def _fused_virtual_topk_ids(
+    topk_ids: torch.Tensor,
+    token_lora_mapping: torch.Tensor,
+    num_experts: int,
+    shared_outer: bool,
+    max_loras: int,
+) -> tuple[torch.Tensor, torch.Tensor, int]:
+    """
+    Returns virtual topk_ids, token_lora_mask, and virtual_num_experts.
+    """
+    M, top_k = topk_ids.shape
+    device = topk_ids.device
+
+    if shared_outer:
+        num_experts_for_weight = 1
+        # For shared_outer, we need topk_ids to be zeros
+        zero_topk = torch.zeros_like(topk_ids)
+        input_topk = zero_topk
+    else:
+        num_experts_for_weight = num_experts
+        input_topk = topk_ids
+
+    virtual_topk_ids = torch.empty_like(topk_ids)
+    token_lora_mask = torch.empty(M, dtype=torch.bool, device=device)
+
+    BLOCK_SIZE = 1024
+    grid = ((M * top_k + BLOCK_SIZE - 1) // BLOCK_SIZE,)
+
+    _fused_virtual_topk_ids_kernel[grid](
+        input_topk,
+        token_lora_mapping,
+        virtual_topk_ids,
+        token_lora_mask,
+        num_experts_for_weight,
+        M,
+        top_k,
+        BLOCK_SIZE,
+    )
+
+    virtual_num_experts = num_experts_for_weight * max_loras
+    return virtual_topk_ids, token_lora_mask, virtual_num_experts
+
+
+@triton.jit
+def _fused_sanitize_expert_ids_kernel(
+    expert_ids_ptr,
+    output_ptr,
+    num_virtual_experts,
+    N,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    valid = offs < N
+
+    eid = tl.load(expert_ids_ptr + offs, mask=valid, other=0)
+    result = tl.where(eid < num_virtual_experts, eid, -1)
+    tl.store(output_ptr + offs, result, mask=valid)
+
+
+def fused_sanitize_expert_ids(
+    expert_ids: torch.Tensor,
+    num_virtual_experts: int,
+) -> torch.Tensor:
+    """
+    Sanitize expert_ids by replacing values >= num_virtual_experts with -1.
+
+    Returns a new tensor with expert_ids >= num_virtual_experts replaced by -1.
+    """
+    N = expert_ids.numel()
+    output = torch.empty_like(expert_ids)
+
+    BLOCK_SIZE = 1024
+    grid = ((N + BLOCK_SIZE - 1) // BLOCK_SIZE,)
+
+    _fused_sanitize_expert_ids_kernel[grid](
+        expert_ids,
+        output,
+        num_virtual_experts,
+        N,
+        BLOCK_SIZE,
+    )
+    return output
+
+
+@triton.jit
+def _moe_lora_shrink_splitk_kernel(
+    # Pointers
+    a_ptr,  # type: ignore  # [num_tokens, K]
+    b_ptr,  # type: ignore  # [num_virtual_experts, N, K]
+    c_ptr,  # type: ignore  # [num_tokens * top_k, N]  (pre-zeroed when SPLIT_K > 1)
+    sorted_token_ids_ptr,  # type: ignore
+    expert_ids_ptr,  # type: ignore
+    num_tokens_post_padded_ptr,  # type: ignore
+    # Dimensions
+    N,  # type: ignore
+    K,  # type: ignore
+    num_valid_tokens,  # type: ignore
+    # Strides
+    stride_am,  # type: ignore
+    stride_ak,  # type: ignore
+    stride_be,  # type: ignore
+    stride_bn,  # type: ignore
+    stride_bk,  # type: ignore
+    stride_cm,  # type: ignore
+    stride_cn,  # type: ignore
+    # Constexprs
+    top_k: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+):
+    """Split-K grouped GEMM for the LoRA A (shrink) stage with few virtual experts."""
+    pid = tl.program_id(0)
+    pid_sk = pid % SPLIT_K
+    pid_mn = pid // SPLIT_K
+
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+    num_pid_m = tl.cdiv(num_tokens_post_padded, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid_mn // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid_mn % num_pid_in_group) % group_size_m)
+    pid_n = (pid_mn % num_pid_in_group) // group_size_m
+
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+
+    # Token routing (same pattern as fused_moe_triton_kernels)
+    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
+    token_mask = offs_token < num_valid_tokens
+
+    off_expert = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+    if off_expert == -1:
+        return
+
+    # Pointers
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
+    offs_k = pid_sk * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
+
+    a_ptrs = a_ptr + (
+        offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
+    )
+    b_ptrs = (
+        b_ptr
+        + off_expert * stride_be
+        + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+    )
+
+    # Accumulate
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)
+    for k in range(0, grid_k):
+        k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K)
+        k_mask = offs_k[:, None] < k_remaining
+        a = tl.load(
+            a_ptrs,
+            mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),
+            other=0.0,
+        )
+        b = tl.load(b_ptrs, mask=k_mask, other=0.0)
+        accumulator += tl.dot(a, b.to(a.dtype))
+        a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk
+
+    accumulator = accumulator.to(c_ptr.dtype.element_ty)
+
+    # Write output
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
+    if SPLIT_K == 1:
+        tl.store(c_ptrs, accumulator, mask=c_mask)
+    else:
+        tl.atomic_add(c_ptrs, accumulator, mask=c_mask, sem="relaxed")
+
+
+def _invoke_moe_lora_shrink_splitk(
+    hidden_states: torch.Tensor,
+    weight: torch.Tensor,
+    output: torch.Tensor,
+    topk_ids: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    top_k: int,
+    config: dict[str, Any],
+) -> None:
+    """Launch split-K shrink kernel for LoRA A with few virtual experts."""
+    N = weight.shape[1]
+    K = weight.shape[2]
+    BLOCK_SIZE_M = config["BLOCK_SIZE_M"]
+    BLOCK_SIZE_N = min(config.get("BLOCK_SIZE_N", 64), max(16, N))
+    BLOCK_SIZE_K = config.get("BLOCK_SIZE_K", 64)
+    GROUP_SIZE_M = config.get("GROUP_SIZE_M", 1)
+
+    num_m_blocks = triton.cdiv(sorted_token_ids.shape[0], BLOCK_SIZE_M)
+    num_n_blocks = triton.cdiv(N, BLOCK_SIZE_N)
+    base_grid = num_m_blocks * num_n_blocks
+    max_split_k = max(1, K // BLOCK_SIZE_K)
+    SPLIT_K = min(max_split_k, max(1, 128 // base_grid)) if base_grid < 128 else 1
+
+    grid = (SPLIT_K * base_grid,)
+
+    _moe_lora_shrink_splitk_kernel[grid](
+        hidden_states,
+        weight,
+        output,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        topk_ids.numel(),
+        hidden_states.stride(0),
+        hidden_states.stride(1),
+        weight.stride(0),
+        weight.stride(1),
+        weight.stride(2),
+        output.stride(0),
+        output.stride(1),
+        top_k=top_k,
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=BLOCK_SIZE_N,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+        GROUP_SIZE_M=GROUP_SIZE_M,
+        SPLIT_K=SPLIT_K,
+        num_warps=config.get("num_warps", 4),
+        num_stages=config.get("num_stages", 4),
+    )
+
+
+def _align_block_size_jit(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """CUDA JIT align_block_size for num_experts > 1024 (up to 8191).
+
+    Uses the v2 kernel from moe_align_kernel.cu which supports large expert
+    counts via per-thread multi-expert processing and a two-level warp scan,
+    replacing the previous pure-PyTorch fallback that had excessive CPU overhead
+    from 15+ individual kernel launches and torch.argsort.
+
+    The JIT kernel uses a +1 offset convention: topk_ids are shifted by +1 so
+    that the EP sentinel value (-1) maps to bucket 0. The kernel internally
+    handles histogram, padded prefix-sum, expert_ids assignment, and token
+    scattering in just 2–3 CUDA kernel launches.
+    """
+    assert num_experts <= 8191, (
+        f"_align_block_size_jit supports at most 8191 experts "
+        f"(num_moe_experts * max_loras), got {num_experts}"
+    )
+
+    device = topk_ids.device
+    flat_topk_ids = topk_ids.reshape(-1)
+    if flat_topk_ids.dtype == torch.int64:
+        flat_topk_ids = flat_topk_ids.to(torch.int32)
+    num_total_tokens = flat_topk_ids.numel()
+
+    if num_total_tokens == 0:
+        empty = torch.empty(0, dtype=torch.int32, device=device)
+        return empty, empty, torch.zeros(1, dtype=torch.int32, device=device)
+
+    # JIT kernel uses +1 offset convention: -1 -> bucket 0 (sentinel),
+    # expert i -> bucket i+1. So pass num_experts + 1 as the bucket count.
+    jit_num_experts = num_experts + 1
+
+    if num_total_tokens < jit_num_experts:
+        max_num_tokens_padded = num_total_tokens * block_size
+    else:
+        max_num_tokens_padded = num_total_tokens + jit_num_experts * (block_size - 1)
+
+    # Align every sub-buffer offset to a multiple of 4 (VEC_SIZE). The CUDA
+    # kernel fills sorted_token_ids with vectorized int4 writes whose last
+    # store can spill up to 3 int32s past the logical end. With a fused
+    # allocation the spill would corrupt the adjacent sub-buffer.
+    _A4 = lambda n: (n + 3) & ~3  # noqa: E731
+    max_num_tokens_padded = _A4(max_num_tokens_padded)
+    max_num_m_blocks = (max_num_tokens_padded + block_size - 1) // block_size
+    max_num_m_blocks_padded = _A4(max_num_m_blocks)
+    num_post_pad_size = _A4(1)  # 1 element, padded to 4
+    cumsum_size = _A4(jit_num_experts + 1)
+
+    # Single allocation sliced into 4 views (zero-copy) to avoid
+    # per-call Python overhead of 4 separate torch.empty calls.
+    total_buf = (
+        max_num_tokens_padded
+        + max_num_m_blocks_padded
+        + num_post_pad_size
+        + cumsum_size
+    )
+    buf = torch.empty(total_buf, dtype=torch.int32, device=device)
+    off = 0
+    sorted_token_ids = buf[off : off + max_num_tokens_padded]
+    off += max_num_tokens_padded
+    expert_ids = buf[off : off + max_num_m_blocks]
+    off += max_num_m_blocks_padded
+    num_tokens_post_padded = buf[off : off + 1]
+    off += num_post_pad_size
+    cumsum_buffer = buf[off : off + jit_num_experts + 1]
+
+    jit_moe_align_block_size(
+        flat_topk_ids,
+        jit_num_experts,
+        block_size,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        cumsum_buffer,
+        True,  # pad_sorted_token_ids
+    )
+
+    return sorted_token_ids, expert_ids, num_tokens_post_padded
+
+
+@torch.compile(dynamic=True)
+def _align_block_size_torch(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Pure-PyTorch align_block_size for num_experts > 1024, compiled via torch.compile.
+
+    Fallback for platforms where the CUDA JIT kernel is unavailable (e.g. AMD/ROCm).
+
+    Out-of-range topk_ids (negative sentinels left by EP dispatch, or virtual-
+    expert IDs >= num_experts produced when those sentinels are combined with
+    a per-adapter offset) are routed into a dedicated sentinel bucket. Without
+    this, indexing ``padded_offsets[sorted_expert_ids]`` would wrap (-1) or
+    OOB-read, and the bad expert ids would propagate into the downstream LoRA
+    GEMM as real expert slots.
+    """
+    device = topk_ids.device
+    flat_topk_ids = topk_ids.reshape(-1).to(torch.int64)
+    num_total_tokens = flat_topk_ids.numel()
+
+    sentinel = num_experts
+    valid_mask = (flat_topk_ids >= 0) & (flat_topk_ids < num_experts)
+    safe_topk_ids = torch.where(
+        valid_mask,
+        flat_topk_ids,
+        torch.full_like(flat_topk_ids, sentinel),
+    )
+
+    bucket_count = num_experts + 1
+    max_total_padded_tokens = (
+        (num_total_tokens + bucket_count * (block_size - 1) + block_size - 1)
+        // block_size
+    ) * block_size
+    max_num_blocks = max_total_padded_tokens // block_size
+
+    sorted_token_ids = torch.full(
+        (max_total_padded_tokens,),
+        num_total_tokens,
+        dtype=torch.int32,
+        device=device,
+    )
+    expert_ids = torch.full(
+        (max_num_blocks,),
+        -1,
+        dtype=torch.int32,
+        device=device,
+    )
+
+    if num_total_tokens == 0:
+        num_tokens_post_padded = torch.zeros((1,), dtype=torch.int32, device=device)
+        return sorted_token_ids, expert_ids, num_tokens_post_padded
+
+    sorted_order = torch.argsort(safe_topk_ids)
+    sorted_expert_ids = safe_topk_ids[sorted_order]
+    expert_range = torch.arange(bucket_count, device=device, dtype=torch.int64)
+    counts_offsets = torch.searchsorted(sorted_expert_ids, expert_range, right=False)
+    counts_end = torch.searchsorted(sorted_expert_ids, expert_range, right=True)
+    counts = counts_end - counts_offsets
+    padded_counts = ((counts + block_size - 1) // block_size) * block_size
+    total_padded_tokens = padded_counts.sum().to(torch.int32).reshape(1)
+    padded_offsets = torch.cumsum(padded_counts, dim=0) - padded_counts
+
+    token_ranks = (
+        torch.arange(num_total_tokens, device=device, dtype=torch.int64)
+        - counts_offsets[sorted_expert_ids]
+    )
+    output_positions = padded_offsets[sorted_expert_ids] + token_ranks
+    sorted_token_ids.scatter_(
+        0,
+        output_positions.to(torch.int64),
+        sorted_order.to(torch.int32),
+    )
+
+    block_counts = padded_counts // block_size
+    real_block_counts = block_counts.clone()
+    real_block_counts[sentinel] = 0
+    actual_num_blocks = real_block_counts.sum()
+
+    if max_num_blocks <= 0:
+        return sorted_token_ids, expert_ids, total_padded_tokens
+
+    block_offsets = torch.cumsum(real_block_counts, dim=0)
+    all_block_positions = torch.arange(max_num_blocks, device=device, dtype=torch.int64)
+    assigned_experts = torch.searchsorted(
+        block_offsets, all_block_positions, right=True
+    ).to(torch.int32)
+    expert_ids.copy_(
+        torch.where(
+            all_block_positions < actual_num_blocks,
+            assigned_experts,
+            torch.full_like(assigned_experts, -1),
+        )
+    )
+
+    return sorted_token_ids, expert_ids, total_padded_tokens
+
+
+def _align_block_size_large(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Dispatch to the CUDA JIT kernel when available, otherwise fall back to
+    the pure-PyTorch torch.compile path (needed on AMD/ROCm or when the JIT
+    module fails to load)."""
+    try:
+        return _align_block_size_jit(topk_ids, block_size, num_experts)
+    except Exception:
+        return _align_block_size_torch(topk_ids, block_size, num_experts)
+
+
+def _merged_experts_fused_moe_lora_add_fake(
+    output: torch.Tensor,
+    hidden_states: torch.Tensor,
+    lora_a: torch.Tensor,
+    lora_b: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    token_lora_mapping: torch.Tensor,
+    mul_routed_weight: bool,
+    experts_shared_outer_loras_a: bool,
+    experts_shared_outer_loras_b: bool,
+) -> None:
+    return
+
+
+def _merged_experts_fused_moe_lora_add_impl(
+    output: torch.Tensor,
+    hidden_states: torch.Tensor,
+    lora_a: torch.Tensor,
+    lora_b: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    token_lora_mapping: torch.Tensor,
+    mul_routed_weight: bool,
+    experts_shared_outer_loras_a: bool,
+    experts_shared_outer_loras_b: bool,
+    routing_cache: dict | None = None,
+) -> None:
+    """
+    1. Prepare virtual expert routing metadata from topk_ids + token_lora_mapping * num_experts.
+    2. Flatten LoRA weights from [max_loras, num_experts, ...] to [max_loras * num_experts, ...].
+    3. Run regular SGLang fused-MoE kernels for LoRA A and LoRA B.
+    4. Mask out tokens with token_lora_mapping == -1 on the add path.
+    """
+    max_loras, _, max_lora_rank, _ = lora_a.shape
+    input_top_k = 1 if hidden_states.shape[0] == topk_ids.numel() else topk_ids.shape[1]
+
+    def _merge_lora_expert_weight(t: torch.Tensor) -> torch.Tensor:
+        # [max_loras, num_experts, x, y] -> [max_loras * num_experts, x, y]
+        return t.reshape(t.shape[0] * t.shape[1], t.shape[2], t.shape[3])
+
+    def _get_stage_config(
+        weight: torch.Tensor,
+        stage_top_k: int,
+    ) -> dict[str, Any]:
+        from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config import (
+            get_config_dtype_str,
+            try_get_optimal_moe_config,
+        )
+
+        config_dtype = get_config_dtype_str(dtype=hidden_states.dtype)
+        get_config_func = functools.partial(
+            try_get_optimal_moe_config,
+            weight.shape,
+            weight.shape,
+            stage_top_k,
+            config_dtype,
+        )
+        try:
+            cfg = get_config_func(token_lora_mapping.shape[0])
+        except ValueError:
+            K_dim = weight.shape[2]
+            N_dim = weight.shape[1]
+            if K_dim >= 1024:
+                default_block_k = 256
+            elif K_dim >= 64:
+                default_block_k = 64
+            else:
+                default_block_k = max(16, K_dim)
+            cfg = {
+                "BLOCK_SIZE_M": 64,
+                "BLOCK_SIZE_N": min(64, max(16, N_dim)),
+                "BLOCK_SIZE_K": min(default_block_k, max(16, K_dim)),
+                "GROUP_SIZE_M": 1,
+                "num_warps": 4,
+                "num_stages": 4,
+            }
+        return cfg
+
+    def _align_block_size(
+        topk_ids: torch.Tensor,
+        block_size: int,
+        num_experts: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # The native align kernel consumes num_experts + 1 internally for its
+        # sentinel bucket, so the 1024-expert boundary must use the fallback path.
+        if num_experts < 1024:
+            from sglang.srt.layers.moe.moe_runner.triton_utils.moe_align_block_size import (
+                moe_align_block_size as native_moe_align_block_size,
+            )
+
+            return native_moe_align_block_size(topk_ids, block_size, num_experts)
+        return _align_block_size_large(topk_ids, block_size, num_experts)
+
+    def _get_routing(
+        topk_ids: torch.Tensor,
+        token_lora_mapping: torch.Tensor,
+        num_experts: int,
+        shared_outer: bool,
+        block_size: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Check routing_cache for cross-call reuse (gate_up and down share routing)
+        cache_key = (num_experts, shared_outer, block_size)
+        if routing_cache is not None:
+            cached = routing_cache.get(cache_key)
+            if cached is not None:
+                return cached
+
+        virtual_topk_ids, token_lora_mask, virtual_num_experts = (
+            _fused_virtual_topk_ids(
+                topk_ids, token_lora_mapping, num_experts, shared_outer, max_loras
+            )
+        )
+        sorted_token_ids, expert_ids, num_tokens_post_padded = _align_block_size(
+            virtual_topk_ids,
+            block_size=block_size,
+            num_experts=virtual_num_experts,
+        )
+        # _align_block_size uses a worst-case padded allocation. Trim the routing buffers
+        # to a tighter upper bound so we keep the real routed work but drop unused padding
+        num_tokens = topk_ids.numel()
+        max_nonempty = min(num_tokens, virtual_num_experts)
+        tight_padded = (
+            triton.cdiv(num_tokens + max_nonempty * (block_size - 1), block_size)
+            * block_size
+        )
+        sorted_token_ids = sorted_token_ids[:tight_padded]
+        expert_ids = expert_ids[: tight_padded // block_size]
+        expert_ids = fused_sanitize_expert_ids(expert_ids, virtual_num_experts)
+        result = (
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            token_lora_mask,
+        )
+
+        if routing_cache is not None:
+            routing_cache[cache_key] = result
+
+        return result
+
+    from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import (
+        invoke_fused_moe_kernel,
+    )
+
+    lora_a_virtual = _merge_lora_expert_weight(lora_a)
+    lora_b_virtual = _merge_lora_expert_weight(lora_b)
+    num_experts_a = lora_a.shape[1]
+    num_experts_b = lora_b.shape[1]
+
+    intermediate = torch.zeros(
+        [token_lora_mapping.shape[0], topk_ids.shape[1], max_lora_rank],
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+
+    a_stage_config = _get_stage_config(lora_a_virtual, input_top_k)
+    (
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        token_lora_mask,
+    ) = _get_routing(
+        topk_ids,
+        token_lora_mapping,
+        num_experts_a,
+        experts_shared_outer_loras_a,
+        a_stage_config["BLOCK_SIZE_M"],
+    )
+
+    _invoke_moe_lora_shrink_splitk(
+        hidden_states,
+        lora_a_virtual,
+        intermediate.view(-1, max_lora_rank),
+        topk_ids,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        input_top_k,
+        a_stage_config,
+    )
+
+    b_stage_config = _get_stage_config(lora_b_virtual, 1)
+    (
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        token_lora_mask,
+    ) = _get_routing(
+        topk_ids,
+        token_lora_mapping,
+        num_experts_b,
+        experts_shared_outer_loras_b,
+        b_stage_config["BLOCK_SIZE_M"],
+    )
+
+    invoke_fused_moe_kernel(
+        intermediate.view(-1, max_lora_rank),
+        lora_b_virtual,
+        None,
+        output,
+        None,
+        None,
+        None,
+        topk_weights,
+        topk_ids,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        mul_routed_weight,
+        1,
+        b_stage_config,
+        tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16,
+        False,
+        False,
+        False,
+        False,
+        False,
+        None,
+        fuse_add_to_output=True,
+        add_output_mask=token_lora_mask,
+        router_topk=topk_ids.shape[1],
+    )
+
+
+def _merged_experts_fused_moe_lora_add_op(
+    output: torch.Tensor,
+    hidden_states: torch.Tensor,
+    lora_a: torch.Tensor,
+    lora_b: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    token_lora_mapping: torch.Tensor,
+    mul_routed_weight: bool,
+    experts_shared_outer_loras_a: bool,
+    experts_shared_outer_loras_b: bool,
+) -> None:
+    _merged_experts_fused_moe_lora_add_impl(
+        output,
+        hidden_states,
+        lora_a,
+        lora_b,
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        mul_routed_weight,
+        experts_shared_outer_loras_a,
+        experts_shared_outer_loras_b,
+    )
+
+
+from sglang.srt.utils.common import direct_register_custom_op
+
+direct_register_custom_op(
+    op_name="merged_experts_fused_moe_lora_add",
+    op_func=_merged_experts_fused_moe_lora_add_op,
+    mutates_args=["output"],
+    fake_impl=_merged_experts_fused_moe_lora_add_fake,
+)
+
+
+def merged_experts_fused_moe_lora_add(
+    output: torch.Tensor,
+    hidden_states: torch.Tensor,
+    lora_a: torch.Tensor,
+    lora_b: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    token_lora_mapping: torch.Tensor,
+    mul_routed_weight: bool,
+    experts_shared_outer_loras_a: bool,
+    experts_shared_outer_loras_b: bool,
+    routing_cache: dict | None = None,
+) -> None:
+    """Public API: wraps the registered op with routing_cache support."""
+    _merged_experts_fused_moe_lora_add_impl(
+        output,
+        hidden_states,
+        lora_a,
+        lora_b,
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        mul_routed_weight,
+        experts_shared_outer_loras_a,
+        experts_shared_outer_loras_b,
+        routing_cache,
+    )