import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, Optional
import math

try:
    import triton
    import triton.language as tl
    TRITON_AVAILABLE = True
except ImportError:
    TRITON_AVAILABLE = False


if TRITON_AVAILABLE:
    @triton.autotune(
        configs=[
            triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32, 'SPLIT_K': 1, 'num_stages': 3, 'num_warps': 4}, ),
            triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32, 'BLOCK_K': 32, 'SPLIT_K': 1, 'num_stages': 3, 'num_warps': 4}, ),
            triton.Config({'BLOCK_M': 32, 'BLOCK_N': 128, 'BLOCK_K': 32, 'SPLIT_K': 1, 'num_stages': 4, 'num_warps': 4}, ),
            triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 32, 'SPLIT_K': 1, 'num_stages': 3, 'num_warps': 8}, ),
            triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32, 'SPLIT_K': 1, 'num_stages': 3, 'num_warps': 8}, ),
        ],
        key=['M', 'N', 'K'],
    )
    @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,
        ACC_TYPE: tl.constexpr,
        BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
        SPLIT_K: tl.constexpr,
    ):
        pid = tl.program_id(axis=0)
        num_pid_n = tl.cdiv(N, BLOCK_N)
        pid_m = pid // num_pid_n
        pid_n = pid % num_pid_n


        if SPLIT_K == 1:
            rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
            rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
            rk = tl.arange(0, BLOCK_K)
            A = A_ptr + rm[:, None] * stride_am + rk[None, :] * stride_ak
            B = B_ptr + rk[:, None] * stride_bk + rn[None, :] * stride_bn
            acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
            for k_loop in range(tl.cdiv(K, BLOCK_K)):
                a = tl.load(A + k_loop * BLOCK_K * stride_ak, mask=(rm[:, None] < M) & (rk[None, :] + k_loop * BLOCK_K < K), other=0.0)
                b = tl.load(B + k_loop * BLOCK_K * stride_bk, mask=(rk[:, None] + k_loop * BLOCK_K < K) & (rn[None, :] < N), other=0.0)
                acc += tl.dot(a, b)
            C = C_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn
            tl.store(C, acc, mask=(rm[:, None] < M) & (rn[None, :] < N))
        else: 
            pass


    def triton_matmul(a: torch.Tensor, b: torch.Tensor, acc_type=torch.float32) -> torch.Tensor:
        M, K = a.shape
        _K, N = b.shape
        assert K == _K, f"Incompatible dimensions for matmul: A ({M},{K}), B ({_K},{N})"
        c = torch.empty((M, N), device=a.device, dtype=a.dtype)

        grid = lambda META: (triton.cdiv(M, META['BLOCK_M']) * triton.cdiv(N, META['BLOCK_N']),)
        
        ACC_TYPE = tl.float32 if acc_type == torch.float32 else tl.float16

        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),
            ACC_TYPE=ACC_TYPE,
        )
        return c

    def expert_ffn_triton_pytorch_layer(x, w_gate, w_up, w_down):
        h_gate = triton_matmul(x, w_gate)
        h_up = triton_matmul(x, w_up)
        gated_h = F.silu(h_gate) * h_up
        output = triton_matmul(gated_h, w_down)
        return output

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 = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_up = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_down = nn.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))
        up_hidden = self.W_up(x)
        out = self.W_down(gate * up_hidden)
        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 = nn.Linear(self.d_hidden, self.num_experts, bias=False)

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

class MoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.num_routed_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.expert_W_gates = []
        self.expert_W_ups = []
        self.expert_W_downs = []
        self.use_triton_experts = TRITON_AVAILABLE

        self.gating_network = MoEGate(config)
        
        shared_expert_config = config.copy() 
        n_shared_experts = config.get("n_shared_experts", 1) 
        d_expert_shared = config["d_expert"] * n_shared_experts
        self.shared_expert = Expert(shared_expert_config, d_expert=d_expert_shared)
        
        self.aux_loss_weight = config.get("aux_loss_weight", 0.01)

        if not self.use_triton_experts:
            self.pytorch_experts = nn.ModuleList([
                Expert(config) for _ in range(self.num_routed_experts)
            ])

    def _expert_forward(self, x_tokens: torch.Tensor, expert_id: int) -> torch.Tensor:
        if self.use_triton_experts and x_tokens.numel() > 0 :
            w_gate = self.expert_W_gates[expert_id]
            w_up = self.expert_W_ups[expert_id]
            w_down = self.expert_W_downs[expert_id]
            return expert_ffn_triton_pytorch_layer(x_tokens, w_gate, w_up, w_down)
        elif hasattr(self, 'pytorch_experts'):
            return self.pytorch_experts[expert_id](x_tokens)
        else:
            if x_tokens.numel() == 0:
                return torch.empty_like(x_tokens) 
            raise RuntimeError("Erorr.")


    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_sorted_by_expert = flat_expert_indices.argsort()
        
        counts_per_expert = flat_expert_indices.bincount(minlength=self.num_routed_experts)
        tokens_per_expert_cumsum = counts_per_expert.cumsum(dim=0)
        padded_tokens_per_expert_cumsum = F.pad(tokens_per_expert_cumsum, (1, 0), "constant", 0)

        original_token_indices_for_assignments = idxs_sorted_by_expert // self.top_k

        for expert_id in range(self.num_routed_experts):
            start_idx_in_sorted_list = padded_tokens_per_expert_cumsum[expert_id]
            end_idx_in_sorted_list = padded_tokens_per_expert_cumsum[expert_id + 1]
            
            if start_idx_in_sorted_list == end_idx_in_sorted_list:
                continue
            
            current_expert_original_token_indices = original_token_indices_for_assignments[start_idx_in_sorted_list:end_idx_in_sorted_list]
            expert_tokens_input = x[current_expert_original_token_indices]
            
            if expert_tokens_input.numel() == 0: 
                continue

            expert_tokens_output = self._expert_forward(expert_tokens_input, expert_id)
            
            weights_for_current_assignments = flat_expert_weights[idxs_sorted_by_expert[start_idx_in_sorted_list:end_idx_in_sorted_list]]
            expert_tokens_output = expert_tokens_output * weights_for_current_assignments
            
            expert_cache.scatter_reduce_(
                dim=0, 
                index=current_expert_original_token_indices.view(-1, 1).expand_as(expert_tokens_output),
                src=expert_tokens_output,
                reduce='sum',
                include_self=True
            )
        return expert_cache

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Dict]:
        batch_size, seq_len, hidden_dim = x.shape
        orig_shape = x.shape

        shared_output = self.shared_expert(x)
        expert_indices, expert_scores, all_router_scores = self.gating_network(x)
        
        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)
        final_output = routed_output + shared_output
        aux_data = self.compute_aux_losses(all_router_scores, expert_indices)
        
        return final_output, aux_data

    def compute_aux_losses(self, router_probs_all_experts: torch.Tensor, 
                         topk_expert_indices: torch.Tensor) -> Dict:
        output_dtype = router_probs_all_experts.dtype
        topk_expert_indices_flat = topk_expert_indices.view(-1)
        counts = torch.bincount(topk_expert_indices_flat, minlength=self.num_routed_experts).to(device=router_probs_all_experts.device, dtype=torch.float32)
        total_assignments = topk_expert_indices.numel() 
        
        if total_assignments == 0:
            load_balance_loss = torch.tensor(0.0, device=counts.device, dtype=torch.float32)
        else:
            fraction_assignments_per_expert = counts / total_assignments
            load_balance_loss = self.num_routed_experts * torch.sum(fraction_assignments_per_expert.pow(2)) - 1.0 

        z_loss = (router_probs_all_experts.pow(2)).mean().to(torch.float32)
        total_aux_loss = self.aux_loss_weight * (load_balance_loss + z_loss)
        
        return {
            "router_probs": router_probs_all_experts,
            "load_balance_loss": load_balance_loss.to(output_dtype),
            "z_loss": z_loss.to(output_dtype),
            "total_aux_loss": total_aux_loss.to(output_dtype)
        }

