¿Cuánta memoria reservar para la caché KV?

Describir lo que queda en memoria

Durante una generación autorregresiva, las claves y valores de los tokens ya procesados pueden conservarse para los pasos siguientes. Esta caché pertenece a las capas de atención. Su volumen depende, por tanto, del modelo y del historial conservado, no solo del número de parámetros. El contexto de una conversación también incluye las instrucciones, los mensajes previos y los documentos añadidos por la aplicación.

Tu primera decisión es operativa: cuántas secuencias deberán permanecer activas, hasta qué longitud. Una cola de veinte solicitudes de las que dos se ejecutan simultáneamente no representa necesariamente veinte cachés residentes. Registra la admisión real de las peticiones y las posibles secuencias adicionales creadas por la generación.

Prepara una ficha con la revisión del modelo, las capas de atención, las cabezas KV, la dimensión de las claves y valores, su dtype y la estrategia de caché. Cuenta los tokens después del tokenizer y de la plantilla de conversación. Un límite en caracteres no describe esta asignación.

Fuentes técnicas: Hugging Face — funcionamiento y forma de las cachés por capa

Emplear las cabezas KV, en particular con GQA

El número de cabezas de consulta Q y el de cabezas KV pueden diferir. En la atención multi-cabeza clásica, coinciden. Con MQA, se comparte una única cabeza KV; GQA agrupa varias cabezas Q alrededor de una cabeza KV. Lee num_key_value_heads en la configuración cuando ese campo exista y verifica su significado para la arquitectura.

Por ejemplo, cuarenta cabezas Q y ocho cabezas KV forman cinco cabezas Q por grupo KV. La fórmula del caché almacenado usa ocho, no cuarenta. Este ratio no describe toda la memoria de la atención: algunas operaciones o conversiones pueden producir temporales.

No modifiques simplemente este número para reducir el requisito de memoria de un modelo ya entrenado. El esquema de atención forma parte de su arquitectura. Dos modelos con distinto número de cabezas no se convierten en variantes equivalentes por un cálculo de capacidad; su calidad debe evaluarse por separado.

Fuentes técnicas: Hugging Face — campos de LlamaConfig y distinción MHA, MQA, GQA · PyTorch — dimensiones Q/K/V y restricciones de la atención GQA

Plantear la fórmula y sus unidades

En el caso uniforme, L es el número de capas, Hkv el número de cabezas KV, D su dimensión, T el número de posiciones conservadas por secuencia, B el número de secuencias y q el número de bytes por valor. El factor dos cuenta K y V. Esta aproximación supone claves y valores de la misma dimensión y el mismo formato, sin compresión ni compartición de prefijo.

Convierte solo el resultado final: un GiB equivale a 1 073 741 824 bytes; un GB decimal equivale a 1 000 000 000 bytes. Mantén los bytes en tu ficha para evitar que un redondeo o un cambio de unidad oculte una diferencia.

Para longitudes diferentes sin padding, sustituye B × T por la suma de las posiciones realmente almacenadas. Para capas heterogéneas, haz una suma capa por capa. Un almacenamiento denso con padding, una asignación por bloques o una reserva estática exigen contar las posiciones asignadas, que pueden superar los tokens útiles.

Fuentes técnicas: Hugging Face — dimensiones de los tensores del caché · NIST — unidades decimales y prefijos binarios

Ejemplo trabajado: seis secuencias, sin medición de GPU

Tomemos una arquitectura ficticia de cuarenta capas, ocho cabezas KV y una dimensión de 128. Supongamos un caché uniforme de dos bytes por valor. Cada secuencia recibe como máximo 3 072 tokens de entrada y una reserva para 1 024 tokens adicionales, es decir, un límite de 4 096 posiciones. Son hipótesis pedagógicas, no la configuración verificada de un modelo.

