🤖 machine learning

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

This paper introduces Continuous Depth Batching (CDB), a novel scheduling framework that enables efficient, end-to-end depth-adaptive inference for looped language models by decoupling loop iterations from boundary stages and overlapping scheduling with GPU computation, thereby achieving up to 1.9x higher throughput and significantly lower latency compared to standard methods.

Original authors: Kristian Schwethelm, Daniel Rueckert, Georgios Kaissis

Published 2026-08-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Kristian Schwethelm, Daniel Rueckert, Georgios Kaissis

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive, super-fast library where a single librarian can answer thousands of questions at once. In the world of Artificial Intelligence, this librarian is a "Large Language Model" (LLM), and the questions are text prompts from users. Usually, to save time and money, the librarian processes all questions in a strict line, giving every single one the exact same amount of attention, no matter how simple or complex the question is. This is efficient, but it's wasteful: answering "What is 2+2?" takes just as much brainpower as solving a complex physics problem, even though the first one is a breeze.

Recently, scientists discovered a clever trick called "looped" models. Instead of a straight line, these models have a "thinking loop." They can check their work, run the answer through a special thinking chamber, and if the answer looks good, they stop. If it looks shaky, they run it through the chamber again. This means easy questions get a quick pass, while hard ones get extra loops of deep thought. It's like a student who glances at a simple math problem and writes the answer immediately, but spends ten minutes scribbling on a whiteboard for a tough calculus problem. This sounds perfect, but it creates a nightmare for the librarian trying to manage the line. If everyone in the line needs a different amount of time in the thinking chamber, the whole line gets stuck. The librarian can't process the next group until the current group is done, and if some people are still thinking while others are finished, the system grinds to a halt. This paper tackles that exact traffic jam.

The authors of this paper introduce a new system called Continuous Depth Batching (CDB) to fix the traffic jam caused by these "thinking loops." They realized that the old way of managing the line (called standard batching) was too rigid. It treated every question as if it needed the same number of loops, which wasted the speed advantage of the new models. To solve this, they built a dynamic scheduler that acts like a super-organized traffic controller. Instead of waiting for the whole group to finish, this controller lets finished questions leave the line immediately and instantly swaps them with new questions waiting to start.

Think of it like a busy restaurant kitchen. In the old system, if a chef is making a simple salad (easy token) and a complex soufflé (hard token) at the same time, the whole kitchen has to wait until the soufflé is done before serving the salad or starting a new order. The new CDB system is different. As soon as the salad is ready, it's plated and sent out, and the chef immediately grabs a new order to start chopping vegetables, keeping the kitchen constantly busy. The paper shows that by doing this, the system can keep the "thinking loop" running at full speed without ever letting the chef's hands go idle.

The researchers tested this on two different AI models, named Ouro and Huginn. They found that CDB works incredibly well. For the Ouro model, the new system achieved up to 99% of the theoretical maximum speed-up possible. In plain English, this means the system is almost as fast as it could possibly be if you could magically remove every single second of wasted time. In real-world tests, this translated to processing 1.5 to 1.9 times more requests per second than the old methods. When the system was under heavy load, like a busy website with thousands of users, the wait time for answers dropped by 45% to 90%.

However, the paper also points out that this magic isn't automatic. The system works best when the "thinking loop" is the main part of the work. If the steps before the loop (like reading the question) and after the loop (like writing the final answer) are too slow or heavy, they can slow down the whole process. The authors found that for models where these extra steps are heavy, the speed-up is still good but not quite as dramatic. They also discovered a "sweet spot" for how many questions to process at once. If you try to process too few questions, the system wastes time reloading the model's "brain" (weights) for every single step. If you process too many, the system gets bogged down. CDB helps find that perfect balance by constantly refilling the line with new questions whenever a slot opens up.

One of the cleverest tricks the authors used is something called a "lookahead gate." Normally, the system has to wait for a question to finish its thinking loop before it knows if the answer is good enough to leave. This causes a tiny pause where the computer just sits there waiting. The new system predicts the answer one step early, allowing the next batch of questions to be prepared while the current one is still finishing. It's like a conductor who cues the next section of the orchestra just before the current section ends, so there is no silence between the notes. This tiny change reduced the computer's "idle time" to nearly zero, making the whole process incredibly smooth.

In short, this paper proves that we can make these smart, self-correcting AI models actually fast enough to use in the real world. By treating the "thinking loops" as a flexible schedule rather than a rigid line, and by constantly swapping finished tasks for new ones, the authors have built a system that is not only smarter but also significantly faster and more efficient. They didn't just suggest it might work; they built it, tested it on real hardware, and showed that it delivers massive speed gains, turning a theoretical idea into a practical tool for the future of AI.

Technical Summary: Continuous Depth Batching for Looped Language Models

Problem Statement

Looped Language Models (LMs) offer a promising architecture for depth-adaptive inference, where a shared recurrent core is iterated a variable number of times per token. This allows the model to allocate more compute to "hard" tokens and less to "easy" ones, improving the compute-accuracy trade-off. However, this adaptivity fundamentally breaks standard batching mechanisms used in modern inference frameworks (e.g., vLLM).

