SRT: Speculative Rollout with Tree-Structured Cache

1 · 动机

1.1 历史脉络:rollout 怎么从"小头"变成"大头"

2024 年之前的 RLHF:reward model + PPO,response 长度多在 256–512 token,generation 在一个 RL step 里只占 30–40%,大家把功夫花在 critic loss、reference model KL、replay buffer 上。

2025 年起 reasoning 模型起飞 (DeepSeek-R1, GLM-4.5, Kimi-1.5, Seed1.5-Thinking, Qwen3),trajectory 长度普遍 4k–32k token,GRPO/DAPO 还要求每个 prompt 同时采 5–16 个 response。论文 Figure 2a 给出的硬数据:在 PPO/GRPO/DAPO/ReTool 四个算法上,rollout 平均吃掉 65% (有时 >70%) 的 RL step 时间。这就是 Amdahl 视角下"必须被攻击的部分"。

更糟的是 long-tail 问题:

• 同一个 batch 内 response 长度高度不均衡 — 多数早早 EOS,少数 long-CoT 一路写到 max_response_length (DAPO 里 8k, ReTool 里 16k)。
• 同步 RL 必须等最后一条结束才能进 train 阶段 → 一堆 GPU 在"bubble"里空转等。

1.2 别人的招数为什么"够呛"

1.3 为什么 tree-structured cache 是"非平凡的解锁"

把"用历史 rollout 当 draft"这件事说出来不难,真正非平凡的是怎么存、怎么查、怎么填。考察以下三个观察:

1. 跨 epoch 相似度高,但不是 verbatim 重复。Figure 2c 实测同一 prompt 在第 t 步和前 t-1 步累积 n-gram 的 overlap 随 epoch 单调上升,后期可达 60%+。意味着子串命中率高,但整段 response 命中率低 — 必须以"路径"为单位检索,不能以"句子"为单位。
2. GRPO/DAPO 同一 prompt 一次产 K=5–16 个 rollout。这 K 条彼此当下就高度相似 (相同 prompt + 相同 policy + 不同 sample,只差 stochastic decoding)。如果只在 epoch 边界更新 cache,会错过同一 batch 内这 K-1 条对第 K 条的帮助 — 这正是 RhymeRL 的盲区。
3. Long-tail bubble 是免费算力。Run-ahead generation 不会增加 wall-clock(它本来就 idle),但能给即将出现的 prompt 把 cache 提前烧热,消除 first-time cold-start。

把这三个观察落到数据结构上,Trie / token tree 是几乎唯一合理的选择:

• 共享前缀压缩存储 (一句话 "the cat sit on the mat" 和 "the cat is on the mat" 共享 "the cat")
• 沿路径累计频率统计 — 父节点 count = 子节点 count 之和,自底向上聚合
• 线性时间插入 (按 token 走一遍即可),CPU 内存就够,与 KV cache 无关
• 查询时:从当前已生成 suffix 的最长匹配出发,沿"概率最高的子节点链"贪心展开 → 一束 draft tokens

2 · 背景速查

2.1 标准 speculative decoding 复习 (一行)

Draft 给 k 个 token → target 一次并行验证 k+1 个位置 → rejection sampling (Leviathan 2023) 保证输出分布与直接从 target 采样等价。SRT 只是把"draft 模型"换成"从 token tree 沿路径抽 token"。

3 · 核心机制:per-prompt token tree + 路径打分

3.1 数据结构

对每个 prompt p,维护一棵 trie T_p:

• 每个节点 u 对应一个 token 子序列 (从 root 到 u 的路径)
• 每个边由一个 next-token 标号
• 每个节点存 count(u) = 这个子序列在历史 rollout 里出现的次数
• 非叶节点的 count = 所有子节点 count 之和 (运行时维护即可,O(depth) 更新)

3.2 查询过程 (drafting)

当前 partial continuation 是 y₁:t,SRT 这样查 T_p:

1. 最长后缀匹配:从 y_t 往前回溯,找最长的 q 使得 y_{t-q+1:t} 是 T_p 上某条路径。这一步决定 "我们站在树的哪个节点 u_q"
2. 如果连 q=1 都没匹配上 (e.g. 前所未见的 token),退化为正常 1-step 自回归,下一步再试
3. 贪心扩展:从 u_q 出发,反复挑选 conditional 概率最高的子节点 v,即:
   C(v) = count(v) / Σ_{w ∈ children(parent(v))} count(w)
   路径分 = 沿路 C 的乘积 (类似 beam 但只保留 top-1)
4. 扩展到预算 B(q) 用尽 (depth 限制 / 或路径分阈值) — 论文没给精确公式,但实现里通常是 fixed budget = 几十个 token
5. 把这一束 draft tokens 一次性丢给 verifier (current policy) 做 spec decoding 标准验证

