¿Cuánta memoria necesita el KV cache de una petición a un LLM?
Traducción automática
Este artículo se tradujo automáticamente a partir de la versión original en inglés.
En un transformer convencional con atención completa, la memoria del KV cache aumenta con el número de tokens almacenados, las capas, las cabezas de clave/valor y los bytes por valor. Multiplica la asignación por secuencia por el número de secuencias simultáneas para estimar el volumen de datos en bruto. Después reserva memoria adicional para los pesos y el trabajo del runtime.
Esta estimación ayuda a los ingenieros de serving a dimensionar una carga de trabajo. La atención latente multicabeza (MLA), las cachés de ventana deslizante, las híbridas y las recurrentes necesitan estimaciones específicas para cada arquitectura.
Usa las cabezas KV en lugar de las cabezas de consulta
La caché almacena las proyecciones de clave y valor de los tokens anteriores para que el modelo no las vuelva a calcular en cada paso de generación. La atención densa sigue leyendo un historial de tokens cada vez mayor. Para cachés de atención completa con cabezas de clave/valor separadas, como en MHA, MQA o GQA, y secuencias de igual longitud:
cache bytes = 2 × layers × KV_heads × head_dimension
× cached_tokens × concurrent_sequences × bytes_per_elementEl factor dos cuenta las claves y los valores. Para secuencias de distinta longitud, usa la suma de sus cantidades de tokens almacenados. Lee el valor explícito de head_dim de la configuración del modelo cuando esté disponible; hidden_size / num_attention_heads es solo una alternativa para arquitecturas que usan esa relación.
La tabla de arquitectura de Llama 3 de Meta proporciona las dimensiones de estos ejemplos, que usan una caché FP16/BF16 de dos bytes por valor:
| Modelo | Capas / cabezas KV / dimensión de cabeza | Tokens por secuencia | Caché en bruto por secuencia |
|---|---|---|---|
| Llama 3.1 8B | 32 / 8 / 128 | 8,192 | 1 GiB |
| Llama 3.1 8B | 32 / 8 / 128 | 131,072 | 16 GiB |
| Llama 3.1 70B | 80 / 8 / 128 | 4,096 | 1.25 GiB |
| Llama 3.1 70B | 80 / 8 / 128 | 131,072 | 40 GiB |
Un presupuesto de caché de 40 GiB permite, según este cálculo, un máximo de 32 secuencias completas de 4,096 tokens del modelo 70B. Esto solo cuenta los datos de la caché; fija el límite de concurrencia después de reservar memoria para el runtime.
Comprueba la asignación real del servidor
Los pesos, los búferes temporales, las activaciones, los bloques parcialmente llenos y la capacidad reservada requieren más memoria. Los servidores distribuidos pueden repartir o replicar las cabezas KV. Una caché de un byte compatible reduce a la mitad este volumen de datos en bruto, pero los factores de escala, la compatibilidad de la implementación y la calidad pueden alterar la capacidad utilizable.
Mide la asignación máxima y el desalojo de secuencias con longitudes representativas de prompt y salida. PagedAttention mejora la asignación de bloques; no elimina los datos de clave/valor almacenados ni hace que la longitud de la secuencia deje de importar.
Guía de ingeniería: memoria del KV cache incluye un cálculo de capacidad ejecutable.