def custom_kernel(data: Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict]) -> torch.Tensor:
    input_tensor, weights_dict, config = data
    n_routed_experts_config = config["n_routed_experts"] 
    
    input_tensor = input_tensor.to(dtype=torch.float16)
    device = input_tensor.device
    
    moe_module = MoE(config)
    
    moe_module.gating_network.W_g.weight = nn.Parameter(weights_dict['router.weight'].to(device, dtype=torch.float16))
    moe_module.shared_expert.W_gate.weight = nn.Parameter(weights_dict['shared_experts.0.weight'].t().to(device, dtype=torch.float16))
    moe_module.shared_expert.W_up.weight = nn.Parameter(weights_dict['shared_experts.1.weight'].t().to(device, dtype=torch.float16))
    moe_module.shared_expert.W_down.weight = nn.Parameter(weights_dict['shared_experts.2.weight'].t().to(device, dtype=torch.float16))

    if moe_module.use_triton_experts:
        for i in range(n_routed_experts_config):
            moe_module.expert_W_gates.append(weights_dict[f'experts.{i}.0.weight'].to(device, dtype=torch.float16))
            moe_module.expert_W_ups.append(weights_dict[f'experts.{i}.1.weight'].to(device, dtype=torch.float16))
            moe_module.expert_W_downs.append(weights_dict[f'experts.{i}.2.weight'].to(device, dtype=torch.float16))
    else:
        for i in range(n_routed_experts_config):
            gate_w_torch = weights_dict[f'experts.{i}.0.weight'].t().to(device, dtype=torch.float16)
            up_w_torch = weights_dict[f'experts.{i}.1.weight'].t().to(device, dtype=torch.float16)
            down_w_torch = weights_dict[f'experts.{i}.2.weight'].t().to(device, dtype=torch.float16)
            moe_module.pytorch_experts[i].W_gate.weight = nn.Parameter(gate_w_torch)
            moe_module.pytorch_experts[i].W_up.weight = nn.Parameter(up_w_torch)
            moe_module.pytorch_experts[i].W_down.weight = nn.Parameter(down_w_torch)

    moe_module = moe_module.to(device=device, dtype=torch.float16)
    
    compiled_moe_module = torch.compile(moe_module, mode="max-autotune")
    output_tensor, _ = compiled_moe_module(input_tensor) 
    
    return output_tensor
