WaveFront Decoding: Parallelized Self-Speculative Decoding for Looped Language Models
Hyeongju Ha, Jae-Joon Kim
cs.LG
2026-09-19
波前解码将循环LM草稿与验证同批执行，Ouro相对AR为2.42×、Huginn为3.54×，KV共享后达4.81×。
Looped language model（循环语言模型）把同一组权重反复套 T 次，用迭代计算换有效深度。Ouro 这类 fullstack 把整栈 Transformer 循环；Huginn 这类 prelude-recurrent-coda（P/R/C）只循环中间块。参数少了，解码没跟着变快：每个新 token 仍要串行跑完 T 次共享块。小 batch 下解码吃内存带宽，每次循环把同一份权重从显存再读一遍，几乎挤不出额外并行，循环延迟大致随 T 线性涨。Huginn 默认 T=32，这件事更硬。
中间循环步经常已经能猜对最终 next-token。浅循环当草稿、满深度当 verifier，是模型自带的 self-speculative decoding。论文把这个想法做成对照基线 Draft-then-Verify（DtV）：先用深度 Td<T 自回归写出 γ 个草稿，再把保留状态从 Td 批处理推到 T。草稿阶段仍一个接一个，验证才成批。两段来回切，浅层草稿和深层验证从来不同时占用那一次共享块调用。
Wavefront Decoding（WFD，波前解码）吃循环架构的两条性质：中间步能当草稿；权重共享让不同位置、不同循环深度的状态可以塞进同一次 batched recurrent-block 调用。调度把混深度状态排成对角波前：新位置在浅层出草稿，更早的位置同时往满深度推进并验证。像一条流水线，最前面的 token 在做期末考，后面几个还在随堂测验，同一节课一起改。
稳态波前宽度 W=⌈T/Td⌉。实验里 Ouro 取 T=4、Td=1，W=4；Huginn 取 T=32、Td=4，W=8。每一步 P、R、C 各自最多一次前向：P 只处理新位置；R 把在飞状态推进一层，到达草稿点 Td 或提交点 T 的交给 C；C 在 Td 出草稿，在 T 做 greedy 验证。拒稿时用满深度预测改掉该 token，冲掉后面所有在飞位置，波前从纠正点重新长回来，最多冲掉 W-1 个半成品。DtV 拒稿会丢掉整段草稿后缀，γ 越大浪费越大。
被接受的位置三种调度都走满 T 步，加速来自把串行 R 收成 batched R。小 batch 下一次 batched R 读权重几乎和单 token 一样贵。全接受时 WFD 每 token 延迟是 LP+Td·LR+LC，等于 DtV 的 γ→∞ 极限，有限波前就摸到。P 和 C 每个提交 token 仍付一次，加速上界严格小于 T/Td。
长上下文会再收一笔 KV 账：混深度位置默认读各自 recurrence 专属 KV slot，流量随 W 涨。跨 recurrence KV sharing 让它们读同一份 cache，但这是近似扩展，波前看到的共享 KV 可能和串行 AR 不一致。Huginn 推理时直接共享还能保住精度；Ouro 没为这训过，事后硬共享保不住基线，实验仍用原来的 4 槽。
单卡 RTX A6000，greedy、bf16、user batch=1，按 Spec-Bench 六类任务量端到端吞吐量，含 prefill。总体接受率 Ouro 为 0.92、Huginn 为 0.94。
| 模型 | 方法 | 吞吐量 | 相对 AR |
| Ouro-2.6B | AR | 6.62 tok/s | 1.00× |
| Ouro-2.6B | DtV | 12.56 tok/s | 1.90× |
| Ouro-2.6B | WFD | 16.04 tok/s | 2.42× |
| Huginn-3.5B | AR | 5.29 tok/s | 1.00× |
| Huginn-3.5B | DtV | 14.77 tok/s | 2.79× |
| Huginn-3.5B | WFD | 18.73 tok/s | 3.54× |
两条模型上 WFD 相对 DtV 都是 1.27×。接受率掉下来时差距更大：Ouro 翻译任务 DtV 在 α=0.78 时只剩 1.06×，WFD 在 α=0.80 时仍有 1.92×。强制接受率从 0.6 扫到 0.95，WFD 全程压过 γ=2 到 32 的全部 DtV 配置。
跨 recurrence KV 共享把 Huginn 的 GSM8K 加速从 2.54× 拉到 4.81×（4 slot，α 从 0.88 升到 0.98），MATH-500 从 2.73× 到 4.46×。同一 KV 共享配置下，Huginn 上 WFD 相对 AR 最多差 0.70 个百分点。Ouro 不能事后共享，GSM8K / MATH-500 上分别 2.50× 和 2.31×。bf16 下 Ouro 的 MATH-500 为 50.20 对 AR 的 48.20，论文归为 batch 形状和 reduction 顺序；fp32、分 recurrence KV 时 token 序列与 AR 完全一致。
没有 KV 共享时，加速随 prefill 变长往下掉，Huginn（W=8）比 Ouro（W=4）掉得更狠。加上共享后，最长测到 64k 时 Ouro 仍约 3.2×、Huginn 约 4.4×。user batch 变大后两边相对 AR 的加速都单调下降。Ouro 的 eager PyTorch 路径里，GPU 空闲占每 token 延迟的 65%，主要是 kernel launch。CUDA graph 把 AR 从 6.74 拉到 16.72 tok/s（+148%），WFD 相对这份优化后的 AR 仍有 2.30×。Huginn 的 AR 只因此快了 3%，它更接近带宽墙。
给已经在用 Ouro、Huginn 这类循环模型的人，这是一份零训练、不改权重、不外挂草稿模型的解码调度。只要实现能把混深度状态打成一批 R，就能把 DtV 还在串行出草稿的那段时间也吃掉。
它不是通用 LLM 的投机解码替代品。固定深度 Transformer 没有「不同深度、同一套权重」这个免费 batch 维度，WFD 吃的就是循环结构。投机解码在大 user batch 下收益本来就缩，这篇同样如此。Ouro 主表上的 2.42× 有一部分来自 eager 启动开销，上 CUDA graph 之后数字会收敛，但 2.30× 还在。
对 serving：P/R/C 里 R 每步最多吃 W 个状态，P 和 C 只吃一两个，附录建议拆到不同硬件。树状草稿和自适应 T、Td 都还没做实验。
论文自己划了几条：只测了单卡和两个公开 checkpoint；长上下文在无 KV 共享时是主瓶颈；带 KV 共享的 WFD 是近似算法；自适应深度和分布式 P/R/C 留作后续。
主结果是 greedy。调度号称兼容任意采样，但没有采样设置下的数字。Ouro 事后 KV 共享直接放弃，4.81× 只属于 Huginn。Table 2 里 bf16 精度偶有 WFD 略高于 AR，论文说是数值噪声，fp32 才对得上；不能把 bf16 精度差当质量收益。
对照只打了 DtV，没有和 LayerSkip、Draft & Verify、外挂草稿模型或 Geiping 等人 2025 年并行 sampler 做端到端对比。主表含 prefill，附录 CUDA graph 是稳态宽度、静态 shape 的受控计时，不能和 Table 1 横比。Td 只在 Huginn GSM8K 上扫过，更深草稿接受率更高，加速并不单调涨。