submissions/byteshark/tree-v2-align-debug-v0/tree_v2_ext.py

"""tree-v2 extension module (chiku-inu) — star-tree machinery validated by
tree-dev/sim_tree_identity.py (240-token greedy identity vs sequential oracle).

Ships alongside sitecustomize.py, which imports this at the end. Registers:
  1. runner patch  — duplicated branch positions AFTER slot mapping + step flag
  2. attention patch — star-tree decode attention (paged, graph-safe) replacing
     the unified op on tree verify steps
  3. star_reject() — fused rejection walk on the pinned vLLM layout
     (target rows 0..KW-1; full-accept next = row K; last-position salvage
     next = the sampler bonus; salvage KV relocation via stashed caches)

REQUIRE philosophy: every assumption is checked at boot/first-use and raises —
no silent fallbacks (the tree-v1 lesson).
"""
from __future__ import annotations

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

SPEC_TREE_SPEC = os.environ.get("SPEC_TREE_SPEC", "1") == "1"
TREE_W = int(os.environ.get("SPEC_TREE_WIDTH", "2"))
TREE_REQUIRE = os.environ.get("TREE_REQUIRE", "1") == "1"
STAR_REJECT_FUSED = os.environ.get("STAR_REJECT_FUSED", "1") == "1"
TREE_STATS_EVERY = int(os.environ.get("TREE_STATS_EVERY", "2048"))
TREE_FAIL_ON_PLACEHOLDER_SPEC = (
    os.environ.get("TREE_FAIL_ON_PLACEHOLDER_SPEC", "1") == "1"
)

def _conf_K() -> int:
    import json
    try:
        cfg = json.loads(os.environ.get("SPECULATIVE_CONFIG", "{}"))
        n = int(cfg.get("num_speculative_tokens", 0))
    except Exception:
        n = 0
    if n <= 0 or n % TREE_W:
        return 0
    return n // TREE_W

TREE_K = _conf_K()    # chain length; scheduler sees K*W drafts

_STATE: dict[str, Any] = {
    "tree_step": False,        # set by runner patch per step (host-side)
    "kv_caches": {},           # layer_name -> kv_cache tensor (stashed by attn patch)
    "layers": {},              # layer_name -> v4 per-layer fast-path cache
    "attn_calls": 0,           # python op-body invocations (dispatch diagnostic)
    "windows": {},             # layer_name -> window int (0=full)
    "seq_lens": None,          # current step's seq_lens tensor (runner patch)
    "steps": 0, "acc_tokens": 0, "salvages": 0, "full_accepts": 0,
}


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


def _fail(msg: str) -> None:
    if TREE_REQUIRE:
        raise RuntimeError(f"[tree-v2] REQUIRE: {msg}")
    _log(f"WARN (require off): {msg}")


# ---------------- star attention kernel (validated: star_attn_paged.py) -----

_KERNELS: dict[str, Any] = {}
_P_CACHE: dict[Any, Any] = {}
_PREWARMED_REJECT: set[tuple[int, int, str]] = set()


