Technical Tutorials

L'inférence distribuée des grands modèles de langage (LLM) est passée d'une tâche de traitement par lots statique à une opération dynamique et avec état. Contrairement aux serveurs web traditionnels qui traitent les requêtes comme des unités indépendantes, les LLM modernes s'appuient fortement sur le cache Key-Value (KV) pour maintenir le contexte sur plusieurs étapes d'inférence. Lors du déploiement dans un cluster distribué, le placement de ces requêtes devient critique. Si une requête ultérieure pour la conversation d'un utilisateur est acheminée vers un nœud GPU différent, le modèle doit soit recalculer l'intégralité du cache KV depuis le début, soit subir des pénalités de latence sévères dues aux surcoûts de transfert de données. C'est ici que l'équilibrage de charge avec état devient un composant essentiel de l'infrastructure IA.

Le problème du routage sans état

Les équilibreurs de charge standard, tels que NGINX ou HAProxy, fonctionnent généralement de manière sans état, en utilisant des algorithmes comme Round Robin ou Least Connections pour distribuer le trafic. Dans le contexte des LLM, cette approche est sous-optimale. Le cache KV agit comme une empreinte mémoire qui augmente avec la longueur de la séquence. Déplacer cet état entre les GPU est coûteux. Si nous perdons l'"affinité", nous perdons en performance.

Considérez une application de chat multi-tours. Le premier prompt (« Raconte-moi une histoire sur un chat ») est traité sur le nœud GPU A. Le cache KV pour ce préfixe est stocké dans la mémoire HBM du nœud A. L'utilisateur répond (« Maintenant, rends-la drôle »). Si un équilibreur de charge sans état achemine cette deuxième requête vers le nœud GPU B, le nœud B n'a aucune connaissance du contexte précédent. Il doit soit :

Récupérer le cache KV depuis le nœud A via le réseau (PCIe/NVLink ou InfiniBand), ce qui introduit des contraintes de bande passante significatives.
Recalculer la phase de préremplissage (prefill) pour l'intégralité du contexte précédent, dupliquant ainsi le travail de calcul.

Mise en œuvre de la persistance de session avec affinité KV

Pour résoudre ce problème, nous mettons en œuvre l'Affinité du Cache KV. Cette stratégie consiste à maintenir une correspondance entre les sessions utilisateur (ou les préfixes de requêtes uniques) et des workers GPU spécifiques. L'équilibreur de charge consulte cette carte avant d'acheminer une nouvelle requête. Si une session est déjà assignée à un worker, la requête est dirigée vers celui-ci. Ce n'est que lorsqu'un worker est saturé ou indisponible que l'équilibreur migre la session vers un nouveau nœud, acceptant le coût de la migration d'état ou du recalcul comme compromis nécessaire.

Considérations architecturales

Un équilibreur avec état efficace nécessite un magasin centralisé ou distribué (comme Redis ou etcd) pour maintenir la carte d'affinité. Les décisions de conception clés impliquent :

Définition de la clé : L'utilisation de l'ID de session est simple mais rigide. L'utilisation d'un hachage du préfixe du prompt offre une meilleure flexibilité pour la mise en cache sémantique, mais augmente les surcoûts de calcul à la porte d'entrée.
Contrôles de santé : L'équilibreur doit surveiller l'utilisation de la mémoire GPU et la température. Si un worker détenant une session « chaude » approche des limites de mémoire, le système doit évincer proactivement la session ou la migrer gracieusement avant qu'une erreur OOM (Out of Memory) ne se produise.
Rejet de charge (Load Shedding) : Lorsque tous les nœuds avec affinité sont surchargés, le système doit prioriser les nouvelles sessions vers des nœuds moins utilisés plutôt que de mettre en file d'attente indéfiniment derrière des charges existantes à haute affinité.

Exemple de code : Logique de routage basée sur l'affinité

Ci-dessous se trouve une implémentation conceptuelle en Python d'un équilibreur de charge avec état qui priorise l'affinité tout en respectant les limites de charge.


import random
from collections import defaultdict

