Bole 阅读笔记:为混合注意力大模型重新设计树状投机解码

笔记日期: 2026-08-09 作者: Zhongzhu Zhou 论文标题: Bole: Efficient Tree Speculation for Hybrid-Attention Language Models 论文作者: Li Wang, Yi Su, Xiabao Wu, Chiran You, Yongchao Liu, Zhan Qiu, Juelu Zhang, Jiajun Zheng, Fangxin Liu, Jie Zhang, Chen Tian, Chengying Huan(南京大学、蚂蚁集团、上海交通大学、北京大学) arXiv: 2608.01651 发表状态: 预印本(cs.DC),2026年8月

1. 为什么混合注意力模型需要一套”不一样”的投机解码方案

如果你这两年一直在关注 LLM 服务系统,树状投机解码(tree speculative decoding)大概已经算是一个”解决了的问题”。EAGLE-2、Medusa 这类方法用一个轻量级”起草模型”(drafter)一次性提出多条候选续写,组织成一棵树而不是一条直链;大模型(target model)用一个”祖先掩码”(ancestor mask)在一次并行前向里给整棵树打分——树上的每个节点只能看到自己到根节点这条路径上的祖先,看不到兄弟分支;然后标准的接受/纠正采样从树里挑出最长的可接受路径,同时保证输出分布和原模型完全一致。SGLang、AdaServe 这类系统已经把这套流程做成了生产级基础设施,报告的加速比动辄 5 倍。

但所有这些系统背后都藏着一个安静的假设:目标模型是一个纯全注意力(full-attention)Transformer。全注意力有一个对树验证极其友好的性质:因为每个 token 关注的 KV cache 只是简单地”拼接”增长,你可以把整棵候选树摊平成一条带祖先掩码的长序列,一次注意力调用就能打完全部分数,掩码矩阵负责挡住兄弟分支之间互相看见。一个节点需要的”状态”,就是”它所有祖先的 key/value 拼在一起”——因为独立向量的拼接没有顺序依赖,这天然就是并行的。

循环式线性注意力(recurrent linear attention)把这个性质彻底打破了。Qwen3.5、Kimi Linear 这类模型把全注意力层和门控 Delta 网络(Gated DeltaNet, GDN)层交替堆叠,GDN 不维护一个不断增长的 KV cache,而是把整段前缀压缩进一个固定大小的状态矩阵 StRdk×dvS_t \in \mathbb{R}^{d_k \times d_v}。这个状态是被”更新”出来的,不是”拼接”出来的:StS_tSt1S_{t-1} 的函数,外加当前 token 的贡献。这里没有一个”注意力矩阵”可以拿来做掩码,因为在经典意义上根本不存在注意力矩阵——递推关系本身就是全部机制。如果候选树里节点 B 是节点 A 的子节点,要算出 B 的线性注意力输出,你就必须先有 A 更新后的状态,而 A 的状态又依赖它父节点的状态,一路回溯到最后一个已提交的 token。这是一条彻头彻尾的串行依赖链——恰恰是全注意力树验证的那套技巧本想消除、却对这种结构完全无能为力的东西。

这篇论文——Bole——正是要补上这道缺口:让混合全注意力/线性注意力模型的树状投机解码,做到和纯全注意力模型一样并行、一样省显存,而且不改变计算内容本身(论文最核心的主张,并有定理支撑:它的重写方式与串行执行完全等价,不是近似)。

前置知识:读懂这篇论文之前你需要知道什么

自回归解码与”显存带宽受限”区间。 LLM 逐 token 生成文本时,每生成一个新 token 都要走完整个模型的一次前向计算。在较小的 batch size 下,这次前向通常是”显存带宽受限”的:GPU 大部分时间花在把模型权重从 HBM 搬到片上内存,真正用这些权重做计算的时间反而不多,因为每次权重加载只对应一个 token 的运算量。这就留出了大量闲置的算力(FLOPs)余量——投机解码正是要吃掉这部分余量。

投机解码,一段话讲清楚。 与其每次前向只生成一个 token,不如先用一个便宜的”起草模型”(常见做法是训练一个多 token 预测头,或者用一个独立的小模型)一次性提出好几个候选 token;然后让昂贵的”目标模型”一次性在所有候选上跑一遍前向——因为在显存带宽受限区间里,验证 kk 个 token 的成本几乎不比验证 1 个 token 贵多少,你相当于花 1 份钱验证了 kk 份 token。一套采样机制(最早由 Leviathan 等人和 Chen 等人提出)会接受一段与目标模型自身分布统计一致的候选前缀,保证输出分布不变,并在第一个被拒绝的位置用目标模型自己的分布重新采样一个纠正 token。

从直链投机到树状投机。 一条单链的候选序列会浪费机会:如果起草模型对两个都还算合理的续写犹豫不决,只提出一个就白白扔掉了另一个。树状投机(Medusa、EAGLE-2、Sequoia)让起草模型提出一棵分支树而不是一条直链,让目标模型用祖先掩码一次性验证整棵树——每个节点只”看得到”自己回到根节点的路径,看不到无关的兄弟分支。因为树上更多路径有机会和目标模型真实会生成的内容重合,树状投机每轮能接受的平均 token 数(mean accepted tokens, MAT)通常明显高于直链投机——EAGLE-2 报告每轮 4–5.5 个接受 token。

门控 Delta 网络(GDN)与门控 delta 规则,简要介绍。 GDN 是一种现代线性注意力变体(出自 Yang、Kautz 与 Hatamizadeh 的论文 “Gated Delta Networks: Improving Mamba2 with Delta Rule”),每个注意力头维护一个固定大小的状态矩阵 StRdk×dvS_t \in \mathbb{R}^{d_k \times d_v},而不是逐 token 的 KV cache。每一步它先施加一个衰减门 γt\gamma_t,算出新 value 与当前状态”预测值”之间的差量修正 utu_t,再用外积把这个修正叠加进状态:

S~t=γtSt1,ut=βt(vtS~tkt),\widetilde{S}_t = \gamma_t S_{t-1}, \qquad u_t = \beta_t (v_t - \widetilde{S}_t^\top k_t), St=S~t+ktut,ot=Stqt.(1)S_t = \widetilde{S}_t + k_t u_t^\top, \qquad o_t = S_t^\top q_t. \tag{1}

这里 qt,ktRdkq_t, k_t \in \mathbb{R}^{d_k} 是 query/key 向量,vt,ut,otRdvv_t, u_t, o_t \in \mathbb{R}^{d_v} 是 value/修正量/输出向量,γt>0\gamma_t > 0 是衰减门(保留多少旧状态),βt\beta_t 是写入门(新修正量施加的力度)。状态 StS_t 是”到 token tt 为止所有信息”的一个压缩、固定大小的摘要:序列长度上是 O(1)O(1),相对地全注意力 KV cache 会随提示长度 LL 线性增长成 O(L)O(L)。这正是混合模型存在的全部理由——用全注意力无限回溯的能力换成一个有界、维护成本低的状态——但也恰恰是为什么在 GDN 层上做树状投机很难:StS_t 字面意义上依赖St1S_{t-1},你无法像全注意力”绕开拼接依赖”那样绕开它。

图1(论文 Fig.2):混合注意力模型中的状态表示

图1(论文 Fig.2):全注意力每个前缀 token 保留一对 KV,随内存以 O(L)O(L) 增长(论文设定下每层每 token 约 2 KB);GDN 则把整段前缀压缩进一个固定大小的状态矩阵,序列长度维度上是 O(1)O(1)(每层约 4 MB)。两者都接入同一个下一 token 预测流水线,但它们的”历史表示方式”——也就是”给候选树多加一个节点”在计算上意味着什么——是根本不同的。

补充一点:为什么“拼接”和“递推”在并行化难易度上差那么多

刚接触这个领域的读者很容易产生一个误解:“不就是维护一个状态吗,拼接和递推不都是一种“先前面的接后面的”操作吗?”事实上两者在代数结构上完全不同。拼接是一个无状态操作:把向量 aabb 拼在一起得到 [a;b][a; b],这个结果只取决于 aabb 本身,跟它们是以什么顺序被计算出来的无关——你可以先算 aa 再算 bb,也可以先算 bb 再算 aa,甚至可以完全并行地两个一起算,结果都一样。递推则是一个有状态操作:St=f(St1,xt)S_t = f(S_{t-1}, x_t) 这个式子里,StS_t 的值字面上就写在自己前一个状态上,你必须先知道 St1S_{t-1} 是什么,才能算出 StS_t。全注意力碥就碥在它本身的歷史表示是拼接式的(KV cache 只是把每个 token 的 (kt,vt)(k_t, v_t) 存下来列表拼在一起,本质上没有任何一个元素依赖另一个元素才能算出),而 GDN 碥就碥在它碾成一个状态,就只能接受递推。这个差异不是实现细节,而是两种信息压缩方式在数学本质上的分差:一个保留历历可查的历史(代价是内存线性增长),一个只保留一个换代式的摘要(代价是失去直接回溯的能力,递推也因此天生带上了一条负载不斑的依赖链)。Bole 的全部技术工作本质上就是在回答一个问题:能不能在不改变递推本质的前提下,把它在一段已知数据上的计算重写成接近拼接那样可并行的形式。公式(2)的回答是:可以,但只能在树形拓扑这个特定结构下,因为只有树形结构才能保证引理3里那个关键的幂零性。

