TreeWY 阅读笔记:用树结构 WY 变换拆掉 Gated DeltaNet 混合模型投机解码里的内存墙

笔记日期: 2026-09-05 论文: TreeWY: Speculative Verification for Gated DeltaNet Hybrids 作者: Sneha Murthy Ghantasala(Thomson Reuters) arXiv: 2608.20961 状态: arXiv preprint,2026-08-21 提交

1. 这篇论文到底想解决什么问题

现在很多主流的开源大模型都是”混合”架构:不是每一层都用 softmax 自注意力,而是把一部分层换成线性注意力层——最常见的就是 Gated DeltaNet(GDN)块,它用一个固定大小的循环状态矩阵代替不断增长的 key/value(KV)缓存。这样做的理由很单纯:GDN 每条序列的显存占用是常数,不随上下文长度增长,而 softmax 层的 KV 缓存却是线性增长,长上下文时会吃掉大部分显存。论文研究的 Qwen3.5 系列就是按大约 3:1 的 GDN:softmax 比例混合,专门为了拿到这个显存优势。

问题在于:这个”整段前缀已经被压成一个固定大小状态”的特性,恰恰让投机解码在 GDN 层上变得很别扭。投机解码的做法是:先用一个便宜的”drafter”提前猜好几个候选 token,再让”真正”的目标模型一次并行验证所有候选,目标模型本来会生成的那段候选前缀被接受,其余的被丢弃,模型从被接受的位置继续往下解码。对普通 softmax 层来说,回滚几乎是免费的:KV 缓存就是一个列表,验证阶段把候选的 key/value 追加进去,被拒绝的部分只需要把指针往回移一下就丢掉了。但 GDN 层的循环状态在验证这一次并行前向传播结束时,已经被整个草稿序列覆写过了——如果不提前保存,根本没法恢复到草稿走到一半时状态长什么样。目前 vLLM 和 SGLang 里默认的做法,论文称之为全量状态快照(full-state snapshotting):在草稿序列的每一个位置都存一份完整的循环状态副本,这样不管最终接受到哪个位置都能直接跳回去。对一条长度为 kk 的链式草稿来说,这就是 k+1k+1 份完整状态副本;而对一棵有 NN 个候选节点的(树越宽,每个位置提供的候选越多,接受率也就越高)来说,就是 N+1N+1 份——而且关键是,这些快照没法在树的不同分支之间共享,因为每条分支的历史从分叉点开始就已经不一样了。快照显存直接随投机的激进程度线性增长,这就限制了草稿树能开多宽,也把 GDN 本该省下来的显存又吃了回去。

TreeWY 的贡献是让这个问题直接消失,而不是做一个显存/效果的权衡取舍。论文观察到,定义 GDN 循环的门控 delta 规则,可以被代数上重新推导为经典数值线性代数技巧——WY 变换(最早来自 Householder QR 分解,在 DeltaNet 的文献里已经被用来把训练阶段沿序列长度并行化)的一个特例:一整串秩一状态更新可以被写成一个三角线性方程组一次求解。TreeWY 把这个想法从”链”推广到完整的草稿,又把它的应用场景从”训练时并行”推广到”验证时并行”:草稿树里不管有多宽,每个节点”本应产生”的循环状态贡献,都通过对一个严格下三角方程组做一次前向替换来算出来;等验证阶段最终决定接受哪个节点,论文只从已经求解好的数值里重建那一个被接受节点的状态——从始至终都不需要为其余的 NN 个候选状态单独存一份。实际效果:论文在 vLLM 里实现了这套方法,并在 B200 GPU 上服务了同一个混合模型家族的两个规模(Qwen3.5 35B-A3B 与 397B-A17B):在同等负载下峰值 KV/状态显存降低 2–3 倍,在显存本来就是瓶颈的场景下,这换来了最高 1.49 倍的吞吐提升和大约 30–40 倍更低的 p99 首字延迟(TTFT),而在显存本不是瓶颈的场景下,只付出百分之几的吞吐代价。

Figure 1(自绘图):TreeWY 架构总览——GDN 投机验证的显存问题出在哪里,树结构 WY 变换又是如何把它去掉的。

2. 前置知识:读懂这篇论文需要先知道什么

2.1 自回归解码,以及为什么它慢

一个 Transformer 语言模型给定一个 prompt 后,是一次生成一个 token:跑一次前向传播,从输出分布里采样出下一个 token,把它接到序列末尾,再重复。生成阶段的每一次前向传播是显存带宽受限的,而不是算力受限的——模型必须把自己的每一个权重(以及 KV 缓存里当前的每一项)都从 GPU 显存读到片上内存才能产出一个新 token,所以 GPU 的绝大部分算力其实是闲着的。这正是投机解码要利用的根本低效之处:如果每次读权重都能顺带产出好几个 token 的有效工作量,而不是只产出一个,吞吐就能白得提升。

2.2 投机解码的正式定义

投机解码(Leviathan et al. 2023;Chen et al. 2023)通过把”提议”和”验证”解耦,解决了算力空闲的问题。一个便宜的drafter——可以是一个独立的小模型,也可以是挂在目标模型上的一个轻量级额外”头”(比如 MTP,multi-token-prediction 头)——自回归地提议 kk 个候选 token,这一步很便宜,因为 drafter 本身很小。目标模型接下来对全部 kk 个候选位置一次性跑一次前向传播(这一步正好用上了本来闲置的算力),同时算出它自己在每个位置本应输出的分布。因为验证是并行完成的,耗时大致跟正常生成一个 token 差不多,但如果所有候选都被接受,一次就能产出多达 kk 个 token。

让这个过程严格等价于目标模型自身采样分布(这就是”无损”的含义)的接受规则是:

paccept(xi)=min(1,ptarget(xix<i)pdraft(xix<i)),(1)p_{\text{accept}}(x_i) = \min\left(1, \frac{p_{\text{target}}(x_i \mid x_{<i})}{p_{\text{draft}}(x_i \mid x_{<i})}\right), \tag{1}

i=1i=1 开始逐位置判断,直到第一次被拒绝为止,此时从一个修正后的残差分布里重新采样一个 token 替换掉被拒绝的那个,后面的全部丢弃。关键是,模型的最终输出分布跟不做投机、直接普通解码得到的分布可证明完全相同——这里的”加速是免费的”意思是不牺牲质量换速度,只是把浪费掉的算力换成了被用上的算力。

这个想法的一个进阶版本是树结构草稿(比如 Medusa,以及本文研究的 NN 节点树):不再只提议一条长度为 kk 的候选链,drafter 提议一棵,某个位置可以分叉出好几个候选 next token,每个候选又可以继续往下分叉。验证一棵树,让目标模型可以接受它认可的任意一条”从根到某节点”的路径,平均下来每次验证能接受的 token 数比同样深度的单链要多,因为树给了 drafter 在每一步多次猜对的机会。代价是:验证一棵有 NN 个节点的树,计算成本(对 NN 个位置做注意力/状态计算,而不是 kk 个)更高——而如果每个节点都要单独存一份状态,存储成本也更高,这正是本文要解决的问题。

2.3 KV 缓存注意力(“正常情况”)

对普通 softmax 注意力层来说,每个生成步都会给当前 token 计算新的 key(ktk_t)和 value(vtv_t)向量,追加到不断增长的缓存里:TT 个 token 之后就是 K,VRT×dK, V \in \mathbb{R}^{T \times d}。对整段前缀的注意力是:

Attn(qt,Kt,Vt)=softmax ⁣(qtKtd)Vt.(2)\mathrm{Attn}(q_t, K_{\le t}, V_{\le t}) = \mathrm{softmax}\!\left(\frac{q_t K_{\le t}^\top}{\sqrt{d}}\right) V_{\le t}. \tag{2}

在投机解码下,这种缓存是天然”追加与截断友好”的:草稿 token 的 key/value 在验证阶段被追加进去,如果某个草稿 token 被拒绝,只需要把指针挪回到被接受的前缀之后,把后面全部丢掉即可——不需要重新计算,不需要额外存储,而且对一棵树来说,共享同一个公共前缀的分支甚至可以直接共享那段前缀的缓存条目。这就是为什么论文说”softmax 层在投机解码下依然便宜”——数据结构本身天生就对回滚友好。

2.4 线性注意力与 Gated DeltaNet(GDN)循环

