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
    
    # Extract dimensions
    m, k = a.shape
    n = b.shape[0]
    block_size = 128

    # Convert to FP32 first to maintain precision
    with torch.amp.autocast('cuda', enabled=False):
        a_fp32 = a.to(torch.float32)
        b_fp32 = b.to(torch.float32)
        
        # Process scales in blocks of 128
        result = torch.zeros((m, n), dtype=torch.float32, device=a.device)
        
        for k_block in range(0, k, block_size):
            k_end = min(k_block + block_size, k)
            k_idx = k_block // block_size
            
            # Get current blocks
            a_block = a_fp32[:, k_block:k_end]
            b_block = b_fp32[:, k_block:k_end]
            
            # Apply scales for current block
            # Handle a_scale: [m, k//128]
            a_scale_block = a_scale[:, k_idx].unsqueeze(1)  # [m, 1]
            
            # Handle b_scale: [n//128, k//128]
            # First expand each scale to cover its 128-element block
            b_scale_temp = b_scale[:, k_idx].unsqueeze(1)  # [n//128, 1]
            b_scale_block = b_scale_temp.repeat_interleave(block_size, dim=0)  # [n, 1]
            
            # Expand scales to match block dimensions
            a_scale_expanded = a_scale_block.expand(-1, k_end - k_block)  # [m, block_size]
            b_scale_expanded = b_scale_block.expand(-1, k_end - k_block)  # [n, block_size]
            
            # Apply scales and compute block multiplication
            a_scaled = a_block * a_scale_expanded
            b_scaled = b_block * b_scale_expanded[:n]  # Ensure we don't exceed n
            
            # Accumulate result
            result.addmm_(a_scaled, b_scaled.t())
        
        # Convert final result to BF16
        c.copy_(result.to(torch.bfloat16))
    
    return c
python eval.py leaderboard /tmp/tmpc41pnpu6