sgl-project--sglang/python/sglang/kernels/ops/moe/virtual_experts.py

"""
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,
    )