In standard continuous batching, a forward pass is a uniform operation where all tokens in a batch undergo the same number of layers. In depth-adaptive looped LMs, tokens exit the recurrent core at different depths. Consequently, a single batch contains tokens requiring different operations: some need another recurrent step, while others that have exited must run the non-looped "coda" (LM head) to generate the next token. Existing token-level schedulers cannot handle this heterogeneity because they cannot route different tokens to different stages within a single forward pass. Furthermore, looped architectures introduce "boundary stages" (prelude and coda) that run once per token but at different frequencies than the loop, adding scheduling complexity. No existing inference system supports end-to-end serving of these models with dynamic depth.

Methodology: Continuous Depth Batching (CDB)

The authors introduce Continuous Depth Batching (CDB), an inference method that schedules at the granularity of individual loop iterations rather than token boundaries. CDB addresses three primary challenges: stage heterogeneity, critical path exits, and ragged Key-Value (KV) caches.

1. Queue-Based Decode Scheduler

CDB decomposes the generation process into four distinct queues:

Prefill: Waiting requests.
Prelude: Tokens awaiting the initial embedding/preprocessing.
Recurrent Core: Tokens ready for their next loop step.
Coda: Tokens that have exited the core and await the final prediction.

The scheduler treats these stages as separate operations, allowing flexible execution orders. The paper implements two modes:

No-Refill: A fixed order (Prelude $\to$ Core $\to$ Coda). When tokens exit, the batch shrinks, and no new tokens enter until the next decode round.
Refill: Stages run in any order based on priority. When a token exits, its slot is immediately filled by a new token from the prefill queue. This maximizes batch size and throughput, leveraging the fact that the recurrent core reuses the same weights regardless of the loop iteration count.

2. Depth-Aware KV Cache

Standard KV caches are indexed by depth, creating gaps when tokens exit early. CDB proposes two layouts:

Last-Exited Cache: Fills empty slots for skipped steps by copying the final state of the exited token. This preserves attention semantics but wastes memory and compute on copying.
Shared Cache: Each physical layer maintains only one KV slot, which is overwritten at every loop step. This reduces memory footprint by a factor of $r_{max}$ but alters attention semantics (tokens attend to the most recent state of previous tokens regardless of their exit depth). The authors note that modern attention kernels (e.g., PagedAttention, FlashAttention-3) are compatible with both.

3. Asynchronous Execution and Lookahead

To minimize the overhead of CPU-GPU synchronization, CDB employs asynchronous execution:

Overlapping: The CPU prepares the next batch (based on exit decisions from step $N-1$ ) while the GPU executes the current step ( $N$ ).
Lookahead Gate: To eliminate GPU idle time caused by waiting for exit decisions, the scheduler uses a "lookahead gate." The exit signal at step $r$ determines participation in step $r+2$ rather than $r+1$ . This gives the CPU a full loop step to prepare the next batch without wasting compute, reducing GPU idle time to nearly zero (0.67% in experiments).

Key Contributions

End-to-End Implementation: The first complete implementation of CDB for depth-adaptive looped LMs, including a serving scheduler, exit gate integration, and KV-cache handling.
Theoretical Bounds: Derivation of a FLOPs-based upper bound on speed-up from adaptive depth and a hardware-aware roofline model to identify the batch-size regimes where the "refill" mode is most beneficial (specifically in memory-bound regimes).
Empirical Evaluation: Comprehensive evaluation on two architectures (Ouro 1.4B and Huginn 3.5B) demonstrating the impact of model structure, loop count, and boundary stage costs on inference efficiency.

Experimental Results

The authors evaluated CDB on Ouro 1.4B (fully looped, cheap boundary stages) and Huginn 3.5B (2-4-2 split, expensive boundary stages) using Alpaca and ShareGPT workloads.

Throughput: CDB achieves up to 99% of the theoretical FLOP-based speed-up bound.
- Ouro 1.4B: CDB (refill) achieves 1.5–1.58× higher offline throughput compared to standard continuous batching (CB).
- Huginn 3.5B: Despite higher boundary stage costs, CDB achieves up to 1.9× higher throughput.
Latency: Under dynamic online serving loads, CDB reduces normalized latency by 45–90% compared to standard CB.
Refill Efficacy: The "refill" mode is most effective in the memory-bound regime (small batch sizes), where early exits would otherwise leave GPU compute idle. In the compute-bound regime (large batches), the advantage diminishes as wall-clock time scales with total work regardless of early exits.
Accuracy: The authors note that the lookahead gate (deciding exit one step early) does not significantly degrade accuracy compared to standard exit gates. However, they observe that KV cache sharing (specifically the single-slot shared layout) can cause accuracy collapse in Ouro unless the model is trained for it, whereas Huginn remains robust.

Significance and Claims

The paper claims that looped LMs served with CDB are not only parameter-efficient but also inference-efficient. By decomposing the variable-depth forward pass into a scheduling problem, CDB enables existing hardware to handle adaptive depth without the massive overhead previously thought to make such models impractical for serving.

The authors emphasize that the efficiency gains are heavily dependent on the co-design of the model architecture and the inference system. Specifically:

Boundary Stages: Models with expensive boundary stages (like Huginn) see reduced speed-ups because tokens must wait for these stages before re-entering the queue for refill.
Memory vs. Compute: The "refill" strategy is critical for overcoming the memory-bound nature of small-batch inference, effectively hiding the weight-reload costs of the recurrent core.
Trade-offs: The work highlights the tension between accuracy and efficiency; while depth-adaptive decoding saves compute, the choice of exit thresholds and KV cache policies directly impacts model performance.

The paper concludes that CDB provides a viable path to serving depth-adaptive models, realizing the theoretical benefits of variable-depth inference in practical, high-throughput scenarios.