#!POPCORN leaderboard amd-mixture-of-experts
#!POPCORN gpu MI300
from task import input_t, output_t

import torch
import torch.nn.functional as F
from typing import Dict, Tuple
import triton
import triton.language as tl
from collections import defaultdict

# --- Debugging Setup ---
DEBUG_MODE = True  # << SET THIS TO True TO ENABLE DEBUGGING FEATURES >>
# To make this easily changeable, you might consider setting it via an environment variable.

if DEBUG_MODE:
    print("=" * 60)
    print("!!!!!!!!!!!!!!!!! DEBUG MODE IS ENABLED !!!!!!!!!!!!!!!!!")
    print("=" * 60)
    # Fixed configurations for debugging
    # Ensure these are simple and use powers of 2 for block sizes.
    # Adjust if these specific sizes cause issues on MI300 (e.g., too small/large leading to other errors)
    DEBUG_CONFIG_MATMUL = triton.Config({
        'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32,
        'num_stages': 2, 'num_warps': 4
    })
    DEBUG_CONFIG_SCATTER = triton.Config({
        'BLOCK_SIZE_N_expert': 64, 'BLOCK_SIZE_D_feature': 64,
        'num_warps': 4
    })

    # Counters to limit debug output to the first few calls per kernel type
    _KERNEL_DEBUG_CALL_COUNTS = defaultdict(int)
    _MAX_DEBUG_CALLS_PER_KERNEL_TYPE = 1 # Print comparison for the first call only

    def compare_outputs_for_debug(kernel_name: str, 
                                  triton_output: torch.Tensor, 
                                  pytorch_output: torch.Tensor, 
                                  inputs_dict: Dict[str, torch.Tensor],
                                  atol=1e-3, rtol=1e-3): # Adjusted tolerances for fp16
        _KERNEL_DEBUG_CALL_COUNTS[kernel_name] += 1
        if _KERNEL_DEBUG_CALL_COUNTS[kernel_name] > _MAX_DEBUG_CALLS_PER_KERNEL_TYPE:
            return True # Skip further checks for this kernel type
        print()
        print(f"--- Debugging Kernel: {kernel_name} (Call #{_KERNEL_DEBUG_CALL_COUNTS[kernel_name]}) ---")
        
        # Print input shapes and a few values
        # print("Inputs:")
        # for name, tensor in inputs_dict.items():
        #     print(f"  {name}: shape={tensor.shape}, dtype={tensor.dtype}, device={tensor.device}")
        #     if tensor.numel() > 0:
        #         print(f"    Values (first 5 if possible): {tensor.flatten()[:5]}")


        triton_output_for_compare = triton_output.to(pytorch_output.dtype) # Match dtypes for comparison
        
        # Basic check
        are_close = torch.allclose(triton_output_for_compare, pytorch_output, atol=atol, rtol=rtol)

        if not are_close:
            print(f"!!!!!! MISMATCH DETECTED IN {kernel_name} !!!!!!")
            print(f"Triton output: shape={triton_output.shape}, dtype={triton_output.dtype}, device={triton_output.device}")
            print(f"PyTorch ref:   shape={pytorch_output.shape}, dtype={pytorch_output.dtype}, device={pytorch_output.device}")

            diff = torch.abs(triton_output_for_compare - pytorch_output)
            max_abs_diff = torch.max(diff).item()
            
            num_mismatched_elements = torch.sum(diff > (atol + rtol * torch.abs(pytorch_output))).item()
            
            print(f"Max absolute difference: {max_abs_diff:.6e}")
            print(f"Number of mismatched elements (atol={atol}, rtol={rtol}): {num_mismatched_elements} / {triton_output.numel()}")

            # Find and print first few mismatched elements
            mismatched_indices = (diff > (atol + rtol * torch.abs(pytorch_output))).nonzero(as_tuple=True)
            num_to_print = min(5, len(mismatched_indices[0]))

            if num_to_print > 0:
                print("First few mismatches (Triton vs PyTorch @ index):")
                for i in range(num_to_print):
                    idx_tuple = tuple(dim_idx[i].item() for dim_idx in mismatched_indices)
                    print(f"  Idx {idx_tuple}: {triton_output_for_compare[idx_tuple]:.4f} vs {pytorch_output[idx_tuple]:.4f} (Diff: {diff[idx_tuple]:.4f})")
            
            # For very detailed debugging, you might save tensors to files here:
            # torch.save(triton_output, f"{kernel_name}_triton_output.pt")
            # torch.save(pytorch_output, f"{kernel_name}_pytorch_output.pt")
            # for name, tensor in inputs_dict.items():
            #    torch.save(tensor, f"{kernel_name}_input_{name}.pt")
            # raise ValueError(f"Debug Mismatch in {kernel_name}. Halting.") # Optionally halt
            return False
        else:
            print(f"Outputs for {kernel_name} MATCH PyTorch reference.")
            return True
