import torch
import triton
import triton.language as tl

@triton.jit
def custom_kernel_triton(
    a_ptr, b_ptr, a_scale_ptr, b_scale_ptr, c_ptr,
    m, n, k,
    stride_am, stride_ak,
    stride_bn, stride_bk,
    stride_asm, stride_ask,
    stride_bsn, stride_bsk,
    stride_cm, stride_cn,
    block_size: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    off_m = pid_m * block_size + tl.arange(0, block_size)
    off_n = pid_n * block_size + tl.arange(0, block_size)
    
    accumulator = tl.zeros((block_size, block_size), dtype=tl.float32)
    
    for block_k in range(0, k, block_size):
        off_k = block_k + tl.arange(0, block_size)
        
        a_ptrs = a_ptr + (off_m[:, None] * stride_am) + (off_k[None, :] * stride_ak)
        a_mask = (off_m[:, None] < m) & (off_k[None, :] < k)
        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
        
        b_ptrs = b_ptr + (off_n[None, :] * stride_bn) + (off_k[:, None] * stride_bk)
        b_mask = (off_n[None, :] < n) & (off_k[:, None] < k)
        b = tl.load(b_ptrs, mask=b_mask, other=0.0)
        
        a_scale_off = (off_m[:, None] * stride_asm) + (off_k[None, :] // block_size * stride_ask)
        a_scale = tl.load(a_scale_ptr + a_scale_off, mask=a_mask, other=1.0)
        a = a.to(tl.float32) * a_scale
        
        b_scale_off = ((off_n[None, :] // block_size) * stride_bsn + (off_k[:, None] // block_size) * stride_bsk)
        b_scale = tl.load(b_scale_ptr + b_scale_off, mask=b_mask, other=1.0)
        b = b.to(tl.float32) * b_scale
        
        accumulator += tl.dot(a, b, allow_tf32=False)
    
    c_ptrs = c_ptr + (off_m[:, None] * stride_cm) + (off_n[None, :] * stride_cn)
    c_mask = (off_m[:, None] < m) & (off_n[None, :] < n)
    c = accumulator.to(tl.bfloat16)
    tl.store(c_ptrs, c, mask=c_mask)

def custom_kernel(data: tuple) -> torch.Tensor:
    a, b, a_scale, b_scale, c = data
    m, k = a.shape
    n = b.shape[0]
    block_size = 128
    
    grid = (triton.cdiv(m, block_size), triton.cdiv(n, block_size))
    
    custom_kernel_triton[grid](
        a, b, a_scale, b_scale, c,
        m, n, k,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        a_scale.stride(0), a_scale.stride(1),
        b_scale.stride(0), b_scale.stride(1),
        c.stride(0), c.stride(1),
        block_size
    )
    return c

python eval.py benchmark /tmp/tmpeuvgkqr5