tree_dev/tree_decode.py

"""Tree speculative decoding — CORRECTNESS CORE + greedy-identity proof.

The challenge's hard constraint: served greedy decode must be TOKEN-IDENTICAL to
plain greedy. Tree spec decode (EAGLE-2/SpecInfer style) accepts more tokens per
verify step by pre-verifying multiple candidate continuations in ONE forward pass
with a tree attention mask, then walking the tree along the target's argmax path.

This module implements the algorithm model-agnostically and PROVES that the
accepted sequence == plain greedy decode for ANY tree (even adversarial/wrong
candidates) — i.e. token-identity is a property of the rejection walk, not of
drafter quality. It also measures E[L] (accepted tokens / verify step), the
speed prize (+1 accepted token ~= +107 TPS at the frontier).

Layout / semantics
------------------
A tree node n has: token tok(n), parent par(n) (an index into the flattened
node list; the root's "parent" is the prefix tail x0), depth(n).
Verify forward: sequence = [prefix ... x0] ++ [node tokens], with a 4D additive
mask where node t attends to (a) ALL prefix causally and (b) its tree ANCESTORS
only (not siblings/cousins). position_id(node at depth d) = len(prefix)-1 + d.
logits at node t = target's argmax for the token that would FOLLOW the path
root->t (the "prediction" pred(t)).

Greedy tree-accept walk (identity-preserving):
  cur = root(x0); pred = argmax(logits at x0)         # target's true next token
  loop: among children(cur), is there c with tok(c)==pred?
        yes -> accept tok(c); cur=c; pred=argmax(logits at c); continue
        no  -> stop; emit bonus = pred
  accepted = path tokens ++ [bonus]
Every accepted token equals the target's greedy argmax given the accepted
prefix => accepted == plain greedy. QED (verified empirically below).
"""
from __future__ import annotations

import argparse
import math
import random
from dataclasses import dataclass, field

import torch


# ----------------------------- tree structure -----------------------------
@dataclass
class Tree:
    """Flattened candidate tree. Node i in [0,N). par[i] in [-1,N): -1 means the
    node is a depth-1 child of the root x0; otherwise par[i] is another node."""
    tokens: list[int]                 # tok(i)
    par: list[int]                    # parent node index, -1 == child of root
    depth: list[int] = field(default_factory=list)

    def __post_init__(self):
        if not self.depth:
            self.depth = [self._compute_depth(i) for i in range(len(self.tokens))]

    def _compute_depth(self, i: int) -> int:
        d = 1
        p = self.par[i]
        while p != -1:
            d += 1
            p = self.par[p]
        return d

    @property
    def N(self) -> int:
        return len(self.tokens)

    def children(self, i: int) -> list[int]:
        """children of node i (i==-1 => children of root x0)."""
        return [j for j in range(self.N) if self.par[j] == i]


def build_spine_sibling_tree(cands: list[list[int]], width: int) -> Tree:
    """Spine = top-1 chain; at each spine depth add top-2..width as LEAF siblings.
    cands[d] = drafter's top-K candidate tokens at spine depth d (d=0..D-1).
    Returns a tree of depth D where the spine has children and siblings are leaves."""
    tokens, par = [], []
    spine_idx = []  # node index of the spine node at each depth
    for d, c in enumerate(cands):
        parent = -1 if d == 0 else spine_idx[d - 1]
        # spine node (top-1) first
        tokens.append(c[0]); par.append(parent); spine_idx.append(len(tokens) - 1)
        # leaf siblings top-2..width (children of the SAME parent => alternatives)
        for tk in c[1:width]:
            tokens.append(tk); par.append(parent)
    return Tree(tokens, par)


