#!POPCORN leaderboard amd-mixture-of-experts
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple
import sys # For printing to stderr

import triton
import triton.language as tl

from task import input_t, output_t # Assuming task.py defines these types

_printed_moe_ffn_config_summaries = {}
_printed_moe_gate_config_summaries = {}

@triton.jit
def silu(x): return (x * tl.sigmoid(x.to(tl.float32))).to(tl.float16)

@triton.autotune(
    configs=[
        # num_stages = 1
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 128, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 4}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 128, 'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 4}, num_warps=8, num_stages=1),

        # num_stages = 2
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 128, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 64,  'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 4}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 128, 'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 4}, num_warps=8, num_stages=2),
        # Larger block sizes with more warps/stages
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N_OUT': 128, 'BLOCK_SIZE_K_IN': 64, 'BLOCK_SIZE_K_INTER': 64, 'GROUP_SIZE_M': 4}, num_warps=8, num_stages=2),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 64, 'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 4}, num_warps=8, num_stages=2),
    ],
    key=['M', 'N_out', 'K_in', 'K_inter'],
)
@triton.jit
def fused_ffn_kernel_v2(
    # Pointers to Matrices
    x_ptr, w_gate_ptr, w_up_ptr, w_down_ptr, output_ptr,
    # Matrix dimensions
    M, N_out, K_in, K_inter,
    # Strides
    stride_x_m, stride_x_k,
    stride_w_gate_k, stride_w_gate_inter,
    stride_w_up_k, stride_w_up_inter,
    stride_w_down_inter, stride_w_down_n,
    stride_out_m, stride_out_n,
    # Meta-parameters from autotuner
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N_OUT: tl.constexpr,
    BLOCK_SIZE_K_IN: tl.constexpr, BLOCK_SIZE_K_INTER: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    # num_warps and num_stages are passed by autotuner to kernel launch
):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N_out, BLOCK_SIZE_N_OUT)
    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 % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n_out = pid_n * BLOCK_SIZE_N_OUT + tl.arange(0, BLOCK_SIZE_N_OUT)
    x_start_ptr = x_ptr + offs_m[:, None] * stride_x_m

    weight_dtype = w_gate_ptr.dtype.element_ty
    accum_dtype = tl.float32

    acc_down = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N_OUT), dtype=accum_dtype)

    for k_inter_start in range(0, tl.cdiv(K_inter, BLOCK_SIZE_K_INTER)):
        offs_k_inter = k_inter_start * BLOCK_SIZE_K_INTER + tl.arange(0, BLOCK_SIZE_K_INTER)
        acc_gate_fp32 = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K_INTER), dtype=accum_dtype)
        acc_up_fp32 = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K_INTER), dtype=accum_dtype)

        w_gate_k_inter_ptr = w_gate_ptr + offs_k_inter[None, :] * stride_w_gate_inter
        w_up_k_inter_ptr = w_up_ptr + offs_k_inter[None, :] * stride_w_up_inter

        for k_in_start in range(0, tl.cdiv(K_in, BLOCK_SIZE_K_IN)):
            offs_k_in = k_in_start * BLOCK_SIZE_K_IN + tl.arange(0, BLOCK_SIZE_K_IN)

            x_ptrs = x_start_ptr + offs_k_in[None, :] * stride_x_k
            x_mask = (offs_m[:, None] < M) & (offs_k_in[None, :] < K_in)
            x_block = tl.load(x_ptrs, mask=x_mask, other=0.0) # Added other=0.0

            w_gate_ptrs = w_gate_k_inter_ptr + offs_k_in[:, None] * stride_w_gate_k
            w_gate_mask = (offs_k_in[:, None] < K_in) & (offs_k_inter[None, :] < K_inter)
            w_gate_block = tl.load(w_gate_ptrs, mask=w_gate_mask, other=0.0) # Added other=0.0
            acc_gate_fp32 += tl.dot(x_block, w_gate_block, out_dtype=accum_dtype)

            w_up_ptrs = w_up_k_inter_ptr + offs_k_in[:, None] * stride_w_up_k
            w_up_mask = (offs_k_in[:, None] < K_in) & (offs_k_inter[None, :] < K_inter)
            w_up_block = tl.load(w_up_ptrs, mask=w_up_mask, other=0.0) # Added other=0.0
            acc_up_fp32 += tl.dot(x_block, w_up_block, out_dtype=accum_dtype)

        acc_gate_fp16 = acc_gate_fp32.to(weight_dtype)
        acc_up_fp16 = acc_up_fp32.to(weight_dtype)
        gate_activated_fp16 = silu(acc_gate_fp16)
        intermediate_block_fp16 = gate_activated_fp16 * acc_up_fp16

        w_down_start_ptr = w_down_ptr + offs_n_out[None, :] * stride_w_down_n
        w_down_ptrs = w_down_start_ptr + offs_k_inter[:, None] * stride_w_down_inter
        w_down_mask = (offs_k_inter[:, None] < K_inter) & (offs_n_out[None, :] < N_out)
        w_down_block = tl.load(w_down_ptrs, mask=w_down_mask, other=0.0) # Added other=0.0
        acc_down += tl.dot(intermediate_block_fp16, w_down_block, out_dtype=accum_dtype)

    output_ptrs = output_ptr + offs_m[:, None] * stride_out_m + offs_n_out[None, :] * stride_out_n
    output_mask = (offs_m[:, None] < M) & (offs_n_out[None, :] < N_out)
    tl.store(output_ptrs, acc_down.to(output_ptr.dtype.element_ty), mask=output_mask)