El coste calculado por posición y por secuencia es 2 × 40 × 8 × 128 × 2 = 163 840 bytes. Una secuencia de 4 096 posiciones representa entonces 671 088 640 bytes, es decir, 0,625 GiB. Seis secuencias dan 4 026 531 840 bytes, es decir, 3,75 GiB solo para el caché.

Duplicar la longitud conservada duplica esta partida en esta fórmula. Sustituir ocho cabezas KV por cuarenta la multiplica por cinco, con todas las demás hipótesis sin cambios. Estas proporciones no anuncian ninguna aceleración, pérdida de calidad ni compatibilidad de un modelo real.

Adaptar el presupuesto a la estrategia de caché

Un caché dinámico crece con las posiciones conservadas. Un caché estático reserva una capacidad máxima: dimensiona esa reserva, no solo la petición corta observada al arrancar. Para una atención de ventana deslizante, algunas capas pueden limitar su historial; las capas con atención completa requieren un cálculo distinto.

La cuantización del caché y su desvío a la CPU son otras estrategias, dependientes del modelo y del software. Modifican las restricciones de almacenamiento, de transferencia o de cálculo. Cuantizar los pesos no demuestra que el caché tenga el mismo formato.

Anota la clase de caché y sus parámetros explícitos. Verifica también la liberación de las posiciones tras el fin o la cancelación de una petición. Para una carga que mezcla secuencias cortas y largas, una reserva máxima uniforme puede pesar más que la sola suma de los contenidos útiles.

Fuentes técnicas: Hugging Face — cachés dinámicos, estáticos, cuantizados y desviados

Verificar una carga representativa por etapas

Construye tres casos: entrada habitual, entrada larga esperada y número máximo de peticiones simultáneas permitido. Fija modelo, tokenizer, plantilla, límite de generación y regla de fin. Varía una dimensión a la vez; una respuesta acortada o un documento truncado cambian el trabajo realizado.

Separa la carga, el procesamiento inicial de la entrada y la generación. Sincroniza el device en torno a las fases cronometradas y conserva los niveles de memoria iniciales y los picos. El método de memoria explica por qué allocated y reserved no se suman y por qué la diferencia de sus máximos no aísla el caché.

Tu control debe dar como resultado una cota probada y unas salidas aceptables: IDs de las solicitudes completadas, errores, longitud producida y criterio de calidad. Un lanzamiento que se sostiene con una entrada corta no valida la concurrencia máxima. Un error de memoria antes del final no constituye una medición de lo que necesita una ejecución completa.

Fuentes técnicas: PyTorch — sincronización del trabajo en el device elegido · PyTorch — alcance de los contadores de memoria

Descartar los errores que distorsionan la elección

No conviertas el caché calculado en capacidad GPU total. Añade un análisis de los pesos, las activaciones temporales, las salidas conservadas y el software. Tampoco dividas ese volumen automáticamente entre el número de tarjetas: la ubicación de las capas o las cabezas debe configurarse y verificarse realmente en cada device.

El notebook IteraGPU Lab es un ejercicio de instrumentación sobre una red densa pequeña sin atención. Ayuda a leer fases de memoria; su parámetro context no valida este cálculo KV. Usa la ficha de carga y el expediente de inferencia para preparar la prueba de tu propio modelo.

Compara las capacidades de las ofertas solo después de distinguir el volumen aritmético, el máximo realmente observado y los casos aún no probados. El plan de alquiler organiza tu ventana de trabajo; no garantiza ninguna longitud de contexto ni cadencia de generación.

• Confundir cabezas Q y cabezas KV: retoma la configuración de la arquitectura.
• Presupuestar solo el prompt: integra la generación y la reserva efectiva.
• Leer «pesos 4 bits» como «caché 4 bits»: registra ambos formatos.
• Olvidar la concurrencia: cuenta las secuencias realmente residentes.
• Prometer una memoria común entre tarjetas: verifica la ubicación real.

Preguntas prácticas