k1_dev/mega_runner.py

"""NVRTC compile + launch + verify + bench for the K1 drafter megakernel (v1).

Adds: per-stage clock64 table, barrier microbench, stream queried at launch
time (so CUDA-graph capture works).
"""

from __future__ import annotations

import ctypes
import time

import torch

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

STAGE_NAMES = [
    "S0 embed+preproj",
    "L0 qrows", "L0 qnorm", "L0 attn", "L0 comb+oproj", "L0 mlp1", "L0 mlp2",
    "L1 qrows", "L1 qnorm", "L1 attn", "L1 comb+oproj", "L1 mlp1", "L1 mlp2",
    "L2 qrows", "L2 qnorm", "L2 attn", "L2 comb+oproj", "L2 mlp1", "L2 mlp2",
    "L3 qrows", "L3 qnorm", "L3 attn", "L3 comb+oproj", "L3 mlp1", "L3 mlp2",
    "S6 fnorm+cent", "S7 topk||postproj", "S8 logits", "S9 argmax",
]


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 compile_kernel(src_path: str, arch: str = "compute_86") -> bytes:
    import os as _os
    src = open(src_path, "rb").read()
    prog = _ck(nvrtc.nvrtcCreateProgram(src, b"megakernel.cu", 0, [], []))
    opts = [f"--gpu-architecture={arch}".encode(), b"--std=c++17", b"-default-device"]
    bns = _os.environ.get("BARSPIN_NS")
    if bns is not None:
        opts.append(f"-DBARSPIN_NS={bns}".encode())
    res = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)
    log_size = _ck(nvrtc.nvrtcGetProgramLogSize(prog))
    log = b" " * log_size
    nvrtc.nvrtcGetProgramLog(prog, log)
    logtxt = log.decode(errors="replace").strip()
    if logtxt and logtxt != "\x00":
        print(f"[mega] nvrtc log:\n{logtxt}")
    err = res[0]
    bad = (err.value != 0) if hasattr(err, "value") else (int(err) != 0)
    if bad:
        raise RuntimeError("NVRTC compilation failed")
    ptx_size = _ck(nvrtc.nvrtcGetPTXSize(prog))
    ptx = b" " * ptx_size
    nvrtc.nvrtcGetPTX(prog, ptx)
    return ptx


def _make_args(ptab, nblk, a, b, c):
    arg_ptab = ctypes.c_void_p(ptab.data_ptr())
    arg1 = ctypes.c_int(nblk)
    arg2 = ctypes.c_int(a)
    arg3 = ctypes.c_int(b)
    arg4 = ctypes.c_int(c)
    holders = [arg_ptab, arg1, arg2, arg3, arg4]
    arr = (ctypes.c_void_p * 5)(
        *[ctypes.cast(ctypes.pointer(x), ctypes.c_void_p) for x in holders]
    )
    return holders, arr


