Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching

1 INTRODUCTION

• CDB forms new batches between loop steps because tokens exiting at different depths cannot share a standard uniform forward pass.Refill replaces exited tokens to keep batch size large; Figure 1 illustrates the schedules with two concurrent slots.Source: §1 · p. 2, Figure 1 · p. 2
• The implementation interleaves prefill, prelude, and coda batches with loop steps, and uses a lookahead gate to overlap batch preparation with GPU execution.It also implements a no-refill mode that only removes tokens.Source: §1 · p. 2
• On Ouro, CDB realizes up to 99% of estimated maximum early-exit speedup; on Huginn, it realizes 83 to 96% of available speedup.Refill is important for Ouro, while costly transformer boundary stages make no-refill perform better on Huginn.Source: §1 · p. 2
• The paper derives a prefill-corrected speedup estimate and a hardware-aware latency model to analyze when refill helps.Its contributions also connect boundary-stage cost, KV-cache design, and loop count to inference efficiency.Source: §1 · p. 2

2 PROBLEM SETUP

• The architecture has a prelude, a shared recurrent core, and a coda; the core runs r_t ≤ r_max times per token, while the other stages run once.A 2-4-2 model executes 2+4r_t+2 transformer layers; a fully looped 0-4-0 model has no transformer layers outside the core.Source: §2 · pp. 2–3
• Depth-adaptive decoding lets each token exit at an input-dependent depth, with the exit decision made at the end of a loop step.Source: §2 · p. 3
• Standard continuous batching treats a complete forward pass as one scheduling unit, but adaptive depth requires routing exited tokens to the coda while others continue looping.Source: §2 · p. 3
• CDB must handle stage heterogeneity, missing KV states for deeper loop steps, and exit decisions that arrive only after the step they affect.Source: §2 · p. 3

3 RELATED WORK

• Prior looped-transformer methods use learned exit heads [Banino et al., 2021; Fu et al., 2026; Song et al., 2026; Tan et al., 2023; Zhu et al., 2025] or training-free convergence checks [Geiping et al., 2025; Kohli et al., 2026; Lee et al., 2026].Other methods exit tokens in fixed-sized groups or exit all tokens at the same depth.Source: §3 · p. 3
• Early exit in standard transformers remains difficult to batch efficiently because unshared layers prevent cross-depth refill and leave missing KV states to recompute or approximate.This paper focuses on looped LMs, whose shared core allows cross-depth refill and KV sharing without recomputation.Source: §3 · p. 3
• Earlier CDB work introduced the method theoretically with simulated throughput estimates, then evaluated ablations excluding prefill, scheduling, and KV-cache updates.This work includes those costs in its end-to-end implementation and evaluation.Source: §3 · p. 3

4.1 STAGE-WISE SCHEDULING

• The scheduler treats prefill, prelude, recurrent items, and exited tokens awaiting coda as separate queues, then orders stages dynamically to refill freed recurrent slots.No-refill removes exited tokens and delays coda; refill prioritizes coda, then prefill, while a minimum coda batch size K trades efficient boundary batches against delayed refill.Source: §4.1 · pp. 3–4

4.2 DEPTH-AWARE KV CACHE

• Depth-indexed caching reserves r_max KV slots per token, while shared caching uses one slot per recurrent layer and overwrites it each loop step.Last-exited caching copies final states into skipped slots and generally preserves accuracy; shared caching reduces recurrent-core KV memory by a factor of r_max but may require specialized training to preserve accuracy.Source: §4.2 · p. 4

4.3 ASYNCHRONOUS SCHEDULING

• A lookahead gate predicts exits one loop step early, giving the CPU time to prepare the next batch while the GPU processes the current one.The exit signal at step r determines whether a token participates in step r+2 instead of r+1; the implementation reduces GPU idle time to nearly zero.Source: §4.3 · pp. 4–5, Figure 3 · p. 5

