k3_dev/k3_attn_patch.py

"""fableous / ultra-kernels K3: specialized verify-window decode attention.

Env-gated (K3_ATTN=1). Hooks TritonAttentionImpl.forward and diverts
decode-shaped calls (batch 1, num_actual_tokens == 8 == K+1 spec tokens) to a
two-phase flash-decoding CUDA kernel specialized for the fixed verify shape:
GQA 8q/2kv, hd 256 (sliding W=512) / hd 512 (full), causal tail, paged bf16
KV. All other calls (prefill, odd shapes, quantized cache) fall through to
the stock unified_attention path.

The KV cache write happens in do_kv_cache_update (separate op) BEFORE this
forward — the hook replaces only the attention computation. Numerics:
f32 accumulation, exact softmax, bf16 output — same contract as the triton
kernel (target greedy tokens must hold).

Requires megakernel-style PTX shipping: k3_attn_vllm.ptx (QG4/NSUB4 for
hd256 and QG2/NSUB2 for hd512 compiled into ONE ptx via two entry... no —
two PTX files: k3_attn_hd256.ptx / k3_attn_hd512.ptx) or NVRTC fallback.

Fail policy: K3_ATTN_REQUIRE=1 raises on setup/launch errors; otherwise
falls back to stock forward permanently (logged).
"""

from __future__ import annotations

import ctypes
import importlib.abc
import importlib.util
import os
import pathlib
import sys
from typing import Any

K3 = os.environ.get("K3_ATTN", "0") == "1"
REQUIRE = os.environ.get("K3_ATTN_REQUIRE", "0") == "1"
QLEN = 8
NCHUNK_S = int(os.environ.get("K3_NCHUNK_S", "4"))
NCHUNK_F = int(os.environ.get("K3_NCHUNK_F", "8"))
NSUB_S, NSUB_F = 4, 2

TARGET_MODULE = "vllm.v1.attention.backends.triton_attn"


def _log(msg: str) -> None:
    print(f"[k3] {msg}", file=sys.stderr, flush=True)


class _K3Runtime:
    """Compiled modules + scratch shared across layers."""

    _inst = None

    @classmethod
    def get(cls):
        if cls._inst is None:
            cls._inst = cls()
        return cls._inst

    def __init__(self) -> None:
        import torch

        try:
            from cuda.bindings import driver as cu
        except ImportError:
            from cuda import cuda as cu  # type: ignore
        self.cu = cu
        self.torch = torch

        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)

        self.ck = ck
        ck(cu.cuInit(0))

        pkg = pathlib.Path(__file__).resolve().parent
        self.funcs = {}
        for tag in ("hd256", "hd512"):
            ptx_file = pkg / f"k3_attn_{tag}.ptx"
            if ptx_file.exists():
                ptx = ptx_file.read_bytes()
                _log(f"{tag}: shipped PTX ({len(ptx)} bytes)")
            else:
                defines = ("-DQG=4", "-DNSUB=4") if tag == "hd256" else ("-DQG=2", "-DNSUB=2")
                ptx = self._nvrtc(str(pkg / "k3_attn.cu"), defines)
                _log(f"{tag}: NVRTC-compiled at serve time")
            mod = ck(cu.cuModuleLoadData(ptx))
            self.funcs[tag] = (
                ck(cu.cuModuleGetFunction(mod, b"attn_partial")),
                ck(cu.cuModuleGetFunction(mod, b"attn_combine")),
            )

        dev = torch.device("cuda")
        # PART scratch sized for the larger geometry (hd512, nch2 = NCHUNK_F*NSUB_F)
        n_part = max(
            8 * NCHUNK_S * NSUB_S * QLEN * (256 + 2),
            8 * NCHUNK_F * NSUB_F * QLEN * (512 + 2),
        )
        self.part = torch.zeros(n_part, dtype=torch.float32, device=dev)
        self.layer_state: dict[int, dict] = {}

    def _nvrtc(self, src_path: str, defines) -> bytes:
        import glob
        import site

        roots = set(site.getsitepackages())
        for p in sys.path:
            if p.endswith("site-packages"):
                roots.add(p)
        for root in roots:
            for so in glob.glob(os.path.join(root, "nvidia", "cuda_nvrtc", "lib", "libnvrtc.so*")):
                try:
                    ctypes.CDLL(so, mode=ctypes.RTLD_GLOBAL)
                except OSError:
                    pass
        try:
            from cuda.bindings import nvrtc
        except ImportError:
            from cuda import nvrtc  # type: ignore

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

        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)
        err = res[0]
        if (err.value != 0) if hasattr(err, "value") else (int(err) != 0):
            raise RuntimeError(f"NVRTC failed: {log.decode(errors='replace')[:500]}")
        psz = ck(nvrtc.nvrtcGetPTXSize(prog))
        ptx = b" " * psz
        nvrtc.nvrtcGetPTX(prog, ptx)
        return ptx

    def layer_plumbing(self, impl, kv_cache, query, output, block_table) -> dict:
        """Per-layer cached pointers/args. Rebuilt if buffer identities change."""
        torch = self.torch
        key = id(impl)
        st = self.layer_state.get(key)
        sig = (kv_cache.data_ptr(), query.data_ptr(), output.data_ptr(),
               block_table.data_ptr())
        if st is not None and st["sig"] == sig:
            return st

        hd = impl.head_size
        nqh = impl.num_heads
        nkv = impl.num_kv_heads
        if hd not in (256, 512) or nqh != 8 or nkv != 2:
            raise RuntimeError(f"unsupported geometry hd={hd} nqh={nqh} nkv={nkv}")
        if kv_cache.dim() != 5 or kv_cache.shape[1] != 2:
            raise RuntimeError(f"cache shape {tuple(kv_cache.shape)}")
        kc, vc = kv_cache[:, 0], kv_cache[:, 1]
        if kc.stride(-1) != 1:
            raise RuntimeError("non-contiguous head dim")
        bs = kc.shape[1]
        if bs & (bs - 1):
            raise RuntimeError(f"block_size {bs} not pow2")

        sw = impl.sliding_window
        left = sw[0] if isinstance(sw, (tuple, list)) else (sw or -1)
        win = (left + 1) if (left is not None and left >= 0) else 0

        tag = "hd256" if hd == 256 else "hd512"
        nchunk = NCHUNK_S if hd == 256 else NCHUNK_F

        ptab = torch.zeros(16, dtype=torch.int64, device="cuda")
        host = torch.zeros(16, dtype=torch.int64)
        host[0] = query.data_ptr()
        host[1] = kc.data_ptr()
        host[2] = vc.data_ptr()
        host[3] = block_table.data_ptr()
        host[4] = self.part.data_ptr()
        host[5] = output.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, non_blocking=True)

        scale = float(impl.scale)
        hold = [
            ctypes.c_void_p(ptab.data_ptr()),
            ctypes.c_int(nqh), ctypes.c_int(nqh // nkv), ctypes.c_int(hd),
            ctypes.c_int(0),   # seq (mutated per call)
            ctypes.c_int(0),   # lo  (mutated per call)
            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]
        )
        st = {
            "sig": sig, "ptab": ptab, "hold": hold, "arr": arr,
            "tag": tag, "win": win, "nchunk": nchunk, "hd": hd, "nqh": nqh,
            "smem": QLEN * hd * 4,
        }
        self.layer_state[key] = st
        _log(f"layer plumbing built: hd={hd} win={win} bs={bs} tag={tag}")
        return st

    def run(self, st, seq: int) -> None:
        win = st["win"]
        lo = max(0, seq - QLEN - win + 1) if win else 0
        st["hold"][4].value = seq
        st["hold"][5].value = lo
        fp, fc = self.funcs[st["tag"]]
        stream = self.torch.cuda.current_stream().cuda_stream
        addr = ctypes.addressof(st["arr"])
        self.ck(self.cu.cuLaunchKernel(
            fp, st["nqh"], st["nchunk"], 1, 256, 1, 1, st["smem"], stream, addr, 0))
        self.ck(self.cu.cuLaunchKernel(
            fc, st["nqh"], 1, 1, 256, 1, 1, 0, stream, addr, 0))


