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

ICLR 2025 Conference Submission2715 Authors

Keywords: LLM inference acceleration, prompt tuning, hardware-aware design

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 two-stage tree pruning algorithm 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 will be open-sourced upon acceptance of the paper.

Primary Area: generative models

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 2715