w3_dev/w4gemm.py

"""Bench the custom int4 W4A16 GEMM (M=16,K=2560,N=10240,g32) vs roofline + torch.
Achieved GB/s tells us if a bespoke GEMM can play with Marlin (~78% roofline)."""
import ctypes, time, torch
try:
    from cuda.bindings import driver as cu, nvrtc
except ImportError:
    from cuda import cuda as cu, nvrtc  # type: ignore
torch.manual_seed(0); DEV="cuda"
M,K,N,G=16,2560,10240,32
def ck(r):
    e=r[0]; ok=(e.value==0) if hasattr(e,'value') else (int(e)==0)
    if not ok: raise RuntimeError(str(r))
    return r[1:] if len(r)>2 else (r[1] if len(r)==2 else None)
def compile():
    src=open("w4gemm.cu","rb").read(); p=ck(nvrtc.nvrtcCreateProgram(src,b"w4gemm.cu",0,[],[]))
    o=[b"--gpu-architecture=compute_86",b"--std=c++17",b"-default-device"]
    r=nvrtc.nvrtcCompileProgram(p,len(o),o); l=ck(nvrtc.nvrtcGetProgramLogSize(p)); b=b" "*l
    nvrtc.nvrtcGetProgramLog(p,b); t=b.decode(errors="replace").strip()
    if t and t!="\x00": print("[nvrtc]",t[:1200])
    e=r[0]
    if (e.value!=0) if hasattr(e,'value') else (int(e)!=0): raise RuntimeError("compile fail")
    s=ck(nvrtc.nvrtcGetPTXSize(p)); ptx=b" "*s; nvrtc.nvrtcGetPTX(p,ptx); return ptx
def main():
    torch.zeros(1,device=DEV); ck(cu.cuInit(0))
    x=(torch.randn(M,K,device=DEV)*0.3).to(torch.bfloat16)
    # int4 weights [N,K] symmetric, group-32 scale
    qint=torch.randint(-8,8,(N,K),device=DEV,dtype=torch.int32)
    scale=(torch.rand(N,K//G,device=DEV)*0.02+0.005).to(torch.bfloat16)
    # pack int4: 8 per u32 along K
    qb=(qint+8).to(torch.int32).reshape(N,K//8,8)
    packed=torch.zeros(N,K//8,dtype=torch.int32,device=DEV)
    for i in range(8): packed|=(qb[:,:,i]<<(4*i))
    packed=packed.to(torch.int32)
    # reference: dequant W -> bf16 matmul
    wdq=(qint.float().reshape(N,K//G,G)*scale.float().unsqueeze(-1)).reshape(N,K).to(torch.bfloat16)
    ref=(x.float()@wdq.float().T).to(torch.bfloat16)
    y=torch.zeros(M,N,dtype=torch.bfloat16,device=DEV)
    mod=ck(cu.cuModuleLoadData(compile())); fn=ck(cu.cuModuleGetFunction(mod,b"w4gemm"))
    ptab=torch.zeros(8,dtype=torch.int64,device=DEV); host=torch.zeros(8,dtype=torch.int64)
    host[0]=x.data_ptr();host[1]=packed.data_ptr();host[2]=scale.data_ptr();host[3]=y.data_ptr()
    ptab.copy_(host)
    n_sm=torch.cuda.get_device_properties(0).multi_processor_count
    grid=n_sm*2
    smem=M*K*2
    a0=ctypes.c_void_p(ptab.data_ptr()); a1=ctypes.c_int(K); a2=ctypes.c_int(N); a3=ctypes.c_int(G)
    arr=(ctypes.c_void_p*4)(*[ctypes.cast(ctypes.pointer(z),ctypes.c_void_p) for z in (a0,a1,a2,a3)])
    ck(cu.cuFuncSetAttribute(fn, cu.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem))
    def launch():
        st=torch.cuda.current_stream().cuda_stream
        ck(cu.cuLaunchKernel(fn,grid,1,1,256,1,1,smem,st,ctypes.addressof(arr),0))
    launch(); torch.cuda.synchronize()
    diff=(y.float()-ref.float()).abs().max().item(); rel=diff/(ref.abs().max().item()+1e-9)
    print(f"[w4] max|diff|={diff:.4e} rel={rel:.4e} {'OK' if rel<0.05 else 'WRONG'}")
    for _ in range(20): launch()
    torch.cuda.synchronize(); t0=time.perf_counter()
    for _ in range(200): launch()
    torch.cuda.synchronize(); us=(time.perf_counter()-t0)/200*1e6
    wbytes=N*(K//2) + N*(K//G)*2   # int4 weights + bf16 scales
    bw=wbytes/(us*1e-6)/1e9
    print(f"[w4] {us:.2f} us/GEMM  achieved {bw:.0f} GB/s ({100*bw/600:.0f}% of 600 roofline; Marlin~78%)")
    print(f"[w4] weight bytes {wbytes/1e6:.1f}MB  roofline {wbytes/600e9*1e6:.2f}us @600GB/s")
if __name__=="__main__": main()