def _apply_k3_patch(module: Any) -> None:
    impl_cls = module.TritonAttentionImpl
    base_forward = impl_cls.forward
    state = {"failed": False, "active_logged": False}

    def forward(self, layer, query, key, value, kv_cache, attn_metadata, output,
                output_scale=None, output_block_scale=None):
        md = attn_metadata
        if (
            state["failed"]
            or md is None
            or md.num_actual_tokens != QLEN
            or output_scale is not None
            or output_block_scale is not None
        ):
            return base_forward(self, layer, query, key, value, kv_cache,
                                attn_metadata, output, output_scale,
                                output_block_scale)
        try:
            if (
                md.query_start_loc.shape[0] != 2     # batch 1
                or getattr(self, "_is_per_token_head_quant", False)
                or getattr(self, "alibi_slopes", None) is not None
                or getattr(self, "logits_soft_cap", None)
                or getattr(self, "sinks", None) is not None
                or query.dtype != _K3Runtime.get().torch.bfloat16
            ):
                return base_forward(self, layer, query, key, value, kv_cache,
                                    attn_metadata, output, output_scale,
                                    output_block_scale)
            rt = _K3Runtime.get()
            st = rt.layer_plumbing(self, kv_cache, query, output, md.block_table)
            rt.run(st, int(md.max_seq_len))
            if not state["active_logged"]:
                state["active_logged"] = True
                _log(f"ACTIVE (pid {os.getpid()})")
            return output
        except Exception as exc:
            state["failed"] = True
            if REQUIRE:
                raise RuntimeError("K3_ATTN_REQUIRE=1 and hook failed") from exc
            _log(f"DISABLED after error: {exc!r}")
            return base_forward(self, layer, query, key, value, kv_cache,
                                attn_metadata, output, output_scale,
                                output_block_scale)

    impl_cls.forward = forward
    _log(f"patched TritonAttentionImpl.forward (require={REQUIRE}) in pid {os.getpid()}")


class _ChainLoader(importlib.abc.Loader):
    def __init__(self, inner, patch_fn):
        self._inner = inner
        self._patch_fn = patch_fn

    def create_module(self, spec):
        return self._inner.create_module(spec)

    def exec_module(self, module):
        self._inner.exec_module(module)
        self._patch_fn(module)


class _ChainFinder(importlib.abc.MetaPathFinder):
    def __init__(self, target, patch_fn):
        self._target = target
        self._patch_fn = patch_fn
        self._busy = False

    def find_spec(self, fullname, path=None, target=None):
        if fullname != self._target or self._busy:
            return None
        self._busy = True
        try:
            spec = importlib.util.find_spec(fullname)
        finally:
            self._busy = False
        if spec is None or spec.loader is None:
            return None
        spec.loader = _ChainLoader(spec.loader, self._patch_fn)
        return spec


if K3:
    sys.meta_path.insert(0, _ChainFinder(TARGET_MODULE, _apply_k3_patch))
    _log("finder registered")