submission 257399

import torch

def custom_kernel(data):
    a, b1, b2, sfa, sfb1, sfb2, sfa_p, sfb1_p, sfb2_p, c = data
    M, K, L = a.shape
    N = b1.shape[0]

    out = torch.empty_like(c)

    # Pre-compute layout transformation on GPU once if possible
    # to_blocked logic: 
    # blocks = input.view(rows//128, 128, cols//4, 4).permute(0, 2, 1, 3)
    # rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
    
    # Let's optimize this function to run entirely on GPU without intermediate copies
    def to_blocked_fast(input_matrix):
        # Input is (M, K) or (N, K)
        rows, cols = input_matrix.shape
        # We assume shapes are aligned (verified by problem statement)
        
        # 1. View as blocks
        # 2. Permute: (RowBlock, ColBlock, 128, 4)
        x = input_matrix.view(rows // 128, 128, cols // 4, 4).permute(0, 2, 1, 3)
        
        # 3. Reshape/Transpose sequence for UE4M3 layout
        # (..., 4, 32, 4) -> (..., 32, 4, 4) -> (..., 32, 16)
        x = x.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
        return x.flatten()

    for l_idx in range(L):
        # We use the raw 'sfa' (not sfa_p) but keep it on GPU.
        # sfa is ALREADY on GPU (passed as cuda tensor).
        # We just need to drive the strides correctly.
        
        scale_a = to_blocked_fast(sfa[:, :, l_idx])
        scale_b1 = to_blocked_fast(sfb1[:, :, l_idx])
        scale_b2 = to_blocked_fast(sfb2[:, :, l_idx])
        
        res1 = torch._scaled_mm(
            a[:, :, l_idx],
            b1[:, :, l_idx].transpose(0, 1),
            scale_a,
            scale_b1,
            out_dtype=torch.float32
        )
        
        res2 = torch._scaled_mm(
            a[:, :, l_idx],
            b2[:, :, l_idx].transpose(0, 1),
            scale_a,
            scale_b2,
            out_dtype=torch.float32
        )
        
        out[:, :, l_idx] = (torch.nn.functional.silu(res1) * res2).to(torch.float16)

    return out