Beam Search、约束解码与 Tree/Self Speculation

三类机制解决不同问题

机制	目标	改变的状态	不保证什么
Beam Search	近似寻找高序列概率输出	多个 beams、累计 log score、每 beam KV	不是从目标分布采样；beam 更宽不保证任务更好
Constrained Decoding	确保输出属于 grammar/FSM/schema 允许语言	parser/FSM state + allowed-token mask	不保证内容事实正确或全局最优
Tree/Self Speculation	一次 target verification 覆盖多个共享前缀候选	candidate tree、tentative KV、accepted path	不自动保持 exact sampling；协议必须另行定义

它们都位于 logits/sampling 下游，但与 exact speculative sampling的概率校正边界不同。Beam Search 通常是确定性/近似搜索；grammar mask 改变可采样支持集；tree verifier 可以承载 greedy、beam-like 或专门训练的候选，是否 exact 取决于 acceptance/correction 设计。

Beam Search：比较累计序列分数，不是逐步 greedy

给定 prefix $y_{1:t}$ ，未归一化累计分：

s(y_{1:t})=\sum_{i=1}^{t}\log p(y_i\mid y_{<i})

每步将 $B_{beam}$ 条 beams 各扩展到 vocabulary，再保留全局 top- $B_{beam}$ candidates。常见 length penalty：

\tilde s(y)=\frac{s(y)}{\operatorname{lp}(|y|)}

其具体函数、EOS 处理、early stopping 都属于 API 语义，不能只写“除以长度”。本页小例子中 greedy 先选 A 后立刻 EOS；beam width 2 保留第二条 B，最终发现 B A EOS 累计 log probability 更高。

Beam KV 不能每步无条件深拷贝全部 prefix。Paged KV 可让共同前缀 blocks 共享引用，分叉尾部 copy-on-write；beam reorder metadata 必须同步 token ids、scores、finished flags 和每层 KV block table。

成本主项近似让 active sequences 从 $B$ 变成 $B\times B_{beam}$ ，attention/KV/MLP 都增加；vocab top-k 还要处理 [B,B_{beam},V] 逻辑 scores。更宽 beam 不保证线性耗时，因为 kernels/batch utilization 变化，但容量通常明显上升。

约束解码：状态决定 allowed-token set

设解析器/FSM 当前状态为 $s_t$ ，允许 next token 集合 $A(s_t)$ ，mask：

M_v(s_t)=\begin{cases}0,&v\in A(s_t)\\-\infty,&v\notin A(s_t)\end{cases}

p_C(v\mid y_{<t},s_t)=\operatorname{softmax}(z_v+M_v(s_t))

选择 token 后更新 $s_{t+1}=\delta(s_t,v)$ 。若用采样，得到的是条件在合法语言上的逐步 masked 分布；它不是原始无约束分布。若 allowed set 为空，需要定义 error/repair/fallback，不能把全 -inf 交给 softmax。

本页 FSM 只允许语言：

A (B|C) EOS

即使每步原始最大 logit 是非法 token，mask 后仍得到 A C EOS。真实 JSON/regex/CFG 还需处理 tokenizer token 可能对应多个字符/bytes：allowed-token 计算不能假定“一 token 一字符”，与 Streaming 的 UTF-8 边界相邻但方向相反——constraint 在生成前限制 token，streaming 在提交后安全解码 bytes。

约束引擎成本包括 parser state、token trie/automaton、allowed bitmap/list、mask kernel 与 batch 中不同 states。若在 CPU 每 token 扫全 vocabulary，可能成为 TPOT 瓶颈；生产实现会缓存 state→allowed tokens、在 GPU mask，或与 sampling kernel 融合。

Tree/Self Speculation：candidate graph 与 committed path 分离

Tree speculation 让多个 candidates 共享 prefix：

base [10,11]
├─ 4 ─ 5 ─ 6
│    └ 7 ─ 8
└─ 9 ─ 0 ─ 1

教学 branches 表示为 [4,5,6]、[4,7,8]、[9,0,1]。Target verification 给出 next path [4,7,3]，只有 [4,7] 在树中连续匹配，因此 committed cache 从 [10,11] 变成 [10,11,4,7]；其他 tentative nodes 和分支上的位置/attention metadata 必须 rollback/release。

Self-speculation 指同一 target model 的早退层、跳层或辅助 heads 产生候选；它减少独立 draft 权重，但候选质量和 target verification 成本仍决定收益。MTP可提供多未来 token heads，但训练目标不等于 tree verification/acceptance protocol。

Exact sampling 若用 tree proposals，需要定义 tree probability、节点访问顺序、接受/校正分布；本页的 prefix match 只证明 greedy/tree commit 与 rollback ownership，不证明保持目标采样分布。

