APE: Faster and Longer Context-Augmented Generation via Adaptive Parallel Encoding

TL;DR

The paper introduces Adaptive Parallel Encoding (APE), a method that significantly speeds up context-augmented generation by efficiently parallelizing context encoding while maintaining high performance.

Contribution

APE aligns parallel encoding with sequential encoding using shared parameters, enabling faster and scalable context-augmented generation without performance loss.

Findings

Achieves 4.5× speedup in inference time.

Preserves 98% and 93% of sequential encoding performance on RAG and ICL tasks.

Scales to hundreds of contexts in parallel.

Abstract

Context-augmented generation (CAG) techniques, including RAG and ICL, require the efficient combination of multiple contexts to generate responses to user queries. Directly inputting these contexts as a sequence introduces a considerable computational burden by re-encoding the combined selection of contexts for every request. To address this, we explore the promising potential of parallel encoding to independently pre-compute and cache each context's KV states. This approach enables the direct loading of cached states during inference while accommodating more contexts through position reuse across contexts. However, due to misalignments in attention distribution, directly applying parallel encoding results in a significant performance drop. To enable effective and efficient CAG, we propose Adaptive Parallel Encoding ($\textbf{APE}$), which brings shared prefix, attention temperature, and scaling factor to align the distribution of parallel encoding with sequential encoding. Results on RAG and ICL tasks demonstrate that APE can preserve 98% and 93% sequential encoding performance using the same inputs while outperforming parallel encoding by 3.6% and 7.9%, respectively. It also scales to many-shot CAG, effectively encoding hundreds of contexts in parallel. Efficiency evaluation shows that APE can achieve an end-to-end 4.5$\times$ speedup by reducing 28$\times$ prefilling time for a 128K-length context.