python eval.py test /tmp/tmp9yi7aqup
/usr/local/lib/python3.10/dist-packages/torch/_inductor/lowering.py:7061: 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(
AUTOTUNE mm(512x7168, 7168x2048)
strides: [7168, 1], [2048, 1]
dtypes: torch.float16, torch.float16
  mm 0.0660 ms 100.0% 
  triton_mm_40 0.0805 ms 82.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=32, BLOCK_N=64, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_47 0.0819 ms 80.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_48 0.0822 ms 80.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_43 0.0829 ms 79.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_45 0.0844 ms 78.2% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_49 0.0850 ms 77.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_55 0.0969 ms 68.2% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_57 0.0993 ms 66.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_52 0.1114 ms 59.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
SingleProcess AUTOTUNE benchmarking takes 0.7057 seconds and 0.0025 seconds precompiling for 37 choices
AUTOTUNE mm(512x2048, 2048x7168)
strides: [2048, 1], [7168, 1]
dtypes: torch.float16, torch.float16
  mm 0.0412 ms 100.0% 
  triton_mm_100 0.0520 ms 79.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_93 0.0534 ms 77.2% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_86 0.0547 ms 75.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_99 0.0554 ms 74.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_98 0.0565 ms 73.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_96 0.0568 ms 72.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_97 0.0571 ms 72.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=1, num_stages=2, num_warps=8
  triton_mm_95 0.0610 ms 67.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_103 0.0616 ms 66.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=256, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
SingleProcess AUTOTUNE benchmarking takes 4.2956 seconds and 0.0002 seconds precompiling for 37 choices
AUTOTUNE mm(512x7168, 7168x4)
strides: [7168, 1], [1, 7168]
dtypes: torch.float16, torch.float16
  triton_mm_111 0.0190 ms 100.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=32, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=2
  triton_mm_109 0.0191 ms 99.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=32, BLOCK_N=16, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=2
  triton_mm_112 0.0237 ms 80.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=32, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=2
  triton_mm_119 0.0271 ms 70.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_122 0.0272 ms 70.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_113 0.0273 ms 69.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_116 0.0273 ms 69.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_120 0.0279 ms 68.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  mm 0.0334 ms 57.1% 
  triton_mm_118 0.0334 ms 57.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
SingleProcess AUTOTUNE benchmarking takes 3.4991 seconds and 0.0002 seconds precompiling for 30 choices
AUTOTUNE mm(1024x7168, 7168x2048)
strides: [7168, 1], [2048, 1]
dtypes: torch.float16, torch.float16
  mm 0.0953 ms 100.0% 
  triton_mm_194 0.1151 ms 82.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_184 0.1193 ms 79.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_189 0.1216 ms 78.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_190 0.1226 ms 77.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_188 0.1287 ms 74.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_195 0.1295 ms 73.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_201 0.1314 ms 72.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_192 0.1345 ms 70.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_193 0.1396 ms 68.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
SingleProcess AUTOTUNE benchmarking takes 4.5947 seconds and 0.0003 seconds precompiling for 37 choices
AUTOTUNE mm(1024x2048, 2048x7168)
strides: [2048, 1], [7168, 1]
dtypes: torch.float16, torch.float16
  mm 0.0692 ms 100.0% 
  triton_mm_243 0.0858 ms 80.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=256, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_239 0.0864 ms 80.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=256, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_223 0.0928 ms 74.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_238 0.0928 ms 74.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=256, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_232 0.0930 ms 74.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_241 0.0938 ms 73.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=256, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_229 0.0947 ms 73.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_242 0.0960 ms 72.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=256, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_230 0.0977 ms 70.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
SingleProcess AUTOTUNE benchmarking takes 4.4238 seconds and 0.0003 seconds precompiling for 37 choices
AUTOTUNE mm(1024x7168, 7168x8)
strides: [7168, 1], [1, 7168]
dtypes: torch.float16, torch.float16
  triton_mm_248 0.0213 ms 100.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=32, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=2
  triton_mm_246 0.0222 ms 95.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=32, BLOCK_N=16, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=2
  triton_mm_249 0.0238 ms 89.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=32, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=2
  triton_mm_259 0.0272 ms 78.2% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_250 0.0273 ms 78.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_253 0.0273 ms 78.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_256 0.0275 ms 77.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_257 0.0279 ms 76.2% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_252 0.0335 ms 63.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_255 0.0338 ms 63.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
SingleProcess AUTOTUNE benchmarking takes 3.5036 seconds and 0.0003 seconds precompiling for 30 choices
AUTOTUNE mm(1024x7168, 7168x2048)
strides: [7168, 1], [1, 7168]
dtypes: torch.float16, torch.float16
  mm 0.0743 ms 100.0% 
  triton_mm_354 0.0866 ms 85.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_359 0.0908 ms 81.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_348 0.0923 ms 80.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_353 0.0933 ms 79.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_365 0.1077 ms 69.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_349 0.1096 ms 67.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_356 0.1161 ms 64.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
[...] 149 lines omitted