WaveFront Decoding: Parallelized Self-Speculative Decoding for Looped Language Models

AI summary

Overview

• Research area: Efficient inference and decoding for looped (recurrent-depth / universal-transformer style) language models, combining self-speculative decoding with batched recurrent-block execution.
• Technical level: Intermediate. Readers need some familiarity with autoregressive decoding, KV caching, and speculative decoding (draft-then-verify) to follow the cost model, but the core idea is explained in accessible terms.
• One-sentence scope: The paper introduces WaveFront Decoding (WFD), a training-free decoding schedule that fuses drafting and verification for looped LMs into a single pipelined stream of batched recurrent-block calls, and evaluates it on Ouro-2.6B and Huginn-3.5B across six Spec-Bench task categories.

What This Paper Is About

Looped language models reuse the same weights across repeated applications of a block, which keeps parameter counts small but forces T sequential recurrent-block calls per generated token, making decoding slow. The paper asks whether drafting and verification in self-speculative decoding can be run concurrently instead of in separate phases, so that different token positions at different recurrence depths share the same batched recurrent-block call. The goal is a training-free schedule that preserves the model's autoregressive output while cutting per-token latency.

Key Contributions

1. WaveFront Decoding (WFD), a training-free decoding schedule that concurrently batches token states across different positions and recurrence depths, removing the phase boundary between drafting and verification in looped LMs. It requires no additional training, auxiliary draft model, or weight modification, and applies to both full-stack and prelude-recurrent-coda (P/R/C) architectures.
2. A decoding cost model explaining WFD's speedup over autoregressive (AR) decoding and the two-phase Draft-then-Verify (DtV) baseline, showing that WFD reaches the γ → ∞ per-token latency limit of DtV using only a finite wavefront of width W = ⌈T/T_d⌉.
3. Identification of recurrence-wise KV traffic as WFD's principal long-context bottleneck, and a demonstration that cross-recurrence KV sharing (treated as an approximate extension) removes this traffic and raises speedup to as much as 4.81×.
4. Validation across models and axes: six Spec-Bench task categories, forced acceptance rates from 0.6 to 0.95, user batch sizes, and context lengths, on both a full-stack model (Ouro-2.6B) and a P/R/C model (Huginn-3.5B).

Main Findings

• End-to-end throughput: Across the six Spec-Bench task categories, WFD achieves 2.42× speedup over AR on Ouro-2.6B and 3.54× on Huginn-3.5B, corresponding to a 1.27× improvement over DtV.
• WFD beats DtV on every task: For Ouro-2.6B overall throughput, AR 6.62 TPS, DtV 12.56 TPS (1.90×), WFD 16.04 TPS (2.42×). For Huginn-3.5B overall, AR 5.29 TPS, DtV 14.77 TPS (2.79×), WFD 18.73 TPS (3.54×).
• Native draft quality is high without training: Task-level acceptance rates range from 0.78 to 0.97, with overall rates of 0.92 for Ouro and 0.94 for Huginn.
• Low-acceptance robustness: On Ouro translation, DtV reaches only 1.06× speedup at α = 0.78, while WFD retains 1.92× at α = 0.80. The paper attributes this to WFD flushing only the in-flight positions currently in the wavefront, bounded by W ≤ T, versus DtV discarding a suffix of its draft block.
• KV sharing is model dependent: Huginn tolerates inference-time cross-recurrence KV sharing, but Ouro was not trained or validated for it, and naive post-hoc sharing could not preserve Ouro's baseline accuracy, so Ouro keeps its original recurrence-wise cache.
• KV sharing raises speedup substantially on Huginn: Reducing 32 recurrence-specific KV slots to 4 or 1 shared slots maintains AR accuracy and increases WFD's acceptance rate from 0.88 to as high as 0.98 on GSM8K and from 0.90 to 0.96 on MATH-500, raising speedup from 2.54× to 4.81× on GSM8K and from 2.73× to 4.46× on MATH-500.
• Accuracy is preserved: On Huginn-3.5B GSM8K, AR scores 30.64 and WFD with 32 slots scores 31.12 (+0.48); with KV sharing, WFD stays within at most 0.70 percentage points of AR under the identical KV-sharing configuration.
• Ouro accuracy and speed: AR GSM8K 86.43, WFD 86.81 (+0.38) at 2.50× speedup; MATH-500 AR 48.20, WFD 50.20 (+2.00) at 2.31× speedup, using the original 4 recurrence-specific KV slots.
• Finite-precision explanation: Under fp32 with recurrence-specific KV cache slots, WFD and AR produce identical token sequences on both models, showing the bf16 accuracy gap stems from finite-precision batch-shape and reduction-order effects rather than the scheduling algorithm. With cross-recurrence KV sharing, accuracy differs by at most 0.2 percentage points with match rates of 0.93 to 0.99.
• Controlled scheduling: With accept/reject decisions sampled independently at forced rates from 0.6 to 0.95, WFD outperforms all evaluated DtV configurations on both models. DtV shows the expected draft-length trade-off: γ = 2 limits wasted computation after rejection, while γ = 32 approaches WFD as α → 1 but loses its advantage as acceptance decreases.
• Context-length behavior: Without KV sharing, WFD's speedup decreases as context length grows, more pronounced on Huginn (W = 8) than Ouro (W = 4), while DtV is less sensitive because its verification positions share one recurrence-specific cache. With KV sharing, WFD sustains approximately 3.2× on Ouro and 4.4× on Huginn at context lengths up to 64k tokens.
• Memory limits: Original recurrence-specific caches run out of memory at 32k tokens on Ouro and 16k tokens on Huginn, whereas shared-KV configurations support contexts up to 64k tokens.
• Batch-size scaling: Speedups of both WFD and DtV decrease monotonically with user batch size b on both models; the decline is less pronounced on Ouro because eager-mode decoding there is dominated by kernel-launch overhead rather than memory bandwidth.
• Draft depth trade-off: Higher acceptance generally comes with deeper draft depths, but does not necessarily yield greater speedup, because larger T_d means more recurrent-block calls per draft and a smaller wavefront width W.

