submission 189837

import torch
from task import input_t, output_t

_tmp1 = None
_tmp2 = None


def custom_kernel(data: input_t) -> output_t:
    a, b1, b2, _, _, _, sfa_p, sfb1_p, sfb2_p, c = data
    m, n, L = c.shape
    device = c.device

    global _tmp1, _tmp2
    if _tmp1 is None or _tmp1.shape != (m, n) or _tmp1.device != device:
        _tmp1 = torch.empty((m, n), device=device, dtype=torch.float16)
        _tmp2 = torch.empty((m, n), device=device, dtype=torch.float16)

    for l_idx in range(L):

        torch._scaled_mm(
            a[:, :, l_idx],
            b1[:, :, l_idx].transpose(0, 1),
            sfa_p[..., l_idx].permute(2, 4, 0, 1, 3).contiguous().view(-1),
            sfb1_p[..., l_idx].permute(2, 4, 0, 1, 3).contiguous().view(-1),
            bias=None,
            out_dtype=torch.float16,
            out=_tmp1,
        )
        torch._scaled_mm(
            a[:, :, l_idx],
            b2[:, :, l_idx].transpose(0, 1),
            sfa_p[..., l_idx].permute(2, 4, 0, 1, 3).contiguous().view(-1),
            sfb2_p[..., l_idx].permute(2, 4, 0, 1, 3).contiguous().view(-1),
            bias=None,
            out_dtype=torch.float16,
            out=_tmp2,
        )

        c[:, :, l_idx] = (torch.nn.functional.silu(_tmp1) * _tmp2)

    return c