# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from math import prod
from typing import Optional

import torch
import triton
import triton.language as tl
from triton.language.extra import libdevice

from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
    per_token_group_quant_fp8)
from vllm.utils import round_up
try:
    from lmslim.layers.gemm.int8_utils import (
        per_token_group_quant_int8, per_token_quant_int8)
except Exception:
    print("INFO: Please install lmslim if you want to use int utils.\n") 
from vllm.utils import cdiv


def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor:
    """
    Shrink the given tensor and apply the given view to it.  This is
    used to resize the intermediate fused_moe caches.
    """
    assert prod(v) <= x.numel(
    ), f"{v} ({prod(v)}) <= {x.shape} ({x.numel()})"  # CUDAGRAPH unfriendly?
    return x.flatten()[:prod(v)].view(*v)


def _fp8_quantize(
    A: torch.Tensor,
    A_scale: Optional[torch.Tensor],
    per_act_token: bool,
    block_shape: Optional[list[int]] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Perform fp8 quantization on the inputs.  If a block_shape
    is provided, the output will be blocked.
    """
    if block_shape is None:
        A, A_scale = ops.scaled_fp8_quant(
            A, A_scale, use_per_token_if_dynamic=per_act_token)
    else:
        assert not per_act_token
        assert len(block_shape) == 2
        _, block_k = block_shape[0], block_shape[1]
        A, A_scale = per_token_group_quant_fp8(A, block_k)
        assert cdiv(A.size(-1), block_k) == A_scale.size(-1)

    return A, A_scale

# @triton.jit
# def _per_token_quant_int8_kernel_opt(
#     x_ptr,
#     xq_ptr,
#     scale_ptr,
#     stride_x,
#     stride_xq,
#     N, 
#     T_dim, 
#     has_tokens_per_expert: tl.constexpr,
#     tokens_per_expert_ptr, 
#     BLOCK: tl.constexpr
# ):
#     row_id = tl.program_id(0)

#     if has_tokens_per_expert:
#         e = row_id // T_dim
#         t = row_id % T_dim
        
#         num_valid_tokens_for_e = tl.load(tokens_per_expert_ptr + e)
        
#         if t >= num_valid_tokens_for_e:
#             return

#     cols = tl.arange(0, BLOCK)
#     mask = cols < N

#     x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask,
#                 other=0.0).to(tl.float32)
#     absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10)
#     scale_x = absmax / 127
#     x_q = x * (127 / absmax)
#     x_q = libdevice.nearbyint(x_q).to(tl.int8)

#     tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask)
#     tl.store(scale_ptr + row_id, scale_x)

# def per_token_quant_int8_triton_opt(x: torch.Tensor,
#                                     tokens_per_expert: Optional[torch.Tensor] = None):
#     """
#     Python wrapper for the Triton kernel.
#     """
#     if x.dim() != 3:
#         raise ValueError(f"Input must be 3D [E, T, H], but got {x.shape}")
#     E, T, H = x.shape
#     M = E * T 
#     N = H
    
#     x_q = torch.empty_like(x, device=x.device, dtype=torch.int8)
#     scales = torch.empty(x.shape[:-1] + (1, ),
#                          device=x.device,
#                          dtype=torch.float32)
#     BLOCK = triton.next_power_of_2(N)
#     num_warps = min(max(BLOCK // 256, 1), 8)

#     grid_opt = M

#     _per_token_quant_int8_kernel_opt[(grid_opt, )](
#         x,
#         x_q,
#         scales,
#         stride_x=x.stride(-2),
#         stride_xq=x_q.stride(-2),
#         N=N,
#         T_dim=T,
#         has_tokens_per_expert= tokens_per_expert is not None,
#         tokens_per_expert_ptr=tokens_per_expert,
#         BLOCK=BLOCK,
#         num_warps=num_warps,
#         num_stages=1,
#     )

#     return x_q, scales

@triton.jit
def _per_token_quant_int8_one_kernel_opt(
    x_ptr,
    xq_ptr,
    scale_ptr,
    stride_x,
    stride_xq,
    N, 
    T_dim, 
    has_tokens_per_expert: tl.constexpr,
    tokens_per_expert_ptr, 
    BLOCK: tl.constexpr
):
    row_id = tl.program_id(0)

    if has_tokens_per_expert:
        e = row_id // T_dim
        t = row_id % T_dim
        
        num_valid_tokens_for_e = tl.load(tokens_per_expert_ptr + e)
        
        if t >= num_valid_tokens_for_e:
            return

    cols = tl.arange(0, BLOCK)
    mask = cols < N

    x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask,
                other=0.0).to(tl.float32)
    absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10)
    scale_x = absmax / 127
    x_q = x * (127 / absmax)
    x_q = libdevice.nearbyint(x_q).to(tl.int8)

    tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask)
    tl.store(scale_ptr + row_id, scale_x)