Medusa 与 EAGLE 改变的是 proposal，不是共同 verifier 不变量

Medusa 在 backbone hidden state 上增加多个 decoding heads，分别预测多个后续位置；从各 head 的 top candidates 组合候选树，再用 tree attention 并行验证。Medusa-1 冻结 backbone、训练 heads；Medusa-2 会共同微调 backbone 和 heads。多 head logits 不是 target 在每个已展开 candidate prefix 下逐步运行得到的同一组 logits，仍需接受规则。
EAGLE 把 draft 自回归放在 target 模型的次顶层 feature space，并输入错开一位的 token sequence 来降低 feature uncertainty；draft feature 经 target LM head 等路径形成 token candidates，最后仍由 target 验证。它不是“另一个更小的 token-level draft model”，也不等同 MTP 的多个监督未来 heads。

两者都可能构造 tree，但 candidate source、训练目标、draft state 和接受协议不同。本页只承载它们共同依赖的 packed tree attention、position/mask、tentative KV 与 commit/rollback；论文中的 speedup 是特定模型、硬件和配置结果，不能移植为全站性能承诺。

Packed tree attention：一次 target 前向怎样隔离 siblings

把树中唯一节点按拓扑序 packed。仍用 branches [4,5,6]、[4,7,8]、[9,0,1]，唯一节点可表示为：

packed node	token	parent	depth	candidate position id	可见 candidate ancestors
0	4	-1	1	2	`[0]`
1	9	-1	1	2	`[1]`
2	5	0	2	3	`[0,2]`
3	7	0	2	3	`[0,3]`
4	0	1	2	3	`[1,4]`
5	6	2	3	4	`[0,2,5]`
6	8	3	3	4	`[0,3,6]`
7	1	4	3	4	`[1,4,7]`

假设 committed base prefix 长 $P=2$ 。packed tensor：

base K/V       [B,H,P,D] = [1,2,2,8]   committed、只读
candidate QKV  [B,H,N,D] = [1,2,8,8]   tentative、节点所有权
packed K/V     [B,H,P+N,D] = [1,2,10,8]
tree mask      [N,P+N] = [8,10]         bool / additive -∞
output         [B,H,N,D] = [1,2,8,8]

每个 candidate Query 可看全部 base prefix 和自己的 ancestor chain，不能看 siblings/cousins。根 nodes 0、1 虽 packed indices 不同，tree depth 相同，所以 RoPE position 都是 $P=2$ ；nodes 2、3、4 都是 position 3。把 packed 顺序直接当 position id 会改变 Q/K rotation。

可运行的 packed-vs-branch attention 实验

实验问题：一个 [N,P+N] ancestry mask 的 packed attention，能否与每个 candidate node 单独收集 base + ancestor path 的 reference 对齐？错误 lower-triangular candidate mask 是否会因 sibling leakage 改变输出？

import math,platform,torch
torch.manual_seed(233)
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE=torch.float64

def rotate_half(x):
    return torch.stack((-x[...,1::2],x[...,0::2]),-1).flatten(-2)

def rope(x,positions,base=10_000.):
    d=x.shape[-1]
    inv=base**(-torch.arange(0,d,2,device=x.device,dtype=x.dtype)/d)
    angles=positions.to(x.dtype)[...,None]*inv
    cos=torch.repeat_interleave(angles.cos(),2,-1)
    sin=torch.repeat_interleave(angles.sin(),2,-1)
    return x*cos+rotate_half(x)*sin

def tree_depths(parents):
    depths=[]
    for node in range(len(parents)):
        depth,cursor=1,parents[node]
        while cursor!=-1: depth+=1; cursor=parents[cursor]
        depths.append(depth)
    return depths

def ancestor_path(node,parents):
    path=[]
    while node!=-1: path.append(node); node=parents[node]
    return list(reversed(path))

def tree_mask(parents,prefix_length):
    n=len(parents); mask=torch.zeros(n,prefix_length+n,device=DEVICE,dtype=torch.bool)
    mask[:,:prefix_length]=True
    for node in range(n):
        path=torch.tensor(ancestor_path(node,parents),device=DEVICE)
        mask[node,prefix_length+path]=True
    return mask

def packed_attention(q,bk,bv,ck,cv,mask):
    k=torch.cat((bk,ck),-2); v=torch.cat((bv,cv),-2)
    scores=q@k.transpose(-2,-1)/math.sqrt(q.shape[-1])
    return torch.softmax(scores.masked_fill(~mask,-torch.inf),-1)@v

