submission 218864

import torch
from task import input_t, output_t

# Use reduce-overhead to enable CUDA Graphs and minimize dispatch latency
@torch.compile(mode="reduce-overhead")
def processing_loop(a, b1, b2, sfa_permuted, sfb1_permuted, sfb2_permuted, c):
    m, K_packed, l = a.shape
    
    # Pre-process Scales INSIDE the graph
    # Optimizes layout: Permute to [L, Mb, Kb, 32, 4, 4] -> Flatten
    # sfa_permuted is [32, 4, Mb, 4, Kb, L]
    # Permute indices: 5, 2, 4, 0, 1, 3
    sfa_p = sfa_permuted.permute(5, 2, 4, 0, 1, 3).contiguous().flatten(1, 5)
    sfb1_p = sfb1_permuted.permute(5, 2, 4, 0, 1, 3).contiguous().flatten(1, 5)
    sfb2_p = sfb2_permuted.permute(5, 2, 4, 0, 1, 3).contiguous().flatten(1, 5)
    
    # Iterate over L
    for i in range(l):
        # Extract slices
        a_slice = a[:, :, i]
        b1_slice = b1[:, :, i]
        b2_slice = b2[:, :, i]
        
        # Scales (contiguous slice due to pre-permute)
        scale_a = sfa_p[i, :]
        scale_b1 = sfb1_p[i, :]
        scale_b2 = sfb2_p[i, :]
        
        # GEMM
        # Use float32 output for precision correctness
        res1 = torch._scaled_mm(
            a_slice,
            b1_slice.t(), 
            scale_a,
            scale_b1,
            bias=None,
            out_dtype=torch.float32
        )
        
        res2 = torch._scaled_mm(
            a_slice,
            b2_slice.t(),
            scale_a,
            scale_b2,
            bias=None,
            out_dtype=torch.float32
        )
        
        # Elementwise
        c[:, :, i] = (torch.nn.functional.silu(res1) * res2).to(torch.float16)
    return c

def custom_kernel(data: input_t) -> output_t:
    a, b1, b2, sfa, sfb1, sfb2, sfa_permuted, sfb1_permuted, sfb2_permuted, c = data
    
    # Directly call the compiled loop with raw permuted inputs
    # The permutation logic is moved inside to be captured by CUDA Graphs
    return processing_loop(a, b1, b2, sfa_permuted, sfb1_permuted, sfb2_permuted, c)