sgl-project--sglang/python/sglang/srt/layers/moe/ep_moe/kernels.py

import logging
from typing import Tuple

import torch
import triton

from sglang.srt.utils import ceil_div, is_cuda, is_musa

logger = logging.getLogger(__name__)

_is_cuda = is_cuda()
_is_musa = is_musa()

if _is_cuda or _is_musa:
    from sglang.srt.layers.quantization.fp8_kernel import (
        sglang_per_token_group_quant_fp8 as per_token_group_quant_fp8,
    )

import triton.language as tl


def _get_launch_config_1d(device, numel):
    MAX_THREADS_PER_BLOCK = 1024
    MIN_THREADS_PER_BLOCK = 512
    MAX_WAVES = 8  # empirical numbers

    props = torch.cuda.get_device_properties(device)
    sm_count = props.multi_processor_count
    max_threads_per_sm = props.max_threads_per_multi_processor
    max_num_blocks = sm_count * max_threads_per_sm // MAX_THREADS_PER_BLOCK

    block_dim = MAX_THREADS_PER_BLOCK

    def get_num_blocks(block_dim):
        return triton.cdiv(numel, block_dim)

    while (
        block_dim > MIN_THREADS_PER_BLOCK
        and get_num_blocks(block_dim // 2) <= max_num_blocks
    ):
        block_dim = block_dim // 2

    num_blocks = get_num_blocks(block_dim)
    grid_dim = min(num_blocks, max_num_blocks * MAX_WAVES)

    return (grid_dim,), block_dim


def _get_launch_config_2d(device, m, n):
    MAX_THREADS_PER_BLOCK = 1024
    MIN_THREADS_PER_BLOCK = 512
    MAX_WAVES = 8  # empirical numbers

    props = torch.cuda.get_device_properties(device)
    sm_count = props.multi_processor_count
    max_threads_per_sm = props.max_threads_per_multi_processor
    max_num_blocks = sm_count * max_threads_per_sm // MAX_THREADS_PER_BLOCK

    block_dim = MAX_THREADS_PER_BLOCK

    def get_num_blocks(block_dim):
        return m * triton.cdiv(n, block_dim)

    while (
        block_dim > MIN_THREADS_PER_BLOCK
        and get_num_blocks(block_dim // 2) <= max_num_blocks
    ):
        block_dim = block_dim // 2

    grid_dim_x = triton.cdiv(n, block_dim)
    grid_dim_y = max(min(m, max_num_blocks * MAX_WAVES // grid_dim_x), 1)

    return (grid_dim_y, grid_dim_x), block_dim


@triton.jit
def deepep_permute_triton_kernel(
    input_ptr,
    gateup_input_ptr,
    src2dst_ptr,
    topk_ids_ptr,
    a1_scales_ptr,
    topk,
    hidden_size,
    BLOCK_SIZE: tl.constexpr,
):
    OutDtype = gateup_input_ptr.dtype.element_ty

    src_idx = tl.program_id(0)
    src2dst_ptr = src2dst_ptr + src_idx * topk
    topk_ids_ptr = topk_ids_ptr + src_idx * topk

    src_ptr = input_ptr + src_idx * hidden_size

    for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
        offset = start_offset + tl.arange(0, BLOCK_SIZE)
        mask = offset < hidden_size
        in_data = tl.load(src_ptr + offset, mask=mask).to(OutDtype)

        for idx in range(topk):
            dst_idx = tl.load(src2dst_ptr + idx)
            if dst_idx >= 0:
                dst_ptr = gateup_input_ptr + dst_idx * hidden_size
                tl.store(dst_ptr + offset, in_data, mask=mask)


@triton.jit
def deepep_post_reorder_triton_kernel(
    down_output_ptr,
    output_ptr,
    src2dst_ptr,
    topk_ids_ptr,
    topk_weights_ptr,
    topk,
    hidden_size,
    BLOCK_SIZE: tl.constexpr,
):
    InDtype = down_output_ptr.dtype.element_ty

    src_idx = tl.program_id(0)
    src2dst_ptr = src2dst_ptr + src_idx * topk
    topk_ids_ptr = topk_ids_ptr + src_idx * topk
    topk_weights_ptr = topk_weights_ptr + src_idx * topk

    store_ptr = output_ptr + src_idx * hidden_size
    for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
        offset = start_offset + tl.arange(0, BLOCK_SIZE)
        mask = offset < hidden_size
        sum_vec = tl.zeros([BLOCK_SIZE], dtype=InDtype)
        for idx in range(topk):
            dst_idx = tl.load(src2dst_ptr + idx)
            if dst_idx >= 0:
                weigh_scale = tl.load(topk_weights_ptr + idx).to(InDtype)
                load_ptr = down_output_ptr + dst_idx * hidden_size
                in_data = tl.load(load_ptr + offset, mask=mask)
                sum_vec += in_data * weigh_scale
        tl.store(store_ptr + offset, sum_vec, mask=mask)


@triton.jit
def compute_src2dst_triton_kernel(
    reorder_ids, src2dst, num_toks, BLOCK_SIZE: tl.constexpr
):
    pid = tl.program_id(axis=0)
    dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = dst_id < num_toks
    src_id = tl.load(reorder_ids + dst_id, mask=mask)
    tl.store(src2dst + src_id, dst_id, mask=mask)


@triton.jit
def deepep_compute_src2dst_triton_kernel(
    reorder_ids, src2dst, num_toks, num_minus_one, BLOCK_SIZE: tl.constexpr
):
    pid = tl.program_id(axis=0)
    dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = dst_id < num_toks
    src_id = tl.load(reorder_ids + dst_id, mask=mask)
    num_invalid = tl.load(num_minus_one)
    tl.store(src2dst + src_id, dst_id - num_invalid, mask=mask)


def deepep_run_moe_deep_preprocess(topk_ids: torch.Tensor, num_experts: int):
    reorder_topk_ids, reorder_ids = torch.sort(topk_ids.view(-1), stable=True)
    seg_indptr = torch.empty(num_experts + 1, device=topk_ids.device, dtype=torch.int64)
    src2dst = torch.empty(topk_ids.numel(), device=topk_ids.device, dtype=torch.int64)

    # Find offset
    expert_ids = torch.arange(
        num_experts + 1, device=topk_ids.device, dtype=reorder_topk_ids.dtype
    )
    torch.searchsorted(reorder_topk_ids, expert_ids, out=seg_indptr)
    num_minus_one = seg_indptr[0]
    seg_indptr = seg_indptr - num_minus_one

    BLOCK_SIZE = 512
    grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),)
    deepep_compute_src2dst_triton_kernel[grid](
        reorder_ids, src2dst, topk_ids.numel(), num_minus_one, BLOCK_SIZE
    )
    reorder_topk_ids = reorder_topk_ids[num_minus_one:]
    return reorder_topk_ids, src2dst, seg_indptr


def cutlass_w4_run_moe_ep_preproess(topk_ids: torch.Tensor):
    _, reorder_ids = torch.sort(topk_ids.view(-1), stable=True)

    BLOCK_SIZE = 512
    grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),)
    src2dst = torch.empty(topk_ids.numel(), device=topk_ids.device, dtype=torch.int32)
    compute_src2dst_triton_kernel[grid](
        reorder_ids, src2dst, topk_ids.numel(), BLOCK_SIZE
    )

    return src2dst


@triton.jit
def pre_reorder_triton_kernel_for_cutlass_moe(
    input_ptr,
    gateup_input_ptr,
    src2dst_ptr,
    topk_ids_ptr,
    a1_scales_ptr,
    num_local_experts,
    topk,
    num_tokens,
    hidden_size,
    BLOCK_SIZE: tl.constexpr,
    NUM_STAGES: tl.constexpr,
):
    OutDtype = gateup_input_ptr.dtype.element_ty

    if a1_scales_ptr is not None:
        a1_scale = 1.0 / tl.load(a1_scales_ptr)
    else:
        a1_scale = 1.0

    offset = BLOCK_SIZE * tl.program_id(1) + tl.arange(0, BLOCK_SIZE)
    mask = offset < hidden_size

    start_src_idx = tl.program_id(0)
    step = tl.num_programs(0)

    for src_idx_int32 in tl.range(
        start_src_idx, num_tokens, step, num_stages=NUM_STAGES
    ):
        src_idx = src_idx_int32.to(tl.int64)
        token_src2dst_ptr = src2dst_ptr + src_idx * topk
        token_topk_ids_ptr = topk_ids_ptr + src_idx * topk

        src_ptr_offs = input_ptr + src_idx * hidden_size + offset
        dst_ptr_offs = gateup_input_ptr + offset
        in_data = tl.load(src_ptr_offs, mask=mask).to(tl.float32)
        out_data = (in_data * a1_scale).to(OutDtype)
        for idx in range(topk):
            expert_id = tl.load(token_topk_ids_ptr + idx)
            if expert_id != num_local_experts:
                dst_idx = tl.load(token_src2dst_ptr + idx)
                tl.store(dst_ptr_offs + dst_idx * hidden_size, out_data, mask=mask)


