PagedAttention

PagedAttention est un algorithme de gestion de la mémoire hautement efficace conçu pour optimiser la vitesse d'inférence et le débit des grands modèles linguistiques (LLM). S'inspirant des concepts de mémoire virtuelle et de pagination des systèmes d'exploitation traditionnels, cette technique résout le problème de la consommation massive de mémoire associée au cache clé-valeur (souvent appelé cache KV) lors de la génération de texte. En divisant les blocs de mémoire continus requis pour le cache en « pages » plus petites et non contiguës, PagedAttention élimine efficacement la fragmentation de la mémoire tant interne qu'externe. Cela permet aux serveurs d'IA de traiter simultanément un nombre nettement plus important de requêtes, optimisant ainsi GPU .

Attention par pages vs attention instantanée

Si ces deux techniques optimisent les performances des réseaux neuronaux, elles s'attaquent à des goulots d'étranglement différents. Flash Attention est une optimisation au niveau du calcul qui accélère le mécanisme d'attention lui-même en minimisant les lectures et écritures lentes en mémoire à travers la GPU . En revanche, PagedAttention est une stratégie d' allocation de mémoire. Elle se concentre uniquement sur la manière dont la mémoire de la fenêtre de contexte est structurée et stockée, permettant ainsi une mise à l'échelle dynamique sans pré-allouer de grands blocs de mémoire inutiles.

Applications concrètes

L'efficacité mémoire offerte par PagedAttention a révolutionné la manière dont les modèles génératifs à grande échelle sont déployés en production.

1. Prestation d'API à haut débit: Les systèmes de production servant des modèles similaires à GPT-4 utilisent PagedAttention via des frameworks tels que vLLM. En partageant des blocs de mémoire entre différentes requêtes utilisateur, les fournisseurs peuvent desservir jusqu'à quatre fois plus d'utilisateurs sur le même matériel, ce qui réduit considérablement le coût d'exploitation des services d'IA basés sur le cloud.
2. Stratégies de décodage complexes: Lorsqu'un modèle d'IA génère plusieurs réponses potentielles simultanément (comme dans le cas de la recherche par faisceaux ou de l'échantillonnage parallèle), PagedAttention permet à ces séquences parallèles de partager en toute sécurité les mêmes pages de mémoire de base. Cela empêche le système de dupliquer de la mémoire redondante, ce qui accélère considérablement les tâches de raisonnement complexes.

Efficacité de la mémoire en vision par ordinateur

Si PagedAttention est principalement utilisé dans le traitement du langage naturel, le principe sous-jacent d’une optimisation rigoureuse de la mémoire est tout aussi crucial en vision par ordinateur (CV). Lors du déploiement de modèles sur des périphériques en périphérie aux ressources matérielles limitées, il est essentiel d’éviter toute surcharge de mémoire. Ultralytics atteint une efficacité d'inférence en temps réel de manière native, en contournant la nécessité d'une gestion lourde du cache grâce à une architecture de bout en bout NMS.

Pour les développeurs qui souhaitent gérer de manière transparente les besoins en mémoire et les exigences d'exportation des chaînes de traitement de détection d'objets, la Ultralytics propose des outils de déploiement automatisés qui packagent les modèles pour une exécution matérielle optimale.

Exemple de code

PagedAttention fonctionne en arrière-plan dans les frameworks de serveurs, en remplaçant les fonctions d'attention standard par Cuda optimisés. Vous trouverez ci-dessous un exemple conceptuel illustrant comment on pourrait définir l'attention standard dans PyTorch, que des systèmes tels que vLLM interceptent automatiquement et optimisent à l'aide de la pagination lors du déploiement du modèle.

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}")

En tirant parti de stratégies avancées d'allocation de mémoire, le secteur de l'IA ne cesse de repousser les limites du possible, garantissant ainsi que les modèles fondamentaux à grande échelle puissent être déployés et consultés efficacement partout dans le monde.