线性注意力层用一个固定大小的矩阵状态 SRdv×dkS \in \mathbb{R}^{d_v \times d_k}(每个 head 一份)代替不断增长的 KV 缓存,每来一个 token 就循环更新一次,读出时是一次矩阵-向量乘法,而不是对整个历史做 softmax。本文所依赖的具体循环——门控 delta 规则(Yang, Kautz & Hatamizadeh;用于 Gated DeltaNet,ICLR 2025)是:

St=αtSt1(Iβtktkt)+βtvtkt,ot=Stqt,(3)S_t = \alpha_t S_{t-1}\left(I - \beta_t k_t k_t^\top\right) + \beta_t v_t k_t^\top, \qquad o_t = S_t q_t, \tag{3}

其中 kt,vt,qtk_t, v_t, q_t 是 token tt 通常意义上的 key/value/query 向量,αt(0,1)\alpha_t \in (0,1) 是标量衰减门(旧状态保留多少),βt(0,1)\beta_t \in (0,1)写入强度。逐项拆开来看:

  • αtSt1\alpha_t S_{t-1}:先把整个旧状态按标量因子衰减一下,再做别的事情。这让旧信息随时间淡出而不是无限累积——对一个固定大小的状态来说,如果没有衰减,固定大小的摘要迟早会饱和,所以这一步是必须的。
  • αtβtSt1ktkt-\alpha_t \beta_t S_{t-1} k_t k_t^\top:“delta”修正项。在写入新内容之前,这一层先移除(衰减后的)状态已经对当前输入 key ktk_t 做出的预测——也就是先算出 St1ktS_{t-1} k_t(根据已存储的一切,对 ktk_t 应该关联的 value 做的一次”猜测”),再减去这个猜测按 βt\beta_t 加权后产生的外积。这正是经典联想记忆学习里的 delta 规则:不是简单地把新信息盖在旧信息上面,而是只写入记忆已有预测和实际应有内容之间的误差。这正是为了防止状态被出现频率最高的某个 key 向量方向主导(这是不带这个修正项的朴素线性注意力的一个已知失效模式)。
  • +βtvtkt+\beta_t v_t k_t^\top:移除旧预测之后,写入与 key ktk_t 关联的真实新 value vtv_t,同样按 βt\beta_t 缩放。
  • ot=Stqto_t = S_t q_t:读出就是对当前 query 做一次矩阵-向量乘法——没有 softmax,没有归一化,每个 token 的成本是 O(dvdk)O(d_v d_k),而不是 O(Td)O(T d)

后面一切推导都依赖的关键结构性事实是:转移算子 Tt=αt(Iβtktkt)T_t = \alpha_t\left(I - \beta_t k_t k_t^\top\right) 是”标量衰减乘以秩一修正”,而且它在不同时间步之间不满足交换律(不同于纯标量衰减的状态空间模型,比如 Mamba2,那里的转移就是纯标量,整个循环可以直接压缩成一个累积乘积/累加和)。这种不可交换性正是为什么专门为 Mamba2 类模型设计的技巧(比如论文引用的先前工作 STree)没法直接搬到 GDN 上——你不能简单地把循环变成累加和,因为这里矩阵乘法的顺序是有讲究的。

2.5 为什么 GDN 会破坏投机解码的回滚

在正常(非投机)解码下,GDN 的显存故事非常好:每条序列一个状态块,不随上下文长度增长(本文研究的 35B 模型是 30 MiB,397B 模型是 90 MiB,而同样模型的 softmax KV 缓存在 32K 上下文时会长到 0.6–0.9 GiB)。但投机验证是对全部 kk(或树的 NN)个草稿位置一次性做并行前向传播,这意味着公式(3)里的循环必须一路算到草稿序列的末尾才行,而这时接受/拒绝的判断(公式 1)甚至还没做出来。如果目标模型最终只接受了 5 个草稿 token 里的前 2 个,你需要的是”第 2 个 token 之后”那一刻的状态才能继续往下解码——但前向传播早就把状态一路算到第 5 个 token 了,中间那些值早被覆写掉了。

今天生产系统(vLLM、SGLang)里默认的修复方案是全量状态快照:在验证这一次前向传播中,给草稿序列的每一个位置都存一份完整的 SS 副本,这样不管最终接受到哪个位置,都能直接索引回那个位置对应的快照。对一条 k=3k=3 的草稿链来说,这就是每条序列每个 GDN 层 k+1=4k+1=4 份完整状态副本——相对于单份已提交状态是 4×4\times 的显存倍数,也就是 35B 模型 120 MiB/序列、397B 模型 360 MiB/序列,而且仅仅是为投机解码这一件事额外付出的开销。更关键的是,这些快照没法在树的不同分支间共享:如果草稿是一棵树而不是一条链,每条分支的历史从分叉点开始就已经不同,所以树里的每一个节点都需要独立的完整快照,即使两个节点共享一段公共祖先前缀也不能复用。这正是本文要移除的那个机制。

2.6 WY 变换(为什么它是解决这个问题的对症工具)

WY 变换是数值线性代数里的一个经典技巧(最早是为了把一串 Householder 反射(用于 QR 分解)表示成单个低秩更新,而不是逐个字面矩阵乘法)。它跟线性注意力的关联,是先前研究普通(非门控)DeltaNet 的工作(Yang et al., “Parallelizing Linear Transformers with the Delta Rule over Sequence Length”)建立起来的:他们证明了一整串像公式(3)这样的秩一 delta 规则更新,可以被重新表达为对一组”有效”或”伪”值求解单个三角线性方程组,而不是字面上逐步重算这个循环。这之所以成立,是因为结构上每一步的更新只依赖更早步骤计算出来的量(这就是”三角”的含义),所以整条更新链可以被展开成一个线性方程组,未知数就是每一步的有效贡献,一次性通过前向替换求解,而不是 TT 步顺序循环。那项先前工作纯粹是用这个技巧来加速训练(希望沿整个序列长度并行,而不是逐步跑循环,后者在 GPU 上很慢)。TreeWY 的贡献是:注意到完全一样的结构性技巧——一串只往回看的依赖关系——在投机验证问题里同样存在,只是”链”要换成”树”,然后把针对树而不是链(这一步不平凡,因为树不是链)、针对门控delta 规则本身(具有 2.4 节描述的衰减乘秩一修正、不可交换的结构),而不是普通无门控版本或纯标量衰减的 Mamba2 所需要的线性代数给做出来。

Figure 2(自绘图):公式(2)所依赖的严格下三角祖先掩码——在 DFS 先序遍历下,父节点总是排在子节点前面,因此整棵草稿树的线性方程组可以一次前向替换求解。

3. 方法:树结构 WY 变换

3.1 第一步——把循环重写成带”伪 value”的衰减加性注意力

第一步纯粹是代数操作:展开公式(3)里的门控 delta 规则,然后重新分组。从

St=αtSt1(Iβtktkt)+βtvtktS_t = \alpha_t S_{t-1}\left(I - \beta_t k_t k_t^\top\right) + \beta_t v_t k_t^\top

出发,把 αtSt1\alpha_t S_{t-1} 分配到括号内的两项上:

St=αtSt1αtβtSt1ktkt+βtvtkt.S_t = \alpha_t S_{t-1} - \alpha_t\beta_t S_{t-1} k_t k_t^\top + \beta_t v_t k_t^\top.

再把带 ktk_t^\top 的两项归并到一起:

St=αtSt1+(βtvtαtβtSt1kt)kt=αtSt1+v~tkt,v~tβt(vtαtSt1kt).(4)S_t = \alpha_t S_{t-1} + \left(\beta_t v_t - \alpha_t\beta_t S_{t-1} k_t\right) k_t^\top = \alpha_t S_{t-1} + \tilde v_t k_t^\top, \qquad \tilde v_t \triangleq \beta_t\left(v_t - \alpha_t S_{t-1} k_t\right). \tag{4}

