#!POPCORN leaderboard amd-mixture-of-experts
#!POPCORN gpu MI300

from task import input_t, output_t
# from reference import check_implementation, generate_input # Assuming these are for testing
import math
import torch
import torch.nn.functional as F
from typing import Dict, Tuple
import triton
import triton.language as tl
from collections import defaultdict

# Existing matmul kernel (unchanged for non-batched, adapted for batched below)
@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)

    a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)

    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k_loop in range(0, K, BLOCK_SIZE_K):
        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

    c_ptrs = c_ptr + (offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn)
    tl.store(c_ptrs, acc.to(tl.float16), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

def triton_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    M, K = a.shape
    _K, N = b.shape
    if K != _K:
        raise ValueError(f"Matrix dimensions mismatch for matmul: A.shape[1]={K}, B.shape[0]={_K}")
    c = torch.empty((M, N), device=a.device, dtype=torch.float16)

    BLOCK_SIZE_M = 64 # Or some other tuned value
    BLOCK_SIZE_N = 64
    BLOCK_SIZE_K = 32
    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))

    matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K
    )
    return c


# --- Batched Kernels ---
@triton.jit
def batched_fused_gate_up_kernel(
    input_ptr, w_gate_ptr, w_up_ptr, output_ptr,
    actual_tokens_per_expert_ptr, # Pointer to a tensor holding actual token counts for each expert
    num_experts, max_tokens_curr_expert_dim, K_dim, N_dim, # K_dim = d_hidden, N_dim = d_expert
    stride_in_exp, stride_in_m, stride_in_k,          # input_ptr strides (expert, token_idx, feature_idx)
    stride_wg_exp, stride_wg_k, stride_wg_n,          # w_gate_ptr strides
    stride_wu_exp, stride_wu_k, stride_wu_n,          # w_up_ptr strides
    stride_out_exp, stride_out_m, stride_out_n,       # output_ptr strides
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr
):
    pid_expert = tl.program_id(0) # Parallelize over experts
    pid_m = tl.program_id(1)      # Parallelize over token blocks for an expert
    pid_n = tl.program_id(2)      # Parallelize over feature blocks (d_expert)

    # actual M for this specific expert
    M = tl.load(actual_tokens_per_expert_ptr + pid_expert)

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k_loop = tl.arange(0, BLOCK_SIZE_K)

    # Pointers for the current expert
    expert_input_ptr = input_ptr + pid_expert * stride_in_exp
    expert_w_gate_ptr = w_gate_ptr + pid_expert * stride_wg_exp
    expert_w_up_ptr = w_up_ptr + pid_expert * stride_wu_exp
    expert_output_ptr = output_ptr + pid_expert * stride_out_exp

    acc_gate = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    acc_up = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    # Mask for valid tokens for the current expert and block
    m_mask = (offs_m[:, None] < M)

    for k_start in range(0, K_dim, BLOCK_SIZE_K):
        offs_k = k_start + offs_k_loop
        k_mask = (offs_k[None, :] < K_dim)

        a_ptrs = expert_input_ptr + (offs_m[:, None] * stride_in_m + offs_k[None, :] * stride_in_k)
        b_gate_ptrs = expert_w_gate_ptr + (offs_k[:, None] * stride_wg_k + offs_n[None, :] * stride_wg_n)
        b_up_ptrs = expert_w_up_ptr + (offs_k[:, None] * stride_wu_k + offs_n[None, :] * stride_wu_n)

        a = tl.load(a_ptrs, mask=m_mask & k_mask, other=0.0)
        b_gate = tl.load(b_gate_ptrs, mask=(offs_k[:, None] < K_dim) & (offs_n[None, :] < N_dim), other=0.0) # K for weights is K_dim
        b_up = tl.load(b_up_ptrs, mask=(offs_k[:, None] < K_dim) & (offs_n[None, :] < N_dim), other=0.0)   # K for weights is K_dim

        acc_gate += tl.dot(a, b_gate)
        acc_up += tl.dot(a, b_up)

    output_val = (acc_gate * tl.sigmoid(acc_gate)) * acc_up

    # Mask for storing output
    output_m_mask = offs_m[:, None] < M
    output_n_mask = offs_n[None, :] < N_dim
    final_mask = output_m_mask & output_n_mask

    c_ptrs = expert_output_ptr + (offs_m[:, None] * stride_out_m + offs_n[None, :] * stride_out_n)
    tl.store(c_ptrs, output_val.to(tl.float16), mask=final_mask)


