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

# --- Configuration for torch.compile ---
if hasattr(torch, '_dynamo') and hasattr(torch._dynamo, 'config'):
    torch._dynamo.config.capture_scalar_outputs = True
    # torch._dynamo.config.suppress_warnings = True # This line caused an AttributeError and is removed.
else:
    print("Warning: torch._dynamo.config not found. Cannot set capture_scalar_outputs.")


# --- Model Classes ---
class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.act_fn = nn.SiLU()
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = d_expert if d_expert is not None else config["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:
        # Gracefully handle empty tensor inputs.
        # If the token dimension (assumed to be dim 0) is 0.
        if x.shape[0] == 0:
            # Construct the expected output shape with the token dimension as 0.
            output_shape = list(x.shape)
            output_shape[-1] = self.d_hidden # Output dimension is d_hidden
            return torch.zeros(output_shape, device=x.device, dtype=x.dtype)

        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]:
        logits = self.W_g(x)
        # Inductor warning about online softmax might appear here if not suppressed.
        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_Optimized(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.num_routed_experts = config["n_routed_experts"]
        self.experts = nn.ModuleList([
            Expert(config)
            for _ in range(self.num_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)

        orig_shape = x.shape

        x_flat = x.view(-1, x.shape[-1])
        flat_expert_indices_router_output = expert_indices.view(-1)
        flat_expert_weights = expert_scores.view(-1, 1)

        routed_output_flat = self.moe_infer(x_flat,
                                            flat_expert_indices_router_output,
                                            flat_expert_weights)

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

    def moe_infer(self,
                  x_flat: torch.Tensor,
                  flat_expert_indices_router_output: torch.Tensor,
                  flat_expert_weights: torch.Tensor
                 ) -> torch.Tensor:
        expert_cache = torch.zeros_like(x_flat)

        # sorted_assignment_original_indices: Original positions of assignments (0 to N*K-1)
        #                                     in flat_expert_indices_router_output / flat_expert_weights,
        #                                     after these assignments are sorted by the expert_id they are routed to.
        sorted_assignment_original_indices = flat_expert_indices_router_output.argsort()

        # sorted_expert_ids: For each assignment in the sorted list, which expert_id it's routed to.
        sorted_expert_ids = flat_expert_indices_router_output[sorted_assignment_original_indices]

        num_per_tok = self.config["n_experts_per_token"] # Python int from config.
        # original_token_indices_for_sorted_assignments: For each assignment in the sorted list,
        # this gives the original token index (row in x_flat) it corresponds to.
        original_token_indices_for_sorted_assignments = sorted_assignment_original_indices // num_per_tok

        for expert_id in range(self.num_routed_experts):
            # Create a boolean mask to identify assignments for the current expert_id
            # from the `sorted_expert_ids` tensor.
            expert_mask = (sorted_expert_ids == expert_id)

            # Use the mask to select the original indices of assignments for this expert.
            # If no assignments for this expert, these will be empty tensors.
            current_expert_assignment_indices_from_sorted = sorted_assignment_original_indices[expert_mask]
            current_expert_token_indices_in_xflat = original_token_indices_for_sorted_assignments[expert_mask]

            expert_tokens_input = x_flat[current_expert_token_indices_in_xflat]

            expert_module = self.experts[expert_id]
            expert_output_unweighted = expert_module(expert_tokens_input) # Handles empty inputs

            weights_for_current_assignments = flat_expert_weights[current_expert_assignment_indices_from_sorted]
            expert_output_weighted = expert_output_unweighted * weights_for_current_assignments # Handles empty

            # The `if current_expert_token_indices_in_xflat.numel() > 0:` condition was REMOVED.
            # We rely on scatter_reduce_ to correctly handle cases where expert_output_weighted
            # (the src) is an empty tensor (e.g., shape [0, hidden_dim]) and
            # scatter_indices (derived from current_expert_token_indices_in_xflat) is also empty
            # (e.g., shape [0, hidden_dim]). In such cases, scatter_reduce_ should effectively be a no-op.

            # scatter_indices will be empty with the correct rank if current_expert_token_indices_in_xflat is empty.
            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,        # index tensor; can be empty
                expert_output_weighted, # src tensor; can be empty
                reduce='sum'
                # include_self=False is default for 'sum', which is fine.
            )

        return expert_cache

# --- custom_kernel: Entry point for the submission ---
input_t = Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict]
output_t_for_custom_kernel = torch.Tensor # Type hint for the return value

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)

    moe_module = MoE_Optimized(config)

    # Weight loading (identical to previous correct versions)
    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()

    if hasattr(torch, 'compile'):
        try:
            # Using "max-autotune" as the primary attempt.
            # print("Attempting to compile MoE module with mode='max-autotune'...") # For local debugging
            moe_module = torch.compile(moe_module, mode="max-autotune")
            # print("Compilation with 'max-autotune' mode successful.") # For local debugging
        except Exception as e_max_autotune:
            print(f"Warning: torch.compile with mode='max-autotune' failed: {e_max_autotune}. Trying 'default'.")
            try:
                moe_module = torch.compile(moe_module, mode="default")
            except Exception as e_default:
                 print(f"Warning: torch.compile with mode='default' also failed: {e_default}. Running uncompiled.")

    with torch.inference_mode():
        output_tensor = moe_module(input_tensor)

    return output_tensor

