import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, Optional, List
from task import input_t, output_t
import triton
import triton.language as tl
from copy import deepcopy
import time

@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_SIZE": 512,
            },
            num_warps=4,
            num_stages=1,
        ),    
    ],
    key=[],
)
@triton.jit
def calculate_expert2token_count_and_within_expert_index(
    topk_indices,
    expert2token_count,
    within_expert_index,
    token_count,
    n_experts_per_token,
    n_routed_experts,
    BLOCK_SIZE: tl.constexpr,
):
    # topk_indices: (token_count, n_experts_per_token) : (n_experts_per_token, 1)
    # expert2token_count: (n_routed_experts,) : (1, )
    # within_expert_index: (token_count, n_experts_per_token) : (n_experts_per_token, 1)
    # grid: (n_routed_experts, )
    expert_idx = tl.program_id(0)
    
    total_count = 0
    for idx in range(0, token_count * n_experts_per_token, BLOCK_SIZE):
        offsets = tl.arange(0, BLOCK_SIZE) + idx
        mask = offsets < token_count * n_experts_per_token
        topk_indices_block = tl.load(topk_indices + offsets, mask=mask, other=-1) # (BLOCK_SIZE, )
        is_current_expert = tl.where(topk_indices_block == expert_idx, 1, 0) # (BLOCK_SIZE, )
        cur_cumsum = tl.cumsum(is_current_expert, axis=0) + total_count - 1 #shift left by 1
        total_count += tl.sum(is_current_expert)
        
        within_expert_index_block = tl.where(is_current_expert, cur_cumsum, -1)
        tl.store(
            within_expert_index + idx + tl.arange(0, BLOCK_SIZE), within_expert_index_block, 
            mask= (idx + tl.arange(0, BLOCK_SIZE) < token_count * n_experts_per_token) and (is_current_expert.to(tl.int1))
        )
        
    tl.store(expert2token_count + expert_idx, total_count)

@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_SIZE": 512,
            },
            num_warps=4,
            num_stages=1,
        ),    
    ],
    key=[],
)
@triton.jit
def calculate_token2index(
    topk_indices,
    expert2token_count,
    within_expert_index,
    token2index,
    expert_token_count_cumsum,
    token_count,
    n_experts_per_token,
    n_routed_experts,
    BLOCK_SIZE: tl.constexpr,
):
    # topk_indices: (token_count * n_experts_per_token, ) : (1, )
    # expert2token_count: (n_routed_experts,) : (1, )
    # within_expert_index: (token_count * n_experts_per_token, ) : (1, )
    # token2index: (token_count * n_experts_per_token, ) : (1, )
    # expert_token_count_cumsum: (n_routed_experts + 1, ) : (1, )
    # grid: (cdiv(token_count * n_experts_per_token, BLOCK_SIZE), )
    
    # tl.device_assert(BLOCK_SIZE >= n_experts_per_token + 1)
    expert2token_count_block = tl.load(
        expert2token_count + tl.arange(0, BLOCK_SIZE),
        mask=tl.arange(0, BLOCK_SIZE) < n_routed_experts,
        other=0,
    ) # (BLOCK_SIZE, )
    expert_token_count_cumsum_block = tl.cumsum(
        expert2token_count_block,
        axis=0,
    ) # (BLOCK_SIZE, )
    if tl.program_id(0) == 0:
        tl.store(expert_token_count_cumsum + tl.arange(0, BLOCK_SIZE) + 1, expert_token_count_cumsum_block, mask=tl.arange(0, BLOCK_SIZE) < n_routed_experts)
    
    topk_indices_block = tl.load(
        topk_indices + tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE),
        mask=tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < token_count * n_experts_per_token,
        other=0,
    ) # (BLOCK_SIZE, )
    within_expert_index_block = tl.load(
        within_expert_index + tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE),
        mask=tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < token_count * n_experts_per_token,
        other=0,
    ) # (BLOCK_SIZE, )
    is_zero = topk_indices_block == 0
    expert_offset = tl.gather(
        expert_token_count_cumsum_block ,
        tl.where(is_zero, 0, topk_indices_block - 1),
        axis=0,
    ) # (BLOCK_SIZE, )
    expert_offset = tl.where(is_zero, 0, expert_offset) # (BLOCK_SIZE, )
    token2index_block = expert_offset + within_expert_index_block # (BLOCK_SIZE, )
    tl.store(
        token2index + tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE),
        token2index_block,
        mask=tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < token_count * n_experts_per_token,
    )
@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_SIZE": 1024,
            },
            num_warps=4,
            num_stages=1,
        ),
    ],
    key=[],
)
@triton.jit
def scatter(
    input_tensor,
    token2index,
    repeated_tokens,
    batch_size,
    seq_len,
    n_experts_per_token,
    d_hidden, 
    BLOCK_SIZE: tl.constexpr,
):
    # input_tensor: (batch_size * seq_len, d_hidden) : (d_hidden, 1)
    # token2index: (batch_size * seq_len, n_experts_per_token) : (n_experts_per_token, 1)
    # repeated_tokens: (batch_size * seq_len * n_experts_per_token, d_hidden) : (d_hidden, 1)
    # grid: (n_experts_per_token, seq_len, batch_size)
    expert_idx = tl.program_id(0)
    seq_idx = tl.program_id(1)
    batch_idx = tl.program_id(2)
    
    token_idx = batch_idx * seq_len + seq_idx
    dst_idx = tl.load(token2index + token_idx * n_experts_per_token + expert_idx)
    for i in range(0, d_hidden, BLOCK_SIZE):
        offset = tl.arange(0, BLOCK_SIZE) + i
        mask = offset < d_hidden
        input_tensor_block = tl.load(
            input_tensor + token_idx * d_hidden + offset,
            mask=mask,
        )
        tl.store(
            repeated_tokens + dst_idx * d_hidden + offset,
            input_tensor_block,
            mask=mask,
        )