def pre_reorder_for_cutlass_moe(
    input,
    gateup_input,
    src2dst,
    topk_ids,
    a1_scales,
    num_local_experts,
    topk,
    num_tokens,
    hidden_size,
):
    grid, block_dim = _get_launch_config_2d(input.device, num_tokens, hidden_size)

    pre_reorder_triton_kernel_for_cutlass_moe[grid](
        input_ptr=input,
        gateup_input_ptr=gateup_input,
        src2dst_ptr=src2dst,
        topk_ids_ptr=topk_ids,
        a1_scales_ptr=a1_scales,
        num_local_experts=num_local_experts,
        topk=topk,
        num_tokens=num_tokens,
        hidden_size=hidden_size,
        BLOCK_SIZE=block_dim,
        NUM_STAGES=3,
    )


# copy from https://github.com/ModelTC/lightllm/blob/a000ab69098654df4731f5b12587dd4e7f0a4f41/lightllm/common/fused_moe/moe_silu_and_mul_mix_quant_ep.py
@triton.jit
def _silu_and_mul_post_quant_kernel(
    input_ptr,
    stride_input_0,
    stride_input_1,
    stride_input_2,
    output_ptr,
    stride_output_0,
    stride_output_1,
    stride_output_2,
    output_scale_ptr,
    stride_output_scale_0,
    stride_output_scale_1,
    stride_output_scale_2,
    masked_m_ptr,
    size_n,
    fp8_max,
    fp8_min,
    QUANT_GROUP_SIZE: tl.constexpr,
    BLOCK_N: tl.constexpr,
    NUM_STAGE: tl.constexpr,
    SCALE_UE8M0: tl.constexpr,
    GEMM1_ALPHA: tl.constexpr,
    GEMM1_CLAMP_LIMIT: tl.constexpr,
):
    expert_id = tl.program_id(2)
    token_id = tl.program_id(1)
    hidden_dim_block_index = tl.program_id(0)

    block_num_per_expert = tl.num_programs(1)

    token_num_cur_expert = tl.load(masked_m_ptr + expert_id)

    stride_input_0 = tl.cast(stride_input_0, dtype=tl.int64)
    stride_output_0 = tl.cast(stride_output_0, dtype=tl.int64)
    stride_input_1 = tl.cast(stride_input_1, dtype=tl.int64)
    stride_output_1 = tl.cast(stride_output_1, dtype=tl.int64)

    N_GROUPS: tl.constexpr = BLOCK_N // QUANT_GROUP_SIZE

    offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N)
    input_ptr_offs = input_ptr + expert_id * stride_input_0 + offs_in_d
    output_ptr_offs = output_ptr + expert_id * stride_output_0 + offs_in_d
    scale_base = (
        output_scale_ptr
        + expert_id * stride_output_scale_0
        + hidden_dim_block_index * N_GROUPS * stride_output_scale_2
    )

    for token_index in tl.range(
        token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE
    ):
        gate = tl.load(
            input_ptr_offs + token_index * stride_input_1,
            mask=offs_in_d < size_n,
            other=0.0,
        ).to(tl.float32)
        up = tl.load(
            input_ptr_offs + token_index * stride_input_1 + size_n,
            mask=offs_in_d < size_n,
            other=0.0,
        )
        if GEMM1_ALPHA > 0:
            gate = tl.minimum(gate, GEMM1_CLAMP_LIMIT)
            up = tl.clamp(up, -GEMM1_CLAMP_LIMIT, GEMM1_CLAMP_LIMIT)
            gate_up = gate * tl.sigmoid(gate * GEMM1_ALPHA) * (up + 1)
        else:
            gate = gate / (1 + tl.exp(-gate))
            gate = gate.to(input_ptr.dtype.element_ty)
            gate_up = up * gate

        gate_up_2d = tl.reshape(gate_up, (N_GROUPS, QUANT_GROUP_SIZE))
        group_absmax = tl.max(tl.abs(gate_up_2d), axis=1)
        group_absmax = tl.maximum(group_absmax, 1e-10)

        output_s = group_absmax / fp8_max
        if SCALE_UE8M0:
            output_s = tl.exp2(tl.ceil(tl.log2(output_s)))

        inv_s = tl.reshape(1.0 / output_s, (N_GROUPS, 1))
        output_q_2d = tl.clamp(gate_up_2d * inv_s, fp8_min, fp8_max)
        output_q = tl.reshape(output_q_2d, (BLOCK_N,)).to(output_ptr.dtype.element_ty)

        tl.store(
            output_ptr_offs + token_index * stride_output_1,
            output_q,
            mask=offs_in_d < size_n,
        )
        scale_offs = scale_base + token_index * stride_output_scale_1
        tl.store(
            scale_offs + tl.arange(0, N_GROUPS) * stride_output_scale_2,
            output_s,
        )


def silu_and_mul_masked_post_quant_fwd(
    input: torch.Tensor,
    output: torch.Tensor,
    output_scale: torch.Tensor,
    quant_group_size: int,
    masked_m: torch.Tensor,
    scale_ue8m0: bool = False,
    gemm1_alpha: float = 0.0,
    gemm1_clamp_limit: float = 0.0,
):
    """
    input shape [expert_num, token_num_padded, hidden_dim]
    output shape [expert_num, token_num_padded, hidden_dim // 2], dtype fp8
    output_scale [expert_num token_num_paddded, hidden_dim // 2 // 128] dtype float32
    quant_group_size  int,
    masked_m shape [expert_num],
    """

    assert input.is_contiguous()
    assert output.dtype == torch.float8_e4m3fn
    assert output.is_contiguous()
    assert len(input.shape) == 3
    assert input.shape[0] == masked_m.shape[0]
    assert input.shape[-1] % 2 == 0

    size_n = input.shape[-1] // 2
    assert size_n % quant_group_size == 0

    expert_num = len(masked_m)

    if expert_num < 4:
        BLOCK_NUM_PER_EXPERT = 64
    else:
        BLOCK_NUM_PER_EXPERT = 32

    groups_total = size_n // quant_group_size
    gpb = 4
    while gpb > 1:
        block_n = quant_group_size * gpb
        if (block_n & (block_n - 1) == 0) and (groups_total % gpb == 0):
            break
        gpb //= 2
    BLOCK_N = quant_group_size * gpb

    num_warps = 1
    NUM_STAGES = 6
    hidden_dim_split_block_num = triton.cdiv(size_n, BLOCK_N)

    grid = (
        hidden_dim_split_block_num,
        BLOCK_NUM_PER_EXPERT,
        expert_num,
    )

    finfo = torch.finfo(torch.float8_e4m3fn)
    fp8_max = finfo.max
    fp8_min = -fp8_max

    _silu_and_mul_post_quant_kernel[grid](
        input,
        *input.stride(),
        output,
        *output.stride(),
        output_scale,
        *output_scale.stride(),
        masked_m,
        size_n,
        fp8_max,
        fp8_min,
        QUANT_GROUP_SIZE=quant_group_size,
        BLOCK_N=BLOCK_N,
        NUM_STAGE=NUM_STAGES,
        num_warps=num_warps,
        SCALE_UE8M0=scale_ue8m0,
        GEMM1_ALPHA=gemm1_alpha if gemm1_alpha is not None else 0.0,
        GEMM1_CLAMP_LIMIT=gemm1_clamp_limit if gemm1_clamp_limit is not None else 0.0,
    )
    return


