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

AI summary

Overview

Research area: Efficient LLM inference systems and looped (recurrent-depth) language model architecture design.

Technical level: Advanced — the paper assumes familiarity with continuous batching, KV caching, FLOP accounting, and GPU memory-bandwidth versus compute-bound regimes.

Scope: The paper presents and evaluates the first end-to-end implementation of continuous depth batching (CDB), a scheduling method that makes per-token, depth-adaptive decoding of looped language models practical, and analyzes which architectural properties determine the achievable speedup.

What This Paper Is About

Looped language models reuse a shared block of layers (the recurrent core) and can, in principle, spend fewer loops on "easy" tokens and more on "hard" ones. The problem is that tokens running different numbers of loops cannot share a single forward pass, so standard batching systems such as vLLM cannot serve them. The paper's goal is to build and measure an inference engine that batches these variable-depth tokens efficiently, and to identify which model architectures benefit most.

Key Contributions

1. First end-to-end CDB implementation. The authors build a scheduler that decomposes decoding into separate stage queues — prefill, prelude, recurrent core, and coda — and forms new batches between loop steps, enabling refill of slots freed by token exits.

2. Theory linking compute savings to wall-clock speedup. They derive a FLOP-based upper bound on decode speedup plus a prefill-corrected end-to-end estimate, and a hardware-aware latency model for a recurrent step that predicts when refill helps.

3. Architecture guidance from measurement. They show that boundary-stage cost, KV-cache layout, and loop count determine inference efficiency, concluding that fully looped models are best suited to depth-adaptive inference.

4. Empirical evaluation across workloads and serving modes. Offline throughput and online serving benchmarks on Ouro 1.4B and Huginn 3.5B over Alpaca, ShareGPT, and ArXiv-Summarization workloads.

Main Findings

• CDB realizes up to 99% of the estimated maximum speedup. Across settings, the implementation captures nearly all of the speedup theoretically available from early exit, leaving further gains dependent on model architecture and exit behavior.

• Ouro 1.4B (fully looped) benefits most from refill. Refill achieves 1.32 to 1.53× the throughput of full-depth continuous batching and realizes 96 to 99% of the estimated end-to-end speedup bound.

• Huginn 3.5B performs better without refill. No-refill reaches 1.30 to 1.37× continuous batching and 83 to 96% of the end-to-end bound. Huginn's 2-4-2 split means 50% of its transformer layers sit outside the recurrent core, making refill costlier.

• Scheduling overhead is small. At full depth, where CDB has no early exits, the implementation still retains 95 to 99% of continuous batching throughput.

• Refill helps most in the memory-bound regime. Below the saturation batch size, weight loading dominates latency, so shrinking a recurrent batch saves FLOPs but little time; refill fills freed slots at little additional cost. Above it, latency scales approximately linearly with batch size and no-refill already benefits from exits.

• Prefill limits the achievable end-to-end gain. Decode-FLOP estimates overestimate speedup because prefill is unaffected by early exit. Prefill accounts for 3.7–36% of GPU time across the workloads, with the largest share on long-context ArXiv.

• Boundary-stage cost constrains efficiency. Models with large non-looped layers (token embedding, LM head, unshared transformer blocks) slow down and complicate scheduling. Removing transformer layers from Huginn's boundary stages makes immediate coda execution more competitive and improves throughput; at a fully looped 0-4-0 split, K=1 becomes near optimal, mirroring Ouro.

• Minimum coda batch size K is architecture-dependent. Huginn performs best at K=W/4 (e.g., K=16 at W=64); Ouro performs best at K=1 because its boundary stages are cheap and keeping a large core batch matters more than coda efficiency. At the largest batch sizes, the choice of K matters less as stages become compute bound.

• Asynchronous scheduling with a lookahead gate removes nearly all GPU idle time. Device idle time drops from 10.27% (synchronous) and 10.05% (async with immediate gate) to 0.56% with the lookahead gate. The lookahead gate constrains minimum exit depth to r_min = 2, which the authors report is consistent with Ouro's and Huginn's exit behavior.

• Offline gains translate to online serving capacity. Under Poisson arrivals over a 10-minute window, CDB sustains higher request rates before normalized latency rises sharply. For Ouro, refill provides the largest capacity gain while no-refill adds about half as much; for Huginn the two modes reach the same serving capacity.

Methodology in Plain English

The authors start from the observation that a looped LM has three parts: a prelude that processes each token once before it enters the loop, a shared recurrent core that can be applied a variable number of times, and a coda that predicts the next token after the token exits. Because tokens exit at different depths, the engine must make scheduling decisions inside the forward pass.