def branch_reference(q,bk,bv,ck,cv,parents):
    outputs=[]
    for node in range(len(parents)):
        path=torch.tensor(ancestor_path(node,parents),device=DEVICE)
        k=torch.cat((bk,ck.index_select(-2,path)),-2)
        v=torch.cat((bv,cv.index_select(-2,path)),-2)
        scores=q[...,node:node+1,:]@k.transpose(-2,-1)/math.sqrt(q.shape[-1])
        outputs.append(torch.softmax(scores,-1)@v)
    return torch.cat(outputs,-2)

def main():
    tokens=[4,9,5,7,0,6,8,1]
    parents=[-1,-1,0,0,1,2,3,4]
    b,h,p,n,d=1,2,2,len(parents),8
    depths=tree_depths(parents)
    positions=torch.tensor([p+depth-1 for depth in depths],device=DEVICE)
    bkraw=torch.randn(b,h,p,d,device=DEVICE,dtype=DTYPE); bv=torch.randn_like(bkraw)
    qraw=torch.randn(b,h,n,d,device=DEVICE,dtype=DTYPE)
    ckraw=torch.randn_like(qraw); cv=torch.randn_like(qraw)
    bk=rope(bkraw,torch.arange(p,device=DEVICE)); q=rope(qraw,positions); ck=rope(ckraw,positions)
    mask=tree_mask(parents,p)
    packed=packed_attention(q,bk,bv,ck,cv,mask)
    reference=branch_reference(q,bk,bv,ck,cv,parents)
    torch.testing.assert_close(packed,reference,rtol=1e-12,atol=1e-12)
    wrong_mask=torch.cat((torch.ones(n,p,device=DEVICE,dtype=torch.bool),
                          torch.tril(torch.ones(n,n,device=DEVICE,dtype=torch.bool))),-1)
    wrong=packed_attention(q,bk,bv,ck,cv,wrong_mask)
    leak=(wrong-reference).abs().max(); assert leak>1e-3
    paths=[ancestor_path(node,parents) for node in range(n)]
    assert mask.sum(-1).tolist()==[3,3,4,4,4,5,5,5]
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={DTYPE} prefix_KV={(b,h,p,d)}")
    print(f"candidate_QKV={(b,h,n,d)} packed_keys={(b,h,p+n,d)}")
    print(f"token_ids={tokens} parents={parents} depths={depths}")
    print(f"candidate_position_ids={positions.tolist()}")
    print(f"ancestor_paths={paths}")
    print(f"tree_mask_visible_counts={mask.sum(-1).tolist()}")
    print(f"packed_vs_branch_max_abs_diff={(packed-reference).abs().max():.3e}")
    print(f"wrong_causal_mask_sibling_leak_max_abs_diff={leak:.6f}")
    print("branch_expanded_nodes=9 unique_tree_nodes=8 shared_nodes_saved=1")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/tree_speculation_attention.py。2026-08-13 在 PyTorch 2.13.0+cu130 上，packed tree 与逐节点 branch reference 最大差为 4.441e-16；错误 lower-triangular mask 的 sibling-leak 最大差为 1.845393。三条 branch 若逐分支复制共有 9 个 candidate positions，去重后 tree 有 8 个 nodes，本例只省一个共享节点；真实收益取决于树宽、共享度与 verifier batch 效率。

代码只验证单层 attention ancestry、position 和 tensor shape，没有实现 Medusa/EAGLE proposal 网络、target logits、sampling correction、paged tentative KV 或 production tree kernel。因此不能用其 Python 速度推断端到端收益。

可运行的三段语义实验

import platform,torch

torch.manual_seed(227)
VOCAB=["A","B","C","EOS"]; EOS=3

def next_logits(prefix):
    table={():torch.tensor([2.,1.7,-2.,-4.]),
        (0,):torch.tensor([-.1,-.2,-.3,0.]),
        (1,):torch.tensor([2.5,-3.,-3.,-1.]),
        (1,0):torch.tensor([-3.,-3.,-3.,3.])}
    return table.get(tuple(prefix),torch.tensor([-4.,-4.,-4.,4.]))

def beam_search(width=2,max_steps=3):
    beams=[([],0.)]
    for _ in range(max_steps):
        candidates=[]
        for prefix,score in beams:
            if prefix and prefix[-1]==EOS: candidates.append((prefix,score)); continue
            log_probs=torch.log_softmax(next_logits(prefix),-1)
            for token in range(len(VOCAB)):
                candidates.append((prefix+[token],score+float(log_probs[token])))
        beams=sorted(candidates,key=lambda item:item[1],reverse=True)[:width]
        if all(prefix[-1]==EOS for prefix,_ in beams): break
    return beams

