SparDA : Attention Découplée Éparse pour une Inférence Efficace de LLM à Long Contexte

Résumé

L’attention clairsemée réduit les besoins en calcul et en bande passante mémoire pour l’inférence de grands modèles de langage (LLM) sur des contextes longs. Cependant, deux défis majeurs subsistent : (1) la capacité du cache KV continue de croître avec la longueur de séquence, et son déchargement vers la mémoire CPU introduit un goulot d’étranglement lié au transfert PCIe ; (2) l’étape de sélection clairsemée conserve elle-même une complexité en O(T²) et peut dominer le coût de l’attention pour les contextes longs. Nous proposons SparDA, une architecture d’attention clairsemée découplée qui introduit une quatrième projection par couche, le Forecast, aux côtés des projections Query, Key et Value. Le Forecast prédit les blocs KV nécessaires à la couche suivante, ce qui permet une sélection anticipée superposant le préchargement CPU-à-GPU à l’exécution de la couche courante. Grâce au découplage du Forecast de la requête d’attention, notre implémentation GQA utilise une tête Forecast par groupe GQA, réduisant ainsi le surcoût de sélection par rapport au sélecteur multi-têtes original. SparDA ajoute moins de 0,5 % de paramètres et n’entraîne que les projections Forecast en les faisant correspondre à la distribution d’attention du sélecteur original. Sur deux modèles de 8B paramètres pré-entraînés avec attention clairsemée, SparDA atteint, voire améliore légèrement la précision, et offre jusqu’à 1,25× d’accélération du préremplissage et 1,7× d’accélération du décodage par rapport à la baseline de déchargement avec attention clairsemée. En permettant des tailles de lots réalisables plus grandes sur un seul GPU, SparDA atteint en outre jusqu’à 5,3× de débit de décodage supérieur à la baseline clairsemée sans déchargement. Notre code source est disponible à l’adresse https://github.com/NVlabs/SparDA.

English

Sparse attention reduces compute and memory bandwidth for long-context LLM inference. However, two key challenges remain: (1) KV cache capacity still grows with sequence length, and offloading to CPU memory introduces a PCIe transfer bottleneck; (2) the sparse selection step itself retains O(T^2) complexity and can dominate attention cost at long contexts. We propose SparDA, a decoupled sparse attention architecture that introduces a fourth per-layer projection, the Forecast, alongside Query, Key, and Value. The Forecast predicts the KV blocks needed by the next layer, enabling lookahead selection that overlaps CPU-to-GPU prefetch with current-layer execution. Because Forecast is decoupled from the attention query, our GQA implementation uses one Forecast head per GQA group, reducing selection overhead versus the original multi-head selector. SparDA adds <0.5% parameters and trains only the Forecast projections by matching the original selector's attention distribution. On two sparse-pretrained 8B models, SparDA matches or slightly improves accuracy and delivers up to 1.25times prefill speedup and 1.7times decode speedup over the sparse-attention offload baseline. By enabling larger feasible batch sizes on a single GPU, SparDA further reaches up to 5.3times higher decode throughput than the non-offload sparse baseline. Our source code is available at https://github.com/NVlabs/SparDA.