submission 188960

import torch
import torch.nn.functional as F


def _to_blocked(input_matrix: torch.Tensor) -> torch.Tensor:
    rows, cols = input_matrix.shape
    n_row_blocks = (rows + 127) // 128
    n_col_blocks = (cols + 3) // 4
    padded = input_matrix.contiguous()
    blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
    return rearranged.flatten()


def custom_kernel(data):
    a, b1, b2, sfa, sfb1, sfb2, sfa_p, sfb1_p, sfb2_p, c = data
    out = c.contiguous()
    _, _, l = out.shape
    for l_idx in range(l):
        scale_a = _to_blocked(sfa[:, :, l_idx])
        scale_b1 = _to_blocked(sfb1[:, :, l_idx])
        scale_b2 = _to_blocked(sfb2[:, :, l_idx])
        g1 = torch._scaled_mm(
            a[:, :, l_idx],
            b1[:, :, l_idx].transpose(0, 1),
            scale_a,
            scale_b1,
            bias=None,
            out_dtype=torch.float16,
        )
        g2 = torch._scaled_mm(
            a[:, :, l_idx],
            b2[:, :, l_idx].transpose(0, 1),
            scale_a,
            scale_b2,
            bias=None,
            out_dtype=torch.float16,
        )
        out[:, :, l_idx] = F.silu(g1) * g2
    return out


__all__ = ["custom_kernel"]