PagedAttention

PagedAttention es un algoritmo de atención y un gestor de memoria para servir modelos de lenguaje grandes (LLM) que almacena la caché KV de cada solicitud en pequeños bloques de tamaño fijo, ubicados en cualquier parte de la memoria de la GPU, y los gestiona mediante una tabla de bloques por solicitud, como hace un sistema operativo al asignar páginas de memoria virtual. Woosuk Kwon, Zhuohan Li y sus colegas de UC Berkeley y universidades colaboradoras lo presentaron junto con el motor de inferencia vLLM en Efficient Memory Management for Large Language Model Serving with PagedAttention, presentado en SOSP 2023. Al eliminar la fragmentación de las asignaciones de caché contiguas, permite que una GPU mantenga muchas más secuencias a la vez, y el artículo informa de un rendimiento entre 2 y 4 veces superior al de FasterTransformer y Orca con la misma latencia.

El problema de memoria que resuelve#

Durante la generación de texto, un Transformer conserva los vectores de claves y valores de cada token anterior para no tener que recalcularlos. Esa caché crece una entrada por token y capa, y una única secuencia del modelo LLaMA de 13.000 millones de parámetros puede ocupar hasta 1,7 GB. Al recibirse una solicitud se desconoce la longitud de salida, por lo que los primeros sistemas de inferencia reservaban para cada solicitud una región contigua dimensionada para la longitud máxima posible. La mayor parte de esa región quedaba vacía y los huecos entre regiones eran demasiado pequeños para reutilizarlos. El equipo de vLLM midió que estos sistemas desperdiciaban entre el 60 y el 80 % de la memoria de la caché KV, lo que limitaba el tamaño del lote mucho antes de agotar la capacidad de cálculo.

Cómo funciona la paginación#

PagedAttention divide la caché en bloques que contienen las claves y los valores de un número fijo de tokens; vLLM utiliza 16 de forma predeterminada. Una solicitud recibe un bloque nuevo solo cuando se llena el anterior, y lo obtiene de cualquier ubicación libre; la tabla de bloques registra el orden. El kernel del mecanismo de atención consulta esa tabla para reunir las claves y los valores bloque a bloque, por lo que la caché de una secuencia no tiene que ser contigua. El desperdicio se limita a la parte sin llenar del último bloque de cada secuencia: menos del 4 % según las mediciones de vLLM.

Este ejemplo esquemático de PyTorch almacena las claves de una solicitud en bloques dispersos de un conjunto compartido y las lee mediante su tabla de bloques:

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

Imprime la tabla de bloques [41, 7, 19] y keys intact: True, y después 8 posiciones sin usar para la solicitud paginada, frente a 2.008 para una región reservada para 2.048 tokens.

Compartir bloques entre secuencias#

Las tablas de bloques también permiten que varias secuencias apunten a los mismos bloques físicos. El muestreo paralelo, la búsqueda en haz y las solicitudes que empiezan con la misma indicación del sistema reutilizan una única copia del prefijo compartido, con un contador de referencias para cada bloque. Cuando una secuencia necesita escribir en un bloque que todavía utilizan otras, el motor primero copia ese bloque, en un paso de copia al escribir tomado de los sistemas operativos.

Los autores de vLLM informan de que este uso compartido reduce hasta un 55 % la memoria empleada por el muestreo paralelo y la búsqueda en haz, lo que permite alcanzar hasta 2,2 veces más rendimiento. El hash de bloques completos amplía la idea a varias solicitudes; así es como vLLM implementa el almacenamiento en caché de indicaciones.

Técnicas relacionadas#

PagedAttention decide dónde se almacenan las claves y los valores en caché, y se combina con métodos que cambian cuánto hay que almacenar o con qué rapidez se lee:

• Atención flash: Divide el cálculo de atención en bloques para no escribir nunca la matriz completa de puntuaciones en la memoria principal de la GPU. Los motores actuales ejecutan kernels de estilo FlashAttention que leen bloques paginados.
• Procesamiento por lotes continuo: Admite una solicitud nueva cada vez que termina otra; la paginación devuelve de inmediato al conjunto libre los bloques de la secuencia finalizada, dejando espacio para la nueva.
• Atención de consultas agrupadas: Comparte los cabezales de claves y valores entre los cabezales de consulta y reduce el contenido de cada bloque.
• Cuantización de la caché FP8: Almacena claves y valores con una precisión de 8 bits, lo que reduce aproximadamente a la mitad la memoria de caché en comparación con los formatos de 16 bits.

La interfaz LLM de Ultralytics reenvía los argumentos de las solicitudes sin modificarlos; así, una llamada a Chat Completions con n=3, enviada a un servidor vLLM local mediante su base_url, devuelve tres muestras que comparten los bloques en caché de la indicación, y solo sus continuaciones divergentes ocupan memoria nueva.