chore: import upstream snapshot with attribution

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