• With F0 ≪ Fr, the FLOP bound approaches rmax/dbar; non-looped transformer layers raise F0, while larger rmax reduces each skipped loop’s fractional savings.Source: §5 · p. 5
• Because prefill is unaffected and can add substantial latency, the decode-only speedup bound is adjusted using the measured prefill-time fraction α.Source: §5 · p. 5
• The recurrent-step model is t(Nc, Nb) = δ + γNb + αNc, separating fixed weight reload from batch-scaled compute and context-scaled KV streaming.Source: §5 · p. 5
• B* marks the shift from nearly constant, weight-dominated small-batch latency to approximately linear growth at larger batches, and falls as context length increases.Source: Figure 4 · p. 6, §5 · p. 6
• Refill helps most in the memory-bound regime, where early exits save FLOPs but shrinking batches barely cuts latency because weight loading dominates.Source: §5 · p. 6

6.1 SETUP

• The evaluation replays Alpaca, ShareGPT, and ArXiv-Summarization workloads on fully looped Ouro and partly unlooped Huginn, using offline throughput and online serving tests.The models run on one 80 GB H100 GPU; workload prompt and output lengths vary the prefill costs, decode durations, and context lengths tested.Source: §6 · p. 6, §6.1 · p. 6

6.2 OFFLINE THROUGHPUT

• 1.32–1.53× CB throughput on Ouro realizes 96–99% of estimated end-to-end speedup; Huginn’s best no-refill results reach 1.30–1.37× CB and 83–96% of that estimate.At full depth, CDB has no early exits and no opportunity to outperform CB.Source: §6.2 · p. 7, Figure 5 · p. 7
• Refill stays within 5% of estimated end-to-end speedup on Ouro and helps most around or below B*, while prefill consumes 3.7–36% of GPU time across workloads.Its advantage narrows as larger batches become compute bound; prefill reduces end-to-end speedup most on long-context ArXiv.Source: §6.2 · p. 7

6.3 ONLINE SERVING

• On ShareGPT, CDB sustains higher request rates before normalized latency rises; refill gives Ouro the largest capacity gain, while Huginn’s refill and no-refill capacities match.At thresholds with fewer exits, the capacity advantage disappears.Source: §6.3 · p. 8, Figure 6 · p. 8

6.4 ABLATIONS

• K=W/4 best balances Huginn’s coda efficiency against delayed refill, while Ouro’s cheap boundary stages favor immediate coda execution at K=1.At the largest decode batch sizes, K matters less as stages become compute bound.Source: §6.4 · p. 8, Figure 7 · p. 8
• Reducing Huginn’s boundary layers makes immediate coda execution more competitive; with a fully looped 0-4-0 split, K=1 is near optimal and the speedup gap closes.This supports boundary-stage cost as a central efficiency constraint.Source: §6.4 · p. 8, Figure 7 · p. 8

7 DISCUSSION

• CDB realizes up to 99% of the estimated maximum speedup, but useful speedups require enough accurate early exits, recurrent compute to matter, and an asynchronous system that converts skipped work into lower latency.Expensive prefill limits end-to-end speedup, while shrinking recurrent batches saves little time in the memory-bound regime without refill.Source: §7 · p. 9
• Keeping non-looped stages small can improve refill efficiency, trading off against accuracy gains from adding layers outside the recurrent core.More loops provide finer-grained exit depths but make each skipped loop a smaller share of recurrent compute; cache layouts must also handle unwritten KV slots after early exits.Source: §7 · p. 9
• The lookahead gate requires rmin=2, which could reduce speedup for models that produce useful outputs after one loop; predicting only the first exit beforehand could allow rmin=1.The constraint does not reduce realistic speedup for Ouro and Huginn, whose exits after one loop are excluded.Source: §7 · p. 9
• Evaluation covers two looped architectures on one H100 GPU and omits production features such as multi-GPU execution, preemption, prefix caching, and speculative decoding.Those features may change the prefill-versus-decoding balance or interact with depth-adaptive scheduling, so CDB’s relative benefit may differ in broader serving systems.Source: §7 · p. 9