They treat each stage as its own queue and run them as separate batches, interleaving prefill, prelude, and coda work with loop steps. In refill mode, slots freed by exiting tokens are filled with new tokens to keep the recurrent batch large; a minimum coda batch size K decides how many exits accumulate before the coda runs. In no-refill mode, the recurrent batch simply shrinks and the coda runs once at the end. They also define two depth-aware KV-cache layouts — a last-exited cache that copies a token's final state into skipped slots, and a shared cache with one slot per recurrent layer per token — and they prepare batches asynchronously, using a lookahead gate so exit decisions are known one loop step earlier.

To interpret results, they derive a FLOP-based decode speedup bound, correct it for measured prefill time, and fit a latency model of a recurrent step with a fixed weight-reload term, a batch-dependent compute term, and a context-dependent KV-streaming term. They then benchmark two architectures — Ouro 1.4B (0-24-0, r_max = 4) and Huginn 3.5B (2-4-2, r_max = 16) — on three request workloads with offline throughput and online serving tests, all on a single 80 GB H100 GPU, using the same codebase, kernels, and admission policy as the continuous batching baseline.

Why This Matters

Research impact. The paper shifts depth-adaptive decoding for looped LMs from a theoretical idea to a measured system, and it reframes the question from "how many FLOPs are saved?" to "how much wall-clock time is saved, and under what conditions?" The negative result for Huginn — no-refill beating refill — is as useful as the positive one for Ouro, because it traces the cause to a concrete architectural property (boundary-stage cost) rather than to scheduler quality. The work also reinforces hardware-conscious architecture design, a theme the authors note is increasingly common in LLM development.

Real-world applications:

• Serving looped LMs behind online APIs, where many sequences must be processed concurrently and per-token latency matters.
• Reinforcement learning rollouts, which the authors cite as a setting requiring many sequences processed at once.
• Benchmarking and evaluation pipelines, also cited as a batch-heavy use case where all requests are known upfront.
• Long-context, prefill-heavy tasks such as the ArXiv-Summarization workload, though the paper shows this is precisely where prefill limits the end-to-end gain.

Industry relevance. The paper gives inference engineers a concrete answer to a deployment question: can a looped model's compute savings be converted into throughput? It also gives model designers a trade-off to manage — extra layers outside the recurrent core may improve accuracy, but they degrade the efficiency of depth-adaptive serving. Limitations are explicit: results cover two architectures on a single H100, and exclude multi-GPU execution, preemption, prefix caching, and speculative decoding. Preemption was disabled in the main experiments to avoid confounding the comparison.

Future Directions

• Predicting exits before the core runs. The lookahead gate constrains minimum exit depth to r_min = 2. The authors suggest predicting the complete exit depth before a token enters the recurrent core, or more promisingly, predicting only the first loop-exit early and using lookahead decisions for later loops, which would allow r_min = 1.

• KV-cache layout co-design. Early exits leave slots for skipped loops unwritten, requiring a policy for missing states. The authors call for understanding when each layout works and for designing models that reliably support shared caches, which would reduce memory footprint and implementation complexity.

• Extending to new model components and serving features. Future looped models may include components such as mixture-of-experts, and production features like multi-GPU execution, preemption, prefix caching, and speculative decoding could shift the balance between prefill and decode or interact with depth-adaptive scheduling.

• More accurate, more frequent early exits. Further speedups are described as depending on future models that allow more frequent exits at high accuracy, and on architectures that keep non-looped stages small despite the accuracy trade-off this may entail.

Target Audience

Inference and serving-systems engineers who need to decide whether variable-depth looped models are worth deploying; machine learning systems researchers working on batching, scheduling, and KV-cache design; and architecture researchers designing looped or recurrent-depth models who want to know which structural choices make depth-adaptive inference practical. Readers without familiarity with batching, KV caching, and GPU memory-versus-compute-bound behavior will need background reading first.

Authors’ abstract

A main promise of looped language models is depth-adaptive inference. By looping a block of shared layers a variable number of times, the model can use less compute for "easy" tokens and more for "hard" ones. However, tokens with different numbers of loops cannot share a uniform forward pass and therefore cannot be handled by standard batching systems such as vLLM. The practical value of depth-adaptive inference thus hinges on whether batching can be made efficient. We introduce the first efficient method for depth-adaptive looped LMs via continuous depth batching (CDB), which forms new batches between loop steps. Our method dynamically schedules looped and non-looped parts of the architecture, manages looped KV-caching, and predicts which tokens will exit the loop in advance so it can prepare batches asynchronously. Experiments on Ouro 1.4B and Huginn 3.5B show that fully looped architectures are best suited to depth-adaptive inference, as large non-looped layers outside the recurrent core (e.g., token embedding, LM head, and unshared transformer blocks) slow down and complicate scheduling. Overall, CDB realizes up to 99% of the estimated maximum speedup available, leaving further gains primarily dependent on model architecture and exit behavior.