这一次重写是整套方法的核心,值得多花点时间理解它的直觉。把 St1ktS_{t-1}k_t 读作”更新之前,状态本身已经对 key ktk_t 预测出的 value”。那么 v~t\tilde v_t不是原始输入的 value vtv_t,而是 βt\beta_t 乘以真实 value 与状态自身预测之间的残差。换句话说,公式(4)说明门控 delta 规则本质上就是衰减加性注意力(St=αtSt1+(某个值)ktS_t = \alpha_t S_{t-1} + (\text{某个值}) k_t^\top,结构上跟最简单的线性注意力循环完全一样),前提是你写入的是正确计算出来的伪 value v~t\tilde v_t,而不是原始 value。delta 规则”不覆盖已有内容”的全部特殊行为,都被吸收进了 v~t\tilde v_t 的计算方式里——这个循环本身的形式完全没有保留任何特殊之处。

这为什么对投机解码重要?因为一串加性更新(St=αtSt1+v~tktS_t = \alpha_t S_{t-1} + \tilde v_t k_t^\top,反复迭代)可以直接展开成一个完全没有循环的闭式解:

St=(j=1tαj)S0+i=1t(j=i+1tαj)v~iki=gtS0+i=1tgtgiv~iki,gtj=1tαj.(5)S_t = \left(\prod_{j=1}^{t} \alpha_j\right) S_0 + \sum_{i=1}^{t} \left(\prod_{j=i+1}^{t}\alpha_j\right) \tilde v_i k_i^\top = g_t S_0 + \sum_{i=1}^t \frac{g_t}{g_i}\, \tilde v_i k_i^\top, \qquad g_t \triangleq \prod_{j=1}^t \alpha_j. \tag{5}

现在每个状态 StS_t 都只是对目前已经算出的伪 value v~1,,v~t\tilde v_1, \dots, \tilde v_t 的加权求和,权重是只依赖标量门 αj\alpha_j 的累积衰减比 gt/gig_t/g_i——指数上没有矩阵乘积,不像原始循环那样。如果不用逐步走循环就能拿到所有 v~i\tilde v_i,就可以直接从公式(5)读出任意 StS_t(或任意下游输出 ot=Stqto_t = S_t q_t)——这正是验证阶段需要的:你想要的是最终被接受的那个节点对应的状态,而不一定是最后算出来的那个。

问题(也是这不是件平凡事情的原因)在于:v~t\tilde v_t 本身依赖 St1S_{t-1}(见公式 4 的定义)——所以要算伪 value 貌似还是得先走一遍循环,这看起来是循环论证。3.2 节会说明论文如何打破这个循环性:不是顺序逐个计算每个 v~t\tilde v_t,而是构造一个单一的线性方程组,一次求解就同时给出所有的 v~t\tilde v_t

3.2 第二步——为整棵草稿树建立一个线性方程组

算法 1——构造并求解树验证方程组(非形式化伪代码)

输入:已提交状态 S0,草稿树的 N 个节点按 DFS 先序排列
      (保证每个节点的父节点索引严格小于它自己),
      每个节点的 (k_t, v_t, q_t, alpha_t, beta_t)
输出:伪 value 矩阵 V_tilde ∈ R^{N x d_v},既用于每个节点的
      输出,也用于最终的状态重建

1. 对每个节点 t = 1..N,计算累积衰减 g_t = 从 S0 到 t 沿祖先链上
   所有 alpha 的乘积          # 每个节点一个标量,前缀积,很便宜
2. 构造严格下三角、按祖先关系加权的 Gram 矩阵 G:
     G[t, i] = (g_t / g_i) * beta_t * (k_t^T k_i)   如果 i 是 t 的严格祖先
     G[t, i] = 0                                     否则(包括 i = t)
3. 构造右端项 R:
     R[t] = beta_t * v_t  -  beta_t * g_t * (S0^T k_t)   # "仅由 S0 本身预测出的部分"
4. 求解三角方程组  (I + diag(beta) G) @ V_tilde = R   得到 V_tilde
   # G 严格下三角意味着这一步只需要一次前向替换,即 O(N) 步顺序的
   # 小规模矩阵-向量运算,而不是 O(N) 次完整状态重算——这才是真正的计算收益所在
5. 每个节点 t 的注意力式输出现在都可以直接通过公式(5)的闭式解
   从 V_tilde 和 g_t 读出,全程不需要为任何节点单独物化状态。

拆解一下为什么这段伪代码是对的:关键结构性事实是,节点 tt 的伪 value v~t\tilde v_t 只依赖 St1S_{t-1},而根据公式(5),St1S_{t-1} 本身只依赖 tt 在树里祖先的伪 value(不是兄弟节点,也不是子孙节点——这正是为什么这是一个”树”问题,而不需要模型考虑全部 N2N^2 个节点对)。把节点按 DFS 先序排列,保证每个祖先的索引都严格小于其后代,这正是让论文公式(2)里的矩阵方程(对应上面伪代码第 4 步)变成严格下三角的条件——索引较大的未知数 v~t\tilde v_t 只依赖索引较小的未知数,反过来永远不成立。这种三角性,正是把”求解 NN 个相互依赖的未知数”从一个通用的 O(N3)O(N^3) 线性求解问题,变成 O(N)O(N) 步顺序前向替换(每一步都是一次小而便宜的更新,而不是一次完整的矩阵求逆)的关键——这正是 WY 变换这个名字所指代的数值线性代数技巧。

拿这个跟(DeltaNet 原始论文处理的 k=1,1,1,k=1,1,1,\dots 单分支情形)对比:在链的情形里,“祖先”就是”排在你前面的所有节点”,所以 GG 就是普通的按 token 顺序排列的下三角矩阵,而论文公式(2)的树结构公式是它的一个严格推广——链只是没有分叉的树,链上的 DFS 先序遍历就是从左到右的 token 顺序,所以当 N=kN = k 且没有分叉时,算法 1 会精确退化为 DeltaNet 原本的链式 WY 变换。这正是为什么说 TreeWY 不是一个专用的新奇技巧,而是对一个已有的链式并行化方法向树的真正推广。

Figure 3(自绘图):验证 → 求解 → 提交的数据流管线。一棵有 N 个节点的草稿树按 DFS 先序排列,构造出按祖先关系加权的 Gram 矩阵,对严格下三角方程组做一次前向替换,得到全部 N 个节点的伪 value 矩阵 V-tilde,最终只重建被接受节点的那一份状态。

3.3 第三步——提交时只重建被接受的那个状态

一旦目标模型的验证决定节点 aa 是其对应根到叶路径上最后一个被接受的节点,论文就直接从已经求解好的 V~\tilde V 出发,把公式(5)的求和限制在 aa 的祖先链上,来重建续接状态:

Sa=gaS0+iagagiv~iki,(6)S_a = g_a S_0 + \sum_{i \preceq a} \frac{g_a}{g_i}\, \tilde v_i k_i^\top, \tag{6}

其中 iai \preceq a 遍历 aa 本身及其全部祖先(即从根到 aa 的单条路径,而不是整棵树)。这就成为下一轮投机的新 S0S_0。这里值得深究的设计选择——为什么存 V~RN×dv\tilde V \in \mathbb{R}^{N \times d_v} 而不是直接存 NN 份完整状态——正是真正显存收益的来源:一份完整状态 SRdv×dkS \in \mathbb{R}^{d_v \times d_k}dv×dkd_v \times d_k 个元素(在 dv=dk=128d_v=d_k=128 时是 16384 个),而单行伪 value v~iRdv\tilde v_i \in \mathbb{R}^{d_v} 只有 dvd_v 个元素(128 个)——单节点存储量缩小了 dk=128d_k = 128 倍,这正是为什么存下全部 NN 个节点的伪 value(总共 N×dvN \times d_v)、外加提交时做一次便宜的重建,最终会比存 NN 份完整快照状态便宜得多。

显而易见的替代方案,以及为什么它效果不如这个: 有人可能会问,为什么不直接沿被接受的路径顺序重跑公式(3)的循环来算出 SaS_a,而要先解整棵树的方程组?答案是:验证本身是对整棵树一次性并行前向传播(这正是投机解码的全部意义所在——利用本来空闲的算力)——在这次并行传播算出每个节点的输出、并且目标模型的拒绝采样(公式 1)对所有节点跑完之前,你根本不知道最终会接受哪个节点。所以树里每个节点的输出 ot=Stqto_t = S_t q_t 必须在验证结果揭晓之前就已经全部算好,这正是公式(5)/算法 1 一次性为全部 NN 个节点提供的东西;只有重建那一步(公式 6,只针对最终被接受的那一个节点)才可以推迟到接受/拒绝判断做出之后。