@torch.no_grad()
def draft_tree(drafter, prefix_ids, device, depth=7, width=4,
               branch_depth=0, branch_width=2):
    """Autoregressive REAL-drafter tree. The spine (top-1 chain) extends to
    `depth`; at each spine node, top-2..width are added as LEAF siblings
    (breadth / first-reject rescue, ~free). Additionally, the first
    `branch_depth` levels branch `branch_width` ways and EACH branch is
    continued autoregressively by the drafter to `depth` (continuation / depth
    rescue — the megakernel's unique capability). Returns a Tree."""
    tokens, par = [], []

    def topk_after(path):
        ids = torch.cat([prefix_ids, torch.tensor(path, device=device, dtype=prefix_ids.dtype)]) \
            if path else prefix_ids
        logits = drafter(input_ids=ids.unsqueeze(0), use_cache=False).logits[0, -1]
        return torch.topk(logits, width + 2).indices.tolist()

    def extend_spine(parent_idx, path, d, continue_branch):
        """Extend a top-1 spine from `path` down to `depth`, adding leaf siblings."""
        if d > depth:
            return
        cands = topk_after(path)
        # spine node (top-1)
        spine_tok = cands[0]
        sidx = len(tokens); tokens.append(spine_tok); par.append(parent_idx)
        # leaf siblings top-2..width (alternatives sharing this parent)
        sibs = []
        for tk in cands[1:width]:
            j = len(tokens); tokens.append(tk); par.append(parent_idx); sibs.append((j, tk))
        # continuation branches: if within branch_depth, continue each sibling too
        if continue_branch and d <= branch_depth:
            for j, tk in sibs[:branch_width - 1]:
                extend_spine(j, path + [tk], d + 1, continue_branch)
        extend_spine(sidx, path + [spine_tok], d + 1, continue_branch)

    extend_spine(-1, [], 1, branch_depth > 0)
    return Tree(tokens, par)


# --------------------------- verify (tree mask) ---------------------------
NEG = torch.finfo(torch.float32).min


def tree_verify(model, prefix_ids: torch.Tensor, tree: Tree, device) -> tuple[list[int], int]:
    """Run the target over [prefix ++ tree] with the tree mask. Returns
    (pred_per_node, pred_root): pred_root = argmax after x0; pred_per_node[i] =
    argmax for the token following node i (given the root->i path)."""
    L = prefix_ids.shape[0]
    N = tree.N
    seq = torch.cat([prefix_ids, torch.tensor(tree.tokens, device=device, dtype=prefix_ids.dtype)])
    S = L + N

    # position ids: prefix 0..L-1 ; node at depth d -> (L-1)+d
    pos = torch.arange(L, device=device)
    node_pos = torch.tensor([L - 1 + tree.depth[i] for i in range(N)], device=device)
    position_ids = torch.cat([pos, node_pos]).unsqueeze(0)

    # 4D additive mask [1,1,S,S]
    mask = torch.full((S, S), NEG, device=device, dtype=torch.float32)
    # prefix block: causal among prefix
    idx = torch.arange(L, device=device)
    mask[:L, :L] = torch.where(idx[:, None] >= idx[None, :], 0.0, NEG)
    # each node attends to ALL prefix (it is causally after x0)
    mask[L:, :L] = 0.0
    # node-node: node t attends to its ancestors (incl self) only
    # precompute ancestor sets
    for t in range(N):
        # self
        mask[L + t, L + t] = 0.0
        p = tree.par[t]
        while p != -1:
            mask[L + t, L + p] = 0.0
            p = tree.par[p]
    mask = mask.view(1, 1, S, S)

    with torch.no_grad():
        out = model(input_ids=seq.unsqueeze(0), position_ids=position_ids,
                    attention_mask=mask, use_cache=False)
    logits = out.logits[0]  # [S, V]
    pred_root = int(logits[L - 1].argmax())               # argmax after x0
    pred_node = [int(logits[L + i].argmax()) for i in range(N)]
    return pred_node, pred_root


def greedy_tree_accept(tree: Tree, pred_node: list[int], pred_root: int) -> list[int]:
    """Walk the tree along target-argmax matches. Returns accepted token list
    (the matched path) PLUS the final bonus token. Identity-preserving."""
    accepted = []
    cur = -1                  # root x0
    pred = pred_root
    while True:
        match = None
        for c in tree.children(cur):
            if tree.tokens[c] == pred:
                match = c
                break
        if match is None:
            accepted.append(pred)         # bonus (target argmax) — always greedy-correct
            return accepted
        accepted.append(tree.tokens[match])
        cur = match
        pred = pred_node[match]


# ------------------------------- references -------------------------------
@torch.no_grad()
def plain_greedy(model, prefix_ids: torch.Tensor, n: int, device) -> list[int]:
    """Vanilla greedy continuation of `n` tokens (the ground truth)."""
    ids = prefix_ids.clone().unsqueeze(0)
    out = []
    for _ in range(n):
        logits = model(input_ids=ids, use_cache=False).logits[0, -1]
        nxt = int(logits.argmax())
        out.append(nxt)
        ids = torch.cat([ids, torch.tensor([[nxt]], device=device)], dim=1)
    return out