else: # Not DEBUG_MODE
    # Define dummy compare function if not in debug mode, or ensure it's not called
    def compare_outputs_for_debug(*args, **kwargs):
        return True

# --- Original Autotune Configs (used if DEBUG_MODE is False) ---
AUTOTUNE_CONFIGS_MATMUL = [
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32,  'num_stages': 2, 'num_warps': 4}),
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32,  'num_stages': 2, 'num_warps': 4}),
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32,  'num_stages': 2, 'num_warps': 4}),
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64,  'num_stages': 3, 'num_warps': 8}),
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64,  'num_stages': 3, 'num_warps': 8}),
    # ... (add all your original configs here)
    triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64,  'num_stages': 3, 'num_warps': 8}),
    triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 128, 'num_stages': 2, 'num_warps': 8}),
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'num_stages': 2, 'num_warps': 8}),
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'num_stages': 2, 'num_warps': 4}),
]
AUTOTUNE_KEY_MATMUL = ['M', 'N', 'K', 'stride_am', 'stride_ak', 'stride_bk', 'stride_bn']

AUTOTUNE_CONFIGS_SCATTER = [
    triton.Config({'BLOCK_SIZE_N_expert': 32,  'BLOCK_SIZE_D_feature': 64,  'num_warps': 4}),
    triton.Config({'BLOCK_SIZE_N_expert': 64,  'BLOCK_SIZE_D_feature': 32,  'num_warps': 4}),
    # ... (add all your original configs here)
    triton.Config({'BLOCK_SIZE_N_expert': 64,  'BLOCK_SIZE_D_feature': 64,  'num_warps': 4}),
    triton.Config({'BLOCK_SIZE_N_expert': 128, 'BLOCK_SIZE_D_feature': 64,  'num_warps': 8}),
    triton.Config({'BLOCK_SIZE_N_expert': 64,  'BLOCK_SIZE_D_feature': 128, 'num_warps': 8}),
    triton.Config({'BLOCK_SIZE_N_expert': 128, 'BLOCK_SIZE_D_feature': 128, 'num_warps': 8}),
    triton.Config({'BLOCK_SIZE_N_expert': 256, 'BLOCK_SIZE_D_feature': 64,  'num_warps': 8}),
    triton.Config({'BLOCK_SIZE_N_expert': 64,  'BLOCK_SIZE_D_feature': 256, 'num_warps': 8}),
]
AUTOTUNE_KEY_SCATTER = ['num_expert_tokens', 'D_feature_dim', 'N_out_tokens']

# --- Kernels ---

@triton.autotune(
    configs=([DEBUG_CONFIG_MATMUL] if DEBUG_MODE else AUTOTUNE_CONFIGS_MATMUL),
    key=(['M', 'N', 'K'] if DEBUG_MODE else AUTOTUNE_KEY_MATMUL), # Simpler key for debug
)
@triton.jit
def matmul_kernel( # ... (kernel code remains the same as your last correct version)
    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_tile = tl.arange(0, BLOCK_SIZE_K) 
    
    a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k_tile[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k_tile[:, None] * stride_bk + offs_n[None, :] * stride_bn)
    
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k_loop_offset in range(0, K, BLOCK_SIZE_K):
        k_remaining_in_tensor = K - k_loop_offset
        mask_k_tile_active = offs_k_tile < k_remaining_in_tensor

        a_load_mask = (offs_m[:, None] < M) & mask_k_tile_active[None, :]
        b_load_mask = mask_k_tile_active[:, None] & (offs_n[None, :] < N)
        
        a = tl.load(a_ptrs, mask=a_load_mask, other=0.0)
        b = tl.load(b_ptrs, mask=b_load_mask, 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)
    
    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,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1)
    )

    if DEBUG_MODE and _KERNEL_DEBUG_CALL_COUNTS.get("triton_matmul", 0) < _MAX_DEBUG_CALLS_PER_KERNEL_TYPE:
        # PyTorch reference for matmul
        # Cast inputs to float32 for PyTorch matmul for potentially more stable accumulation, then cast back
        a_fp32 = a.to(torch.float32)
        b_fp32 = b.to(torch.float32)
        c_ref = torch.matmul(a_fp32, b_fp32).to(torch.float16)
        compare_outputs_for_debug("triton_matmul", c, c_ref, {'a': a, 'b': b})
    return c