这个设计仍然要付出一些代价的地方(是权衡,不是缺陷): 构造 GG 需要一个 N×NN \times N 的祖先掩码,以及所有祖先对之间的 NN 次成对点积 ktkik_t^\top k_i——这比一个只支持链的验证内核需要更多计算(朴素循环下,一条链每一步只需要跟紧邻前一个节点做一次点积,不过 DeltaNet 自己的链式 WY 变换出于同样的三角求解原因,本身也要付出类似的 O(N2)O(N^2) 成本)。论文第 4 节自己讨论的实现成本(见下文第 4 节)正好反映了这一点:一个非因果的祖先掩码目前没法被捕获进 CUDA graph(不像简单的因果/链式掩码那样),所以真正的树(w>1w>1)会退化到更慢、非 graph-捕获的执行路径——闭式数学从链到树的推广很干净,但让树跟链一样快的系统工程还没有解决,论文对此说得很明确,这是未来工作,而不是号称白得的午餐。

4. 实现与实验设置

TreeWY 是作为 vLLM 主分支的一个 fork 实现的(写作本文时尚未上游合并),通过两个 SpeculativeConfig 选项暴露出来:mamba_state_commit="reconstruct" 把默认的 "store_all" 快照策略换成 3.3 节的提交时重建策略,draft_tree_widths 通过指定每一层的分叉因子,把一条普通链变成一棵树。一条链(或者一棵每层分叉因子都是 1 的”树”,本质上就是伪装成树的链)会在一个融合、可被 CUDA graph 捕获的 Triton 内核里完成验证与提交——这正是论文报告其头条吞吐/TTFT 数字的场景。而一棵真正的树(w>1w>1)需要 3.2 节描述的非因果祖先掩码,目前没法被 CUDA graph 重放;vLLM 会在这种情况下让整个模型退化到”分段(piecewise)“(非 graph)执行,论文特别指出这个代价远大于掩码计算本身应有的开销,因为这同时把 GDN 混合器内核也从 graph 捕获里挤了出去。这正是为什么论文明确把树宽度的结果(论文表 2,复述于下文 5.3 节)报告为”已启用且正确,而非提速”——这是一次诚实、少见的坦诚披露,表明系统工程尚未跟上数学推导的进度。

树还有一个调度上的小麻烦:一棵树的 DFS 前缀跟完整的树是不同的拓扑结构(不像链,链的任意前缀本身仍然是一条更短的合法链),所以某个请求的完整树如果放不进当前步的 token 预算,这一步会整个跳过投机,而不是被截断成一棵更小的树。所有实验都用贪心(greedy)草拟与验证,正确性用两种方式检查:闭式求解跟字面逐节点循环相比,fp64 下相对误差约 101510^{-15},fp32 下约 10710^{-7};生产环境的 bf16 内核跟这个 fp64/fp32 参考在 bf16 容差范围内一致——作者明确指出这意味着输出 token 流跟快照基线不是逐位相同的(不同计算顺序下的浮点非结合性),所以他们通过对比接受长度是否跟一个共享的无投机参考一致来判定正确性,而不要求两种投机实现之间逐位输出完全相同。

模型与硬件。 同一个混合模型家族的两个规模:Qwen3.5-35B-A3B(30 层 GDN + 10 层 softmax,张量并行度 1)与 Qwen3.5-397B-A17B(45 层 GDN + 15 层 softmax,张量并行度 8),都跑在 B200 GPU 上(每卡 178 GiB 显存),主扫描实验都固定用深度为 3 的 MTP 链式草稿。基线:storeall(vLLM 当前默认的全量状态快照)是主要对比对象;ReplaySSM(延迟物化 + 一份秩一缓存)在论文附录 D 里单独用同一套 35B 扫描做了对比。工作负载: 一共六种——ShareGPT、spec-bench、BurstGPT,以及三种合成场景(balanced-chat、generation-heavy、summarize-heavy)——在 gpu_memory_utilization(gmu){0.6,0.75,0.9}\in \{0.6, 0.75, 0.9\} 上扫描,其中五种额外在最大并发数 {1,8,32,64,128,256}\in \{1,8,32,64,128,256\} 上扫描(BurstGPT 换成扫描泊松到达率)。为什么这样设计是合理的: 同时扫描显存预算(gmu)和实际负载(并发数)恰好是暴露”依赖场景”效应所需要的——论文最核心的实证主张是,TreeWY 的收益取决于显存是否已经是瓶颈,而只有同时展示出显存受限区间(收益很大)和非显存受限区间(付出百分之几的代价),而不是报告一个平均掉一切的单一数字,才能有说服力地证明这一点。

5. 实验结果与分析

5.1 核心机制:释放出来的 KV/状态余量

Figure 4(论文原图 Fig.1):峰值 KV 缓存占用率(占池比例)随负载变化,按模型和显存利用率区分。store-all 基线的显存池很早就被打满,只能把打不进的请求排队;TreeWY 在相同负载下能多留出 2-3 倍的余量。

上面的 Figure 4(论文原图 Fig. 1)是后续所有结果都建立在其上的机制:在所有测试的并发水平下,storeall 的峰值 KV 缓存占用都比 TreeWY 更早地爬到显存池 100% 的上限,因为 storeall 每一个 N+1N+1 份的快照状态都在跟 KV 缓存、正在处理的批次请求争抢同一块 HBM 池。具体来说,在测得的最紧张预算下(35B 在 gmu 0.6,397B 在 gmu 0.75),TreeWY 在五种工作负载上的几何平均 KV 缓存降幅,从低并发时约 1.0 倍(此时谁都还没被显存卡住)到高并发时大约 2.4–5.6 倍不等——高并发正是释放出来的显存真正开始起作用的地方。

5.2 吞吐与延迟:一个真正依赖场景的收益,而不是固定倍数

论文最见诚意的一步是它的 Figure 3(下面用论文自己的分解数据来概念性复述),把每一个比率指标都按 storeall 基线在该具体(工作负载、并发数、gmu)点上是否真的已经打光了 KV 余量来拆分,而不是报告一个整体平均数。在 35B 上,105 个测量点里有 31 个是显存受限的;在那里,TreeWY 达到 1.15 倍吞吐、2.94 倍更低的 p99 TTFT、1.17 倍更低的平均端到端延迟,代价是每 token(TPOT)效率变成 0.83 倍——也就是说,单个 token 产出确实慢了一点,但系统能接纳进多得多的并发请求、批次变得更大,整体吞吐和延迟反而都变好了。在 397B 上,70 个点里有 17 个显存受限,展现出同样的定性形状但幅度更小(1.06 倍吞吐,1.66 倍更低 TTFT)。在剩下的”余量充足”的点上(基线本来就没被显存卡住),由构造方式决定,每个指标都在个位数百分比误差内跟基线持平——TreeWY 在那里没赢什么,因为原本就没什么可赢的,但也没有任何实质性的损失。

为什么这种拆分很重要,如果省略掉会掩盖什么: 如果论文把全部 175 个(工作负载、并发数、gmu)点的结果平均成一个数字,显存受限区间里那些漂亮的收益(高达论文表 1 里的 30–40 倍 TTFT 提升)就会被数量更多的”余量充足”点稀释掉,得到的头条数字会既低估这个技术在它真正设计针对的场景里的价值,同时又让它看起来像是一个比实际情况更大的无条件收益。论文自己的 p99 TTFT 数字把这个量级体现得很直观:在 128 并发、35B、gmu 0.6 时,TreeWY 大约快 40 倍(683ms vs 27489ms)——但如果不知道这恰好是 storeall 的显存池已经打满、正在排队请求,而 TreeWY 还没有的那个点,这个数字本身是没有意义的。

Figure 5(论文原图 Fig.2,概念性复述):p99 首字延迟随负载变化。随着并发升高,storeall 的 KV 池打满,TTFT 在对数坐标上急剧爬升;TreeWY 在同样负载下维持较低 TTFT,只有在显存确实充裕时两者才会趋同。

5.3 树宽度:变得可负担,而非变得更快

Figure 6(论文表 2,重绘为图):固定深度为 3 时,树宽度(分叉形状)与接受长度的关系。store-all 基线每请求的块开销随树规模 N 线性增长(从 N=3 时的 4 块到 N=39 时的 40 块),而 TreeWY 无论宽度多少都固定为一块。接受长度随宽度上升而持续增长,但增速逐渐趋于深度 3 带来的上限。