Methodology in Plain English

The researchers begin from a property of looped models: because the same weights are applied repeatedly, an intermediate recurrence output can already act as a usable next-token prediction. A shallow output therefore serves as a draft, and the full-depth output serves as its verifier, with no separate draft model and no extra training.

They first formalize the straightforward use of this property as a two-phase baseline called Draft-then-Verify. DtV drafts γ tokens autoregressively at shallow depth T_d < T, retains each draft's hidden state, then stacks those retained states into a batch and resumes them from T_d to full depth T. Verification is batched, but drafting and verification never overlap.

WFD instead keeps several token positions in flight at once, at different recurrence depths, and advances them together in a single batched call to the shared recurrent block. Each call does double duty: it moves older positions closer to full-depth verification while simultaneously pushing newer positions to the shallow draft point where they emit new draft tokens. In steady state the active positions form a diagonal of width W = ⌈T/T_d⌉. A thin scheduler routes tokens between three blocks — the prelude P, the recurrent block R, and the coda C — forwarding tokens that reach the draft point (t = T_d) or the commit point (t = T) to the coda. Full-depth logits verify drafted tokens; on a mismatch the corrected token is committed, in-flight positions behind it are flushed, and the wavefront regrows. Greedy verification is used in the main results, though the scheduler is compatible with any sampling method and with adaptive-depth variants.

The analytical argument is that WFD does not reduce the recurrence computation any committed token needs — every committed position still traverses all T steps — but converts serial recurrent-block applications into batched ones. Because decoding is memory-bandwidth bound, a batched call over several positions streams the shared weights once and costs nearly the same latency as a single-position call. The per-token latency with all drafts accepted is L_P + T_d·L_R + L_C for WFD, matching DtV's γ → ∞ limit but with finite width. A roofline arithmetic-intensity analysis on the NVIDIA RTX A6000 (768 GB/s bandwidth, 154.8 TFLOP/s dense bf16, machine balance ≈ 202 FLOP/B) supports the bandwidth-bound assumption.

Experiments use the largest public checkpoint from each looped-model family: Ouro-2.6B (full-stack, T = 4, T_d = 1, W = 4) and Huginn-3.5B (P/R/C, T = 32, T_d = 4, W = 8), on a single NVIDIA RTX A6000 with greedy decoding and bf16 precision, user batch size 1, up to 512 output tokens, and prompts limited to 1,024 tokens. DtV uses draft block length γ = 8. Throughput counts total committed output tokens over wall-clock time including prefill and decoding.

