Hardware-Aware Parallel Prompt Decoding for Memory-Efficient Acceleration of LLM Inference

Hao Mark Chen, Wayne Luk, Ka Fai Cedric Yiu, Rui Li, Konstantin Mishchenko +2 more
5/28/2024
cs.LGcs.CL

Official code: hmarkc/parallel-prompt-decoding (50 stars, Apache-2.0, last commit 2024-10-22).

Abstract

The auto-regressive decoding of Large Language Models (LLMs) results in significant overheads in their hardware performance. While recent research has investigated various speculative decoding techniques for multi-token generation, these efforts have primarily focused on improving processing speed such as throughput. Crucially, they often neglect other metrics essential for real-life deployments, such as memory consumption and training cost. To overcome these limitations, we propose a novel parallel prompt decoding that requires only $0.0002$% trainable parameters, enabling efficient training on a single A100-40GB GPU in just 16 hours. Inspired by the human natural language generation process, $PPD$ approximates outputs generated at future timesteps in parallel by using multiple prompt tokens. This approach partially recovers the missing conditional dependency information necessary for multi-token generation, resulting in up to a 28% higher acceptance rate for long-range predictions. Furthermore, we present a hardware-aware dynamic sparse tree technique that adaptively optimizes this decoding scheme to fully leverage the computational capacities on different GPUs. Through extensive experiments across LLMs ranging from MobileLlama to Vicuna-13B on a wide range of benchmarks, our approach demonstrates up to 2.49$\times$ speedup and maintains a minimal runtime memory overhead of just $0.0004$%. More importantly, our parallel prompt decoding can serve as an orthogonal optimization for synergistic integration with existing speculative decoding, showing up to $1.22\times$ further speed improvement. Our code is available at https://github.com/hmarkc/parallel-prompt-decoding.

View on arXivView PDF

Code Implementations(1)

Apache-2.0

Efficient LLM Inference Acceleration using Prompting

502Shell, PythonMay 24, 20241 years agoApache-2.0

Similar Papers

Cite this paper

@article{chen2024hardwareaware,
  title  = {Hardware-Aware Parallel Prompt Decoding for Memory-Efficient Acceleration of LLM Inference},
  author = {Hao Mark Chen and Wayne Luk and Ka Fai Cedric Yiu and Rui Li and Konstantin Mishchenko and Stylianos I. Venieris and Hongxiang Fan},
  year   = {2024},
  eprint = {2405.18628},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2405.18628v3}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2405.18628.svg)](https://codeofpaper.com/paper/2405.18628)

Discussion