import time
from typing import Dict, Tuple, List, Optional

import torch
import torch.nn.functional as F
import triton
import triton.language as tl

from eval import _clone_data
from task import input_t, output_t
from utils import verbose_allclose
from reference import generate_input, ref_kernel


BLOCK_SIZE_M: tl.constexpr = 256
BLOCK_SIZE_N: tl.constexpr = 256
BLOCK_SIZE_K: tl.constexpr = 64
GROUP_SIZE_M: tl.constexpr = 8


# matmul

@triton.jit
def matmul_kernel(
        # Pointers to matrices
        a_ptr, b_ptr, c_ptr,
        # Matrix dimensions
        M, N, K,
        # The stride variables represent how much to increase the ptr by when moving by 1
        # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr`
        # by to get the element one row down (A has M rows).
        stride_am, stride_ak,  #
        stride_bk, stride_bn,  #
        stride_cm, stride_cn,
        # Meta-parameters
        BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,  #
        GROUP_SIZE_M: tl.constexpr,  #
        ACTIVATION: tl.constexpr  #
):
    """Kernel for computing the matmul C = A x B.
    A has shape (M, K), B has shape (K, N) and C has shape (M, N)
    """
    # -----------------------------------------------------------
    # Map program ids `pid` to the block of C it should compute.
    # This is done in a grouped ordering to promote L2 data reuse.
    # See above `L2 Cache Optimizations` section for details.
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, 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 // 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 % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    # ----------------------------------------------------------
    # Create pointers for the first blocks of A and B.
    # We will advance this pointer as we move in the K direction
    # and accumulate
    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
    # See above `Pointer Arithmetic` section for details
    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

    # -----------------------------------------------------------
    # Iterate to compute a block of the C matrix.
    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
    # of fp32 values for higher accuracy.
    # `accumulator` will be converted back to fp16 after the loop.
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        # Load the next block of A and B, generate a mask by checking the K dimension.
        # If it is out of bounds, set it to 0.
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
        # We accumulate along the K dimension.
        accumulator = tl.dot(a, b, accumulator)
        # Advance the ptrs to the next K block.
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk
    # You can fuse arbitrary activation functions here
    # while the accumulator is still in FP32!
    if ACTIVATION == "leaky_relu":
        accumulator = leaky_relu(accumulator)
    c = accumulator.to(tl.float16)

    # -----------------------------------------------------------
    # Write back the block of the output matrix C with masks.
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

def matmul(a, b, activation=""):
    # Check constraints.
    K = a.size(-1)
    M = a.numel() // K
    N = b.size(1)
    assert a.is_contiguous(), "Matrix A must be contiguous"
    assert K == b.shape[0], "Incompatible dimensions"
    # Allocates output.
    c = torch.empty((M, N), device=a.device, dtype=torch.float16)
    # 1D launch kernel where each block gets its own program.
    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
    matmul_kernel[grid](
        a, b, c,  #
        M, N, K,  #
        K, 1,  #
        b.stride(0), b.stride(1),  #
        N, 1,  #
        BLOCK_SIZE_M=64,
        BLOCK_SIZE_N=64,
        BLOCK_SIZE_K=128,
        GROUP_SIZE_M=4,
        ACTIVATION=activation  #
    )
    return c


@torch.no_grad()
def moe_expert_kernel(x: torch.Tensor, w_gate: torch.Tensor, w_up: torch.Tensor, w_down: torch.Tensor):
    x_gate = matmul(x, w_gate)
    x_gate = F.silu(x_gate)
    x_up = matmul(x, w_up)
    out = x_gate * x_up
    out = matmul(out, w_down)
    return out

# moe gate

@torch.no_grad()
def moe_gate_kernel(x: torch.Tensor, w_g: torch.Tensor, topk: int):
    logits = F.linear(x, w_g)
    scores = logits.softmax(dim=-1)
    topk_scores, topk_indices = torch.topk(scores, k=topk, dim=-1, sorted=False)

    return topk_indices, topk_scores


# moe align

def ceil_div(a, b):
    return (a + b - 1) // b


