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

2026-08-10

AI Summary

This work addresses the inefficiency of conventional batching in recurrent language models, where varying inference depths across tokens hinder effective parallel execution. To overcome this, we propose Continuous Depth Batching (CDB), a novel scheduling approach that treats a single recurrent iteration as the fundamental scheduling unit. CDB decouples boundary and recurrent phases, enables early exit decisions, and overlaps scheduling with GPU computation. For the first time, we achieve end-to-end recurrence-level scheduling by managing the two phases through independent priority queues, thereby breaking the deployment efficiency bottleneck of depth-adaptive inference. Evaluated on Ouro 1.4B and Huginn 3.5B models, CDB attains 99% of the theoretical maximum speedup, improves offline throughput by 1.5–1.9×, and reduces normalized latency by 45–90% under dynamic workloads.

Abstract

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.

Problem

• depth-adaptive inference
• looped language models
• batching
• scheduling
• efficient inference

Innovation

• depth-adaptive inference
• looped language models
• continuous depth batching
• dynamic scheduling
• efficient inference