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

Continuous depth batching (CDB) is an inference method that schedules at the granularity of individual loop iterations, routing tokens between stages within the forward pass. We build CDB in two modes (Fig. 1, right). No-refill only lets the recurrent batch shrink, so exited tokens wait until the entire batch finishes before running the coda together. Refill maximizes throughput by scheduling stages independently via a priority queue system, which interleaves boundary stages with recurrent steps and fills freed recurrent batch slots with new tokens (Bae et al., 2025a). Both modes require new KV-cache layouts, as tokens at different exit depths leave gaps in the standard depth-indexed cache. Finally, to minimize per-step scheduling overhead, we overlap batch preparation with GPU execution by introducing a lookahead gate that makes exit decisions one loop step in advance.

We analyze the maximum speed-up from depth-adaptive inference by deriving a FLOPs-based bound. We also fit a hardware-aware roofline model that identifies the batch-size regimes where the refill mode is most beneficial. On Ouro 1.4B (Zhu et al., 2025) and Huginn 3.5B (Geiping et al., 2025), CDB realizes up to 99% of the theoretical maximum speed-up, which translates to 1.5–1.9× higher offline throughput and 45–90% lower normalized latency under online serving.

Contributions.

1. We implement continuous depth batching (CDB) end to end for depth-adaptive looped LMs, with a high-throughput refill and a simpler no-refill execution mode.

2. We derive an upper bound on the speed-up from adaptive depth and identify the batching regime in which refill is beneficial.

3. We evaluate CDB on two looped LM architectures and show how their structure, the loop count and the cost of boundary stages, enables efficient inference.

Improvements to AI systems:

1. Implement dynamic token-routing inference schedulers that treat each loop iteration as a schedulable unit, allowing tokens to exit early based on their individual confidence or task difficulty. This reduces wasted compute on easy tokens while preserving quality on hard ones.

2. Add a priority-queue-based stage interleaver that overlaps boundary computations (e.g., embedding/output projections) with recurrent steps, filling freed batch slots with new tokens immediately. This maximizes GPU utilization under variable token exit rates.

3. Introduce a lookahead exit-decision gate that predicts token exits one loop step in advance, enabling batch preparation to be overlapped with GPU execution. This hides scheduling latency and keeps the GPU busy during dynamic depth changes.

4. Redesign KV-cache layouts to support non-contiguous token depths, using sparse or segmented indexing that avoids gaps when tokens exit at different iterations. This enables efficient memory access for variable-depth batches.

5. Build a hardware-aware roofline model that predicts optimal batch-size regimes for refill vs. no-refill modes, allowing runtime systems to automatically switch between throughput-optimized and latency-optimized scheduling based on current load and hardware.

What the improved AI system can do:

• Serve looped language models (e.g., Ouro, Huginn) with 1.5–1.9× higher offline throughput and 45–90% lower normalized latency under online serving, by dynamically adapting compute per token.

• Achieve up to 99% of the theoretical maximum speed-up from depth-adaptive inference, as derived from a FLOPs-based bound.

• Handle mixed workloads (e.g., short queries vs. long reasoning tasks) without manual batching policies, automatically interleaving tokens of different exit depths.

• Operate efficiently on both high-throughput batch servers and low-latency interactive endpoints, by selecting the appropriate scheduling mode (refill or no-refill) via the roofline model.

• Reduce energy consumption per token by skipping unnecessary recurrent steps for easy inputs, while maintaining accuracy on complex ones.

A main promise of looped language models (LMs) is depth-adaptive inference. By iterating 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, this adaptivity breaks standard batching: tokens in the same batch now require a different number of loops, so there is no unified forward pass, making efficient inference difficult. Standard inference frameworks like vLLM schedule on the token level and cannot handle this because tokens need to be removed from the batch within the forward pass. Loop-level scheduling has been proposed as a solution, but never implemented end to end. The key challenge is that looped architectures also contain non-looped boundary stages (e.g., token embedding and LM head) that must be scheduled at different frequencies than the loop. We introduce continuous depth batching (CDB), which schedules at the granularity of individual loop iterations. CDB handles boundary stages and loop steps in separate priority queues, makes exit decisions one step ahead, and overlaps all scheduling work with GPU computation. On Ouro 1.4B and Huginn 3.5B, CDB can realize up to 99% of the theoretical maximum speed-up from adaptive-depth, translating to 1.5 - 1.9 times higher offline throughput and 45 - 90% lower normalized latency under dynamic serving load.