Source-linked AI summary
WaveFront Decoding: Parallelized Self-Speculative Decoding for Looped Language Models
Hyeongju Ha, Jae-Joon Kim
TL;DR
Looped language models require sequential recurrent-block calls for each generated token, increasing decoding latency. WFD uses a diagonal wavefront to co-batch drafting and verification, and evaluations report higher throughput than autoregressive decoding and draft-then-verify. Cross-recurrence KV sharing further raises the reported speedup on Huginn.
Problem
Each generated token requires T sequential recurrent-block applications, repeatedly streaming shared weights with little additional parallelism in memory-bandwidth-bound decoding.
Source: §1 · pp. 1–2Method
WFD diagonally co-batches token positions at different recurrence depths, drafting new positions shallowly as earlier positions advance toward full-depth verification, without additional training.
Source: §1 · p. 2Results
Takeaways & Limitations
Cross-recurrence KV sharing reduces wavefront KV overhead and raises WFD’s speedup to as much as 4.81× while maintaining accuracy comparable to autoregressive decoding under the same sharing.
Source: §6 · p. 9With cross-recurrence KV sharing, WFD is not guaranteed to reproduce recurrence-wise KV-cached autoregressive decoding because wavefront positions can observe different shared-KV update states.
Source: §3.3 · p. 6
Abstract
from arXiv · showhide
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 co-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.
1 INTRODUCTION
Looped language models reduce parameter count through weight sharing, but their sequential recurrent calls increase decoding latency. Wavefront Decoding (WFD) co-batches drafting and verification; evaluations report higher throughput than autoregressive decoding and draft-then-verify.
1 INTRODUCTION
- Generating each token requires T sequential recurrent-block calls, repeatedly streaming shared weights while exposing little parallelism in memory-bandwidth-bound decoding.The recurrent portion's latency consequently grows approximately linearly with T, reducing the practical benefit of a smaller parameter footprint.Source: §1 · pp. 1–2
- WFD diagonally co-batches token positions at different recurrence depths, drafting new positions shallowly while advancing earlier positions toward full-depth verification.Full-depth predictions correct speculative errors, preserving autoregressive output by construction; WFD requires no additional training, draft model, or model modification.Source: §1 · p. 2
- 2.42× and 3.54× higher throughput than autoregressive decoding on Ouro and Huginn, respectively; cross-recurrence KV sharing raises Huginn speedup to 4.81×.WFD also achieves 1.27× higher throughput than draft-then-verify on both models.Source: §1 · p. 2
- WFD’s cost model predicts that a finite wavefront reaches the asymptotic per-token latency of an infinitely long draft-then-verify draft, with validation across acceptance rates, batch sizes, and context lengths.The analysis covers both full-stack and P/R/C looped models.Source: §1 · p. 2
- Cross-recurrence KV sharing addresses recurrence-wise KV traffic, identified as WFD’s principal long-context bottleneck, and yields up to 4.81× speedup while maintaining evaluated task accuracy.Without sharing, distinct recurrence-specific KV slots cause KV traffic to scale with wavefront width and gradually erode speedup.Source: §1 · p. 2
2 BACKGROUND
Looped language models reuse a shared block to increase effective depth without increasing unique parameter count, and their architectures include full-stack and P/R/C families. The background also describes self-speculative decoding in looped models and a two-phase draft-then-verify baseline.
2.1 LOOPED LANGUAGE MODELS
- Modern looped language models repeatedly apply a shared layer stack, increasing effective depth while keeping unique parameters fixed; repeated latent updates are also framed as iterative reasoning.Reported pretraining work shows competitive language modeling and reasoning performance (Zhu et al., 2025; Prairie et al., 2026; Yang et al., 2026c;a).Source: §2.1 · pp. 2–3
- P/R/C models repeat only the recurrent block between a fixed prelude and coda, while full-stack models repeat the entire transformer stack; WFD uses this unified abstraction for both.For full-stack models, the recurrent block contains the whole stack and the prelude and coda reduce to token embedding and LM head, respectively.Source: §2.1 · p. 3
- Distinct KV states at each recurrence step preserve the original recurrent computation but increase cache capacity and traffic with recurrence depth; sharing can reduce this overhead.Prior work reports native sharing for trained full-stack models (Vendrell et al., 2026) and inference-time sharing tolerated with limited evaluated-task accuracy degradation for pretrained P/R/C models (Geiping et al., 2025).Source: §2.1 · pp. 3–4
2.2 SELF-SPECULATIVE DECODING
- Looped models provide a native self-speculative pair: shallow recurrence outputs draft next tokens, full-depth outputs verify them, and retained shallow states can be resumed without extra training or an auxiliary model.Earlier self-speculative methods for fixed-depth transformers use layer skipping or specialized training because intermediate layers are not trained to feed the output head.Source: §2.2 · p. 4
- Draft-then-verify (DtV) autoregressively drafts γ tokens at shallow depth, then batches their retained states for the remaining recurrence steps and accepts the longest matching prefix.At the first mismatch it commits the full-depth target token and discards the remaining draft suffix; if all drafts are accepted, it commits the next target token.Source: §2.2 · p. 4
3 WAVEFRONT DECODING
Wavefront Decoding batches token positions at different recurrence depths so drafting and verification proceed together rather than in separate phases. Its cost model explains the resulting latency advantage and identifies recurrence-wise KV traffic as a long-context bottleneck that cross-recurrence sharing can reduce.
3.1 BATCHING DRAFTING AND VERIFICATION VIA WAVEFRONT DECODING
- WFD advances mixed-depth token positions in one batched recurrent-block call, drafting new positions while earlier ones approach full-depth verification.On rejection, the corrected token is committed, in-flight positions behind it are flushed, and the wavefront regrows.Source: §3.1 · pp. 4–5
3.2 COST MODEL AND SPEEDUP ANALYSIS
- With all drafts accepted, WFD reaches DtV’s γ → ∞ latency limit at finite wavefront width W = ⌈T/Td⌉.Each Td batched recurrent-block applications commits one token; WFD rejections flush at most W − 1 in-flight positions, while DtV’s wasted suffix grows with γ.Source: §3.2 · p. 6
3.3 REDUCING WAVEFRONT KV TRAFFIC WITH CROSS-RECURRENCE KV SHARING
- Mixed-depth WFD positions access distinct recurrence-specific caches, increasing KV traffic with wavefront width and potentially offsetting batching gains at long contexts.Cross-recurrence KV sharing lets those positions access a common cache, but the combination is approximate because positions may observe different shared-cache update stages than sequential AR.Source: §3.3 · p. 6
4 EXPERIMENTS
Across Spec-Bench tasks, WFD outperforms draft-then-verify and autoregressive decoding, with cross-recurrence KV sharing increasing speed on Huginn while maintaining comparable accuracy against matched AR. Controlled acceptance-rate measurements also show WFD outperforming the evaluated DtV configurations across the tested range.
4.1 SETUP
- Experiments use Ouro-2.6B and Huginn-3.5B at default recurrence depths T = 4 and T = 32, with draft depths Td = 1 and Td = 4, respectively.Evaluation uses one RTX A6000, bf16, greedy decoding, batch size 1 by default, and end-to-end throughput including prefill and decoding.Source: §4.1 · pp. 6–7
4.2 END-TO-END PERFORMANCE ON SPEC-BENCH
- 2.42× speedup over AR on Ouro and 3.54× on Huginn; WFD also delivers 1.27× higher throughput than DtV on both models.WFD outperforms DtV on every task for both models across the six Spec-Bench categories.Source: Table 1 · p. 7, §4.2 · p. 7
4.3 EFFECT OF CROSS-RECURRENCE KV SHARING
- 4.81× speedup over matched AR on Huginn GSM8K with KV sharing, up from 2.54× without sharing.On Huginn MATH-500, speedup rises from 2.73× to 4.46×; WFD remains within 0.70 percentage points of matched AR accuracy.Source: Table 2 · p. 8, §4.3 · p. 8
- With its original recurrence-specific KV cache, Ouro WFD achieves 2.50× and 2.31× speedup over AR on GSM8K and MATH-500, respectively.The slightly higher bf16 accuracy reflects finite-precision differences; WFD and AR produce identical token sequences in fp32.Source: §4.3 · p. 8, Table 2 · p. 8
4.4 CONTROLLED SCHEDULING ANALYSIS
- Across forced acceptance rates from 0.6 to 0.95, WFD outperforms all evaluated DtV configurations on both models.WFD closely tracks or exceeds the best DtV configuration without a separate draft-block-length parameter γ.Source: §4.4 · p. 8
5 RELATED WORK
Related work covers looped language models and KV-sharing approaches, alongside methods that parallelize decoding for looped models. The cited parallel approaches either tolerate approximation or require training from scratch, and do not preserve the original sequential-loop behavior in all model types.
Other looped language models
- Related work spans weight-shared looped models, retrofitting fixed-depth LLMs into recurrent architectures, and cross-recurrence KV sharing.The cited KV-sharing approaches include native training and inference-time sharing on pretrained looped models.Source: §5 · pp. 8–9
Parallel decoding for looped models
- Prior parallel-decoding approaches for looped models either tolerate inference-time approximation or require training from scratch, and apply only to certain model types.The cited approaches do not preserve the behavior of the original sequential looped model.Source: §5 · p. 9
6 CONCLUSION
The conclusion presents WFD as a training-free schedule that co-batches token positions at different recurrence depths, reporting speedups over AR and further gains with cross-recurrence KV sharing. The accompanying wavefront-state description explains how the schedule fills, verifies, and recovers its pipeline.
6 CONCLUSION
- WFD achieves up to 2.42× and 3.54× speedup over AR on Ouro-2.6B and Huginn-3.5B, respectively, and cross-recurrence KV sharing raises speedup to as much as 4.81×.Evaluations are limited to greedy decoding, single-GPU execution, and two public checkpoints; shared-KV WFD maintains accuracy comparable to AR under the same sharing configuration.Source: §6 · p. 9
A.1 WAVEFRONT STATE
- After warm-up, the wavefront settles at width W = ⌈T/Td⌉; a rejected draft commits the corrected token, flushes in-flight positions behind it, and regrows the wavefront.Tokens reaching draft depth spawn the next position, while a token reaching full depth verifies and commits its draft.Source: §A.1 · p. 12
B.1 ARITHMETIC INTENSITY ESTIMATION
The analysis estimates arithmetic intensity for recurrent-block decoding with and without KV sharing, then compares both regimes with the GPUs’ machine balance. It concludes that decoding remains memory-bound at feasible operating points.
B.1 ARITHMETIC INTENSITY ESTIMATION
- The estimate counts recurrent-block weights as one stream per step and KV traffic as b k without sharing or k with sharing.The analysis uses bf16 decoding and the attention projections, MLP, KV load, and corresponding matmul and attention operations.Source: §B.1 · p. 12
- 13.3 and 37.6 are the no-sharing intensity ceilings for Ouro-2.6B and Huginn-3.5B at N = 1,024, falling to 2.5 and 5.6 at N = 8,192.These ceilings sit 5–22× below either GPU’s machine balance; only contexts shorter than ∼200 tokens would bring the ceiling to machine balance.Source: §B.1 · p. 12
- With sharing, intensity equals batch size b, but machine balance requires b ≈202 on A6000 or 292 on B200, versus wavefront widths capped at T = 4 or 32.Both sharing regimes therefore remain below the compute-bound region on either GPU.Source: §B.1 · p. 12
B.2 ANALYSIS OF EXPERIMENTAL TPS VALUES
The roofline analysis compares measured autoregressive throughput with a memory-bandwidth prediction, then examines why the models differ. It attributes Ouro’s low attainment to kernel-launch overhead and checks throughput under CUDA graph capture.
B.2 ANALYSIS OF EXPERIMENTAL TPS VALUES
- 74% of roofline is attained by Huginn-3.5B, versus 18% by Ouro-2.6B, whose measured throughput is too low for a genuinely memory-bandwidth-limited decoder.The predicted bound uses TPS ≈ β/n (w + k) and is compared with the measured autoregressive throughput.Source: §B.2 · p. 14, Table 5 · p. 15
- 65% of Ouro’s ≈150 ms per-token latency is GPU idle time dominated by kernel-launch overhead, with 192 layer applications per token under eager execution.The active operations are memory-bandwidth bound, but host overhead between short kernels consumes wall-clock time.Source: §B.2 · p. 14, Table 6 · p. 15
- +148% is Ouro’s AR throughput change with CUDA graph capture, while Huginn’s is +3%; WFD speedup changes from 2.78× to 2.30× on Ouro.On Huginn, WFD changes from 4.08× to 5.21× without sharing and from 5.86× to 6.22× with sharing.Source: §B.2 · p. 14
C EVALUATIONS
The evaluations examine how batch size and context length affect WFD, and discuss tree drafting and adaptive recurrence depth. KV sharing preserves WFD’s advantage at larger measured batches and long contexts.
C.1 USER BATCH SIZE
- WFD and DtV speedups both decline monotonically as user batch size grows, while KV sharing lets WFD retain its advantage over DtV throughout the measured range.Without cross-recurrence KV sharing, WFD’s KV traffic grows as bW rather than b, and its speedup falls fastest as batch size increases.Source: §C.1 · p. 15, Figure 5 · p. 16
C.2 CONTEXT LENGTH
- With KV sharing, WFD sustains approximately 3.2× speedup on Ouro and 4.4× on Huginn through 64k-token contexts, remaining faster than DtV at every measured length.Without sharing, WFD speedup declines with context length as mixed-depth positions access distinct recurrence-specific caches; the shared-KV configurations reach 64k tokens.Source: §C.2 · p. 15
C.3.1 DRAFTING TREE STRUCTURE
- WFD can accommodate tree-structured drafting without structural changes, but suitable tree shapes and branching rules for looped models remain empirically unexamined.The figure caption identifies speedup over AR by batch size, with and without KV sharing.Source: §C.3.1 · p. 16, Figure 5 · p. 16
C.3.2 ADAPTIVE RECURRENCE DEPTH T
- Adaptive recurrence-depth schemes can reduce WFD’s speedup slightly because both approaches reduce recurrent-block applications per committed token, shrinking WFD’s remaining headroom.Reducing average depth contracts the effective wavefront width.Source: §C.3.2 · p. 16
C.3.3 ADAPTIVE DRAFT DEPTH Td
The draft-depth analysis finds a trade-off between acceptance rate and wavefront width. Under fp32, WFD matches AR token sequences with recurrence-specific KV caches in the evaluated prompts, but shared-KV sequence matches are imperfect despite unchanged task accuracy.
C.3.3 ADAPTIVE DRAFT DEPTH Td
- Higher draft depth raises acceptance but does not guarantee higher speedup, because the resulting narrower wavefront reduces effective batch size.Figure 7 plots speedup against measured acceptance rate across draft depths; the authors suggest adapting draft depth to this trade-off.Source: §C.3.3 · pp. 16–17, Figure 7 · p. 17
C.4 MATHEMATICAL EXACTNESS WITH FP32
- With recurrence-specific KV caches, fp32 WFD and AR produce identical token sequences on the first 100 GSM8K prompts, while shared-KV WFD has imperfect matches but unchanged task accuracy.The reported fp32 speedups are not a faithful evaluation because fp32 forces attention onto a different backend kernel.Source: §C.4 · pp. 17–18, Table 8 · p. 18