k3_dev/k3_dev.py

"""fableous / ultra-kernels K3 dev harness.

Verify-window decode attention microbenchmark: batch 1, q_len=8, causal tail,
GQA 8q/2kv, paged bf16 KV. Two variants: sliding (hd 256, window 512) and
full (hd 512). Correctness vs an exact f32 torch reference; timing vs torch
SDPA as a rough anchor (production triton anchor = ~34us/call from kprof).
"""

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:
    from cuda import cuda as cu  # type: ignore
    from cuda import nvrtc  # type: ignore

torch.manual_seed(0)
DEV = "cuda"
DT = torch.bfloat16

QLEN = 8
NQH = 8
NKV = 2
GQA = NQH // NKV
BS = 16
CTX = int(__import__("os").environ.get("CTX", "1024"))
SEQ = CTX + QLEN


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, defines=()) -> bytes:
    src = open(src_path, "rb").read()
    prog = _ck(nvrtc.nvrtcCreateProgram(src, b"k3_attn.cu", 0, [], []))
    opts = [b"--gpu-architecture=compute_86", b"--std=c++17", b"-default-device"]
    opts += [d.encode() for d in defines]
    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(f"[k3] nvrtc log:\n{t[:2000]}")
    err = res[0]
    if (err.value != 0) if hasattr(err, "value") else (int(err) != 0):
        raise RuntimeError("NVRTC failed")
    psz = _ck(nvrtc.nvrtcGetPTXSize(prog))
    ptx = b" " * psz
    nvrtc.nvrtcGetPTX(prog, ptx)
    return ptx


def make_args(ptab, nqh, gqa, hd, seq, lo, nchunk, scale):
    hold = [
        ctypes.c_void_p(ptab.data_ptr()),
        ctypes.c_int(nqh), ctypes.c_int(gqa), ctypes.c_int(hd),
        ctypes.c_int(seq), ctypes.c_int(lo), ctypes.c_int(nchunk),
        ctypes.c_float(scale),
    ]
    arr = (ctypes.c_void_p * len(hold))(
        *[ctypes.cast(ctypes.pointer(x), ctypes.c_void_p) for x in hold]
    )
    return hold, arr


def ref_attention(q, kc, vc, bt, hd, win, scale):
    """Exact f32 reference. q: [QLEN, NQH*hd] bf16; kc/vc: [NB, BS, NKV, hd]."""
    out = torch.empty(QLEN, NQH * hd, device=DEV, dtype=torch.float32)
    # flatten paged cache to [SEQ, NKV, hd]
    flat_k = kc.reshape(-1, NKV, hd)[:SEQ].float()
    flat_v = vc.reshape(-1, NKV, hd)[:SEQ].float()
    for qh in range(NQH):
        kvh = qh // GQA
        for qp in range(QLEN):
            pos = SEQ - QLEN + qp
            lo_q = max(0, pos - win + 1) if win else 0
            ks = flat_k[lo_q:pos + 1, kvh]
            vs = flat_v[lo_q:pos + 1, kvh]
            qv = q[qp, qh * hd:(qh + 1) * hd].float() * scale
            sc = ks @ qv
            p = torch.softmax(sc, dim=0)
            out[qp, qh * hd:(qh + 1) * hd] = p @ vs
    return out