一个经常被问到的问题:能不能把树展开成一条“伪序列”再用普通注意力思路处理?

一个自然会想到的“接口拕展”思路是:既然全注意力可以把树展开成一条带掩码的长序列,为什么不把线性注意力也展开成一条伪序列,然后用普通递推就行了?问题在于:展开本身没有改变递推的本质——你只是把树的节点按某种遍历顺序排成一行(比如广度优先或深度优先),但每个伪序列位置仍然必须等它的真实父节点先被处理完才能处理——展开并没有删除任何依赖关系,只是换了个坐标系。因为线性注意力本身没有注意力矩阵可以捕获“哪些位置彼此相关”,展开之后你依然需要一个显式的机制来处理依赖。这恰好就是引理2、引理3在做的事——但它们做的不是重新排列节点,而是把递推本身的代数结构,重写成一个可以用有限阶多项式表达的线性系统。这个区别解释了为什么“直接把全注意力的技巧搬过来”这个直觉想法行不通,而必须对线性递推本身推一套完全新的代数。

2. 三个具体瓶颈,量化呈现

在给出解决方案之前,论文花了相当篇幅去测量现有系统在混合模型上到底是怎么失败的,而不是直接跳到自己的方案。这一点很重要,因为它告诉你哪些设计选择是承重结构,哪些只是附带效果。

图2(论文 Fig.1):面向混合模型的树验证动机

图2(论文 Fig.1):(a)在朴素串行递推下,线性注意力验证时间占前向总时间的比例随树从 8 节点扩展到 64 节点,从 4% 涨到 27%,而 Bole 的并行求解始终低于 5%。(b)逐节点状态快照占用几十 GiB 的 HBM,挤占 KV cache/在跑请求的容量;Bole 的因子化表示只占用一小块固定大小的空间。(c)“多验证几个 token 几乎不用额外花钱”的 roofline 区间是硬件相关的——在 A100 上转折点在约 128 行附近,在 GB10(NVIDIA DGX Spark,GB10 Grace Blackwell 超级芯片)上则要到 256 行——所以针对某一款 GPU 调好的固定验证预算,换一款 GPU 就可能是错的。

L1——串行递推。 现有引擎(例如 SGLang 原生的 GDN 树路径)逐节点应用公式 (1),因为每个节点都需要先拿到父节点更新后的状态才能算自己的。一棵有 TT 个候选节点的树因此需要 TT 步串行递推——即便每个候选 token 的输入(q,k,vq, k, v、门控值)在目标前向开始前就已经全部就位。也就是说,数据是并行可用的,但经典递推公式硬生生把计算变成了串行。在 Qwen3.5-9B/A100 上测得,这会让线性注意力层占前向总时间的比例从 4%(8 节点树)涨到 27%(64 节点树)——这种随规模上涨的行为直接和树状投机的核心目标背道而驰,树状投机本该让”更大的树(每轮验证更多候选 token)“接近免费。

L2——状态快照爆炸。 在采样确定哪个分支被接受之前,现有系统不知道哪个候选节点的状态最终会有用,于是保守地为每个树节点都物化一份完整的 GDN 更新后状态——每个头一份完整的 dk×dvd_k \times d_v 矩阵,每个节点都要一份。这是 Θ(THdkdv)\Theta(T \cdot H \cdot d_k \cdot d_v) 的空间开销,同时随树大小 TT 和(因为要对 batch 里每个请求都这么做)batch size BB 线性增长。具体来说,对 Qwen3.5-122B-A10B、树大小 T=32T=32,这种快照存储在 batch size 8 时占 36 GiB,batch size 16 时占 72 GiB——挤占了本可以用来扩大服务吞吐量的 KV cache 和在跑请求容量。

L3——硬件相关、不可迁移的验证容量。 论文用标准的 roofline 模型来解释:在验证行数 nn 较小时,一次 GEMM(MLP/注意力层内部的主要开销来源)是带宽受限的,所以多验证几个树 token 几乎”免费”——反正权重传输的钱已经花了,多加几行几乎不会拖慢总时间。超过某个硬件相关的阈值后,同一个 GEMM 会变成计算受限,延迟开始随行数近似线性增长。这个阈值不是模型单独决定的——它是模型、GPU 计算/带宽比、以及(对完整混合前向而言)当前 KV 长度和树形状共同决定的联合属性。一个一次性调好、写死在代码里的验证预算(大多数现有树投机系统的实际做法),在换一款 GPU、甚至在同一款 GPU 换一个 batch/KV 长度区间下,都很可能系统性地错误。

这三个限制归根到底来自三个横跨算法、内核、运行时的真实挑战:(1)祖先掩码无法消除的树耦合递推依赖(因为根本没有注意力矩阵可以掩),(2)在运行时知道哪个分支会被接受之前就已经出现的状态分歧,以及(3)依赖硬件和当前服务条件、而不是一个固定常数的单请求验证成本。Bole 用一套连贯的内核—运行时协同设计同时解决了这三个问题。

图3(论文 Fig.4):Bole 架构与章节路线图

图3(论文 Fig.4):Bole 的端到端流水线。一个硬件感知、批次共享的验证预算(图上方)反馈进每轮的批处理准备、起草模型前向、目标模型树前向和投机采样。在目标前向内部,Bole 的并行线性注意力树验证(闭式求解 + 值域分块内核)取代了串行递推,而因子化线性状态存储/提交机制取代了逐节点完整快照。整套流程集成在 SGLang 中。

再多说一句:为什么这不只是“用内存换时间”这么简单

在刚接触这篇论文时,我开始一度以为 Bole 的思路可能只是“把递推改成先算好再批量处理”这个常见思路——毕竟很多并行化技巧本质上都是“先预计算、再一次性应用”。但仔细看完引理1–3 之后会发现,这里真正难的地方不在于“提前算好什么”,而在于“提前能算好什么”。具体来说,如果没有引理2把递推重写成一个线性方程组,你根本无从知道“提前算好什么”;而如果没有引理3发现这个线性系统的系数矩阵 GG 在树形拓扑下幂零阶数有界,你就必须真正去求一个 T×TT \times T 矩阵的逆,这在 TT 很大时代价高得很难接受。换个角度看,这三个引理的关系很像一个推理链:引理1 告诉你“状态可以分解成因子”,引理2 告诉你“这些因子可以堆成一个共享的线性方程组”,引理3 告诉你“因为这个方程组碥巧是树形的,求解就变得很便宜”。三步号缺一不可,任何一个环节发生变形,整个“闭式解”就会堆回去,变成一个看似向前一步、但实际上比平面地递推还要昂贵的东西。

3. 推导闭式解:把串行递推变成一次矩阵求解

这是这篇论文的数学核心,值得慢慢展开,因为一旦看懂原理会觉得这个结果相当漂亮。

3.1 记号约定

考虑一棵有 TT 个节点的候选树,记 π(i)\pi(i) 为节点 ii 的父节点(π(i)=0\pi(i) = 0 表示”已提交根节点的子节点”,其树前状态记为 SpreS_{\mathrm{pre}})。对节点 ii,记 A(i)A(i) 为它的严格草稿祖先集合(根到 ii 路径上除 ii 本身之外的所有节点),记 dd 为所有节点 A(i)A(i) 大小的最大值——也就是树的最大深度。把每个节点的操作数堆叠成矩阵 Q,KRT×dkQ, K \in \mathbb{R}^{T \times d_k}VRT×dvV \in \mathbb{R}^{T \times d_v},第 ii 行分别是 qi,ki,viq_i^\top, k_i^\top, v_i^\top;同样把门控值堆叠成向量 β,γRT\beta, \gamma \in \mathbb{R}^T

定义”沿路径的累积衰减”:Pi=rA(i){i}γrP_i = \prod_{r \in A(i) \cup \{i\}} \gamma_r——本质上就是把序贯递推的衰减门沿从根到节点 ii 的整条链应用一遍。把它们堆叠成 P=(P1,,PT)P = (P_1, \ldots, P_T)^\top,令 DP=diag(P)D_P = \mathrm{diag}(P)Dβ=diag(β)D_\beta = \mathrm{diag}(\beta)

3.2 引理1:每条路径的状态都有闭式解

第一个关键步骤是发现:如果沿着一条从根到某节点的路径展开公式(1)的序贯更新,祖先的依赖关系会叠加成一个求和,而不是嵌套递归。具体来说:

S~i=PiSpre+jA(i)PiPjkjuj,Si=PiSpre+jA(i){i}PiPjkjuj.(3)\widetilde{S}_i = P_i S_{\mathrm{pre}} + \sum_{j \in A(i)} \frac{P_i}{P_j} k_j u_j^\top, \qquad S_i = P_i S_{\mathrm{pre}} + \sum_{j \in A(i) \cup \{i\}} \frac{P_i}{P_j} k_j u_j^\top. \tag{3}