@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_M": 128,
                "BLOCK_N": 128,
                "BLOCK_K": 64,
            },
            num_warps=8,
            num_stages=2,
        ),
    ],
    key=[],
)
@triton.jit
def group_gemm(
    a, b_base_ptrs, c,
    m_cumsum,
    max_m,
    n,
    k,
    group_size,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    # a: (total_m, k) : (k, 1)
    # b_base_ptrs: (group_size, ) : (1, )
    # b: (k, n) : (n, 1)
    # c: (total_m, n) : (n, 1)
    # m_cumsum: (group_size + 1, ) : (1, )
    # grid: ( triton.cdiv(n, BLOCK_N), triton.cdiv(max_m, BLOCK_M), group_size )
    group_id = tl.program_id(2)
    n_start = tl.program_id(0) * BLOCK_N
    m_id = tl.program_id(1)
    
    group_m_start = tl.load(m_cumsum + group_id)
    group_m_end = tl.load(m_cumsum + group_id + 1)
    group_b_base_ptr = tl.load(b_base_ptrs + group_id)
    if group_m_start + m_id * BLOCK_M >= group_m_end:
        return
    a_block_ptr = tl.make_block_ptr(
        a + group_m_start * k,
        shape=(group_m_end - group_m_start, k),
        strides=(k, 1),
        offsets=(m_id * BLOCK_M, 0),
        block_shape=(BLOCK_M, BLOCK_K),
        order=(1, 0),
    )
    b_block_ptr = tl.make_block_ptr(
        base=group_b_base_ptr.to(tl.pointer_type(tl.float16)),
        shape=(k, n),
        strides=(n, 1),
        offsets=(0, n_start),
        block_shape=(BLOCK_K, BLOCK_N),
        order=(1, 0),
    )
    c_block_ptr = tl.make_block_ptr(
        base=c + group_m_start * n,
        shape=(group_m_end - group_m_start, n),
        strides=(n, 1),
        offsets=(m_id * BLOCK_M, n_start),
        block_shape=(BLOCK_M, BLOCK_N),
        order=(1, 0),
    )
    accum = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # (BLOCK_M, BLOCK_N)
    for i in range(0, k, BLOCK_K):
        a_block = tl.load(a_block_ptr, boundary_check=(0, 1)) # (BLOCK_M, BLOCK_K)
        b_block = tl.load(b_block_ptr, boundary_check=(0, 1)) # (BLOCK_K, BLOCK_N)
        accum += tl.dot(a_block, b_block) # (BLOCK_M, BLOCK_N)
        a_block_ptr = tl.advance(a_block_ptr, offsets=(0, BLOCK_K))
        b_block_ptr = tl.advance(b_block_ptr, offsets=(BLOCK_K, 0))
    tl.store(c_block_ptr, accum.to(tl.float16), boundary_check=(0, 1), )
        
@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_SIZE": 1024 * 8,
            },
            num_warps=4,
            num_stages=1,
        ),
    ],
    key=[],
)           
@triton.jit
def weighted_sum_gather(
    ffn1_out,
    token2index,
    topk_scores,
    shared_expert_ffn1,
    moe_out,
    d_hidden,
    batch_size,
    seq_len,
    n_experts_per_token,
    BLOCK_SIZE: tl.constexpr,
):
    # ffn1_out: (batch_size * seq_len * n_experts_per_token, d_hidden) : (d_hidden, 1)
    # token2index: (batch_size * seq_len, n_experts_per_token) : (n_experts_per_token, 1)
    # topk_scores: (batch_size * seq_len, n_experts_per_token) : (n_experts_per_token, 1)
    # shared_expert_ffn1: (batch_size * seq_len, d_hidden) : (d_hidden, 1)
    # moe_out: (batch_size * seq_len, d_hidden) : (d_hidden, 1)
    # grid: (triton.cdiv(d_hidden, BLOCK_SIZE), batch_size * seq_len)
    
    segment_idx = tl.program_id(0)
    token_idx = tl.program_id(1)
    
    offsets = tl.arange(0, BLOCK_SIZE) + segment_idx * BLOCK_SIZE
    mask = offsets < d_hidden
    
    data = tl.load(
        shared_expert_ffn1 + token_idx * d_hidden + offsets,
        mask=mask,
    )
    for i in range(n_experts_per_token):
        weight = tl.load(topk_scores + token_idx * n_experts_per_token + i)
        index = tl.load(token2index + token_idx * n_experts_per_token + i)
        ffn1_out_block = tl.load(
            ffn1_out + index * d_hidden + offsets,
            mask=mask,
        )
        data += weight * ffn1_out_block
    tl.store(
        moe_out + token_idx * d_hidden + offsets,
        data,
        mask=mask,
    )
   
# config: {'d_hidden': 7168, 'd_expert': 2048, 'n_routed_experts': 32, 'n_shared_experts': 1, 'n_experts_per_token': 4, 'batch_size': 1, 'seq_len': 2048}
# input_tensor layout: (batch_size, seq_len, d_hidden) : (d_hidden * seq_len, d_hidden, 1)
# router.weight layout: (n_routed_experts,  d_hidden) : (d_hidden, 1)
# experts.{i}.0.weight layout: (d_hidden, d_expert) : (d_expert, 1)
# experts.{i}.1.weight layout: (d_hidden, d_expert) : (d_expert, 1)
# experts.{i}.2.weight layout: (d_expert, d_hidden) : (d_hidden, 1)
# shared_experts.{i}.0.weight layout: (d_hidden, d_expert * n_shared_experts) : (d_expert * n_shared_experts, 1)
# shared_experts.{i}.1.weight layout: (d_hidden, d_expert * n_shared_experts) : (d_expert * n_shared_experts, 1)
# shared_experts.{i}.2.weight layout: (d_expert * n_shared_experts, d_hidden) : (d_hidden, 1)