论文表 2(上面 Figure 6)大概是更具新意的实证结果,因为这是一个快照式基线在超过某个宽度之后根本负担不起、而不仅仅是”更贵”的情形:从一条链(形状 (1,1,1),N=3N=3)变成一棵完全分叉的深度 3 树(形状 (3,3,3),N=39N=39),接受长度从 3.24 上升到 3.58(一个真实的质量收益——每次验证接受的 token 更多,因为树给了 drafter 在每一深度多次猜中的机会),而 storeall 的快照开销在同样的区间里增长了 10 倍(每请求每个 GDN 层从 4 块到 40 块),TreeWY 的存储开销则不管 NN 是多少都保持在一块。这是”可负担性”论证,跟 5.1 节”同样负载、更省显存”的论证不同:不是说 TreeWY 让一个固定工作负载变便宜了,而是说 TreeWY 让原本根本负担不起的配置(宽树)变得可行。论文谨慎地补充说明,呼应上文第 4 节:这种可负担性目前还不是吞吐上的收益,因为更宽的树每一步要往目标模型里多推 N+1N+1 个 token,目前又跑在更慢的非 graph 捕获路径上——宽度扫描展示的是能力解锁,而不是延迟改善,把两者混为一谈会夸大这个结果。

5.4 与 ReplaySSM 的对比

附录 D 里跟 ReplaySSM(最接近的、只支持链的并发工作,它把状态物化推迟到周期性 flush,而不是解另一个三角方程组——回顾 2.6 节/前述先前工作讨论,在它的验证窗口内部,它解的其实是跟 TreeWY 相同的三角方程组)的正面对比,是论文最自我批判的一节:在同一套 35B 扫描下,两种方法释放出的 KV 余量相当(TreeWY 是 ReplaySSM 峰值占用的 0.96 倍——基本打平),但 ReplaySSM 在测试的每个并发水平上原始吞吐都更高(比如相对各自基线是 1.12–1.20 倍,TreeWY 是 0.99–1.08 倍)。论文自己对这个差距的归因——认为差别在于状态何时被物化(ReplaySSM 延迟写入,TreeWY 每次验证都提交),而不是共享的三角求解数学本身——被作者明确标注为”一种归因,而非已验证的结论”,因为他们自己并没有实现一个延迟写入版本的提交策略来验证这个差距是否真能被补上。这是一个相当坦诚的承认:TreeWY 目前的实现,并不是在每个维度上都简单地优于最接近的竞争方案——它在可负担性和树支持上占优,但目前在原始链式吞吐上,输给了一个写入延迟更激进的系统。

5.5 一个数值小例子:三节点链上手工过一遍公式(4)-(6)

抽象公式看多了容易失去手感,这里用一条最简单的三 token 链(没有分叉,方便手算)把公式(4)-(6)走一遍,建立一点直觉。假设 dv=dk=1d_v = d_k = 1(标量简化版,方便纯手算,真实实现里 kt,vtk_t, v_t 都是向量,SS 是矩阵,但标量版的逻辑完全一致),取 S0=0S_0 = 0(序列刚开始,状态是空的),三个 token 的 (αt,βt,kt,vt)(\alpha_t, \beta_t, k_t, v_t) 分别是 (0.9,0.5,1,2)(0.9, 0.5, 1, 2)(0.9,0.5,1,3)(0.9, 0.5, 1, 3)(0.9,0.5,1,1)(0.9, 0.5, 1, 1)

逐步走原始循环(公式 3)作为对照组:

  • t=1t=1: S1=α1S0(1β1k12)+β1v1k1=0+0.5×2×1=1S_1 = \alpha_1 S_0 (1-\beta_1 k_1^2) + \beta_1 v_1 k_1 = 0 + 0.5\times2\times1 = 1
  • t=2t=2: S2=α2S1(1β2k22)+β2v2k2=0.9×1×0.5+0.5×3=0.45+1.5=1.95S_2 = \alpha_2 S_1(1-\beta_2 k_2^2) + \beta_2 v_2 k_2 = 0.9\times1\times0.5 + 0.5\times3 = 0.45+1.5=1.95
  • t=3t=3: S3=α3S2(1β3k32)+β3v3k3=0.9×1.95×0.5+0.5×1=0.8775+0.5=1.3775S_3 = \alpha_3 S_2(1-\beta_3 k_3^2) + \beta_3 v_3 k_3 = 0.9\times1.95\times0.5+0.5\times1=0.8775+0.5=1.3775

再走一遍公式(4)-(5)的闭式解法: 先算伪 value:v~1=β1(v1α1S0k1)=0.5×(20)=1\tilde v_1 = \beta_1(v_1-\alpha_1 S_0 k_1)=0.5\times(2-0)=1v~2=β2(v2α2S1k2)=0.5×(30.9×1)=0.5×2.1=1.05\tilde v_2 = \beta_2(v_2-\alpha_2 S_1 k_2)=0.5\times(3-0.9\times1)=0.5\times2.1=1.05v~3=β3(v3α3S2k3)=0.5×(10.9×1.95)=0.5×(11.755)=0.3775\tilde v_3=\beta_3(v_3-\alpha_3 S_2 k_3)=0.5\times(1-0.9\times1.95)=0.5\times(1-1.755)=-0.3775。累积衰减 g1=0.9,g2=0.81,g3=0.729g_1=0.9, g_2=0.81, g_3=0.729。代入公式(5)算 S3S_3:S3=g3S0+g3g1v~1k1+g3g2v~2k2+g3g3v~3k3=0+0.81×1+0.9×1.05+1×(0.3775)=0.81+0.9450.3775=1.3775S_3=g_3 S_0+\frac{g_3}{g_1}\tilde v_1 k_1+\frac{g_3}{g_2}\tilde v_2 k_2+\frac{g_3}{g_3}\tilde v_3 k_3=0+0.81\times1+0.9\times1.05+1\times(-0.3775)=0.81+0.945-0.3775=1.3775。跟原始循环算出的 S3=1.3775S_3=1.3775 完全一致,验证了公式(4)-(5)确实是原始循环的等价重写,而不是近似。

这个例子也直观展示了公式(2)/算法 1 描述的三角方程组是怎么一次性把 v~1,v~2,v~3\tilde v_1,\tilde v_2,\tilde v_3 都解出来的:注意 v~2\tilde v_2 依赖 S1S_1,S1S_1 又依赖 v~1\tilde v_1;v~3\tilde v_3 依赖 S2S_2,S2S_2 又依赖 v~1,v~2\tilde v_1,\tilde v_2——这条依赖链正是矩阵 GG 里非零元素排布的样子(严格下三角),前向替换一步步解出 v~1v~2v~3\tilde v_1 \to \tilde v_2 \to \tilde v_3,跟上面手算的顺序完全对应,只是在真实实现里这是通过一次矩阵方程求解完成的,而不是像这里为了展示直觉而写成的分步算术。

5.6 工程落地时值得关注的几个点

如果有团队想在自己的服务栈上评估或复现这套方法,这里是几条从论文实现细节(第 4 节)里能提炼出来、值得在自己的落地评估中重点关注的实践要点:第一,mamba_state_commit="reconstruct"draft_tree_widths 这两个配置项在语义上是独立正交的——即使暂时不想承担树验证的 graph-capture 工程代价,单独把 reconstruct 打开、draft_tree_widths 保持链式(全 1),已经能拿到论文报告的全部链式吞吐/TTFT收益,这是一个风险更低的渐进式落地路径。第二,论文强调正确性检查是”接受长度是否匹配”,而不是”token 是否逐位相同”——这意味着任何要复现或迁移这套方法的团队,在自己的回归测试体系里也需要采用同样的相对性正确性标准,而不能简单套用要求逐位输出相同的传统单测框架,否则会得到大量误报的”失败”。第三,论文所有实验都关闭了前缀缓存(见第 6 节的局限性讨论)——如果目标生产环境高度依赖前缀缓存(比如多轮对话、共享系统 prompt 的场景),在评估这套方法之前应该先验证 reconstruct 策略跟前缀缓存开启的组合是否会产生额外的状态一致性问题,这是论文完全没有覆盖的场景。

5.7 跟 STree 的关系再讲清楚一点

