k3_dev/mma_test.py

"""Validate the m16n8k16 bf16 MMA fragment layout against a torch reference."""
import ctypes
import torch

try:
    from cuda.bindings import driver as cu
    from cuda.bindings import nvrtc
except ImportError:
    from cuda import cuda as cu  # type: ignore
    from cuda import nvrtc  # type: ignore

torch.manual_seed(0)
DEV = "cuda"


def ck(res):
    err = res[0]
    ok = (err.value == 0) if hasattr(err, "value") else (int(err) == 0)
    if not ok:
        raise RuntimeError(f"CUDA error: {res}")
    return res[1:] if len(res) > 2 else (res[1] if len(res) == 2 else None)


def main():
    torch.zeros(1, device=DEV)
    ck(cu.cuInit(0))
    src = open("mma_test.cu", "rb").read()
    prog = ck(nvrtc.nvrtcCreateProgram(src, b"mma_test.cu", 0, [], []))
    opts = [b"--gpu-architecture=compute_86", b"--std=c++17", b"-default-device"]
    res = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)
    lsz = ck(nvrtc.nvrtcGetProgramLogSize(prog))
    log = b" " * lsz
    nvrtc.nvrtcGetProgramLog(prog, log)
    t = log.decode(errors="replace").strip()
    if t and t != "\x00":
        print("nvrtc log:", t[:1500])
    err = res[0]
    if (err.value != 0) if hasattr(err, "value") else (int(err) != 0):
        raise RuntimeError("compile failed")
    psz = ck(nvrtc.nvrtcGetPTXSize(prog))
    ptx = b" " * psz
    nvrtc.nvrtcGetPTX(prog, ptx)
    mod = ck(cu.cuModuleLoadData(ptx))
    fn = ck(cu.cuModuleGetFunction(mod, b"mma_test"))

    # bf16-round inputs so the reference matches what the MMA actually sees
    A = (torch.randn(16, 16, device=DEV) * 0.5).to(torch.bfloat16).float()
    B = (torch.randn(16, 8, device=DEV) * 0.5).to(torch.bfloat16).float()
    D = torch.zeros(16, 8, device=DEV, dtype=torch.float32)
    ref = (A @ B)  # A[16,16] @ B[16,8]

    Ac, Bc = A.contiguous(), B.contiguous()
    ptab = [Ac.data_ptr(), Bc.data_ptr(), D.data_ptr()]
    hold = [ctypes.c_void_p(p) for p in ptab]
    arr = (ctypes.c_void_p * 3)(*[ctypes.cast(ctypes.pointer(x), ctypes.c_void_p) for x in hold])
    stream = torch.cuda.current_stream().cuda_stream
    ck(cu.cuLaunchKernel(fn, 1, 1, 1, 32, 1, 1, 0, stream, ctypes.addressof(arr), 0))
    torch.cuda.synchronize()

    diff = (D - ref).abs().max().item()
    rel = diff / (ref.abs().max().item() + 1e-9)
    print(f"[mma] max|diff|={diff:.4e} rel={rel:.4e}")
    print(f"[mma] D[0,:4]   = {D[0, :4].tolist()}")
    print(f"[mma] ref[0,:4] = {ref[0, :4].tolist()}")
    print(f"[mma] D[9,:4]   = {D[9, :4].tolist()}")
    print(f"[mma] ref[9,:4] = {ref[9, :4].tolist()}")
    print("[mma] LAYOUT OK" if rel < 1e-2 else "[mma] LAYOUT WRONG")


if __name__ == "__main__":
    main()