8 CONCLUSION

• CDB decomposes decode steps into separate stage queues, solving the batching problem created by variable-depth forward passes.The paper describes this as its first efficient inference method for depth-adaptive looped LMs.Source: §8 · p. 9
• The implementation is near-optimal and realizes most of the speedup available from early exits, showing that depth-adaptive decoding can be served efficiently.Source: §8 · p. 9
• Co-designing looped architectures and serving systems can lead to substantially faster inference, including through more accurate early exits and less expensive boundary stages.Source: §8 · p. 9

Chunked prefill and decode batching

• Prefill packs up to M prompt tokens and chunks prompts that do not fit; decode advances each active sequence by one token, up to W sequences.The paged KV cache allocates fixed-size blocks as sequences grow, avoiding reservations for their maximum possible lengths; if cache capacity is exhausted, the system can offload or recompute state, though preemption is disabled in the main experiments.Source: Figure 8 · p. 15, §A.1 · pp. 15–16

A.1 INFERENCE SYSTEM OVERVIEW

• The engine separates prefill and decode batches, and paged KV allocation assigns cache blocks as sequences grow to use GPU capacity efficiently.Preemption can offload KV state to CPU or discard it for later recomputation; the main experiments disable preemption to avoid its costs and state changes.Source: §A.1 · pp. 15–16

A.2 BATCHING AND ADMISSION PARAMETERS

• The sequence limit W caps active sequences and decode batch size, while token limit M caps prefill batches; S and KV-cache headroom ρKV tune admission.Waiting for free slots can combine short prompts into more efficient prefill batches, while headroom adapts admission to actual cache use; larger M can require more activation memory and interrupt decode longer.Source: §A.2 · pp. 16–17

A.3 SCHEDULING ALGORITHMS

• With refill, CDB independently schedules prefill, recurrent-core, and coda work; exited tokens proceed through coda and prelude to refill the recurrent queue.No-refill instead lets the recurrent batch shrink and waits for the cohort to exit before running coda; the algorithms also use admission limits and a minimum coda batch size.Source: §A.3 · pp. 17–18

A.4 SCHEDULING EXAMPLES

• With refill and K=1, coda and prelude run at each exit, letting successor tokens refill the recurrent batch to W=64 in the shown traces.Shared recurrent weights allow tokens at different loop steps to share a forward pass; larger K delays coda and refill, trading fewer coda executions against temporarily smaller recurrent batches.Source: §A.4 · p. 18, Table 2 · p. 19, Table 3 · p. 19

A.5 EFFICIENT EXECUTION

• CUDA graphs replay fixed-shape operations using preallocated buffers and batch-size buckets; CDB captures prelude, core, and coda separately.Prefill graphs are broken at attention operations; the implementation uses FlashAttention-3 and computes logits only for final prompt tokens.Source: §A.5 · pp. 19–20

B RECURRENT-STEP LATENCY DETAILS

• R² ≥0.999 confirms decode latency is linear in depth; Ouro’s τ0 is ≈0.4 ms, while Huginn’s is 2.8 ms at L = 128 and 3.3 ms at L = 8192.The latency curves are shown for Ouro 1.4B, Ouro 2.6B, and Huginn 3.5B; τ0 is the prelude and coda overhead.Source: Figure 9 · p. 20, B RECURRENT-STEP LATENCY DETAILS · p. 20
• B* falls with context length as KV-cache streaming cost grows; fitted B* values are mild overestimates because the slopes mildly underestimate per-token cost.The latency model separates context-independent weight reload cost δ, matrix-multiplication cost γB, and KV-cache streaming cost αLB. For Ouro 1.4B, B* falls from about 140 at L = 128 to single digits at L = 8192; Huginn follows a similar trajectory from about 80.Source: B RECURRENT-STEP LATENCY DETAILS · pp. 20–21