论文在第 2 节把 STree 列为”树结构但只支持 Mamba2”的先前工作,这一点值得展开一下,因为它直接解释了 TreeWY 真正”新”在哪里。STree 处理的是 Mamba2 风格的状态空间模型,那里的转移算子纯粹是一个标量(没有 2.4 节描述的秩一修正项),所以整条链或整棵树上的转移可以直接压缩成标量的累积乘积——这是一种远比 TreeWY 需要的三角求解简单得多的运算(标量乘法满足交换律,顺序无所谓)。论文原文写道 STree 的”严格祖先掩码”被 TreeWY 直接复用了,TreeWY 真正添加的部分,是把这个掩码从”标量衰减”推广到”标量衰减 × 秩一修正”这个不满足交换律的转移算子上——这正是 2.4 节强调的、GDN 与 Mamba2 之间那个看似很小、实则决定了整套方法能不能直接套用的结构差异。换句话说,TreeWY 不是从零发明了树结构验证这件事,而是把 STree 已经解决的”树”这一半问题,和 DeltaNet 已经解决的”门控 delta 规则”这一半问题,组合并重新推导成一个能同时处理两者的统一方案——这也是为什么论文摘要里特意强调”推导只依赖门控 delta 规则本身,不依赖任何其他架构细节”,这是在明确划定它相对于 STree 和 DeltaNet 各自贡献的边界。

6. 局限性与边界条件

作者明确指出的局限。 (1) 树验证路径目前还没法被 graph 捕获,所以更宽的树跑在更慢的执行路径上,目前这个代价超过了额外接受长度带来的收益——树宽度是”已启用且正确,而非提速”。(2) 该方法目前只在一个模型家族(Qwen3.5,两个规模)上验证过——论文明确把”扩展到第二个模型家族”列为下一步未来工作,这意味着在不同 GDN 超参数、衰减/写入强度取值范围、乃至完全不同的线性注意力变体(比如非 GDN 的门控循环)上的通用性尚未测试。(3) 这套推导专门依赖门控 delta 规则的精确代数形式(衰减乘秩一修正)——它不会自动推广到,比如说,一个具有真正不同(非秩一)更新结构的状态空间模型,不过作者认为同样的底层技巧(重写为带修正伪 value 的衰减加性注意力)很可能推广到其他具有类似结构性质的线性注意力循环。(4) 所有基准测试只用贪心草拟/验证——生产环境常见的基于采样(非贪心)投机解码没有被评估,接受长度匹配这种正确性论证是否能干净地推广到随机接受规则,尚不清楚。(5) 所有评测都关闭了前缀缓存,论文没有解释或说明原因,尽管前缀缓存是常见的生产优化,可能会跟任意一种基线的显存核算产生交互。

论文没有完全说清楚,但仔细读会注意到的问题。 构造祖先加权 Gram 矩阵 GG(一个 N×NN\times N 掩码加上所有祖先对之间的 NN 次成对点积)的 O(N2)O(N^2) 成本意味着,随着树宽度增长,TreeWY 自身在验证阶段的计算成本增长得比朴素逐节点循环(每个节点只需对其直接前驱做一次点积,总计只是 O(N)O(N))要快。论文 5.3 节的宽度结果衡量了接受长度收益和显存可负担性,但没有报告随树宽度扩展到几十个节点时,三角求解本身的实际耗时——读者只能自行推断、而不是被直接展示,O(N2)O(N^2) 的构造成本会在什么时候开始超过论文确实讨论过的分段执行开销。类似地,5.4 节的 ReplaySSM 对比只在 35B 上跑过——论文没有说明这个吞吐差距在 397B 规模下会变大、变小还是维持不变,而 TreeWY 自身的头条数字在 397B 上已经明显比 35B 小(1.06 倍吞吐 vs 1.15 倍,1.66 倍 TTFT vs 2.94 倍)——这个模式本身值得比论文给出的篇幅更多的讨论,因为一个收益随模型规模缩小的技术,跟一个收益规模不变的技术,在实践中是意义完全不同的两回事。

7. 批判性分析

(a) 这篇论文自身特有的弱点与缺陷。 首先,第 3 节的核心机制跟论文最亮眼的数字(最高约 40 倍 TTFT 降低)之间其实是正交的——那些数字来自”显存压力被缓解后能接纳更多并发请求进批次”这一系统/调度效应,而这个效应本可以由任何释放出相当显存量的方法产生,并不是闭式三角求解本身独有的。论文在这一点上相当谨慎(5.2 节的场景拆分说得很明白),但如果不那么仔细地读摘要里”1.49 倍吞吐、约 30–40 倍 TTFT”这样的头条数字,很容易高估这些收益里有多少能归功于 WY 变换推导本身的巧妙,而低估了它有多少其实是”每次投机尝试用更少的显存”——任何能达到类似显存节省效果的替代方案(包括 5.4 节提到、通过完全不同的延迟写入机制实现相当显存降低的 ReplaySSM)也都能解锁类似的收益。其次,树宽度结果——大概是这篇论文最具新意的实证贡献,因为链模式的显存降低跟 ReplaySSM 在概念上是共享的——在当前实现里明确不是一个提速,这意味着论文最有意思的主张(宽树变得可负担)目前还带着一个脚注:收益尚未在实际耗时上兑现,这明显削弱了读者今天能获得的实际价值,直到论文推迟到未来工作的 graph 捕获工程完成为止。第三,并发工作综述(第 2 节,“现有方案”)把 Bole 列为”架构上最接近”的工作,但明确说明没有做正面对比,因为 Bole 的论文没有附代码链接——这让本文跟自己最直接的竞争对手之间的定位完全没有被验证过,只能依据 Bole 自己报告的数字(82–99 倍瞬时显存降低,3.4–7.7 倍更快的树验证),这明显是比 TreeWY 自己报告的树场景数字更大的收益,而论文没有给出任何解释,说明一个表面上类似的技术为什么会报告出如此不同量级的结果。

(b) 作者低估或省略的局限。 这项工作是单一作者、单一机构(Thomson Reuters)完成的,加上一个尚未上游合并的 vLLM fork,意味着外部实践者当前的可复现门槛相当高——没有任何公开代码发布的迹象,“实现为 vLLM 主分支的一个 fork(写作本文时尚未上游合并)“是关于如何实际获取和运行这份代码的唯一实现细节。论文报告了 GPU 小时总量(三次扫描合计约 85 GPU 小时),这是个不错的透明化举措,但没有报告在得到最终报告的这几次扫描之前经历了多少次调参或失败的配置尝试,这对任何想估算采用这个方案真实工程成本的人来说很重要——尤其是相比于直接等 ReplaySSM(已经进入 vLLM RFC 和 TensorRT-LLM)或 Bole(已经在 SGLang 里上线)成熟起来,这两者今天看起来都比 TreeWY 自己这个未合并的 fork 更接近生产可用状态。此外,虽然论文很仔细地把自己的结果分成”显存受限”和”余量充足”两类,却没有讨论在真实生产流量里,到底有多大比例真正落在显存受限区间里——整个价值主张都建立在一个假设上(显存压力足够常见、值得重视),而论文没有用比如真实部署 trace 的利用率分布来独立佐证这一点,而只是用专门为了产生显存受限点而构造的并发扫描来说明。

(c) 具体、可操作的改进建议。 (1) 把三角求解本身的实际耗时,明确报告为树宽度 NN 的函数(而不只是接受长度和块数),这样读者就能独立估算出,在 graph 捕获工程完成之前,O(N2)O(N^2) 的 Gram 矩阵构造成本会在什么时候抵消掉接受长度带来的收益——这只是一项测量上的补充,不需要额外的工程投入。(2) 在 397B 规模上跑同一套 35B 的 ReplaySSM 对比,因为论文自己的数字已经暗示 TreeWY 的收益会随规模缩小,值得知道这个 ReplaySSM 对比结果是否也会出现同样的趋势(即 ReplaySSM 在 397B 上对 TreeWY 的吞吐领先会缩小、维持还是扩大?)。(3) 既然延迟写入是 ReplaySSM 吞吐优势的一个明确未经验证的假设,原型实现一个最简单的延迟写入版本重建提交策略(比如每隔 mm 次验证步骤才把被接受的状态写入持久存储一次,而不是每次都写)就能直接检验作者自己提出的这个假设,把一个推测性的归因变成一个经过验证的发现。(4) 既然 Bole 在一个名义上类似的技术上报告了明显更大的显存降低和加速数字,即使是一次不完整的、非正面对比的合理性检查——例如从 Bole 论文陈述的方法重新推导它报告的显存降低公式,检查它衡量的是否跟 TreeWY 衡量的是同一件事(在其他条件完全相同的服务条件下的峰值占用降低,还是某种定义上更有利的指标)——也会大大加强本文在这个定位主张上的说服力,而不是把这个差异搁置不处理。