@torch.no_grad()
def linear_accept_len(model, prefix_ids: torch.Tensor, spine: list[int], device) -> int:
    """Standard linear chain accept: how many of `spine` the target greedily
    accepts (matching prefix) + 1 bonus. The E[L] baseline."""
    L = prefix_ids.shape[0]
    seq = torch.cat([prefix_ids, torch.tensor(spine, device=device, dtype=prefix_ids.dtype)])
    logits = model(input_ids=seq.unsqueeze(0), use_cache=False).logits[0]
    m = 0
    for i, s in enumerate(spine):
        t = int(logits[L - 1 + i].argmax())   # target argmax given prefix+spine[:i]
        if t == s:
            m += 1
        else:
            break
    return m + 1   # + bonus


# --------------------------------- tests ----------------------------------
def stress_identity(model, tok, device, rounds=40, depth=7, width=4, vocab=None, seed=0):
    """For RANDOM/adversarial trees, the tree-accept output MUST be a prefix of
    plain greedy. This proves identity is algorithmic, not drafter-dependent."""
    rng = random.Random(seed)
    V = vocab or model.config.vocab_size
    prompt = "The history of the Roman empire began when"
    base = tok(prompt, return_tensors="pt").input_ids[0].to(device)
    fails = 0
    for r in range(rounds):
        # ground-truth greedy continuation (enough to cover any accept length)
        gt = plain_greedy(model, base, depth + 2, device)
        # build an adversarial spine+sibling tree: mix correct / wrong / random
        cands = []
        for d in range(depth):
            row = []
            mode = rng.random()
            if mode < 0.5:        # spine correct at this depth
                row.append(gt[d])
            else:                 # spine wrong
                row.append(rng.randrange(V))
            # siblings: sometimes include the correct token, sometimes random
            while len(row) < width:
                if rng.random() < 0.4:
                    row.append(gt[d])           # correct as a sibling
                else:
                    row.append(rng.randrange(V))
            row = list(dict.fromkeys(row))[:width] or [rng.randrange(V)]
            cands.append(row)
        tree = build_spine_sibling_tree(cands, width)
        pred_node, pred_root = tree_verify(model, base, tree, device)
        acc = greedy_tree_accept(tree, pred_node, pred_root)
        # identity check: acc must equal gt[:len(acc)]
        ok = acc == gt[:len(acc)]
        if not ok:
            fails += 1
            print(f"  [FAIL r{r}] acc={acc} gt={gt[:len(acc)]}")
    print(f"[identity] {rounds-fails}/{rounds} adversarial trees produced "
          f"greedy-identical output  ({'PASS' if fails==0 else 'FAIL'})")
    return fails == 0


PROMPTS = [
    "The history of the Roman empire began when",
    "In computer science, a binary search tree is",
    "The recipe calls for two cups of flour and",
    "Quantum entanglement is a phenomenon where",
    "She opened the door and saw that",
    "The stock market fell sharply today after",
    "To prove the theorem, we first assume that",
    "The weather forecast for tomorrow predicts",
    "The capital of France is a city that",
    "def fibonacci(n):\n    if n <= 1:\n        return",
]


@torch.no_grad()
def draft_spine(drafter, prefix_ids, device, depth, width):
    """Real drafter top-(width) at each depth along the drafter's own greedy
    spine. Returns cands[d] (top-K at depth d) — the leaf-sibling tree input."""
    ids = prefix_ids.clone().unsqueeze(0)
    cands = []
    for _ in range(depth):
        logits = drafter(input_ids=ids, use_cache=False).logits[0, -1]
        row = torch.topk(logits, width).indices.tolist()
        cands.append(row)
        ids = torch.cat([ids, torch.tensor([[row[0]]], device=device)], dim=1)
    return cands