def run_mega(w, kv, cs_caches, first, hidden0, pos, seq, kcent, ref_tokens,
             ref_hidden, k_spec=7):
    dev = torch.device("cuda")
    _ck(cu.cuInit(0))

    ptx = compile_kernel("megakernel.cu")
    module = _ck(cu.cuModuleLoadData(ptx))
    func = _ck(cu.cuModuleGetFunction(module, b"drafter_megakernel"))
    func_bar = _ck(cu.cuModuleGetFunction(module, b"bench_barrier"))

    max_per_sm = _ck(cu.cuOccupancyMaxActiveBlocksPerMultiprocessor(func, 256, 0))
    n_sm = torch.cuda.get_device_properties(0).multi_processor_count
    nblk = min(n_sm, max_per_sm * n_sm)
    import os as _os
    if _os.environ.get("MEGA_NBLK"):
        nblk = min(int(_os.environ["MEGA_NBLK"]), max_per_sm * n_sm)
    print(f"[mega] SMs={n_sm} maxBlocksPerSM={max_per_sm} -> launching {nblk} blocks "
          f"(BARSPIN_NS={_os.environ.get('BARSPIN_NS','64')})")
    if nblk < 8:
        raise RuntimeError("not enough co-resident blocks for grid barrier")

    scratch_f = torch.zeros(70000, dtype=torch.float32, device=dev)
    scratch_i = torch.zeros(32768, dtype=torch.int32, device=dev)
    out_tokens = torch.zeros(k_spec, dtype=torch.int64, device=dev)
    out_bb = torch.zeros(2560, dtype=torch.bfloat16, device=dev)
    barrier = torch.zeros(2, dtype=torch.int32, device=dev)
    clocks = torch.zeros(256, dtype=torch.int64, device=dev)

    ptab = torch.zeros(128, dtype=torch.int64, 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]
        base = 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[base + off] = lw[key].data_ptr()
        host[base + 11] = kv[li][0].data_ptr()
        host[base + 12] = kv[li][1].data_ptr()
    host[60] = first.data_ptr()
    host[61] = hidden0.data_ptr()
    host[62] = out_tokens.data_ptr()
    host[63] = scratch_f.data_ptr()
    host[64] = scratch_i.data_ptr()
    host[65] = out_bb.data_ptr()
    host[66] = barrier.data_ptr()
    host[67] = clocks.data_ptr()
    for li in range(4):
        host[108 + li] = cs_caches[li].data_ptr()
    host[112] = 0  # f32 caches in the microbench
    ptab.copy_(host)

    holders, arg_array = _make_args(ptab, nblk, pos, seq, kcent)

    def launch():
        stream = torch.cuda.current_stream().cuda_stream
        _ck(cu.cuLaunchKernel(func, nblk, 1, 1, 256, 1, 1, 0, stream,
                              ctypes.addressof(arg_array), 0))

    # correctness
    launch()
    torch.cuda.synchronize()
    mt = out_tokens.tolist()
    rt = ref_tokens.tolist()
    match = sum(int(a == b) for a, b in zip(mt, rt))
    hdiff = (out_bb.float() - ref_hidden.float()).abs().max().item()
    print(f"[mega] tokens ref =[{', '.join(map(str, rt))}]")
    print(f"[mega] tokens mega=[{', '.join(map(str, mt))}]  match={match}/{k_spec}")
    print(f"[mega] backbone hidden max|diff|={hdiff:.4e}")

    # stage clock table (mean over iters 1..k_spec-1, skip warm iter 0)
    rate_khz = _ck(cu.cuDeviceGetAttribute(
        cu.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_CLOCK_RATE, 0))
    cl = clocks.cpu().tolist()
    n_st = len(STAGE_NAMES)
    sums = [0.0] * n_st
    cnt = 0
    for it in range(1, k_spec):
        base = it * 32
        if cl[base + n_st] == 0:
            continue
        cnt += 1
        for s in range(n_st):
            sums[s] += (cl[base + s + 1] - cl[base + s]) / (rate_khz / 1000.0)
    if cnt:
        print(f"[mega] per-stage us (mean over {cnt} iters, clock={rate_khz}kHz):")
        tot = 0.0
        for s in range(n_st):
            v = sums[s] / cnt
            tot += v
            print(f"[mega]   {STAGE_NAMES[s]:<20s} {v:8.2f} us")
        print(f"[mega]   {'TOTAL':<20s} {tot:8.2f} us/iter")

    # barrier microbench: 1000 syncs
    holders_b, arr_b = _make_args(ptab, nblk, 1000, 0, 0)
    stream = torch.cuda.current_stream().cuda_stream
    _ck(cu.cuLaunchKernel(func_bar, nblk, 1, 1, 256, 1, 1, 0, stream,
                          ctypes.addressof(arr_b), 0))
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(5):
        _ck(cu.cuLaunchKernel(func_bar, nblk, 1, 1, 256, 1, 1, 0, stream,
                              ctypes.addressof(arr_b), 0))
    torch.cuda.synchronize()
    per_sync = (time.perf_counter() - t0) / 5 / 1000 * 1e6
    print(f"[mega] grid_sync cost: {per_sync:.3f} us/sync ({nblk} blocks)")

    # bench
    for _ in range(10):
        launch()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    n_iter = 50
    for _ in range(n_iter):
        launch()
    torch.cuda.synchronize()
    t = (time.perf_counter() - t0) / n_iter * 1e3
    print(f"[mega] mega      : {t:8.3f} ms/loop  {t/k_spec:7.3f} ms/iter")

    # graph capture (stream now resolved inside launch())
    try:
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            launch()
        g.replay()
        torch.cuda.synchronize()
        mt2 = out_tokens.tolist()
        for _ in range(10):
            g.replay()
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(n_iter):
            g.replay()
        torch.cuda.synchronize()
        tg = (time.perf_counter() - t0) / n_iter * 1e3
        print(f"[mega] mega_graph: {tg:8.3f} ms/loop  {tg/k_spec:7.3f} ms/iter "
              f"(replay tokens match={sum(int(a==b) for a,b in zip(mt2, rt))}/{k_spec})")
    except Exception as exc:
        print(f"[mega] graph capture failed (non-fatal): {exc!r}")

    return mt, t