APE: Faster and Longer Context-Augmented Generation via Adaptive
  Parallel Encoding
Authors: name: Xinyu Yang, name: Tianqi Chen, name: Beidi Chen

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
(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.5times
speedup by reducing 28times prefilling time for a 128K-length context.