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

Dear Reviewer 3h4G,

Thanks for your valuable feedback. Here is our response.

Q1: Speedup benefit from tree attention technique and effectiveness of PPD without tree attention as compared to Medusa.

We would like to highlight that almost all speculative/parallel decoding methods benefit significantly from the use of tree attention. As shown in Figure 7b, PPD outperforms all other parallel decoding methods, including Medusa, without the use of tree attention.

Q2: Effect of hardware changes on workload and speedup ratio

In Figure 5, we present speedup ratios across different GPUs, demonstrating the effectiveness of PPD in adapting to various hardware configurations.

Figures 6b and 6c provide a more detailed analysis of how PPD adjusts the tree size based on the hardware's capabilities. For GPUs with lower computational power, the optimal speedup ratio is achieved with a smaller tree size, as larger tree sizes would incur higher latency overheads, reducing the overall efficiency.

Dear Reviewer F2bw,

Thanks for your detailed and valuable feedback. Here is our response.

Q1.1: Limited improvement over previous methods.

Our approach has key advantages beyond speedups, like memory saving and cheaper training cost as mentioned in Figure 7a and Table 2. More free memory enables the use of larger batch sizes, resulting in the potential of improved throughput. Memory efficiency is also vital for long-context tasks, where a large KV cache size could otherwise lead to OOM errors. Additionally, reducing training time significantly lowers training costs, making PPD more accessible for local deployments. Faster training is particularly valuable in online learning scenarios, where the target model distribution may shift over time, requiring frequent retraining.

Q2: Further evaluation on larger models

First, we would like to emphasize that our primary focus, as highlighted in the paper, is on edge and mobile settings. The design of our method, particularly features such as memory savings, is specifically tailored to address the constraints and requirements inherent to these scenarios.

Second, the decrease in acceptance length observed in Table 1 as model sizes increase is attributable to the smaller optimal tree size, rather than a decline in the capability of PPD. In fact, as demonstrated in Figure 6a, the prediction accuracy of PPD improves for larger models due to their greater depth and wider token embeddings, which enhance its overall performance.

Q3: Performance under high GPU utilisation

First, given that our primary use case is the mobile/edge setting (e.g., AIPC), a single-batch scenario is common. In such settings, idle compute resources are often present due to the memory-bound nature of LLM inference, making our method particularly efficient.

Second, in high-utilization scenarios (e.g., large batch sizes), memory availability becomes a critical concern, as highlighted by FlexGen. Unlike previous speculative/parallel decoding methods, PPD incurs negligible memory overhead, which helps mitigate these memory constraints and supports efficient inference in such scenarios.

Lastly, we provide analysis of the effect of batch size on the speedup ratio in Appendix E, showing that significant speedups can still be achieved with increased batch sizes. Furthermore, our method is orthogonal to recent advancements in batched speculative decoding optimization, making it compatible with such approaches.

Q4: Unnecessary computation for incorrect guess tokens.

It is important to note that PPD adopts the same guess-and-verification procedure as previous parallel decoding methods such as Medusa, LookAhead Decoding, and Blockwise Parallel Decoding. Therefore, the additional computation observed is not introduced by our proposed method.

Q5: Details on EPT training and usage.

Please refer to Appendix D for the ablation studies on the design choices of EPT, including the attention mechanism, hyperparameters, and other relevant factors.

In summary, each EPT is associated with its own trainable embedding, which is optimized via backpropagation during training. This ensures an effective initialization that supports accurate multi-token generation. A single prompt token is composed of multiple EPTs. To produce a single output token prediction, the logits generated by these multiple EPTs are averaged, functioning similarly to an ensemble approach. For example, if three EPTs are used per prompt token, the next-next word prediction is obtained by averaging the logits produced by the three EPTs.

Q6: Explanation on x-axis in Figure 8.

It refers to the tree size as discussed in the paper. The acceptance ratio of 2.6 for a tree size of 500 is reasonable, given that only 3 prompt tokens are used, which effectively bounds the acceptance ratio to 3.

Q7: How to combine PPD with speculative decoding.

We apply PPD to the draft model to accelerate its draft generation, with the draft model running on the same GPU. In this setup, there exist 2 tiers of guess&verification processes. The guess-and-verification process is combined for PPD acceleration of the draft mode, while, at the top level, the guess and verification processes are separated between the draft and target models.

Dear reviewer e6Zd,

We thank you for your time and feedback on our submission. Please see our responses below.

Unlike some other submission, this one has no implementation code.

The reviewer’s comment is untrue. We did submit our code as supplementary material in our original submission. Additionally, we also attached the model weights for PPD for the convenience of experiments. Everyone could download the codebase and easily reproduce our results.