8. 可复现性与实践笔记

截至本文写作时没有公开代码发布,实现只被描述为”vLLM 主分支的一个 fork(写作本文时尚未上游合并)“,所以实践者目前没法直接跑 TreeWY,只能要么直接根据论文公式(2)-(3)(对应本文公式 4-6)针对某个具体推理引擎的 GDN 内核重新实现这个两步算法(构造祖先加权 Gram 矩阵、前向替换、提交时重建),要么等作者把 fork 上游合并。不过推导本身跟具体引擎无关,只依赖门控 delta 规则的代数形式——任何已经有 GDN/门控线性注意力层的服务栈(vLLM、SGLang,或者自己的定制栈)原则上都可以实现同样的三角求解技巧,因为论文明确说”这个推导只依赖门控 delta 规则,不依赖任何其他架构细节”。对想现在就着手做原型的人来说,最容易入手的起点大概是只支持链(不带树)的情形,因为这直接对应到已经发表、已经实现的 DeltaNet 链式 WY 变换(Yang et al., arXiv:2406.06484)——在一个已有的链式 WY 实现上加入公式(2)的树祖先掩码,比从零开始搭建整套系统要小得多的工作量,也完全绕开了 graph 捕获这个工程难题(因为链本来就是可以被 graph 捕获的,见第 4 节)。复现论文结果所需的算力也不小:论文自己核算的是三次扫描合计约 85 GPU 小时(B200,每卡 178 GiB 显存),不含一次性模型加载开销,而且需要同时能访问 Qwen3.5 模型家族的两个规模(35B-A3B 单卡 B200 即可,397B-A17B 需要 TP8,即 8 张 B200)。

8b. 相关工作定位一览表

为了把论文讨论的四种替代方案(论文第 2 节)放在一起方便对比,汇总如下:

方案支持树?支持 GDN 不可交换转移?状态何时物化有代码吗?
全量状态快照(vLLM/SGLang 默认)支持支持(不需要数学技巧,天然支持)每个草稿位置都物化有(已上线的默认方案)
ReplaySSM(Dao AI Lab / NVIDIA)不支持(实现仅支持链)支持(覆盖 Mamba2 和 GDN)延迟,周期性 flush有(vLLM RFC、TensorRT-LLM PR)
STree支持不支持(仅限 Mamba2 纯标量衰减)通过累积和重算部分(研究性代码)
Bole(SGLang)支持支持(闭式解,通用混合注意力)树结构闭式解无(只有论文,没有附代码)
TreeWY(本文)支持支持(树结构 WY/UT 变换)提交时只重建被接受节点无(未合并的 vLLM fork)

对照 5.4 节的 ReplaySSM 对比和第 2 节的 Bole 讨论来读这张表,TreeWY 是唯一一个同时做到支持树、原生支持 GDN、又建立在严谨推导闭式解基础上的方案——但它也是”支持/支持”这一行里,连同 Bole 一起,唯一没有任何可运行、可链接的实现的方案,这正是第 6 节和第 7(b) 节指出的可复现性缺口。

8c. 关键数字对照表(方便速查)

论文正文和附录里分散了不少具体数字,这里把跟”内存墙”这个核心论点关系最直接的几组数字集中放在一起,方便对照阅读:

场景35B-A3B397B-A17B
GDN 固定状态显存(每序列,永久)30 MiB90 MiB
softmax KV 缓存(32K 上下文)0.625 GiB0.94 GiB
store-all 快照开销(k=3 链)120 MiB(4× 状态)360 MiB(4× 状态)
显存受限点占比(全部测量点)31/10517/70
显存受限区间吞吐提升1.15×1.06×
显存受限区间 p99 TTFT 降低2.94×1.66×
最大观测 TTFT 提升(单点)~40×(128 并发,gmu 0.6)~26×(256 并发,gmu 0.75)
树宽度 (1,1,1)→(3,3,3) 接受长度变化3.24 → 3.58未测(仅 35B 做了宽度扫描)
树宽度 (1,1,1)→(3,3,3) store-all 块数变化4 → 40(10×)未测

这张表最值得注意的一点是:显存受限的点在两个规模下都只占全部测量点的一小部分(35B 约 30%,397B 约 24%),这也再次印证了 6 节和 7(b) 节提出的疑问——论文没有独立验证真实生产流量里显存受限的比例是否真的有这么高,整个价值主张的现实意义在很大程度上取决于这个未被验证的假设。

9. 总结

TreeWY 的核心贡献是一段干净的代数操作:认识到门控 delta 规则,经过项的重新分组之后,本质上就是对一组正确计算出的伪 value 做衰减加性注意力,而这个事实——已经被普通 DeltaNet 用来做训练阶段的并行化——同样解决了 GDN 层的投机验证回滚问题,前提是把”链”的推广工作做到”树”上。实际收益是真实的,但有条件:在服务系统的 KV/状态显存池确实是瓶颈的地方,TreeWY 把释放出的 HBM 转化成明显更高的吞吐和显著更低的尾延迟(在论文最显存紧张的配置下最高约 30–40 倍 p99 TTFT);在显存不紧张的地方,这个技术付出一个诚实报告出来的小额吞吐代价。树宽度这个结果在科学意义上更具新意(更宽、接受率更高的草稿变得可以存下,而不只是更便宜),但它的系统工程收益——一个可被 graph 捕获的树验证内核——目前仍是未来工作,而不是已经交付的结果。再加上跟一个紧密竞争的并发方案(ReplaySSM)之间坦诚而尚未解决的对比——后者目前通过不同机制在原始链式吞吐上占优——这是一篇数学贡献比它的系统工程与竞争定位故事更成熟、更清晰确立的论文,而后者被诚实地呈现为仍在进行中。

8d. 逐条对照本文各章节与论文对应内容

为了方便手边同时对照论文原文阅读的读者,这里列出本文各主要章节对应论文的哪个部分,以及本文相对原文额外补充了什么:

本文章节对应论文位置本文额外补充的内容
第 1 节(问题动机)论文 Abstract + Section 1”Background”显式区分”softmax 层为什么天然回滚友好”与”GDN 为什么不是”,并补充自绘架构总览图
第 2 节(前置知识)论文 Section 1(简略带过)从零展开自回归解码、投机解码接受规则(公式 1)、KV 缓存机制、GDN 循环逐项拆解,论文假设读者已经熟悉这些背景
第 3 节(方法)论文 Section 3”Method: TreeWY”论文公式(1)(2)只给出最终形式,本文补充了从原始递推到公式(4)的完整代数展开步骤、算法 1 伪代码、以及”为什么不能反过来先重跑循环”的显式论证
第 4 节(实现)论文 Section 4”Implementation”补充为什么这个设计选择(先解全树再重建单节点)是必要的,而不只是复述论文的实现细节
第 5 节(结果)论文 Section 5 + Appendix A/B/D補充关键数字对照表(§8c)、数值小例子(§5.5)、跟 ReplaySSM 对比的解读
第 6-7 节(局限与批判)论文没有专门的批判性小节完全是本文原创的分析,论文本身只在正文散落处提及部分局限

这张表也顺带说明了本文相对论文原文的定位:论文本身写得很紧凑(10 页左右,含大量附录表格),本文的价值主要在于把公式背后隐藏的代数步骤、以及论文没有明说的设计权衡和风险点,尽可能显式地摊开讲清楚。

8e. 一句话版本(如果只想记住一件事)