@triton.jit
def _per_token_quant_int8_kernel_opt(
    x_ptr,
    xq_ptr,
    scale_ptr,
    stride_x,
    stride_xq,
    N, 
    E_dim,
    T_dim,
    has_tokens_per_expert: tl.constexpr,
    tokens_per_expert_ptr, 
    BLOCK: tl.constexpr
):
    token_idx_start = tl.program_id(0)
    grid_size = tl.num_programs(0) 
    num_total_tokens = E_dim * T_dim

    for token_idx in range(token_idx_start, num_total_tokens, grid_size):
        
        is_valid_token = True
        if has_tokens_per_expert:
            e = token_idx // T_dim
            t = token_idx % T_dim
            
            num_valid_tokens_for_e = tl.load(tokens_per_expert_ptr + e)
            
            if t >= num_valid_tokens_for_e:
                is_valid_token = False 

        if is_valid_token:
            cols = tl.arange(0, BLOCK)
            mask = cols < N

            x = tl.load(x_ptr + token_idx * stride_x + cols, mask=mask,
                        other=0.0).to(tl.float32)
            absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10)
            scale_x = absmax / 127
            x_q = x * (127 / absmax)
            x_q = libdevice.nearbyint(x_q).to(tl.int8)

            tl.store(xq_ptr + token_idx * stride_xq + cols, x_q, mask=mask)
            tl.store(scale_ptr + token_idx, scale_x)


