WaveFront Decoding: Parallelized Self-Speculative Decoding for Looped Language Models
Hyeongju Ha, Jae-Joon Kim
cs.LG
2026-09-19
Training-free Wavefront Decoding co-batches mixed-depth looped-LM states, reaching 2.42× on Ouro-2.6B and 3.54× on Huginn-3.5B vs AR; KV sharing lifts Huginn to 4.81×.
Looped language models reuse one weight-shared block T times so effective depth grows without growing the parameter count. Full-stack models such as Ouro loop the entire transformer; prelude-recurrent-coda (P/R/C) models such as Huginn loop only the middle block. The smaller footprint leaves decode latency untouched. Each new token still needs T serial applications of that shared block. At small batch, decode is memory-bandwidth bound: every recurrence streams the same weights again and adds almost no extra parallelism. Recurrent latency therefore scales roughly linearly with T. Huginn's default T=32 makes this painful.
Intermediate recurrence outputs often already match the final next-token prediction, so a shallow pass can draft and the full depth can verify. The paper turns that observation into a concrete baseline, Draft-then-Verify (DtV): draft γ tokens autoregressively at depth Td<T, then resume the saved states as a batch from Td to T. Drafting stays serial. Verification is the only batched phase. Shallow drafts and deep verification never share a single call to the recurrent block.
Wavefront Decoding (WFD) uses two facts about looped models. Intermediate steps are usable drafts. Weight sharing lets states at different positions and different recurrence depths ride in one batched recurrent-block call. The scheduler packs those mixed-depth states onto a diagonal wavefront: new positions keep drafting at shallow depth while earlier positions deepen toward full-depth verification. The picture is a pipeline. The lead token is taking the final exam while later tokens are still on the quiz, and both get graded in the same class period.
Steady-state width is W=⌈T/Td⌉. The runs use T=4, Td=1 (W=4) on Ouro and T=32, Td=4 (W=8) on Huginn. Each iteration, P, R, and C fire at most once. P initializes new positions. R advances every in-flight state by one step and hands positions at Td or T to C. C drafts at Td and greedily verifies at T. On a reject, the full-depth prediction replaces the bad token, every in-flight position behind it is flushed, and the wavefront grows back from the correction. At most W-1 partial states die. DtV discards the unaccepted suffix of a draft block of length γ, so longer γ costs more on a miss.
WFD does not cut FLOPs on an accepted token. AR, DtV, and WFD all run all T steps for that token. In the bandwidth-bound regime a batched R costs nearly as much as a single-token R, so wall-clock time tracks serial R calls per committed token. When every draft is accepted, WFD's per-token latency is LP+Td·LR+LC, the γ→∞ limit of DtV, reached here with a finite wavefront. P and C still fire once per committed token, so speedup stays strictly below T/Td.
Long context opens a second bill. Mixed-depth positions normally read distinct recurrence-specific KV slots, so KV traffic grows with W. Cross-recurrence KV sharing points them at one cache. That is an approximate extension: as the shared cache is updated, wavefront positions can see a different snapshot than sequential AR. Huginn tolerates this at inference time. Ouro was never trained for it; naive post-hoc sharing did not keep baseline accuracy, so Ouro keeps its original 4 recurrence-wise slots.
All main numbers are on one RTX A6000, greedy, bf16, user batch=1, Spec-Bench six-task throughput including prefill. Overall acceptance is 0.92 on Ouro and 0.94 on Huginn.
| Model | Method | Throughput | vs 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 is 1.27× DtV on both models. The gap widens when acceptance drops: on Ouro translation, DtV is 1.06× at α=0.78 while WFD holds 1.92× at α=0.80. With accept/reject forced from 0.6 to 0.95, WFD stays above every DtV draft length from γ=2 to 32.
Cross-recurrence KV sharing lifts Huginn GSM8K speedup from 2.54× to 4.81× (4 slots, α 0.88→0.98) and MATH-500 from 2.73× to 4.46×. Under the same KV-sharing setup, WFD stays within 0.70 percentage points of AR task accuracy on Huginn. Ouro cannot share after the fact; WFD is 2.50× on GSM8K and 2.31× on MATH-500. In bf16, Ouro MATH-500 is 50.20 versus AR 48.20; the paper blames batch shape and reduction order. In fp32 with recurrence-wise KV, token sequences match AR exactly.
Without KV sharing, speedup falls as prefill grows, and the drop is steeper on Huginn (W=8) than on Ouro (W=4). With sharing, the paper reports about 3.2× on Ouro and 4.4× on Huginn out to 64k. Larger user batches shrink the speculative headroom; both schedules slow relative to AR.
On Ouro's eager PyTorch path, GPU idle is 65% of per-token latency, mostly kernel launch. CUDA graphs move AR from 6.74 to 16.72 tok/s (+148%). WFD is still 2.30× that optimized AR. Huginn's AR gains only 3% from graphs; it already sits closer to the bandwidth wall.
For anyone already serving Ouro or Huginn, this is a zero-training decode schedule: no weight edits, no extra draft model. If the runtime can pack mixed-depth states into one R call, the serial drafting window that DtV still pays becomes useful work.
This is not a drop-in speculative decoder for ordinary fixed-depth LLMs. Those models do not get a free batch dimension across recurrence depth, and that is the resource WFD spends. Speculative decoding also fades at large user batch; the same curve shows up here. Part of Ouro's headline 2.42× is eager launch overhead. After CUDA graphs the number compresses, and 2.30× remains.
Serving-wise, R in a P/R/C model processes up to W states per step while P and C see one or two. The appendix suggests pinning the three blocks on separate hardware. Tree drafts and adaptive T or Td are described, not measured.
The paper's own list: single GPU, two public checkpoints, long context without KV sharing as the main bottleneck, KV-shared WFD as an approximation, adaptive depth and distributed P/R/C left open.
Main tables use greedy decoding. The scheduler is said to work with any sampler; no acceptance or speedup numbers under sampling are given. Naive KV sharing was dropped for Ouro, so 4.81× is Huginn-only. Table 2 sometimes shows WFD slightly above AR in bf16; the paper calls that numerical noise, and fp32 is the check. Those bf16 gaps are not a quality win.
The only decode baseline is DtV. There is no end-to-end throughput comparison with LayerSkip, Draft & Verify, an external draft model, or the 2025 parallel sampler from Geiping et al. Table 1 includes prefill. The CUDA-graph numbers use a static-shape, steady-width microbenchmark and should not be lined up against Table 1. Td was swept only on Huginn GSM8K: deeper drafts raise acceptance, narrow the wavefront, and do not monotonically raise speedup.