def _get_star_attn_kernel() -> Any:
    if "attn" in _KERNELS:
        return _KERNELS["attn"]
    import triton
    import triton.language as tl

    @triton.jit
    def _k_star_gqa(
        q_ptr, kc_ptr, vc_ptr, o_ptr, P_ptr, seqlen_ptr, bt_ptr, scale,
        sq_r, sq_h, skc_b, skc_s, skc_h, so_r, so_h,
        G: tl.constexpr, D: tl.constexpr, BK: tl.constexpr,
        BS: tl.constexpr, WINDOW: tl.constexpr, R: tl.constexpr, BG: tl.constexpr,
    ):
        # GQA KV-head sharing: one program per (row, KV-head) scores all G sibling
        # query heads against each KV tile loaded once (bf16 tensor-core tl.dot).
        r = tl.program_id(0)
        kvh = tl.program_id(1)
        d = tl.arange(0, D)
        g = tl.arange(0, BG)
        real = g < G
        qh = kvh * G + g
        q_g = tl.load(q_ptr + r * sq_r + qh[:, None] * sq_h + d[None, :],
                      mask=real[:, None], other=0.0)        # [BG, D] bf16
        ctx = tl.load(seqlen_ptr) - R          # seq_lens includes this step's R tokens
        Pr = tl.load(P_ptr + r)
        end = ctx + Pr
        if WINDOW > 0:
            start = tl.maximum(end - (WINDOW - 1), 0)
        else:
            start = 0
        start = (start // BK) * BK
        m_i = tl.full([BG], -float("inf"), tl.float32)
        l_i = tl.zeros([BG], tl.float32)
        acc = tl.zeros([BG, D], tl.float32)
        nblocks = tl.cdiv(end - start, BK)
        for b in range(nblocks):
            base = start + b * BK
            offs = base + tl.arange(0, BK)
            if WINDOW > 0:
                kmask = (offs < end) & (offs > (end - WINDOW))
            else:
                kmask = offs < end
            phys = tl.load(bt_ptr + offs // BS, mask=kmask, other=0)
            tok = phys * skc_b + (offs % BS) * skc_s + kvh * skc_h
            Kt = tl.load(kc_ptr + tok[:, None] + d[None, :], mask=kmask[:, None], other=0.0)
            Vt = tl.load(vc_ptr + tok[:, None] + d[None, :], mask=kmask[:, None], other=0.0)
            s = tl.dot(q_g, tl.trans(Kt)) * scale            # [BG, BK] f32
            s = tl.where(kmask[None, :], s, -float("inf"))
            m_new = tl.maximum(m_i, tl.max(s, axis=1))
            alpha = tl.exp(m_i - m_new)
            p = tl.exp(s - m_new[:, None])
            acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), Vt)
            l_i = l_i * alpha + tl.sum(p, axis=1)
            m_i = m_new
        soff = ctx + r
        sphys = tl.load(bt_ptr + soff // BS)
        sslot = soff - (soff // BS) * BS
        ks = tl.load(kc_ptr + sphys * skc_b + sslot * skc_s + kvh * skc_h + d).to(tl.float32)
        ss = tl.sum(q_g.to(tl.float32) * ks[None, :], axis=1) * scale
        m_new = tl.maximum(m_i, ss)
        alpha = tl.exp(m_i - m_new)
        ps = tl.exp(ss - m_new)
        vs = tl.load(vc_ptr + sphys * skc_b + sslot * skc_s + kvh * skc_h + d).to(tl.float32)
        acc = acc * alpha[:, None] + ps[:, None] * vs[None, :]
        l_i = l_i * alpha + ps
        tl.store(o_ptr + r * so_r + qh[:, None] * so_h + d[None, :],
                 acc / l_i[:, None], mask=real[:, None])

    _KERNELS["attn"] = _k_star_gqa
    return _k_star_gqa


def _star_P_tensor(K: int, W: int, device: Any) -> Any:
    import torch
    key = (K, W, device)
    if key not in _P_CACHE:
        # device-side construction — a host-list tensor() is an H2D copy,
        # illegal during CUDA-graph capture (run3 lesson)
        R = 1 + K * W
        P = torch.zeros(R, dtype=torch.int32, device=device)
        ar = torch.arange(1, K + 1, dtype=torch.int32, device=device)
        P[1 : K + 1] = ar
        P[K + 1 :] = ar.repeat_interleave(W - 1)
        _P_CACHE[key] = P
    return _P_CACHE[key]


# ---------------- fused rejection walk (pinned vLLM layout) -----------------

def _get_star_reject_kernel() -> Any:
    if "reject" in _KERNELS:
        return _KERNELS["reject"]
    import triton
    import triton.language as tl

    @triton.jit
    def _k_star_reject(
        ta_ptr,        # int32[K*W] argmax of TARGET rows
        draft_ptr,     # int32[K] main-path draft ids
        topk_ptr,      # int32[K, W]
        bonus_ptr,     # int32[1] sampler bonus (= b_K's continuation row)
        out_ptr,       # int32[K+1]
        valid_ptr,     # int32[1]
        salv_ptr,      # int32[1]: salvaged physical tree node or 0
        K: tl.constexpr, W: tl.constexpr,
    ):
        NT: tl.constexpr = K * W
        n_acc = 0
        nxt = -1
        salv_node = 0
        done = False
        for i in tl.static_range(K):
            if not done:
                pred = tl.load(ta_ptr + i)
                main = tl.load(draft_ptr + i)
                if pred == main:
                    tl.store(out_ptr + i, pred)
                    n_acc += 1
                else:
                    offs = tl.arange(0, W)
                    cand = tl.load(topk_ptr + i * W + offs, mask=offs > 0, other=-2)
                    hitmask = cand == pred
                    anyhit = tl.sum(hitmask.to(tl.int32), axis=0) > 0
                    if anyhit:
                        j = tl.min(tl.where(hitmask, offs, W * 2), axis=0)
                        node = K + 1 + i * (W - 1) + (j - 1)
                        tl.store(out_ptr + i, pred)
                        n_acc += 1
                        salv_node = node
                        if node < NT:
                            nxt = tl.load(ta_ptr + node)
                        else:
                            nxt = tl.load(bonus_ptr)   # last-position salvage
                    done = True
        if nxt == -1:
            if n_acc < K:
                nxt = tl.load(ta_ptr + n_acc)          # plain rejection
            else:
                nxt = tl.load(ta_ptr + K)              # full accept (row K)
        tl.store(out_ptr + n_acc, nxt)
        tl.store(valid_ptr, n_acc + 1)
        tl.store(salv_ptr, salv_node)

    _KERNELS["reject"] = _k_star_reject
    return _k_star_reject


def star_reject(target_argmax, draft_main, topk, bonus_token, K, placeholder, device,
                out_width=None):
    """Run the fused walk. out_width: vLLM expects (1, max_spec_len+1) = K*W+1
    PLACEHOLDER-filled columns; the walk writes at most K+1 of them."""
    import torch
    out = torch.full((1, out_width or (K + 1)), placeholder, dtype=torch.int32, device=device)
    valid = torch.zeros(1, dtype=torch.int32, device=device)
    salv = torch.zeros(1, dtype=torch.int32, device=device)
    kern = _get_star_reject_kernel()
    kern[(1,)](
        target_argmax.to(torch.int32), draft_main.to(torch.int32).contiguous(),
        topk.to(torch.int32).contiguous(), bonus_token.to(torch.int32).reshape(1),
        out.view(-1), valid, salv, K=K, W=TREE_W,
    )
    return out, valid, salv


def star_reject_from_drafts(target_argmax, draft_token_ids, bonus_token, K, placeholder, device):
    """draft_token_ids = the K*W scheduled drafts in [main(K); branches(K*(W-1),
    grouped by position)] order — reconstruct main/topk and run the fused walk.
    No reliance on side-channel topk state."""
    main = draft_token_ids[:K]
    branches = draft_token_ids[K:].reshape(K, TREE_W - 1)
    import torch
    topk = torch.cat([main.reshape(K, 1), branches], dim=1)
    return star_reject(target_argmax, main, topk, bonus_token, K, placeholder, device,
                       out_width=K * TREE_W + 1)


def prewarm_star_reject(K: int, W: int, device: Any, placeholder: int = -1) -> None:
    """Force Triton to compile the star rejection kernel before real inference."""
    import torch

    key = (K, W, str(device))
    if key in _PREWARMED_REJECT or K <= 0 or W <= 1:
        return
    target = torch.arange(K * W, dtype=torch.int64, device=device)
    draft = torch.arange(K * W, dtype=torch.int64, device=device)
    bonus = torch.zeros(1, dtype=torch.int64, device=device)
    star_reject_from_drafts(target, draft, bonus, K, placeholder, device)
    if device.type == "cuda" and not torch.cuda.is_current_stream_capturing():
        torch.cuda.synchronize()
    _PREWARMED_REJECT.add(key)
    _log(f"star rejection prewarmed (K={K}, W={W}, device={device})")


def relocate_salvaged_kv(salv_node: int, K: int) -> None:
    """Copy each layer's salvaged-branch KV (slot ctx+K+i) into the accepted
    slot (ctx+i). Called host-side only on salvage steps (~30%)."""
    import torch
    seq_lens = _STATE["seq_lens"]
    if seq_lens is None:
        _fail("relocate: seq_lens not stashed")
        return
    R = 1 + K * TREE_W
    # seq_lens already includes this step's R tokens
    base = seq_lens[0:1].to(torch.long) - R
    # salv_node is the physical input row: K+1 + pos0*(W-1) + branch0.
    # The previous W=2-only relocation copied base+K+pos, which is wrong for
    # W>2 when branch 2/3 wins and corrupts the accepted chain's KV.
    pos = ((int(salv_node) - (K + 1)) // (TREE_W - 1)) + 1
    src = base + int(salv_node)
    dst = base + pos
    for name, kv in _STATE["kv_caches"].items():
        kc, vc = kv.unbind(1)        # [num_blocks, bs, Hkv, D] each
        fk = kc.reshape(-1, *kc.shape[2:])
        fv = vc.reshape(-1, *vc.shape[2:])
        fk[dst] = fk[src]
        fv[dst] = fv[src]


def record_step(valid_count: int, salv: int, K: int) -> None:
    s = _STATE
    s["steps"] += 1
    s["acc_tokens"] += valid_count
    s["salvages"] += 1 if salv else 0
    s["full_accepts"] += 1 if valid_count == K + 1 and not salv else 0
    if s["steps"] % TREE_STATS_EVERY == 0:
        _log(f"stats steps={s['steps']} tok/step={s['acc_tokens']/s['steps']:.3f} "
             f"salvages={s['salvages']} full={s['full_accepts']} "
             f"attn_py_calls/step={s['attn_calls']/max(1,s['steps']):.1f} "
             f"(~37=EAGER dispatch, ~0=FULL-graph replay)")


# ---------------- runner patch: positions + step flag -----------------------

def _apply_runner_tree_patch(module: Any) -> None:
    import torch

    runner_cls = module.GPUModelRunner
    orig_prepare = runner_cls._prepare_inputs

    def _prepare_inputs(self: Any, scheduler_output: Any, *a: Any, **kw: Any) -> Any:
        result = orig_prepare(self, scheduler_output, *a, **kw)
        spec = scheduler_output.scheduled_spec_decode_tokens
        tree = False
        if SPEC_TREE_SPEC and len(spec) == 1:
            (toks,) = spec.values()
            ksch = len(toks)
            if (
                TREE_FAIL_ON_PLACEHOLDER_SPEC
                and ksch
                and any(int(tok) < 0 for tok in toks)
            ):
                _fail(
                    "scheduler handed negative tree draft tokens before "
                    f"GPU verify: len={ksch} head={list(toks)[:8]}"
                )
            if ksch and ksch % TREE_W == 0:
                K = ksch // TREE_W
                R = 1 + ksch
                pos = getattr(self.positions, "gpu", self.positions)
                # duplicate branch positions onto their main siblings.
                # branch rows are grouped BY POSITION (topk[:,1:W].flatten()):
                # row 1+K+t belongs to position 1 + t // (W-1)
                import torch as _t
                t = _t.arange(K * (TREE_W - 1), device=pos.device)
                pos[1 + K : R] = pos[1 + (t // (TREE_W - 1))]
                _STATE["seq_lens"] = self.seq_lens[:1]
                tree = True
        _STATE["tree_step"] = tree
        return result

    runner_cls._prepare_inputs = _prepare_inputs

    orig_dummy = runner_cls._dummy_run

    def _dummy_run(self: Any, num_tokens: int, *a: Any, **kw: Any) -> Any:
        # With MAX_NUM_SEQS=1 + spec K*W, a (1+K*W)-token dummy run is uniquely
        # the tree shape. Do NOT gate on uniform_decode: call sites pass it
        # inconsistently (run5: star path missed warmup AND capture), and a
        # star-valued dummy attention never affects piecewise replay semantics.
        R = 1 + TREE_K * TREE_W
        is_tree_shape = SPEC_TREE_SPEC and TREE_K > 0 and num_tokens == R
        prev = _STATE["tree_step"]
        if is_tree_shape:
            _STATE["tree_step"] = True
            _STATE["seq_lens"] = self.seq_lens[:1]
        try:
            return orig_dummy(self, num_tokens, *a, **kw)
        finally:
            _STATE["tree_step"] = prev

    runner_cls._dummy_run = _dummy_run
    _log(f"runner positions+dummy patch installed (pid {os.getpid()})")


# ---------------- star attention dispatch (called INSIDE the op body) -------
# Inserted into vllm's unified_attention_with_output impl by a serve.py disk
# patch (the stack's proven mechanism for compiled-region changes). Op bodies
# are opaque to Dynamo: this python runs at capture time / eagerly, never
# traced — is_current_stream_capturing is legal here (run2 lesson).

def maybe_star_attention(query, key, value, output, layer, kv_cache, md) -> bool:
    """Return True if the star-tree path handled this attention call."""
    import torch

    if md is None or not SPEC_TREE_SPEC or TREE_K <= 0:
        return False
    # EXCLUDE the drafter's own attention (agent-smith's documented gotcha):
    # draft_model.* has a different KV-share/layout; applying the star kernel
    # there caused the W=4 run's illegal memory access.
    if "draft" in getattr(layer, "layer_name", ""):
        return False
    num = getattr(md, "num_actual_tokens", -1)
    R = 1 + TREE_K * TREE_W
    if num != R or getattr(md, "max_query_len", -1) != R:
        return False
    if not _STATE["tree_step"]:
        return False

    # v4 fast path: everything layer-constant cached on first call; the
    # per-call python is one dict hit + cache-write + one kernel launch.
    # Calls are COUNTED: attn_calls/steps ~37 => eager dispatch; ~0 => replay.
    _STATE["attn_calls"] = _STATE.get("attn_calls", 0) + 1
    lc = _STATE["layers"].get(layer.layer_name)
    if lc is None:
        from vllm.v1.attention.backends.flash_attn import reshape_and_cache_flash

        impl = layer.impl
        if getattr(impl, "logits_soft_cap", None):
            _fail("logits_soft_cap unsupported in star kernel")
        key_cache, value_cache = kv_cache.unbind(1)
        sw = getattr(impl, "sliding_window", (-1, -1))
        window = sw[0] + 1 if isinstance(sw, tuple) and sw[0] >= 0 else 0
        H, Hkv, D = impl.num_heads, impl.num_kv_heads, impl.head_size
        dev = query.device
        lc = {
            "rcf": reshape_and_cache_flash,
            "kc": key_cache, "vc": value_cache,
            "dtype": layer.kv_cache_dtype,
            "ks": layer._k_scale, "vs": layer._v_scale,
            "H": H, "Hkv": Hkv, "D": D, "G": H // Hkv,
            "win": window, "scale": float(impl.scale),
            "BS": key_cache.shape[1],
            "P": _star_P_tensor(TREE_K, TREE_W, dev),
            "out32": torch.empty(num, H, D, dtype=torch.float32, device=dev),
            "kern": _get_star_attn_kernel(),
            "kcs": (key_cache.stride(0), key_cache.stride(1), key_cache.stride(2)),
        }
        _STATE["layers"][layer.layer_name] = lc
        _STATE["kv_caches"][layer.layer_name] = kv_cache
        prewarm_star_reject(TREE_K, TREE_W, dev, -1)
        import torch as _t
        _log(f"star attention layer-cache built ({layer.layer_name}, "
             f"capturing={_t.cuda.is_current_stream_capturing()}, num={num}) pid={os.getpid()}")

    lc["rcf"](
        key, value, lc["kc"], lc["vc"],
        md.slot_mapping[:num], lc["dtype"], lc["ks"], lc["vs"],
    )
    H, D = lc["H"], lc["D"]
    Hkv = lc["Hkv"]
    q = query[:num].reshape(num, H, D)
    out32 = lc["out32"]
    lc["kern"][(num, Hkv)](                       # grid (rows, KV-heads) for GQA sharing
        q, lc["kc"], lc["vc"], out32, lc["P"],
        md.seq_lens[:1], md.block_table[0],
        lc["scale"],
        q.stride(0), q.stride(1),
        lc["kcs"][0], lc["kcs"][1], lc["kcs"][2],
        out32.stride(0), out32.stride(1),
        G=lc["G"], D=D, BK=32,
        BS=lc["BS"], WINDOW=lc["win"], R=num, BG=16,
        num_warps=4, num_stages=1,
    )
    output[:num].copy_(out32.to(output.dtype).reshape(output[:num].shape))
    return True


# ---------------- registration ----------------------------------------------

class _Loader(importlib.abc.Loader):
    def __init__(self, inner: Any, apply_fn: Any) -> None:
        self._inner = inner
        self._apply = apply_fn

    def create_module(self, spec: Any) -> Any:
        return self._inner.create_module(spec)

    def exec_module(self, module: Any) -> None:
        self._inner.exec_module(module)
        self._apply(module)


class _Finder(importlib.abc.MetaPathFinder):
    def __init__(self, target: str, apply_fn: Any) -> None:
        self._target = target
        self._apply = apply_fn
        self._busy = False

    def find_spec(self, fullname: str, path: Any = None, target: Any = None) -> Any:
        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 = _Loader(spec.loader, self._apply)
        return spec


if SPEC_TREE_SPEC:
    sys.meta_path.insert(0, _Finder("vllm.v1.worker.gpu_model_runner", _apply_runner_tree_patch))
    # star attention is inserted into the unified op IMPL by a serve.py disk
    # patch; no meta-path wrapper (Dynamo traces Attention.forward — run2).
    _log(f"tree-v2 ext armed: W={TREE_W}, fused_reject={STAR_REJECT_FUSED}, require={TREE_REQUIRE}")