Stratified Sparse Attention accélère le décodage d'un facteur 2,96 sur un contexte de 1M

Une équipe de chercheurs propose Stratified Sparse Attention (SSA), une méthode qui n'accède qu'à une petite fraction du cache KV tout en perdant nettement moins de qualité que les schémas parcimonieux précédents. Sur RULER-HARD avec Ministral-3-8B, à 50 fois de parcimonie de l'attention, SSA fait passer la précision de 81,8% à 84,2% sur un contexte de 32K et de 63,8% à 69,3% sur 128K. Cela représente 94% et 85% du résultat de l'attention dense complète. Sur un contexte d'un million de tokens, des noyaux CUDA optimisés accélèrent le décodage jusqu'à 2,96 fois par rapport à FlashInfer.

Où l'attention parcimonieuse perd en qualité

Avec un contexte long, le goulot d'étranglement de l'inférence n'est plus tant le calcul que la lecture du cache KV, qui grossit à chaque étape de décodage. L'attention parcimonieuse (sparse attention) réduit ce coût en n'accédant qu'à une partie du cache. Deux voies s'opposent alors.

La première est la sélection déterministe. Top-k ou top-p retiennent les tokens aux scores les plus élevés. Cette approche capte de façon fiable l'essentiel, mais tout le reste de l'attention, négligeable token par token et pourtant potentiellement significatif en somme, est tout simplement jeté. La seconde voie est l'échantillonnage. Il couvre mieux la queue de la distribution, mais au prix d'une variance : le résultat fluctue davantage d'une exécution à l'autre.

Les approches hybrides associent un top-k à un échantillonnage uniforme. Les auteurs montrent que, à forte parcimonie, l'écart avec l'attention dense reste important pour ces méthodes. L'échantillonnage uniforme est aveugle : il dépense son budget sur des tokens qui influencent à peine la réponse.

Comment SSA répartit le cache en strates

SSA construit un schéma unique en deux parties. Le modèle lit de façon déterministe les tokens aux scores élevés. Le reste du cache n'est ni écarté ni tiré au hasard : ses tokens sont échantillonnés avec des probabilités proportionnelles à des scores proxy d'attention. Pour que la contribution des tokens sélectionnés ne fausse pas le résultat, SSA applique une correction fondée sur leur probabilité d'inclusion dans l'échantillon. Cela permet de restituer la masse d'attention située en dehors de l'ensemble retenu.

L'idée s'apparente à l'échantillonnage stratifié classique : la zone importante est couverte en totalité, et le reste est estimé à partir d'un sous-échantillon pondéré. Les auteurs soulignent que SSA est une généralisation. Top-k, top-p et l'hybride top-k plus échantillonnage en sont des cas particuliers, si bien que la méthode peut servir de cadre commun à toute cette famille.

Chiffres et accélération

L'évaluation porte sur quatre benchmarks, RULER, LongBench, LOFT et AIME, et sur cinq modèles de 3B à 27B paramètres. À parcimonie égale, SSA surpasse systématiquement les méthodes de référence. Principaux résultats :

• RULER-HARD, Ministral-3-8B, 32K : 81,8% → 84,2% face à la meilleure baseline pratique, soit 94% de l'attention dense ;
• RULER-HARD, Ministral-3-8B, 128K : 63,8% → 69,3%, soit 85% de l'attention dense ;
• régime : parcimonie de l'attention de 50 fois ;
• vitesse : jusqu'à 2,96x en décodage par rapport à FlashInfer sur un contexte de 1M ;
• code ouvert : github.com/skylight-org/SSA-Paper-Code.

Le second volet du résultat compte aussi. Une méthode plus précise exige généralement des opérations plus complexes, et le gain de vitesse peut s'évaporer. Les auteurs montrent ici, avec des noyaux CUDA optimisés, que le gain en qualité n'est pas absorbé par le surcoût. Le gain de précision est plus marqué sur 128K que sur 32K (5,5 points de pourcentage contre 2,4), ce qui signifie que la méthode aide davantage lorsque le contexte s'allonge et que la parcimonie fait plus mal.

Ce que cela change pour les praticiens

Pour les équipes qui servent du contexte long, le cache KV pèse directement sur la mémoire GPU et la bande passante. Si une réduction de 50 fois des accès au cache préserve 85–94% de la qualité de l'attention dense sur des tâches difficiles, les ingénieurs disposent d'un levier supplémentaire entre coût et précision. C'est particulièrement intéressant pour les scénarios agentiques et le travail sur de grandes bases de code, où le contexte grossit vite.

Les limites sont également visibles. Une partie des chiffres concerne un seul modèle, Ministral-3-8B, et à 128K l'écart avec le mode dense reste net : 85% est un bon résultat, mais pas la parité. En production, il faudra tester la méthode sur ses propres tâches et vérifier la compatibilité des noyaux avec la pile utilisée. Le code ouvert abaisse toutefois le seuil d'entrée : il est possible de reproduire les expériences et de comparer SSA aux schémas top-k actuels sans attendre une implémentation dans les frameworks populaires.