KVzap de código abierto con IA de NVIDIA: un método de poda de caché SOTA KV que ofrece una compresión 2x-4x casi sin pérdidas

A medida que la longitud del contexto pasa a decenas y cientos de miles de tokens, la caché de valores clave en los decodificadores de transformadores se convierte en el principal cuello de botella de la implementación. El caché almacena claves y valores para cada capa y cabeza con forma (2, L, H, T, D). Para un transformador básico como Llama1-65B, el caché alcanza aproximadamente 335 GB con 128 000 tokens en bfloat16, lo que limita directamente el tamaño del lote y aumenta el tiempo hasta el primer token.

https://arxiv.org/pdf/2601.07891

La compresión arquitectónica deja intacto el eje de secuencia

Los modelos de producción ya comprimen la caché a lo largo de varios ejes. La atención de consultas agrupadas comparte claves y valores en múltiples consultas y produce factores de compresión de 4 en Llama3, 12 en GLM 4.5 y hasta 16 en Qwen3-235B-A22B, a lo largo del eje principal. DeepSeek V2 comprime la dimensión de clave y valor a través de atención latente de cabezales múltiples. Los modelos híbridos combinan atención con atención de ventana deslizante o capas de espacio de estado para reducir la cantidad de capas que mantienen un caché completo.

Estos cambios no se comprimen a lo largo del eje de secuencia. La atención al estilo disperso y de recuperación recupera solo un subconjunto del caché en cada paso de decodificación, pero todos los tokens aún ocupan memoria. Por lo tanto, el servicio práctico de contexto largo necesita técnicas que eliminen las entradas de la caché, lo que tendrá un efecto insignificante en los tokens futuros.

El proyecto KVpress de NVIDIA recopila más de veinte métodos de poda en una base de código y los expone a través de una tabla de clasificación pública en Hugging Face. Métodos como H2O, Expected Attention, DuoAttention, Compactor y KVzip se evalúan de forma coherente.

KVzip y KVzip plus como oráculo de puntuación

KVzip es actualmente la base de eliminación de caché más sólida en la tabla de clasificación de KVpress. Define una puntuación de importancia para cada entrada de caché mediante una tarea de pretexto de copiar y pegar. El modelo se ejecuta en un mensaje extendido donde se le pide que repita exactamente el contexto original. Para cada posición del token en el mensaje original, la puntuación es el peso de atención máximo que cualquier posición en el segmento repetido asigna a ese token, entre cabezas del mismo grupo cuando se utiliza atención de consulta agrupada. Las entradas con puntuaciones bajas se desalojan hasta que se alcanza un presupuesto global.

KVzip+ refina esta puntuación. Multiplica el peso de la atención por la norma de la contribución de valor al flujo residual y lo normaliza según la norma del estado oculto de recepción. Esto coincide mejor con el cambio real que induce un token en el flujo residual y mejora la correlación con la precisión posterior en comparación con la puntuación original.

Estas puntuaciones de Oracle son efectivas pero costosas. KVzip requiere un llenado previo en el mensaje extendido, lo que duplica la longitud del contexto y lo hace demasiado lento para la producción. Tampoco puede ejecutarse durante la decodificación porque el procedimiento de puntuación supone un mensaje fijo.

https://arxiv.org/pdf/2601.07891

KVzap, un modelo sustituto de estados ocultos

KVzap reemplaza la puntuación del oráculo con un pequeño modelo sustituto que opera directamente en estados ocultos. Para cada capa de transformador y cada posición de secuencia t, el módulo recibe el vector oculto hₜ y genera puntuaciones de registro predichas para cada cabeza de valor clave. Se consideran dos arquitecturas, una sola capa lineal (KVzap Linear) y un MLP de dos capas con GELU y un ancho oculto igual a un octavo del tamaño oculto del modelo (KVzap MLP).

El entrenamiento utiliza indicaciones del ejemplo del conjunto de datos de preentrenamiento de Nemotron. El equipo de investigación filtró 27.000 mensajes a longitudes de entre 750 y 1250 tokens, tomó muestras de hasta 500 mensajes por subconjunto y luego tomó muestras de 500 posiciones de tokens por mensaje. Para cada cabeza de valor clave, obtienen alrededor de 1,2 millones de pares de entrenamiento y un conjunto de validación de 23.000 pares. El sustituto aprende a regresar del estado oculto a la puntuación log KVzip+. En todos los modelos, la correlación de Pearson al cuadrado entre las predicciones y las puntuaciones de Oracle alcanza entre aproximadamente 0,63 y 0,77, y la variante MLP supera consistentemente a la variante lineal.

https://arxiv.org/pdf/2601.07891

Umbral, ventana corredera y gastos generales insignificantes

Durante la inferencia, el modelo KVzap procesa estados ocultos y produce puntuaciones para cada entrada de caché. Las entradas con puntuaciones por debajo de un umbral fijo se eliminan, mientras que siempre se mantiene una ventana deslizante de los 128 tokens más recientes. El equipo de investigación proporciona una función concisa de estilo PyTorch que aplica el modelo, establece puntuaciones de la ventana local en infinito y devuelve tensores de valores y claves comprimidos. En todos los experimentos la poda se aplica después de la operación de atención.