为什么成立,逐步来看。Sπ(i)S_{\pi(i)} 展开一层,再乘上 γi\gamma_i:你会得到一个 γi\gamma_i 乘以 Sπ(i)S_{\pi(i)} 原本值的因子。沿着祖先链一路递归下去回到根节点,这些门控因子会连乘复合——Pi/PjP_i / P_j 恰好捕捉到了这一点:也就是”节点 jj 刚处理完时存在的状态”与”节点 ii 祖先链继续衰减到最后时存在的状态”之间的残余衰减量。按父节点先于子节点的拓扑顺序处理,用归纳法就能严格证明这一点。用一个查询 qiq_i 读出这个状态,得到:

oi=PiqiSpre+jA(i){i}PiPj(qikj)uj.(4)o_i^\top = P_i\, q_i^\top S_{\mathrm{pre}} + \sum_{j \in A(i) \cup \{i\}} \frac{P_i}{P_j}(q_i^\top k_j)\, u_j^\top. \tag{4}

直觉是:节点 ii 的输出等于”共享树前状态的一次衰减读出”,加上”每个祖先修正项 uju_j 的加权和”,权重恰好是一个类似注意力机制的、经过残余衰减比例缩放的 key 相似度项 (qikj)(q_i^\top k_j)

3.3 引理2:把修正量堆叠成一个线性系统

公式(4)里 uju_j 的系数正是论文里所称的 CijC_{ij}——所以把公式(4)沿全部 TT 个节点堆叠,得到 O=DPQSpre+CUO = D_P Q S_{\mathrm{pre}} + CU,其中 UU 堆叠了每个节点的修正向量 uiu_i^\top。但 UU 本身是递归定义的(通过 ui=βi(viS~iki)u_i = \beta_i(v_i - \widetilde{S}_i^\top k_i),而 S~i\widetilde{S}_i 又通过公式(3)依赖祖先的 uju_j)——所以还没有真正消除递归,只是把递归转移到了求解 UU 这一步。把公式(3)代入 uiu_i 的定义并整理,论文证明这个递归会变成一个线性系统

(I+G)U=R,Gij=βiPiPjkikj    (j 是 i 的严格祖先),R=Dβ(VDPKSpre).(5)(I + G)\, U = R, \qquad G_{ij} = \beta_i \frac{P_i}{P_j} k_i^\top k_j \;\; (j \text{ 是 } i \text{ 的严格祖先}), \qquad R = D_\beta\left(V - D_P K S_{\mathrm{pre}}\right). \tag{5}

RR 恰好是每个节点相对于”纯粹按树前状态衰减读出”会计算出的局部差量——也就是”如果没有任何祖先额外贡献任何东西,节点 ii 会施加的修正”。GG 捕捉的正是祖先额外流入的修正量:GijG_{ij} 只有在 jjii 的严格祖先时才非零(这是通过祖先掩码 MM^- 构造保证的),所以在父节点先于子节点的排序下,GG严格下三角的——对角线及以上的每一项都恰好是零。

3.4 引理3:为什么树状下三角矩阵求逆特别容易

这是让整套方法从”理论上优雅”变成”实际可用”的关键技巧。一个限制在树结构上的严格下三角矩阵(而不是一般的下三角矩阵,一般情况需要 O(T3)O(T^3) 高斯消元或 O(T2)O(T^2) 前向替换)具有极小的幂零阶数GmG^m 的非零项对应树上一条”mm 条边的祖先路径”,而这样的路径不可能超过树的最大深度 dd。所以 Gd+1=0G^{d+1} = 0,根据幂零矩阵的标准等比级数恒等式:

(I+G)1=m=0d(G)m.(6)(I+G)^{-1} = \sum_{m=0}^{d} (-G)^m. \tag{6}

由于典型的投机解码树都很浅(深度 d4d \approx 488,因为起草模型一般只往前看几个 token),这个求和只有寥寥几项——把看起来很昂贵的一般矩阵求逆变成了一个只有 d+1d{+}1 项的有限多项式,每一项都只是一次矩阵—向量乘法 GZGZ。这就是承重的核心洞察:不是说求解 (I+G)U=R(I+G)U=R 在一般情况下渐近上便宜(一般三角求解本来就是 O(T2)O(T^2))——而是对于树形GG,多项式次数是被树深度而不是树大小 TT 限定的,所以这个求解会并行化成 dd 轮独立的矩阵乘法,而不是 TT 轮串行更新。

3.5 定理1:完整闭式解

把三个引理拼起来:

O=DPQSpre+C(I+G)1Dβ(VDPKSpre).(2)O = D_P Q S_{\mathrm{pre}} + C(I+G)^{-1} D_\beta\big(V - D_P K S_{\mathrm{pre}}\big). \tag{2}

论文通过上面这条引理链证明并声明:这个等式是代数上精确的——不是对序贯递推的近似,而是把同样的计算重新排列成一种没有逐节点遍历的形式。等式右边的每一项(DP,Q,Spre,C,G,Dβ,V,KD_P, Q, S_{\mathrm{pre}}, C, G, D_\beta, V, K)在任何节点的输出被算出之前就是可知的,因为它只依赖(已经完全已知的)树结构、门控值以及 Q/K/VQ/K/V——都不依赖其他任何节点的输出。这正是把一个本质上串行的递推转化为一个可以同时对全部 TT 个节点求解的共享线性系统的关键所在。

3.6 Bole 并行验证算法,逐步展开

把这些内容整理成一个明确的编号流程(按第15a条要求——论文把这称为”设计1”,我们把其中隐含的步骤显式化):

算法1:Bole 闭式并行树验证(每层、每个头执行一次)

输入:候选树结构(父节点映射 π,祖先掩码 M⁻),
     堆叠的 Q, K, V ∈ R^{T×d},门控 β, γ ∈ R^T,树前状态 S_pre
输出:节点输出 O ∈ R^{T×d_v},修正因子 U ∈ R^{T×d_v}(供后续提交使用)

1.  计算路径累积衰减:P_i ← ∏_{r ∈ A(i)∪{i}} γ_r  (对每个节点 i)
    (沿每条根到节点路径走一次的简单前缀积;可以用倍增法
     在 O(d) 并行轮次内对全部节点算完,或用一次拓扑遍历 O(T) 算完
     ——都不是瓶颈)。
2.  构造 D_P ← diag(P),D_β ← diag(β)。
3.  计算 Gram 矩阵  KKᵗ ← K Kᵗ,  QKᵗ ← Q Kᵗ   (批量矩阵乘法)。
4.  应用祖先掩码,构造树交互矩阵:
        G ← D_β D_P [ (KKᵗ) ⊙ M⁻ ] D_P⁻¹        (严格下三角)
        C ← D_P     [ (QKᵗ) ⊙ M⁺ ] D_P⁻¹        (M⁺ = M⁻ + I)
5.  计算树前状态读出与局部差量:
        B₀ ← D_P (Q S_pre)
        R  ← D_β ( V − D_P (K S_pre) )
6.  通过有限 Neumann 级数(引理3)求解 (I + G) U = R:
        Z⁽⁰⁾ ← R;  U ← R
        对 m = 0 到 d−1:
            Z⁽ᵗ⁺¹⁾ ← −G Z⁽ᵗ⁾          # 每轮一次矩阵乘法,全部T个节点同时算
            U ← U + Z⁽ᵗ⁺¹⁾
        # d 轮之后,U 精确等于 (I+G)⁻¹R(引理3保证这是精确值,
        # 不是近似,因为 G^{d+1} = 0)
7.  计算最终读出:  O ← B₀ + C U
8.  返回 O(用于采样)和 U, P, K(用于第5节的因子化状态提交
    机制——不是完整的 T 节点状态快照)

上面每个步骤都是同时对全部 TT 个节点操作的——步骤 1、3–7 都是矩阵/张量操作,没有逐节点的串行循环(步骤1的前缀积是唯一存在拓扑依赖的地方,但它只有 O(d)O(d) 层深,不是 O(T)O(T) 层深,相对 Gram 矩阵和 Neumann 级数矩阵乘法而言开销很小)。这就是”解决 L1”具体意味着什么:朴素实现里的 TT 步串行链变成了一个固定的 dd 轮矩阵乘法循环,而 dd(树深度,一般 ≤ 8)远小于 TT(树大小,论文实验中可达 64+)。

接上一个小面前面公式推导中的遗漏点:“因子”到底算的是什么

另一个容易被跳过的细节是:公式(5)里的 UU 得到之后,它本身就是公式(2)里的那个 UU,也就是公式(9)里被叫作“树因子”的部分。这并不是一个巧合,而是整个设计的一个关键拉拉链——因为 UU 同时是(a)计算输出 OO 需要的中间量,以及(b)重构接受分支的真实状态时需要的全部信息,两个目的用的是同一份数据。这意味着设计1(并行求解)和设计2(因子化存储)并不是两个独立可以任意搭配的模块,而是建立在同一个代数对象上的两种使用方式。这也解释了为什么图4里“并行线性注意力树验证”和“因子化线性状态存储与提交”被并排放在目标模型树前向内部的同一个时间片:验证内核一旦输出 UU,这份 UU 就可以直接被状态提交机制复用,中间没有任何额外的格式转换或重新计算。