python eval.py test /tmp/tmpm0ew5x0u
/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.0637 ms 100.0% 
  triton_mm_69 0.0787 ms 80.9% 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_72 0.0806 ms 79.1% 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_77 0.0811 ms 78.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_76 0.0818 ms 77.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_74 0.0842 ms 75.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_78 0.0850 ms 74.9% 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_86 0.0987 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_84 0.1012 ms 63.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
  triton_mm_81 0.1115 ms 57.2% 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.4416 seconds and 0.0014 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.0190 ms 97.7% 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_8 0.0274 ms 67.8% 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_14 0.0274 ms 67.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_4 0.0275 ms 67.6% 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_5 0.0276 ms 67.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_11 0.0277 ms 66.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_12 0.0281 ms 66.0% 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_7 0.0330 ms 56.3% 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_10 0.0332 ms 55.8% 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 0.2700 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.0474 ms 100.0% 
  triton_mm_122 0.0516 ms 92.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
  triton_mm_129 0.0516 ms 91.9% 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.0526 ms 90.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_126 0.0549 ms 86.3% 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_125 0.0564 ms 84.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=0, num_stages=2, num_warps=8
  triton_mm_128 0.0566 ms 83.7% 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_127 0.0567 ms 83.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_124 0.0590 ms 80.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_121 0.0592 ms 80.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
SingleProcess AUTOTUNE benchmarking takes 4.2975 seconds and 0.0003 seconds precompiling for 37 choices
AUTOTUNE mm(1024x7168, 7168x2048)
strides: [7168, 1], [2048, 1]
dtypes: torch.float16, torch.float16
  mm 0.0951 ms 100.0% 
  triton_mm_331 0.1146 ms 83.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
  triton_mm_326 0.1159 ms 82.1% 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_321 0.1187 ms 80.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_327 0.1266 ms 75.1% 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_329 0.1275 ms 74.6% 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.1293 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_338 0.1299 ms 73.2% 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.1314 ms 72.4% 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.1392 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.6207 seconds and 0.0997 seconds precompiling for 37 choices
AUTOTUNE mm(1024x7168, 7168x8)
strides: [7168, 1], [1, 7168]
dtypes: torch.float16, torch.float16
  triton_mm_248 0.0194 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.0221 ms 87.7% 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 81.3% 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.0276 ms 70.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.0278 ms 69.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_256 0.0278 ms 69.7% 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.0279 ms 69.3% 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_257 0.0281 ms 69.0% 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_255 0.0336 ms 57.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
  triton_mm_252 0.0336 ms 57.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
SingleProcess AUTOTUNE benchmarking takes 3.4954 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.0673 ms 100.0% 
  triton_mm_380 0.0842 ms 80.0% 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.0854 ms 78.9% 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_369 0.0887 ms 75.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=2, num_stages=2, num_warps=4
  triton_mm_366 0.0903 ms 74.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_375 0.0916 ms 73.5% 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_378 0.0926 ms 72.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_379 0.0954 ms 70.6% 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_360 0.0957 ms 70.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_367 0.0981 ms 68.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
SingleProcess AUTOTUNE benchmarking takes 4.4422 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.0550 ms 100.0% 
  triton_mm_422 0.0819 ms 67.2% 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_425 0.0835 ms 65.9% 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_427 0.0854 ms 64.4% 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_429 0.0854 ms 64.4% 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_430 0.0864 ms 63.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_431 0.0891 ms 61.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
[...] 61 lines omitted