@triton.autotune(
    configs=([DEBUG_CONFIG_MATMUL] if DEBUG_MODE else AUTOTUNE_CONFIGS_MATMUL),
    key=(['M', 'N', 'K'] if DEBUG_MODE else AUTOTUNE_KEY_MATMUL),
)
@triton.jit
def matmul_silu_kernel( # ... (kernel code remains the same as your last correct version)
    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_tile = tl.arange(0, BLOCK_SIZE_K) 
    
    a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k_tile[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k_tile[:, None] * stride_bk + offs_n[None, :] * stride_bn)
    
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k_loop_offset in range(0, K, BLOCK_SIZE_K):
        k_remaining_in_tensor = K - k_loop_offset
        mask_k_tile_active = offs_k_tile < k_remaining_in_tensor
        
        a_load_mask = (offs_m[:, None] < M) & mask_k_tile_active[None, :]
        b_load_mask = mask_k_tile_active[:, None] & (offs_n[None, :] < N)
        
        a = tl.load(a_ptrs, mask=a_load_mask, other=0.0)
        b = tl.load(b_ptrs, mask=b_load_mask, other=0.0)
        
        acc += tl.dot(a, b)
        
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk
    
    acc = acc * tl.sigmoid(acc) # Using tl.sigmoid for SiLU: x * sigmoid(x)
    
    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_silu(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_silu: A.shape[1]={K}, B.shape[0]={_K}")
    c = torch.empty((M, N), device=a.device, dtype=torch.float16)
    
    grid = lambda META: (
        triton.cdiv(M, META['BLOCK_SIZE_M']),
        triton.cdiv(N, META['BLOCK_SIZE_N'])
    )
    
    matmul_silu_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)
    )

    if DEBUG_MODE and _KERNEL_DEBUG_CALL_COUNTS.get("triton_matmul_silu", 0) < _MAX_DEBUG_CALLS_PER_KERNEL_TYPE:
        a_fp32 = a.to(torch.float32)
        b_fp32 = b.to(torch.float32)
        mm_result_fp32 = torch.matmul(a_fp32, b_fp32)
        c_ref = (mm_result_fp32 * torch.sigmoid(mm_result_fp32)).to(torch.float16)
        compare_outputs_for_debug("triton_matmul_silu", c, c_ref, {'a': a, 'b': b})
    return c

@triton.autotune(
    configs=([DEBUG_CONFIG_SCATTER] if DEBUG_MODE else AUTOTUNE_CONFIGS_SCATTER),
    key=(['num_expert_tokens', 'D_feature_dim'] if DEBUG_MODE else AUTOTUNE_KEY_SCATTER),
)
@triton.jit
def scatter_add_kernel( # ... (kernel code remains the same as your last correct version)
    output_ptr, input_ptr, indices_ptr, weights_ptr,
    N_out_tokens, D_feature_dim, num_expert_tokens,
    stride_out_n, stride_out_d,
    stride_in_n, stride_in_d,
    BLOCK_SIZE_N_expert: tl.constexpr,
    BLOCK_SIZE_D_feature: tl.constexpr
):
    pid_n_expert = tl.program_id(0) 
    pid_d_feature = tl.program_id(1) 
    
    offs_n_expert = pid_n_expert * BLOCK_SIZE_N_expert + tl.arange(0, BLOCK_SIZE_N_expert)
    offs_d_feature = pid_d_feature * BLOCK_SIZE_D_feature + tl.arange(0, BLOCK_SIZE_D_feature)
    
    mask_is_valid_expert_token_offset = offs_n_expert < num_expert_tokens
    mask_is_valid_feature_offset = offs_d_feature < D_feature_dim
    
    loaded_output_row_index = tl.load(indices_ptr + offs_n_expert, mask=mask_is_valid_expert_token_offset, other=0)
    weight_val = tl.load(weights_ptr + offs_n_expert, mask=mask_is_valid_expert_token_offset, other=0.0)
    
    mask_is_valid_output_row = (loaded_output_row_index < N_out_tokens) & (loaded_output_row_index >= 0)
    mask_can_process_expert_token = mask_is_valid_expert_token_offset & mask_is_valid_output_row
    final_op_mask = mask_can_process_expert_token[:, None] & mask_is_valid_feature_offset[None, :]
    
    current_input_ptrs = input_ptr + (offs_n_expert[:, None] * stride_in_n + offs_d_feature[None, :] * stride_in_d)
    input_val = tl.load(current_input_ptrs, mask=final_op_mask, other=0.0) 
    
    scaled_val = input_val * weight_val[:, None]
    
    current_output_ptrs = output_ptr + (loaded_output_row_index[:, None] * stride_out_n + offs_d_feature[None, :] * stride_out_d)
    tl.atomic_add(current_output_ptrs, scaled_val.to(tl.float16), mask=final_op_mask)


