Summary of "Depth-Adaptive Inference of Looped Language Models via Continuous Depth Batching"

Summary (Overview)

• First end-to-end implementation of Continuous Depth Batching (CDB) for looped language models, enabling efficient depth-adaptive inference where tokens exit the recurrent core at different depths.
• CDB achieves up to 99% of the estimated maximum speedup on fully looped architectures (Ouro 1.4B) and 83–96% on models with costly boundary stages (Huginn 3.5B), compared to standard continuous batching at full depth.
• Key technical contributions: stage-wise scheduling with refill/no-refill modes, depth-aware KV cache designs, and asynchronous lookahead scheduling that reduces GPU idle time to nearly zero.
• Architectural insights: fully looped models (minimal non-looped layers) are best suited for depth-adaptive inference; boundary-stage cost, KV cache design, and loop count are the primary determinants of inference efficiency.
• Theoretical framework: FLOP-based speedup bounds and hardware-aware latency models that predict when refill is beneficial (memory-bound vs. compute-bound regimes).

Introduction and Theoretical Foundation

Background and Motivation

Looped language models repeatedly apply a block of shared layers (the recurrent core) to refine latent representations before predicting a token. This architecture offers two key advantages:

1. Parameter efficiency: Sharing layers reduces parameter memory relative to an equally deep standard transformer.
2. Depth-adaptive decoding: Tokens can exit the recurrent core after different numbers of loops, spending less compute on "easy" tokens and more on "hard" ones.

However, the compute savings from depth-adaptive decoding do not automatically translate into faster inference. The central challenge: tokens with different loop counts cannot share a uniform forward pass and therefore cannot be handled by standard batching systems like vLLM.

Architecture Definition

The looped-LM architecture is split into three stages:

• Prelude $P_\theta$: processes each token before entering the loop (runs once)
• Recurrent core $R_\theta$: shared layers applied $r_t \leq r_{max}$ times per token
• Coda $C_\theta$: predicts the next token after exit (runs once)

The architecture is denoted as prelude-core-coda layer counts. For example, a 2-4-2 model executes $2 + 4r_t + 2$ layers, giving effective depth 20 at $r_t = 4$ loops. A fully looped 0-4-0 model has only embedding and LM head outside the core.

Three Main Challenges for CDB

1. Stage heterogeneity: CDB must schedule prefill, prelude, recurrent core, and coda separately, each at different batch sizes and frequencies.
2. Missing KV states: When a token exits at depth $r_t$, it writes no KV states for deeper loop steps, leaving later tokens without states to attend to.
3. Critical-path exits: Exit decisions are available only after a loop step finishes, but determine which tokens enter the next step, potentially leaving the GPU idle.

Methodology

Stage-Wise Scheduling

The scheduler maintains four queues corresponding to inference stages: inactive requests (prefill), decode tokens (prelude), recurrent items (loop steps), and exited tokens (coda). Two execution modes are implemented:

No-refill mode: Fixed order (prelude → recurrent steps → coda). Exited tokens are removed, reducing batch size on subsequent loop steps.

Refill mode: Dynamic ordering where freed slots are filled by new tokens. Priority rules:

1. Prelude runs directly after prefill or coda (no routing decision needed)
2. Coda has highest priority (predicts next tokens for refill)
3. Prefill comes next (admitting prompts early enables larger decode batches)
4. Recurrent core runs otherwise

A minimum coda batch size K controls the trade-off: larger K makes coda/prelude batches more efficient but delays refill and temporarily reduces recurrent batch size.

Depth-Aware KV Cache

Two designs for handling missing KV states from early exits:

• Last-exited cache: Copies a token's final KV state into skipped slots upon exit. Preserves accuracy, avoids recomputation, but retains $r_{max}$ slots per token and adds copying overhead.
• Shared cache: Gives each recurrent layer a single KV slot per token (reduces memory by factor $r_{max}$). Each loop step overwrites the slot; later tokens attend to the most recent state. May require specialized training.

Asynchronous Scheduling with Lookahead Gate

