submission 424644

# submission.py - Simple correct version
import torch
from task import input_t, output_t
from reference import to_blocked

def custom_kernel(data: input_t) -> output_t:
    """
    Simple implementation that matches the reference exactly.
    Uses torch._scaled_mm for correctness.
    """
    abc_tensors, sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes = data
    
    result_tensors = []
    for i, ((a, b, c), (sfa_ref, sfb_ref), (m, n, k, l)) in enumerate(
        zip(abc_tensors, sfasfb_tensors, problem_sizes)
    ):
        for l_idx in range(l):
            # Convert scale factors to blocked format
            scale_a = to_blocked(sfa_ref[:, :, l_idx])
            scale_b = to_blocked(sfb_ref[:, :, l_idx])
            
            # Use torch's scaled_mm
            res = torch._scaled_mm(
                a[:, :, l_idx].view(torch.float4_e2m1fn_x2),
                b[:, :, l_idx].transpose(0, 1).view(torch.float4_e2m1fn_x2),
                scale_a.cuda(),
                scale_b.cuda(),
                bias=None,
                out_dtype=torch.float16,
            )
            c[:, :, l_idx] = res
        
        result_tensors.append(c)
    
    return result_tensors