class StatefulLLMLoadBalancer:
    def __init__(self, max_workers_per_session=1):
        # Maps session_id -> worker_id
        self.affinity_map = {}
        # Maps worker_id -> list of active session_ids
        self.worker_sessions = defaultdict(list)
        # Simulated worker load (e.g., number of active requests or memory usage)
        self.worker_load = defaultdict(int)

    def assign_worker(self, session_id: str, current_loads: dict) -> str:
        """
        Assigns a worker for a given session ID.
        Prioritizes existing affinity if the worker is under the load threshold.
        """
        threshold = 10  # Maximum concurrent sessions per worker before migration

        # 1. Check if we have an existing affinity
        if session_id in self.affinity_map:
            current_worker = self.affinity_map[session_id]
            
            # If the current worker is not overloaded, stick to it
            if current_loads.get(current_worker, 0) < threshold:
                return current_worker
            
            # If overloaded, we must migrate. Remove from old worker context.
            print(f"Warning: Worker {current_worker} overloaded. Migrating session {session_id}.")
            self.worker_sessions[current_worker].remove(session_id)

        # 2. Select a new worker
        # Strategy: Least Loaded Worker among those not at capacity
        available_workers = [
            w for w, load in current_loads.items() 
            if load < threshold
        ]
        
        if not available_workers:
            raise Exception("System Capacity Full: No available workers under threshold.")

        # Pick the worker with the lowest current load
        new_worker = min(available_workers, key=lambda w: current_loads[w])

        # 3. Update Affinity Maps
        self.affinity_map[session_id] = new_worker
        if session_id not in self.worker_sessions[new_worker]:
            self.worker_sessions[new_worker].append(session_id)
        
        return new_worker

    def simulate_batch(self, sessions: list[str], initial_loads: dict):
        """Simulates routing a batch of requests."""
        routing_decisions = []
        for session in sessions:
            worker = self.assign_worker(session, initial_loads)
            initial_loads[worker] += 1
            routing_decisions.append((session, worker))
        return routing_decisions

# Usage Example
lb = StatefulLLMLoadBalancer()
# Simulate initial state: Worker 1 has 9 sessions, Worker 2 has 2
current_loads = {"worker_1": 9, "worker_2": 2}

# Session "abc" was previously assigned to worker_1 (simulated by pre-populating map)
lb.affinity_map["abc"] = "worker_1"
lb.worker_sessions["worker_1"].append("abc")

# New batch of incoming requests
new_requests = ["abc", "def", "ghi"]

print("Routing Decisions:")
for session, worker in lb.simulate_batch(new_requests, current_loads):
    print(f"Session: {session:5s} -> Routed to: {worker}")

# Output Explanation:
# Session 'abc' should route to worker_1 if load < 10. 
# Since worker_1 load is 9, it stays there. Load becomes 10.
# Session 'def' has no affinity, goes to least loaded (worker_2). Load becomes 3.
# Session 'ghi' has no affinity, goes to least loaded (worker_2). Load becomes 4.

Stratégies avancées : Affinité sémantique et migration préemptive

Pour les déploiements à grande échelle, la simple persistance par ID de session peut ne pas suffire. Les systèmes avancés emploient l'affinité sémantique, où l'équilibreur de charge hache les premiers jetons du prompt. Si deux utilisateurs différents commencent une conversation avec le même prompt système (par exemple, « Vous êtes un assistant de codage utile »), leurs caches KV pour ce préfixe peuvent être partagés ou maintenus sur le même nœud. Cela est particulièrement efficace pour les applications RAG (Retrieval-Augmented Generation) où les fenêtres de contexte sont grandes et statiques.

De plus, la mise en œuvre d'une migration préemptive est cruciale. Au lieu d'attendre qu'un worker plante en raison d'un OOM, l'équilibreur doit surveiller la pression mémoire. Lorsqu'un worker dépasse 85 % d'utilisation de la mémoire, il peut signaler à l'équilibreur de « démarrer à froid » de nouvelles sessions ailleurs et, si possible, sérialiser et décharger les caches KV des sessions moins actives vers un stockage plus lent mais de plus grande capacité (comme la RAM CPU ou les SSD NVMe) pour libérer de la mémoire GPU pour les requêtes de haute priorité.

Conclusion

L'équilibrage de charge avec état n'est plus une fonctionnalité optionnelle, mais un requirement central pour une infrastructure LLM efficace. En gérant l'affinité du cache KV et la persistance des sessions, nous pouvons réduire significativement le temps jusqu'au premier jeton (TTFT) et améliorer l'utilisation globale du cluster. La clé réside dans l'équilibre entre le coût de la migration d'état et le bénéfice de la localité de la cache. À mesure que les LLM continuent de croître en taille et en longueur de fenêtre de contexte, la sophistication de ces algorithmes de routage n'augmentera que, nous amenant des sessions collantes simples vers une orchestration de ressources intelligente et consciente de la sémantique. Pour les développeurs construisant des backends IA, l'intégration de la logique avec état tôt dans la couche d'équilibrage de charge est critique pour passer à l'échelle efficacement sans engager de coûts de calcul prohibitifs.