¿Qué es Flash-Decoding y cuándo acelera los LLMs?

Traducción automática

Este artículo se tradujo automáticamente a partir de la versión original en inglés.

Flash-Decoding paraleliza la atención durante la decodificación a lo largo de la secuencia de claves y valores almacenada en caché. Distintos bloques de la GPU procesan distintas partes de un historial largo y después combinan los resultados parciales de atención con una normalización numéricamente correcta. Puede ayudar a generar con lotes pequeños cuando una caché larga ofrece más trabajo del que la planificación actual de la atención ejecuta de forma eficiente.

Acelera una operación de atención. La mejora resultante en el serving también depende del tiempo empleado en otras partes del modelo.

Por qué la decodificación necesita otra planificación

Durante el prefill, muchas posiciones de consulta pueden procesarse en paralelo. Durante la decodificación habitual, cada secuencia activa aporta una nueva posición de consulta. Por tanto, un lote pequeño puede ofrecer demasiado pocos bloques independientes para aprovechar la GPU, aunque cada consulta deba leer un historial largo.

Flash-Decoding crea unidades de trabajo adicionales al dividir ese historial. Cada unidad calcula un resultado parcial y un valor de normalización log-sum-exp. Una reducción los combina para obtener el resultado que requiere la atención sobre el historial completo. Consulta la explicación de Flash-Decoding de los autores de Stanford.

Con un historial corto, el trabajo adicional de los resultados parciales y la reducción puede superar el beneficio. Un lote mayor puede aportar ya suficiente trabajo paralelo. Ninguno de los dos casos garantiza una mejora al seguir dividiendo la secuencia.

Cómo interpretar el resultado publicado

Los autores comunicaron una aceleración de hasta ocho veces de extremo a extremo en su evaluación de CodeLlama-34B con un lote de tamaño uno en cuatro GPUs A100, con longitudes de secuencia de 512 a 64K. Esa comparación utilizó la referencia y la configuración de ejecución que eligieron. No predice una mejora de ocho veces respecto a un servidor actual con un backend de decodificación ya optimizado.

Al probar tu servidor:

  • comprueba que el backend de atención seleccionado utiliza realmente el método;
  • compara el mismo modelo, disposición de GPUs, precisión de caché y longitudes de secuencia;
  • prueba tanto lotes de tamaño uno como una concurrencia realista;
  • registra el tiempo del kernel de atención por separado del paso completo de decodificación;
  • mide la latencia de los tokens de salida y la corrección numérica.

Un kernel de atención más rápido puede dejar la lectura de pesos, el cálculo de la red feed-forward o la comunicación como el mayor coste restante. Utiliza el resultado de extremo a extremo para decidir si cambias la configuración de serving.

Guía de ingeniería: Flash-Decoding relaciona la división de la secuencia con la distinción entre prefill y decodificación.