def allowed_tokens(state):
    return {"start":[0],"after_A":[1,2],"after_middle":[EOS]}[state]

def transition(state,token):
    return {("start",0):"after_A",("after_A",1):"after_middle",
            ("after_A",2):"after_middle",("after_middle",EOS):"done"}[(state,token)]

def constrained_greedy(logits_per_step):
    state="start"; output=[]
    for logits in logits_per_step:
        allowed=allowed_tokens(state); mask=torch.full_like(logits,-torch.inf); mask[allowed]=0
        token=int(torch.argmax(logits+mask)); output.append(token); state=transition(state,token)
        if state=="done": break
    return output,state

def verify_tree(base_cache,branches,target_next):
    tentative={tuple(branch[:d]) for branch in branches for d in range(1,len(branch)+1)}
    committed=list(base_cache); accepted=[]
    for depth,target in enumerate(target_next):
        matching=[branch for branch in branches if branch[:depth]==accepted and branch[depth]==target]
        if not matching: break
        accepted.append(target); branches=matching
    committed.extend(accepted)
    accepted_nodes={tuple(accepted[:d]) for d in range(1,len(accepted)+1)}
    return committed,accepted,len(tentative-accepted_nodes)

def main():
    greedy=[]
    while len(greedy)<3 and (not greedy or greedy[-1]!=EOS):
        greedy.append(int(next_logits(greedy).argmax()))
    beams=beam_search(); assert greedy==[0,EOS] and beams[0][0]==[1,0,EOS]
    logits=[torch.tensor([-2.,4.,3.,1.]),torch.tensor([5.,1.,2.,0.]),
            torch.tensor([3.,2.,1.,-1.])]
    constrained,state=constrained_greedy(logits)
    assert constrained==[0,2,EOS] and state=="done"
    branches=[[4,5,6],[4,7,8],[9,0,1]]
    committed,accepted,rejected=verify_tree([10,11],branches,[4,7,3])
    assert committed==[10,11,4,7] and accepted==[4,7] and rejected==6
    names=lambda tokens:[VOCAB[t] if t<len(VOCAB) else str(t) for t in tokens]
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"greedy={names(greedy)}")
    print(f"beam_best={names(beams[0][0])} score={beams[0][1]:.6f} beam_second={names(beams[1][0])} score={beams[1][1]:.6f}")
    print(f"constrained={names(constrained)} final_state={state}")
    print(f"tree_branches={branches} target_next={[4,7,3]}")
    print(f"accepted_prefix={accepted} committed_cache={committed} rejected_tentative_tokens={rejected}")

if __name__=="__main__": main()

完整文件位于 examples/decoding_extensions.py。实际 greedy 为 A EOS，beam 最优为 B A EOS；约束输出为 A C EOS；tree verifier 接受 [4,7]、提交后的 cache 为 [10,11,4,7]，释放 6 个非 accepted tentative nodes。

小例子没有模型 KV tensor、length penalty、batched beams 或 exact tree sampling，因此不做性能结论。它固定了三类机制最容易混淆的状态不变量。

成本、失败场景与组合

Beam Search：beam reorder 与 EOS/length penalty 错位会选错序列；beam 多样性不足时更宽只保留相似候选。
Constrained Decoding：grammar 与 tokenizer bytes 不一致会误禁/误放；allowed set 为空必须可诊断；mask 后要重新稳定 softmax。
Tree Speculation：attention mask 未表达 parent ancestry、tentative KV 泄漏、rollback refcount 错误、未验证 token 提前 streaming 都会破坏语义。
packed node index 当成 position id，会让同 depth siblings 使用不同 RoPE angles；depth/parent metadata 必须与 Q/K/KV slots 同步。
将 grammar mask 放入 speculative proposal 后，target verifier也必须使用相同约束目标；否则 draft/target 分布不是同一任务。
Beam + speculative/tree 会把 beam×branches 状态乘起来，容量和 verifier batch 可能抵消 target calls 的减少。

Tree 的 target 串行调用减少并不自动等于端到端加速。每轮还要支付 proposal heads/draft feature network、tree construction/top-k、packed verifier tokens、mask/position metadata、tentative KV 分配和 rejected-node cleanup。接受深度低、树过宽、target verification 已 compute-bound、batch 被其他请求填满，或 custom tree kernel fallback 时都可能不划算。

Prefix Cache共享的是请求间已提交的相同 prefix；tree tentative nodes 是同一请求内尚未提交的候选，identity/lifetime 不同。PagedAttention可用 block refcount/copy-on-write承载 beam/tree 分叉，但 page table 本身不定义搜索或接受协议。