import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, List, Optional # Keep List from reference
import math

# --- Start of Reference Classes (Expert, MoEGate are used as is) ---
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))
        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 = nn.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
# --- End of Reference Classes (Expert, MoEGate) ---


# MoE class, based on reference, but with optimized moe_infer
class MoE_Optimized(nn.Module): # Renamed to avoid clash if original MoE is needed elsewhere
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.num_routed_experts = config["n_routed_experts"] # Used in moe_infer
        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)
        
        # Storing original shape for view operation later, as in reference
        orig_shape = x.shape 
        
        x_flat = x.view(-1, x.shape[-1]) # Use x.shape[-1] for 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) # Use orig_shape as in reference
        return routed_output + shared_output

    # @torch.no_grad() # torch.compile handles this when in inference mode.
                      # Keep if not compiling or for clarity if preferred.
                      # For compiled version, usually not needed on the method itself.
    def moe_infer(self,
                  x_flat: torch.Tensor, # Renamed from x in reference for clarity
                  flat_expert_indices: torch.Tensor,
                  flat_expert_weights: torch.Tensor
                 ) -> torch.Tensor:
        expert_cache = torch.zeros_like(x_flat)
        idxs = flat_expert_indices.argsort() # Indices to sort assignments by expert
        
        # --- OPTIMIZED GPU-based counts and cumsum ---
        # Use minlength for robustness, ensuring counts_gpu has num_routed_experts elements.
        # This handles cases where some experts might not receive any tokens.
        counts_gpu = flat_expert_indices.bincount(minlength=self.num_routed_experts)
        tokens_per_expert_gpu_cumsum = counts_gpu.cumsum(dim=0) # Cumulative sum on GPU
        # --- End of OPTIMIZATION ---
        
        num_per_tok = self.config["n_experts_per_token"]
        # original_token_indices_for_sorted_assignments: for each item in `idxs` (sorted assignments),
        # this gives the original token index in x_flat.
        original_token_indices_for_sorted_assignments = idxs // num_per_tok

        current_assignment_idx_start = 0 # Python int for defining slice ranges
        for expert_id in range(self.num_routed_experts): # Loop over all potential experts
            # Get the end index for the current expert's assignments from the cumsum tensor.
            # .item() converts the 0-dim tensor to a Python scalar.
            end_idx_for_assignments = tokens_per_expert_gpu_cumsum[expert_id].item()

            if current_assignment_idx_start == end_idx_for_assignments: # No tokens for this expert
                continue # current_assignment_idx_start remains end_idx_for_assignments for next iter check implicitly

            expert = self.experts[expert_id] # nn.ModuleList access
            
            # Slice `idxs` to get the original indices (into flat_expert_weights) of assignments for this expert
            current_expert_assignment_original_indices = idxs[current_assignment_idx_start:end_idx_for_assignments]
            
            # Slice `original_token_indices_for_sorted_assignments` to get actual token indices in x_flat
            current_expert_token_indices_in_xflat = original_token_indices_for_sorted_assignments[current_assignment_idx_start:end_idx_for_assignments]
            
            expert_tokens_input = x_flat[current_expert_token_indices_in_xflat]
            expert_output_weighted = expert(expert_tokens_input) # Expert forward pass
            
            weights_for_current_assignments = flat_expert_weights[current_expert_assignment_original_indices]
            expert_output_weighted.mul_(weights_for_current_assignments) # Apply router scores
            
            # Prepare indices for scatter_reduce_ operation
            scatter_indices = current_expert_token_indices_in_xflat.view(-1, 1).repeat(1, x_flat.shape[-1])
            
            expert_cache.scatter_reduce_(
                0, # dim
                scatter_indices,
                expert_output_weighted,
                reduce='sum'
                # include_self=False is default for 'sum'
            )
            
            current_assignment_idx_start = end_idx_for_assignments # Update start for the next expert
            
        return expert_cache

# --- custom_kernel: Entry point for the submission ---
# Define input_t and output_t_for_custom_kernel based on the problem description's usage
input_t = Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict]
output_t_for_custom_kernel = torch.Tensor # ref_kernel code returns a single tensor

