submission 117927

import torch
from task import input_t, output_t

def custom_kernel(data: input_t) -> output_t:
    a, b, _, _, sfa_permuted, sfb_permuted, c = data
    
    # Use select() for potentially faster single-dimension indexing
    sfa_0 = sfa_permuted.select(5, 0)
    sfb_0 = sfb_permuted.select(5, 0)
    
    # Permute
    sfa_perm = sfa_0.permute(2,4,0,1,3)
    sfb_perm = sfb_0.permute(2,4,0,1,3)
    
    # Reshape
    sa = sfa_perm.reshape(-1)
    sb = sfb_perm.reshape(-1)
    
    # All matrix operations with select()
    a_0 = a.select(2, 0)
    b_0_t = b.select(2, 0).t()
    c_0 = c.select(2, 0)
    
    # GEMM
    torch._scaled_mm(a_0, b_0_t, sa, sb, out=c_0, out_dtype=torch.float16)
    
    return c