class Timer:
    def __init__(self, label: str):
        self.label = label
        self.start_time = None
        self.end_time = None
        self.interval = None
    def __enter__(self):
        self.start = time.time()
        return self
    def __exit__(self, *args):
        torch.cuda.synchronize()
        self.end = time.time()
        self.interval = (self.end - self.start) * 1000 # in ms
        print(f"{self.label}: {self.interval:.4f} ms")
        return False
def moe_forward(
    input_tensor: torch.Tensor,
    weights: Dict[str, torch.Tensor],
    d_hidden: int,
    d_expert: int,
    n_routed_experts: int,
    n_shared_experts: int,
    n_experts_per_token: int,
    batch_size: int,
    seq_len: int
) -> torch.Tensor:
    input_tensor = input_tensor.reshape(-1, d_hidden) # (batch_size * seq_len, d_hidden)
    
    with Timer(label=f"({batch_size}, {seq_len}):  allocate tensor timer") as t:
    # Allocate tensor 
        expert2token_count = torch.empty(size=(n_routed_experts,), device=input_tensor.device, dtype=torch.int32)
        within_expert_index = torch.empty(size=(batch_size * seq_len, n_experts_per_token), device=input_tensor.device, dtype=torch.int32)
        token2index = torch.empty(size=(batch_size * seq_len, n_experts_per_token), device=input_tensor.device, dtype=torch.int32)
        expert_token_count_cumsum = torch.zeros(size=(n_routed_experts + 1,), device=input_tensor.device, dtype=torch.int32)
        
        repeated_tokens = torch.empty(size=(batch_size * seq_len * n_experts_per_token, d_hidden), device=input_tensor.device, dtype=input_tensor.dtype)

        ffn0_lhs = torch.empty(size=(batch_size * seq_len * n_experts_per_token, d_expert), device=input_tensor.device, dtype=input_tensor.dtype)
        ffn0_rhs = torch.empty(size=(batch_size * seq_len * n_experts_per_token, d_expert), device=input_tensor.device, dtype=input_tensor.dtype)
        ffn1_out = torch.empty_like(repeated_tokens)
        moe_out = torch.empty_like(input_tensor)

    with Timer(label=f"({batch_size}, {seq_len}):  allocate ptr array") as t:
        b_ptrs0 = torch.tensor(
            [weights[f'experts.{i}.0.weight'].data_ptr()  for i in range(n_routed_experts)],
            device=repeated_tokens.device,
            dtype=torch.uint64,
        )
        b_ptrs1 = torch.tensor(
            [weights[f'experts.{i}.1.weight'].data_ptr()  for i in range(n_routed_experts)],
            device=repeated_tokens.device,
            dtype=torch.uint64,
        )
        b_ptrs2 = torch.tensor(
            [weights[f'experts.{i}.2.weight'].data_ptr()  for i in range(n_routed_experts)],
            device=repeated_tokens.device,
            dtype=torch.uint64,
        )

    
    router_weight = weights['router.weight'] # (n_routed_experts, d_hidden)
    # TODO: fused matmul + softmax + top k ? or implement matmul? or use torch.matmul?
    logits = torch.matmul(input_tensor, router_weight.t()) # (batch_size * seq_len, n_routed_experts)
    # TODO: fused softmax + top k https://github.com/sgl-project/sglang/blob/fba8eccd7ebe41bbdbf70ab3b6a2df1835f8b532/sgl-kernel/csrc/torch_extension_rocm.cc#L65
    scores = F.softmax(logits, dim=-1) # (batch_size * seq_len, n_routed_experts)
    topk_scores, topk_indices = torch.topk(scores, n_experts_per_token, dim=-1, sorted=False) # (batch_size * seq_len, n_experts_per_token)    

    shared_expert_weight0 = weights['shared_experts.0.weight'] # (d_hidden, d_expert * n_shared_experts)
    shared_expert_weight1 = weights['shared_experts.1.weight'] # (d_hidden, d_expert * n_shared_experts)
    shared_expert_weight2 = weights['shared_experts.2.weight'] # (d_expert * n_shared_experts, d_hidden)
    # TODO: fused matmul + silu + matmul
    # TODO: implement matmul? or use torch.matmul?
    # TODO: shared expert in a separate stream?
    with Timer(label=f"({batch_size}, {seq_len}):  shared expert timer") as t:
        shared_expert_ffn0 = F.silu(input_tensor @ shared_expert_weight0) * (input_tensor @ shared_expert_weight1) # (batch_size * seq_len, d_expert * n_shared_experts)
        shared_expert_ffn1 = shared_expert_ffn0 @ shared_expert_weight2 # (batch_size * seq_len, d_hidden)

    with Timer(label=f"({batch_size}, {seq_len}):  calculate index timer") as t:
        # expert2token_count (n_routed_experts,)
        # within_expert_index (batch_size * seq_len, n_experts_per_token)
        calculate_expert2token_count_and_within_expert_index[ (n_routed_experts, ) ](
            topk_indices,
            expert2token_count,
            within_expert_index,
            batch_size * seq_len,
            n_experts_per_token,
            n_routed_experts,
        )
        # token2index (batch_size * seq_len, n_experts_per_token)
        calculate_token2index[ lambda meta : (triton.cdiv(batch_size * seq_len * n_experts_per_token, meta["BLOCK_SIZE"]),) ](
            topk_indices,
            expert2token_count,
            within_expert_index,
            token2index,
            expert_token_count_cumsum,
            batch_size * seq_len,
            n_experts_per_token,
            n_routed_experts,
        )

    with Timer(label=f"({batch_size}, {seq_len}):  scatter timer") as t:
        scatter[ (n_experts_per_token, seq_len, batch_size) ](
            input_tensor,
            token2index,
            repeated_tokens,
            batch_size,
            seq_len,
            n_experts_per_token,
            d_hidden,
        )
    
    max_m = batch_size * seq_len
    with Timer(label=f"({batch_size}, {seq_len}):  ffn0 timer") as t:
        group_gemm[ lambda meta: (triton.cdiv(d_expert, meta["BLOCK_N"]), triton.cdiv(max_m, meta["BLOCK_M"]), n_routed_experts) ](
            a = repeated_tokens,
            b_base_ptrs = b_ptrs0,
            c = ffn0_lhs,
            m_cumsum = expert_token_count_cumsum,
            max_m = batch_size * seq_len * n_experts_per_token,
            n = d_expert,
            k = d_hidden,
            group_size = n_routed_experts,
            # BLOCK_M=BLOCK_M,
            # BLOCK_N=BLOCK_N,
            # BLOCK_K=BLOCK_K,
        )
        
        group_gemm[ lambda meta: (triton.cdiv(d_expert, meta["BLOCK_N"]), triton.cdiv(max_m, meta["BLOCK_M"]), n_routed_experts) ](
            a = repeated_tokens,
            b_base_ptrs = b_ptrs1,
            c = ffn0_rhs,
            m_cumsum = expert_token_count_cumsum,
            max_m = batch_size * seq_len * n_experts_per_token,
            n = d_expert,
            k = d_hidden,
            group_size = n_routed_experts,
            # BLOCK_M=BLOCK_M,
            # BLOCK_N=BLOCK_N,
            # BLOCK_K=BLOCK_K,
        )
    
    with Timer(label=f"({batch_size}, {seq_len}):  ffn0 Silu timer") as t:
        # TODO: fused silu with group gemm ?
        # TODO: implement F.silu(ffn0_lhs) * ffn0_rhs using triton?
        # TODO: torch.matmul(ffn0_lhs, ffn0_rhs) ?
        ffn0_out = F.silu(ffn0_lhs) * ffn0_rhs # (batch_size * seq_len * n_experts_per_token, d_expert)
    
    with Timer(label=f"({batch_size}, {seq_len}):  ffn1 timer") as t:
        group_gemm[ lambda meta: (triton.cdiv(d_hidden, meta["BLOCK_N"]), triton.cdiv(max_m, meta["BLOCK_M"]), n_routed_experts) ](
            a = ffn0_out,
            b_base_ptrs = b_ptrs2,
            c = ffn1_out,
            m_cumsum = expert_token_count_cumsum,
            max_m = batch_size * seq_len * n_experts_per_token,
            n = d_hidden,
            k = d_expert,
            group_size = n_routed_experts,
            # BLOCK_M=BLOCK_M,
            # BLOCK_N=BLOCK_N,
            # BLOCK_K=BLOCK_K,
        )
    
    with Timer(label=f"({batch_size}, {seq_len}):  weighted sum gather timer") as t:
        weighted_sum_gather[ lambda meta: (triton.cdiv(d_hidden, meta["BLOCK_SIZE"]), batch_size * seq_len)](
            ffn1_out,
            token2index,
            topk_scores,
            shared_expert_ffn1,
            moe_out,
            d_hidden,
            batch_size,
            seq_len,
            n_experts_per_token,
            # BLOCK_SIZE=BLOCK_SIZE,
        )
    # sleep 8 ms 
    # benchmark sucks !
    # time.sleep(0.008)
    return moe_out.reshape(batch_size, seq_len, d_hidden) # (batch_size, seq_len, d_hidden)