def custom_kernel(data: input_t) -> output_t_for_custom_kernel:
    input_tensor, weights, config = data
    
    input_tensor = input_tensor.to(device='cuda', dtype=torch.float16)

    # Use the optimized MoE implementation
    moe_module = MoE_Optimized(config) 
    
    # Weight loading (identical to reference ref_kernel)
    moe_module.gating_network.W_g.weight = nn.Parameter(
        weights['router.weight'].to(device='cuda', dtype=torch.float16)
    )
    num_routed_experts = config["n_routed_experts"]
    for i in range(num_routed_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']
        moe_module.experts[i].W_gate.weight = nn.Parameter(gate_proj_weight.t().to(device='cuda', dtype=torch.float16))
        moe_module.experts[i].W_up.weight = nn.Parameter(up_proj_weight.t().to(device='cuda', dtype=torch.float16))
        moe_module.experts[i].W_down.weight = nn.Parameter(down_proj_weight.t().to(device='cuda', dtype=torch.float16))
        
    moe_module.shared_expert.W_gate.weight = nn.Parameter(
        weights['shared_experts.0.weight'].t().to(device='cuda', dtype=torch.float16)
    )
    moe_module.shared_expert.W_up.weight = nn.Parameter(
        weights['shared_experts.1.weight'].t().to(device='cuda', dtype=torch.float16)
    )
    moe_module.shared_expert.W_down.weight = nn.Parameter(
        weights['shared_experts.2.weight'].t().to(device='cuda', dtype=torch.float16)
    )
    
    moe_module = moe_module.to(device='cuda', dtype=torch.float16)
    moe_module.eval()

    # Re-introduce torch.compile for performance
    # Using "max-autotune" for best potential speed given the large performance gap.
    # Fallback to "reduce-overhead" or "default" if max-autotune causes issues or is too slow to compile.
    if hasattr(torch, 'compile'):
        # Check for PyTorch version if specific options are version-dependent
        # For PyTorch 2.8.0.dev, max-autotune should be available.
        try:
            # print("Attempting to compile MoE module with mode='max-autotune'...") # For local debugging
            moe_module = torch.compile(moe_module, mode="max-autotune")
            # print("Compilation successful.") # For local debugging
        except Exception as e:
            # Fallback if compilation fails
            print(f"Warning: torch.compile with mode='max-autotune' failed: {e}. Trying 'reduce-overhead'.")
            try:
                moe_module = torch.compile(moe_module, mode="reduce-overhead")
            except Exception as e2:
                 print(f"Warning: torch.compile with mode='reduce-overhead' also failed: {e2}. Running uncompiled.")
    
    # Perform inference. Gradient tracking is disabled by torch.compile in inference mode,
    # or by torch.no_grad() if not compiled / compile fails.
    # Using torch.inference_mode() is often preferred over torch.no_grad() for inference.
    with torch.inference_mode(): # More modern context manager for inference
        output_tensor = moe_module(input_tensor) 
    
    return output_tensor
python eval.py test /tmp/tmpw0vnexoi
/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.0643 ms 100.0% 
  triton_mm_69 0.0815 ms 79.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_77 0.0817 ms 78.7% 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_74 0.0823 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_76 0.0825 ms 78.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_72 0.0840 ms 76.6% 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_78 0.0852 ms 75.5% 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_84 0.0997 ms 64.5% 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_86 0.0998 ms 64.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_81 0.1118 ms 57.6% 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 1.2375 seconds and 0.0013 seconds precompiling for 37 choices
AUTOTUNE mm(512x7168, 7168x4)
strides: [7168, 1], [1, 7168]
dtypes: torch.float16, torch.float16
  triton_mm_1 0.0186 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_3 0.0231 ms 80.9% 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_11 0.0273 ms 68.2% 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_8 0.0274 ms 68.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=0, num_stages=2, num_warps=4
  triton_mm_5 0.0275 ms 67.7% 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_4 0.0275 ms 67.7% 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_14 0.0276 ms 67.5% 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_12 0.0282 ms 66.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
  triton_mm_10 0.0332 ms 56.2% 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
  triton_mm_7 0.0333 ms 56.0% 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
SingleProcess AUTOTUNE benchmarking takes 0.2749 seconds and 0.0013 seconds precompiling for 30 choices
AUTOTUNE mm(512x2048, 2048x7168)
strides: [2048, 1], [7168, 1]
dtypes: torch.float16, torch.float16
  mm 0.0459 ms 100.0% 
  triton_mm_122 0.0520 ms 88.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_129 0.0529 ms 86.7% 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_115 0.0560 ms 81.9% 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_128 0.0560 ms 81.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
  triton_mm_126 0.0560 ms 81.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=1, num_stages=2, num_warps=8
  triton_mm_127 0.0566 ms 81.1% 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_125 0.0575 ms 79.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=0, num_stages=2, num_warps=8
  triton_mm_121 0.0579 ms 79.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
  triton_mm_124 0.0604 ms 76.0% 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
SingleProcess AUTOTUNE benchmarking takes 4.2941 seconds and 0.0003 seconds precompiling for 37 choices
W0518 21:58:35.461000 514 torch/_dynamo/variables/tensor.py:936] [2/0] Graph break from `Tensor.item()`, consider setting:
W0518 21:58:35.461000 514 torch/_dynamo/variables/tensor.py:936] [2/0]     torch._dynamo.config.capture_scalar_outputs = True
W0518 21:58:35.461000 514 torch/_dynamo/variables/tensor.py:936] [2/0] or:
W0518 21:58:35.461000 514 torch/_dynamo/variables/tensor.py:936] [2/0]     env TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1
W0518 21:58:35.461000 514 torch/_dynamo/variables/tensor.py:936] [2/0] to include these operations in the captured graph.
W0518 21:58:35.461000 514 torch/_dynamo/variables/tensor.py:936] [2/0] 
W0518 21:58:35.461000 514 torch/_dynamo/variables/tensor.py:936] [2/0] Graph break: from user code at:
W0518 21:58:35.461000 514 torch/_dynamo/variables/tensor.py:936] [2/0]   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 101, in torch_dynamo_resume_in_moe_infer_at_88
W0518 21:58:35.461000 514 torch/_dynamo/variables/tensor.py:936] [2/0]     end_idx_for_assignments = tokens_per_expert_gpu_cumsum[expert_id].item()
W0518 21:58:35.461000 514 torch/_dynamo/variables/tensor.py:936] [2/0] 
W0518 21:58:35.461000 514 torch/_dynamo/variables/tensor.py:936] [2/0] 
AUTOTUNE mm(1024x7168, 7168x2048)
strides: [7168, 1], [2048, 1]
dtypes: torch.float16, torch.float16
  mm 0.0971 ms 100.0% 
  triton_mm_331 0.1135 ms 85.6% 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_326 0.1162 ms 83.6% 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_327 0.1222 ms 79.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_321 0.1224 ms 79.3% 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_329 0.1293 ms 75.1% 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_332 0.1306 ms 74.4% 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_338 0.1315 ms 73.8% 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_325 0.1319 ms 73.6% 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_324 0.1423 ms 68.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