@triton.jit
def _silu_and_mul_post_quant_packed_kernel(
    input_ptr,
    stride_input_0,
    stride_input_1,
    stride_input_2,
    output_ptr,
    stride_output_0,
    stride_output_1,
    stride_output_2,
    scale_ptr,  # int32 [E, G//4, m_max] (MN-major: packed scale stored token-minor)
    stride_scale_e,
    stride_scale_g4,
    stride_scale_m,
    masked_m_ptr,
    num_experts,
    size_n,
    fp8_max,
    fp8_min,
    QUANT_GROUP_SIZE: tl.constexpr,
    BLOCK_N: tl.constexpr,  # == 4 * QUANT_GROUP_SIZE (one packed int32 per block)
    GEMM1_ALPHA: tl.constexpr,
    GEMM1_CLAMP_LIMIT: tl.constexpr,
    E_PADDED: tl.constexpr,
):
    # Flat work dim rides axis 0 (x): num_real_tokens*topk can exceed the 65535 grid.y/z limit.
    work_id = tl.program_id(0)
    hidden_dim_block_index = tl.program_id(1)

    e_off = tl.arange(0, E_PADDED)
    mm = tl.load(masked_m_ptr + e_off, mask=e_off < num_experts, other=0)
    incl = tl.cumsum(mm)
    total = tl.sum(mm)
    if work_id >= total:
        return
    excl = incl - mm  # first global slot of each expert
    owner = (excl <= work_id) & (work_id < incl)
    expert_id = tl.sum(tl.where(owner, e_off, 0))
    token_index = work_id - tl.sum(tl.where(owner, excl, 0))

    stride_input_0 = tl.cast(stride_input_0, tl.int64)
    stride_input_1 = tl.cast(stride_input_1, tl.int64)
    stride_output_0 = tl.cast(stride_output_0, tl.int64)
    stride_output_1 = tl.cast(stride_output_1, tl.int64)

    N_GROUPS: tl.constexpr = BLOCK_N // QUANT_GROUP_SIZE

    offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N)
    mask_d = offs_in_d < size_n
    in_base = input_ptr + expert_id * stride_input_0 + token_index * stride_input_1
    out_base = output_ptr + expert_id * stride_output_0 + token_index * stride_output_1

    gate = tl.load(in_base + offs_in_d, mask=mask_d, other=0.0).to(tl.float32)
    up = tl.load(in_base + offs_in_d + size_n, mask=mask_d, other=0.0)
    if GEMM1_ALPHA > 0:
        gate = tl.minimum(gate, GEMM1_CLAMP_LIMIT)
        up = tl.clamp(up, -GEMM1_CLAMP_LIMIT, GEMM1_CLAMP_LIMIT)
        gate_up = gate * tl.sigmoid(gate * GEMM1_ALPHA) * (up + 1)
    else:
        gate = gate / (1 + tl.exp(-gate))
        gate = gate.to(input_ptr.dtype.element_ty)
        gate_up = up * gate

    gate_up_2d = tl.reshape(gate_up, (N_GROUPS, QUANT_GROUP_SIZE))
    group_absmax = tl.max(tl.abs(gate_up_2d), axis=1)
    group_absmax = tl.maximum(group_absmax, 1e-10)
    output_s = group_absmax / fp8_max
    # UE8M0: round to a power of two so the (>>23)&0xFF exponent extraction below is valid.
    output_s = tl.exp2(tl.ceil(tl.log2(output_s)))

    inv_s = tl.reshape(1.0 / output_s, (N_GROUPS, 1))
    output_q_2d = tl.clamp(gate_up_2d * inv_s, fp8_min, fp8_max)
    output_q = tl.reshape(output_q_2d, (BLOCK_N,)).to(output_ptr.dtype.element_ty)
    tl.store(out_base + offs_in_d, output_q, mask=mask_d)

    # Pack 4 UE8M0 exponent bytes little-endian into one int32, matching deep_gemm's
    # get_mn_major_tma_aligned_packed_ue8m0_tensor (fp32>>23 -> uint8 -> 4-group int32 view).
    s_bits = output_s.to(tl.int32, bitcast=True)
    expo = (s_bits >> 23) & 0xFF
    shifts = tl.arange(0, N_GROUPS) * 8
    packed = tl.sum(expo << shifts)
    scale_off = (
        expert_id * stride_scale_e
        + hidden_dim_block_index * stride_scale_g4
        + token_index * stride_scale_m
    )
    tl.store(scale_ptr + scale_off, packed)


def silu_and_mul_masked_post_quant_packed_fwd(
    input: torch.Tensor,
    output: torch.Tensor,
    output_scale_packed: torch.Tensor,
    quant_group_size: int,
    masked_m: torch.Tensor,
    num_real_tokens: int,
    topk: int,
    gemm1_alpha: float = 0.0,
    gemm1_clamp_limit: float = 0.0,
):
    assert input.is_contiguous()
    assert output.dtype == torch.float8_e4m3fn
    assert output.is_contiguous()
    assert input.dim() == 3 and input.shape[-1] % 2 == 0

    E, m_max, _ = input.shape
    size_n = input.shape[-1] // 2
    assert size_n % quant_group_size == 0
    G = size_n // quant_group_size
    assert G % 4 == 0, "packed UE8M0 path requires num_groups % 4 == 0"
    BLOCK_N = quant_group_size * 4
    assert size_n % BLOCK_N == 0, "packed UE8M0 path requires size_n % (4*group) == 0"
    hidden_dim_split = size_n // BLOCK_N
    assert tuple(output_scale_packed.shape) == (E, hidden_dim_split, m_max)
    assert output_scale_packed.dtype == torch.int32

    finfo = torch.finfo(torch.float8_e4m3fn)
    fp8_max = finfo.max

    grid = (num_real_tokens * topk, hidden_dim_split)
    _silu_and_mul_post_quant_packed_kernel[grid](
        input,
        *input.stride(),
        output,
        *output.stride(),
        output_scale_packed,
        *output_scale_packed.stride(),
        masked_m,
        E,
        size_n,
        fp8_max,
        -fp8_max,
        QUANT_GROUP_SIZE=quant_group_size,
        BLOCK_N=BLOCK_N,
        GEMM1_ALPHA=gemm1_alpha if gemm1_alpha is not None else 0.0,
        GEMM1_CLAMP_LIMIT=gemm1_clamp_limit if gemm1_clamp_limit is not None else 0.0,
        E_PADDED=triton.next_power_of_2(E),
        num_warps=1,
    )
    return


@triton.jit
def _silu_and_mul_kernel(
    input_ptr,
    stride_input_0,
    stride_input_1,
    stride_input_2,
    output_ptr,
    stride_output_0,
    stride_output_1,
    stride_output_2,
    masked_m_ptr,
    size_n,
    BLOCK_N: tl.constexpr,
    NUM_STAGE: tl.constexpr,
):
    expert_id = tl.program_id(2)
    token_id = tl.program_id(1)
    hidden_dim_block_index = tl.program_id(0)

    block_num_per_expert = tl.num_programs(1)

    token_num_cur_expert = tl.load(masked_m_ptr + expert_id)

    stride_input_0 = tl.cast(stride_input_0, dtype=tl.int64)
    stride_output_0 = tl.cast(stride_output_0, dtype=tl.int64)
    stride_input_1 = tl.cast(stride_input_1, dtype=tl.int64)
    stride_output_1 = tl.cast(stride_output_1, dtype=tl.int64)

    offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N)
    input_ptr_offs = input_ptr + expert_id * stride_input_0 + offs_in_d
    output_ptr_offs = output_ptr + expert_id * stride_output_0 + offs_in_d

    for token_index in tl.range(
        token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE
    ):
        gate = tl.load(
            input_ptr_offs + token_index * stride_input_1,
            mask=offs_in_d < size_n,
            other=0.0,
        ).to(tl.float32)
        up = tl.load(
            input_ptr_offs + token_index * stride_input_1 + size_n,
            mask=offs_in_d < size_n,
            other=0.0,
        ).to(tl.float32)
        gate = gate / (1 + tl.exp(-gate))
        gate_up = up * gate
        # Compute SiLU in fp32 for better precision, then cast back to the
        # input dtype.
        gate_up = gate_up.to(input_ptr.dtype.element_ty)
        tl.store(
            output_ptr_offs + token_index * stride_output_1,
            gate_up,
            mask=offs_in_d < size_n,
        )


def silu_and_mul_masked_fwd(
    input: torch.Tensor,
    output: torch.Tensor,
    masked_m: torch.Tensor,
):
    """
    input shape [expert_num, token_num_padded, hidden_dim], dtype bf16
    output shape [expert_num, token_num_padded, hidden_dim // 2], dtype bf16
    masked_m shape [expert_num]
    """

    assert input.is_contiguous()
    assert output.dtype == torch.bfloat16
    assert input.dtype == torch.bfloat16
    assert output.is_contiguous()
    assert len(input.shape) == 3
    assert input.shape[0] == masked_m.shape[0]
    assert input.shape[-1] % 2 == 0

    size_n = input.shape[-1] // 2
    expert_num = len(masked_m)

    if expert_num < 4:
        BLOCK_NUM_PER_EXPERT = 64
    else:
        BLOCK_NUM_PER_EXPERT = 32

    BLOCK_N = 128
    num_warps = 4
    NUM_STAGES = 4

    hidden_dim_split_block_num = triton.cdiv(size_n, BLOCK_N)

    grid = (
        hidden_dim_split_block_num,
        BLOCK_NUM_PER_EXPERT,
        expert_num,
    )

    _silu_and_mul_kernel[grid](
        input,
        *input.stride(),
        output,
        *output.stride(),
        masked_m,
        size_n,
        BLOCK_N=BLOCK_N,
        NUM_STAGE=NUM_STAGES,
        num_warps=num_warps,
    )
    return output