4. 把数学映射到硬件上:值域分块内核

一个闭式解只有能高效映射到 GPU 执行模型上才有用。这是”设计1”的后半部分,也是这篇论文真正体现系统功底的地方。

图4(论文 Fig.5):面向并行树验证的值域分块 CTA

图4(论文 Fig.5):一个协作线程组(CTA)处理一个宽度为 bvb_v 的值域分块 rr,覆盖整棵树的固定宽度 T~\widetilde{T}。它一次性从 HBM/L2 加载 SpreS_{\mathrm{pre}}Q/KQ/KV/β/PV/\beta/P 以及预先构造好的因子 G,CG, C,形成 QSpreQS_{\mathrm{pre}}KSpreKS_{\mathrm{pre}},在片上工作集内完整跑完有限 Neumann 递推,只把最终的 UUOO 分块写回 HBM。

为什么是值域分解,而不是树域分解? 朴素做法是把整棵树的固定宽度 T~\widetilde{T} 分配给一个 CTA,同时处理全部 dvd_v 个值通道。但这需要每个 CTA 占用 Θ(T~2+T~dv+dkdv)\Theta(\widetilde{T}^2 + \widetilde{T} d_v + d_k d_v) 的片上存储,而且每个注意力头只能调度一个 CTA——对现代 GPU 上几十个流多处理器(SM)而言,这样的并行度远远不够。论文抓住的关键点是:树交互因子 G,CRT×TG, C \in \mathbb{R}^{T \times T} 只依赖树结构、Q/KQ/K 和门控值——跟正在算哪个值通道无关。所以 Bole 把 dvd_v 个值列切分成 Nv=dv/bvN_v = \lceil d_v / b_v \rceil 个宽度为 bvb_v 的分块,G,CG, C 每层/每头只算一次(把 O(T~2dk)O(\widetilde{T}^2 d_k) 的 Gram 矩阵构造成本摊到全部 NvN_v 个分块上),再把每个值分块分配给一个独立可调度的 CTA。这直接把可用并行度乘以 NvN_v,而且没有重复计算昂贵的那部分。

为什么还要沿 key 维度 dkd_k 再分块? 单靠值域分块并没有约束住另一个大的中间量:每个值分块仍然需要跟完整的 key 维度 dkd_k 收缩,才能算出 QSpreQS_{\mathrm{pre}}KSpreKS_{\mathrm{pre}}。Bole 进一步把这个维度切成 bkb_k 宽的小块,通过单阶段软件流水线依次流过,这样任意时刻只需要当前的 Q/KQ/K 小块和对应的 SpreS_{\mathrm{pre}} 切片保持在片上,而不需要整个 dk×bvd_k \times b_v 的状态分块都常驻。

设计选择讨论——(bk,bv)(b_k, b_v) 的权衡。 这正是第15c条要求明确讨论的那种非平凡设计选择:

  • 更大的 bvb_v:把加载 G,CG, C、启动 kernel 这类固定的 per-CTA 设置成本摊到更多值通道上,但会扩大常驻的 R/Z/UR/Z/U 和累加器寄存器/共享内存分块,独立 CTA 的数量也会变少——在 SM 数量多的 GPU 上会拉低占用率。
  • 更小的 bvb_v:降低每个 CTA 的寄存器/共享内存占用,暴露更多可调度并行度,但会把固定的设置/因子加载开销在更多、更窄的分块上重复付出——如果 NvN_v 相对 SM 数量增长得太大,这部分开销就白白浪费了。
  • 更大的 bkb_k:张量核心上的 MMA(矩阵乘加)工作更粗粒度,累加轮次更少,但单次小块的瞬时共享内存占用更大。
  • 更小的 bkb_k:相反的权衡——占用更低,小块更多,循环开销更大。

论文没有声称存在一个通用最优解;它从一小组编译期配置中选取 (bk,bv)(b_k, b_v),通过实测在资源占用和并发度之间权衡,把单个 CTA 的工作集上界限定在 MCTA=Θ ⁣(T~2+T~bv+bk(T~+bv))M_{\mathrm{CTA}} = \Theta\!\big(\widetilde{T}^2 + \widetilde{T} b_v + b_k(\widetilde{T}+b_v)\big),同时每个请求可以创建 HNvH N_v 个独立可调度的 CTA(HH 是注意力头数)。这是一个经典的 GPU 内核张力——单个 thread block 的效率 vs. SM 占用率——针对这种具体的递推结构给出的具体化答案。

片上有限 Neumann 执行(把算法1第6步映射到硬件)。 一旦进入某个 CTA 内部,算法1第6步的 dd 轮 Neumann 级数完全在片上寄存器/共享内存中执行:Z(0)=R(r)Z^{(0)} = R^{(r)},然后 Z(m+1)=GZ(m)Z^{(m+1)} = -GZ^{(m)}U(r)+=Z(m+1)U^{(r)} \mathrel{+}= Z^{(m+1)},对 m=0,,d1m = 0, \ldots, d-1 循环,其中每个 GZ(m)GZ^{(m)} 沿节点归约维度分解成张量核心 MMA 片段。因为 dd(树深度)在编译期已知且很小,这个循环会展开成一个短而规整的张量核心操作序列,中间没有任何 HBM 流量——B0(r)B_0^{(r)}R(r)R^{(r)} 以及每个 Z(m)Z^{(m)} 都完全生于片上、死于片上;只有最终的 O(r)O^{(r)} 分块(供采样使用)和 U(r)U^{(r)} 分块(供状态提交使用)会离开 CTA。论文报告这套设计相比朴素串行内核,把实测 GPU 占用率提高了 3.6–7.9 倍,L1/共享内存吞吐提高了 1.8–4.1 倍——直接转化为第6节里测到的 3.4–7.7 倍完整验证核心加速。

5. 问题的另一半:因子化投机状态生命周期

解决了 L1(串行递推)并不会自动解决 L2(状态快照爆炸)——你完全可以想象一个并行求解器,最后仍然给每个节点物化一份完整状态。Bole 的第二个设计(“设计2”)通过观察引理1具体是如何表示状态这一点,规避了这个问题。

5.1 为什么完整快照是浪费的,形式化说明

回顾引理1:Si=PiSpre+jA(i){i}PiPjkjujS_i = P_i S_{\mathrm{pre}} + \sum_{j \in A(i) \cup \{i\}} \frac{P_i}{P_j} k_j u_j^\top。这说明每个节点的完整状态 SiRdk×dvS_i \in \mathbb{R}^{d_k \times d_v} 完全由(a)单一的共享 SpreS_{\mathrm{pre}},加上(b)它自己路径上一小组秩一更新因子 (Pj,kj,uj)(P_j, k_j, u_j) 决定。物化 TT 个完整状态需要 Θ(THdkdv)\Theta(T H d_k d_v) 的空间——如果 dk=dv=dsd_k = d_v = d_s,这在每头状态维度 dsd_s 上是平方级的。但只存储生成这些状态所需的因子,成本只有:

Θ(Hdkdv)+Θ(TH(dk+dv+1)).(9)\Theta(H d_k d_v) + \Theta\big(T H (d_k + d_v + 1)\big). \tag{9} X已提交状态XXXXXXXXXXXXXXXXXXX树因子,与 ds 线性相关而非平方相关\underbrace{\phantom{X}}_{\text{已提交状态}} \qquad \underbrace{\phantom{XXXXXXXXXXXXXXXXXXX}}_{\text{树因子,与 } d_s \text{ 线性相关而非平方相关}}

dk=dv=dsd_k = d_v = d_s 时,一个完整的单分支状态占用 Θ(Hds2)\Theta(H d_s^2) 个元素,而一个因子化节点只占用 Θ(2Hds)\Theta(2 H d_s)——节省比例会随着状态维度本身增长,而不依赖于任何具体树形状。这就是为什么表I里测到的节省幅度(T=100T=100 时,Bole 的因子只占 57–151 MB,而完整快照要 4.7–14.1 GB)会这么大:现代混合模型的 dsd_s 并不小。

5.2 状态生命周期,作为一个明确流程

图5(论文 Fig.4 架构图的一部分):整体流水线中的因子化生命周期

图5:因子化状态生命周期(§IV-D,架构图中”因子化线性状态存储与提交”这个框)位于目标模型树前向内部、与并行验证内核并列,并在投机采样确定接受路径之后,接入一次跨层批量提交步骤。

算法2:因子化投机状态生命周期(每请求、每轮执行一次)

阶段A——验证(对已提交状态只读,不写):
1. 从该请求/该层的 HBM 状态槽(标准 [H, d_v, d_k] 布局)中读取
   唯一的、不可变的已提交状态 S_pre。不要拷贝它。
2. 运行算法1,对每个候选节点 i 得到:
       - 输出 O_i               (供目标模型采样使用)
       - 修正因子 U_i           (保留,不展开成完整的 S_i)
   同时保留逐节点的 (P_i, K_i)——(P, K, U) 三者一起构成
   公式(9)所说的"树因子"。任何节点的完整状态都从未被构造出来。

