Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching
2026-08-10 • Machine Learning
Machine LearningComputation and LanguageDistributed, Parallel, and Cluster Computing
AI summaryⓘ
The authors study a special type of language model that can spend different amounts of effort on each word, doing less work on easy words and more on hard ones. This creates a problem because it makes it hard to process many words together quickly using usual methods. They propose a new method called continuous depth batching (CDB) that smartly schedules parts of the model to work more efficiently by handling different stages separately and planning ahead. Their results show that CDB achieves almost the best possible speed improvements, making the model faster and reducing waiting times during processing.
looped language modelsadaptive inferencebatchingschedulingGPU computationthroughputlatencytoken embeddinglanguage model headpriority queues
Authors
Kristian Schwethelm, Daniel Rueckert, Georgios Kaissis
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.