C.1 INPUT AND OUTPUT TOKENS

• ArXiv prompts average 6,989 tokens, versus 267 for ShareGPT and 18 for Alpaca; ShareGPT is long-tailed, while ArXiv is prefill-heavy.The workloads use replayed prompt/reference-output pairs tokenized for each model; the offline threshold sweeps use 6,000 ShareGPT, 10,000 Alpaca, and 1,000 ArXiv requests.Source: Figure 10 · p. 21, §C.1 · p. 21

C.2 EARLY-EXIT DISTRIBUTION

• At a given threshold, exit-depth distributions remain similar across workloads with modest shifts, while Huginn’s gate spreads exits over a wider range of depths than Ouro’s.The distributions are based on deterministic exits recorded in advance and replayed during experiments.Source: §C.2 · p. 21, Figure 11 · p. 22

D COMPLETE ONLINE SERVING RESULTS

• Figures 12–13 extend ShareGPT online-serving comparisons to Alpaca and ArXiv across more thresholds, with the general observations holding for all workloads.They plot normalized latency against request rate; columns progress from aggressive to conservative thresholds for Ouro and the reverse for Huginn.Source: D COMPLETE ONLINE SERVING RESULTS · p. 22, Figure 12 · p. 22, Figure 13 · p. 23

E.1 KV CACHE SHARING

• The cache comparison covers full depth-indexed, single-slot, and first-then-shared layouts, with cache footprint reported per token and per 4096-token sequence.The layouts use rmax, 1, and 2 KV slots per recurrent layer, respectively.Source: §E.1 · p. 23, Table 4 · p. 24
• Huginn is insensitive to KV-sharing layout, while Ouro loses about 6 accuracy points with first-then-shared and collapses under single-slot sharing.The authors could not reproduce Ouro’s reported near-lossless sharing; they suggest its prompt may have used a depth-indexed cache.Source: §E.1 · p. 23
• At Huginn depths 6 and 9, single-slot sharing yields 11.45% and 25.25% accuracy, versus 3.71% and 12.43% with the full depth-indexed cache.The layouts converge to within about two points at depth 16.Source: §E.1 · p. 24
• Single-slot sharing shrinks the KV cache 4× for Ouro and 8.5× for Huginn, making reduced sensitivity to sharing a development target.KV sharing also enables simpler continuous depth batching.Source: §E.1 · p. 24
• At truncated Huginn loop depths, Table 5 compares GSM8K accuracy across recurrent-depth KV-cache layouts.The caption identifies fixed recurrent depth as a comparison dimension.Source: Table 5 · p. 24

E.2 EARLY-EXIT GATING

• Ouro uses a trained exit gate with a minimum depth of 2, while Huginn uses a convergence threshold and a maximum depth of 16.The Huginn threshold direction differs from Ouro’s, and single-slot KV caching is used for performance.Source: §E.2 · p. 24
• At mean recurrent depths down to about 2.5 for Ouro and 8 for Huginn, depth-adaptive decoding Pareto-dominates fixed-depth decoding in accuracy.Figure 14 shows accuracy–depth frontiers, mean exit depth by threshold, and each model’s exit-depth distribution.Source: §E.2 · p. 24, Figure 14 · p. 25
• The lookahead gate matches the teacher gate’s exit behavior and approximately its accuracy and threshold curves; Huginn’s delayed decision needs only threshold recalibration.The pre-loop gate performs poorly, while delayed exits incur no accuracy cost for either model.Source: §E.2 · p. 25
• Minimum loop counts of 2 for Ouro’s trained gate and 3 for Huginn’s convergence criterion limit exit distributions and compute savings, especially for shallower looped models.The convergence criterion compares consecutive latent states, and the delayed scheduler adds a step.Source: §E.2 · p. 25