阶段B——采样(决定哪个分支"获胜"):
3. 标准的树上拒绝采样(基于目标/起草模型的概率)选出一个终端
   接受节点 a,定义接受路径 P(a) = A(a) ∪ {a}。
4. 被拒绝的分支:直接丢弃它们的 (P, K, U) 因子。没有回滚,因为
   从来没有为它们分配过状态缓存槽(没有东西需要释放)。

阶段C——提交(只物化接受路径):
5. 按拓扑顺序压缩接受路径上的因子:
       K_a ∈ R^{|P(a)|×d_k},  U_a ∈ R^{|P(a)|×d_v},
       D_a = diag(P_a / P_j)_{j ∈ P(a)}
6. 通过引理1重构出唯一需要持久化的状态:
       S_new = P_a S_pre + (D_a K_a)^T U_a          # 公式(10)
   这恰好是一次批量张量核心矩阵乘法。
7. 用 S_new 就地覆写该请求的 HBM 状态槽(旧的 S_pre 在这次写入完成
   之前,仍然可以被任何正在进行的消费者读取——不需要额外的快照
   或写时复制槽位)。
8. 如果这一轮没有任何草稿 token 被接受,直接跳过步骤5–7:状态槽
   原样保留 S_pre 不变。
9. 把这次逐层提交折叠进*一次*跨全部循环层的设备端启动(因为接受
   路径 P(a) 是所有层共享的),与全注意力层的普通 KV cache 提交
   (由 SGLang 原生处理)并行重叠。

让这一套流程正确工作的关键性质是论文明确指出的一点:验证过程中 SpreS_{\mathrm{pre}} 从未被修改。每个候选分支读取的都是同一个不可变的 SpreS_{\mathrm{pre}},各自产生独立的因子;不需要为任何分支”分叉”出一份私有状态拷贝,也不需要回滚机制,因为压根就没有任何临时写入需要撤销。这正是第4步中”被拒绝分支可以直接丢弃”这一点在结构上成立的原因——没有可变状态需要撤销。

实测效果。 论文表I报告,在四种模型规模和两种树深度(T=100,200T=100, 200)下,这一机制相比完整快照把瞬时单请求状态显存降低了 82–99 倍。论文的消融实验(§6G,下面第8节讨论)把 Bole 端到端在线服务收益里最大的单一贡献,具体归功于这个机制——比并行求解器或硬件感知预算单独的贡献都大——因为释放出的几十 GiB HBM 直接转化为更多 KV cache 容量,进而提升前缀缓存命中率、减少连续批处理下的重复预填充。

回到一个比方:因子化跟数据库的“日志而非快照”思路很像

如果你精通数据库系统,可能会觉得这个因子化设计很熟悉。传统数据库的“write-ahead log”思路就是:不要每次修改都把整个表或页冒一份完整备份(那太昂贵),而是记录“与上一个已知状态相比发生了什么变化”,在需要时才回放日志恢复到目标状态。Bole 的因子 (P,K,U)(P, K, U) 本质上就是一条“日志记录”:它不存“结果”(完整状态),只存“每个候选分支相对于 SpreS_{\mathrm{pre}} 发生了什么变化”,只在采样确定接受路径后,才一次性“回放日志”得到真正需要持久化的新状态。两者的共同点是:日志/因子的体积远小于完整状态,因为它只记录“变化量”而不是“结果量”;不同点在于数据库的 WAL 是为了容灾(宕机后可重放),而 Bole 的因子是为了避免在接受路径还未确定时提前“提交”任何一个分支,两个场景都是在“计算/存储的开销”和“改变的不可逆性”之间寻找平衡,只不过一个面向磁盘,一个面向 GPU 寄存器/显存。

6. 设计3:把 L3 变成一个校准过的、批次共享的预算

设计1和设计2让验证变得便宜;设计3决定要买多少验证量,以及哪些候选节点有资格消耗这份预算。这正是直接解决 L3(硬件相关的容量)的部分。

为什么不能直接选一个固定的树大小? 论文把验证批次中全部 N=iniN = \sum_i n_i 个已选节点(每个请求的节点数 nin_i、已提交 KV 长度 LiL_i、树深度 did_i)所对应的完整目标前向延迟,建模为四项之和:

Tver(n,L,d)=Tfull(n,L)+Tlinear(n,d)+Tdense(N)+Tother(N).(11)T_{\mathrm{ver}}(\mathbf{n}, \mathbf{L}, \mathbf{d}) = T_{\mathrm{full}}(\mathbf{n}, \mathbf{L}) + T_{\mathrm{linear}}(\mathbf{n}, \mathbf{d}) + T_{\mathrm{dense}}(N) + T_{\mathrm{other}}(N). \tag{11}

每一项的伸缩方式都不一样:全注意力层的 KV 开销同时依赖节点数KV长度;线性注意力层(有了设计1–2之后已经变便宜)依赖节点数和树深度;共享的 MLP/投影层把全部 NN 个节点一起批处理,可能会在中途从带宽受限跨越到计算受限,跨越点由 GPU 决定。这里的设计张力在于:一个只看节点数的成本模型(“每个请求最多验证 kk 个节点”)实现简单,但忽略了同样的节点数 nin_i 在不同的 LiL_idid_i 和 GPU roofline 拐点下,实际延迟可能天差地别——这正是上面图2(c)量化出来的失败模式。

算法3:硬件感知的批次共享验证预算

离线阶段(一次性校准,针对每种执行配置 c
         = {模型、GPU/并行方式、batch size 区间桶,
            KV长度区间桶、树深度/模板区间桶}):
1. 在静态支持的 CUDA Graph 容量集合 G 上扫描候选总节点数 N。
2. 对每个 N ∈ G,测量 T_ver(N | c):完整混合目标前向的延迟(全
   注意力 + 线性注意力 + MLP + TP 通信 + graph capture 效应)
   ——不是某个组件的估算值,是真实的端到端数字。
3. 测量 T_dec(c):同一配置下普通逐 token 解码的延迟
   ("什么都不多做"的基线)。
4. 选出保持在可接受延迟开销容差 ε 之内的最大校准容量:
       B_ver(c) = max { N ∈ G : T_ver(N | c) ≤ (1+ε) T_dec(c) }   (12)
5. 把 B_ver(c) 存入一张按配置区间桶索引的查找表。

在线阶段(每一轮解码执行):
6. 根据批处理元数据(batch size、KV长度区间、正在使用的树模板)
   识别当前所处的配置区间桶 c。
7. 取出预先校准好的批次共享容量 B_ver(c)。
8. 对每个候选节点 v,沿其根到 v 的路径 P(v),计算其累积草稿概率:
       ρ(v) = ∏_{u ∈ P(v)} p_draft(u | π(u))
   (起草模型赋予"沿整条路径抵达 v"这件事的概率)
9. 保留每个在跑请求树的根子节点(公平性下限——不让任何请求
   完全得不到验证)。
10. 用固定形状的设备端 top-k + scatter 操作(保持在已捕获的 GPU
    迭代内部,不需要回到主机端),用整个批次中 ρ(v) 最高的节点
    填满*剩余*容量。
11. 因为一个子节点的累积概率不可能超过它父节点的累积概率
    (ρ 沿任意路径单调不增),每个请求被选中的节点集合会
    *自动*保持前缀连通——top-k 选择之后不需要额外的树修复步骤。

设计选择讨论。 步骤2的显而易见的替代方案是一个逐组件的解析成本模型——分别从 FLOPs/字节数估算 TfullT_{\mathrm{full}}TlinearT_{\mathrm{linear}}TdenseT_{\mathrm{dense}} 再求和。这样校准成本会低得多(不需要做扫描性能剖析),但论文实质上否定了这个做法:内核融合、张量并行通信重叠以及 CUDA Graph 分桶效应在实际中是非加性交互的,所以一个解析求和很可能系统性地预测错拐点的位置。这里的权衡是校准成本(一次性、离线,摊到整个部署生命周期里)vs. 预测精度(需要针对每个硬件/模型/并行方式组合实测)——对于一次部署、长期服务的生产系统来说,这是个合理的选择;但如果你需要在很多不同硬件型号之间频繁、临时地重新配置,而不想每次都重新做性能剖析,这个选择就没那么合理了。

值得关注的边界条件:公式(12)里的 ϵ\epsilon(可接受的延迟开销容差)是一个由用户设置的旋钮,直接在解码延迟开销和平均接受 token 数(MAT)收益之间做权衡——论文没有单独报告对 ϵ\epsilon 本身的敏感性扫描(只扫描了总树 token 数,图11),所以在不同部署 SLO 下,系统对 ϵ\epsilon 选得不好这件事有多敏感,论文没有直接给出答案。

对“为什么需要一个公平性下限”的补充说明

算法3第9步里“始终保留每个在跑请求树的根子节点”这一条県而未决的小细节,实际上在服务系统设计里非常重要。如果完全按累积概率 ρ(v)\rho(v) 排序选取 top-k,在高并发負载下完全可能出现某些请求的所有候选节点都排在全批次后面,一个 token 也得不到验证机会——这就退化成了普通自回归解码,对那个请求而言反而是一种递推延迟上的退步。固定保留根子节点,本质上是在全局最优化和每个请求的最低保障之间掉一个安全缸,代价是少量浪费一点预算在低概率分支上,换取对所有请求公平性的硬保证。这类“先保底、再优化”的设计模式在访少买升、排队系统里比比皮是,但在投机解码这类“根据概率选候补节点”的场景里并不总是被明确讨论,论文能把这一点写进算法里,是一个值得肯定的工程细节。

