Atención por páginas

PagedAttention es un algoritmo de gestión de memoria altamente eficiente diseñado para optimizar la velocidad de inferencia y el rendimiento de los modelos de lenguaje a gran escala (LLM). Inspirada en los conceptos de memoria virtual y paginación de los sistemas operativos tradicionales, esta técnica aborda el enorme consumo de memoria asociado a la caché de clave-valor (a menudo denominada caché KV) durante la generación de texto. Al dividir los bloques de memoria continuos necesarios para la caché en «páginas» más pequeñas y no contiguas, PagedAttention elimina de forma eficaz la fragmentación de la memoria tanto interna como externa. Esto permite a los servidores de IA procesar simultáneamente un número significativamente mayor de solicitudes, maximizando así GPU .

Atención por páginas frente a atención instantánea

Aunque ambas técnicas optimizan el rendimiento de las redes neuronales, se centran en distintos cuellos de botella. Flash Attention es una optimización a nivel de computación que acelera el propio mecanismo de atención minimizando las lecturas y escrituras lentas en la memoria a lo largo de GPU . Por el contrario, PagedAttention es una estrategia de asignación de memoria. Se centra exclusivamente en cómo se estructura y almacena la memoria para la ventana de contexto, lo que permite un escalado dinámico sin necesidad de preasignar bloques de memoria grandes y poco eficientes.

Aplicaciones en el mundo real

La eficiencia en el uso de la memoria que ofrece PagedAttention ha transformado la forma en que se implementan los modelos generativos a gran escala en producción.

1. Servicio de API de alto rendimiento: Los sistemas de producción que prestan servicios a modelos similares a GPT-4 utilizan PagedAttention a través de marcos como vLLM. Al compartir bloques de memoria entre las diferentes solicitudes de los usuarios, los proveedores pueden atender hasta cuatro veces más usuarios con el mismo hardware, lo que reduce drásticamente el coste de ejecución de los servicios de IA en la nube.
2. Estrategias de decodificación complejas: Cuando un modelo de IA genera varias respuestas potenciales a la vez (como en la búsqueda por haces o el muestreo paralelo), PagedAttention permite que estas secuencias paralelas compartan de forma segura las mismas páginas de memoria fundamentales. Esto evita que el sistema duplique memoria redundante, lo que agiliza considerablemente las tareas de razonamiento complejas.

Eficiencia de la memoria en la visión artificial

Aunque PagedAttention se utiliza principalmente en el procesamiento del lenguaje natural, el principio subyacente de la optimización rigurosa de la memoria es igualmente crucial en la visión artificial (CV). Al implementar modelos en dispositivos periféricos con limitaciones de hardware, es esencial evitar el exceso de memoria. Ultralytics logra una eficiencia de inferencia en tiempo real de forma nativa, evitando la necesidad de una gestión pesada de la caché mediante el uso de una arquitectura de extremo a extremo NMS.

Para los desarrolladores que deseen gestionar de forma fluida los requisitos de memoria y exportación de los flujos de trabajo de detección de objetos, la Ultralytics ofrece herramientas de implementación automatizadas que empaquetan los modelos para una ejecución óptima en hardware.

Ejemplo de código

PagedAttention funciona en segundo plano en los marcos de trabajo de servidores, sustituyendo las funciones de atención estándar por Cuda optimizados. A continuación se muestra un ejemplo conceptual que ilustra cómo se podría definir la atención estándar en PyTorch, que sistemas como vLLM interceptan automáticamente y optimizan mediante paginación durante la implementación del modelo.

import torch
import torch.nn.functional as F

# Simulated Key, Query, and Value tensors for a standard attention block
batch_size, num_heads, sequence_length, head_dim = 1, 8, 1024, 64
query = torch.randn(batch_size, num_heads, sequence_length, head_dim)
key = torch.randn(batch_size, num_heads, sequence_length, head_dim)
value = torch.randn(batch_size, num_heads, sequence_length, head_dim)

# Standard attention computation (often replaced by PagedAttention kernels in production LLM servers)
attention_output = F.scaled_dot_product_attention(query, key, value)

print(f"Computed attention shape: {attention_output.shape}")

Gracias al uso de estrategias avanzadas de asignación de memoria, el sector de la inteligencia artificial sigue ampliando los límites de lo posible, garantizando que los modelos fundamentales a gran escala puedan ampliarse y consultarse de manera eficiente en todo el mundo.