3.3 验证过程 (一次 forward, classical spec decoding)

Verifier (current π_θ) 把 [y₁:t, t̂₁, t̂₂, ..., t̂_K] 全部一次 forward,得到每个位置的 p_target。然后:

• 对每个 t̂_i,以概率 min(1, p_target(t̂_i | context) / p_draft(t̂_i)) 接受
• 第一次 reject 时,从 (p_target − p_draft)+ 重采样得到一个新 token,停止

这里有个 subtle 点:p_draft 是什么? SRT 的 draft 来自 cache 频率,不是一个真正的 distribution。论文 (隐含) 用 C(v) 作为 p_draft;这样 rejection sampling 的数学保证依然成立 — 因为 rejection sampling 只要求 p_draft 是某个合法分布,无所谓它有多差 (差只意味着接受率低,不意味着输出分布偏)。这就是"lossless"成立的关键。

4 · Cache 维护:on-the-fly + run-ahead

树本身只是数据结构;怎么把树喂饱才是 SRT 的真正贡献。这一节讲两条互补的 maintenance 策略,以及它们各自解决什么。

4.1 On-the-fly insert (跨 group 内即时受益)

当前 batch 的 K 条 rollout 在生成的同时就把 token 流式插进 T_p:

4.2 Run-ahead generation (用 long-tail bubble 预热未来)

另一条:当 batch 内只剩零星 long-tail rollout 还在跑、其它 GPU 已经空,SRT 不让 GPU 闲着,而是从 data loader 的 lookahead 窗口 / active prompt queue 里取下一个待训练的 prompt p',对它跑一段 partial generation:

• 这些 token 插进 T_{p'}(给未来用),
• 但不会被当成 learning targets — 真正的 rollout 还是要等到那个 prompt 被排到时,由当时的 policy 重新生成。

论文 §4.2 给出消融:enable run-ahead 后 mean accepted tokens / step 进一步上升,而单独 history-only baseline (RhymeRL 风格) 由于 cold-start 在前几个 epoch 几乎没收益。

5 · Worked example:一个 prompt 跨多 step 的完整生命周期

把上面所有机制串起来,我们追踪一个具体 prompt p = "Q: 求 ∫x² dx" 在三步 RL 训练中 cache 是怎么生长、查询、命中的。

Step t=1 (cold start)

• T_p 是空树。GRPO 采 K=5 条 rollout,全部走标准自回归。
• 5 条 rollout 完成,token 序列例如:
  • r1: "Let me compute. ∫x² dx = x³/3 + C."
  • r2: "积分 ∫x² dx 等于 x³/3 + C 。"
  • r3: "We have ∫x² dx = x³/3 + C."
  • r4: "Let me think. The integral is x³/3 + C."
  • r5: "积分等于 x³/3 + C."
• 5 条全部插进 T_p。注意 "x³/3 + C" 这个尾巴在所有 5 条里都出现 → cache 里这条路径 count=5,极强 signal。

Step t=2 (on-the-fly 受益)

• 同一 prompt p 又被采样 (DAPO 里 K=16,假设这次也是 5)。
• 第一条 rollout 自回归生成,产到 "积分 ∫x² dx",在 T_p 上找到这个 prefix → 抽 draft "等于 x³/3 + C 。"。
• Verifier 一次 forward 验证:假设 "等于 x³/3 + C 。" 全部接受 — 这一步本来要 12 token 自回归,现在 1 次 forward 完成 → α ≈ 12,巨大加速。
• 这条 rollout 一边生成一边把刚验证的 token 也插进 T_p (count++),后续 rollout 拿到更新过的更胖的树。

Step t=2.5 (run-ahead 预热)

• step t=2 batch 里有 1 条 long-tail rollout 在 8k token 上挣扎,其它 GPU 空闲。
• 调度器从 lookahead window 取下一个 prompt p' = "Q: 求 ∫sin(x) dx",用当前 policy 跑一段 partial gen。
• partial 产出 "积分 ∫sin(x) dx 等于 -cos(x) + C" → 插进 T_{p'},不计 loss,不留在 batch。
• step t=3 的 batch 里真正轮到 p' 时,T_{p'} 已经是热的。

Step t=3 (持续累积)

• policy π_θ 已经被 update 了两步,但跟 t=1 时差距不大 — Figure 2c 实测显示同 prompt 跨 epoch overlap 单调上升,所以 cache 里那些"x³/3 + C" 路径依然适用。
• 新的 5 条 rollout 与 cache 高度匹配 → α 持续高 → 速度持续优于 baseline。

6 · 实验关键结果

6.1 主表 (Table 1, Qwen2.5-1.5B, on-policy)

读法重点:

• SuffixDecoding 在 DAPO 上反而变慢了 (103 > 81.7) — 注意这个细节。SuffixDecoding 用全局后缀树作 draft,在 K=16/prompt 的 DAPO 设置下可能因为 draft 命中无关 token 浪费 verify 算力,或者后缀树太大导致 lookup overhead。SRT 因为 per-prompt tree 没这个问题。
• N-gram 几乎没用 — 全语料 n-gram 与"这个 prompt 该生成什么"几乎不相关。
• SRT generation time 在 ReTool 上 49→28.7 (1.7×) — 多轮 + 长 trajectory 是 SRT 的甜区:cache 累积更厚。

6.2 Ablation: K 越大 SRT 越甜 (Table 3)

这是 SRT 论点的核心实验性证据:每个 prompt 同时采的 rollout 越多,on-the-fly insert 提供的早期 hint 越多 → 后产生的 rollout 命中率越高。意味着 DAPO (K=16) / GSPO (K=64) 这种"高分组数"的现代算法上 SRT 收益会进一步放大。

6.3 Ablation: run-ahead 的边际收益 (Figure 5)

论文用仿真比较:

• history-only (RhymeRL 风格): mean accepted tokens 上升缓慢,前期 cold-start 严重
• + on-the-fly insert: 立刻拉高
• + run-ahead: 进一步抬升,曲线最早期就站在高位

7 · 与同类工作对比

SRT 在这条线里的独特位置

• 它是唯一完全 model-free 的 RL spec decoding。EAGLE-3 / SPEC-RL / ReSpec 都要训一个 draft 模型并维护它的对齐;SRT 一行额外参数都不需要。
• 它对算法侧零侵入 — 直接接 Verl + vLLM 就能跑,GRPO/DAPO/PPO/ReTool 全适配。NeMo-RL 的方案需要把 EAGLE-3 head 嵌进 vLLM 的 inference path,系统集成成本高。
• 但它的加速天花板取决于"prompt 重复出现 + 跨 epoch 相似性"。如果数据集每 epoch 只过一遍 (no replay) 或 prompt 极多样,cache 就喂不饱。EAGLE 系不依赖这个假设。

8 · 局限 / 个人 take / 待验证问题

论文承认的局限 / 我观察到的

• 实测只跑了 Qwen2.5-1.5B,8× Hopper。在更大模型 (8B+) 上 verify 的 forward 本身就更贵,SRT 的相对加速可能稀释。
• Cache 大小没量化分析 — 大数据集 + 长 prompt + 多 epoch,CPU 主存占用可能爆。论文没给上限管理策略 (LRU? 频率阈值剪枝?)。
• 对 K (responses/prompt) 高度依赖。PPO (K=1) 上仍有效但收益最小;GRPO/DAPO 是甜区。
• Run-ahead 的实现细节 (调度器怎么决定何时启动、怎么避免占用即将到来的 batch 的 KV cache) 论文几乎没写。
• 没和 EAGLE 系做正面 head-to-head 比较 (NeMo-RL/ReSpec/SPEC-RL),只比了 model-free baselines (N-gram, SuffixDecoding)。读者很难知道 model-free 路线相对带参 draft 的真实差距。

我的疑问 (落地后想验证)

1. Cache 漂移问题:RL 训到几百 step 后 policy 已经显著演变,早期插入的 token 可能"过期"。SRT 用 frequency 但没有 recency decay。在 long horizon RL (1k+ step) 上 acceptance 是否会衰退?
2. 探索 vs 利用矛盾:如果 cache 里某条路径出现频率极高,SRT 会不断提议它 — 而 rejection sampling 虽然分布等价,但方差不一样;高方差 prompt 上 effective sample 数会变少,对 GRPO 的 advantage 估计可能引入偏置 (虽不破坏期望)。
3. 多轮 (ReTool) 场景下,树是 per-prompt 还是 per-(prompt, turn-history)?论文没说清楚。如果 per-prompt 但 turn 之间状态差异大,cache 会很乱。
4. 与 EAGLE-3 / ReSpec 组合:cache 命中率高时直接取,低时 fallback 到 EAGLE 模型。系统上是否值得这复杂度?ablation 缺位。
5. Run-ahead 的"未来 prompt"可能已经不是 next batch 的 prompt:大规模分布式训练里,sampler 从全局打乱,run-ahead 怎么准确预测?
6. KV cache 与 paged-attention 的交互:verify 一束 draft tokens 时,K+1 个位置一起进 attention,在 vLLM 的 page table 里要预留连续 page。SRT 在 long-context (32k+) 下的 OOM 风险?

记忆点

精读笔记 v1 · 2026-05-07 · 配套论文 PDF 在 /data/szhang967/papers/spec-rl/SRT_2601.09083.pdf