"""
LoRA Virtual Experts Triton Ops.
"""

import functools
from typing import Any

import torch
import triton
import triton.language as tl

from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
    get_pdl_launch_metadata,
)
from sglang.kernels.ops.moe.moe_align import (
    moe_align_block_size as jit_moe_align_block_size,
)
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs


@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,
    local_expert_offset: tl.constexpr,
    local_num_experts: tl.constexpr,
    EP_LOCAL: tl.constexpr,
    ENABLE_PDL: tl.constexpr = False,
):
    """
    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

    # GDC wait: topk_ids/token_lora_mapping come from the immediately
    # preceding routing kernels; wait before consuming them.
    if ENABLE_PDL:
        tl.extra.cuda.gdc_wait()

    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)
    if EP_LOCAL:
        # EP: drop experts this rank does not own to the -1 sentinel (handled just
        # below); owned experts keep their GLOBAL id so the global contiguous weight
        # buffer indexes with a single stride and the merged-weight reshape stays a
        # free view. Fused in here -> no separate where()/elementwise launch.
        owned = (base >= local_expert_offset) & (
            base < local_expert_offset + local_num_experts
        )
        base = tl.where(owned, base, -1)
    # 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)
    result = tl.where(mask_val, result, -1)
    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)

    # All work is done; hint the runtime to launch the dependent kernel.
    if ENABLE_PDL:
        tl.extra.cuda.gdc_launch_dependents()


def _fused_virtual_topk_ids(
    topk_ids: torch.Tensor,
    token_lora_mapping: torch.Tensor,
    num_experts: int,
    shared_outer: bool,
    max_loras: int,
    local_expert_offset: int = 0,
    local_num_experts: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor, int]:
    """
    Returns virtual topk_ids, token_lora_mask, and virtual_num_experts.

    `local_num_experts` (< `num_experts`) enables the EP mask on the per-expert
    path: non-owned `topk_ids` slots are dropped to the -1 sentinel in-kernel.
    The shared-outer path routes by lora id, not expert, so it is never masked
    (else non-owned-but-valid tokens would be dropped on ranks > 0).
    """
    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
        ep_local = False
    else:
        num_experts_for_weight = num_experts
        input_topk = topk_ids
        ep_local = local_num_experts is not None and local_num_experts < num_experts

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

    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
    _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,
        local_expert_offset,
        local_num_experts if local_num_experts is not None else 0,
        ep_local,
        ENABLE_PDL=enable_pdl,
        **pdl_kwargs,
    )

    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,
    ENABLE_PDL: tl.constexpr = False,
):
    pid = tl.program_id(0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    valid = offs < N

    # GDC wait: expert_ids comes from the immediately preceding align kernel;
    # wait before consuming it.
    if ENABLE_PDL:
        tl.extra.cuda.gdc_wait()

    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)

    # All work is done; hint the runtime to launch the dependent kernel.
    if ENABLE_PDL:
        tl.extra.cuda.gdc_launch_dependents()


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

    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
    _fused_sanitize_expert_ids_kernel[grid](
        expert_ids,
        output,
        num_virtual_experts,
        N,
        BLOCK_SIZE,
        ENABLE_PDL=enable_pdl,
        **pdl_kwargs,
    )
    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,
    ENABLE_PDL: tl.constexpr = False,
):
    """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

    # GDC wait: the routing buffers (sorted_token_ids/expert_ids/
    # num_tokens_post_padded) come from the immediately preceding align and
    # sanitize kernels; wait before consuming them.
    if ENABLE_PDL:
        tl.extra.cuda.gdc_wait()

    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

    # All input reads are done; hint the runtime to launch the dependent kernel.
    if ENABLE_PDL:
        tl.extra.cuda.gdc_launch_dependents()

    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(128, triton.next_power_of_2(N))
    BLOCK_SIZE_K = 256
    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
    # Single source of truth shared with the caller's zero-intermediate decision:
    # split-K accumulation REQUIRES a pre-zeroed output, so the predicted and
    # launched SPLIT_K must never diverge.
    SPLIT_K = _get_moe_lora_shrink_split_k(weight, sorted_token_ids, config)

    grid = (SPLIT_K * base_grid,)

    enable_pdl, pdl_kwargs = get_pdl_launch_metadata()
    _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,
        ENABLE_PDL=enable_pdl,
        num_warps=config.get("num_warps", 4),
        num_stages=config.get("num_stages", 4),
        **pdl_kwargs,
    )


