BUDDY: BUdget-Driven DYnamic Depth Routing for Adaptive Large Language Model Inference
2026-06-08 • Machine Learning
Machine Learning
AI summaryⓘ
The authors address the problem of high costs when using large language models by introducing Buddy, a method that dynamically skips some layers of the model to save computation based on user-set limits. Buddy uses a small decision tool to pick which parts of the model to run depending on the input and adapts during text generation by using cached information efficiently. If no specific budget is given, Buddy can estimate a good balance between speed and accuracy on its own. Testing on popular models shows Buddy matches or improves on existing methods while offering more flexible control over computation.
large language modelsinference costdepth pruningTransformer blocksdynamic routingkey-value cachedecode-time adaptationcompute budgetLlama modelsQwen models
Authors
Yuhua Zhou, Shaoqi Yu, Shichao Weng, Changhai Zhou, Mingze Yin, Fei Yang, Aimin Pan
Abstract
Large language models (LLMs) incur high inference cost due to their depth and parameter scale. Depth pruning can reduce latency by skipping redundant Transformer blocks, but existing methods (i) provide limited control under user-specific compute budgets and (ii) typically fix the routing path, failing to adapt as the context grows during decoding. We propose Buddy, a budget-driven dynamic depth routing framework. Buddy uses a lightweight Decision Module to score intermediate layers conditioned on the input and deterministically executes the top-k layers to satisfy a given budget. To support decode-time adaptation, Buddy reuses the first-layer KV cache as a low-overhead global context source and pools it together with the newest token representation before each routing decision. When no explicit budget is provided, an optional Budget Predictor estimates an input-dependent compute level to balance quality and efficiency. Experiments on Llama-family and Qwen models show that Buddy is competitive with strong static pruning baselines and often improves the accuracy-compute trade-off, while uniquely supporting strict budget control, decode-time rerouting, and multiple budgets within a single trained model.