def per_token_quant_int8_triton_opt(x: torch.Tensor,
                                    tokens_per_expert: Optional[torch.Tensor] = None):
    if x.dim() != 3:
        raise ValueError(f"Input must be 3D [E, T, H], but got {x.shape}")
    E, T, H = x.shape
    N = H
    
    x_q = torch.empty_like(x, device=x.device, dtype=torch.int8)
    scales = torch.empty(x.shape[:-1] + (1, ),
                         device=x.device,
                         dtype=torch.float32)
    BLOCK = triton.next_power_of_2(N)
    num_warps = min(max(BLOCK // 256, 1), 8)
    if T >= 4096:
        num_warps = 1
    
    num_tokens = E * T
    grid_opt = num_tokens
    if E == 16 and T >= 1024 :
        grid_opt = max(1, num_tokens // (T // 256))
        _per_token_quant_int8_kernel_opt[(grid_opt, )](
            x,
            x_q,
            scales,
            stride_x=x.stride(-2),
            stride_xq=x_q.stride(-2),
            N=N,
            E_dim=E,
            T_dim=T,
            has_tokens_per_expert=tokens_per_expert is not None,
            tokens_per_expert_ptr=tokens_per_expert,
            BLOCK=BLOCK,
            num_warps=num_warps,
            num_stages=1,
        )
    else:
        _per_token_quant_int8_one_kernel_opt[(grid_opt, )](
            x,
            x_q,
            scales,
            stride_x=x.stride(-2),
            stride_xq=x_q.stride(-2),
            N=N,
            T_dim=T,
            has_tokens_per_expert=tokens_per_expert is not None,
            tokens_per_expert_ptr=tokens_per_expert,
            BLOCK=BLOCK,
            num_warps=num_warps,
            num_stages=1,
        )
    return x_q, scales


def _int8_quantize(
    A: torch.Tensor,
    A_scale: Optional[torch.Tensor],
    per_act_token: bool,
    block_shape: Optional[list[int]] = None,
    expert_num_tokens: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Perform int8 quantization on the inputs.  If a block_shape
    is provided, the output will be blocked.
    """

    # If weights are per-channel (per_channel_quant=True), then
    # activations apply per-token quantization. Otherwise, assume
    # activation tensor-wise fp8/int8 quantization, dynamic or static
    if block_shape is None:
        # assert per_act_token, \
        #     "int8 quantization only supports block or channel-wise"
        if expert_num_tokens is None:   
            A, A_scale = per_token_quant_int8(A)
        else:
            A, A_scale = per_token_quant_int8_triton_opt(A, expert_num_tokens)
    else:
        assert not per_act_token
        assert len(block_shape) == 2
        _, block_k = block_shape[0], block_shape[1]
        A, A_scale = per_token_group_quant_int8(A, block_k)
        assert cdiv(A.size(-1), block_k) == A_scale.size(-1)

    return A, A_scale


def moe_kernel_quantize_input(
    A: torch.Tensor,
    A_scale: Optional[torch.Tensor],
    quant_dtype: Optional[torch.dtype],
    per_act_token_quant: bool,
    block_shape: Optional[list[int]] = None,
    expert_num_tokens: Optional[torch.Tensor]= None,
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
    if quant_dtype == torch.float8_e4m3fn:
        return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
    elif quant_dtype == torch.int8:
        return _int8_quantize(A, A_scale, per_act_token_quant, block_shape, expert_num_tokens)
    else:
        return A, A_scale


def _fp8_perm(m: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
    """
    A permutation routine that works on fp8 types.
    """
    if torch.is_floating_point(m) and m.dtype.itemsize == 1:
        return m.view(dtype=torch.uint8)[idx, ...].view(dtype=m.dtype)
    else:
        return m[idx, ...]


def normalize_scales_shape(
        scales: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
    if scales is not None:
        if scales.numel() == 1:
            scales = scales.view(1, 1)
        else:
            scales = scales.view(-1, scales.size(-1))
    return scales


def normalize_batched_scales_shape(
    scales: Optional[torch.Tensor],
    num_experts: int,
) -> Optional[torch.Tensor]:
    if scales is not None and scales.ndim < 3:
        if scales.numel() == 1:
            scales = scales.view(1)
            scales = torch.repeat_interleave(scales, num_experts,
                                             dim=0).view(num_experts, 1, 1)
        else:
            scales = scales.view(num_experts, -1, scales.size(-1))

    return scales


def _validate_scale_shape(
    a: torch.Tensor,
    a_scale: Optional[torch.Tensor],
    per_act_token_quant: bool,
    block_shape: Optional[list[int]],
) -> None:
    if a_scale is None:
        return

    if not per_act_token_quant and block_shape is None:
        assert a_scale.numel() == 1, f"{a_scale.shape}"
    elif per_act_token_quant:
        assert a_scale.shape[0] == a.shape[0] and a_scale.shape[1] == 1, (
            f"{a_scale.shape[0]} == {a.shape[0]} and {a_scale.shape[1]} == 1")
    else:
        assert block_shape is not None
        expected = (a.shape[0], cdiv(a.shape[1], block_shape[1]))
        assert a_scale.shape == expected, f"{a_scale.shape} == {expected}"

@triton.jit
def _count_expert_num_tokens(
    topk_ids_ptr,
    expert_num_tokens_ptr,
    num_experts,
    topk_numel,
    expert_map,
    HAS_EXPERT_MAP: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    curr_expert = tl.program_id(0)

    offsets = tl.arange(0, BLOCK_SIZE)
    topk_ids_ptrs = topk_ids_ptr + offsets

    acc = tl.zeros((BLOCK_SIZE,), dtype=tl.int32)
    for x in range(tl.cdiv(topk_numel, BLOCK_SIZE)):
        mask = offsets < (topk_numel - x * BLOCK_SIZE)
        expert_ids = tl.load(topk_ids_ptrs, mask=mask, other=-1)
        if HAS_EXPERT_MAP:
            expert_map_ptrs = expert_map + expert_ids
            expert_map_mask = expert_ids >= 0
            expert_ids = tl.load(expert_map_ptrs, mask=expert_map_mask, other=-1)

        has_curr_expert = tl.where(expert_ids == curr_expert, 1, 0)
        acc = acc + has_curr_expert
        topk_ids_ptrs += BLOCK_SIZE

    if curr_expert < num_experts:
        tl.store(expert_num_tokens_ptr + curr_expert, tl.sum(acc))


def count_expert_num_tokens(
    topk_ids: torch.Tensor, num_local_experts: int, expert_map: torch.Tensor | None
) -> torch.Tensor:
    """
    Count the number to tokens assigned to each expert.

    Parameters:
    - topk_ids (torch.Tensor): Tensor mapping each token to its
    list of experts.
    - num_local_experts (int): Number of experts in this rank.
    - expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
    from the global expert space to the local expert space of the expert
    parallel shard.

    Returns:
    A tensor of size num_local_experts, where tensor[i] holds the number
    of tokens assigned to the ith expert.
    """
    assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids"
    expert_num_tokens = torch.empty(
        (num_local_experts), device=topk_ids.device, dtype=torch.int32
    )

    grid = num_local_experts
    BLOCK_SIZE = min(topk_ids.numel(), 1024)
    BLOCK_SIZE = triton.next_power_of_2(BLOCK_SIZE)

    _count_expert_num_tokens[(grid,)](
        topk_ids,
        expert_num_tokens,
        num_local_experts,
        topk_ids.numel(),
        expert_map,
        HAS_EXPERT_MAP=expert_map is not None,
        BLOCK_SIZE=BLOCK_SIZE,
    )

    return expert_num_tokens

def expert_num_tokens_round_up_and_sum(
    expert_num_tokens: torch.Tensor, alignment: int
) -> int:
    # Round up each element in expert_num_tokens to the nearest multiple of
    # alignment.
    ent = (expert_num_tokens.to(torch.int64) + (alignment - 1)) // alignment * alignment
    return torch.sum(ent).item()

def compute_aligned_M(
    M: int,
    num_topk: int,
    local_num_experts: int,
    alignment: int,
    expert_num_tokens_cpu: Optional[torch.Tensor] = None,
):
    if expert_num_tokens_cpu is not None:
        return expert_num_tokens_round_up_and_sum(
            expert_num_tokens_cpu, alignment=alignment
        )

    # expert_num_tokens information is not available on the cpu.
    # compute the max required size.
    M_sum = (M * num_topk) + local_num_experts * (alignment - 1)
    M_sum = round_up(M_sum, alignment)
    return M_sum

@triton.jit
def apply_expert_map(expert_id, expert_map):
    if expert_id != -1:
        expert_id = tl.load(expert_map + expert_id).to(expert_id.dtype)
    return expert_id

@triton.jit
def round_up_256(x: int) -> int:
    y = 256
    return ((x + y - 1) // y) * y

@triton.jit
def round_up_128(x: int) -> int:
    y = 128
    return ((x + y - 1) // y) * y

@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,
    )
    #tokens_per_expert = round_up_128(tokens_per_expert)
    tokens_per_expert = round_up_256(tokens_per_expert)
    cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert

    #if cur_expert == 0:
    tl.store(expert_start_loc + offset_cumsum, cumsum, mask=offset_cumsum < num_experts)

    tl.debug_barrier()

    #cur_expert_start = cumsum[cur_expert]
    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,
            mask=start_m + off_expert < cur_expert_token_num
        )


@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,
    expert_map,
    HAS_EXPERT_MAP: tl.constexpr,
    HIDDEN_SIZE: tl.constexpr,
    HIDDEN_SIZE_PAD: tl.constexpr,
    SCALE_HIDDEN_SIZE: tl.constexpr,
    SCALE_HIDDEN_SIZE_PAD: 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

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

    # for token_id in range(start_token_id, total_token_num, grid_num):
    #     to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask)
    #     to_copy_s = tl.load(
    #         recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s, mask=mask_s
    #     )

    #     for topk_index in tl.range(0, topk_num, 1, num_stages=4):
    #         expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index)

    #         if HAS_EXPERT_MAP:
    #             expert_id = apply_expert_map(expert_id, expert_map)

    #         if expert_id >= 0:
    #             dest_token_index = tl.atomic_add(expert_start_loc + expert_id, 1)
    #             tl.store(
    #                 output_index + token_id * output_index_stride0 + topk_index,
    #                 dest_token_index,
    #             )
    #             output_tensor_ptr = (
    #                 output_tensor + dest_token_index * output_tensor_stride0
    #             )
    #             output_tensor_scale_ptr = (
    #                 output_tensor_scale + dest_token_index * output_tensor_scale_stride0
    #             )
    #             tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
    #             tl.store(output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s)
    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)
        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 HAS_EXPERT_MAP:
                expert_id = apply_expert_map(expert_id, expert_map)
            if expert_id >= 0:
                dest_token_index_int32 = tl.atomic_add(expert_start_loc + expert_id, 1)
                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
                )
                output_tensor_scale_ptr = (
                    output_tensor_scale + dest_token_index * output_tensor_scale_stride0
                )
                tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
                tl.store(
                    output_tensor_scale_ptr + index_in_s * output_tensor_scale_stride1,
                    to_copy_s,
                    mask=mask_s,
                )

@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_map: torch.Tensor | None,
    expert_start_loc: torch.Tensor,
    output_tensor: torch.Tensor,
    output_tensor_scale: torch.Tensor,
    m_indices: torch.Tensor,
    output_index: torch.Tensor,
):
    #BLOCK_E = 128  # token num of per expert is aligned to 128
    #BLOCK_D = 128  # block size of quantization
    BLOCK_E = 256  # token num of per expert is aligned to 256
    num_warps = 8
    num_experts = num_recv_tokens_per_expert.shape[0]
    hidden_size = recv_x.shape[1]
    scale_hidden_size = recv_x_scale.shape[-1]
    # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts)
    grid = num_experts

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

    _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),
        recv_x_scale.stride(1),
        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),
        output_tensor_scale.stride(1),
        output_index,
        output_index.stride(0),
        output_index.stride(1),
        topk_num=recv_topk.shape[1],
        expert_map=expert_map,
        HAS_EXPERT_MAP=expert_map is not None,
        num_warps=num_warps,
        HIDDEN_SIZE=hidden_size,
        HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size),
        SCALE_HIDDEN_SIZE=scale_hidden_size,#hidden_size // BLOCK_D,
        SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size)#triton.next_power_of_2(hidden_size // BLOCK_D),
    )

    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,
    expert_map,
    HAS_EXPERT_MAP: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    # cur_block = tl.program_id(0)
    # start_cur_token = tl.program_id(1)
    # grid_num = tl.num_programs(1)

    # for cur_token in range(start_cur_token, total_token_num, grid_num):
    #     off_d = tl.arange(0, BLOCK_D)
    #     accumulator = tl.zeros([BLOCK_D], dtype=tl.float32)
    #     for topk_index in range(0, topk_num):
    #         expert_id = tl.load(
    #             recv_topk_ids + cur_token * recv_topk_ids_stride0 + topk_index
    #         )

    #         if HAS_EXPERT_MAP:
    #             expert_id = apply_expert_map(expert_id, expert_map)

    #         if expert_id >= 0:
    #             source_token_index = tl.load(
    #                 input_index + cur_token * input_index_stride0 + topk_index
    #             )
    #             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),
    #     )
    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 HAS_EXPERT_MAP:
                expert_id = apply_expert_map(expert_id, expert_map)

            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,
    expert_map: torch.Tensor | None,
    output_tensor: torch.Tensor,
):
    num_warps = 2
    num_tokens = output_tensor.shape[0]
    hidden_size = input_tensor.shape[1]
    BLOCK_D = min(hidden_size, 1024)
    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],
        expert_map=expert_map,
        HAS_EXPERT_MAP=expert_map is not None,
        num_warps=num_warps,
        BLOCK_D=BLOCK_D,
    )
    return

def deepgemm_moe_permute(
    aq: torch.Tensor,
    aq_scale: torch.Tensor,
    topk_ids: torch.Tensor,
    local_num_experts: int,
    expert_map: torch.Tensor | None,
    block_shape: list[int],
    expert_num_tokens: Optional[torch.Tensor] = None,
    expert_num_tokens_cpu: Optional[torch.Tensor] = None,
    aq_out: torch.Tensor | None = None,
):
    assert aq.ndim == 2
    assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids"
    H = aq.size(1)
    device = aq.device

    block_m = block_shape[0]

    M_sum = compute_aligned_M(
        M=topk_ids.size(0),
        num_topk=topk_ids.size(1),
        local_num_experts=local_num_experts,
        alignment=block_m,
        expert_num_tokens_cpu=expert_num_tokens_cpu,
    )

    expert_start_loc = torch.empty(
        (local_num_experts), device=device, dtype=torch.int32
    )

    assert aq_out is None or aq_out.shape == (M_sum, H)
    if aq_out is None:
        aq_out = torch.empty((M_sum, H), device=device, dtype=aq.dtype)

    aq_scale_out = torch.empty(
        (M_sum, aq_scale.shape[-1]), device=device, dtype=torch.float32
        #(M_sum, H // block_k), device=device, dtype=torch.float32
    )

    # maybe_has_empty_blocks = expert_num_tokens_cpu is None
    # expert_ids_init = torch.zeros# if maybe_has_empty_blocks else torch.empty

    # expert_ids = expert_ids_init((M_sum), device=device, dtype=torch.int32)

    expert_ids = torch.full(
        (M_sum,), -1, dtype=torch.int32, device=device
    )
    inv_perm = torch.empty(topk_ids.shape, device=device, dtype=torch.int32)

    if expert_num_tokens is None:
        expert_num_tokens = count_expert_num_tokens(
            topk_ids, local_num_experts, expert_map
        )

    ep_scatter(
        recv_x=aq,
        recv_x_scale=aq_scale,
        recv_topk=topk_ids,
        num_recv_tokens_per_expert=expert_num_tokens,
        expert_start_loc=expert_start_loc,
        expert_map=expert_map,
        output_tensor=aq_out,
        output_tensor_scale=aq_scale_out,
        m_indices=expert_ids,
        output_index=inv_perm,
    )

    return aq_out, aq_scale_out, expert_ids, inv_perm

def deepgemm_unpermute_and_reduce(
    a: torch.Tensor,  # Grouped gemm output
    topk_ids: torch.Tensor,
    topk_weights: torch.Tensor,
    inv_perm: torch.Tensor,
    expert_map: torch.Tensor | None,
    output: torch.Tensor,
):  
    return ep_gather(
        input_tensor=a,
        recv_topk_ids=topk_ids,
        recv_topk_weight=topk_weights,
        input_index=inv_perm,
        expert_map=expert_map,
        output_tensor=output,
    )
