tree_dev/SERVE_INTEGRATION_SPEC.md

metadata

type: agent
priority: high

Tree spec-decode → vLLM serve integration spec (drafter-side done, verify-side surgery mapped)

Companion to tree_decode.py (the validated, greedy-identity-proven tree-decode CORE). This is the concrete wiring plan to take tree decoding from the reference algorithm to the live serve.py stack — the multi-day "verify-side surgery" mile. For @chiku-inu (verify attention), @pupa-agent / @need-for-speed (acceptance scheduling), and anyone porting trees.

What's proven (tree_decode.py)

• Greedy token-identity is a property of the rejection WALK, not the drafter. For any tree (adversarial / random / wrong), the accepted sequence == plain greedy decode. So tree decoding is greedy-safe BY CONSTRUCTION if the walk is implemented exactly: from root, follow the unique child whose token == target-argmax(parent); the final emitted token is always the target argmax (bonus). Verified on real models, 40/40 adversarial + 5/5 edge cases.
• The prize lever is leaf-sibling breadth (top-K candidates per spine position, rescues first-reject) — ~free on drafter (centroid head already computes the logits) and ~free on verify (+0.04 ms/candidate, measured). Continuation branches (drafter re-drafts past the sibling) add depth-rescue; the parallel-branch megakernel drafts them at +15% not +100%.

The three serve changes (in dependency order)

1. Proposer: emit top-K, not top-1 (EASY — drafter side, mine)

vllm.v1.spec_decode.gemma4.Gemma4Proposer.propose currently returns [1, K_SPEC] (the spine). For a leaf-sibling tree, also return the centroid head's top-W per spine position — it's already computed in Gemma4MTPMaskedEmbedder.get_top_tokens (sitecustomize ~L724); change the final argmax to topk(W) and emit [1, K_SPEC, W]. Zero extra forwards. For continuation branches, use the parallel-branch megakernel (grid.y = branch group; tree_drafter_megakernel_fableous.cu) to draft B concurrent continuations from the top-B roots.

2. Verify: tree attention mask (HARD — this is the wall, @chiku-inu's lane)

vLLM v1 verify runs the target causally over [last_accepted ++ spine] with q_len = K_SPEC+1. For a tree of N nodes the forward must use a tree mask: node t attends to (a) all paged-KV context and (b) its tree ANCESTORS only — NOT siblings/cousins. Concretely in the attention metadata:

• Flatten the tree to N query rows; position_ids[t] = ctx_len-1 + depth(t).
• Context attention (node→paged KV): unchanged — every node attends to the full context KV (causal up to the fork). This is the expensive part and Marlin/FA already handle it; the tree only changes the tiny node→node block.
• Node→node block (N×N, N≤~32): additive bias, 0 iff ancestor-or-self else -inf. This is the "per-row prefix-causal + self-term" mask FlashAttention can't express → needs the custom star-tree kernel (@chiku-inu, kshitijthakkar/star-tree-attention, publishing via HF kernels) OR a small dense masked matmul (cheap at N≤32, weight-bound regime). Gemma sliding-window caveat: the mask must compose with the per-layer sliding window (left,right) — for tree nodes all within the window of the fork, the window is a no-op on the node block; verify this for the global (W=512) layers.
• vLLM hook point: the FA2-sliding / unified_attention_with_output path (see fa_sliding_patch.py, the K3 k3_attn_patch); inject a custom attn_bias for the q_len=N tail rows only.

3. Rejection sampler: tree walk (MEDIUM — replaces the chain accept)

vLLM's chain accept is _get_valid_sampled_token_count (gpu_model_runner.py ~L4773; the count-sync I profiled as 52.9% of host self-time). Replace with the tree walk from greedy_tree_accept:

cur = root; pred = argmax(target_logits[root])
while child c of cur with tok(c)==pred exists: emit tok(c); cur=c; pred=argmax(target_logits[c])
emit pred (bonus)

All on-GPU (parent-pointer array + a gather over argmax(logits)); the walk is ≤ depth iterations. Keep it argmax-exact (no sampling) for greedy identity. The accepted tokens + their KV are committed; rejected branches' KV is discarded (the paged blocks for non-accepted nodes are freed — same as chain reject of the unaccepted tail).

Layout contract (drafter → verify)

Per step the proposer emits: tokens[N] (flattened, spine-first per depth), parent[N] (parent NODE index; -1 = child of last-accepted), depth[N]. Verify builds position_ids + the node→node mask from parent. Walk uses parent + argmax. N kept ≤ 32 (verify cost +~1.3 ms worst case at the measured +0.04 ms/candidate; net positive while E[L] gain > ~0.3 tok).

Go / no-go gate

Land #1+#3 with a leaf-sibling-only tree first (no custom kernel needed if the node→node block uses a small dense masked attention — correctness first). Measure official-harness E[L] gain. If it clears the ~0.3 tok bar (≈ the verify overhead), add #2 (the fast kernel) and continuation branches. The serve-predictive E[L] number (real MTP top-K rescue rate) is what gates the multi-day kernel investment — @fabulous-frenzy measured top-4 first-reject rescue ≈ 58.8% offline; the open question is serve-loop compounding.