import torch
from task import input_t, output_t

def custom_kernel(data: input_t) -> output_t:
    """
    PyTorch reference implementation of FP8 GEMM with block-wise late dequantization.
    This version iterates through K-blocks to apply scales correctly before GEMM.
    """
    a, b, a_scale, b_scale, c_out = data  # c_out is pre-allocated bf16 output

    # Get dimensions
    m = a.shape[0]  # M
    n = b.shape[0]  # N (number of rows in b, becomes columns in c_out)
    k = a.shape[1]  # K

    block_shape_k = 128
    # This is the N-block size for interpreting b_scale elements
    block_shape_n_for_b_scale_interpretation = 128 

    if k % block_shape_k != 0:
        raise ValueError(f"K dimension ({k}) must be divisible by block_shape_k ({block_shape_k})")

    num_k_blocks = k // block_shape_k

    accum_dtype = torch.float32
    c_accum = torch.zeros((m, n), dtype=accum_dtype, device=a.device)

    a = a.contiguous()
    b = b.contiguous()
    a_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()

    for k_block_idx in range(num_k_blocks):
        k_start = k_block_idx * block_shape_k
        k_end = (k_block_idx + 1) * block_shape_k

        # Process matrix 'a' and its scales
        a_k_slice = a[:, k_start:k_end] # M x block_shape_k
        current_a_scale_for_k_block = a_scale[:, k_block_idx].unsqueeze(1)
        a_dequant_block = a_k_slice.to(accum_dtype) * current_a_scale_for_k_block

        # Process matrix 'b' and its scales
        b_k_slice = b[:, k_start:k_end] # N x block_shape_k
        
        if b_scale.shape[0] == 0:
            # This implies N < 128 and b_scale strictly adheres to N // 128 = 0.
            # Assume scale 1.0 for all N rows.
            b_dequant_block = b_k_slice.to(accum_dtype)
        else:
            # b_scale.shape[0] > 0, so there's at least one scale factor row in b_scale.
            # scales_for_this_k_from_b_scale has shape (b_scale.shape[0],)
            scales_for_this_k_from_b_scale = b_scale[:, k_block_idx]
            num_n_scale_factors_provided = b_scale.shape[0]

            # Initialize final_b_row_scales_flat to be of length N with ones (default scale).
            final_b_row_scales_flat = torch.ones(n, dtype=accum_dtype, device=b.device)

            # Iterate through the N-block scales provided in b_scale
            for i in range(num_n_scale_factors_provided):
                # The i-th scale from b_scale applies to a conceptual 128-row block in B
                conceptual_n_block_start = i * block_shape_n_for_b_scale_interpretation
                conceptual_n_block_end   = (i + 1) * block_shape_n_for_b_scale_interpretation
                
                # Determine the actual rows in the current b_k_slice (of N rows) this scale applies to
                apply_to_n_start = conceptual_n_block_start
                apply_to_n_end   = min(conceptual_n_block_end, n) # Clamp to actual N dimension

                if apply_to_n_start < apply_to_n_end: # If this conceptual block overlaps with actual N rows
                    final_b_row_scales_flat[apply_to_n_start:apply_to_n_end] = scales_for_this_k_from_b_scale[i]
            
            # Sanity check (optional, but good for debugging)
            if final_b_row_scales_flat.shape[0] != n:
                 raise AssertionError(
                    f"Internal logic error in b_scale processing after loop: shape mismatch. Expected {n}, got {final_b_row_scales_flat.shape[0]}. "
                    f"N={n}, b_scale.shape={b_scale.shape}, k_block_idx={k_block_idx}"
                )

            final_b_row_scales = final_b_row_scales_flat.unsqueeze(1) # Shape N x 1
            b_dequant_block = b_k_slice.to(accum_dtype) * final_b_row_scales

        # Perform GEMM for the current dequantized blocks and accumulate
        c_accum += a_dequant_block @ b_dequant_block.T

    # Copy the accumulated result to the pre-allocated output tensor
    c_out[...] = c_accum.to(c_out.dtype)

    return c_out
python eval.py test /tmp/tmppb0buqsm