Prompt Cache: Reutilização Modular de Atenção para Inferência de Baixa Latência

Prompt Cache: Modular Attention Reuse for Low-Latency Inference

November 7, 2023

Autores: In Gim, Guojun Chen, Seung-seob Lee, Nikhil Sarda, Anurag Khandelwal, Lin Zhong

cs.AI

Resumo

Apresentamos o Prompt Cache, uma abordagem para acelerar a inferência em modelos de linguagem de grande escala (LLM) por meio da reutilização de estados de atenção entre diferentes prompts de LLM. Muitos prompts de entrada possuem segmentos de texto sobrepostos, como mensagens de sistema, modelos de prompt e documentos fornecidos como contexto. Nossa principal percepção é que, ao pré-computar e armazenar os estados de atenção desses segmentos de texto que ocorrem frequentemente no servidor de inferência, podemos reutilizá-los de forma eficiente quando esses segmentos aparecem em prompts de usuários. O Prompt Cache emprega um esquema para definir explicitamente esses segmentos de texto reutilizáveis, chamados de módulos de prompt. O esquema garante a precisão posicional durante a reutilização dos estados de atenção e fornece aos usuários uma interface para acessar os estados armazenados em cache em seus prompts. Utilizando uma implementação protótipo, avaliamos o Prompt Cache em vários LLMs. Demonstramos que o Prompt Cache reduz significativamente a latência no tempo para o primeiro token, especialmente para prompts mais longos, como respostas a perguntas baseadas em documentos e recomendações. As melhorias variam de 8x para inferência baseada em GPU a 60x para inferência baseada em CPU, tudo isso mantendo a precisão da saída e sem a necessidade de modificações nos parâmetros do modelo.

English

We present Prompt Cache, an approach for accelerating inference for large language models (LLM) by reusing attention states across different LLM prompts. Many input prompts have overlapping text segments, such as system messages, prompt templates, and documents provided for context. Our key insight is that by precomputing and storing the attention states of these frequently occurring text segments on the inference server, we can efficiently reuse them when these segments appear in user prompts. Prompt Cache employs a schema to explicitly define such reusable text segments, called prompt modules. The schema ensures positional accuracy during attention state reuse and provides users with an interface to access cached states in their prompt. Using a prototype implementation, we evaluate Prompt Cache across several LLMs. We show that Prompt Cache significantly reduce latency in time-to-first-token, especially for longer prompts such as document-based question answering and recommendations. The improvements range from 8x for GPU-based inference to 60x for CPU-based inference, all while maintaining output accuracy and without the need for model parameter modifications.