@triton.jit
def batched_matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    actual_tokens_per_expert_ptr, # Pointer to actual token counts for each expert
    num_experts, max_tokens_curr_expert_dim, N_dim, K_dim, # M is max_tokens_curr_expert_dim, N is d_hidden (output), K is d_expert
    stride_a_exp, stride_a_m, stride_a_k,
    stride_b_exp, stride_b_k, stride_b_n,
    stride_c_exp, stride_c_m, stride_c_n,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr
):
    pid_expert = tl.program_id(0)
    pid_m = tl.program_id(1)
    pid_n = tl.program_id(2)

    M = tl.load(actual_tokens_per_expert_ptr + pid_expert) # Actual M for this specific expert

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k_loop = tl.arange(0, BLOCK_SIZE_K)

    expert_a_ptr = a_ptr + pid_expert * stride_a_exp
    expert_b_ptr = b_ptr + pid_expert * stride_b_exp
    expert_c_ptr = c_ptr + pid_expert * stride_c_exp
    
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    m_mask = (offs_m[:, None] < M)

    for k_start in range(0, K_dim, BLOCK_SIZE_K):
        offs_k = k_start + offs_k_loop
        k_mask = (offs_k[None, :] < K_dim)

        a_ptrs = expert_a_ptr + (offs_m[:, None] * stride_a_m + offs_k[None, :] * stride_a_k)
        b_ptrs = expert_b_ptr + (offs_k[:, None] * stride_b_k + offs_n[None, :] * stride_b_n)
        
        a = tl.load(a_ptrs, mask=m_mask & k_mask, other=0.0)
        # For B, K is K_dim, N is N_dim
        b = tl.load(b_ptrs, mask=(offs_k[:, None] < K_dim) & (offs_n[None, :] < N_dim), other=0.0)
        
        acc += tl.dot(a, b)

    output_m_mask = offs_m[:, None] < M
    output_n_mask = offs_n[None, :] < N_dim # N_dim is the output feature dimension
    final_mask = output_m_mask & output_n_mask
    
    c_ptrs = expert_c_ptr + (offs_m[:, None] * stride_c_m + offs_n[None, :] * stride_c_n)
    tl.store(c_ptrs, acc.to(tl.float16), mask=final_mask)


