Résumé Technique : FlashSpec

Problématique

L'inférence des grands modèles de langage (LLM) est fondamentalement limitée par la bande passante mémoire, car la génération autorégressive nécessite le chargement de tous les paramètres du modèle pour chaque nouveau jeton (token). Le décodage spéculatif remédie à cela en utilisant un modèle « draft » plus petit pour proposer $\gamma$ jetons candidats, qu'un modèle « cible » plus grand vérifie lors d'une seule passe avant (forward pass). Bien que théoriquement fondé, les implémentations open-source existantes souffrent de deux lacunes pratiques critiques :

1. Goulot d'étranglement de la vérification côté CPU : La décision d'acceptation/rejet par jeton est généralement exécutée en Python/PyTorch sur le CPU. Cela nécessite une synchronisation dispositif-hôte à chaque étape de décodage, créant un blocage de pipeline proportionnel à la taille du vocabulaire ($V$), qui varie de 32k à 128k jetons dans les LLM modernes.
2. Sélection de draft statique : Les systèmes actuels sélectionnent un modèle draft hors ligne et le maintiennent fixe tout au long du déploiement. Cela ignore la nature non stationnaire des charges de travail d'inférence, où les taux d'acceptation varient considérablement selon les domaines de prompts (ex: texte conversationnel vs code source), entraînant un débit sous-optimal.

Méthodologie

FlashSpec est un moteur de décodage spéculatif adaptatif open-source conçu pour résoudre ces deux goulots d'étranglement simultanément, tout en préservant la distribution de sortie exacte du modèle cible.

1. Noyau de vérification sur dispositif (On-Device)

L'innovation centrale est un noyau GPU Triton qui effectue le test d'acceptation/rejet entièrement sur le dispositif.

• Mécanisme : Au lieu de collecter l'intégralité des logits du vocabulaire, le noyau ne lit que deux scalaires de log-probabilité par jeton candidat (un provenant du modèle draft, un du modèle cible).
• Complexité : L'empreinte SRAM du noyau est de $O(1)$ par rapport à la taille du vocabulaire, indépendamment du fait que $V$ soit de 32k ou 128k. Le trafic mémoire par étape est réduit de $O(B\gamma V)$ à $O(B\gamma)$, où $B$ est la taille du batch.
• Stabilité numérique : L'implémentation utilise l'arithmétique en espace logarithmique pour éviter le sous-dépassement (underflow) et applique correctement l'échelle de température avant l'opération log-softmax, une source courante de bugs subtils dans les implémentations précédentes.

2. Sélection de Draft par Bandit en ligne

FlashSpec cadre la sélection du modèle draft comme un problème de bandit multi-bras ($K$-armed bandit).

• Adaptation : À chaque étape de décodage, le système sélectionne un modèle draft (un « bras ») basé sur ses performances récentes. La récompense est définie comme la fraction de jetons draft acceptés ($\alpha$).
• Algorithmes : Le système prend en charge UCB1 (Upper Confidence Bound) et l'Échantillonnage de Thompson (Thompson Sampling). Ces algorithmes s'adaptent aux taux d'acceptation non stationnaires sans réajustement humain, minimisant le regret cumulé au fil du temps.
• Garantie théorique : La politique du bandit dépend uniquement de l'historique des acceptations passées, et non des jetons draft actuels ou de la distribution cible, garantissant que la distribution jointe des jetons générés reste identique à l'échantillonnage autorégressif standard du modèle cible.

3. Correction et Tests

Le système inclut une suite de tests rigoureux pour garantir la correction mathématique :

• Équivalence de distribution : Un test de Kolmogorov–Smirnov (KS) est imposé comme une barrière stricte dans le pipeline CI, comparant la sortie de l'échantillonneur spéculatif à la distribution cible sur 10 000 échantillons avec un seuil de signification de $\alpha=0,01$.
• Tests basés sur les propriétés : La bibliothèque Hypothesis est utilisée pour tester les cas limites (ex: tailles de batch et longueurs de spéculation spécifiques) que les tests unitaires standards pourraient manquer.

Résultats Clés

Les mesures initiales ont été effectuées sur un GPU Tesla T4 utilisant TinyLlama-1.1B-Chat (quantification NF4 4-bits) avec une longueur de spéculation $\gamma=4$.

• Débit (Throughput) : Le système a atteint 44,2 jetons/seconde avec un taux d'acceptation moyen ($\bar{\alpha}$) de 0,75.
• Latence : La latence par étape p50 a été mesurée à 22,1 ms.
• Efficacité du noyau : Sur le T4, le noyau Triton a montré une accélération de 1,6$\times$ par rapport à une implémentation de référence en pur PyTorch pour une taille de batch de 8, bien qu'il ait été plus lent pour un batch de 1 en raison de la bande passante mémoire limitée du T4 (320 Go/s). Les auteurs notent que les bénéfices du noyau devraient augmenter significativement sur du matériel à haute bande passante comme le H100 (3,35 To/s).
• Performance du Bandit : Dans une simulation de $T=10\,000$ tours avec $K=3$ bras, le sélecteur UCB1 a atteint un regret cumulé de 100,2, bien en deçà de la limite supérieure théorique d'environ 526. L'échantillonnage de Thompson a convergé encore plus rapidement avec un regret de 18,9.

Signification et Revendications

L'article positionne FlashSpec comme le premier système traitant simultanément le goulot d'étranglement de la vérification CPU et la limitation de la sélection statique de draft dans un cadre open-source prouvablement correct.

• Contribution théorique : Les auteurs fournissent une preuve formelle (Théorème 1) que la distribution de sortie de FlashSpec est identique à l'échantillonnage autorégressif standard, quelle que soit la politique de sélection du bandit.
• Impact pratique : En déplaçant la vérification sur le GPU et en adaptant la sélection de draft en ligne, FlashSpec vise à maximiser le débit à travers diverses distributions de prompts sans intervention manuelle.
• Rigueur d'ingénierie : L'article souligne des leçons d'ingénierie spécifiques, telles que la nécessité d'appliquer la température avant le log-softmax et l'importance de grands échantillons (10 000) dans les tests d'équivalence de distribution pour détecter des décalages subtils.

Les auteurs notent que les benchmarks complets sur des modèles plus larges (Llama-3-8B/70B) et les comparaisons avec Medusa et EAGLE sur du matériel H100 sont actuellement en cours et seront publiés dans le dépôt associé. Le logiciel est publié sous licence Apache 2.0.