7. 生产系统集成:让理论在 SGLang 里真正跑起来

定理1加算法1–3描述的是一个理想化的计算过程;论文§V-B(约6.2千行 Python/Triton 代码)覆盖了让这套东西在一个真正持续批处理的生产级服务引擎里、没有主机端停顿地跑起来所需的工程工作:

  • 原生支持 CUDA Graph 的打包森林执行。 大小可变的每请求候选树被打包进一个设备端驻留的扁平化不规则森林(请求偏移量 + 父节点索引 + 紧凑的祖先元数据),而不是补齐成固定形状。一个请求感知的块映射(block map)随后只发出每个请求内部真正需要的验证分块——这不仅是为了内存布局方便,同时也减少了送进循环层内核的实际工作量。
  • 统一的跨层状态提交。 每个循环层把它的 (P,K,U)(P, K, U) 因子写入一个共享的、预分配好的 GPU 缓冲区的固定切片;采样为每个请求生成一份接受路径描述后,一次批量启动就消费全部循环层的描述符,一并计算出公式(10),而 SGLang 原生机制单独提交对应的全注意力 KV 条目。一次启动,而不是每层一次。
  • 无气泡的 GPU 执行。 选择、森林打包、目标模型执行、采样、路径压缩和状态提交被捕获成一条完整的 GPU 执行流,消除了这些阶段之间由主机端引发的同步停顿,让 CPU 调度器可以在 GPU 完成当前一轮的同时准备下一轮的工作。

这一节提醒我们:一个正确的闭式算法加一个高效的内核,是获得真正加速的必要条件,但不是充分条件——没有兼容 CUDA Graph 的打包方式和跨层批量提交,主机端开销和同步气泡很可能轻易吞掉设计1–2带来的大部分算法层面收益。

7b. 一个小例子:在3节点树上看闭式解如何工作

抽象的指标记号很容易掩盖一个事实:这个计算到底有多机械化。值得慢下来走一个小例子。假设已提交根刚刚产生了树前状态 SpreS_{\mathrm{pre}},起草模型提出一棵深度为2、三个节点的树:节点1(根的子节点),节点2和节点3(都是节点1的子节点,也就是彼此为兄弟)。那么 A(1)=A(1) = \emptysetA(2)=A(3)={1}A(2) = A(3) = \{1\},最大深度为 d=1d = 1(节点2、3都只有一个严格祖先)。

在朴素串行递推下,你会依次计算:从 SpreS_{\mathrm{pre}} 和节点1的 (k1,u1)(k_1, u_1) 得到 S1S_1;再从 S1S_1 和节点2的 (k2,u2)(k_2, u_2) 得到 S2S_2;接下来单独从同一个 S1S_1(不是 S2S_2——节点2和3是兄弟)和节点3的 (k3,u3)(k_3, u_3) 得到 S3S_3。一棵只有三个节点、深度为1的树,就已经需要三次序贯递推应用,而且节点3的计算甚至要等 S1S_1 完整成形才能开始。

在 Bole 的闭式解下,GG 是一个 3×33\times 3 严格下三角矩阵,只有一种结构模式:G21=β2(P2/P1)k2k1G_{21} = \beta_2 (P_2/P_1) k_2^\top k_1(节点2的祖先是节点1),G31=β3(P3/P1)k3k1G_{31} = \beta_3(P_3/P_1) k_3^\top k_1(节点3的祖先也是节点1),其余所有项——包括 G32G_{32}G23G_{23},因为2和3是兄弟,不是祖孙关系——都恰好为零。因为 d=1d=1,引理3的 Neumann 级数只有两项:(I+G)1=IG(I+G)^{-1} = I - G(因为 G2=0G^2 = 0——你可以写直接验证 G2G^2 需要一条两条边的祖先链,但这里最长的链就一条边)。所以 U=(IG)RU = (I - G)R 一次矩阵—向量乘法就能一次性算完三个节点——包括节点2、3。它们是兄弟关系,本来谁也不依赖谁,但串行写法就是因为记账方便把他们一前一后地算了,并不是因为任何真实的数据依赖。这就是设计1买到的核心价值:不是给出一个不一样的答案,而是让计算日程真正与树的真实依赖结构(父先子后)对齐,而不是盲目地把它抓平成任何一个固定遍历顺序。

7c. 与以往树投机和混合服务工作的关系

值得把 Bole 准确地放回它所池汲的两条研究脉络里看一下,因为论文的相关工作部分(§VII)写得相对紧凑。

与全注意力树投机(SpecInfer、DeFT、Medusa、EAGLE-2、Sequoia、AdaServe)相比。 这些系统解决的是一个完全不同的问题:对全注意力而言,任何节点需要的”状态”就是”它所有祖先的 KV 拼在一起”,而拼接本身没有任何计算依赖需要解决——你可以字面上直接构建好带祖先掩码的注意力矩阵,一次矩阵乘法解决。这一条研究脉络的全部贡献都在讨论选哪些 token(树拓扑、起草策略)和如何高效安排注意力计算(DeFT 的 flash-tree-attention 内核、AdaServe 的 SLO 感知预算)——而不是在解决一个序贯依赖,因为在状态层面根本不存在这种依赖。Bole 面对的问题是正交的:这一条研究脉络里选拓扑的思路(累积起草概率、SLO 感知预算)可以直接搬过来(Bole 的算法3就借用了这一框架),但验证内核必须完全重新发明,因为线性注意力的状态完全没有基于拼接的表示可以利用。

与不含投机机制的混合模型服务工作(Marconi、HLX、Pimba)相比。 Marconi 关注混合模型的前缀缓存——在共享相同 prompt 前缀的请求之间复用已经算好的线性注意力状态,这是一个跨请求复用问题。HLX 和 Pimba 是硬件/体系结构论文,专门为 Transformer-Mamba 混合模型设计加速器。这三篇都没有处理树形投机验证问题;它们处理的是相邻但完全不同的问题(跨请求状态复用、硬件专化),不是单个请求内部的分枝。

与最接近的前人工作 STree 相比。 STree 是唯一一篇在做结构上相似的事情——在候选树上组合对角状态空间模型(SSM)转移。但对角矩阵的转移可以交换并通过逐元素乘法进行合并(因为对角矩阵的”状态转移”在每个通道上彼此独立),这是一个比 GDN 使用的非对角、依赖 token 内容的门控 delta 规则弱得很多的代数结构(公式(1)里 ktutk_t u_t^\top 那项是秩一、跨通道的更新,而不是对角逐通道缩放)。论文直接声明 STree 的组合代数”对现代门控 delta 递推不适用”——这正是 Bole 必须从头推导一个新闭式解(定理1)、而不是直接移用 STree 方法的精确技术原因,这是一个真实、实质性的区别,而不是一个表面差异。

一张总览表:四个系统的定位对比

在进入详细实验数字之前,用一张表把四个被对比的系统在设计目标上的定位梳理一下,对理解后面的对比实验很有帮助:

系统递推处理方式状态管理方式验证预算决定方式对混合模型的支持
SGLang-AR无需处理(逐 token)普通单步递推不适用原生支持
SGLang-Tree逐节点串行递推逐节点完整快照固定预算原生支持但未专门优化
AdaServe逐节点串行递推(本论文移植)逐节点完整快照SLO 感知动态预算原开源不支持,作者自行移植
Bole闭式并行求解(定理1)因子化,只存接受路径硬件校准批次共享预算专门针对门控 delta 递推设计

这张表格最直接地说明了一个事实:前三个系统在“递推处理方式”和“状态管理方式”两栏里都一样——它们对待线性注意力的方式本质上是把全注意力的思路硬套过来,只是在预算分配上用了不同的策略。只有 Bole 在前两栏上都采用了专门针对线性递推重新设计的方法,这也是为什么它能在 L1/L2 两个瓶颈上同时取得量级提升、而不仅仅是线性改善的根本原因。

8. 实验评估:理论能不能兑现为吸吐量?

作者在两种计算/内存比例差异很大的平台上评估——四张 NVIDIA A100 80GB SXM4(每卡 312 FP16/BF16 TFLOP/s,2.04 TB/s HBM2e)和一台 NVIDIA DGX Spark / GB10 Grace Blackwell 超级芯片(256 FP16/BF16 TFLOP/s,128 GB 统一内存 LPDDR5x,273 GB/s)——两种平台上共四个 Qwen3.5 混合模型(4B 到 122B-A10B),对比 SGLang-AR(普通自回归解码)、SGLang-Tree(SGLang 原生树投机路径)以及 AdaServe(一个 state-of-the-art SLO 感知树投机系统,因为其开源实现本身不支持混合模型,作者不得不自己泡制一个支持版本)。

图6(论文 Fig.7):A100 上各 Qwen3.5 模型、各 batch size 的 MBPP 吸吐量

