让大模型按需提前退出,加速推理且不改模型。
TIDE: Token-Informed Depth Execution for Per-Token Early Exit in LLM Inference
- 在模型中间层加小路由器,动态决定每字是否提前结束计算。
- 预填充阶段95%以上字词提前退出,延迟降7.2%,吞吐升6.6%。
- 无需重训练,支持主流模型和多种精度,适合部署优化场景。
大型语言模型对每个输入字元均执行全部层的计算,无论难易。本文提出TIDE,一种后训练系统,在周期性检查点层添加微小可学习的路由模块,推理时为每个字元选择隐藏状态已收敛的最早层进行退出。TIDE无需模型重训练,兼容任意HuggingFace因果语言模型,自动识别GPU架构,支持float32、float16与bfloat16通过融合CUDA内核实现。在NVIDIA A100上运行DeepSeek R1 Distill 8B时,预填充阶段实现100%退出率(5%字元于第11层退出,其余在第31层),预填充延迟降低7.2%,单批次吞吐提升6.6%。自回归解码中,98-99%字元提前退出,模型仍能正确完成包含95个唯一输出字元的多步数学题。在Qwen3 8B(36层)上,批量大小为8时吞吐提升8.1%。基于2,000条WikiText样本的校准耗时不足3分钟,生成约4MB的路由器检查点。系统代码共1,308行Python与1,081行CUDA/C++,含74个通过测试。代码开源:https://github.com/RightNow-AI/TIDE
原文摘要 · Abstract (English)
Large language models run every token through every layer, regardless of difficulty. We present TIDE, a post-training system that attaches tiny learned routers at periodic checkpoint layers and, at inference time, selects the earliest layer whose hidden state has converged for each token. TIDE requires no model retraining, works with any HuggingFace causal LM, auto-detects GPU architecture, and supports float32, float16, and bfloat16 through fused CUDA kernels. On an NVIDIA A100 with DeepSeek R1 Distill 8B, TIDE achieves 100% prefill exit rate (5% of tokens exit at layer 11, the remaining at layer 31), reduces prefill latency by 7.2%, and increases single-batch throughput by 6.6%. During autoregressive decoding, 98-99% of tokens exit early while the model correctly solves a multi-step math problem with 95 unique output tokens. On Qwen3 8B (36 layers), throughput improves by 8.1% at batch size 8. Calibration on 2,000 WikiText samples takes under 3 minutes and produces a ~4 MB router checkpoint. The system comprises 1,308 lines of Python and 1,081 lines of CUDA/C++ with 74 passing tests. Code: https://github.com/RightNow-AI/TIDE
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。