submission 119501

import torch
import triton
import triton.language as tl
from task import input_t, output_t
from utils import make_match_reference

sf_vec_size = 16

def ceil_div(a, b):
    return (a + b - 1) // b

@triton.jit
def to_blocked_kernel(in_ptr, out_ptr, cols, n_col_blocks, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    row_block = pid // n_col_blocks
    col_block = pid % n_col_blocks
    
    offs = tl.arange(0, BLOCK_SIZE)
    i32 = offs // 16
    temp = offs % 16
    i4 = temp // 4
    k4 = temp % 4
    
    row_in_block = i4 * 32 + i32
    row = row_block * 128 + row_in_block
    col = col_block * 4 + k4
    
    in_idx = row * cols + col
    val = tl.load(in_ptr + in_idx)
    
    out_idx = pid * BLOCK_SIZE + offs
    tl.store(out_ptr + out_idx, val)

@triton.jit
def to_blocked_dual_kernel(
    in_a_ptr, out_a_ptr, rows_a, cols_a, n_col_blocks_a, n_blocks_a,
    in_b_ptr, out_b_ptr, rows_b, cols_b, n_col_blocks_b,
    BLOCK_SIZE: tl.constexpr
):
    pid = tl.program_id(0)
    
    if pid < n_blocks_a:
        row_block = pid // n_col_blocks_a
        col_block = pid % n_col_blocks_a
        in_ptr = in_a_ptr
        out_ptr = out_a_ptr
        cols = cols_a
    else:
        local_pid = pid - n_blocks_a
        row_block = local_pid // n_col_blocks_b
        col_block = local_pid % n_col_blocks_b
        in_ptr = in_b_ptr
        out_ptr = out_b_ptr
        cols = cols_b
        pid = local_pid
    
    offs = tl.arange(0, BLOCK_SIZE)
    i32 = offs // 16
    temp = offs % 16
    i4 = temp // 4
    k4 = temp % 4
    
    row_in_block = i4 * 32 + i32
    row = row_block * 128 + row_in_block
    col = col_block * 4 + k4
    
    in_idx = row * cols + col
    val = tl.load(in_ptr + in_idx)
    
    out_idx = pid * BLOCK_SIZE + offs
    tl.store(out_ptr + out_idx, val)

def to_blocked(input_matrix):
    rows, cols = input_matrix.shape
    n_row_blocks = (rows + 127) // 128
    n_col_blocks = (cols + 3) // 4
    n_blocks = n_row_blocks * n_col_blocks
    output = torch.empty(n_blocks * 512, dtype=input_matrix.dtype, device=input_matrix.device)
    to_blocked_kernel[(n_blocks,)](input_matrix, output, cols, n_col_blocks, BLOCK_SIZE=512, num_warps=8)
    return output

def to_blocked_dual(input_a, input_b):
    rows_a, cols_a = input_a.shape
    rows_b, cols_b = input_b.shape
    n_row_blocks_a = (rows_a + 127) // 128
    n_col_blocks_a = (cols_a + 3) // 4
    n_blocks_a = n_row_blocks_a * n_col_blocks_a
    n_row_blocks_b = (rows_b + 127) // 128
    n_col_blocks_b = (cols_b + 3) // 4
    n_blocks_b = n_row_blocks_b * n_col_blocks_b
    output_a = torch.empty(n_blocks_a * 512, dtype=input_a.dtype, device=input_a.device)
    output_b = torch.empty(n_blocks_b * 512, dtype=input_b.dtype, device=input_b.device)
    to_blocked_dual_kernel[(n_blocks_a + n_blocks_b,)](
        input_a, output_a, rows_a, cols_a, n_col_blocks_a, n_blocks_a,
        input_b, output_b, rows_b, cols_b, n_col_blocks_b,
        BLOCK_SIZE=512, num_warps=8
    )
    return output_a, output_b

def ref_kernel(data: input_t) -> output_t:
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data
    _, _, l = c_ref.shape
    for l_idx in range(l):
        scale_a = to_blocked(sfa_ref_cpu[:, :, l_idx])
        scale_b = to_blocked(sfb_ref_cpu[:, :, l_idx])
        res = torch._scaled_mm(
            a_ref[:, :, l_idx],
            b_ref[:, :, l_idx].transpose(0, 1),
            scale_a.cuda(),
            scale_b.cuda(),
            bias=None,
            out_dtype=torch.float16,
        )
        c_ref[:, :, l_idx] = res
    return c_ref

def custom_kernel(data: input_t) -> output_t:
    a_ref, b_ref, sfa_ref, sfb_ref, _, _, c_ref = data
    _, _, l = c_ref.shape
    for l_idx in range(l):
        scale_a, scale_b = to_blocked_dual(sfa_ref[:, :, l_idx], sfb_ref[:, :, l_idx])
        torch._scaled_mm(
            a_ref[:, :, l_idx],
            b_ref[:, :, l_idx].T,
            scale_a,
            scale_b,
            bias=None,
            out_dtype=torch.float16,
            out=c_ref[:, :, l_idx],
        )
    return c_ref

def generate_input(m: int, n: int, k: int, l: int, seed: int):
    torch.manual_seed(seed)
    a_ref = torch.randint(-6, 6, (l, m, k // 2), dtype=torch.int8, device="cuda").permute(1, 2, 0)
    b_ref = torch.randint(-6, 6, (l, n, k // 2), dtype=torch.int8, device="cuda").permute(1, 2, 0)
    a_ref = a_ref.view(torch.float4_e2m1fn_x2)
    b_ref = b_ref.view(torch.float4_e2m1fn_x2)
    c_ref = torch.empty((m, n, l), dtype=torch.float16, device="cuda")
    def create_scale_factor_tensors(l, mn, sf_k):
        ref_shape = (l, mn, sf_k)
        ref_permute_order = (1, 2, 0)
        ref_f8_random_int = torch.randint(-3, 3, ref_shape, dtype=torch.int8, device="cuda")
        ref_f8_torch_tensor = ref_f8_random_int.to(dtype=torch.float8_e4m3fn)
        ref_f8_torch_tensor_permuted = ref_f8_torch_tensor.permute(*ref_permute_order)
        atom_m = (32, 4)
        atom_k = 4
        mma_shape = (l, ceil_div(mn, atom_m[0] * atom_m[1]), ceil_div(sf_k, atom_k), atom_m[0], atom_m[1], atom_k)
        mma_permute_order = (3, 4, 1, 5, 2, 0)
        rand_int_tensor = torch.randint(-3, 3, mma_shape, dtype=torch.int8, device="cuda")
        reordered_f8_torch_tensor = rand_int_tensor.to(dtype=torch.float8_e4m3fn)
        reordered_f8_torch_tensor = reordered_f8_torch_tensor.permute(*mma_permute_order)
        i_idx = torch.arange(mn, device="cuda")
        j_idx = torch.arange(sf_k, device="cuda")
        b_idx = torch.arange(l, device="cuda")
        i_grid, j_grid, b_grid = torch.meshgrid(i_idx, j_idx, b_idx, indexing="ij")
        mm = i_grid // (atom_m[0] * atom_m[1])
        mm32 = i_grid % atom_m[0]
        mm4 = (i_grid % 128) // atom_m[0]
        kk = j_grid // atom_k
        kk4 = j_grid % atom_k
        reordered_f8_torch_tensor[mm32, mm4, mm, kk4, kk, b_grid] = ref_f8_torch_tensor_permuted[i_grid, j_grid, b_grid]
        return ref_f8_torch_tensor_permuted.cpu(), reordered_f8_torch_tensor
    sf_k = ceil_div(k, sf_vec_size)
    sfa_ref_cpu, sfa_ref_permuted = create_scale_factor_tensors(l, m, sf_k)
    sfb_ref_cpu, sfb_ref_permuted = create_scale_factor_tensors(l, n, sf_k)
    return (a_ref, b_ref, sfa_ref_cpu.to("cuda"), sfb_ref_cpu.to("cuda"), sfa_ref_permuted, sfb_ref_permuted, c_ref)

check_implementation = make_match_reference(ref_kernel, rtol=1e-03, atol=1e-03)