tree_dev/SERVE_INTEGRATION_SPEC.md

---
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.