图6(论文 Fig.7):A100 上在 batch size 1/2/4/8、四个模型上的 MBPP(代码生成)离线解码吸吐量。Bole(红色)在所有配置下领先,且优势随 batch size 变大而拉大——值得特别注意的是,在 Qwen3.5-27B、batch size 8 时,两个投机基线都显存溢出,而 Bole 没有——这是设计2(82–99 倍瞬时状态节省)直接带来的结果。

头条数字。 在全館 28 个评估配置上,Bole 相对 SGLang-AR 取得幾何平均 2.74 倍加速,相对最强投机基线 1.26 倍。峰值加速在 GB10 上达到 4.72 倍(对 AR)和 2.06 倍(对 AdaServe),A100 上分别为 3.62 倍/1.39 倍。表III 还显示 Bole 在四个模型上都比两个投机基线拿到更高的平均接受 token 数(MAT,例如 Qwen3.5-4B 上为 6.38 vs. 6.35/6.29)——说明 Bole 的批次层最大化收益选节点机制(算法3步骤8–11)真的在做更好的候选树剪枝,而不仅仅是把同样的提案执行得更快。

为什么差差在 GB10 上更大。 相对最强投机基线的加速在 GB10 上从 batch size 1 时的约 1.11–1.14 倍升到 batch size 8 时的 1.70–2.03 倍,升幅比 A100 大。这跟论文的 roofline 论据(图2c)一致:GB10 计算/带宽比更低、内存是统一的 LPDDR5x,意味着状态物化开销在串行验证总时间里占的比例更大(论文单独测得在某个设定下,GB10 上是 86%,A100 上只是 38%),所以消除这块开销(设计2)在 GB10 上的回报自然比例更大。

图7(论文 Fig.9):Qwen3.5-27B 在各个生成任务上的吸吐量

图7(论文 Fig.9):在四种预测难度差异很大的任务上——代码(MBPP)、数学(GSM8K)、对话(ShareGPT)、摘要(CNN/DailyMail)——batch size 4 下的吸吐量。平均接受 token 数从 4.98(CNN/DailyMail,最难预测)到 7.08(GSM8K,最模板化/最易预测)不等,Bole 相对 SGLang-AR 的优势(A100 上 1.98–3.36 倍,GB10 上 2.85–4.12 倍)在这个整个难度区间上都成立,说明收益不是某一个简单任务的偶然结果。

真实 Agent 在线服务场景。 用泊松过程(Poisson process)重放 OpenHands 多轮编码 agent 会话(来自 NVIDIA 的 Open-SWE-Traces),Bole 相对 SGLang-AR 将平均 TTFT(首 token 延迟)降低 15.8%–64.3%,平均 TPOT(单 token 解码延迟)降低 37.6%–67.6%;相对最强投机基线 AdaServe,降低幅度分别为 61.3%–73.3%(TTFT)和 28.9%–49.9%(TPOT)。TTFT 提升的背后机制很微妙但重要:因为 agent 会话会一轮又一轮地延伸同一段共享前缀,驱逐一个已缓存的前缀(为临时的树验证状态腾地方)会迫使下一轮重新预填充,代价很高。Bole 因子化状态释放出的内存足以让它在 A100/GB10 上达到 90.8%/92.6% 的前缀缓存命中率——几乎追上普通自回归解码的 89.8%/93.3%——而 SGLang-Tree 和 AdaServe 只能达到 58%–70%,因为它们的状态快照挤掉了本可以用于缓存前缀的容量。TTFT 和 TPOT 在持续批处理下互相强化:更低的 TPOT 让批次槽位更早释放(缩短排队,帮助 TTFT),而更少的强制重新预填充减少了对在跑解码的干扰(帮助 TPOT)。

单独隴离内核的性能。 在 Qwen3.5-9B 上将值域分块验证器与 SGLang-Tree 的串行 delta 规则验证器单独对比(论文表VII),完整验证核心(不仅是算法1的矩阵乘法,而是整个融合后的 Gated DeltaNet 层)在 batch size 1 时加速 3.4 倍,到 batch size 16 时加速 7.7 倍,实测 GPU 占用率上升 3.6–7.9 倍,L1/共享内存吞吐上升 1.8–4.1 倍。值得注意的是,在某个测量设置下,单独状态物化这项就占了串行验证总时间的 38%(A100)和 86%(GB10)——生动地说明了同一个算法层面的瓶颈(逐节点 HBM 快照写入)在不同 GPU 上因为 HBM 带宽相对计算能力的不同而体现出差异很大的影响。

组件消融实验。 在在线 agent 任务上,将 Bole 的三个设计逐一叠加到 SGLang-Tree 基线上:单靠硬件感知预算就带来 +5.3%(A100)/ +10.7%(GB10)——这一步没改变串行验证器本身,只是把批次不再盲目扩张,一旦额外候选的成本超过其期望接受收益就停下来。加上并行闭式解求解器后,累加收益升到 1.21 倍(A100)/ 1.30 倍(GB10)——比单靠预算带来的提升大,因为所选的候选预算内仍然包含很多互相独立的分支,求解器移除了它们的父子执行链。最后加上因子化状态管理,累加加速比升到 1.59 倍(A100)/ 2.23 倍(GB10)——这是单项贡献最大的一步,在 GB10 上不成比例地更大,因为其内存容量更紧、命中率提升更明显。这个贡献大小的排序(预算 < 求解器 < 因子化)对后续混合注意力投机解码的工作如果从零开始应该优先优化什么,是一个有用的实证信号。

对验证预算的敏感性。 在固定树宽度/深度/batch size 的情况下(Qwen3.5-9B/MBPP,batch size 8)扫描每轮目标模型验证的树 token 总预算,吸吐量刚开始随预算上升(验证更多可能的分支,MAT 上升速度快于成本上升速度),达到顶峰后开始下降,因为附加验证工作开始超出其接受收益。关键是:两个平台的顶峰出现在不同的绝对预算上——A100 上是 128 个总树 token,GB10 上是 256 个——这直接实测存在时确认了:现有大多数系统实际上采用的单一、硬件无关预算,不可能在两种 GPU 上同时最优,验证了设计3”逐配置校准”方案的必要性。

8b. 一个常被忽视的细节:请求级创新与层级创新的叠加

值得为读者拉直一下本文中经常容易混淆的两个“层面”。引理3的幂零阶数上限是在单一层、单一个请求内的候选树上报告的;而真实部署中,一个混合模型往往有十几层交替的全注意力/线性注意力层,而服务引擎一般同时处理十几个到几十个并发请求。因子化状态提交(第5节)的“单次批量启动”设计,正是为了避免在层维度上也重复递推的开销——因为接受路径 P(a)P(a) 在所有循环层上是共享的,提交公式(10)只需要知道“哪个路径”,不需要知道“哪一层”,所以可以把所有层的提交折叠进一个 launch。与此同时,在请求维度上,算法3的批次共享预算保证了即使并发请求数很大,验证工作量仍然受到一个全局上限的控制。换句话说,Bole 的三个设计分别在三个不同维度上发挥作用:算法1 处理的是单层单请求内的节点并行度,算法2 处理的是跨层的提交开销,算法3 处理的是跨请求的全局预算——三者叠加在一起,才能真正消化一个生产级 SGLang 实例中真实存在的全部三维开销。

8c. 一个值得单列的数值对比:因子化到底节省了多少

为了让表I里的数字更有手感,值得具体推一下计算。对 Qwen3.5-9B(假设 HH 个头,每头 dk=dv=dsd_k = d_v = d_s),完整快照占用 THds22 bytes(FP16)T \cdot H \cdot d_s^2 \cdot 2\text{ bytes(FP16)},因子化表示占用大约 TH(2ds+1)2 bytesT \cdot H \cdot (2d_s + 1) \cdot 2\text{ bytes}。两者的比值约等于 ds/2d_s / 2。这意味着当 dsd_s(每头状态维度)达到几十到一百时,因子化表示就能带来几十倍的存储节省——这与表I报告的 82–99 倍实测数值在量级上相互印证,也说明了这个节省比例并不是一个固定常数,而是会随现代模型越来越大的每头状态维度而进一步扩大的——这对未来可能出现的、具有更大容量、更强回溯能力的线性注意力变体而言,是一个比现在更值得期望的技术方向。

9. 论文坦承的局限——以及一些被轻化的地方

这篇论文在一些边界条件上相对坦承,这一点值得肯定:它明确指出 Neumann 级数(公式(6))中使用的最大深度 dd 是在编译期固定的,(bk,bv)(b_k, b_v) 是从小组编译期配置中选取而不是连续调优的,以及设计3的离线校准在每次模型/GPU/并行方式发生变化时都需要重新做。