def triton_scatter_add(output: torch.Tensor, input_tensor: torch.Tensor, 
                       indices: torch.Tensor, weights: torch.Tensor):
    num_expert_tokens_arg, D_feature_dim_arg = input_tensor.shape # Renamed to avoid conflict with kernel args
    N_out_tokens_arg = output.shape[0]

    # For debugging, clone output before modification
    output_before_triton_call = None
    if DEBUG_MODE and _KERNEL_DEBUG_CALL_COUNTS.get("triton_scatter_add", 0) < _MAX_DEBUG_CALLS_PER_KERNEL_TYPE:
        output_before_triton_call = output.clone()

    grid = lambda META: (
        triton.cdiv(num_expert_tokens_arg, META['BLOCK_SIZE_N_expert']),
        triton.cdiv(D_feature_dim_arg, META['BLOCK_SIZE_D_feature'])
    )
    
    scatter_add_kernel[grid](
        output, input_tensor, indices, weights,
        N_out_tokens_arg, D_feature_dim_arg, num_expert_tokens_arg,
        output.stride(0), output.stride(1),
        input_tensor.stride(0), input_tensor.stride(1)
    )

    if DEBUG_MODE and output_before_triton_call is not None: # Check if we should compare
        # PyTorch reference for scatter_add (index_add_ with prior weighting)
        # Ensure indices are int64 for index_add_
        indices_long = indices.to(torch.int64) 
        
        # Weight the input tensor (expert_out_unweighted in the main logic)
        # input_tensor: (num_tokens_for_expert, d_feature)
        # weights: (num_tokens_for_expert) -> unsqueeze to (num_tokens_for_expert, 1)
        weighted_input = input_tensor.to(torch.float32) * weights.to(torch.float32).unsqueeze(-1)
        
        # output_ref should be based on the state *before* the Triton call
        output_ref = output_before_triton_call.clone().to(torch.float32)
        # Perform index_add_ which sums contributions
        output_ref.index_add_(0, indices_long, weighted_input)
        output_ref = output_ref.to(torch.float16) # Cast back to original dtype for comparison

        compare_outputs_for_debug("triton_scatter_add", output, output_ref, 
                                  {'output_before': output_before_triton_call, 
                                   'input_tensor': input_tensor, 'indices': indices, 'weights': weights})


