import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'

from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
import torch

def fp8_gemm_highprec(a_fp8: torch.Tensor,
                      b_fp8: torch.Tensor,
                      a_scale: torch.Tensor,
                      b_scale: torch.Tensor) -> torch.Tensor:
    """
    FP8→FP32 decoding, FP32 GEMM with per-128 scaling, final cast to BF16.
    """
    assert a_fp8.is_cuda and b_fp8.is_cuda

    M, K = a_fp8.shape
    N    = b_fp8.shape[0]

    # 1) Decode FP8→FP32
    a_f32 = a_fp8.to(torch.float32)
    b_f32 = b_fp8.to(torch.float32)

    # 2) Expand scales and apply in FP32
    #    (all on device)
    a_scale_full = a_scale.repeat_interleave(128, dim=1)[:, :K]
    b_scale_full = (
        b_scale
        .repeat_interleave(128, dim=1)       # (N//128, K//128) → (N//128, K)
        .repeat_interleave(128, dim=0)[:N, :] # → (N, K)
    )

    a_f32 = a_f32 * a_scale_full
    b_f32 = b_f32 * b_scale_full

    # 3) FP32 GEMM
    #    Use torch.matmul which will pick the fast FP32 path
    c_f32 = torch.matmul(a_f32, b_f32.transpose(0, 1))

    # 4) Final cast to BF16
    c_bf16 = c_f32.to(torch.bfloat16)

    return c_bf16

# In your custom_kernel:
def custom_kernel(data: input_t) -> output_t:
    a_fp8, b_fp8, a_scale, b_scale, c = data
    c.copy_(fp8_gemm_highprec(a_fp8, b_fp8, a_scale, b_scale))
    return c
python eval.py leaderboard /tmp/tmpp1l5rj5d