def bench_variant(name, hd, win, func_p, func_c, nchunk, nsub):
    nb = (SEQ + BS - 1) // BS + 2
    kc = (torch.randn(nb, BS, NKV, hd, device=DEV) * 0.5).to(DT)
    vc = (torch.randn(nb, BS, NKV, hd, device=DEV) * 0.5).to(DT)
    bt = torch.arange(nb, dtype=torch.int32, device=DEV)
    q = (torch.randn(QLEN, NQH * hd, device=DEV) * 0.5).to(DT)
    part = torch.zeros(NQH * nchunk * nsub * QLEN * (hd + 2), dtype=torch.float32, device=DEV)
    out = torch.zeros(QLEN, NQH * hd, dtype=DT, device=DEV)
    scale = hd ** -0.5

    ptab = torch.zeros(16, dtype=torch.int64, device=DEV)
    host = torch.zeros(16, dtype=torch.int64)
    host[0] = q.data_ptr()
    host[1] = kc.data_ptr()
    host[2] = vc.data_ptr()
    host[3] = bt.data_ptr()
    host[4] = part.data_ptr()
    host[5] = out.data_ptr()
    host[6] = kc.stride(0)
    host[7] = kc.stride(1)
    host[8] = kc.stride(2)
    host[9] = BS.bit_length() - 1
    host[10] = win
    ptab.copy_(host)

    lo = max(0, SEQ - QLEN - win + 1) if win else 0  # loosest lower bound
    smem = QLEN * hd * 4  # v7: CTA-staged queries
    _, arr = make_args(ptab, NQH, GQA, hd, SEQ, lo, nchunk, scale)

    def launch():
        stream = torch.cuda.current_stream().cuda_stream
        _ck(cu.cuLaunchKernel(func_p, NQH, nchunk, 1, 256, 1, 1, smem, stream,
                              ctypes.addressof(arr), 0))
        _ck(cu.cuLaunchKernel(func_c, NQH, 1, 1, 256, 1, 1, 0, stream,
                              ctypes.addressof(arr), 0))

    launch()
    torch.cuda.synchronize()
    ref = ref_attention(q, kc, vc, bt, hd, win, scale)
    diff = (out.float() - ref).abs().max().item()
    print(f"[k3] {name}: max|diff| vs ref = {diff:.4e}")

    for _ in range(50):
        launch()
    torch.cuda.synchronize()
    # GPU-only timing: capture 20 call-pairs in a CUDA graph, replay
    def time_graph(fn, n_inner=20, n_rep=50):
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            for _ in range(n_inner):
                fn()
        g.replay()
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(n_rep):
            g.replay()
        torch.cuda.synchronize()
        return (time.perf_counter() - t0) / (n_rep * n_inner) * 1e6

    stream_holder = []

    def only_partial():
        stream = torch.cuda.current_stream().cuda_stream
        _ck(cu.cuLaunchKernel(func_p, NQH, nchunk, 1, 256, 1, 1, smem, stream,
                              ctypes.addressof(arr), 0))

    def only_combine():
        stream = torch.cuda.current_stream().cuda_stream
        _ck(cu.cuLaunchKernel(func_c, NQH, 1, 1, 256, 1, 1, 0, stream,
                              ctypes.addressof(arr), 0))

    us = time_graph(launch)
    us_p = time_graph(only_partial)
    us_c = time_graph(only_combine)
    print(f"[k3] {name}: {us:7.2f} us/call GPU (partial {us_p:.2f} + combine {us_c:.2f})")

    # torch SDPA anchor (contiguous KV, same math shape)
    flat_k = kc.reshape(-1, NKV, hd)[:SEQ]
    flat_v = vc.reshape(-1, NKV, hd)[:SEQ]
    kq = flat_k.permute(1, 0, 2).unsqueeze(0).repeat_interleave(GQA, dim=1)
    vq = flat_v.permute(1, 0, 2).unsqueeze(0).repeat_interleave(GQA, dim=1)
    qq = q.view(QLEN, NQH, hd).permute(1, 0, 2).unsqueeze(0)
    mask = torch.zeros(QLEN, SEQ, device=DEV, dtype=DT)
    for qp in range(QLEN):
        pos = SEQ - QLEN + qp
        mask[qp, pos + 1:] = float("-inf")
        if win:
            mask[qp, :max(0, pos - win + 1)] = float("-inf")

    def sdpa():
        return torch.nn.functional.scaled_dot_product_attention(
            qq, kq, vq, attn_mask=mask, scale=scale)

    sdpa()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(200):
        sdpa()
    torch.cuda.synchronize()
    us_sdpa = (time.perf_counter() - t0) / 200 * 1e6
    print(f"[k3] {name}: {us_sdpa:7.2f} us/call torch-SDPA anchor "
          f"(production triton ~34us from kprof)")
    return us


def main():
    print(f"[k3] device={torch.cuda.get_device_name(0)} CTX={CTX} SEQ={SEQ}")
    torch.zeros(1, device=DEV)  # materialize torch's primary CUDA context
    _ck(cu.cuInit(0))
    ptx_s = compile_kernel("k3_attn.cu", ("-DQG=4", "-DNSUB=4"))
    mod_s = _ck(cu.cuModuleLoadData(ptx_s))
    fp_s = _ck(cu.cuModuleGetFunction(mod_s, b"attn_partial"))
    fc_s = _ck(cu.cuModuleGetFunction(mod_s, b"attn_combine"))
    ptx_f = compile_kernel("k3_attn.cu", ("-DQG=2", "-DNSUB=2"))
    mod_f = _ck(cu.cuModuleLoadData(ptx_f))
    fp_f = _ck(cu.cuModuleGetFunction(mod_f, b"attn_partial"))
    fc_f = _ck(cu.cuModuleGetFunction(mod_f, b"attn_combine"))

    us_s = bench_variant("sliding hd256 W512", 256, 512, fp_s, fc_s, nchunk=4, nsub=4)
    us_f = bench_variant("full    hd512     ", 512, 0, fp_f, fc_f, nchunk=8, nsub=2)

    per_step = 30 * us_s + 7 * us_f
    print(f"[k3] projected/step: 30*sliding + 7*full = {per_step/1000:.3f} ms "
          f"(production ~1.2-1.7ms) -> potential saving "
          f"{(1200 - per_step)/1000:.2f}-{(1700 - per_step)/1000:.2f} ms/step")


if __name__ == "__main__":
    main()