# --- custom_kernel ---
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"]
    n_experts_per_token = config_dict["n_experts_per_token"]
    batch_size = config_dict["batch_size"]
    seq_len = config_dict["seq_len"]
    num_experts_to_loop = n_routed_experts 
    
    x_input = input_tensor_arg.view(-1, d_hidden) 
    
    if DEBUG_MODE:
        print()
        print(f"[Debug custom_kernel] Initial x_input: shape={x_input.shape}, dtype={x_input.dtype}")
        if x_input.numel() > 0 : print(f"  x_input[0,:5]: {x_input[0,:min(5,x_input.shape[1])]}")

    # Router computation
    W_g = weights_dict['router.weight'].t().contiguous() # Ensure contiguous
    if DEBUG_MODE: print(f"[Debug custom_kernel] W_g: shape={W_g.shape}, dtype={W_g.dtype}")
    
    logits = triton_matmul(x_input, W_g)
    if DEBUG_MODE: print(f"[Debug custom_kernel] Router logits: shape={logits.shape}, dtype={logits.dtype}")

    scores = F.softmax(logits.float(), dim=-1).to(logits.dtype) # softmax in fp32 for stability
    topk_scores, topk_indices = torch.topk(scores, k=n_experts_per_token, dim=-1)
    topk_scores_fp16 = topk_scores.to(torch.float16) # Already float16 if logits is fp16
    
    flat_expert_indices = topk_indices.view(-1) 
    flat_expert_weights = topk_scores_fp16.view(-1)
    
    idxs_argsort = flat_expert_indices.argsort()
    token_idxs_dispatch = idxs_argsort // n_experts_per_token 
    
    num_router_experts = W_g.shape[1]
    counts = flat_expert_indices.bincount(minlength=num_router_experts) 
    tokens_per_expert_cumsum = counts.cumsum(dim=0)
    
    # Shared expert
    W_gate_shared = weights_dict['shared_experts.0.weight'].contiguous()
    W_up_shared = weights_dict['shared_experts.1.weight'].contiguous()
    W_down_shared = weights_dict['shared_experts.2.weight'].contiguous()
        
    gate_shared = triton_matmul_silu(x_input, W_gate_shared)
    up_shared = triton_matmul(x_input, W_up_shared)
    intermediate_shared = gate_shared * up_shared 
    output_combined = triton_matmul(intermediate_shared, W_down_shared)
    
    if DEBUG_MODE:
        print(f"[Debug custom_kernel] After shared expert, output_combined: shape={output_combined.shape}")
        if output_combined.numel() > 0: print(f"  output_combined[0,:5]: {output_combined[0,:min(5,output_combined.shape[1])]}")

    # Expert processing
    for expert_id_loop in range(num_experts_to_loop):
        start_idx = 0 if expert_id_loop == 0 else tokens_per_expert_cumsum[expert_id_loop - 1].item()
        end_idx = tokens_per_expert_cumsum[expert_id_loop].item()
        
        if DEBUG_MODE and (start_idx != end_idx or expert_id_loop < 2) : # Print for active experts or first few
             print(f"[Debug custom_kernel] Expert Loop {expert_id_loop}: start_idx={start_idx}, end_idx={end_idx}")

        if start_idx == end_idx:
            continue
        
        current_expert_token_dispatch_indices = idxs_argsort[start_idx:end_idx]
        exp_token_idxs = token_idxs_dispatch[start_idx:end_idx] 
        exp_weights = flat_expert_weights[current_expert_token_dispatch_indices]
        expert_tokens = x_input[exp_token_idxs].contiguous() 
        
        if DEBUG_MODE and expert_tokens.numel() > 0:
            print(f"  Expert {expert_id_loop} processing {expert_tokens.shape[0]} tokens.")
            # print(f"    exp_token_idxs[:5]: {exp_token_idxs[:5]}")
            # print(f"    exp_weights[:5]: {exp_weights[:5]}")
            # print(f"    expert_tokens[0,:5]: {expert_tokens[0,:min(5,expert_tokens.shape[1])]}")


        W_gate = weights_dict[f'experts.{expert_id_loop}.0.weight'].contiguous()
        W_up = weights_dict[f'experts.{expert_id_loop}.1.weight'].contiguous()
        W_down = weights_dict[f'experts.{expert_id_loop}.2.weight'].contiguous()
        
        gate = triton_matmul_silu(expert_tokens, W_gate)
        up = triton_matmul(expert_tokens, W_up)
        intermediate = gate * up 
        expert_out_unweighted = triton_matmul(intermediate, W_down)
        
        if DEBUG_MODE and expert_out_unweighted.numel() > 0:
            print(f"  Expert {expert_id_loop} expert_out_unweighted[0,:5]: {expert_out_unweighted[0,:min(5,expert_out_unweighted.shape[1])]}")

        triton_scatter_add(output_combined, expert_out_unweighted, 
                           exp_token_idxs.to(output_combined.device, dtype=torch.int64), # Ensure int64 for indices
                           exp_weights)
            
    if DEBUG_MODE:
        print(f"[Debug custom_kernel] Final output_combined before reshape: shape={output_combined.shape}")
        if output_combined.numel() > 0: print(f"  output_combined[0,:5]: {output_combined[0,:min(5,output_combined.shape[1])]}")

    # Shape adjustment
    if output_combined.shape[-1] != d_hidden:
        # This block should ideally not be hit if d_expert and d_hidden are configured correctly for the model
        if DEBUG_MODE: print(f"[Debug custom_kernel] WARNING: Output shape adjustment needed. output_combined.shape[-1]={output_combined.shape[-1]}, d_hidden={d_hidden}")
        if output_combined.shape[-1] < d_hidden:
            padding_amount = d_hidden - output_combined.shape[-1]
            output_combined = F.pad(output_combined, (0, padding_amount), "constant", 0.0)
        else:
            output_combined = output_combined[..., :d_hidden]
            
    return output_combined.view(batch_size, seq_len, d_hidden)