SingleProcess AUTOTUNE benchmarking takes 4.6451 seconds and 0.0995 seconds precompiling for 37 choices
AUTOTUNE mm(1024x7168, 7168x8)
strides: [7168, 1], [1, 7168]
dtypes: torch.float16, torch.float16
  triton_mm_248 0.0197 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.0219 ms 89.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 82.7% 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_250 0.0273 ms 72.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_256 0.0274 ms 71.9% 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_253 0.0275 ms 71.7% 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_259 0.0275 ms 71.7% 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_257 0.0283 ms 69.6% 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.0336 ms 58.7% 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.0336 ms 58.6% 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.4838 seconds and 0.0002 seconds precompiling for 30 choices
AUTOTUNE mm(1024x2048, 2048x7168)
strides: [2048, 1], [7168, 1]
dtypes: torch.float16, torch.float16
  mm 0.0694 ms 100.0% 
  triton_mm_380 0.0855 ms 81.2% 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_376 0.0861 ms 80.7% 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_360 0.0919 ms 75.5% 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_369 0.0930 ms 74.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=2, num_stages=2, num_warps=4
  triton_mm_366 0.0945 ms 73.5% 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_378 0.0953 ms 72.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_375 0.0960 ms 72.3% 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_379 0.0967 ms 71.8% 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_371 0.1003 ms 69.2% 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
SingleProcess AUTOTUNE benchmarking takes 4.4626 seconds and 0.0003 seconds precompiling for 37 choices
AUTOTUNE mm(536x7168, 7168x2048)
strides: [7168, 1], [2048, 1]
dtypes: torch.float16, torch.float16
  mm 0.0571 ms 100.0% 
  triton_mm_425 0.0875 ms 65.3% 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
[...] 66 lines omitted