@triton.jit
def moe_align_block_size_stage1(
    topk_ids_ptr,
    tokens_cnts_ptr,
    num_experts: tl.constexpr,
    numel: tl.constexpr,
    tokens_per_thread: tl.constexpr,
):
    pid = tl.program_id(0)

    start_idx = pid * tokens_per_thread

    off_c = (pid + 1) * num_experts

    for i in range(tokens_per_thread):
        if start_idx + i < numel:
            idx = tl.load(topk_ids_ptr + start_idx + i)
            token_cnt = tl.load(tokens_cnts_ptr + off_c + idx)
            tl.store(tokens_cnts_ptr + off_c + idx, token_cnt + 1)


@triton.jit
def moe_align_block_size_stage2(
    tokens_cnts_ptr,
    num_experts: tl.constexpr,
):
    pid = tl.program_id(0)

    last_cnt = 0
    for i in range(1, num_experts + 1):
        token_cnt = tl.load(tokens_cnts_ptr + i * num_experts + pid)
        last_cnt = last_cnt + token_cnt
        tl.store(tokens_cnts_ptr + i * num_experts + pid, last_cnt)


@triton.jit
def moe_align_block_size_stage3(
    total_tokens_post_pad_ptr,
    tokens_cnts_ptr,
    cumsum_ptr,
    num_experts: tl.constexpr,
    block_size: tl.constexpr,
):
    last_cumsum = 0
    off_cnt = num_experts * num_experts
    for i in range(1, num_experts + 1):
        token_cnt = tl.load(tokens_cnts_ptr + off_cnt + i - 1)
        last_cumsum = last_cumsum + tl.cdiv(token_cnt, block_size) * block_size
        tl.store(cumsum_ptr + i, last_cumsum)
    tl.store(total_tokens_post_pad_ptr, last_cumsum)


