k1_dev/tree_bench.py

"""Validate + time the megakernel drafting 2 branches CONCURRENTLY (grid.y=2).
Each group runs an independent K_SPEC chain from its own root, shared hidden0.
Checks each branch == the reference linear chain from that root, and times
2-concurrent-branches vs 1 linear chain. Reuses k1_dev's model + reference."""
import ctypes, time, torch
import k1_dev as K
try:
    from cuda.bindings import driver as cu, nvrtc
except ImportError:
    from cuda import cuda as cu, nvrtc  # type: ignore

DEV="cuda"; HID=K.HID; BB=K.BB; K_SPEC=K.K_SPEC; KCENT=K.KCENT; SEQ=K.SEQ
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_ptx():
    src=open("megakernel.cu","rb").read()
    p=ck(nvrtc.nvrtcCreateProgram(src,b"megakernel.cu",0,[],[]))
    opts=[b"--gpu-architecture=compute_86",b"--std=c++17",b"-default-device"]
    res=nvrtc.nvrtcCompileProgram(p,len(opts),opts)
    l=ck(nvrtc.nvrtcGetProgramLogSize(p)); buf=b" "*l; nvrtc.nvrtcGetProgramLog(p,buf)
    t=buf.decode(errors="replace").strip()
    if t and t!="\x00": print("[nvrtc]",t[:800])
    e=res[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))
    w=K.make_weights(); K.add_oproj(w); kv=K.make_kv(); cs=K.make_cs_caches()
    pos=SEQ-1; hidden0=(torch.randn(1,BB,device=DEV)*0.5).to(K.DT)
    # two distinct roots (simulating a top-2 fork at position 0)
    rootA=torch.randint(0,K.VOCAB,(1,),device=DEV); rootB=torch.randint(0,K.VOCAB,(1,),device=DEV)
    refA,_=K.run_loop(w,kv,cs,rootA,hidden0,pos)
    refB,_=K.run_loop(w,kv,cs,rootB,hidden0,pos)
    print(f"[tree] refA={refA.tolist()}")
    print(f"[tree] refB={refB.tolist()}")

    mod=ck(cu.cuModuleLoadData(compile_ptx())); fn=ck(cu.cuModuleGetFunction(mod,b"drafter_megakernel"))
    n_sm=torch.cuda.get_device_properties(0).multi_processor_count
    per_group=40
    first=torch.tensor([int(rootA),int(rootB)],dtype=torch.int64,device=DEV)
    out=torch.zeros(2*K_SPEC,dtype=torch.int64,device=DEV)
    Sf=torch.zeros(2*62000,dtype=torch.float32,device=DEV)
    SIi=torch.zeros(2*20000,dtype=torch.int32,device=DEV)
    outbb=torch.zeros(2*BB,dtype=torch.bfloat16,device=DEV)
    bar=torch.zeros(4,dtype=torch.int32,device=DEV)
    host=torch.zeros(128,dtype=torch.int64)
    host[0]=w["embed"].data_ptr();host[1]=w["pre_proj"].data_ptr();host[2]=w["post_proj"].data_ptr()
    host[3]=w["final_norm"].data_ptr();host[4]=w["cent_w"].data_ptr();host[5]=w["lm_head"].data_ptr();host[6]=w["token_ordering"].data_ptr()
    for li in range(4):
        lw=w["layers"][li]; b=8+li*13
        for off,key in enumerate(["in_norm","q_proj","q_norm","o_proj","post_attn_norm","pre_ffn_norm","gate","up","down","post_ffn_norm","layer_scalar"]):
            host[b+off]=lw[key].data_ptr()
        host[b+11]=kv[li][0].data_ptr(); host[b+12]=kv[li][1].data_ptr()
    host[60]=first.data_ptr();host[61]=hidden0.data_ptr();host[62]=out.data_ptr()
    host[63]=Sf.data_ptr();host[64]=SIi.data_ptr();host[65]=outbb.data_ptr();host[66]=bar.data_ptr();host[67]=0
    for li in range(4): host[108+li]=cs[li].data_ptr()
    host[112]=0
    ptab=torch.zeros(128,dtype=torch.int64,device=DEV); ptab.copy_(host)

    def args(nblk):
        a0=ctypes.c_void_p(ptab.data_ptr()); a1=ctypes.c_int(nblk); a2=ctypes.c_int(pos); a3=ctypes.c_int(SEQ); a4=ctypes.c_int(KCENT)
        h=[a0,a1,a2,a3,a4]
        return (ctypes.c_void_p*5)(*[ctypes.cast(ctypes.pointer(x),ctypes.c_void_p) for x in h])
    arr=args(per_group)
    def launch_tree():
        st=torch.cuda.current_stream().cuda_stream
        ck(cu.cuLaunchKernel(fn, per_group,2,1, 256,1,1, 0, st, ctypes.addressof(arr),0))
    launch_tree(); torch.cuda.synchronize()
    mA=out[:K_SPEC].tolist(); mB=out[K_SPEC:].tolist()
    okA=sum(int(a==b) for a,b in zip(mA,refA.tolist())); okB=sum(int(a==b) for a,b in zip(mB,refB.tolist()))
    print(f"[tree] branchA mega={mA} match={okA}/{K_SPEC}")
    print(f"[tree] branchB mega={mB} match={okB}/{K_SPEC}")

    for _ in range(10): launch_tree()
    torch.cuda.synchronize(); t0=time.perf_counter()
    for _ in range(50): launch_tree()
    torch.cuda.synchronize(); t_tree=(time.perf_counter()-t0)/50*1e3

    # linear baseline: 1 chain on 80 blocks (grid.y=1)
    arr1=args(80)
    def launch_lin():
        st=torch.cuda.current_stream().cuda_stream
        ck(cu.cuLaunchKernel(fn, 80,1,1, 256,1,1, 0, st, ctypes.addressof(arr1),0))
    launch_lin(); torch.cuda.synchronize()
    for _ in range(10): launch_lin()
    torch.cuda.synchronize(); t0=time.perf_counter()
    for _ in range(50): launch_lin()
    torch.cuda.synchronize(); t_lin=(time.perf_counter()-t0)/50*1e3
    print(f"[tree] 2-branch CONCURRENT = {t_tree:.3f} ms | 1-chain linear(80blk) = {t_lin:.3f} ms")
    print(f"[tree] overhead for 2x the tree: {100*(t_tree-t_lin)/t_lin:+.1f}% (onegraph would be +100%)")

if __name__=="__main__": main()