def custom_kernel(data: input_t) -> output_t:
    input_tensor_arg, weights_dict, config_dict = data

    d_hidden = config_dict["d_hidden"]
    d_expert = config_dict["d_expert"]
    n_routed_experts = config_dict["n_routed_experts"] # This is num_experts
    n_experts_per_token = config_dict["n_experts_per_token"]
    batch_size = config_dict["batch_size"]
    seq_len = config_dict["seq_len"]
    num_experts = n_routed_experts

    x_input = input_tensor_arg.view(-1, d_hidden)
    total_tokens_x_experts_per_token = x_input.shape[0] * n_experts_per_token # For flattened view

    # Router computation
    W_g = weights_dict['router.weight'].t() # (d_hidden, num_experts)
    logits = triton_matmul(x_input, W_g) # (total_tokens, num_experts)
    scores = F.softmax(logits, dim=-1)
    topk_scores, topk_indices = torch.topk(scores, k=n_experts_per_token, dim=-1) # (total_tokens, n_experts_per_token)
    topk_scores_fp16 = topk_scores.to(torch.float16)

    # --- Start of dispatch logic (largely from original, for preparing batched inputs) ---
    # flat_expert_indices: For each token, which expert it's routed to (0 to num_experts-1)
    # e.g., if token0 -> exp2, exp5; token1 -> exp0, exp2
    # topk_indices = [[2,5], [0,2], ...]
    # flat_expert_indices = [2,5,0,2,...]
    flat_expert_indices = topk_indices.view(-1) # Shape: (total_tokens * n_experts_per_token)

    # idxs_argsort: gives the new positions of elements if flat_expert_indices were sorted.
    # Example: flat_expert_indices = [1,0,1,0], idxs_argsort = [1,3,0,2] (so sorted is [0,0,1,1])
    idxs_argsort = flat_expert_indices.argsort()

    # token_idxs_dispatch: original token index (0 to total_tokens-1) for each entry in the sorted list.
    # If token_i is routed to k experts, it will appear k times.
    # idxs_argsort // n_experts_per_token correctly maps back to the original token index
    # because the first n_experts_per_token entries in idxs_argsort belong to token_0 (if sorted), etc.
    # No, this is: for the i-th element in the sorted expert assignments, which original token (0..total_tokens-1) did it come from.
    token_idxs_dispatch = idxs_argsort // n_experts_per_token # Shape: (total_tokens * n_experts_per_token)

    # k_dispatch: for the i-th element in the sorted expert assignments, which k-th choice (0..n_experts_per_token-1) was it for its original token.
    k_dispatch = idxs_argsort % n_experts_per_token # Shape: (total_tokens * n_experts_per_token)

    # counts: how many tokens are routed to each expert (sum of all choices)
    # e.g. counts[0] = number of times expert 0 was chosen (can be > total_tokens if n_experts_per_token > 1)
    counts = flat_expert_indices.bincount(minlength=num_experts) # Shape: (num_experts)
    tokens_per_expert_cumsum = counts.cumsum(dim=0) # Shape: (num_experts)
    max_tokens_per_expert = counts.max().item()
    if max_tokens_per_expert == 0: # Handle case where no tokens are routed to any expert (e.g. for testing)
        max_tokens_per_expert = 1 # Avoid zero-sized dimensions for tensors

    # --- Prepare batched inputs for experts ---
    # Padded tensors for batched expert processing
    # (num_experts, max_tokens_per_expert, d_hidden)
    batched_expert_inputs = torch.zeros(
        (num_experts, max_tokens_per_expert, d_hidden),
        device=x_input.device, dtype=x_input.dtype
    )
    # To store original token indices and k-values for scattering results later
    # (num_experts, max_tokens_per_expert)
    original_token_indices_for_scatter = torch.zeros(
        (num_experts, max_tokens_per_expert),
        device=x_input.device, dtype=torch.long
    )
    k_values_for_scatter = torch.zeros(
        (num_experts, max_tokens_per_expert),
        device=x_input.device, dtype=torch.long
    )
    # Actual number of tokens processed by each expert (for masking in kernels)
    actual_tokens_per_expert_list = [0] * num_experts

    current_pos_in_expert_batch = torch.zeros(num_experts, dtype=torch.long, device=x_input.device)

    # This loop iterates through the sorted expert assignments
    # token_idxs_dispatch[i] is the original token index
    # flat_expert_indices[idxs_argsort[i]] is the expert_id for this assignment
    # k_dispatch[i] is the k-th choice for that original token
    for i in range(token_idxs_dispatch.shape[0]):
        original_token_idx = token_idxs_dispatch[i]
        expert_id = flat_expert_indices[idxs_argsort[i]] # This is the expert this token is routed to
        k_val = k_dispatch[i]

        pos = current_pos_in_expert_batch[expert_id].item()
        if pos < max_tokens_per_expert: # Should always be true if max_tokens_per_expert is calculated correctly from 'counts'
            batched_expert_inputs[expert_id, pos, :] = x_input[original_token_idx, :]
            original_token_indices_for_scatter[expert_id, pos] = original_token_idx
            k_values_for_scatter[expert_id, pos] = k_val
            current_pos_in_expert_batch[expert_id] += 1
    
    actual_tokens_per_expert_tensor = counts.to(device=x_input.device, dtype=torch.int32)


    # Stack expert weights
    # W_gate and W_up: (d_hidden, d_expert) -> (num_experts, d_hidden, d_expert)
    # W_down: (d_expert, d_hidden) -> (num_experts, d_expert, d_hidden)
    
    # Note on weight layout for kernels:
    # fused_gate_up: input (M,K), weight (K,N)
    #   Here, input is (max_tokens, d_hidden), W_gate/W_up is (d_hidden, d_expert)
    #   So K_dim = d_hidden, N_dim = d_expert
    # matmul for W_down: input (M,K), weight (K,N)
    #   Here, input is (max_tokens, d_expert), W_down is (d_expert, d_hidden)
    #   So K_dim = d_expert, N_dim = d_hidden
    
    W_gate_list = [weights_dict[f'experts.{i}.0.weight'] for i in range(num_experts)] # List of (d_hidden, d_expert)
    W_up_list = [weights_dict[f'experts.{i}.1.weight'] for i in range(num_experts)]   # List of (d_hidden, d_expert)
    W_down_list = [weights_dict[f'experts.{i}.2.weight'] for i in range(num_experts)] # List of (d_expert, d_hidden)

    W_gate_stacked = torch.stack(W_gate_list, dim=0) # (num_experts, d_hidden, d_expert)
    W_up_stacked = torch.stack(W_up_list, dim=0)     # (num_experts, d_hidden, d_expert)
    W_down_stacked = torch.stack(W_down_list, dim=0) # (num_experts, d_expert, d_hidden)

    # Batched computation
    # (num_experts, max_tokens_per_expert, d_expert)
    batched_intermediate = torch.empty(
        (num_experts, max_tokens_per_expert, d_expert),
        device=x_input.device, dtype=x_input.dtype
    )

    BLOCK_SIZE_M_FUSED = 64  # Tune these
    BLOCK_SIZE_N_FUSED = 64
    BLOCK_SIZE_K_FUSED = 32

    grid_fused = (
        num_experts,
        triton.cdiv(max_tokens_per_expert, BLOCK_SIZE_M_FUSED),
        triton.cdiv(d_expert, BLOCK_SIZE_N_FUSED)
    )
    
    # K_dim for fused is d_hidden, N_dim for fused is d_expert
    batched_fused_gate_up_kernel[grid_fused](
        batched_expert_inputs, W_gate_stacked, W_up_stacked, batched_intermediate,
        actual_tokens_per_expert_tensor,
        num_experts, max_tokens_per_expert, d_hidden, d_expert, # M_dim, K_dim, N_dim
        batched_expert_inputs.stride(0), batched_expert_inputs.stride(1), batched_expert_inputs.stride(2),
        W_gate_stacked.stride(0), W_gate_stacked.stride(1), W_gate_stacked.stride(2),
        W_up_stacked.stride(0), W_up_stacked.stride(1), W_up_stacked.stride(2),
        batched_intermediate.stride(0), batched_intermediate.stride(1), batched_intermediate.stride(2),
        BLOCK_SIZE_M=BLOCK_SIZE_M_FUSED,
        BLOCK_SIZE_N=BLOCK_SIZE_N_FUSED,
        BLOCK_SIZE_K=BLOCK_SIZE_K_FUSED
    )
    
    # (num_experts, max_tokens_per_expert, d_hidden)
    batched_expert_out_unweighted = torch.empty(
        (num_experts, max_tokens_per_expert, d_hidden),
        device=x_input.device, dtype=x_input.dtype
    )

    BLOCK_SIZE_M_MM = 64 # Tune these
    BLOCK_SIZE_N_MM = 64
    BLOCK_SIZE_K_MM = 32

    grid_matmul_down = (
        num_experts,
        triton.cdiv(max_tokens_per_expert, BLOCK_SIZE_M_MM),
        triton.cdiv(d_hidden, BLOCK_SIZE_N_MM) # Output dim is d_hidden
    )
    
    # For W_down matmul: input is batched_intermediate (max_tokens, d_expert)
    # weight is W_down_stacked (d_expert, d_hidden)
    # K_dim is d_expert, N_dim is d_hidden
    batched_matmul_kernel[grid_matmul_down](
        batched_intermediate, W_down_stacked, batched_expert_out_unweighted,
        actual_tokens_per_expert_tensor,
        num_experts, max_tokens_per_expert, d_hidden, d_expert, # M_dim, N_dim, K_dim
        batched_intermediate.stride(0), batched_intermediate.stride(1), batched_intermediate.stride(2),
        W_down_stacked.stride(0), W_down_stacked.stride(1), W_down_stacked.stride(2),
        batched_expert_out_unweighted.stride(0), batched_expert_out_unweighted.stride(1), batched_expert_out_unweighted.stride(2),
        BLOCK_SIZE_M=BLOCK_SIZE_M_MM,
        BLOCK_SIZE_N=BLOCK_SIZE_N_MM,
        BLOCK_SIZE_K=BLOCK_SIZE_K_MM
    )

    # Scatter results back
    # expert_outputs is (total_tokens, n_experts_per_token, d_hidden)
    expert_outputs = torch.zeros(
        (x_input.shape[0], n_experts_per_token, d_hidden),
        device=x_input.device, dtype=x_input.dtype
    )

    # Create a mask for valid, non-padded entries in batched outputs
    # (num_experts, max_tokens_per_expert)
    valid_entry_mask = torch.arange(max_tokens_per_expert, device=x_input.device)[None, :] < actual_tokens_per_expert_tensor[:, None]

    # Get the data from valid entries
    # These will be 1D tensors of length sum(actual_tokens_per_expert_tensor)
    valid_data_to_scatter = batched_expert_out_unweighted[valid_entry_mask]
    valid_original_token_indices = original_token_indices_for_scatter[valid_entry_mask]
    valid_k_values = k_values_for_scatter[valid_entry_mask]
    
    expert_outputs[valid_original_token_indices, valid_k_values, :] = valid_data_to_scatter
    
    # Combine with scores
    # topk_scores_fp16 is (total_tokens, n_experts_per_token)
    # expert_outputs is (total_tokens, n_experts_per_token, d_hidden)
    routed_output = (expert_outputs * topk_scores_fp16.unsqueeze(-1)).sum(dim=1) # (total_tokens, d_hidden)

    # --- Shared expert computation (original fused kernel can be reused if d_expert is the same) ---
    # For simplicity, reusing the original non-batched fused kernel structure
    # Assuming 'shared_experts' has a similar structure to one routed expert
    W_gate_shared = weights_dict['shared_experts.0.weight'] # (d_hidden, d_expert_shared)
    W_up_shared = weights_dict['shared_experts.1.weight']   # (d_hidden, d_expert_shared)
    W_down_shared = weights_dict['shared_experts.2.weight'] # (d_expert_shared, d_hidden_shared)
    
    # Determine d_expert_shared, may or may not be same as d_expert for routed ones
    d_expert_shared = W_gate_shared.shape[1] # Or W_up_shared.shape[1]
    d_hidden_shared_output = W_down_shared.shape[1]


    intermediate_shared = torch.empty((x_input.shape[0], d_expert_shared), device=x_input.device, dtype=x_input.dtype)
    
    # Reusing the original fused_gate_up_kernel definition (if it exists standalone, or define one)
    # For this, we need a non-batched fused kernel. Your original one was fine but not extracted.
    # Let's define a simple one here for clarity or adapt your original one.
    @triton.jit
    def single_fused_gate_up_kernel( # Same as your original fused_gate_up_kernel
        input_ptr, w_gate_ptr, w_up_ptr, output_ptr,
        M, K_dim, N_dim, # M=num_tokens, K_dim=d_hidden, N_dim=d_expert_shared
        stride_im, stride_ik,
        stride_gk, stride_gn, # W_gate is (K_dim, N_dim)
        stride_uk, stride_un, # W_up is (K_dim, N_dim)
        stride_om, stride_on,
        BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr
    ):
        pid_m = tl.program_id(0)
        pid_n = tl.program_id(1)
        
        offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
        offs_k_loop = tl.arange(0, BLOCK_SIZE_K)
        
        acc_gate = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        acc_up = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        
        for k_start in range(0, K_dim, BLOCK_SIZE_K):
            offs_k = k_start + offs_k_loop
            a_ptrs = input_ptr + (offs_m[:, None] * stride_im + offs_k[None, :] * stride_ik)
            b_gate_ptrs = w_gate_ptr + (offs_k[:, None] * stride_gk + offs_n[None, :] * stride_gn)
            b_up_ptrs = w_up_ptr + (offs_k[:, None] * stride_uk + offs_n[None, :] * stride_un)
            
            a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < K_dim), other=0.0)
            b_gate = tl.load(b_gate_ptrs, mask=(offs_k[:, None] < K_dim) & (offs_n[None, :] < N_dim), other=0.0)
            b_up = tl.load(b_up_ptrs, mask=(offs_k[:, None] < K_dim) & (offs_n[None, :] < N_dim), other=0.0)
            
            acc_gate += tl.dot(a, b_gate)
            acc_up += tl.dot(a, b_up)
            
        output_val = (acc_gate * tl.sigmoid(acc_gate)) * acc_up
        
        tl.store(output_ptr + (offs_m[:, None] * stride_om + offs_n[None, :] * stride_on),
                 output_val.to(tl.float16),
                 mask=(offs_m[:, None] < M) & (offs_n[None, :] < N_dim))

    grid_shared = (triton.cdiv(x_input.shape[0], BLOCK_SIZE_M_FUSED), triton.cdiv(d_expert_shared, BLOCK_SIZE_N_FUSED))
    single_fused_gate_up_kernel[grid_shared](
        x_input, W_gate_shared, W_up_shared, intermediate_shared,
        x_input.shape[0], d_hidden, d_expert_shared, # M, K_dim, N_dim
        x_input.stride(0), x_input.stride(1),
        W_gate_shared.stride(0), W_gate_shared.stride(1),
        W_up_shared.stride(0), W_up_shared.stride(1),
        intermediate_shared.stride(0), intermediate_shared.stride(1),
        BLOCK_SIZE_M=BLOCK_SIZE_M_FUSED, BLOCK_SIZE_N=BLOCK_SIZE_N_FUSED, BLOCK_SIZE_K=BLOCK_SIZE_K_FUSED
    )
    
    shared_output = triton_matmul(intermediate_shared, W_down_shared) # (total_tokens, d_hidden_shared_output)
    
    # Pad or truncate shared_output to match d_hidden if necessary
    if shared_output.shape[-1] != d_hidden:
        if shared_output.shape[-1] < d_hidden:
            padding_amount = d_hidden - shared_output.shape[-1]
            shared_output = F.pad(shared_output, (0, padding_amount), "constant", 0.0)
        else:
            shared_output = shared_output[..., :d_hidden]
            
    output_combined = routed_output + shared_output
    return output_combined.view(batch_size, seq_len, d_hidden)
python eval.py benchmark /tmp/tmpt0sydvgb