def measure_el_pair(drafter, target, tok, device, depth=7, width=4,
                    branch_depth=2, branch_width=2, n_prompts=10):
    """REAL spec-decode E[L]: small `drafter` proposes, big `target` verifies.
    Compares per-verify-step accepted tokens for: linear chain, leaf-sibling
    tree (breadth), and continuation tree (breadth+depth). One step per prompt
    (the first verify step from the prompt). Also asserts identity per step."""
    lin = sib = cont = 0.0
    steps = 0
    ident_ok = True
    for p in PROMPTS[:n_prompts]:
        base = tok(p, return_tensors="pt").input_ids[0].to(device)
        gt = plain_greedy(target, base, depth + 2, device)   # target's true greedy
        cands = draft_spine(drafter, base, device, depth, width)
        spine = [c[0] for c in cands]

        # linear chain
        lin += linear_accept_len(target, base, spine, device)

        # leaf-sibling tree (breadth)
        t_sib = build_spine_sibling_tree(cands, width)
        pn, pr = tree_verify(target, base, t_sib, device)
        a_sib = greedy_tree_accept(t_sib, pn, pr)
        ident_ok &= (a_sib == gt[:len(a_sib)])
        sib += len(a_sib)

        # continuation tree (breadth + depth): drafter re-runs along branches
        t_cont = draft_tree(drafter, base, device, depth=depth, width=width,
                            branch_depth=branch_depth, branch_width=branch_width)
        pn, pr = tree_verify(target, base, t_cont, device)
        a_cont = greedy_tree_accept(t_cont, pn, pr)
        ident_ok &= (a_cont == gt[:len(a_cont)])
        cont += len(a_cont)
        steps += 1
    print(f"[E[L] real-pair] {steps} steps, w={width}, branch={branch_width}x{branch_depth}")
    print(f"  linear chain        = {lin/steps:.3f} tok/step")
    print(f"  leaf-sibling tree   = {sib/steps:.3f} tok/step  ({sib/steps-lin/steps:+.3f})")
    print(f"  continuation tree   = {cont/steps:.3f} tok/step  ({cont/steps-lin/steps:+.3f})")
    print(f"  identity over all tree steps: {'PASS' if ident_ok else 'FAIL'}")
    # frontier sensitivity: +1 accepted tok ~= +107 TPS
    print(f"  @+107 TPS/tok: sibling ~ {107*(sib/steps-lin/steps):+.0f} TPS, "
          f"continuation ~ {107*(cont/steps-lin/steps):+.0f} TPS (illustrative pair)")


def edge_cases(model, tok, device):
    """Identity must hold for degenerate trees too."""
    base = tok("The quick brown fox", return_tensors="pt").input_ids[0].to(device)
    gt = plain_greedy(model, base, 8, device)
    V = model.config.vocab_size
    cases = {
        "all-correct spine": [[gt[d]] for d in range(7)],
        "all-wrong":         [[(gt[d] + 1) % V] for d in range(7)],
        "correct-as-sibling":[[(gt[d] + 1) % V, gt[d]] for d in range(7)],
        "single-node":       [[gt[0]]],
        "wrong-then-right":  [[(gt[0]+1)%V], [gt[1]], [gt[2]]],
    }
    allok = True
    for name, cands in cases.items():
        w = max(len(c) for c in cands)
        tree = build_spine_sibling_tree(cands, w)
        pn, pr = tree_verify(model, base, tree, device)
        acc = greedy_tree_accept(tree, pn, pr)
        ok = acc == gt[:len(acc)]
        allok &= ok
        print(f"  [{'ok' if ok else 'FAIL'}] {name:20s} accepted {len(acc)} tok")
    print(f"[edge] {'PASS' if allok else 'FAIL'}")
    return allok


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="HuggingFaceTB/SmolLM-135M")
    ap.add_argument("--target", default=None, help="bigger model to verify (real pair)")
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--depth", type=int, default=7)
    ap.add_argument("--width", type=int, default=4)
    ap.add_argument("--branch-depth", type=int, default=2)
    ap.add_argument("--branch-width", type=int, default=2)
    ap.add_argument("--rounds", type=int, default=40)
    args = ap.parse_args()
    from transformers import AutoModelForCausalLM, AutoTokenizer
    dev = args.device

    def load(name):
        print(f"[load] {name} on {dev}")
        m = AutoModelForCausalLM.from_pretrained(
            name, dtype=torch.float32, attn_implementation="eager").to(dev).eval()
        print(f"[load]   vocab={m.config.vocab_size} layers={m.config.num_hidden_layers}")
        return m

    tok = AutoTokenizer.from_pretrained(args.model)
    drafter = load(args.model)

    # correctness proof (model-agnostic): adversarial + edge-case identity
    ok = stress_identity(drafter, tok, dev, rounds=args.rounds, depth=args.depth, width=args.width)
    ok &= edge_cases(drafter, tok, dev)

    # E[L] prize: real draft/target pair if --target given, else self
    target = load(args.target) if args.target else drafter
    measure_el_pair(drafter, target, tok, dev, depth=args.depth, width=args.width,
                    branch_depth=args.branch_depth, branch_width=args.branch_width)
    print("[done] identity", "PASS" if ok else "FAIL")


if __name__ == "__main__":
    main()