@triton.jit
def silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe(
    input_ptr,
    output_ptr,
    scale_ptr,
    num_tokens_tensor_ptr,
    intermediate_size,
    BLOCK_SIZE: tl.constexpr,
    NUM_STAGES: tl.constexpr,
):
    OutDtype = output_ptr.dtype.element_ty

    num_tokens = tl.load(num_tokens_tensor_ptr)
    numel = num_tokens * intermediate_size
    gate_ptr = input_ptr
    up_ptr = input_ptr + intermediate_size
    scale = 1.0 / tl.load(scale_ptr)

    start_idx = tl.program_id(0) * BLOCK_SIZE
    step = tl.num_programs(0) * BLOCK_SIZE

    for id in tl.range(start_idx, numel, step, num_stages=NUM_STAGES):
        ids = id + tl.arange(0, BLOCK_SIZE)
        token_ids = ids // intermediate_size
        mask = ids < numel

        offs = ids + token_ids * intermediate_size
        gate = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32)
        up = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32)
        output = gate / (1 + tl.exp(-gate)) * up * scale
        tl.store(output_ptr + ids, output.to(OutDtype), mask=mask)


def silu_mul_static_tensorwise_quant_for_cutlass_moe(
    input: torch.Tensor,
    output: torch.Tensor,
    scale: torch.Tensor,
    num_tokens_tensor: torch.Tensor,
    expected_num_tokens: int,
    intermediate_size: int,
):
    grid, block_dim = _get_launch_config_1d(
        input.device, expected_num_tokens * intermediate_size
    )

    silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe[grid](
        input_ptr=input,
        output_ptr=output,
        scale_ptr=scale,
        num_tokens_tensor_ptr=num_tokens_tensor,
        intermediate_size=intermediate_size,
        BLOCK_SIZE=block_dim,
        NUM_STAGES=3,
    )


@triton.jit
def post_reorder_triton_kernel_for_cutlass_moe(
    down_output_ptr,
    output_ptr,
    src2dst_ptr,
    topk_ids_ptr,
    topk_weights_ptr,
    num_local_experts,
    topk,
    num_tokens,
    hidden_size,
    routed_scaling_factor: float,
    BLOCK_SIZE: tl.constexpr,
    NUM_STAGES: tl.constexpr,
):
    OutDtype = output_ptr.dtype.element_ty

    offset = BLOCK_SIZE * tl.program_id(1) + tl.arange(0, BLOCK_SIZE)
    mask = offset < hidden_size

    down_output_ptr_offs = down_output_ptr + offset
    output_ptr_offs = output_ptr + offset

    start_src_idx = tl.program_id(0)
    step = tl.num_programs(0)

    for src_idx_int32 in tl.range(
        start_src_idx, num_tokens, step, num_stages=NUM_STAGES
    ):
        src_idx = src_idx_int32.to(tl.int64)
        token_src2dst_ptr = src2dst_ptr + src_idx * topk
        token_topk_ids_ptr = topk_ids_ptr + src_idx * topk
        token_topk_weights_ptr = topk_weights_ptr + src_idx * topk

        sum_vec = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
        for idx in range(topk):
            expert_id = tl.load(token_topk_ids_ptr + idx)
            if expert_id != num_local_experts:
                dst_idx_int32 = tl.load(token_src2dst_ptr + idx)
                dst_idx = dst_idx_int32.to(tl.int64)
                dst_idx = dst_idx
                weight_scale = tl.load(token_topk_weights_ptr + idx).to(tl.float32)
                load_ptr_offs = down_output_ptr_offs + dst_idx * hidden_size
                in_data = tl.load(load_ptr_offs, mask=mask).to(tl.float32)
                sum_vec += in_data * weight_scale
        sum_vec *= routed_scaling_factor
        store_ptr_offs = output_ptr_offs + src_idx * hidden_size
        tl.store(store_ptr_offs, sum_vec.to(OutDtype), mask=mask)


def post_reorder_for_cutlass_moe(
    down_output,
    output,
    src2dst,
    topk_ids,
    topk_weights,
    num_local_experts,
    topk,
    num_tokens,
    hidden_size,
    routed_scaling_factor: float,
):
    grid, block_dim = _get_launch_config_2d(down_output.device, num_tokens, hidden_size)

    post_reorder_triton_kernel_for_cutlass_moe[grid](
        down_output_ptr=down_output,
        output_ptr=output,
        src2dst_ptr=src2dst,
        topk_ids_ptr=topk_ids,
        topk_weights_ptr=topk_weights,
        num_local_experts=num_local_experts,
        topk=topk,
        num_tokens=num_tokens,
        hidden_size=hidden_size,
        routed_scaling_factor=routed_scaling_factor,
        BLOCK_SIZE=block_dim,
        NUM_STAGES=3,
    )


@triton.jit
def post_reorder_deepgemm_triton_kernel(
    down_output_ptr,
    output_ptr,
    src2dst_ptr,
    topk_ids_ptr,
    topk_weights_ptr,
    topk,
    num_tokens,
    hidden_size,
    routed_scaling_factor: float,
    BLOCK_SIZE: tl.constexpr,
    NUM_STAGES: tl.constexpr,
):
    """`expert_id >= 0` includes the shared expert at num_experts (padding=-1); don't
    switch to the cutlass `!= num_local_experts` gate. routed_scaling_factor is folded into the store.
    """
    OutDtype = output_ptr.dtype.element_ty

    offset = BLOCK_SIZE * tl.program_id(1) + tl.arange(0, BLOCK_SIZE)
    mask = offset < hidden_size

    down_output_ptr_offs = down_output_ptr + offset
    output_ptr_offs = output_ptr + offset

    start_src_idx = tl.program_id(0)
    step = tl.num_programs(0)

    for src_idx_int32 in tl.range(
        start_src_idx, num_tokens, step, num_stages=NUM_STAGES
    ):
        src_idx = src_idx_int32.to(tl.int64)
        token_src2dst_ptr = src2dst_ptr + src_idx * topk
        token_topk_ids_ptr = topk_ids_ptr + src_idx * topk
        token_topk_weights_ptr = topk_weights_ptr + src_idx * topk

        sum_vec = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
        for idx in range(topk):
            expert_id = tl.load(token_topk_ids_ptr + idx)
            if expert_id >= 0:
                dst_idx = tl.load(token_src2dst_ptr + idx).to(tl.int64)
                weight_scale = tl.load(token_topk_weights_ptr + idx).to(tl.float32)
                load_ptr_offs = down_output_ptr_offs + dst_idx * hidden_size
                in_data = tl.load(load_ptr_offs, mask=mask).to(tl.float32)
                sum_vec += in_data * weight_scale
        sum_vec *= routed_scaling_factor
        store_ptr_offs = output_ptr_offs + src_idx * hidden_size
        tl.store(store_ptr_offs, sum_vec.to(OutDtype), mask=mask)


def post_reorder_deepgemm(
    down_output,
    output,
    src2dst,
    topk_ids,
    topk_weights,
    topk,
    num_tokens,
    hidden_size,
    routed_scaling_factor: float,
):
    grid, block_dim = _get_launch_config_2d(down_output.device, num_tokens, hidden_size)
    post_reorder_deepgemm_triton_kernel[grid](
        down_output,
        output,
        src2dst,
        topk_ids,
        topk_weights,
        topk,
        num_tokens,
        hidden_size,
        float(routed_scaling_factor),
        BLOCK_SIZE=block_dim,
        NUM_STAGES=3,
    )


