Elamur Research
☰
home / papers / 2024 / Break the Sequential Dependency of LLM Inference Using Lookahead Decoding

Break the Sequential Dependency of LLM Inference Using Lookahead Decoding

Written for the Research library (claude-opus-5-5) from the paper text, in our own words: the license "arXiv — non-exclusive license to distribute" does not allow redistributing the text itself. Original — arxiv.org/pdf/2402.02057

Summary

Generating text with LLMs one token at a time is slow. Each step is limited by memory bandwidth, so most of the GPU's compute sits idle. Speculative decoding helps, but it needs a draft model, which is hard to train and often fails to transfer across base models or datasets. The authors propose Lookahead Decoding, which needs no extra model, training or datastore. It treats decoding as a nonlinear system solved by Jacobi iteration. A lookahead branch uses a 2D window over recent iteration history to produce many disjoint n-grams at once. These n-grams are cached in a pool, and a verification branch checks candidates in the same forward pass, so the output distribution stays exactly the same. The method works with FlashAttention, supports several sampling schemes and scales across GPUs ("Lookahead Parallelism"). On LLaMA-2 it reaches about 1.8x speedup on MT-Bench and up to 4x on code completion with 8 GPUs. This matters because it turns idle compute into lower latency without any auxiliary model.

Key points

  • Speedups: about 1.8x on MT-Bench, and up to 4x on code tasks (HumanEval, ClassEval) using Lookahead Parallelism on 8 GPUs. Results are for LLaMA-2 models, with Python and CUDA implementations available at github.com/hao-ai-lab/LookaheadDecoding.
  • Scaling law: the number of decoding steps falls linearly as per-step log(FLOPs) grows. In other words, each further step reduction requires exponentially more compute per step.
  • Tuning knobs: W is the lookahead window, N is the n-gram size and G is the verification size. Near-optimal settings on A100 (with G = W) are:
    • 7B: W=15, N=5
    • 13B: W=10, N=5
    • 34B: W=7, N=5
  • Main limitation: the configurations above cost roughly 120x, 80x and 56x extra FLOPs per step. Gains therefore shrink in compute-bound settings, such as serving large batches. Gains are also smaller on weaker GPUs: about 30% on RTX 3090 versus more than 50% on A100 for MT-Bench.
  • Ablation: both branches are needed, and balanced lookahead and verification branches beat lopsided setups. The method was compared against the prompt lookup implementation in transformers v4.37. Adding prompt-based n-grams to the pool raises the speedup further.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Computing & GPUs

Cited by

109 papers of our library cite it (of the 140673 whose bibliographies we read).

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2402.02057 · arXiv 2402.02057

@misc{fu2024break,
  title = {{Break the Sequential Dependency of LLM Inference Using Lookahead Decoding}},
  author = {Fu, Yichao and Bailis, Peter and Stoica, Ion and Zhang, Hao},
  year = {2024},
  doi = {10.48550/arxiv.2402.02057},
  eprint = {2402.02057},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2402.02057}
}

References

