Alors que les grands modèles de langage (LLM) deviennent de plus en plus essentiels aux applications de production, le goulot d'étranglement de la latence d'inférence s'est déplacé de l'entraînement du modèle vers l'efficacité du déploiement. Bien que les GPU modernes comme le H100 offrent un débit massif, la bande passante mémoire reste une contrainte critique. Chaque jeton généré nécessite la lecture de la fenêtre de contexte entière, ce qui entraîne des pics de latence significatifs à mesure que la longueur des conversations augmente. Cet article explore un schéma architectural robuste pour atténuer ce problème : le déchargement du cache Key-Value (KV) vers un magasin de mémoire distribué tel que Redis.
Le problème : une inférence limitée par la mémoire
Dans les architectures Transformer standard, le mécanisme d'attention calcule les relations entre tous les jetons du contexte. Pour optimiser les calculs répétés, les systèmes utilisent le cache KV pour stocker les clés et les valeurs des jetons précédents. Cependant, ce cache réside dans la VRAM du GPU. À mesure que les fenêtres de contexte s'agrandissent pour atteindre 128 000 jetons ou plus, la consommation de VRAM explose, limitant les tailles de lot et augmentant les temps d'attente des requêtes. Déplacer ces données statiques hors du GPU libère de la VRAM précieuse pour le calcul et permet une concurrence plus élevée.
Pourquoi Redis ?
Redis n'est pas seulement un simple magasin clé-valeur ; sa capacité à gérer des données binaires, à fournir une latence inférieure à la milliseconde et à prendre en charge des structures de données avancées en fait un outil idéal pour le déchargement du cache KV. Contrairement aux bases de données basées sur le disque, Redis conserve les données en mémoire, ce qui garantit que le goulot d'étranglement de la « paroi mémoire » est déplacé de la VRAM GPU coûteuse vers la RAM système ou le stockage réseau à moindre coût, sans sacrifier la vitesse.
Mise en œuvre architecturale
L'intégration implique un pipeline dans lequel le moteur LLM (tel que vLLM ou Triton) vérifie l'existence d'entrées KV dans Redis avant la génération. Si le cache existe, il est récupéré et chargé dans le contexte du GPU. Sinon, les nouveaux jetons sont calculés, et leurs paires KV sont ensuite écrites dans Redis.
Configuration du client Redis
Premièrement, assurez-vous d'avoir le client redis-py installé et une instance Redis en cours d'exécution. Nous utiliserons la structure de données HASH pour stocker efficacement les matrices KV, car elles constituent des blocs de mémoire contigus.
import redis
import numpy as np
import json
class KVCacheManager:
def __init__(self, host='localhost', port=6379, db=0):
self.r = redis.Redis(host=host, port=port, db=db, decode_responses=False)
def save_kv_cache(self, session_id: str, key_tensor: np.ndarray, value_tensor: np.ndarray):
"""
Enregistre les tenseurs du cache KV dans Redis.
"""
# Sérialisation des tableaux numpy en octets
key_bytes = key_tensor.tobytes()
value_bytes = value_tensor.tobytes()
# Utilisation d'un pipeline pour l'atomicité
pipe = self.r.pipeline()
pipe.hset(session_id, "keys", key_bytes)
pipe.hset(session_id, "values", value_bytes)
pipe.expire(session_id, 3600) # Définir la durée de vie à 1 heure
pipe.execute()
def load_kv_cache(self, session_id: str):
"""
Récupère les tenseurs du cache KV depuis Redis.
"""
key_bytes = self.r.hget(session_id, "keys")
value_bytes = self.r.hget(session_id, "values")
if not key_bytes:
return None, None
# Déterminer les formes en fonction de la configuration de votre modèle
# Exemple : en supposant batch_size=1, num_heads=32, seq_len=current_context
# Vous devez gérer la logique de reconstruction de forme spécifique à votre LLM
return key_bytes, value_bytesIntégration avec le moteur d'inférence
Dans votre boucle d'inférence, vous mettriez en œuvre une vérification avant la génération de jetons :
def generate_token(model, session_id, prompt):
# 1. Vérifier Redis pour le cache existant
cached_keys, cached_values = load_kv_cache(session_id)
if cached_keys:
# 2. Charger les paires KV mises en cache dans la mémoire GPU
load_to_gpu(cached_keys, cached_values)
else:
# 3. Exécuter la passe avant initiale et calculer le cache KV
run_initial_forward(model, prompt)
# 4. Générer le jeton suivant
next_token = model.generate()
# 5. Mettre à jour Redis avec les nouvelles entrées KV
new_keys, new_values = model.get_new_kv_cache()
save_kv_cache(session_id, new_keys, new_values)
return next_tokenConsidérations pratiques et compromis
Bien que le déchargement vers Redis introduise une surcharge réseau, le coût en latence de la sérialisation et de la transmission des données du cache KV est souvent nettement inférieur au coût de bande passante mémoire de leur relecture depuis la VRAM du GPU lors des étapes d'attention suivantes. Cependant, vous devez surveiller votre débit réseau. Pour des exigences de latence ultra-faible, envisagez de colocaliser Redis avec vos nœuds d'inférence ou d'utiliser RDMA (Remote Direct Memory Access) pour contourner le CPU.
De plus, mettez en œuvre une politique d'éviction robuste. Toutes les conversations ne doivent pas être mises en cache indéfiniment. Utilisez des stratégies de type « moins récemment utilisé » (LRU) ou des paramètres de durée de vie (TTL) pour gérer la pression mémoire dans votre cluster Redis.
Conclusion
L'optimisation de l'inférence LLM ne consiste plus seulement à utiliser des modèles plus grands ; il s'agit d'un mouvement efficace des données. En tirant parti de Redis pour le déchargement du cache KV, les développeurs peuvent découpler le stockage mémoire du calcul, permettant ainsi des déploiements évolutifs, rentables et à faible latence. Cette approche vous permet de servir des fenêtres de contexte plus longues avec moins de GPU, impactant directement votre rentabilité et l'expérience utilisateur. Commencez petit, mesurez les améliorations de latence et itérez sur votre stratégie de mise en cache pour trouver l'équilibre parfait pour votre charge de travail spécifique.