def _get_moe_lora_shrink_split_k(
    weight: torch.Tensor,
    sorted_token_ids: torch.Tensor,
    config: dict[str, Any],
) -> int:
    """Choose split-K from rank and available occupancy.

    Skinny output ranks benefit from more K splits because each output tile
    carries less work. Block sizes must mirror _invoke_moe_lora_shrink_splitk.

    """
    N = weight.shape[1]
    K = weight.shape[2]
    block_size_m = config["BLOCK_SIZE_M"]
    block_size_n = min(128, triton.next_power_of_2(N))
    block_size_k = 256
    num_m_blocks = triton.cdiv(sorted_token_ids.shape[0], block_size_m)
    base_grid = num_m_blocks * triton.cdiv(N, block_size_n)
    target = 512 if N <= 16 else 384 if N <= 32 else 256
    max_split_k = max(1, K // block_size_k)
    return max(1, min(triton.cdiv(target, base_grid), max_split_k, 8))


# Rank-specialized LoRA-B expand kernel lives in lora/trtllm_lora_temp/.
# Re-export so existing call sites (and any external imports) keep working.
from sglang.srt.lora.trtllm_lora_temp.specialized_expand import (  # noqa: E402,F401
    _invoke_moe_lora_expand_add,
    _moe_lora_expand_add_kernel,
)


def _align_block_size_jit(
    topk_ids: torch.Tensor,
    block_size: int,
    num_experts: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """CUDA JIT alignment for up to 8191 experts.

    Expert IDs are shifted by one so ``-1`` maps to a sentinel bucket. The
    fused allocation stays int4-aligned for the kernel's vectorized clear.
    """
    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,
    fuse_add_to_output: bool = True,
    fuse_sum_all_reduce: bool = False,
    use_direct_expand_add: bool = False,
    local_expert_offset: int = 0,
    local_num_experts: int | None = None,
    stage: str = "all",
    intermediate_buffer: torch.Tensor | None = None,
    expand_wait_event: "torch.cuda.Event | None" = None,
    broadcast_intermediate: bool = False,
    prewarm_a_routing: bool = True,
    prewarm_b_routing: bool = True,
    zero_intermediate: bool = False,
) -> "torch.Tensor | 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.

    ``stage`` splits the call for cross-stream overlap (down-LoRA/finalize overlap):
    - ``"all"`` (default): shrink + expand, returns None — unchanged behavior.
    - ``"shrink"``: routing-A + LoRA-A shrink only (also pre-warms the routing-B cache so the
      expand stage launches no routing kernels). Returns the shrink ``intermediate`` tensor;
      pass a main(consumer)-stream ``intermediate_buffer`` to control its allocation stream.
    - ``"expand"``: routing-B + LoRA-B expand/add only; requires ``intermediate_buffer`` =
      the tensor produced by the ``"shrink"`` stage.

    ``prewarm_a_routing`` and ``prewarm_b_routing`` let a staged caller skip
    routing for a weight that it replaces with a dense operation. The B flag
    also controls the automatic expand-route prewarm performed by ``"shrink"``.
    ``broadcast_intermediate`` is an expand-only mode where one rank vector per
    token is reused for every routed expert.

    EP accepts either global weights/IDs with a local range or already-localized
    weights/IDs from the standard dispatcher.
    """
    max_loras, _, max_lora_rank, _ = lora_a.shape
    # Global per-expert dim of the LoRA weights. lora_a may be shared-outer (expert
    # dim 1) while lora_b is per-expert, so take the max for the true global count.
    per_expert_dim = max(lora_a.shape[1], lora_b.shape[1])
    ep_local = local_num_experts is not None and local_num_experts < per_expert_dim
    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 _get_shrink_stage_config(weight: torch.Tensor, M: int) -> dict[str, Any]:
        """Heuristic launch config for the LoRA A (shrink) stage.

        Only BLOCK_SIZE_M / num_warps / num_stages are chosen here (a simple
        decode/prefill rule). SPLIT_K is intentionally NOT set: the launcher
        derives it from the post-alignment grid, see _invoke_moe_lora_shrink_splitk.
        M is the number of input tokens; N (= weight.shape[1]) is the LoRA rank.
        """
        N = weight.shape[1]
        if M < 512:  # decode / small batch: pin a small block, latency-bound
            return {
                "BLOCK_SIZE_M": 16,
                "num_warps": 4 if (M <= 4 or N >= 32) else 2,
                "num_stages": 3,
            }
        return {"BLOCK_SIZE_M": 32, "num_warps": 2, "num_stages": 2}  # prefill

    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

        # Shared-outer routing has one bucket per adapter, so the same merged
        # align remains valid for multi-LoRA. Compact EP routing stays single-slot.
        compact_merged = ep_local and not shared_outer and max_loras == 1
        bucket_experts = (
            local_num_experts
            if compact_merged
            else (1 if shared_outer else num_experts) * max_loras
        )
        if (
            lora_envs.SGLANG_OPT_LORA_FUSED_MERGED_ALIGN.get()
            and (shared_outer or compact_merged)
            and bucket_experts + 1 <= 1024
            and topk_ids.shape[0] < 512
        ):
            from sglang.kernels.ops.moe.trtllm_lora_temp.moe_lora_merged_align import (
                moe_lora_merged_align,
            )

            (
                sorted_token_ids,
                expert_ids,
                num_tokens_post_padded,
                token_lora_mask,
                _vne,
            ) = moe_lora_merged_align(
                topk_ids,
                token_lora_mapping,
                num_experts,
                shared_outer,
                max_loras,
                block_size,
                local_expert_offset,
                local_num_experts,
                do_skip=True,
                compact=compact_merged,
            )
            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

        virtual_topk_ids, token_lora_mask, virtual_num_experts = (
            _fused_virtual_topk_ids(
                topk_ids,
                token_lora_mapping,
                num_experts,
                shared_outer,
                max_loras,
                local_expert_offset,
                local_num_experts,
            )
        )
        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.
        # Under EP only this rank's experts are populated, so the per-expert stage bounds
        # the non-empty buckets by `local_num_experts * max_loras` owned (expert, lora)
        # combos + 1 sentinel bucket (all non-owned -> -1). The +1 matters: each non-empty
        # bucket adds up to block-1 padding, so omitting it could truncate real tokens.
        num_tokens = topk_ids.numel()
        if ep_local and not shared_outer and local_num_experts < num_experts:
            populated_buckets = local_num_experts * max_loras + 1
        else:
            populated_buckets = virtual_num_experts
        max_nonempty = min(num_tokens, populated_buckets)
        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]
        # Single adapter (max_loras == 1): the per-adapter virtual-expert offset
        # (safe_lora * num_experts) is always 0, so align never emits an id >=
        # virtual_num_experts and fused_sanitize_expert_ids is an identity -> skip it.
        if max_loras != 1:
            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.kernels.ops.moe.fused_moe_triton_kernels import (
        invoke_fused_moe_kernel,
    )

    assert stage in (
        "all",
        "shrink",
        "expand",
        "routing",
    ), f"invalid stage {stage!r}"
    if broadcast_intermediate:
        assert stage == "expand"
        assert use_direct_expand_add
        assert intermediate_buffer is not None
        assert intermediate_buffer.ndim == 2
        assert intermediate_buffer.shape[0] == token_lora_mapping.shape[0]
    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]
    b_stage_config = _get_stage_config(lora_b_virtual, 1)

    if stage == "routing":
        # Pre-warm the routing cache on the CALLER'S (main) stream so the
        # side-stream chain performs no allocations. Tensors allocated inside a
        # side-stream context during cuda-graph capture can be pool-reused by
        # later allocations on other streams with no cross-stream guard (the
        # allocator's stream tracking is disabled while capturing), corrupting
        # replays. Routing needs only topk_ids + token_lora_mapping, which are
        # both ready before the side-stream fork, so it can run on main.
        a_cfg = _get_shrink_stage_config(lora_a_virtual, token_lora_mapping.shape[0])
        if lora_envs.SGLANG_OPT_LORA_SHRINK_TUNE.get():
            a_cfg = {**a_cfg, "BLOCK_SIZE_M": 16}
        # Match the actual shrink-stage override below. Without this, callers
        # that admit prefill-shaped batches into a side stream prewarm block-32
        # routing here, then miss the cache when shrink switches to the B-stage
        # block size. The miss allocates routing buffers on the side stream
        # during capture, violating the allocation guarantee of stage='routing'.
        if (
            lora_envs.SGLANG_OPT_LORA_PREFILL_ROUTING_REUSE.get()
            and token_lora_mapping.shape[0] >= 512
        ):
            a_cfg["BLOCK_SIZE_M"] = b_stage_config["BLOCK_SIZE_M"]
        if prewarm_a_routing:
            _get_routing(
                topk_ids,
                token_lora_mapping,
                num_experts_a,
                experts_shared_outer_loras_a,
                a_cfg["BLOCK_SIZE_M"],
            )
        if prewarm_b_routing:
            _get_routing(
                topk_ids,
                token_lora_mapping,
                num_experts_b,
                experts_shared_outer_loras_b,
                b_stage_config["BLOCK_SIZE_M"],
            )
        return None

    intermediate = intermediate_buffer
    if stage != "expand":
        a_stage_config = _get_shrink_stage_config(
            lora_a_virtual, token_lora_mapping.shape[0]
        )
        if lora_envs.SGLANG_OPT_LORA_SHRINK_TUNE.get():
            # Test-only override; the launcher fixes N and K block sizes.
            a_stage_config = {
                **a_stage_config,
                "BLOCK_SIZE_M": 16,
                "num_warps": 4,
                "num_stages": 4,
            }
        # Match routing block sizes so prefill stages can share cached alignment.
        if (
            lora_envs.SGLANG_OPT_LORA_PREFILL_ROUTING_REUSE.get()
            and token_lora_mapping.shape[0] >= 512
        ):
            a_stage_config["BLOCK_SIZE_M"] = b_stage_config["BLOCK_SIZE_M"]
        (
            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"],
        )
        intermediate_shape = [
            token_lora_mapping.shape[0],
            topk_ids.shape[1],
            max_lora_rank,
        ]
        intermediate_split_k = _get_moe_lora_shrink_split_k(
            lora_a_virtual, sorted_token_ids, a_stage_config
        )
        # EP leaves non-owned [token, k] shrink slots unwritten. A per-expert expand skips
        # non-owned blocks (never reads them), but a shared-outer expand routes by lora id
        # and would read them into the real (all-reduced) output -> must zero. split_k > 1
        # also needs a zeroed buffer for its accumulation.
        must_zero_intermediate = (
            zero_intermediate
            or intermediate_split_k > 1
            or (ep_local and experts_shared_outer_loras_b)
        )
        if intermediate is None:
            intermediate = (
                torch.zeros(
                    intermediate_shape,
                    dtype=hidden_states.dtype,
                    device=hidden_states.device,
                )
                if must_zero_intermediate
                else torch.empty(
                    intermediate_shape,
                    dtype=hidden_states.dtype,
                    device=hidden_states.device,
                )
            )
        elif must_zero_intermediate:
            # Caller-provided buffer (allocated on the consumer stream): zero it in-stream.
            intermediate.zero_()

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

        if stage == "shrink":
            # Pre-warm the routing-B cache on this (side) stream so the later "expand" stage
            # launches no routing kernels — they overlap finalize together with the shrink.
            if routing_cache is not None and prewarm_b_routing:
                _get_routing(
                    topk_ids,
                    token_lora_mapping,
                    num_experts_b,
                    experts_shared_outer_loras_b,
                    b_stage_config["BLOCK_SIZE_M"],
                )
            return intermediate

    assert intermediate is not None, "stage='expand' requires intermediate_buffer"
    (
        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"],
    )

    intermediate_flat = intermediate.view(-1, max_lora_rank)
    # Shared-add overlap: the expand below adds into the SAME output buffer the
    # overlapped shared-expert add (enqueued on the main stream by the LoRA
    # dispatch) is writing — wait for it here, right before the expand launch,
    # so the shrink + stage-B routing above still overlap the add.
    if expand_wait_event is not None:
        torch.cuda.current_stream().wait_event(expand_wait_event)
    # The rank-specialized direct expand-add doesn't support the shared-outer
    # LoRA-B layout; fall back to the generic kernel for that case instead of
    # asserting, so adapters with experts_shared_outer_loras still work.
    if use_direct_expand_add and not experts_shared_outer_loras_b:
        assert not fuse_add_to_output
        _invoke_moe_lora_expand_add(
            intermediate_flat,
            lora_b_virtual,
            output,
            topk_weights,
            topk_ids,
            sorted_token_ids,
            expert_ids,
            num_tokens_post_padded,
            b_stage_config,
            mul_routed_weight,
            fuse_sum_all_reduce,
            broadcast_intermediate=broadcast_intermediate,
        )
    else:
        assert not broadcast_intermediate, (
            "broadcasted LoRA-A intermediates require the rank-specialized "
            "direct expand kernel"
        )
        invoke_fused_moe_kernel(
            intermediate_flat,
            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=fuse_add_to_output,
            fuse_sum_all_reduce=fuse_sum_all_reduce,
            lora_preserve_base=True,
            add_output_mask=token_lora_mask,
            mask_output=not fuse_add_to_output and not fuse_sum_all_reduce,
            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,
    fuse_add_to_output: bool = True,
    fuse_sum_all_reduce: bool = False,
    use_direct_expand_add: bool = False,
    local_expert_offset: int = 0,
    local_num_experts: int | None = None,
    stage: str = "all",
    intermediate_buffer: torch.Tensor | None = None,
    expand_wait_event: "torch.cuda.Event | None" = None,
    broadcast_intermediate: bool = False,
    prewarm_a_routing: bool = True,
    prewarm_b_routing: bool = True,
    zero_intermediate: bool = False,
) -> "torch.Tensor | None":
    """Public API: wraps the registered op with routing_cache support."""
    return _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,
        fuse_add_to_output=fuse_add_to_output,
        fuse_sum_all_reduce=fuse_sum_all_reduce,
        use_direct_expand_add=use_direct_expand_add,
        local_expert_offset=local_expert_offset,
        local_num_experts=local_num_experts,
        stage=stage,
        intermediate_buffer=intermediate_buffer,
        expand_wait_event=expand_wait_event,
        broadcast_intermediate=broadcast_intermediate,
        prewarm_a_routing=prewarm_a_routing,
        prewarm_b_routing=prewarm_b_routing,
        zero_intermediate=zero_intermediate,
    )