@triton.jit
def post_reorder_triton_kernel(
    down_output_ptr,
    output_ptr,
    src2dst_ptr,
    topk_ids_ptr,
    topk_weights_ptr,
    topk,
    hidden_size,
    BLOCK_SIZE: tl.constexpr,
):
    InDtype = down_output_ptr.dtype.element_ty

    src_idx_int32 = tl.program_id(0)
    src_idx = src_idx_int32.to(tl.int64)
    src2dst_ptr = src2dst_ptr + src_idx * topk
    topk_ids_ptr = topk_ids_ptr + src_idx * topk
    topk_weights_ptr = topk_weights_ptr + src_idx * topk

    store_ptr = output_ptr + src_idx * hidden_size

    vec = tl.arange(0, BLOCK_SIZE)

    for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
        offset = start_offset + vec
        mask = offset < hidden_size

        sum_vec = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
        for idx in range(topk):
            expert_id = tl.load(topk_ids_ptr + idx)
            if expert_id >= 0:
                dst_idx_int32 = tl.load(src2dst_ptr + idx)
                dst_idx = dst_idx_int32.to(tl.int64)
                weigh_scale = tl.load(topk_weights_ptr + idx).to(tl.float32)
                load_ptr = down_output_ptr + dst_idx * hidden_size
                # accumulate expert outputs in fp32 for better precision
                # before casting to the final output dtype.
                in_data = tl.load(load_ptr + offset, mask=mask).to(tl.float32)
                sum_vec += in_data * weigh_scale
        tl.store(store_ptr + offset, sum_vec.to(InDtype), mask=mask)


@triton.jit
def _fwd_kernel_ep_scatter_1(
    num_recv_tokens_per_expert,
    expert_start_loc,
    m_indices,
    num_experts: tl.constexpr,
    BLOCK_E: tl.constexpr,
    BLOCK_EXPERT_NUM: tl.constexpr,
):
    cur_expert = tl.program_id(0)

    offset_cumsum = tl.arange(0, BLOCK_EXPERT_NUM)
    tokens_per_expert = tl.load(
        num_recv_tokens_per_expert + offset_cumsum,
        mask=offset_cumsum < num_experts,
        other=0,
    )
    cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert
    tl.store(expert_start_loc + offset_cumsum, cumsum, mask=offset_cumsum < num_experts)

    cur_expert_start = tl.load(expert_start_loc + cur_expert)
    cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert)

    m_indices_start_ptr = m_indices + cur_expert_start
    off_expert = tl.arange(0, BLOCK_E)

    for start_m in tl.range(0, cur_expert_token_num, BLOCK_E, num_stages=4):
        tl.store(
            m_indices_start_ptr + start_m + off_expert,
            cur_expert,
        )


@triton.jit
def _fwd_kernel_ep_scatter_2(
    total_token_num,
    expert_start_loc,
    recv_x,
    recv_x_stride0,
    recv_x_stride1,
    recv_x_scale,
    recv_x_scale_stride0,
    recv_x_scale_stride1,
    recv_topk,
    recv_topk_stride0,
    recv_topk_stride1,
    output_tensor,
    output_tensor_stride0,
    output_tensor_stride1,
    output_tensor_scale,
    output_tensor_scale_stride0,
    output_tensor_scale_stride1,
    output_index,
    output_index_stride0,
    output_index_stride1,
    topk_num: tl.constexpr,
    HIDDEN_SIZE: tl.constexpr,
    HIDDEN_SIZE_PAD: tl.constexpr,
    SCALE_HIDDEN_SIZE: tl.constexpr,
    SCALE_HIDDEN_SIZE_PAD: tl.constexpr,
    # Platform-specific semaphore for atomic_add performance tuning
    ATOMIC_ADD_SEM: tl.constexpr,
    IS_FP8: tl.constexpr,
):
    start_token_id = tl.program_id(0)
    grid_num = tl.num_programs(0)

    offset_in = tl.arange(0, HIDDEN_SIZE_PAD)
    mask = offset_in < HIDDEN_SIZE

    index_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
    mask_s = index_in_s < SCALE_HIDDEN_SIZE

    for token_id_int32 in range(start_token_id, total_token_num, grid_num):
        token_id = token_id_int32.to(tl.int64)
        to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask)
        if IS_FP8:
            to_copy_s = tl.load(
                recv_x_scale
                + token_id * recv_x_scale_stride0
                + index_in_s * recv_x_scale_stride1,
                mask=mask_s,
            )

        for topk_idx_int32 in tl.range(0, topk_num, 1, num_stages=4):
            topk_index = topk_idx_int32.to(tl.int64)
            expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index)
            if expert_id >= 0:
                dest_token_index_int32 = tl.atomic_add(
                    expert_start_loc + expert_id, 1, sem=ATOMIC_ADD_SEM
                )
                dest_token_index = dest_token_index_int32.to(tl.int64)

                tl.store(
                    output_index + token_id * output_index_stride0 + topk_index,
                    dest_token_index_int32,
                )
                output_tensor_ptr = (
                    output_tensor + dest_token_index * output_tensor_stride0
                )
                tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
                if IS_FP8:
                    output_tensor_scale_ptr = (
                        output_tensor_scale
                        + dest_token_index * output_tensor_scale_stride0
                    )
                    tl.store(
                        output_tensor_scale_ptr
                        + index_in_s * output_tensor_scale_stride1,
                        to_copy_s,
                        mask=mask_s,
                    )


# copy from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/deepep_scatter_gather.py
@torch.no_grad()
def ep_scatter(
    recv_x: torch.Tensor,
    recv_x_scale: torch.Tensor,
    recv_topk: torch.Tensor,
    num_recv_tokens_per_expert: torch.Tensor,
    expert_start_loc: torch.Tensor,
    output_tensor: torch.Tensor,
    output_tensor_scale: torch.Tensor,
    m_indices: torch.Tensor,
    output_index: torch.Tensor,
    scale_ue8m0: bool = False,
    quant_block_size: int = 128,
):
    BLOCK_E = 128  # token num of per expert is aligned to 128
    BLOCK_D = quant_block_size  # block size of quantization
    num_warps = 8
    num_experts = num_recv_tokens_per_expert.shape[0]
    hidden_size = recv_x.shape[1]
    # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts)
    grid = num_experts

    scale_hidden_size = hidden_size // BLOCK_D
    if scale_ue8m0:
        # ue8m0 scales are packed here (4 scales per int32),
        # hence the effective size of this dimension is divided by 4.
        scale_hidden_size = ceil_div(scale_hidden_size, 4)

    assert m_indices.shape[0] % BLOCK_E == 0

    is_fp8 = recv_x_scale is not None and recv_x.dtype != torch.bfloat16
    if is_fp8:
        assert (
            recv_x_scale.dtype == output_tensor_scale.dtype
        ), f"recv_x_scale.dtype: {recv_x_scale.dtype}, output_tensor_scale.dtype: {output_tensor_scale.dtype}"
        assert (
            recv_x_scale.shape[1] == output_tensor_scale.shape[1] == scale_hidden_size
        )

    _fwd_kernel_ep_scatter_1[(grid,)](
        num_recv_tokens_per_expert,
        expert_start_loc,
        m_indices,
        num_experts=num_experts,
        num_warps=num_warps,
        BLOCK_E=BLOCK_E,
        BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts),
    )

    grid = min(recv_topk.shape[0], 1024 * 8)

    _fwd_kernel_ep_scatter_2[(grid,)](
        recv_topk.shape[0],
        expert_start_loc,
        recv_x,
        recv_x.stride(0),
        recv_x.stride(1),
        recv_x_scale,
        recv_x_scale.stride(0) if is_fp8 else 0,
        recv_x_scale.stride(1) if is_fp8 else 0,
        recv_topk,
        recv_topk.stride(0),
        recv_topk.stride(1),
        output_tensor,
        output_tensor.stride(0),
        output_tensor.stride(1),
        output_tensor_scale,
        output_tensor_scale.stride(0) if is_fp8 else 0,
        output_tensor_scale.stride(1) if is_fp8 else 0,
        output_index,
        output_index.stride(0),
        output_index.stride(1),
        topk_num=recv_topk.shape[1],
        num_warps=num_warps,
        HIDDEN_SIZE=hidden_size,
        HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size),
        SCALE_HIDDEN_SIZE=scale_hidden_size,
        SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size),
        # XXX (MUSA): Atomic add with "relaxed" semaphore on musa backend for better performance
        ATOMIC_ADD_SEM=None if not _is_musa else "relaxed",
        IS_FP8=is_fp8,
    )
    return