如果只允许带走这篇论文的一句话总结:GDN 混合模型的循环状态”不可分割”这个性质,本来会强制投机验证时给每个候选草稿节点都存一份完整状态副本;TreeWY 证明门控 delta 规则的递推可以被重写成一个纯粹的加权求和,于是这个”必须存所有候选状态”的约束就消失了——你只需要解一次三角方程组,再在验证结束后只重建被真正接受的那一个状态。 这句话里的每一个词都对应着前面某一节的详细论证:“不可分割”对应 2.5 节,“重写成加权求和”对应 3.1 节的公式(4)-(5),“三角方程组”对应 3.2 节的算法 1,“只重建被接受的那一个”对应 3.3 节的公式(6)。如果读到这里发现哪个环节还没完全理解,可以直接跳回对应小节再读一遍。

附:常见疑问快答

问:TreeWY 是不是意味着 GDN 混合模型的投机解码”免费”变快了? 不完全是。真正免费的部分是”用同样多的显存能做更多事情”——它没有让单次三角求解本身比原来的逐节点循环更便宜(事实上前面 §3.3 提到构造 Gram 矩阵是 O(N2)O(N^2),比朴素逐节点循环的 O(N)O(N) 更贵),它节省的是存储,而不是计算。计算上的收益(尤其是链式场景下融合进一个 CUDA graph)是间接的、来自于工程实现层面的优化,不是闭式解本身自带的。

问:如果我的服务场景显存从来不紧张,还有必要上这套方案吗? 按论文自己的数字,在”余量充足”的点上,TreeWY 大概会付出百分之几的每 token 效率代价(TPOT 比值 0.97-0.99 左右),换来的收益(更多显存余量)在那种场景下并没有被真正兑现成任何吞吐或延迟提升。换句话说,如果你确信自己的生产负载几乎不会碰到显存瓶颈,这套方案在当前实现下大概率是净成本而非净收益——这也是为什么论文反复强调这是一个”看场景”的技术,而不是一个无条件的升级。

问:这跟”我们已经在用的 ReplaySSM 或者 Bole”比,该怎么选? 论文本身没有给出一个简单的结论。跟 ReplaySSM 比,TreeWY 在显存节省上打平,但目前原始链式吞吐上落后(§5.4);跟 Bole 比,论文完全没做过正面对比,只能对照双方各自报告的数字,而且这些数字量级差异很大,原因不明(§7a)。如果你的服务栈已经用了其中一个并且工作正常,现阶段更合理的做法可能是继续观察这三个方案后续的成熟度和上游合并进度,而不是急于替换。

问:为什么论文标题里强调是”Gated DeltaNet”而不是更泛化地说”所有线性注意力”? 因为整套推导(§3.1 的公式 4)严格依赖门控 delta 规则那个特定的代数形式——“标量衰减 × 秩一修正”。论文作者自己也说得很明确,这个推导”只依赖门控 delta 规则,不依赖任何其他架构细节”,这句话反过来的意思是:它确实依赖门控 delta 规则本身。如果换成一个转移算子形式不同的线性注意力变体(比如秩大于 1 的修正,或者非乘法形式的门控),公式(4)的那次代数重组能不能同样成立,论文没有讨论,§6 也把”是否能推广到其他线性注意力变体”列为明确未验证的开放问题。

问:公式(4)里的”伪 value”v~t\tilde v_t 跟原始 value vtv_t 到底差多少,量级上直观吗? 回看 §5.5 的手算例子:第三个 token 的真实 value 是 v3=1v_3=1,但对应的伪 value 却是 v~3=0.3775\tilde v_3=-0.3775,不但数值差得远,连符号都反了。这不是算错了,而恰恰体现了 delta 规则的本质——因为状态在第 3 步之前已经积累了对 k3k_3 这个方向相当强的预测(α3S2k3=1.755\alpha_3 S_2 k_3 = 1.755,已经超过真实 value v3=1v_3=1),所以”残差”是负的,意味着这一步的更新其实是在纠正一个过度预测,而不是简单地叠加新信息。理解这一点有助于打消”伪 value 是某种近似/损失精度的手段”的误解——它是一个精确的代数量,只是在直觉上不像原始 value 那样容易望文生义。

问:如果我的模型不是 Qwen3.5,而是别的 GDN 混合模型,这些数字还成立吗? 论文的定量结论(2-3 倍显存降低、1.49 倍吞吐、30-40 倍 TTFT 等)是在 Qwen3.5 的两个特定规模、特定的 3:1 GDN:softmax 层比例、特定的 dk=dv=128d_k=d_v=128 维度设置下测出来的具体数字,不应该被当成对任意 GDN 混合模型都成立的普适常数。但定性结论——“snapshotting 的显存开销跟草稿规模成正比,而 TreeWY 的开销跟草稿规模无关”——是从公式(4)-(6)的代数结构直接推出来的,不依赖具体的层数比例或维度设置,所以更换到另一个 GDN 混合模型上,这个定性趋势大概率仍然成立,只是具体的倍数会因为该模型的 GDN 层占比、dk/dvd_k/d_v 大小、以及实际的显存/带宽约束而变化。这也是论文明确列为”扩展到第二个模型家族”这项未来工作的原因所在。

问:论文摘要里”树宽度带来的收益,还不是吞吐上的收益”这句话,具体应该怎么理解? 这是论文里最容易被读者一带而过、但其实很关键的一句免责声明。它想表达的是:§5.3 展示的接受长度提升(从链的 3.24 一路涨到宽树的 3.58)是真实的质量收益,而且 TreeWY 让存储这棵更宽的树在显存上变得可行——这两点都不是问题。但”接受长度更长”不等于”每秒吞吐更高”,因为验证一棵更宽的树本身需要更多算力(要往目标模型里多推 N+1N+1 个 token),而且目前这条路径还跑在没有 CUDA graph 加速的慢速执行分支上(见 §4)。所以今天如果你只看”吞吐”这一个指标,把树开得更宽实际上可能是净负——真正等到 graph-capture 工程完成之后,§5.3 的这个收益才会同时体现在吞吐数字上。读这段内容时,不要把”能力解锁”和”性能提升”混为一谈,这也是论文自己反复强调、生怕读者误读的地方。

问:这套方法对推理服务的运维/观测有什么新的要求吗? 有一点值得留意:因为正确性判据从”逐位 token 相同”换成了”接受长度匹配参考实现”(见第 4 节),如果生产环境里要上线这套方案,监控体系里原本用来抓”输出是否跟基线一致”的告警逻辑需要相应调整为对比接受长度分布,而不是简单地做输出哈希比对——否则每次 bf16 数值误差导致的正常的、预期内的 token 差异,都会被现有监控错误地标记为回归。这算是一个论文完全没有涉及、但真要落地时会立刻碰到的运维细节。

问:公式(2)里的 G[t,i]G[t,i] 到底是在算什么,能不能再直白一点? 可以把 G[t,i]G[t,i] 理解成”节点 ii 对节点 tt 的伪 value 有多大影响”的一个耦合系数。它由三部分相乘而成:gt/gig_t/g_i 是从 iitt 之间衰减了多少(时间隔得越远,这个比值越小,影响也越弱);βt\beta_t 是节点 tt 自己的写入强度;ktkik_t^\top k_i 是两个节点 key 向量的相似度(如果 ii 的 key 方向跟 tt的 key 方向根本不像,这一项就接近零,说明 iitt 几乎没有影响)。把这三部分乘起来,再限定只有当 iitt 的祖先时才非零,就得到了整个 N×NN\times N 矩阵——它本质上是一个”按祖先关系和相似度加权的注意力矩阵”,这也是为什么论文说这个整体操作本质上就是”衰减加性注意力”的原因:GG 就是这套注意力机制里的权重矩阵。

参考文献

  1. Y. Leviathan, M. Kalman, Y. Matias. Fast Inference from Transformers via Speculative Decoding. ICML, 2023.
  2. S. Yang, J. Kautz, A. Hatamizadeh. Gated Delta Networks: Improving Mamba2 with Delta Rule. ICLR, 2025.
  3. S. Yang et al. Parallelizing Linear Transformers with the Delta Rule over Sequence Length. arXiv:2406.06484, 2024.
  4. Y. Wu et al. STree: Speculative Tree Decoding for Hybrid State-Space Models. arXiv:2505.14969, 2025.
  5. L. Wang et al. Bole: Efficient Tree Speculation for Hybrid-Attention Language Models. arXiv:2608.01651, 2026.
  6. Dao AI Lab and NVIDIA. ReplaySSM: Cache SSM Inputs, Not State. Blog / vLLM RFC #47572, 2026.
  7. W. Kwon et al. Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM). SOSP, 2023.