43 references, 9 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Automatic tensor parallelism for huggingface models, 2023. URL https://www.deepspeed.ai/tutoria ls/automatic-tensor-parallelism. ↗
  2. Aminabadi, R. Y., Rajbhandari, S., Awan, A. A., Li, C., Li, D., Zheng, E., Ruwase, O., Smith, S., Zhang, M., Rasley, J., et al. Deepspeed-inference: enabling efficient inference of transformer models at unprecedented scale. In SC22: International Conference for High Performance Computing, Networking… find
  3. Austin, J., Odena, A., Nye, M., Bosma, M., Michalewski, H., Dohan, D., Jiang, E., Cai, C., Terry, M., Le, Q., and Sutton, C. Program synthesis with large language models, 2021. find our card
  4. Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint arXiv:1607.06450, 2016. find our card
  5. Ben Allal, L., Muennighoff, N., Kumar Umapathi, L., Lipkin, B., and von Werra, L. A framework for the evaluation of code generation models. https://github.com /bigcode-project/bigcode-evaluation-h arness, 2022. ↗
  6. Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., and Dao, T. Medusa: Simple llm inference acceleration framework with multiple decoding heads, 2024. find our card
  7. Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023. find
  8. Chen, M., Tworek, J., Jun, H., Yuan, Q., de Oliveira Pinto, H. P., Kaplan, J., Edwards, H., Burda, Y., Joseph, N., Brockman, G., Ray, A., Puri, R., Krueger, G., Petrov, M., Khlaaf, H., Sastry, G., Mishkin, P., Chan, B., Gray, S., Ryder, N., Pavlov, M., Power, A., Kaiser, L., Bavarian, M., Winter, C.… find our card
  9. Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021. find our card
  10. Dao, T. FlashAttention-2: Faster attention with better parallelism and work partitioning. 2023. find
  11. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, 2022. find
  12. Du, X., Liu, M., Wang, K., Wang, H., Liu, J., Chen, Y., Feng, J., Sha, C., Peng, X., and Lou, Y. Classeval: A manually-crafted benchmark for evaluating llms on class-level code generation, 2023. find
  13. Gugger, S., Debut, L., Wolf, T., Schmid, P., Mueller, Z., Mangrulkar, S., Sun, M., and Bossan, B. Accelerate: Training and inference at scale made simple, efficient and adaptable. https://github.com/huggingface /accelerate, 2022. ↗
  14. He, Z., Zhong, Z., Cai, T., Lee, J. D., and He, D. Rest: Retrieval-based speculative decoding. arXiv preprint arXiv:2311.08252, 2023. find
  15. Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y. The curious case of neural text degeneration, 2020. find
  16. Jiang, A. Q., Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., Casas, D. d. l., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023. find
  17. Kool, W., van Hoof, H., and Welling, M. Ancestral gumbel-top-k sampling for sampling without replacement. Journal of Machine Learning Research, 21(47):1–36, 2020. URL http://jmlr.org/papers/v21/19-985 .html. ↗
  18. Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pp. 19274– 19286. PMLR, 2023. find
  19. Li, Y., Zhang, C., and Zhang, H. Eagle: Lossless acceleration of llm decoding by feature extrapolation, December 2023. URL https://sites.google.com/vie w/eagle-llm. ↗
  20. Lin, C.-Y. ROUGE: A package for automatic evaluation of summaries. In Text Summarization Branches Out, pp. 74–81, Barcelona, Spain, July 2004. Association for Computational Linguistics. URL https://www.ac lweb.org/anthology/W04-1013. ↗
  21. Liu, X., Hu, L., Bailis, P., Stoica, I., Deng, Z., Cheung, A., and Zhang, H. Online speculative decoding, 2023. find
  22. Miao, X., Oliaro, G., Zhang, Z., Cheng, X., Wang, Z., Wong, R. Y. Y., Zhu, A., Yang, L., Shi, X., Shi, C., Chen, Z., Arfeen, D., Abhyankar, R., and Jia, Z. Specinfer: Accelerating generative large language model serving with speculative inference and token tree verification, 2023. find
  23. Narayan, S., Cohen, S. B., and Lapata, M. Don’t give me the details, just the summary! Topic-aware convolutional neural networks for extreme summarization. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, 2018. find
  24. Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Patwary, M., Korthikanti, V. A., Vainbrand, D., Kashinkunti, P., Bernauer, J., Catanzaro, B., Phanishayee, A., and Zaharia, M. Efficient large-scale language model training on gpu clusters using megatron-lm, 2021. find
  25. OpenAI. Gpt-4 technical report, 2023. find our card
  26. Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35:27730–27744, 2022. find
  27. Roziere, B., Gehring, J., Gloeckle, F., Sootla, S., Gat, I., Tan, X. E., Adi, Y., Liu, J., Remez, T., Rapin, J., et al. Code llama: Open foundation models for code. arXiv preprint arXiv:2308.12950, 2023. find our card
  28. Ruan, J. T., Sabir, F., and Chopra, P. Best prompting practices for using the llama 2 chat llm through amazon sagemaker jumpstart, November 2023. URL https: //aws.amazon.com/cn/blogs/machine-lea rning/best-prompting-practices-for-u sing-the-llama-2-chat-llm-through-a mazon-sagemaker-jumpstart/. find
  29. Santilli, A., Severino, S., Postolache, E., Maiorca, V., Mancusi, M., Marin, R., and Rodola, E. Accelerating transformer inference for translation via parallel decoding. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 12336–1235… find
  30. Saxena, A. Prompt lookup decoding, November 2023. URL https://github.com/apoorvumang/promp t-lookup-decoding/. ↗
  31. See, A., Liu, P. J., and Manning, C. D. Get to the point: Summarization with pointer-generator networks. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1073–1083, 2017. find
  32. Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019. find
  33. Song, Y., Meng, C., Liao, R., and Ermon, S. Accelerating feedforward computation via parallel nonlinear equation solving, 2021. find
  34. Stern, M., Shazeer, N., and Uszkoreit, J. Blockwise parallel decoding for deep autoregressive models, 2018. find
  35. Team, G., Anil, R., Borgeaud, S., Wu, Y., Alayrac, J.-B., Yu, J., Soricut, R., Schalkwyk, J., Dai, A. M., Hauth, A., et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805, 2023. find our card
  36. Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023a. find
  37. Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023b. find
  38. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need, 2023. find
  39. Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Scao, T. L., Gugger, S., Drame, M., Lhoest, Q., and Rush, A. M. Transformers: State-ofthe-art natural l… find our card
  40. Yang, N., Ge, T., Wang, L., Jiao, B., Jiang, D., Yang, L., Majumder, R., and Wei, F. Inference with reference: Lossless acceleration of large language models, 2023. find
  41. Zhang, B. and Sennrich, R. Root mean square layer normalization, 2019. find
  42. Zheng, L., Chiang, W.-L., Sheng, Y., Zhuang, S., Wu, Z., Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing, E. P., Zhang, H., Gonzalez, J. E., and Stoica, I. Judging llm-as-a-judge with mt-bench and chatbot arena, 2023. find
  43. A. Algorithms Algorithm 1 Jacobi decoding find