Prompt Cache: Reutilización Modular de Atención para Inferencia de Baja Latencia

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

Resumen

Presentamos Prompt Cache, un enfoque para acelerar la inferencia en modelos de lenguaje de gran escala (LLM) mediante la reutilización de estados de atención entre diferentes prompts de LLM. Muchos prompts de entrada contienen segmentos de texto superpuestos, como mensajes del sistema, plantillas de prompts y documentos proporcionados como contexto. Nuestra idea clave es que, al precalcular y almacenar los estados de atención de estos segmentos de texto frecuentes en el servidor de inferencia, podemos reutilizarlos eficientemente cuando estos segmentos aparecen en los prompts de los usuarios. Prompt Cache emplea un esquema para definir explícitamente estos segmentos de texto reutilizables, llamados módulos de prompt. El esquema garantiza la precisión posicional durante la reutilización de los estados de atención y proporciona a los usuarios una interfaz para acceder a los estados almacenados en su prompt. Utilizando una implementación prototipo, evaluamos Prompt Cache en varios LLM. Demostramos que Prompt Cache reduce significativamente la latencia en el tiempo hasta el primer token, especialmente en prompts más largos como respuestas a preguntas basadas en documentos y recomendaciones. Las mejoras van desde 8x en inferencia basada en GPU hasta 60x en inferencia basada en CPU, todo ello manteniendo la precisión de la salida y sin necesidad de modificar los parámetros del 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.