@triton.jit
def _fwd_kernel_ep_gather(
    total_token_num,
    input_tensor,
    input_tensor_stride0,
    input_tensor_stride1,
    recv_topk_ids,
    recv_topk_ids_stride0,
    recv_topk_ids_stride1,
    recv_topk_weight,
    recv_topk_weight_stride0,
    recv_topk_weight_stride1,
    input_index,
    input_index_stride0,
    input_index_stride1,
    output_tensor,
    output_tensor_stride0,
    output_tensor_stride1,
    topk_num: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    cur_block_int32 = tl.program_id(0)
    cur_block = cur_block_int32.to(tl.int64)

    start_cur_token_int32 = tl.program_id(1)

    grid_num = tl.num_programs(1)

    for cur_token_int32 in range(start_cur_token_int32, total_token_num, grid_num):
        cur_token = cur_token_int32.to(tl.int64)

        off_d = tl.arange(0, BLOCK_D)
        accumulator = tl.zeros([BLOCK_D], dtype=tl.float32)

        for topk_index_int32 in range(0, topk_num):
            topk_index = topk_index_int32.to(tl.int64)

            expert_id = tl.load(
                recv_topk_ids + cur_token * recv_topk_ids_stride0 + topk_index
            )
            if expert_id >= 0:
                source_token_index_int32 = tl.load(
                    input_index + cur_token * input_index_stride0 + topk_index
                )
                source_token_index = source_token_index_int32.to(tl.int64)

                acc_weight = tl.load(
                    recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index
                )
                tmp = tl.load(
                    input_tensor
                    + source_token_index * input_tensor_stride0
                    + cur_block * BLOCK_D
                    + off_d
                )
                accumulator += tmp.to(tl.float32) * acc_weight

        tl.store(
            output_tensor
            + cur_token * output_tensor_stride0
            + cur_block * BLOCK_D
            + off_d,
            accumulator.to(output_tensor.dtype.element_ty),
        )


@torch.no_grad()
def ep_gather(
    input_tensor: torch.Tensor,
    recv_topk_ids: torch.Tensor,
    recv_topk_weight: torch.Tensor,
    input_index: torch.Tensor,
    output_tensor: torch.Tensor,
):
    num_warps = 2
    num_tokens = output_tensor.shape[0]
    hidden_size = input_tensor.shape[1]
    BLOCK_D = 128 if hidden_size % 1024 != 0 else 1024  # block size of quantization
    assert hidden_size % BLOCK_D == 0
    grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024))
    _fwd_kernel_ep_gather[grid](
        num_tokens,
        input_tensor,
        input_tensor.stride(0),
        input_tensor.stride(1),
        recv_topk_ids,
        recv_topk_ids.stride(0),
        recv_topk_ids.stride(1),
        recv_topk_weight,
        recv_topk_weight.stride(0),
        recv_topk_weight.stride(1),
        input_index,
        input_index.stride(0),
        input_index.stride(1),
        output_tensor,
        output_tensor.stride(0),
        output_tensor.stride(1),
        topk_num=recv_topk_ids.shape[1],
        num_warps=num_warps,
        BLOCK_D=BLOCK_D,
    )
    return


# copy from
# https://github.com/deepseek-ai/DeepGEMM/blob/bd2a77552886b98c205af12f8d7d2d61247c4b27/deep_gemm/jit_kernels/utils.py#L58
def get_tma_aligned_size(x: int, element_size: int) -> int:
    """
    Global memory address of TMA must be 16-byte aligned.
    Since we use column-major layout for the LHS scaling tensor,
        the M-axis of the LHS scaling tensor needs to be padded to a multiple of 16 bytes.

    Arguments:
        x: original M-axis shape of the LHS scaling tensor.
        element_size: element size of the LHS scaling tensor.

    Returns:
        M-axis shape of the LHS scaling tensor after padding.
    """
    tma_alignment_bytes = 16
    assert tma_alignment_bytes % element_size == 0
    alignment = tma_alignment_bytes // element_size
    return ceil_div(x, alignment) * alignment