Extra RAM and training time may often not be the limited factors.

In practical applications, RAM usage is a critical factor when deploying AI models on mobile devices and AI PCs with limited memory capacity. Leading companies such as Apple and Samsung are actively pursuing memory-efficient and training-efficient solutions for deploying LLM in their products. While increasing speedup from 2–3 times (as seen with Medusa) to 3–4 times (as seen with Eagle) may or may not enhance the user experience, reducing memory usage to levels compatible with deployment on mobile or IoT devices is the decisive factor that determines whether users can practically utilize LLMs for everyday applications. Therefore, significant memory saving is actually more important than minor speedup improvement.

It uses the same core idea of learnable tokens fed to transformer model to generate speculative continuation.

We would like to highlight that our main idea is fundamentally different from BiTA in the following ways:

1. Key Difference-1: Although both papers use the name “prompt tokens”, they are referring to different concepts. Prompt token in BiTA refers to the prompt embeddings appended to each layer as part of the model, which introduces extra computation, memory usage, and complexity in training. In contrast, prompt token in PPD uses the idea of ensemble tokens as part of the input sequence, which decouples the number of trainable parameters from the number of prompt tokens used. Our prompt tokens are virtualized and correspond to arbitrary actual tokens, which increases the acceptance ratio. Note that our approach does not prepend any tokens.
2. Key Difference-2: We also propose a novel tree pruning algorithm to minimise the latency overhead introduced by using tree attention with prompt tokens. This is also not mentioned anywhere in the BiTA paper. We do not appreciate the fact that the reviewer claim this is a “ secondary contributions”. The pruning algorithm, as shown in the paper, leads to huge speedup ratio improvement.
3. Key Difference-3: We consider the compute and memory capabilities of different GPUs, and propose hardware-aware optimizations to further enhance the flexibility and efficiency of our approaches on diverse hardware.

Still, we thank the reviewer for bringing up the BiTA paper. We acknowledge the good work in this paper. We will cite the paper in the revised version and make comparison with our method, to avoid any similar confusion.

Lack of proper comparison with Eagle

We presented the speedup comparison here.

Our proposed method serves as a cost-effective and memory-efficient alternative to speculative decoding, specifically designed for edge and mobile deployments. We emphasize that both memory usage and training time are critical metrics in this context.

Memory Efficiency

Training Cost

More free memory enables the use of larger batch sizes, resulting in improved throughput. Memory efficiency is also vital for long-context tasks, where a large KV cache size could otherwise lead to OOM errors. Additionally, reducing training time significantly lowers training costs, making PPD more accessible for local deployments. Faster training is particularly valuable in online learning scenarios, where the target model distribution may shift over time, requiring frequent retraining.

Could you provide results of your method for more modern models, like Qwen and Llama-3?

It is our intention to compare with SOTA parallel decoding methods like Medusa. Thus, for experiments, we choose the models that most parallel decoding methods also support so that we could make a fair comparison. However, Llama-3 is only released recently and it is hard for us to make a comparison with other parallel decoding methods. If the reviewer insists on requesting this evaluation and provides valid reasons to ensure a fair comparison with other approaches, please do not hesitate to let us know.

Dear Reviewer KH3T,

Thank you for your time and valuable feedback. Please check our responses below.

Q1: Limited speedup over previous methods.

Our approach has key advantages beyond speedups, like memory saving and cheaper training cost as mentioned in Figure 7a and Table 2. More free memory enables the use of larger batch sizes, resulting in the potential of improved throughput. Memory efficiency is also vital for long-context tasks, where a large KV cache size could otherwise lead to OOM errors. Additionally, reducing training time significantly lowers training costs, making PPD more accessible for local deployments. Faster training is particularly valuable in online learning scenarios, where the target model distribution may shift over time, requiring frequent retraining.

Q2: The need of fine-tuning for deployment

Thank you for the valuable feedback. We would like to point out that Eagle, a SOTA speculative decoding method, also relies on model-specific features during draft generation. As a result, it encounters similar challenges when maximizing the speedup ratio. PPD reduces the training cost significantly compared to Eagle.

We agree that exploring the transferability of trained PPD tokens across models and enabling training-free deployment are promising directions for future work, and we are excited to investigate these possibilities further.

Q3: The relative memory saving is smaller for larger models.

We would like to emphasize that, as mentioned in the paper, our primary target use case is edge and mobile devices, where memory savings and reduced training costs are critical considerations. In such scenarios, quantized models are commonly employed to further optimize resource efficiency, making the memory saving in the draft model more pronounced.