#!POPCORN leaderboard amd-mixture-of-experts
import torch
from torch import Tensor
import torch.nn as nn
from torch.nn import functional as F
from typing import Dict, Tuple, Optional
from task import input_t, output_t
import math
class Linear(nn.Module):
    __constants__ = ["in_features", "out_features"]
    in_features: int
    out_features: int
    weight: Tensor

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        device=None,
        dtype=None,
    ) -> None:
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(
            torch.empty((out_features, in_features), **factory_kwargs)
        )
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs))
        else:
            self.register_parameter("bias", None)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with
        # uniform(-1/sqrt(in_features), 1/sqrt(in_features)). For details, see
        # https://github.com/pytorch/pytorch/issues/57109
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, input: Tensor) -> Tensor:
        return linear(input, self.weight, self.bias)

    def extra_repr(self) -> str:
        return f"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}"


class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.config = config
        self.act_fn = nn.SiLU()
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = config["d_expert"] if d_expert is None else d_expert

        self.W_gate = Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_up = Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_down = Linear(self.d_expert, self.d_hidden, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = self.act_fn(self.W_gate(x))
        out = self.W_down(gate * self.W_up(x))
        return out


class MoEGate(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.top_k: int = config["n_experts_per_token"]
        self.num_experts: int = config["n_routed_experts"]
        self.d_hidden: int = config["d_hidden"]

        self.W_g = Linear(self.d_hidden, self.num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = self.W_g(x)
        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


class MoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([
            Expert(config)
            for _ in range(config["n_routed_experts"])
        ])
        self.gating_network = MoEGate(config)
        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert(config=config, d_expert=shared_expert_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shared_output = self.shared_expert(x)
        expert_indices, expert_scores = self.gating_network(x)
        batch_size, seq_len, hidden_dim = x.shape
        orig_shape = x.shape
        x_flat = x.view(-1, hidden_dim)
        flat_expert_indices = expert_indices.view(-1)
        flat_expert_weights = expert_scores.view(-1, 1)
        routed_output_flat = self.moe_infer(x_flat,
                                            flat_expert_indices,
                                            flat_expert_weights)

        routed_output = routed_output_flat.view(*orig_shape)
        return routed_output + shared_output

    @torch.no_grad()
    def moe_infer(self,
                  x: torch.Tensor,
                  flat_expert_indices: torch.Tensor,
                  flat_expert_weights: torch.Tensor
                 ) -> torch.Tensor:
        expert_cache = torch.zeros_like(x)
        idxs = flat_expert_indices.argsort()
        counts = flat_expert_indices.bincount().cpu().numpy()
        tokens_per_expert = counts.cumsum()
        num_per_tok = self.config["n_experts_per_token"]
        token_idxs = idxs // num_per_tok
        for expert_id, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if expert_id == 0 else tokens_per_expert[expert_id - 1]
            if start_idx == end_idx:
                continue

            expert = self.experts[expert_id]
            exp_token_idxs = token_idxs[start_idx:end_idx]
            expert_tokens = x[exp_token_idxs]
            expert_out    = expert(expert_tokens)
            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
            expert_cache.scatter_reduce_(
                0,
                exp_token_idxs.view(-1, 1).repeat(1, x.shape[-1]),
                expert_out,
                reduce='sum'
            )

        return expert_cache


import triton
import triton.language as tl
from triton import Config

DEVICE = triton.runtime.driver.active.get_active_torch_device()

fp8_gemm_configs = [
    Config({'BLOCK_SIZE_M': block_m, 'BLOCK_SIZE_N': block_n, 'BLOCK_SIZE_K': 128}, num_stages=num_stages, num_warps=8)
    for block_m in [16, 32, 64] for block_n in [32, 64, 128] for num_stages in [3, 4, 5, 6]
]

@triton.autotune(configs=fp8_gemm_configs, key=['N', 'K'])
@triton.jit
def fp8_gemm_kernel(a_ptr, b_ptr, c_ptr,
                    a_s_ptr, b_s_ptr,
                    M, N: tl.constexpr, K: tl.constexpr,
                    BLOCK_SIZE_M: tl.constexpr,
                    BLOCK_SIZE_N: tl.constexpr,
                    BLOCK_SIZE_K: tl.constexpr):

    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)
    k = tl.cdiv(K, BLOCK_SIZE_K)
    offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
    b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
    a_s_ptrs = a_s_ptr + offs_m * k
    b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_K) * k

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for i in range(k):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - i * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - i * BLOCK_SIZE_K, other=0.0)
        a_s = tl.load(a_s_ptrs)
        b_s = tl.load(b_s_ptrs)
        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K
        a_s_ptrs += 1
        b_s_ptrs += 1
    c = accumulator.to(c_ptr.dtype.element_ty)
    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)
    c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=mask)

