import torch
from task import input_t, output_t

def custom_kernel(data: input_t) -> output_t:
    """
    Optimized implementation of block-scale FP8 GEMM.
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float8_e4m3fnuz] of shape [m, k], 
            b: torch.Tensor[float8_e4m3fnuz] of shape [n, k], 
            a_scale: torch.Tensor[float32] of shape [m, k // 128], 
            b_scale: torch.Tensor[float32] of shape [n // 128, k // 128], 
            c: torch.Tensor[bfloat16] of shape [m, n]
    Returns:
        Tensor containing output in bf16
    """
    a, b, a_scale, b_scale, c = data

    with torch.no_grad():
        # Extract dimensions
        m, k = a.shape
        n = b.shape[0]
        block_shape_n = 128
        block_shape_k = 128

        # Process a_scale
        scale_k = a_scale.shape[1]
        a_scale = a_scale.view(m, scale_k, 1)  # Reshape to [m, scale_k, 1]
        a_scale = a_scale.expand(-1, -1, block_shape_k)  # Expand to [m, scale_k, block_shape_k]
        a_scale = a_scale.reshape(m, -1)[:, :k]  # Flatten and slice to match 'k'

        # Process b_scale
        scale_n, scale_k = b_scale.shape
        b_scale = b_scale.view(scale_n, scale_k, 1, 1)  # Reshape to [scale_n, scale_k, 1, 1]
        b_scale = b_scale.expand(-1, -1, block_shape_n, block_shape_k)  # Expand dimensions
        b_scale = b_scale.permute(0, 2, 1, 3)  # Permute to [scale_n, block_shape_n, scale_k, block_shape_k]
        b_scale = b_scale.reshape(scale_n * block_shape_n, scale_k * block_shape_k)  # Flatten
        b_scale = b_scale[:n, :k]  # Slice to match [n, k]

        # Explicit type conversion
        a_fp32 = a.to(torch.float32) * a_scale
        b_fp32 = b.to(torch.float32) * b_scale

        # Perform matrix multiplication
        result = torch.matmul(a_fp32, b_fp32.T)  # FP32 matrix multiplication
        c.copy_(result.to(torch.bfloat16))  # Convert result to BF16 and copy to pre-allocated memory

    return c
python eval.py leaderboard /tmp/tmpkpihpgfc