@triton.jit
def moe_align_block_size_stage4(
    topk_ids_ptr,
    sorted_token_ids_ptr,
    expert_ids_ptr,
    tokens_cnts_ptr,
    cumsum_ptr,
    num_experts: tl.constexpr,
    block_size: tl.constexpr,
    numel: tl.constexpr,
    tokens_per_thread: tl.constexpr,
):
    pid = tl.program_id(0)
    start_idx = tl.load(cumsum_ptr + pid)
    end_idx = tl.load(cumsum_ptr + pid + 1)

    for i in range(start_idx, end_idx, block_size):
        tl.store(expert_ids_ptr + i // block_size, pid)

    start_idx = pid * tokens_per_thread
    off_t = pid * num_experts

    for i in range(start_idx, tl.minimum(start_idx + tokens_per_thread,
                                         numel)):
        expert_id = tl.load(topk_ids_ptr + i)
        token_cnt = tl.load(tokens_cnts_ptr + off_t + expert_id)
        rank_post_pad = token_cnt + tl.load(cumsum_ptr + expert_id)
        tl.store(sorted_token_ids_ptr + rank_post_pad, i)
        tl.store(tokens_cnts_ptr + off_t + expert_id, token_cnt + 1)


def invoke_moe_align_block_size(
    topk_ids: torch.Tensor,
    num_experts: int,
    block_size: int,
) -> None:
    """
    - topk_ids: A tensor of shape [total_tokens, top_k] representing the
        top-k expert indices for each token.
    - block_size: The block size used in block matrix multiplication.
    - num_experts: The total number of experts.

    Returns:
    - sorted_token_ids: A tensor containing the sorted token indices according
        to their allocated expert.
    - expert_ids: A tensor indicating the assigned expert index for each block.
    - num_tokens_post_padded: The total number of tokens after padding,
        ensuring divisibility by block_size.

    This function pads the number of tokens that each expert needs to process
    so that it is divisible by block_size.
    Padding ensures that during block matrix multiplication, the dimensions
    align correctly.

    Example:
    Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
    block_size = 4, and num_experts = 4:
    - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
        with each expert needing to process 3 tokens.
    - As block_size is 4, we pad 1 token for each expert.
    - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
    - Then append padding tokens [12, 12, 12, 12] for each block.
    - After sorting by expert index, we obtain token_ids
        [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
        Tokens 12 are non-existent (padding) and are ignored in
        the subsequent matrix multiplication.
    - The padding ensures that the total number of tokens is now divisible
        by block_size for proper block matrix operations.
    """
    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
    sorted_ids = torch.empty((max_num_tokens_padded, ),
                             dtype=torch.int32,
                             device=topk_ids.device)
    sorted_ids.fill_(topk_ids.numel())
    max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
    # Expert ids must be zeroed out to prevent index out of bounds error while
    # mapping global expert ids to local expert ids in expert parallelism.
    expert_ids = torch.zeros((max_num_m_blocks, ),
                             dtype=torch.int32,
                             device=topk_ids.device)
    num_tokens_post_pad = torch.empty((1),
                                      dtype=torch.int32,
                                      device=topk_ids.device)

    numel = topk_ids.numel()
    grid = (num_experts, )
    tokens_cnts = torch.zeros((num_experts + 1, num_experts),
                              dtype=torch.int32,
                              device=topk_ids.device)
    cumsum = torch.zeros((num_experts + 1, ),
                         dtype=torch.int32,
                         device=topk_ids.device)
    tokens_per_thread = ceil_div(numel, num_experts)

    # torch.cuda.synchronize()
    # start = time.perf_counter_ns()
    moe_align_block_size_stage1[grid](
        topk_ids,
        tokens_cnts,
        num_experts,
        numel,
        tokens_per_thread,
    )
    # torch.cuda.synchronize()
    # end = time.perf_counter_ns()
    # print(f"moe_align_block_size_stage1: {end - start}")

    # torch.cuda.synchronize()
    # start = time.perf_counter_ns()
    moe_align_block_size_stage2[grid](
        tokens_cnts,
        num_experts,
    )
    # torch.cuda.synchronize()
    # end = time.perf_counter_ns()
    # print(f"moe_align_block_size_stage2: {end - start}")

    # torch.cuda.synchronize()
    # start = time.perf_counter_ns()
    moe_align_block_size_stage3[(1, )](
        num_tokens_post_pad,
        tokens_cnts,
        cumsum,
        num_experts,
        block_size,
    )
    # torch.cuda.synchronize()
    # end = time.perf_counter_ns()
    # print(f"moe_align_block_size_stage3: {end - start}")

    # torch.cuda.synchronize()
    # start = time.perf_counter_ns()
    moe_align_block_size_stage4[grid](
        topk_ids,
        sorted_ids,
        expert_ids,
        tokens_cnts,
        cumsum,
        num_experts,
        block_size,
        numel,
        tokens_per_thread,
    )
    # torch.cuda.synchronize()
    # end = time.perf_counter_ns()
    # print(f"moe_align_block_size_stage4: {end - start}")

    return sorted_ids, expert_ids, num_tokens_post_pad


# silu and mul


# @triton.jit
# def silu_and_mul_kernel(input_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
#     pid = tl.program_id(axis=0)
#     a_block_start = pid * BLOCK_SIZE
#     b_block_start = a_block_start + n_elements
#     a_offsets = a_block_start + tl.arange(0, BLOCK_SIZE)
#     b_offsets = b_block_start + tl.arange(0, BLOCK_SIZE)
#     a_mask = a_offsets < n_elements
#     b_mask = b_offsets < n_elements * 2

#     a = tl.load(input_ptr + a_offsets, mask=a_mask)
#     b = tl.load(input_ptr + b_offsets, mask=b_mask)
#     output = a * (1 / (1 + tl.exp(-a.to(tl.float32)))) * b
#     tl.store(output_ptr + a_offsets, output, mask=a_mask)

# grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
# silu_and_mul_kernel[grid](input_tensor, output_tensor, n_elements, BLOCK_SIZE)

@triton.jit
def silu_and_mul_kernel(input_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    a_row_ptr = input_ptr + pid * n_elements * 2
    b_row_ptr = a_row_ptr + n_elements
    c_row_ptr = output_ptr + pid * n_elements
    offs = tl.arange(0, BLOCK_SIZE)
    mask = offs < n_elements

    for i in range(0, n_elements, BLOCK_SIZE):
        offs = offs + i
        a = tl.load(a_row_ptr + offs, mask=mask)
        b = tl.load(b_row_ptr + offs, mask=mask)
        c = a * (1 / (1 + tl.exp(-a.to(tl.float32)))) * b
        tl.store(c_row_ptr + offs, c, mask=mask)


def invoke_silu_and_mul_kernel(input: torch.Tensor, output: torch.Tensor):
    assert input.is_contiguous()
    assert output.is_contiguous()
    assert len(input.shape) > 1
    assert len(input.shape) == len(output.shape)
    assert input.shape[-1] == output.shape[-1] * 2
    assert input.numel() / input.shape[-1] == output.numel() / output.shape[-1]

    n_elements = input.shape[-1] // 2
    grid = (triton.cdiv(input.numel(), input.shape[-1]),)
    silu_and_mul_kernel[grid](input, output, n_elements, 1024 if n_elements > 1024 else n_elements)


def test_silu_and_mul_kernel():
    num_tokens = 8
    num_expert = 4
    hidden_size = 2048
    a = torch.rand((num_tokens, num_expert, hidden_size), dtype=torch.float16, device="cuda")
    b = torch.ones((num_tokens, num_expert, hidden_size,), dtype=torch.float16, device="cuda")
    # b = b * 2
    input_tensor = torch.concat((a, b), dim=-1)
    # print(input_tensor)

    output_tensor = torch.empty((num_tokens, num_expert, hidden_size,), dtype=torch.float16, device="cuda")
    invoke_silu_and_mul_kernel(input_tensor, output_tensor)
    # print(output_tensor)

    ref_output = torch.nn.functional.silu(a) * b
    # print(ref_output)

    print(verbose_allclose(output_tensor, ref_output, rtol=1e-2, atol=1e-2))


# fused experts

@triton.jit
def fused_moe_kernel(
    # Pointers to matrices
    a_ptr,
    b_ptr,
    c_ptr,
    topk_weights_ptr,
    sorted_token_ids_ptr,
    expert_ids_ptr,
    num_tokens_post_padded_ptr,
    # Matrix dimensions
    N,
    K,
    EM,
    num_valid_tokens,
    # The stride variables represent how much to increase the ptr by when moving by 1
    # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr`
    # by to get the element one row down (A has M rows).
    stride_am,
    stride_ak,
    stride_be,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    stride_weight,
    stride_token_id,
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    MUL_ROUTED_WEIGHT: tl.constexpr,
    top_k: tl.constexpr,
    compute_type: tl.constexpr,
):
    """
    Implements the fused computation for a Mixture of Experts (MOE) using token and expert matrices.

    Key Parameters:
    - A: The input tensor representing tokens with shape (*, K), where '*' can be any shape representing batches and K is the feature dimension of each token.
    - B: The stacked MOE weight tensor with shape (E, N, K), where E is the number of experts, K is the input feature dimension, and N is the output feature dimension.
    - C: The output cache tensor with shape (M, topk, N), where M is the total number of tokens post padding, topk is the number of times each token is repeated,
        and N is the output feature dimension.
    - sorted_token_ids: A tensor containing the sorted indices of tokens, repeated topk times and arranged by the expert index they are assigned to.
    - expert_ids: A tensor containing the indices of the expert for each block. It determines which expert matrix from B should be used for each block in A.
    This kernel performs the multiplication of a token by its corresponding expert matrix as determined by `expert_ids`. The sorting of `sorted_token_ids`
    by expert index and padding ensures divisibility by BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix multiplication across different blocks processed by the same expert.
    """
    # -----------------------------------------------------------
    # Map program ids `pid` to the block of C it should compute.
    # This is done in a grouped ordering to promote L2 data reuse.
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(EM, 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 // 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 % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    # ----------------------------------------------------------
    # Create pointers for the first blocks of A and B.
    # We will advance this pointer as we move in the K direction
    # and accumulate
    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
        return

    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
    token_mask = offs_token < num_valid_tokens

    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak)

    off_experts = tl.load(expert_ids_ptr + pid_m)
    b_ptrs = b_ptr + off_experts * stride_be + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

    # -----------------------------------------------------------
    # Iterate to compute a block of the C matrix.
    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
    # of fp32 values for higher accuracy.
    # `accumulator` will be converted back to fp16 after the loop.
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        # Load the next block of A and B, generate a mask by checking the K dimension.
        a = tl.load(a_ptrs,
                    mask=token_mask[:, None] &
                    (offs_k[None, :] < K - k * BLOCK_SIZE_K),
                    other=0.0)
        b = tl.load(b_ptrs,
                    mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
                    other=0.0)

        # We accumulate along the K dimension.
        accumulator += tl.dot(a, b)
        # Advance the ptrs to the next K block.
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

    if MUL_ROUTED_WEIGHT:
        moe_weight = tl.load(topk_weights_ptr + offs_token * stride_weight,
                             mask=token_mask,
                             other=0)
        accumulator = accumulator * moe_weight[:, None]

    accumulator = accumulator.to(compute_type)
    # -----------------------------------------------------------
    # Write back the block of the 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)
    tl.store(c_ptrs, accumulator, mask=c_mask)


def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
                            topk_weights: torch.Tensor, topk_ids: torch.Tensor,
                            sorted_token_ids: torch.Tensor,
                            expert_ids: torch.Tensor,
                            num_tokens_post_padded: torch.Tensor,
                            mul_routed_weight: bool, top_k: int):

    grid = lambda META: (triton.cdiv(sorted_token_ids.shape[0], META['BLOCK_SIZE_M']) * triton.cdiv(B.shape[1], META['BLOCK_SIZE_N']), )
    fused_moe_kernel[grid](
        A,
        B,
        C,
        topk_weights,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        B.shape[1],
        B.shape[2],
        sorted_token_ids.shape[0],
        topk_ids.numel(),
        A.stride(0),
        A.stride(1),
        B.stride(0),
        B.stride(2),
        B.stride(1),
        C.stride(1),
        C.stride(2),
        topk_weights.stride(2),
        sorted_token_ids.stride(0),
        MUL_ROUTED_WEIGHT=mul_routed_weight,
        top_k=top_k,
        compute_type=tl.bfloat16 if A.dtype == torch.bfloat16 else tl.float16,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K,
        GROUP_SIZE_M=GROUP_SIZE_M,
        num_warps=8,
        num_stages=2,
        waves_per_eu=0,
        kpack=1,
        matrix_instr_nonkdim=16,
    )


# main kernel

@torch.no_grad()
def moe_expert_kernel(x: torch.Tensor, w_gate: torch.Tensor, w_up: torch.Tensor, w_down: torch.Tensor):
    # x_gate = x @ w_gate
    x_gate = F.linear(x, w_gate.T)
    x_gate = F.silu(x_gate)
    # x_up = x @ w_up
    x_up = F.linear(x, w_up.T)
    out = x_gate * x_up
    # out = out @ w_down
    out = F.linear(out, w_down.T)
    return out


@triton.jit
def transpose_kernel(
    M: tl.constexpr,
    N: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    X,
    stride_xm: tl.constexpr,
    stride_xn: tl.constexpr,
    Y,
    stride_ym: tl.constexpr,
    stride_yn: tl.constexpr,
    IDX: tl.constexpr,
):
    pid_m = tl.program_id(0).to(tl.int64)
    pid_n = tl.program_id(1).to(tl.int64)

    m_exact: tl.constexpr = (M % BLOCK_M) == 0
    n_exact: tl.constexpr = (N % BLOCK_N) == 0

    off_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    off_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)

    mask_m = off_m < M
    mask_n = off_n < N

    if m_exact:
        if n_exact:
            mask = None
            other = None
        else:
            mask = mask_n[None, :]
            other = 0
    else:
        if n_exact:
            mask = mask_m[:, None]
            other = 0
        else:
            mask = mask_m[:, None] & mask_n[None, :]
            other = 0

    # X_ptrs = X + off_m[:, None] * 1 + off_n[None, :] * M
    # Y_ptrs = Y + (E_ID * 2 + W_ID) * M * N + off_m[:, None] * N + off_n[None, :] * 1
    X_ptrs = X + off_m[:, None] * stride_xm + off_n[None, :] * stride_xn
    Y_ptrs = Y + IDX * M * N + off_m[:, None] * stride_ym + off_n[None, :] * stride_yn

    tile = tl.load(X_ptrs, mask=mask, other=other)
    tl.store(Y_ptrs, tile, mask=mask)