def blockwise_quantize_a(a_fp16, block_size=128):
    # a_fp16: [m, k], float16
    m, k = a_fp16.shape
    num_blk = k // block_size

    a_f8 = torch.empty_like(a_fp16, dtype=torch.float8_e4m3fnuz)
    a_scale = torch.empty((m, num_blk), dtype=torch.float32, device=a_fp16.device)

    for bi in range(num_blk):
        blk = a_fp16[:, bi*block_size:(bi+1)*block_size]
        absmax = blk.abs().amax(dim=1, keepdim=True) # [m, 1]
        scale = absmax / 127
        scale = scale.clamp(min=1e-6)
        blk_norm = blk / scale

        blk_f8 = blk_norm.to(torch.float8_e4m3fnuz)
        a_f8[:, bi*block_size:(bi+1)*block_size] = blk_f8
        a_scale[:, bi] = scale.squeeze(-1)
    return a_f8, a_scale

def blockwise_quantize_b(b_fp16, block_size=128):
    # b_fp16: [n, k], float16
    n, k = b_fp16.shape
    assert n % block_size == 0 and k % block_size == 0
    n_blk, k_blk = n // block_size, k // block_size
    b_f8 = torch.empty_like(b_fp16, dtype=torch.float8_e4m3fnuz)
    b_scale = torch.empty((n_blk, k_blk), dtype=torch.float32, device=b_fp16.device)

    for i in range(n_blk):
        for j in range(k_blk):
            blk = b_fp16[i*block_size:(i+1)*block_size, j*block_size:(j+1)*block_size]
            absmax = blk.abs().max()
            scale = absmax / 127
            scale = max(scale, 1e-6)
            blk_norm = blk / scale
            blk_f8 = blk_norm.to(torch.float8_e4m3fnuz)
            b_f8[i*block_size:(i+1)*block_size, j*block_size:(j+1)*block_size] = blk_f8
            b_scale[i, j] = scale
    return b_f8, b_scale