class TritonExpert(nn.Module):
    def __init__(self, config: Dict, W_gate: torch.Tensor, W_up: torch.Tensor, W_down: torch.Tensor, d_expert_actual: int):
        super().__init__()
        self.config = config
        self.d_hidden: int = config["d_hidden"]
        self.d_expert_actual = d_expert_actual
        self.W_gate = W_gate
        self.W_up = W_up
        self.W_down = W_down

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        M, K_in = x.shape
        if M == 0: # Handle empty input for an expert
            return torch.empty((0, self.d_hidden), device=x.device, dtype=x.dtype)

        K_inter = self.d_expert_actual
        N_out = self.d_hidden
        x = x.contiguous() # Ensure x is contiguous
        output = torch.empty((M, N_out), device=x.device, dtype=x.dtype)

        grid = lambda META: (
            triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N_out, META['BLOCK_SIZE_N_OUT']),
        )
        
        # Kernel launch - num_warps and num_stages come from autotuner config
        fused_ffn_kernel_v2[grid](
            x, self.W_gate, self.W_up, self.W_down, output,
            M, N_out, K_in, K_inter,
            x.stride(0), x.stride(1),
            self.W_gate.stride(0), self.W_gate.stride(1),
            self.W_up.stride(0), self.W_up.stride(1),
            self.W_down.stride(0), self.W_down.stride(1),
            output.stride(0), output.stride(1)
            # BLOCK_SIZES, GROUP_SIZE_M, num_warps, num_stages are set by autotuner
        )

        # Config printing
        config_key = (M, N_out, K_in, K_inter)
        global _printed_moe_ffn_config_summaries
        if config_key not in _printed_moe_ffn_config_summaries and hasattr(fused_ffn_kernel_v2, 'best_config'):
            best_config = fused_ffn_kernel_v2.best_config
            if best_config:
                try:
                    cfg = best_config.kwargs
                    m_bs = cfg.get('BLOCK_SIZE_M', 'NA')
                    n_bs = cfg.get('BLOCK_SIZE_N_OUT', 'NA')
                    k_in_bs = cfg.get('BLOCK_SIZE_K_IN', 'NA')
                    k_inter_bs = cfg.get('BLOCK_SIZE_K_INTER', 'NA')
                    gm = cfg.get('GROUP_SIZE_M', 'NA')
                    warps = best_config.num_warps if hasattr(best_config, 'num_warps') else 'NA'
                    stages = best_config.num_stages if hasattr(best_config, 'num_stages') else 'NA'
                    
                    summary = f"FFNCo@{M}x{N_out}x{K_in}x{K_inter}: {m_bs}-{n_bs}-{k_in_bs}-{k_inter_bs}-{gm}w{warps}s{stages}"
                    print(summary, file=sys.stderr)
                    _printed_moe_ffn_config_summaries[config_key] = True
                except Exception as e:
                    print(f"Error summarizing FFN config for {config_key}: {e}", file=sys.stderr)
                    _printed_moe_ffn_config_summaries[config_key] = True
        return output

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_K': 64}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_K': 128}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_K': 256}, num_warps=4, num_stages=1),
        triton.Config({'BLOCK_K': 64}, num_warps=8, num_stages=1),
        triton.Config({'BLOCK_K': 128}, num_warps=8, num_stages=1),
        triton.Config({'BLOCK_K': 256}, num_warps=8, num_stages=1),
        triton.Config({'BLOCK_K': 64}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_K': 128}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_K': 256}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_K': 64}, num_warps=8, num_stages=2),
        triton.Config({'BLOCK_K': 128}, num_warps=8, num_stages=2),
        triton.Config({'BLOCK_K': 256}, num_warps=8, num_stages=2),
    ],
    key=['M', 'K', 'N', 'top_k'], # N and top_k are tl.constexpr, but good to include if they could change behavior
)
@triton.jit
def _triton_gate_kernel(
    X_ptr, Wg_ptr, Indices_ptr, Scores_ptr,
    M, N: tl.constexpr, K,
    stride_xm, stride_xk, stride_wn, stride_wk,
    stride_indices_m, stride_indices_k, stride_scores_m, stride_scores_k,
    top_k: tl.constexpr, BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    x_row_ptr = X_ptr + pid_m * stride_xm
    offs_k_loop = tl.arange(0, BLOCK_K)
    offs_n_all = tl.arange(0, N)
    accumulator = tl.zeros((N,), dtype=tl.float32)

    for k_start in range(0, tl.cdiv(K, BLOCK_K)):
        current_offs_k = k_start * BLOCK_K + offs_k_loop
        x_mask = (pid_m < M) & (current_offs_k < K)
        # Load x_chunk as 1D for element-wise product
        x_chunk_1d = tl.load(x_row_ptr + current_offs_k * stride_xk, mask=x_mask, other=0.0) # Shape [BLOCK_K]
        
        # Load Wg_chunk for all N experts and current K block
        wg_block_ptrs = Wg_ptr + offs_n_all[:, None] * stride_wn + current_offs_k[None, :] * stride_wk
        wg_block_mask = (offs_n_all[:, None] < N) & (current_offs_k[None, :] < K)
        wg_chunk_block = tl.load(wg_block_ptrs, mask=wg_block_mask, other=0.0) # Shape [N, BLOCK_K]
        
        # Accumulate: sum(wg_chunk_block[i, k_val] * x_chunk_1d[k_val] for k_val in BLOCK_K) for each i
        accumulator += tl.sum(wg_chunk_block * x_chunk_1d[None, :], axis=1)


    if pid_m < M:
        current_scores_fp16 = tl.softmax(accumulator.to(tl.float16)).to(X_ptr.dtype.element_ty)
        scores_row_ptr = Scores_ptr + pid_m * stride_scores_m
        indices_row_ptr = Indices_ptr + pid_m * stride_indices_m
        loop_scores = current_scores_fp16
        neg_inf_fp16 = -float('inf')
        for k_iter in tl.static_range(top_k):
            max_val_fp16, max_idx = tl.max(loop_scores, axis=0, return_indices=True)
            tl.store(indices_row_ptr + k_iter * stride_indices_k, max_idx.to(tl.int64))
            tl.store(scores_row_ptr + k_iter * stride_scores_k, max_val_fp16)
            loop_scores = tl.where(tl.arange(0, N) == max_idx, neg_inf_fp16, loop_scores)

class MoEGate(nn.Module):
    def __init__(self, config: Dict, W_g_weight: torch.Tensor):
        super().__init__()
        self.top_k: int = config["n_experts_per_token"]
        self.n_experts: int = config["n_routed_experts"]
        self.register_buffer('W_g', W_g_weight)

    @torch.compile
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = F.linear(x, self.W_g)
        scores = logits.softmax(dim=-1)
        topk_scores, topk_indices = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
        return topk_indices, topk_scores

    def forward_triton(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        x = x.contiguous()
        M, K = x.shape
        N = self.n_experts
        topk_indices = torch.empty((M, self.top_k), device=x.device, dtype=torch.int64)
        topk_scores = torch.empty((M, self.top_k), device=x.device, dtype=x.dtype)
        grid = (M,)
        
        _triton_gate_kernel[grid](
            x, self.W_g, topk_indices, topk_scores,
            M, N, K,
            x.stride(0), x.stride(1),
            self.W_g.stride(0), self.W_g.stride(1),
            topk_indices.stride(0), topk_indices.stride(1),
            topk_scores.stride(0), topk_scores.stride(1),
            top_k=self.top_k # BLOCK_K, num_warps, num_stages from autotuner
        )
        
        config_key = (M, K, N, self.top_k) # Using N and top_k from self
        global _printed_moe_gate_config_summaries
        if config_key not in _printed_moe_gate_config_summaries and hasattr(_triton_gate_kernel, 'best_config'):
            best_config = _triton_gate_kernel.best_config
            if best_config:
                try:
                    cfg = best_config.kwargs
                    bk = cfg.get('BLOCK_K', 'NA')
                    warps = best_config.num_warps if hasattr(best_config, 'num_warps') else 'NA'
                    stages = best_config.num_stages if hasattr(best_config, 'num_stages') else 'NA'
                    summary = f"Gate@{M}x{K}(N{N},tk{self.top_k}): BK{bk}-w{warps}s{stages}"
                    print(summary, file=sys.stderr)
                    _printed_moe_gate_config_summaries[config_key] = True
                except Exception as e:
                    print(f"Error summarizing Gate config for {config_key}: {e}", file=sys.stderr)
                    _printed_moe_gate_config_summaries[config_key] = True
        return topk_indices, topk_scores

class TritonMoE(nn.Module):
    def __init__(self, config: Dict, weights: Dict[str, torch.Tensor]):
        super().__init__()
        self.config = config
        self.num_experts = config["n_routed_experts"]
        self.top_k = config["n_experts_per_token"]
        self.d_hidden = config["d_hidden"]
        self.d_expert = config["d_expert"]
        self.gating_network = MoEGate(config, weights['router.weight'])
        self.experts = nn.ModuleList()
        for i in range(self.num_experts):
            w_gate = weights[f'experts.{i}.0.weight']
            w_up = weights[f'experts.{i}.1.weight']
            w_down = weights[f'experts.{i}.2.weight']
            self.experts.append(TritonExpert(config, w_gate, w_up, w_down, self.d_expert))
        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        w_gate_shared = weights['shared_experts.0.weight']
        w_up_shared = weights['shared_experts.1.weight']
        w_down_shared = weights['shared_experts.2.weight']
        self.shared_expert = TritonExpert(config, w_gate_shared, w_up_shared, w_down_shared, shared_expert_dim)

    @torch.compile
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.view(-1, self.d_hidden)
        num_tokens = x_flat.shape[0]

        expert_indices, expert_scores = self.gating_network(x_flat)
        # expert_indices, expert_scores = self.gating_network.forward_triton(x_flat) # To use Triton gate

        flat_expert_indices = expert_indices.view(-1)
        flat_expert_scores = expert_scores.view(-1, 1)
        token_ids_for_flat_assignments = torch.arange(num_tokens, device=x.device).repeat_interleave(self.top_k)
        sort_perm = flat_expert_indices.argsort()
        perm_expert_indices = flat_expert_indices[sort_perm]
        perm_token_ids = token_ids_for_flat_assignments[sort_perm]
        perm_expert_scores = flat_expert_scores[sort_perm]
        perm_x_flat = x_flat[perm_token_ids]
        expert_boundaries = torch.searchsorted(perm_expert_indices, torch.arange(self.num_experts + 1, device=x.device))

        all_expert_outputs_list = []
        for expert_id in range(self.num_experts):
            # Use tensor indexing directly, avoid .item()
            start_idx = expert_boundaries[expert_id]
            end_idx = expert_boundaries[expert_id + 1]
            
            # Create slice on GPU. TritonExpert will handle M=0 if start_idx == end_idx
            current_expert_inputs = perm_x_flat[start_idx:end_idx]
            
            expert_module = self.experts[expert_id]
            current_expert_outputs = expert_module(current_expert_inputs)
            all_expert_outputs_list.append(current_expert_outputs)

        if all_expert_outputs_list: # Ensure list is not empty before cat
            all_expert_outputs_perm = torch.cat(all_expert_outputs_list, dim=0)
        else: # Handle case where no tokens are routed to any expert (e.g. num_tokens=0 or top_k=0)
            all_expert_outputs_perm = torch.empty((0, self.d_hidden), device=x.device, dtype=x.dtype)

        weighted_expert_outputs_perm = all_expert_outputs_perm * perm_expert_scores
        final_output_flat = self.shared_expert(x_flat)
        scatter_indices = perm_token_ids.unsqueeze(-1).expand_as(weighted_expert_outputs_perm)
        
        final_output_flat.scatter_reduce_(
            dim=0, index=scatter_indices, src=weighted_expert_outputs_perm,
            reduce="sum", include_self=True
        )
        return final_output_flat.view(orig_shape)

def custom_kernel(data: input_t) -> output_t:
    input_tensor, weights, config = data
    triton_moe_model = TritonMoE(config, weights)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    input_tensor = input_tensor.to(device)
    triton_moe_model.to(device)

    input_dtype = next(iter(weights.values())).dtype
    for k, v in weights.items():
        weights[k] = v.to(device=device, dtype=input_dtype).contiguous()
    input_tensor = input_tensor.to(input_dtype)

    output = triton_moe_model(input_tensor)
    return output


'''
from reference import generate_input, ref_kernel, generate_input
from utils import make_match_reference
args = {"dhidden": 7168, "dexpert": 2048, "nroutedexperts": 8, "nsharedexperts": 1, "nexpertspertoken": 4, "bs": 1, "seqlen": 8192, "seed": 1212}
inp = generate_input(**args)
out = custom_kernel(generate_input(**args))
ref = make_match_reference(ref_kernel)
print(ref(inp, out))
'''

python eval.py test /tmp/tmpk7km3god
/usr/local/lib/python3.10/dist-packages/torch/_inductor/lowering.py:7078: UserWarning: 
Online softmax is disabled on the fly since Inductor decides to
split the reduction. Cut an issue to PyTorch if this is an
important use case and you want to speed it up with online
softmax.

  warnings.warn(
W0527 18:03:31.610000 743 torch/_dynamo/convert_frame.py:992] [1/8] torch._dynamo hit config.recompile_limit (8)
W0527 18:03:31.610000 743 torch/_dynamo/convert_frame.py:992] [1/8]    function: 'forward' (/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py:129)
W0527 18:03:31.610000 743 torch/_dynamo/convert_frame.py:992] [1/8]    last reason: 1/7: tensor 'x' size mismatch at index 0. expected 546, actual 500
W0527 18:03:31.610000 743 torch/_dynamo/convert_frame.py:992] [1/8] To log all recompilation reasons, use TORCH_LOGS="recompiles".
W0527 18:03:31.610000 743 torch/_dynamo/convert_frame.py:992] [1/8] To diagnose recompilation issues, see https://pytorch.org/docs/main/torch.compiler_troubleshooting.html.
FFNCo@500x7168x7168x2048: 128-128-128-128-4w8s1
FFNCo@1024x7168x7168x2048: 256-128-64-64-4w8s2
FFNCo@4024x7168x7168x2048: 256-128-64-64-4w8s2
FFNCo@4131x7168x7168x2048: 128-128-128-128-4w8s1
FFNCo@4015x7168x7168x2048: 256-128-64-64-4w8s2
FFNCo@4064x7168x7168x2048: 256-128-64-64-4w8s2
FFNCo@4107x7168x7168x2048: 128-128-128-128-4w8s1
FFNCo@4135x7168x7168x2048: 128-128-128-128-4w8s1
FFNCo@4161x7168x7168x2048: 128-128-128-128-4w8s1
FFNCo@8192x7168x7168x2048: 256-128-64-64-4w8s2