def invoke_transpose(X, Y, IDX=0, BLOCK_M=128, BLOCK_N=128, num_warps=16) -> None:
    if X.dtype.itemsize == 1:
        X = X.view(torch.int8)
    if Y.dtype.itemsize == 1:
        Y = Y.view(torch.int8)

    # check compatibility:
    assert X.dtype == Y.dtype

    M = X.shape[0]
    N = X.shape[1]

    stride_xm = X.stride(0)
    stride_xn = X.stride(1)

    stride_ym = Y.stride(-2)
    stride_yn = Y.stride(-1)

    grid = ((M + BLOCK_M - 1) // BLOCK_M, (N + BLOCK_N - 1) // BLOCK_N)
    transpose_kernel[grid](
        M,
        N,
        BLOCK_M,
        BLOCK_N,
        X,
        stride_xm,
        stride_xn,
        Y,
        stride_ym,
        stride_yn,
        IDX=IDX,
        num_warps=num_warps,
    )


w1: torch.Tensor = torch.empty((32, 4096, 7168), dtype=torch.float16, device="cuda")
w2: torch.Tensor = torch.empty((32, 7168, 2048), dtype=torch.float16, device="cuda")


@torch.no_grad()
def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of DeepSeek-style Mixture of Experts using PyTorch.
    
    Args:
        data: Tuple of (input: torch.Tensor, weights: Dict[str, torch.Tensor], config: Dict)
            - input: Input tensor of shape [batch_size, seq_len, hidden_dim]
            - weights: Dictionary containing model weights
            - config: Dictionary containing model configuration parameters
            
    Returns:
        Tuple containing:
            - output: Processed tensor [batch_size, seq_len, d_model]
            - aux_data: Dictionary with auxiliary data
    """
    time_list = []

#######################################################################################################################
# prepare weight
#######################################################################################################################

    input_tensor: torch.Tensor
    weights: Dict[str, torch.Tensor]
    config: Dict[str, int]
    input_tensor, weights, config = data

    M = config["batch_size"] * config["seq_len"]
    N = config["d_expert"]
    K = config["d_hidden"]

    for i in range(config["n_routed_experts"]):
        for j in range(2):
            invoke_transpose(weights[f"experts.{i}.{j}.weight"].T, w1, i * 2 + j)
    for i in range(config["n_routed_experts"]):
        invoke_transpose(weights[f"experts.{i}.2.weight"].T, w2, i)

    intermediate_cache1 = torch.ones((M, config["n_experts_per_token"], N * 2),
                                      device=input_tensor.device,
                                      dtype=input_tensor.dtype)
    intermediate_cache2 = torch.empty((M * config["n_experts_per_token"], N),
                                      device=input_tensor.device,
                                      dtype=input_tensor.dtype)
    intermediate_cache3 = torch.empty((M, config["n_experts_per_token"], K),
                                      device=input_tensor.device,
                                      dtype=input_tensor.dtype)

#######################################################################################################################
# moe gate
#######################################################################################################################

    expert_indices, expert_scores = moe_gate_kernel(
        input_tensor,
        weights['router.weight'],
        config["n_experts_per_token"])

#######################################################################################################################
# shared_experts
#######################################################################################################################

    shared_output = moe_expert_kernel(
        input_tensor,
        weights['shared_experts.0.weight'],
        weights['shared_experts.1.weight'],
        weights['shared_experts.2.weight'])

#######################################################################################################################
# moe align
#######################################################################################################################

    sorted_token_ids, expert_ids, num_tokens_post_padded = invoke_moe_align_block_size(
        expert_indices, config["n_routed_experts"], BLOCK_SIZE_M)

#######################################################################################################################
# fused_moe_1
#######################################################################################################################

    invoke_fused_moe_kernel(input_tensor.view(-1, K), w1, intermediate_cache1,
                            expert_scores, expert_indices, sorted_token_ids,
                            expert_ids, num_tokens_post_padded, False,
                            expert_indices.shape[2])

#######################################################################################################################
# silu_and_mul
#######################################################################################################################

    invoke_silu_and_mul_kernel(intermediate_cache1.view(-1, 2 * N), intermediate_cache2)

#######################################################################################################################
# fused_moe_2
#######################################################################################################################

    invoke_fused_moe_kernel(intermediate_cache2, w2, intermediate_cache3,
                            expert_scores, expert_indices, sorted_token_ids,
                            expert_ids, num_tokens_post_padded, True, 1)

#######################################################################################################################
# sum
#######################################################################################################################

    routed_output = torch.sum(
        intermediate_cache3.view(*intermediate_cache3.shape), dim=1).view(config["batch_size"], config["seq_len"], K)

#######################################################################################################################
# add
#######################################################################################################################

    output = routed_output + shared_output

    return output

python eval.py leaderboard /tmp/tmps20x8kec