def custom_kernel(data: input_t) -> output_t:
    input_tensor, weights, config = data
    output = moe_forward(input_tensor, weights, **config)
    return output

python eval.py benchmark /tmp/tmp8e4zzfjq
/usr/local/lib/python3.10/dist-packages/triton/language/semantic.py:1687: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int32
  warnings.warn(
/usr/local/lib/python3.10/dist-packages/triton/language/semantic.py:1687: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int32
  warnings.warn(

(1, 2048):  allocate tensor timer: 5.2683 ms
(1, 2048):  allocate ptr array: 7.8835 ms
(1, 2048):  shared expert timer: 296.9646 ms
(1, 2048):  calculate index timer: 826.1268 ms
(1, 2048):  scatter timer: 149.6131 ms
(1, 2048):  ffn0 timer: 466.0213 ms
(1, 2048):  ffn0 Silu timer: 0.1636 ms
(1, 2048):  ffn1 timer: 1.6501 ms
(1, 2048):  weighted sum gather timer: 185.1501 ms
(1, 2048):  allocate tensor timer: 0.0961 ms
(1, 2048):  allocate ptr array: 0.2117 ms
(1, 2048):  shared expert timer: 0.5550 ms
(1, 2048):  calculate index timer: 0.2015 ms
(1, 2048):  scatter timer: 0.1099 ms
(1, 2048):  ffn0 timer: 3.1500 ms
(1, 2048):  ffn0 Silu timer: 0.1178 ms
(1, 2048):  ffn1 timer: 1.7619 ms
(1, 2048):  weighted sum gather timer: 0.1066 ms
(1, 2048):  allocate tensor timer: 0.0706 ms
(1, 2048):  allocate ptr array: 0.1664 ms
(1, 2048):  shared expert timer: 0.4489 ms
(1, 2048):  calculate index timer: 0.1218 ms
(1, 2048):  scatter timer: 0.1073 ms
(1, 2048):  ffn0 timer: 3.0284 ms
(1, 2048):  ffn0 Silu timer: 0.0722 ms
(1, 2048):  ffn1 timer: 1.5626 ms
(1, 2048):  weighted sum gather timer: 0.0999 ms
(1, 2048):  allocate tensor timer: 0.0610 ms
(1, 2048):  allocate ptr array: 0.1574 ms
(1, 2048):  shared expert timer: 0.4387 ms
(1, 2048):  calculate index timer: 0.1192 ms
(1, 2048):  scatter timer: 0.1421 ms
(1, 2048):  ffn0 timer: 2.9881 ms
(1, 2048):  ffn0 Silu timer: 0.0746 ms
(1, 2048):  ffn1 timer: 1.5900 ms
(1, 2048):  weighted sum gather timer: 0.1097 ms
(1, 2048):  allocate tensor timer: 0.0589 ms
(1, 2048):  allocate ptr array: 0.1507 ms
(1, 2048):  shared expert timer: 0.4520 ms
(1, 2048):  calculate index timer: 0.1144 ms
(1, 2048):  scatter timer: 0.1130 ms
(1, 2048):  ffn0 timer: 3.0086 ms
(1, 2048):  ffn0 Silu timer: 0.0756 ms
(1, 2048):  ffn1 timer: 1.5066 ms
(1, 2048):  weighted sum gather timer: 0.1016 ms
(1, 2048):  allocate tensor timer: 0.0579 ms
(1, 2048):  allocate ptr array: 0.1459 ms
(1, 2048):  shared expert timer: 0.4358 ms
(1, 2048):  calculate index timer: 0.1137 ms
(1, 2048):  scatter timer: 0.1051 ms
(1, 2048):  ffn0 timer: 2.9554 ms
(1, 2048):  ffn0 Silu timer: 0.0741 ms
(1, 2048):  ffn1 timer: 1.5028 ms
(1, 2048):  weighted sum gather timer: 0.0956 ms
(1, 2048):  allocate tensor timer: 0.0594 ms
(1, 2048):  allocate ptr array: 0.1514 ms
(1, 2048):  shared expert timer: 0.4504 ms
(1, 2048):  calculate index timer: 0.1075 ms
(1, 2048):  scatter timer: 0.1042 ms
(1, 2048):  ffn0 timer: 2.9240 ms
(1, 2048):  ffn0 Silu timer: 0.0713 ms
(1, 2048):  ffn1 timer: 1.5385 ms
(1, 2048):  weighted sum gather timer: 0.1016 ms
(1, 2048):  allocate tensor timer: 0.0587 ms
(1, 2048):  allocate ptr array: 0.1476 ms
(1, 2048):  shared expert timer: 0.4573 ms
(1, 2048):  calculate index timer: 0.1080 ms
(1, 2048):  scatter timer: 0.1049 ms
(1, 2048):  ffn0 timer: 2.9581 ms
(1, 2048):  ffn0 Silu timer: 0.0725 ms
(1, 2048):  ffn1 timer: 1.5085 ms
(1, 2048):  weighted sum gather timer: 0.0980 ms
(1, 2048):  allocate tensor timer: 0.0560 ms
(1, 2048):  allocate ptr array: 0.1466 ms
(1, 2048):  shared expert timer: 0.4475 ms
(1, 2048):  calculate index timer: 0.1066 ms
(1, 2048):  scatter timer: 0.1028 ms
(1, 2048):  ffn0 timer: 2.9714 ms
(1, 2048):  ffn0 Silu timer: 0.0720 ms
(1, 2048):  ffn1 timer: 1.4997 ms
(1, 2048):  weighted sum gather timer: 0.0994 ms
(1, 2048):  allocate tensor timer: 0.0565 ms
(1, 2048):  allocate ptr array: 0.1457 ms
(1, 2048):  shared expert timer: 0.4401 ms
(1, 2048):  calculate index timer: 0.1144 ms
(1, 2048):  scatter timer: 0.1075 ms
(1, 2048):  ffn0 timer: 2.9135 ms
(1, 2048):  ffn0 Silu timer: 0.1459 ms
(1, 2048):  ffn1 timer: 1.5435 ms
(1, 2048):  weighted sum gather timer: 0.0980 ms
(1, 2048):  allocate tensor timer: 0.0570 ms
(1, 2048):  allocate ptr array: 0.1476 ms
(1, 2048):  shared expert timer: 0.4640 ms
(1, 2048):  calculate index timer: 0.1082 ms
(1, 2048):  scatter timer: 0.1042 ms
(1, 2048):  ffn0 timer: 3.0470 ms
(1, 2048):  ffn0 Silu timer: 0.1070 ms
(1, 2048):  ffn1 timer: 1.5447 ms
(1, 2048):  weighted sum gather timer: 0.0978 ms
(1, 2048):  allocate tensor timer: 0.0565 ms
(1, 2048):  allocate ptr array: 0.1450 ms
(1, 2048):  shared expert timer: 0.4666 ms
(1, 2048):  calculate index timer: 0.1097 ms
(1, 2048):  scatter timer: 0.0994 ms
(1, 2048):  ffn0 timer: 3.0072 ms
(1, 2048):  ffn0 Silu timer: 0.0696 ms
(1, 2048):  ffn1 timer: 1.5466 ms
(1, 2048):  weighted sum gather timer: 0.0954 ms
(1, 2048):  allocate tensor timer: 0.0570 ms
(1, 2048):  allocate ptr array: 0.1450 ms
(1, 2048):  shared expert timer: 0.4539 ms
(1, 2048):  calculate index timer: 0.1135 ms
(1, 2048):  scatter timer: 0.0978 ms
(1, 2048):  ffn0 timer: 3.0036 ms
(1, 2048):  ffn0 Silu timer: 0.0694 ms
(1, 2048):  ffn1 timer: 1.6150 ms
(1, 2048):  weighted sum gather timer: 0.1016 ms
(1, 2048):  allocate tensor timer: 0.0608 ms
(1, 2048):  allocate ptr array: 0.1531 ms
(1, 2048):  shared expert timer: 0.4447 ms
(1, 2048):  calculate index timer: 0.1121 ms
(1, 2048):  scatter timer: 0.1080 ms
(1, 2048):  ffn0 timer: 2.9113 ms
(1, 2048):  ffn0 Silu timer: 0.0720 ms
(1, 2048):  ffn1 timer: 1.5514 ms
(1, 2048):  weighted sum gather timer: 0.1013 ms
(1, 2048):  allocate tensor timer: 0.0579 ms
(1, 2048):  allocate ptr array: 0.1483 ms
(1, 2048):  shared expert timer: 0.4566 ms
(1, 2048):  calculate index timer: 0.1111 ms
(1, 2048):  scatter timer: 0.1068 ms
(1, 2048):  ffn0 timer: 2.9683 ms
(1, 2048):  ffn0 Silu timer: 0.0703 ms
(1, 2048):  ffn1 timer: 1.4839 ms
(1, 2048):  weighted sum gather timer: 0.1004 ms
(1, 2048):  allocate tensor timer: 0.0560 ms
(1, 2048):  allocate ptr array: 0.1483 ms
(1, 2048):  shared expert timer: 0.4370 ms
(1, 2048):  calculate index timer: 0.1190 ms
(1, 2048):  scatter timer: 0.1049 ms
(1, 2048):  ffn0 timer: 2.9655 ms
(1, 2048):  ffn0 Silu timer: 0.0756 ms
(1, 2048):  ffn1 timer: 1.5159 ms
(1, 2048):  weighted sum gather timer: 0.0985 ms
(1, 2048):  allocate tensor timer: 0.0584 ms
(1, 2048):  allocate ptr array: 0.1481 ms
(1, 2048):  shared expert timer: 0.4420 ms
(1, 2048):  calculate index timer: 0.1082 ms
(1, 2048):  scatter timer: 0.1073 ms
(1, 2048):  ffn0 timer: 2.9292 ms
(1, 2048):  ffn0 Silu timer: 0.0715 ms
(1, 2048):  ffn1 timer: 1.5116 ms
(1, 2048):  weighted sum gather timer: 0.1006 ms
(1, 2048):  allocate tensor timer: 0.0572 ms
(1, 2048):  allocate ptr array: 0.1483 ms
(1, 2048):  shared expert timer: 0.4640 ms
(1, 2048):  calculate index timer: 0.1073 ms
(1, 2048):  scatter timer: 0.1047 ms
(1, 2048):  ffn0 timer: 2.9819 ms
(1, 2048):  ffn0 Silu timer: 0.0727 ms
(1, 2048):  ffn1 timer: 1.4849 ms
(1, 2048):  weighted sum gather timer: 0.0956 ms
(1, 2048):  allocate tensor timer: 0.0570 ms
(1, 2048):  allocate ptr array: 0.1476 ms
(1, 2048):  shared expert timer: 0.4413 ms
(1, 2048):  calculate index timer: 0.1075 ms
(1, 2048):  scatter timer: 0.1037 ms
(1, 2048):  ffn0 timer: 2.9387 ms
(1, 2048):  ffn0 Silu timer: 0.0710 ms
(1, 2048):  ffn1 timer: 1.5156 ms
(1, 2048):  weighted sum gather timer: 0.1023 ms
(1, 2048):  allocate tensor timer: 6.9921 ms
(1, 2048):  allocate ptr array: 0.1774 ms
(1, 2048):  shared expert timer: 0.5012 ms
(1, 2048):  calculate index timer: 0.1595 ms
(1, 2048):  scatter timer: 0.1001 ms
(1, 2048):  ffn0 timer: 3.1157 ms
(1, 2048):  ffn0 Silu timer: 0.0689 ms
(1, 2048):  ffn1 timer: 1.5500 ms
(1, 2048):  weighted sum gather timer: 0.1001 ms
(1, 2048):  allocate tensor timer: 0.0803 ms
(1, 2048):  allocate ptr array: 0.2017 ms
(1, 2048):  shared expert timer: 0.5760 ms
(1, 2048):  calculate index timer: 0.1881 ms
(1, 2048):  scatter timer: 0.1080 ms
(1, 2048):  ffn0 timer: 3.0959 ms
(1, 2048):  ffn0 Silu timer: 0.0722 ms
(1, 2048):  ffn1 timer: 1.6110 ms
(1, 2048):  weighted sum gather timer: 0.1047 ms
(1, 2048):  allocate tensor timer: 0.0665 ms
(1, 2048):  allocate ptr array: 0.1612 ms
(1, 2048):  shared expert timer: 0.4327 ms
(1, 2048):  calculate index timer: 0.1113 ms
(1, 2048):  scatter timer: 0.0982 ms
(1, 2048):  ffn0 timer: 2.9359 ms
(1, 2048):  ffn0 Silu timer: 0.1130 ms
(1, 2048):  ffn1 timer: 1.6150 ms
(1, 2048):  weighted sum gather timer: 0.0973 ms
(1, 2048):  allocate tensor timer: 0.0587 ms
(1, 2048):  allocate ptr array: 0.1502 ms
(1, 2048):  shared expert timer: 0.4492 ms
(1, 2048):  calculate index timer: 0.1144 ms
(1, 2048):  scatter timer: 0.1016 ms
(1, 2048):  ffn0 timer: 2.9631 ms
(1, 2048):  ffn0 Silu timer: 0.0725 ms
(1, 2048):  ffn1 timer: 1.6005 ms
(1, 2048):  weighted sum gather timer: 0.1030 ms
(1, 2048):  allocate tensor timer: 0.0608 ms
(1, 2048):  allocate ptr array: 0.1552 ms
(1, 2048):  shared expert timer: 0.4528 ms
(1, 2048):  calculate index timer: 0.1137 ms
(1, 2048):  scatter timer: 0.1042 ms
(1, 2048):  ffn0 timer: 3.0584 ms
(1, 2048):  ffn0 Silu timer: 0.0718 ms
(1, 2048):  ffn1 timer: 1.5955 ms
(1, 2048):  weighted sum gather timer: 0.0963 ms
(1, 2048):  allocate tensor timer: 0.0577 ms
(1, 2048):  allocate ptr array: 0.1495 ms
(1, 2048):  shared expert timer: 0.4582 ms
(1, 2048):  calculate index timer: 0.1094 ms
(1, 2048):  scatter timer: 0.0966 ms
(1, 2048):  ffn0 timer: 3.0427 ms
(1, 2048):  ffn0 Silu timer: 0.0725 ms
(1, 2048):  ffn1 timer: 1.5125 ms
(1, 2048):  weighted sum gather timer: 0.1016 ms
(1, 2048):  allocate tensor timer: 0.0582 ms
(1, 2048):  allocate ptr array: 0.1507 ms
(1, 2048):  shared expert timer: 0.4449 ms
(1, 2048):  calculate index timer: 0.1173 ms
(1, 2048):  scatter timer: 0.1042 ms
(1, 2048):  ffn0 timer: 2.9597 ms
(1, 2048):  ffn0 Silu timer: 0.0741 ms
(1, 2048):  ffn1 timer: 1.5523 ms
(1, 2048):  weighted sum gather timer: 0.1028 ms
(1, 2048):  allocate tensor timer: 0.0594 ms
(1, 2048):  allocate ptr array: 0.1526 ms
(1, 2048):  shared expert timer: 0.4449 ms
(1, 2048):  calculate index timer: 0.1116 ms
(1, 2048):  scatter timer: 0.1128 ms
(1, 2048):  ffn0 timer: 2.9202 ms
(1, 2048):  ffn0 Silu timer: 0.0756 ms
(1, 2048):  ffn1 timer: 1.5442 ms
(1, 2048):  weighted sum gather timer: 0.1006 ms
(1, 2048):  allocate tensor timer: 0.0570 ms
(1, 2048):  allocate ptr array: 0.1497 ms
(1, 2048):  shared expert timer: 0.4606 ms
(1, 2048):  calculate index timer: 0.1085 ms
(1, 2048):  scatter timer: 0.1061 ms
(1, 2048):  ffn0 timer: 3.0453 ms
(1, 2048):  ffn0 Silu timer: 0.0706 ms
(1, 2048):  ffn1 timer: 1.5430 ms
(1, 2048):  weighted sum gather timer: 0.0944 ms
(1, 2048):  allocate tensor timer: 0.0567 ms
(1, 2048):  allocate ptr array: 0.1457 ms
(1, 2048):  shared expert timer: 0.4447 ms
(1, 2048):  calculate index timer: 0.1104 ms
(1, 2048):  scatter timer: 0.0958 ms
(1, 2048):  ffn0 timer: 3.0828 ms
(1, 2048):  ffn0 Silu timer: 0.0737 ms
(1, 2048):  ffn1 timer: 1.4989 ms
(1, 2048):  weighted sum gather timer: 0.0987 ms
(1, 2048):  allocate tensor timer: 0.0565 ms
(1, 2048):  allocate ptr array: 0.1502 ms
(1, 2048):  shared expert timer: 0.4578 ms
(1, 2048):  calculate index timer: 0.1171 ms
(1, 2048):  scatter timer: 0.1085 ms
(1, 2048):  ffn0 timer: 2.9376 ms
(1, 2048):  ffn0 Silu timer: 0.0749 ms
(1, 2048):  ffn1 timer: 1.5538 ms
(1, 2048):  weighted sum gather timer: 0.0978 ms
(1, 2048):  allocate tensor timer: 0.0565 ms
(1, 2048):  allocate ptr array: 0.1538 ms
(1, 2048):  shared expert timer: 0.4530 ms
(1, 2048):  calculate index timer: 0.1118 ms
(1, 2048):  scatter timer: 0.1075 ms
(1, 2048):  ffn0 timer: 3.0627 ms
(1, 2048):  ffn0 Silu timer: 0.0710 ms
(1, 2048):  ffn1 timer: 1.5559 ms
(1, 2048):  weighted sum gather timer: 0.0970 ms
(1, 2048):  allocate tensor timer: 0.0563 ms
(1, 2048):  allocate ptr array: 0.1440 ms
(1, 2048):  shared expert timer: 0.4668 ms
(1, 2048):  calculate index timer: 0.1075 ms
(1, 2048):  scatter timer: 0.0994 ms
(1, 2048):  ffn0 timer: 3.0174 ms
(1, 2048):  ffn0 Silu timer: 0.1118 ms
(1, 2048):  ffn1 timer: 1.5566 ms
(1, 2048):  weighted sum gather timer: 0.0949 ms
(1, 2048):  allocate tensor timer: 0.0563 ms
(1, 2048):  allocate ptr array: 0.1433 ms
(1, 2048):  shared expert timer: 0.4652 ms
(1, 2048):  calculate index timer: 0.1147 ms
(1, 2048):  scatter timer: 0.0970 ms
(1, 2048):  ffn0 timer: 2.9879 ms
(1, 2048):  ffn0 Silu timer: 0.1128 ms
(1, 2048):  ffn1 timer: 1.5211 ms
(1, 2048):  weighted sum gather timer: 0.0966 ms
(1, 2048):  allocate tensor timer: 0.0567 ms
(1, 2048):  allocate ptr array: 0.1454 ms
(1, 2048):  shared expert timer: 0.4487 ms
(1, 2048):  calculate index timer: 0.1097 ms
(1, 2048):  scatter timer: 0.0975 ms
(1, 2048):  ffn0 timer: 2.9469 ms
(1, 2048):  ffn0 Silu timer: 0.0744 ms
(1, 2048):  ffn1 timer: 1.5936 ms
(1, 2048):  weighted sum gather timer: 0.0949 ms
(1, 2048):  allocate tensor timer: 0.0553 ms
(1, 2048):  allocate ptr array: 0.1445 ms
(1, 2048):  shared expert timer: 0.4449 ms
(1, 2048):  calculate index timer: 0.1087 ms
(1, 2048):  scatter timer: 0.0958 ms
(1, 2048):  ffn0 timer: 2.9712 ms
(1, 2048):  ffn0 Silu timer: 0.0706 ms
(1, 2048):  ffn1 timer: 1.5960 ms
(1, 2048):  weighted sum gather timer: 0.0963 ms
(1, 2048):  allocate tensor timer: 0.0587 ms
(1, 2048):  allocate ptr array: 0.1469 ms
(1, 2048):  shared expert timer: 0.4656 ms
(1, 2048):  calculate index timer: 0.1106 ms
(1, 2048):  scatter timer: 0.0989 ms
(1, 2048):  ffn0 timer: 3.0234 ms
(1, 2048):  ffn0 Silu timer: 0.0699 ms
(1, 2048):  ffn1 timer: 1.5335 ms
(1, 2048):  weighted sum gather timer: 0.0951 ms
(1, 2048):  allocate tensor timer: 0.0570 ms
(1, 2048):  allocate ptr array: 0.1466 ms
(1, 2048):  shared expert timer: 0.4621 ms
(1, 2048):  calculate index timer: 0.1075 ms
(1, 2048):  scatter timer: 0.1271 ms
(1, 2048):  ffn0 timer: 2.9917 ms
(1, 2048):  ffn0 Silu timer: 0.1113 ms
(1, 2048):  ffn1 timer: 1.5900 ms
(1, 2048):  weighted sum gather timer: 0.0939 ms
(1, 2048):  allocate tensor timer: 0.0565 ms
(1, 2048):  allocate ptr array: 0.1481 ms
(1, 2048):  shared expert timer: 0.4511 ms
(1, 2048):  calculate index timer: 0.1094 ms
(1, 2048):  scatter timer: 0.1001 ms
(1, 2048):  ffn0 timer: 2.9390 ms
(1, 2048):  ffn0 Silu timer: 0.0699 ms
(1, 2048):  ffn1 timer: 1.5764 ms
(1, 2048):  weighted sum gather timer: 0.0944 ms
(1, 2048):  allocate tensor timer: 0.0582 ms
(1, 2048):  allocate ptr array: 0.1507 ms
(1, 2048):  shared expert timer: 0.4582 ms
(1, 2048):  calculate index timer: 0.1087 ms
(1, 2048):  scatter timer: 0.0992 ms
(1, 2048):  ffn0 timer: 2.9604 ms
(1, 2048):  ffn0 Silu timer: 0.1082 ms
(1, 2048):  ffn1 timer: 1.5731 ms
(1, 2048):  weighted sum gather timer: 0.0968 ms
(1, 2048):  allocate tensor timer: 0.0563 ms
(1, 2048):  allocate ptr array: 0.1452 ms
(1, 2048):  shared expert timer: 0.4580 ms
(1, 2048):  calculate index timer: 0.1082 ms
(1, 2048):  scatter timer: 0.0951 ms
(1, 2048):  ffn0 timer: 3.0150 ms
(1, 2048):  ffn0 Silu timer: 0.0699 ms
(1, 2048):  ffn1 timer: 1.5597 ms
(1, 2048):  weighted sum gather timer: 0.0973 ms
(1, 2048):  allocate tensor timer: 0.0570 ms
(1, 2048):  allocate ptr array: 0.1478 ms
(1, 2048):  shared expert timer: 0.4449 ms
(1, 2048):  calculate index timer: 0.1106 ms
(1, 2048):  scatter timer: 0.0992 ms
(1, 2048):  ffn0 timer: 2.9600 ms
(1, 2048):  ffn0 Silu timer: 0.0706 ms
(1, 2048):  ffn1 timer: 1.5595 ms
(1, 2048):  weighted sum gather timer: 0.0980 ms
(1, 2048):  allocate tensor timer: 0.0579 ms
(1, 2048):  allocate ptr array: 0.1454 ms
(1, 2048):  shared expert timer: 0.4535 ms
(1, 2048):  calculate index timer: 0.1106 ms
(1, 2048):  scatter timer: 0.1054 ms
(1, 2048):  ffn0 timer: 2.9562 ms
(1, 2048):  ffn0 Silu timer: 0.0691 ms
(1, 2048):  ffn1 timer: 1.6255 ms
(1, 2048):  weighted sum gather timer: 0.0963 ms
(1, 2048):  allocate tensor timer: 0.0567 ms
(1, 2048):  allocate ptr array: 0.1483 ms
(1, 2048):  shared expert timer: 0.4551 ms
(1, 2048):  calculate index timer: 0.1101 ms
(1, 2048):  scatter timer: 0.0999 ms
(1, 2048):  ffn0 timer: 2.9850 ms
(1, 2048):  ffn0 Silu timer: 0.1087 ms
(1, 2048):  ffn1 timer: 1.5366 ms
(1, 2048):  weighted sum gather timer: 0.0930 ms
(1, 2048):  allocate tensor timer: 0.0565 ms
(1, 2048):  allocate ptr array: 0.1452 ms
(1, 2048):  shared expert timer: 0.4644 ms
(1, 2048):  calculate index timer: 0.1128 ms
(1, 2048):  scatter timer: 0.1290 ms
(1, 2048):  ffn0 timer: 2.9669 ms
(1, 2048):  ffn0 Silu timer: 0.0741 ms
(1, 2048):  ffn1 timer: 1.5099 ms
(1, 2048):  weighted sum gather timer: 0.0968 ms
(1, 2048):  allocate tensor timer: 0.0601 ms
[...] 1170 lines omitted