KVzap utiliza umbrales de puntuación en lugar de una selección k superior fija. Un único umbral produce diferentes relaciones de compresión efectivas en diferentes puntos de referencia e incluso entre indicaciones dentro del mismo punto de referencia. El equipo de investigación informa hasta un 20 por ciento de variación en la relación de compresión entre indicaciones en un umbral fijo, lo que refleja diferencias en la densidad de la información.

La sobrecarga informática es pequeña. Un análisis a nivel de capa muestra que el costo adicional de KVzap MLP es como máximo alrededor del 1,1 por ciento de los FLOP de proyección lineal, mientras que la variante lineal agrega alrededor del 0,02 por ciento. La sobrecarga relativa de memoria sigue los mismos valores. En regímenes de contexto prolongado, domina el costo cuadrático de la atención, por lo que los FLOP adicionales son efectivamente insignificantes.

https://arxiv.org/pdf/2601.07891

Resultados en RULER, LongBench y AIME25

KVzap se evalúa en contextos largos y puntos de referencia de razonamiento utilizando Qwen3-8B, Llama-3.1-8B Instruct y Qwen3-32B. El comportamiento del contexto largo se mide en RULER y LongBench. RULER utiliza tareas sintéticas en longitudes de secuencia de 4k a 128k tokens, mientras que LongBench utiliza documentos del mundo real de múltiples categorías de tareas. AIME25 proporciona una carga de trabajo de razonamiento matemático con 30 problemas de nivel Olimpiada evaluados con aprobación 1 y aprobación 4.

En RULER, KVzap coincide con la línea base de caché completa dentro de un pequeño margen de precisión y elimina una gran fracción del caché. Para Qwen3-8B, la mejor configuración de KVzap logra una fracción eliminada superior a 0,7 en RULER 4k y 16k mientras mantiene la puntuación promedio dentro de unas pocas décimas de punto del caché completo. Se produce un comportamiento similar para Llama-3.1-8B Instruct y Qwen3-32B.

En LongBench, los mismos umbrales conducen a índices de compresión más bajos porque los documentos son menos repetitivos. KVzap permanece cerca de la línea base de caché completa hasta aproximadamente 2 o 3 veces la compresión, mientras que los métodos de presupuesto fijo como la atención esperada se degradan más en varios subconjuntos una vez que aumenta la compresión.

En AIME25, KVzap MLP mantiene o mejora ligeramente la precisión de paso a 4 con una compresión cercana a 2 veces y sigue siendo utilizable incluso cuando se descarta más de la mitad del caché. Las configuraciones extremadamente agresivas, por ejemplo variantes lineales con umbrales altos que eliminan más del 90 por ciento de las entradas, colapsan el rendimiento como se esperaba.

https://arxiv.org/pdf/2601.07891

En general, la tabla anterior muestra que la mejor configuración de KVzap por modelo ofrece una compresión de caché promedio entre aproximadamente 2,7 y 3,5 y, al mismo tiempo, mantiene las puntuaciones de las tareas muy cercanas a la línea base de caché completa en RULER, LongBench y AIME25.

Conclusiones clave

KVzap es una aproximación adaptativa de entrada de KVzip+ que aprende a predecir puntuaciones de importancia de KV de Oracle a partir de estados ocultos utilizando modelos sustitutos pequeños por capa, ya sea una capa lineal o un MLP poco profundo, y luego elimina los pares de KV de puntuación baja. El entrenamiento utiliza indicaciones de preentrenamiento de Nemotron donde KVzip+ proporciona supervisión, produciendo alrededor de 1,2 millones de ejemplos por cabeza y logrando una correlación al cuadrado en el rango de 0,6 a 0,8 entre las puntuaciones previstas y de Oracle, lo cual es suficiente para una clasificación fiel de la importancia del caché. KVzap aplica un umbral de puntuación global con una ventana deslizante fija de tokens recientes, por lo que la compresión se adapta automáticamente a la densidad de información de los mensajes, y el equipo de investigación informa hasta un 20 por ciento de variación en la compresión lograda entre mensajes en el mismo umbral. En Qwen3-8B, Llama-3.1-8B Instruct y Qwen3-32B en RULER, LongBench y AIME25, KVzap alcanza entre 2 y 4 veces la compresión de caché KV mientras mantiene una precisión muy cercana a la caché completa, y logra compensaciones de última generación en la tabla de clasificación NVIDIA KVpress. El cálculo adicional es pequeño, como máximo alrededor de un 1,1 por ciento de FLOP adicionales para la variante MLP, y KVzap se implementa en el marco de código abierto kvpress con puntos de control listos para usar en Hugging Face, lo que hace que sea práctico integrarlo en pilas de servicios LLM de contexto largo existentes.

Consulte el repositorio de Paper y GitHub. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 100.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.

Asif Razzaq es el director ejecutivo de Marktechpost Media Inc.. Como empresario e ingeniero visionario, Asif está comprometido a aprovechar el potencial de la inteligencia artificial para el bien social. Su esfuerzo más reciente es el lanzamiento de una plataforma de medios de inteligencia artificial, Marktechpost, que se destaca por su cobertura en profundidad del aprendizaje automático y las noticias sobre aprendizaje profundo que es técnicamente sólida y fácilmente comprensible para una amplia audiencia. La plataforma cuenta con más de 2 millones de visitas mensuales, lo que ilustra su popularidad entre el público.