还有一些论文本身没有正面提及、但评估中未充分展开的点,值得提出来讨论:

  • 起草模型始终是模型自带的 MTP 头。 每一个实验都用的是每个 Qwen3.5 checkpoint 自带的同一种多 token 预测起草模型。这是个合理的默认选择,但回避了一个实际中很重要的问题:MTP 起草模型和单独训练的 EAGLE 风格起草模型在质量/延迟上的特征不一样,很难确认 Bole 的收益(特别是算法3里归功给收益最大化调度器的 MAT 提升)在一个更弱、校准方式不一样的起草模型上是否仍然稳定——那时累积路径概率 ρ(v)\rho(v) 可能是一个噴声更大的信号。
  • 所有实验都用未量化的权重。 造于量化(INT4/INT8 仅量化权重,或 KV cache 量化)如今已经是生产环境中的标配,而 Bole 的一大部分价值主张恰恰就是“释放 HBM 容量”,本来很值得展示一下这些内存节省会不会和权重/KV cache 量化已经节出来的内存叠加(也就是部分冗余)——也就是说,在一个已经量化过的部署环境中,Bole 还能保持多大优势,论文没有直接回答。
  • 只评估了 Qwen3.5 一个模型家族。 引言里提到的 Kimi Linear 作为另一个混合注意力模型,却从未被实际评估过。因为 Bole 的闭式解是专门针对门控 delta 规则(公式(1))推导的,而不同混合架构可能使用不同的循环方式(不同的门控结构、全注意力/线性层的交替比例不同),论文对”混合注意力 LLM”普适的声明,相对实际验证范围而言有些偏宽——定理本身针对门控 delta 递推,实验则仅限于建在这个递推上的一个模型家族。
  • 树深度固定为 8、top-k 固定为 4,自始至终。 这是个与 EAGLE-2 常见配置匹配的合理默认值,但论文自己的引理3(幂零阶数受树深度 dd 限定)意味着并行求解器的相对优势应该会随 dd 变大而变小(更多 Neumann 轮次,每一轮仍然跨节点并行,但轮次多了)。论文没有报告当最大深度远超过 8 时 Bole 的加速比如何衰减——而这恰好也是串行递推基线自己也会明显变差的区间,在大深度下的相对对比就这样被留下了空白。
  • **公式(12)里的延迟容差 ϵ\epsilon 从未被扫描,**正如第6节指出——只扫描了总 token 预算(图11)。既然 ϵ\epsilon 直接控制解码延迟开销和 MAT 收益之间的权衡,不同部署 SLO(交互式对话 vs. 批量代码审查)具体想要的 ϵ\epsilon 可能很不一样,对任何想在特定延迟 SLO 下部署 Bole 的人而言,这是一个值得关注的缺口。
  • 没有报告离线校准本身的时间开销。 设计3的 Design 3 需要对每个配置区间桶 cc 扫描完整目标前向在不同候选容量 NN 上的延迟,这个离线校准本身需要占用实际 GPU 时间。论文没有报告对一个新模型/新 GPU 组合运行完整校准具体需要多久,这对需要频繁重新部署、换硬件型号的团队而言是个实际运维成本,但论文并未定量。

具体改进建议

  1. 报告一个深度敏感性消融实验。 既然设计1收益的理论论据明确与深度相关(引理3),一个“加速比-vs-最大树深度 dd”的图表(保持 TT 近似不变)可以直接验证或限定论文的核心主张,而不是把它隐含地留在固定的 d=8d=8 里。
  2. 至少评估一个非 Qwen 混合模型(Kimi Linear,或如果 Kimi Linear 权重不容易拿到则选一个小型开源深合模型),即使规模缩小,也能实质性地支撑“混合注意力 LLM”的普适性声明。
  3. 至少在一个实验中与权重/KV 量化结合,明确 Bole 的内存节省是与量化已经节出的内存互补还是重叠——这直接关乎实践者在一个已经量化的生产部署中应当期望多大的 2–4 倍吸吐量提升。
  4. 直接扫描 ϵ\epsilon,或至少报告论文不同实验中片实际使用了什么范围的 ϵ\epsilon 值,让在特定延迟 SLO 下部署的读者有个可以参照的起点,而不需要自己从头重新跑一遍离线校准。
  5. 报告离线校准本身的 GPU-小时开销,以及重新校准的触发条件(每次换硬件型号、换模型、换并行方式都需要重新扫描吗,还是可以对已知配置进行轻量微调),这对需要在异构集集中轮换部署、或频繁改变批处理配置的团队而言是个实际运维成本问题。

对实践部署的一些启发

如果你目前在实际部署一个包含 GDN 层的混合模型(或在跟进 Qwen3.5/Kimi Linear 之类的开源权重),这篇论文最值得带走的三个具体结论是:第一,如果你已经在用 SGLang 的树投机路径服务混合模型,当前的递推处理很可能正在以逐节点串行的方式消耗你无谎可充分早就消耗掉的带宽/计算余量,值得把线性层逐节点递推列为优化安排的高优先项。第二,如果你发现树投机在高 batch/大树下容易 OOM,请先确认内存占用里有多少是候选节点状态快照——这往往是可以直接站在因子化思路上解决的、不需要换硬件的问题。第三,如果你同时在不同型号的 GPU 上部署(例如既有 A100 又有新一代 GB 系统),不要假设一个预算能在两边都最优——这篇论文的 Fig.11 实测结果已经直接验证了这个直觉。

10. 可复现性笔记

论文报告其实现集成在 SGLang release/v0.5.12 上,共≈ 6.2 千行 Python 和 Triton 代码,是一个实实在在、工程量不小的产品。本文讨论时,论文正文本身未标明开源代码仓库地址,所以想复现需要根据论文公式自行实现闭式解内核(上述算法1/图4)和因子化状态提交路径(算法2),或等待/申请作者(南京大学/蚂蚁集团)开源发布。评估硬件(四张 A100 80GB SXM4、一台 NVIDIA DGX Spark GB10)专业但并不稀有;评估目标是 4B 到 122B-A10B 参数量的 Qwen3.5 checkpoint。工作负载数据集(MBPP、GSM8K、ShareGPT、CNN/DailyMail,以及在线 agent 评估用的 Open-SWE-Traces)都是公开标准数据集。最难复现的一步大概是设计3里的离线校准扫描(公式(12)),因为 Bver(c)B_{\mathrm{ver}}(c) 是针对每一个硬件/模型/并行方式组合实测得到的,不是一个可以从论文里直接拄过来用的常数——在不同硬件上复现这套系统的人,需要自己重新完成一遍校准扫描,而不能直接经化论文报告的预算值(Fig.11 中 A100 上 128、GB10 上 256 tokens 都只针对那个特定设定)。

最后一个视角:这篇论文对“投机解码下一步”的启发

足够抽象地看,Bole 默默地提供了一个可能比它自己具体结果更持久的方法论贡献:面对一个看似内在串行的递推机制,先问“这个递推能否展开成一个闭式求和?”,再问“这个求和能否因为一些结构性稀疏性(比如树形拓扑导致的下三角幂零)而变成一个有限项的展开?”。这个两步提问模式很可能不仅适用于 GDN,也适用于其他带有“固定大小状态 + 时间递推”结构的序列建模方法(比如其他 SSM 变体、某些带状态压缩的基础模型变体)。下一步工作很可能就是问:对于不同的门控/衰减结构,是否存在类似的幂零递推关系,以及当不存在时(例如非对角、非秩一的状态转移),能否找到一个精度可控的近似闭式解来代替精确闭式解。

还有一个很值得提的观察角度是:这篇论文其实隐含地提供了一个对“什么样的递推结构值得这么么烦地重写”的反面案例:如果一个循环层本身计算廉价(相对于它节省下来的内存/带宽),那么把它重写成并行形式的性价比就会下降——引理3 的整套讨价过程(Gram 矩阵构造、多轮矩阵乘法)本身不是免费的,只是比逐节点 TT 步串行递推便宜。GDN 孤供处于这个“值得重写”的区间,是因为它的固定状态大小(dk×dvd_k \times d_v,通常比较大)使得单次递推的开销本身就不小,若再乘上树大小 TT 和批次大小 BB,串行开销就会失控得很快,重写成并行方式的回报自然很高。

11. 总结

Bole 是一个把生产瓶颈认真到愿意为底层计算推导一个精确重写、而不是仅仅诉于近似的很好例子。其核心洞察——一个树形递推依赖结构,一旦被正确重排成一个线性系统,其幂零阶数不是由树的大小决定而是由树的深度决定——真的相当优雅,且它与另一个洞察(线性注意力状态的路径因子化表示——引理1,使得精确但便宜的状态管理可行,根本不需要为每个候选分支物化完整快照)搭配得很好。实测结果(离线吸吐量相对自回归解码最大 4.72 倍,相对现有最强树投机基线最大 2.03 倍,在真实 agent 工作负载下 TTFT/TPOT 最大降低 67.6%/49.9%)具体够大,对任何真正在生产中部署 Qwen3.5、Kimi Linear 这类混合全注意力/线性注意力模型的人来说都很重要。论文坦承的局限——评估只限于基于同一递推的一个模型家族、没有与量化的交互实验、以及一个影响 SLO 敏感部署但从未扫描的延迟容差旋钮——并不会削弱其核心算法贡献,但确实意味着标题里的”混合注意力 LLM”现阶段应该被理解为”使用门控 delta 规则的模型”,而向其他递推形式(不同门控方式的 SSM、其他 delta 规则变体)推广,仍是一步自然的下一阶段工作,而不是这篇论文已经完成的内容。