Pregunta:
Está implementando un LLM en producción. Generar los primeros tokens es rápido, pero a medida que la secuencia crece, cada token adicional tarda progresivamente más en generarse, aunque la arquitectura del modelo y el hardware siguen siendo los mismos.
Si la computación no es el principal cuello de botella, ¿qué ineficiencia está causando esta desaceleración y cómo rediseñaría el proceso de inferencia para hacer que la generación de tokens sea significativamente más rápida?
¿Qué es KV Caching y cómo acelera la generación de tokens?
El almacenamiento en caché KV es una técnica de optimización que se utiliza durante la generación de texto en modelos de lenguaje grandes para evitar cálculos redundantes. En la generación autorregresiva, el modelo produce texto, un token a la vez, y en cada paso normalmente vuelve a calcular la atención sobre todos los tokens anteriores. Sin embargo, las claves (K) y los valores (V) calculados para tokens anteriores nunca cambian.
Con el almacenamiento en caché KV, el modelo almacena estas claves y valores la primera vez que se calculan. Al generar el siguiente token, reutiliza los K y V almacenados en caché en lugar de recalcularlos desde cero y solo calcula la consulta (Q), la clave y el valor del nuevo token. Luego, la atención se calcula utilizando la información almacenada en caché más el nuevo token.
Esta reutilización de cálculos anteriores reduce significativamente el trabajo redundante, lo que hace que la inferencia sea más rápida y eficiente (especialmente para secuencias largas) a costa de memoria adicional para almacenar el caché. Consulta el Cuaderno de Prácticas aquí
Evaluación del impacto del almacenamiento en caché de KV en la velocidad de inferencia
En este código, comparamos el impacto del almacenamiento en caché de KV durante la generación de texto autorregresivo. Ejecutamos el mismo mensaje en el modelo varias veces, una con el almacenamiento en caché KV habilitado y otra sin él, y medimos el tiempo promedio de generación. Al mantener constantes el modelo, el mensaje y la duración de la generación, este experimento aísla cómo la reutilización de claves y valores almacenados en caché reduce significativamente el cálculo de atención redundante y acelera la inferencia. Consulta el Cuaderno de Prácticas aquí
para _ en rango(5): inicio = time.time() model.generate( **inputs, use_cache=use_cache, max_new_tokens=1000 ) times.append(time.time() – start) print( f”{‘con’ if use_cache else ‘sin’} Almacenamiento en caché de KV: ” f”{round(np.mean(times), 3)} ± {round(np.std(times), 3)} segundos” )
Los resultados demuestran claramente el impacto del almacenamiento en caché de KV en la velocidad de inferencia. Con el almacenamiento en caché de KV habilitado, generar 1000 tokens toma alrededor de 21,7 segundos, mientras que deshabilitar el almacenamiento en caché de KV aumenta el tiempo de generación a más de 107 segundos, casi una desaceleración de 5 veces. Esta marcada diferencia se produce porque, sin el almacenamiento en caché de KV, el modelo vuelve a calcular la atención sobre todos los tokens generados previamente en cada paso, lo que genera un crecimiento cuadrático en el cálculo. Consulta el Cuaderno de Prácticas aquí
Con el almacenamiento en caché KV, las claves y los valores anteriores se reutilizan, lo que elimina el trabajo redundante y mantiene el tiempo de generación casi lineal a medida que crece la secuencia. Este experimento resalta por qué el almacenamiento en caché de KV es esencial para una implementación eficiente y en el mundo real de modelos de lenguaje autorregresivos.
Consulta el Cuaderno de Prácticas aquí
Soy graduado en ingeniería civil (2022) de Jamia Millia Islamia, Nueva Delhi, y tengo un gran interés en la ciencia de datos, especialmente las redes neuronales y su aplicación en diversas áreas.