def linear(input: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
    """
    Reference implementation of block-scale fp8 gemm
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float8_e4m3fnuz] of shape [m, k],
            b: torch.Tensor[float8_e4m3fnuz] of shape [n, k],
            a_scale: torch.Tensor[float32] of shape [m, k // 128],
            b_scale: torch.Tensor[float32] of shape [n // 128, k // 128],
            c: torch.Tensor[bfloat16] of shape [m, n]
    Returns:
        Tensor containing output in bf16
    """
    print(f"linear: {input.shape=}, {weight.shape=}")
    # a, a_scale = blockwise_quantize_a(weight, block_size=128)
    # b, b_scale = blockwise_quantize_b(input, block_size=128)

    # a.contiguous()
    # b.contiguous()
    # a_scale.contiguous()
    # b_scale.contiguous()

    # K = a.size(-1)
    # M = a.numel() // K
    # N = b.size(0)
    # c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype())
    # grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']), triton.cdiv(N, META['BLOCK_SIZE_N']))
    # fp8_gemm_kernel[grid](a, b, c, a_scale, b_scale, M, N, K)

    c = F.linear(input, weight)
    print(f"{c.dtype=}")

    return c.to(torch.float16)

# def linear(input: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
#     print(f"linear: {input.dtype=}, {weight.dtype=}, bias.dtype={bias and bias.dtype} device={weight.device}")

#     return F.linear(input, weight, bias)

def custom_kernel(data: input_t) -> output_t:
    """
    Submission template for DeepSeek-style Mixture of Experts using PyTorch.

    Args:
        data: Tuple of (input: torch.Tensor, weights: Dict[str, torch.Tensor], config: Dict)
            - input: Input tensor of shape [batch_size, seq_len, hidden_size]
            - weights: Dictionary containing model weights
            - config: Dictionary containing model configuration parameters

    Returns:
        Tuple containing:
            - output: Processed tensor [batch_size, seq_len, d_model]
            - aux_data: Dictionary with auxiliary data
    """
    input_tensor, weights, config = data
    num_experts = config["n_routed_experts"]
    moe = MoE(config)

    # Fill in the given weights of the model
    moe.gating_network.W_g.weight = nn.Parameter(weights['router.weight'])

    for i in range(num_experts):
        gate_proj_weight = weights[f'experts.{i}.0.weight']
        up_proj_weight = weights[f'experts.{i}.1.weight']
        down_proj_weight = weights[f'experts.{i}.2.weight']

        # Transpose weights to match expected shape for Linear
        moe.experts[i].W_gate.weight = nn.Parameter(gate_proj_weight.t())
        moe.experts[i].W_up.weight = nn.Parameter(up_proj_weight.t())
        moe.experts[i].W_down.weight = nn.Parameter(down_proj_weight.t())

    moe.shared_expert.W_gate.weight = nn.Parameter(weights['shared_experts.0.weight'].t())
    moe.shared_expert.W_up.weight = nn.Parameter(weights['shared_experts.1.weight'].t())
    moe.shared_expert.W_down.weight = nn.Parameter(weights['shared_experts.2.weight'].t())

    # Run the model
    output = moe(input_tensor)

    return output
python eval.py benchmark /tmp/tmpbx6zqbhj
linear: input.shape=torch.Size([1, 2048, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([1, 2048, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([1, 2048, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([1, 2048, 7168]), weight.shape=torch.Size([32, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([240, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([240, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([240, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([268, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([268, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([268, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([265, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([265, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([265, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([251, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([251, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([251, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([239, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([239, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([239, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([244, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([244, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([244, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([230, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([230, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([230, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([278, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([278, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([278, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([261, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([261, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([261, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([236, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([236, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([236, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([270, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([270, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([270, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([259, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([259, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([259, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([290, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([290, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([290, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([243, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([243, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([243, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([275, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([275, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([275, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([256, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([256, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([256, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([249, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([249, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([249, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([269, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([269, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([269, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([231, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([231, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([231, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([250, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([250, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([250, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([218, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([218, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([218, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([292, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([292, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([292, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([266, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([266, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([266, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([245, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([245, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([245, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([255, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([255, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([255, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([258, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([258, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([258, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([232, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([232, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([232, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([1, 2048, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([1, 2048, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([1, 2048, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([1, 2048, 7168]), weight.shape=torch.Size([32, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([240, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([240, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([240, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([268, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([268, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([268, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([265, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([265, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([265, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([251, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([251, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([251, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([239, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([239, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([239, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([244, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([244, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([244, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([230, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([230, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([230, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([278, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([278, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([278, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([261, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([261, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([261, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([236, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([236, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([236, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([270, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([270, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([270, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([272, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([259, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([259, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([259, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([290, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([290, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([290, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([243, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([243, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([243, 2048]), weight.shape=torch.Size([7168, 2048])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 7168]), weight.shape=torch.Size([2048, 7168])
c.dtype=torch.float16
linear: input.shape=torch.Size([253, 2048]), weight.shape=torch.Size([7168, 2048])
[...] 2092 lines omitted