@triton.jit
def _tma_align_input_scale_kernel(
    input_scale_ptr,
    output_ptr,
    m,
    k_div_block_size,
    input_scale_stride_m,
    input_scale_stride_k,
    output_stride_m,
    output_stride_k,
    BLOCK_SIZE_K: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    grid_m = tl.num_programs(0)
    k_offsets = tl.arange(0, BLOCK_SIZE_K)

    for m_base in range(pid_m, m, grid_m):
        input_offset = (
            input_scale_ptr
            + m_base * input_scale_stride_m
            + k_offsets * input_scale_stride_k
        )
        input_data = tl.load(input_offset, mask=k_offsets < k_div_block_size)

        output_offset = (
            output_ptr + k_offsets * output_stride_k + m_base * output_stride_m
        )
        tl.store(output_offset, input_data, mask=k_offsets < k_div_block_size)


# copy from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/quantization/triton_quant/fp8/fp8act_quant_kernel.py
def tma_align_input_scale(input_scale: torch.Tensor):
    assert input_scale.dim() == 2
    m, k_div_block_size = input_scale.shape
    padd_m = get_tma_aligned_size(m, input_scale.element_size())
    output = torch.empty(
        (k_div_block_size, padd_m), dtype=input_scale.dtype, device=input_scale.device
    )

    grid_m = min(m, 8192)
    BLOCK_SIZE_K = triton.next_power_of_2(k_div_block_size)

    _tma_align_input_scale_kernel[(grid_m,)](
        input_scale_ptr=input_scale,
        output_ptr=output,
        m=m,
        k_div_block_size=k_div_block_size,
        input_scale_stride_m=input_scale.stride(0),
        input_scale_stride_k=input_scale.stride(1),
        output_stride_m=output.stride(1),  # Note: these are swapped
        output_stride_k=output.stride(0),  # for column-major
        BLOCK_SIZE_K=BLOCK_SIZE_K,
    )
    return output.t()[:m]


@triton.jit
def fused_moe_dispatch_index_triton_kernel(
    topk_ids_ptr,  # flat (num_toks,) int32; -1 = padding (drives the `expert >= 0` gate)
    src2dst_ptr,
    masked_m_ptr,
    m_max,
    num_toks,
    num_experts,
    BLOCK_SIZE: tl.constexpr,
    ZERO_INIT: tl.constexpr,
):
    # Each token picks top_k distinct experts, so per-expert count <= num_tokens < m_max;
    # dst = expert*m_max + offset never spills into the next expert's region.
    pid = tl.program_id(0)
    if ZERO_INIT:
        # Zero the cursor in-kernel and barrier before any atomic_add (single-block path).
        e_off = tl.arange(0, BLOCK_SIZE)
        tl.store(
            masked_m_ptr + e_off,
            tl.zeros((BLOCK_SIZE,), dtype=tl.int32),
            mask=e_off < num_experts,
        )
        tl.debug_barrier()
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < num_toks
    expert = tl.load(topk_ids_ptr + offs, mask=mask, other=-1)
    valid = mask & (expert >= 0)
    # Clamp masked lanes to bin 0 so the masked atomic's pointer stays in-bounds.
    expert_safe = tl.where(valid, expert, 0)
    offset = tl.atomic_add(masked_m_ptr + expert_safe, 1, mask=valid)
    dst = expert_safe * m_max + offset
    tl.store(src2dst_ptr + offs, dst, mask=valid)


def fused_moe_dispatch_index(
    topk_ids: torch.Tensor,
    num_local_experts: int,
    m_max: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    num_toks = topk_ids.numel()
    src2dst = torch.empty(num_toks, device=topk_ids.device, dtype=torch.int32)
    # masked_m doubles as the atomic cursor and the final per-expert count; must be zeroed before any atomic_add.
    single_block = max(num_toks, num_local_experts) <= 1024
    if single_block:
        BLOCK_SIZE = triton.next_power_of_2(max(num_toks, num_local_experts))
        masked_m = torch.empty(
            num_local_experts, device=topk_ids.device, dtype=torch.int32
        )
        grid = (1,)
    else:
        BLOCK_SIZE = 256
        masked_m = torch.zeros(
            num_local_experts, device=topk_ids.device, dtype=torch.int32
        )
        grid = (triton.cdiv(num_toks, BLOCK_SIZE),)
    fused_moe_dispatch_index_triton_kernel[grid](
        topk_ids.view(-1),
        src2dst,
        masked_m,
        m_max,
        num_toks,
        num_local_experts,
        BLOCK_SIZE=BLOCK_SIZE,
        ZERO_INIT=single_block,
    )
    return masked_m, src2dst


@triton.jit
def fill_gateup_input_triton_kernel(
    input_ptr,
    scale_ptr,
    gateup_input_ptr,
    gateup_input_scale_ptr,
    src2dst_ptr,
    topk_ids_ptr,
    topk,
    hidden_size,
    scale_size,
    m_max,
    scale_row_stride,
    scale_col_stride,
    BLOCK_SIZE: tl.constexpr,
    IS_FP8: tl.constexpr = True,
    SCALE_MN_MAJOR: tl.constexpr = False,
):

    src_idx_int32 = tl.program_id(0)
    src_idx = src_idx_int32.to(tl.int64)
    src2dst_ptr = src2dst_ptr + src_idx * topk
    topk_ids_ptr = topk_ids_ptr + src_idx * topk
    src_ptr = input_ptr + src_idx * hidden_size
    if IS_FP8:
        scale_src_ptr = scale_ptr + src_idx * scale_row_stride

    vec = tl.arange(0, BLOCK_SIZE)
    for idx in range(topk):
        expert_id = tl.load(topk_ids_ptr + idx)
        if expert_id >= 0:
            dst_idx_int32 = tl.load(src2dst_ptr + idx)
            dst_idx = dst_idx_int32.to(tl.int64)
            dst_ptr = gateup_input_ptr + dst_idx * hidden_size
            for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
                offset = start_offset + vec
                mask = offset < hidden_size
                in_data = tl.load(src_ptr + offset, mask=mask)
                tl.store(dst_ptr + offset, in_data, mask=mask)

            if IS_FP8:
                if SCALE_MN_MAJOR:
                    expert = dst_idx // m_max
                    m = dst_idx % m_max
                    scale_dst_ptr = (
                        gateup_input_scale_ptr + expert * scale_size * m_max + m
                    )
                    for start_offset in tl.range(0, scale_size, BLOCK_SIZE):
                        offset = start_offset + vec
                        mask = offset < scale_size
                        in_scale = tl.load(
                            scale_src_ptr + offset * scale_col_stride, mask=mask
                        )
                        tl.store(scale_dst_ptr + offset * m_max, in_scale, mask=mask)
                else:
                    scale_dst_ptr = gateup_input_scale_ptr + dst_idx * scale_size
                    for start_offset in tl.range(0, scale_size, BLOCK_SIZE):
                        offset = start_offset + vec
                        mask = offset < scale_size
                        in_scale = tl.load(
                            scale_src_ptr + offset * scale_col_stride, mask=mask
                        )
                        tl.store(scale_dst_ptr + offset, in_scale, mask=mask)


def moe_ep_deepgemm_preprocess(
    topk_ids: torch.Tensor,
    num_local_experts: int,
    hidden_states: torch.Tensor,
    top_k: int,
    block_shape,
    output_dtype: torch.dtype = torch.float8_e4m3fn,
    use_mxfp8: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    # For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m}) https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/jit_kernels/m_grouped_gemm.py#L165
    m_max = (hidden_states.size(0) // 256 + 1) * 256
    expected_m = (topk_ids.numel() - 1) // num_local_experts + 1

    masked_m, src2dst = fused_moe_dispatch_index(topk_ids, num_local_experts, m_max)

    gateup_input = torch.empty(
        (num_local_experts, m_max, hidden_states.size(1)),
        device=hidden_states.device,
        dtype=output_dtype,
    )

    if block_shape is None:
        block_shape = [128, 128]
    assert len(block_shape) == 2
    block_n, block_k = block_shape[0], block_shape[1]
    is_fp8 = output_dtype == torch.float8_e4m3fn
    if is_fp8 and use_mxfp8:
        from sglang.jit_kernel.minimax_quant_ue8m0 import (
            per_token_quant_fp8_ue8m0_scatter,
        )

        num_groups = hidden_states.size(1) // block_k
        gateup_input_scale = torch.empty(
            (gateup_input.size(0), num_groups // 4, m_max),
            device=hidden_states.device,
            dtype=torch.int32,
        )
        per_token_quant_fp8_ue8m0_scatter(
            hidden_states,
            gateup_input,
            gateup_input_scale,
            src2dst,
            topk_ids,
            top_k,
            m_max,
            group_size=block_k,
        )
        gateup_input_scale = gateup_input_scale.transpose(1, 2)
    elif is_fp8:
        hidden_states, scale = per_token_group_quant_fp8(hidden_states, block_k)
        gateup_input_scale = torch.empty(
            (gateup_input.size(0), gateup_input.size(1), scale.size(1)),
            device=hidden_states.device,
            dtype=scale.dtype,
        )
        fill_gateup_input_triton_kernel[(hidden_states.shape[0],)](
            hidden_states,
            scale,
            gateup_input,
            gateup_input_scale,
            src2dst,
            topk_ids,
            top_k,
            hidden_states.size(1),
            scale.size(1),
            m_max,
            scale.stride(0),
            scale.stride(1),
            BLOCK_SIZE=1024,
            IS_FP8=True,
            SCALE_MN_MAJOR=False,
        )
    else:
        scale = None
        gateup_input_scale = None
        fill_gateup_input_triton_kernel[(hidden_states.shape[0],)](
            hidden_states,
            scale,
            gateup_input,
            gateup_input_scale,
            src2dst,
            topk_ids,
            top_k,
            hidden_states.size(1),
            0,
            m_max,
            0,
            0,
            BLOCK_SIZE=1024,
            IS_FP8=False,
            SCALE_MN_MAJOR=False,
        )

    return (
        masked_m,
        expected_m,
        src2dst,
        gateup_input,
        gateup_input_scale,
    )


@triton.jit
def compute_identity_kernel(
    top_k,
    hidden_states_ptr,
    expert_scales_ptr,
    num_tokens,
    output_ptr,
    hidden_dim,
    scales_stride,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)

    batch_id = pid // (hidden_dim // BLOCK_SIZE)
    dim_offset = pid % (hidden_dim // BLOCK_SIZE) * BLOCK_SIZE

    if batch_id >= num_tokens or dim_offset >= hidden_dim:
        return

    h = tl.load(
        hidden_states_ptr
        + batch_id * hidden_dim
        + dim_offset
        + tl.arange(0, BLOCK_SIZE),
        mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim,
    )

    result = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
    for i in range(top_k):
        scale = tl.load(expert_scales_ptr + batch_id * scales_stride + i)
        result += h * scale

    tl.store(
        output_ptr + batch_id * hidden_dim + dim_offset + tl.arange(0, BLOCK_SIZE),
        result,
        mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim,
    )


def zero_experts_compute_triton(
    expert_indices, expert_scales, num_experts, zero_expert_type, hidden_states
):
    N = expert_indices.numel()
    top_k = expert_indices.size(-1)
    grid = lambda meta: (triton.cdiv(N, meta["BLOCK_SIZE"]),)

    if zero_expert_type == "identity":
        zero_expert_mask = expert_indices < num_experts
        zero_expert_scales = expert_scales.clone()
        zero_expert_scales[zero_expert_mask] = 0.0

    normal_expert_mask = expert_indices >= num_experts
    # Keep a valid routed-expert id for MoE kernels that do not accept negative
    # ids. The zero scale below still removes the routed-expert contribution.
    expert_indices[normal_expert_mask] = 0
    expert_scales[normal_expert_mask] = 0.0

    output = torch.zeros_like(hidden_states).to(hidden_states.device)
    hidden_dim = hidden_states.size(-1)
    num_tokens = hidden_states.size(0)

    grid = lambda meta: (num_tokens * (hidden_dim // meta["BLOCK_SIZE"]),)
    compute_identity_kernel[grid](
        top_k,
        hidden_states,
        zero_expert_scales,
        num_tokens,
        output,
        hidden_dim,
        zero_expert_scales.stride(0),
        BLOCK_SIZE=256,
    )

    return output


@triton.jit
def compute_problem_sizes_w4a8_kernel(
    masked_m_ptr,
    problem_sizes1_ptr,
    problem_sizes2_ptr,
    n,
    k,
    num_experts,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = pid < num_experts
    final_occurrences = tl.load(masked_m_ptr + pid, mask=mask, other=0)

    ps1_idx_0 = pid * 3
    ps1_idx_1 = ps1_idx_0 + 1
    ps1_idx_2 = ps1_idx_0 + 2

    ps2_idx_0 = pid * 3
    ps2_idx_1 = ps2_idx_0 + 1
    ps2_idx_2 = ps2_idx_0 + 2

    ps1_mask_0 = ps1_idx_0 < num_experts * 3
    ps1_mask_1 = ps1_idx_1 < num_experts * 3
    ps1_mask_2 = ps1_idx_2 < num_experts * 3
    ps2_mask_0 = ps2_idx_0 < num_experts * 3
    ps2_mask_1 = ps2_idx_1 < num_experts * 3
    ps2_mask_2 = ps2_idx_2 < num_experts * 3

    tl.store(problem_sizes1_ptr + ps1_idx_0, 2 * n, mask=ps1_mask_0)
    tl.store(problem_sizes1_ptr + ps1_idx_1, final_occurrences, mask=ps1_mask_1)
    tl.store(problem_sizes1_ptr + ps1_idx_2, k, mask=ps1_mask_2)

    tl.store(problem_sizes2_ptr + ps2_idx_0, k, mask=ps2_mask_0)
    tl.store(problem_sizes2_ptr + ps2_idx_1, final_occurrences, mask=ps2_mask_1)
    tl.store(problem_sizes2_ptr + ps2_idx_2, n, mask=ps2_mask_2)


def compute_problem_sizes_w4a8(
    masked_m, problem_sizes1, problem_sizes2, n, k, num_experts
):
    BLOCK_SIZE = 256
    grid = lambda meta: (triton.cdiv(num_experts, meta["BLOCK_SIZE"]),)
    compute_problem_sizes_w4a8_kernel[grid](
        masked_m,
        problem_sizes1,
        problem_sizes2,
        n,
        k,
        num_experts,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    return problem_sizes1, problem_sizes2


def deepep_ll_get_cutlass_w4a8_moe_mm_data(
    masked_m,
    problem_sizes1,
    problem_sizes2,
    num_experts,
    n,
    k,
):
    problem_sizes1, problem_sizes2 = compute_problem_sizes_w4a8(
        masked_m, problem_sizes1, problem_sizes2, n, k, num_experts
    )
    return (
        problem_sizes1.to(torch.int32),
        problem_sizes2.to(torch.int32),
    )


@triton.jit
def _silu_and_mul_post_per_tensor_quant_kernel(
    input_ptr,
    stride_input_expert,
    stride_input_token,
    stride_input_dim,
    output_ptr,
    stride_output_expert,
    stride_output_token,
    stride_output_dim,
    scale_ptr,
    masked_m_ptr,
    inner_dim,
    fp8_max,
    fp8_min,
    BLOCK_N: tl.constexpr,
    NUM_STAGE: tl.constexpr,
):
    """
    Triton kernel: fused SiLU(gate) * up + per-tensor FP8 quantization.

    Shape:
        input:  [E, T_padded, 2*D]  -> gate: [:,:,D], up: [:,:,D]
        output: [E, T_padded, D], dtype=float8_e4m3fn
    """
    expert_id = tl.program_id(2)
    block_id_token = tl.program_id(1)
    block_id_dim = tl.program_id(0)

    num_token_blocks = tl.num_programs(1)

    token_num_cur_expert = tl.load(masked_m_ptr + expert_id)

    scale = 1.0 / tl.load(scale_ptr).to(tl.float32)

    stride_input_expert = tl.cast(stride_input_expert, tl.int32)
    stride_output_expert = tl.cast(stride_output_expert, tl.int32)
    stride_input_token = tl.cast(stride_input_token, tl.int32)
    stride_output_token = tl.cast(stride_output_token, tl.int32)

    offset_d = block_id_dim * BLOCK_N + tl.arange(0, BLOCK_N)
    mask_d = offset_d < inner_dim

    # base pointers for current expert and dim block
    input_base_offs = input_ptr + expert_id * stride_input_expert + offset_d
    output_base_offs = output_ptr + expert_id * stride_output_expert + offset_d

    for token_idx in tl.range(
        block_id_token, token_num_cur_expert, num_token_blocks, num_stages=NUM_STAGE
    ):
        gate_ptr = input_base_offs + token_idx * stride_input_token
        up_ptr = gate_ptr + inner_dim
        gate = tl.load(gate_ptr, mask=mask_d, other=0.0).to(tl.float32)
        up = tl.load(up_ptr, mask=mask_d, other=0.0).to(tl.float32)

        # SiLU: x * sigmoid(x)
        gate = gate / (1 + tl.exp(-gate))
        gate = gate.to(input_ptr.dtype.element_ty)
        gate_up = up * gate

        scaled = gate_up * scale
        output_q = tl.clamp(scaled, fp8_min, fp8_max).to(output_ptr.dtype.element_ty)
        out_ptr = output_base_offs + token_idx * stride_output_token
        tl.store(out_ptr, output_q, mask=mask_d)


def silu_and_mul_masked_post_per_tensor_quant_fwd(
    input: torch.Tensor,
    output: torch.Tensor,
    masked_m: torch.Tensor,
    scale: torch.Tensor,
) -> torch.Tensor:
    """
    Fused SiLU + Mul + Per-Tensor Quantization to FP8.

    Args:
        input: [expert_num, token_num_padded, 2 * inner_dim]
        output: [expert_num, token_num_padded, inner_dim], dtype=torch.float8_e4m3fn
        masked_m: [expert_num], actual token count for each expert
        scale: [1] or [expert_num], quantization scale (per-tensor or per-expert)

    Returns:
        output tensor
    """
    assert input.is_contiguous()
    assert output.is_contiguous()
    assert output.dtype == torch.float8_e4m3fn
    assert input.ndim == 3
    assert input.shape[0] == masked_m.shape[0]
    assert input.shape[-1] % 2 == 0
    assert scale.numel() == 1 or scale.shape[0] == input.shape[0]

    expert_num = input.shape[0]
    #  3584
    inner_dim = input.shape[-1] // 2

    BLOCK_N = 256
    BLOCK_M = 64 if expert_num < 4 else 32
    NUM_STAGES = 3
    hidden_dim_split_block_num = triton.cdiv(inner_dim, BLOCK_N)

    grid = (hidden_dim_split_block_num, BLOCK_M, expert_num)
    finfo = torch.finfo(torch.float8_e4m3fn)
    fp8_max = finfo.max
    fp8_min = -fp8_max

    _silu_and_mul_post_per_tensor_quant_kernel[grid](
        input,
        *input.stride(),
        output,
        *output.stride(),
        scale,
        masked_m,
        inner_dim,
        fp8_max,
        fp8_min,
        BLOCK_N=BLOCK_N,
        NUM_STAGE=NUM_STAGES,
    )
    return output


@triton.jit
def _fp8_per_token_quant_to_per_tensor_quant_kernel(
    x_ptr,
    x_scale_ptr,
    x_scale_stride0,
    x_scale_stride1,
    x_scale_stride2,
    masked_m_ptr,
    output_scale_ptr,
    output_ptr,
    m,
    k,
    K_SCALE_BLOCK_SIZE: tl.constexpr,
    K_BLOCK_SIZE: tl.constexpr,
):
    pid_k, pid_m, pid_e = (
        tl.program_id(axis=0),
        tl.program_id(axis=1),
        tl.program_id(axis=2),
    )
    pid_m_dim = tl.num_programs(1)

    token_id = pid_m
    last_effective_id = tl.load(masked_m_ptr + pid_e)

    if token_id >= last_effective_id:
        return
    output_scale_val_inv = 1.0 / tl.load(output_scale_ptr).to(tl.float32)
    k_offsets = pid_k * K_BLOCK_SIZE + tl.arange(0, K_BLOCK_SIZE)
    scale_offsets = (k_offsets // K_SCALE_BLOCK_SIZE) * x_scale_stride2

    x_ptrs = x_ptr + pid_e * m * k + k_offsets
    output_ptrs = output_ptr + pid_e * m * k + k_offsets
    x_scale_ptrs = x_scale_ptr + pid_e * x_scale_stride0 + scale_offsets

    for tok_idx in tl.range(token_id, last_effective_id, pid_m_dim):
        hidden = tl.load(x_ptrs + tok_idx * k).to(tl.float32)
        scale_fp32 = tl.load(x_scale_ptrs + tok_idx * x_scale_stride1).to(tl.float32)
        hidden = hidden * scale_fp32 * output_scale_val_inv
        tl.store(output_ptrs + tok_idx * k, hidden.to(output_ptr.dtype.element_ty))


def fp8_per_token_to_per_tensor_quant_triton(
    x: torch.Tensor,
    x_scale: torch.Tensor,
    masked_m: torch.Tensor,
    output_scale: torch.Tensor,
    output: torch.Tensor,
):
    K_SCALE_BLOCK_SIZE = 128
    assert len(x.shape) == 3 and x.size(2) % K_SCALE_BLOCK_SIZE == 0
    assert x.is_contiguous()
    assert x_scale.size(2) == x.size(2) // K_SCALE_BLOCK_SIZE
    assert output_scale.numel() == 1

    K_BLOCK_SIZE = 1024
    assert x.size(2) % K_BLOCK_SIZE == 0
    grid = (x.size(2) // K_BLOCK_SIZE, 32, x.size(0))
    _fp8_per_token_quant_to_per_tensor_quant_kernel[grid](
        x,
        x_scale,
        *x_scale.stride(),
        masked_m,
        output_scale,
        output,
        x.size(1),
        x.size(2),
        K_SCALE_BLOCK_SIZE=K_SCALE_BLOCK_SIZE,
        K_BLOCK_SIZE=K_BLOCK_SIZE,
        num_warps=8,
    )