python eval.py test /tmp/tmp3fau08fl
============================================================
!!!!!!!!!!!!!!!!! DEBUG MODE IS ENABLED !!!!!!!!!!!!!!!!!
============================================================

[Debug custom_kernel] Initial x_input: shape=torch.Size([512, 7168]), dtype=torch.float16
  x_input[0,:5]: tensor([-0.5322, -1.1123, -0.2883, -0.0271,  0.3738], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] W_g: shape=torch.Size([7168, 4]), dtype=torch.float16

--- Debugging Kernel: triton_matmul (Call #1) ---
Outputs for triton_matmul MATCH PyTorch reference.
[Debug custom_kernel] Router logits: shape=torch.Size([512, 4]), dtype=torch.float16

--- Debugging Kernel: triton_matmul_silu (Call #1) ---
Outputs for triton_matmul_silu MATCH PyTorch reference.
[Debug custom_kernel] After shared expert, output_combined: shape=torch.Size([512, 7168])
  output_combined[0,:5]: tensor([-0.4399, -0.0402,  1.3936, -0.5723,  2.2578], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 0: start_idx=0, end_idx=512
  Expert 0 processing 512 tokens.
  Expert 0 expert_out_unweighted[0,:5]: tensor([ 2.6367, -1.5596,  0.9580, -1.4619,  0.3562], device='cuda:0',
       dtype=torch.float16)

--- Debugging Kernel: triton_scatter_add (Call #1) ---
Outputs for triton_scatter_add MATCH PyTorch reference.
[Debug custom_kernel] Expert Loop 1: start_idx=512, end_idx=1024
  Expert 1 processing 512 tokens.
  Expert 1 expert_out_unweighted[0,:5]: tensor([-0.6128,  1.4990,  0.1814, -0.9136,  0.6226], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 2: start_idx=1024, end_idx=1536
  Expert 2 processing 512 tokens.
  Expert 2 expert_out_unweighted[0,:5]: tensor([-0.2098, -1.7939, -0.4055, -1.4336,  0.9663], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 3: start_idx=1536, end_idx=2048
  Expert 3 processing 512 tokens.
  Expert 3 expert_out_unweighted[0,:5]: tensor([-1.9092,  1.4287, -1.6963, -1.8545, -0.2158], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Final output_combined before reshape: shape=torch.Size([512, 7168])
  output_combined[0,:5]: tensor([-1.2383,  0.0635,  0.6353, -2.0977,  2.6348], device='cuda:0',
       dtype=torch.float16)

[Debug custom_kernel] Initial x_input: shape=torch.Size([1024, 7168]), dtype=torch.float16
  x_input[0,:5]: tensor([-1.5762, -0.1316,  0.3083,  2.3398,  0.1781], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] W_g: shape=torch.Size([7168, 8]), dtype=torch.float16
[Debug custom_kernel] Router logits: shape=torch.Size([1024, 8]), dtype=torch.float16
[Debug custom_kernel] After shared expert, output_combined: shape=torch.Size([1024, 7168])
  output_combined[0,:5]: tensor([ 0.0171,  0.5586, -2.7715, -0.9951,  0.0585], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 0: start_idx=0, end_idx=536
  Expert 0 processing 536 tokens.
  Expert 0 expert_out_unweighted[0,:5]: tensor([ 0.3064, -0.6240,  2.2051,  0.0976,  1.3076], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 1: start_idx=536, end_idx=1013
  Expert 1 processing 477 tokens.
  Expert 1 expert_out_unweighted[0,:5]: tensor([-1.1797, -0.0807, -2.1543,  1.2080, -0.4177], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 2: start_idx=1013, end_idx=1527
  Expert 2 processing 514 tokens.
  Expert 2 expert_out_unweighted[0,:5]: tensor([-1.9658,  1.0215,  0.6025, -1.6211,  0.5625], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 3: start_idx=1527, end_idx=2022
  Expert 3 processing 495 tokens.
  Expert 3 expert_out_unweighted[0,:5]: tensor([-0.3965, -0.1456, -0.7554, -1.6807,  1.6650], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 4: start_idx=2022, end_idx=2537
  Expert 4 processing 515 tokens.
  Expert 4 expert_out_unweighted[0,:5]: tensor([ 2.5293, -0.0104,  0.0611, -0.2130,  0.5215], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 5: start_idx=2537, end_idx=3050
  Expert 5 processing 513 tokens.
  Expert 5 expert_out_unweighted[0,:5]: tensor([ 0.2822,  1.0000,  1.5918, -0.3838, -0.7905], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 6: start_idx=3050, end_idx=3596
  Expert 6 processing 546 tokens.
  Expert 6 expert_out_unweighted[0,:5]: tensor([-1.4756, -2.3164,  0.2109, -2.6309,  0.0307], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 7: start_idx=3596, end_idx=4096
  Expert 7 processing 500 tokens.
  Expert 7 expert_out_unweighted[0,:5]: tensor([ 0.7900, -0.4412, -0.6152,  0.8677,  0.2788], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Final output_combined before reshape: shape=torch.Size([1024, 7168])
  output_combined[0,:5]: tensor([ 1.4414,  0.7295, -2.4238, -1.3438,  0.4404], device='cuda:0',
       dtype=torch.float16)

[Debug custom_kernel] Initial x_input: shape=torch.Size([8192, 7168]), dtype=torch.float16
  x_input[0,:5]: tensor([ 0.1689, -1.8906, -0.9482,  0.1375, -0.7471], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] W_g: shape=torch.Size([7168, 8]), dtype=torch.float16
[Debug custom_kernel] Router logits: shape=torch.Size([8192, 8]), dtype=torch.float16
[Debug custom_kernel] After shared expert, output_combined: shape=torch.Size([8192, 7168])
  output_combined[0,:5]: tensor([0.1072, 0.7661, 1.3818, 0.9028, 0.6992], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 0: start_idx=0, end_idx=4024
  Expert 0 processing 4024 tokens.
  Expert 0 expert_out_unweighted[0,:5]: tensor([-0.1057,  1.0771, -0.5527, -0.5151, -0.6680], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 1: start_idx=4024, end_idx=8155
  Expert 1 processing 4131 tokens.
  Expert 1 expert_out_unweighted[0,:5]: tensor([-0.5825,  2.7617, -1.6221,  1.2109, -0.9062], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 2: start_idx=8155, end_idx=12286
  Expert 2 processing 4131 tokens.
  Expert 2 expert_out_unweighted[0,:5]: tensor([-0.1826, -0.9888, -0.7031,  0.2927, -0.0348], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 3: start_idx=12286, end_idx=16301
  Expert 3 processing 4015 tokens.
  Expert 3 expert_out_unweighted[0,:5]: tensor([ 1.1719,  0.5083, -1.7129,  0.1741,  2.4160], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 4: start_idx=16301, end_idx=20365
  Expert 4 processing 4064 tokens.
  Expert 4 expert_out_unweighted[0,:5]: tensor([-0.1887, -0.3384,  0.2007, -0.4990,  0.3066], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 5: start_idx=20365, end_idx=24472
  Expert 5 processing 4107 tokens.
  Expert 5 expert_out_unweighted[0,:5]: tensor([ 0.9448, -0.8672,  1.9727, -0.8643,  0.0194], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 6: start_idx=24472, end_idx=28607
  Expert 6 processing 4135 tokens.
  Expert 6 expert_out_unweighted[0,:5]: tensor([ 0.9502, -1.6348, -0.7529, -0.1379, -0.3169], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Expert Loop 7: start_idx=28607, end_idx=32768
  Expert 7 processing 4161 tokens.
  Expert 7 expert_out_unweighted[0,:5]: tensor([ 0.6460,  0.2427,  1.8154, -0.1492, -0.0764], device='cuda:0',
       dtype=torch.float16)
[Debug custom_kernel] Final output_combined before reshape: shape=torch.Size([8192, 7168])
  output_combined[0,:5]: tensor([0.5815, 0.6299, 1.5312, 0.9185, 0.4392], device='cuda:0',
       dtype=torch.float16)