To avoid GPU idle time between stages, batches are prepared and transferred asynchronously while the GPU processes the current batch. The lookahead gate solves the critical-path problem: the exit signal at step $r$ determines whether a token participates in step $r+2$ instead of $r+1$, giving the CPU a full loop step to prepare the next batch. This constrains minimum exit depth to $r_{min} = 2$ but does not reduce accuracy (verified empirically).

Empirical Validation / Results

Efficiency Limits: Theoretical Bounds

The FLOP-based upper bound on decode speedup is:

$$S_{dec} = \frac{F_0 + r_{max} F_r}{F_0 + \bar{d} F_r} \leq \frac{r_{max}}{\bar{d}} \tag{1}$$

where $F_r$ = FLOPs of one recurrent-core application, $F_0$ = boundary-stage FLOPs, and $\bar{d}$ = mean exit depth.

Accounting for prefill (fraction $\alpha$ of time):

$$S_{e2e} = (\alpha + (1 - \alpha) / S_{dec})^{-1} \tag{2}$$

The recurrent step latency model:

$$t(N_c, N_b) = \delta + \gamma N_b + \alpha N_c$$

where $\delta$ = fixed weight-reload cost, $\gamma N_b$ = core computation cost, and $\alpha N_c$ = KV cache streaming cost.

Offline Throughput Results

Key findings:

• Refill is most effective in the memory-bound regime (around/below saturation batch size $B^*$), where shrinking batches saves little time.
• At full depth, CDB retains 95–99% of CB throughput, showing minimal scheduling overhead.
• Prefill accounts for 3.7–36% of GPU time, reducing end-to-end speedup below decode-only estimates.

Online Serving

CDB sustains higher request rates before latency rises sharply, translating offline throughput gains into greater serving capacity. For Ouro, refill provides the largest capacity gain; for Huginn, refill and no-refill perform similarly.

Ablations: Minimum Coda Batch Size K

• Huginn: $K = W/4$ best balances coda efficiency vs. refill delay.
• Ouro: $K = 1$ (immediate coda execution) performs best due to cheap boundary stages.
• Removing transformer layers from Huginn's boundary stages (making it fully looped) makes $K=1$ near-optimal, confirming boundary-stage cost is a central efficiency constraint.

Theoretical and Practical Implications

When Does Depth-Adaptive Inference Help?

Three conditions must hold for meaningful speedup:

1. Sufficient early exits at acceptable accuracy (model must be trained appropriately)
2. Recurrent computation must dominate total inference time (expensive prefill limits gains)
3. The inference system must convert skipped FLOPs into wall-clock savings (requires refill in memory-bound regimes and asynchronous scheduling)

Architecture Co-Design Guidance

• Keep boundary stages small: Costly non-looped layers (prelude/coda) reduce refill opportunities and complicate scheduling. Fully looped architectures (0-4-0) are best suited for depth-adaptive inference.
• Consider loop count trade-offs: More loops offer finer-grained exit control but each skipped loop saves a smaller fraction of total recurrent compute.
• KV cache design matters: Shared caches reduce memory footprint and implementation complexity but may require specialized training.

Design of Exit Gates

The lookahead gate constrains minimum exit depth to $r_{min} = 2$. For models that produce useful outputs after a single loop, a promising alternative is to predict only the first loop-exit before the recurrent core, then use lookahead decisions for later loops.

Conclusion

The paper introduces the first efficient inference method for depth-adaptive looped LMs, demonstrating that:

• CDB realizes most of the available speedup from early exits (up to 99%), making depth-adaptive decoding practically viable.
• Fully looped architectures with minimal boundary stages are best suited for this paradigm.
• Co-designing looped architectures and serving systems can lead to substantially faster inference.

Future directions:

• Models that allow more frequent exits at high accuracy
• Reliable prediction of complete exit depth before entering the recurrent core
• Understanding when shared KV caches work without specialized training
• Extension to production features: multi-GPU execution, preemption, prefix caching, speculative decoding

Limitations: Evaluation covers only two looped architectures on a single H100 GPU, without production serving features like multi-GPU execution or speculative decoding.