Why This Matters

• Impact on research: The work shows that the draft-verifier decomposition native to looped LMs can be exploited as a continuous pipeline rather than as alternating phases, and that a cost model based on serial recurrent-block calls (rather than FLOPs) explains where the gains come from. It also isolates recurrence-wise KV traffic as a concrete long-context bottleneck and connects it to the emerging cross-recurrence KV sharing line of work.
• Real-world applications:
  • Latency-sensitive conversational assistants and multi-turn dialogue serving, where the paper measures the largest absolute TPS gains (Ouro WFD 14.85 TPS vs DtV 10.88; Huginn WFD 17.12 vs DtV 13.34).
  • Retrieval-augmented generation over long documents, where long context is the norm and KV sharing is needed to sustain speedup and avoid out-of-memory failures.
  • Translation and summarization services, where the paper shows WFD's advantage over DtV is largest when acceptance is relatively low (Ouro translation: DtV 1.06× vs WFD 1.92×).
  • Mathematical reasoning workloads such as GSM8K and MATH-500, where the paper reports accuracy comparable to AR while still accelerating decoding.
• Industry relevance: Because WFD is training-free, requires no auxiliary draft model, and leaves model weights unmodified, it can be applied to existing looped-model checkpoints as a serving-schedule change. Its compatibility with cross-recurrence KV sharing also speaks directly to memory-constrained and long-context deployment, where the paper reports shared-KV configurations supporting contexts up to 64k tokens versus out-of-memory at 32k (Ouro) and 16k (Huginn) with original caches.

Future Directions

• Going beyond a single GPU: The paper states that evaluations are limited to single-GPU execution, leaving distributed P/R/C execution an open direction.
• Broadening the model evidence: Only two public checkpoints were evaluated, one from each architectural family, so whether the scheduling advantages generalize across looped-model sizes and designs is not established here.
• Adaptive draft depths: The scheduler is described as agnostic to how T and T_d are set and thus applicable to adaptive-depth variants, which the paper flags as future work rather than demonstrating.
• Making KV sharing safe for more models: Ouro could not tolerate naive post-hoc cross-recurrence KV sharing while preserving baseline accuracy, and WFD with KV sharing is explicitly treated as an approximate extension because mixed-depth positions can observe shared KV states at different update stages than sequential AR; the paper notes the approximation's effect is limited in the evaluated settings but does not resolve it in general.

Target Audience

Researchers and engineers working on efficient LLM inference, particularly those interested in speculative and self-speculative decoding, memory-bandwidth-bound decoding analysis, or looped and recurrent-depth architectures. It is also relevant to practitioners deploying looped models who need throughput improvements without retraining or architectural modification, and to readers tracking cross-recurrence KV sharing as a memory-reduction technique.

Authors’ abstract

Looped language models repeatedly apply a weight-shared block to increase effective depth without increasing parameter count, but the resulting T sequential recurrent-block calls per generated token substantially increase decoding latency. To address the issue, we introduce Wavefront Decoding (WFD), a training-free self-speculative decoding framework designed for looped language models. WFD exploits two properties of these architectures: intermediate recurrence outputs provide effective draft predictions, and weight sharing allows token states at different positions and recurrence depths to be processed in one batched recurrent-block call. WFD organizes these mixed-depth states into a diagonal wavefront, continuously drafting new positions at shallow depth while advancing earlier positions toward full-depth verification. Unlike the phase-separated draft-then-verify schedule, WFD therefore concurrently batches drafting and verification within the same recurrent calls, while rejected drafts are corrected using full-depth predictions. Across six Spec-Bench task categories, WFD achieves 2.42x speedup on Ouro-2.6B and 3.54x on Huginn-3.5B over autoregressive decoding, consistently outperforming draft-then-verify. Cross-recurrence KV sharing further reduces wavefront KV traffic and increases WFD's speedup to 4.81x on Huginn-3.5B. The code is available at https://github.com/summerbro-hhj/wavefront-decoding.