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


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.experts = nn.ModuleList([
            Expert(config) for _ in range(self.num_routed_experts)
        ])
        
        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)

    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

            expert_module = self.experts[expert_id]
            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 = expert_module(expert_tokens_input)
            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))

    for i in range(n_routed_experts_config):
        gate_w = weights_dict[f'experts.{i}.0.weight'].t().to(device, dtype=torch.float16)
        up_w = weights_dict[f'experts.{i}.1.weight'].t().to(device, dtype=torch.float16)
        down_w = weights_dict[f'experts.{i}.2.weight'].t().to(device, dtype=torch.float16)

        moe_module.experts[i].W_gate.weight = nn.Parameter(gate_w)
        moe_module.experts[i].W_up.weight = nn.Parameter(up_w)
        moe_module.experts[i].W_down.weight = nn.Parameter(down_w)

    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 benchmark /tmp/tmpq07rwt_c
AUTOTUNE mm(2048x7168, 7168x2048)
strides: [7168, 1], [2048, 1]
dtypes: torch.float16, torch.float16
  mm 0.1185 ms 100.0% 
  triton_mm_64 0.1591 ms 74.5% 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_57 0.1685 ms 70.4% 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_50 0.1716 ms 69.1% 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_56 0.1772 ms 66.9% 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_61 0.1777 ms 66.7% 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_63 0.1794 ms 66.1% 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_60 0.1829 ms 64.8% 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_62 0.1893 ms 62.6% 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_67 0.1926 ms 61.6% 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 1.4152 seconds and 0.0020 seconds precompiling for 37 choices
AUTOTUNE mm(2048x2048, 2048x7168)
strides: [2048, 1], [7168, 1]
dtypes: torch.float16, torch.float16
  mm 0.1164 ms 100.0% 
  triton_mm_104 0.1255 ms 92.7% 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_101 0.1304 ms 89.2% 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_95 0.1501 ms 77.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_96 0.1560 ms 74.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=0, num_stages=2, num_warps=8
  triton_mm_97 0.1574 ms 73.9% 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_107 0.1604 ms 72.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=256, 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_93 0.1605 ms 72.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_100 0.1608 ms 72.4% 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_99 0.1667 ms 69.8% 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
SingleProcess AUTOTUNE benchmarking takes 4.7528 seconds and 0.0003 seconds precompiling for 37 choices
AUTOTUNE mm(2048x7168, 7168x32)
strides: [7168, 1], [1, 7168]
dtypes: torch.float16, torch.float16
  triton_mm_109 0.0237 ms 100.0% 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_111 0.0327 ms 72.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=32, 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_120 0.0333 ms 71.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=32, 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_121 0.0334 ms 70.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=64, BLOCK_N=32, 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_125 0.0336 ms 70.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=32, 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
  mm 0.0340 ms 69.7% 
  triton_mm_113 0.0363 ms 65.1% 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_117 0.0406 ms 58.3% 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_124 0.0484 ms 48.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=32, 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_116 0.0489 ms 48.4% 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=8
SingleProcess AUTOTUNE benchmarking takes 4.5286 seconds and 0.0003 seconds precompiling for 33 choices
AUTOTUNE mm(240x7168, 7168x2048)
strides: [7168, 1], [2048, 1]
dtypes: torch.float16, torch.float16
  mm 0.0501 ms 100.0% 
  triton_mm_179 0.0653 ms 76.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=32, 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_185 0.0671 ms 74.7% 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_184 0.0691 ms 72.5% 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=8
  triton_mm_180 0.0699 ms 71.7% 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_188 0.0752 ms 66.6% 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_183 0.0771 ms 65.0% 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_181 0.0793 ms 63.2% 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_187 0.0806 ms 62.1% 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.0837 ms 59.8% 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
SingleProcess AUTOTUNE benchmarking takes 4.4068 seconds and 0.1354 seconds precompiling for 37 choices
AUTOTUNE mm(240x2048, 2048x7168)
strides: [2048, 1], [7168, 1]
dtypes: torch.float16, torch.float16
  triton_mm_233 0.0349 ms 100.0% 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
  mm 0.0361 ms 96.7% 
  triton_mm_223 0.0383 ms 91.1% 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_228 0.0384 ms 90.9% 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_231 0.0403 ms 86.7% 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_229 0.0418 ms 83.5% 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_234 0.0420 ms 83.2% 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_240 0.0423 ms 82.5% 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_227 0.0426 ms 82.0% 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_232 0.0427 ms 81.7% 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.3251 seconds and 0.0002 seconds precompiling for 37 choices
AUTOTUNE mm(268x7168, 7168x2048)
strides: [7168, 1], [2048, 1]
dtypes: torch.float16, torch.float16
  mm 0.0433 ms 100.0% 
  triton_mm_287 0.0671 ms 64.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=32, 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_288 0.0704 ms 61.5% 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_296 0.0768 ms 56.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_291 0.0779 ms 55.5% 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_293 0.0791 ms 54.7% 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_295 0.0816 ms 53.0% 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_297 0.0843 ms 51.3% 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_303 0.0894 ms 48.4% 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_305 0.0985 ms 43.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.4280 seconds and 0.1074 seconds precompiling for 37 choices
AUTOTUNE mm(268x2048, 2048x7168)
strides: [2048, 1], [7168, 1]
dtypes: torch.float16, torch.float16
  mm 0.0348 ms 100.0% 
  triton_mm_331 0.0411 ms 84.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_336 0.0415 ms 83.8% 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_335 0.0452 ms 76.9% 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_348 0.0455 ms 76.5% 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_337 0.0461 ms 75.5% 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_334 0.0478 ms 72.7% 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_341 0.0492 ms 70.7% 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_332 0.0494 ms 70.4% 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_347 0.0524 ms 66.3% 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
SingleProcess AUTOTUNE benchmarking takes 4.3041 seconds and 0.0002 seconds precompiling for 37 choices
AUTOTUNE mm(268x7168, 7168x2048)
[...] 151 lines omitted