Comment OCRia a atteint une accélération de l'inférence par 10 sur les GPU B200

L'équipe d'ingénierie derrière OCRia-1 a détaillé comment elle a obtenu une augmentation quasi décuplée de la vitesse de génération de tokens, passant d'environ 285 tokens par seconde à 2800 tokens par seconde sur un seul GPU NVIDIA B200. Publiée le 3 novembre 2026, cette analyse technique révèle que ce bond de performance n'est pas le fruit d'une percée unique, mais plutôt d'une réduction systématique des goulots d'étranglement mémoire et de la surcharge computationnelle à travers toute la pile d'inférence.

Ce qui s'est passé

La latence a longtemps été une contrainte majeure pour les flux de travail d'analyse de documents, où les pages denses nécessitent des sorties structurées longues. L'équipe OCRia a identifié que leur pipeline d'inférence était limité par la bande passante mémoire plutôt que par la puissance de calcul, ce qui signifie que le GPU passait plus de temps à déplacer les poids du modèle et les données d'état qu'à effectuer des opérations mathématiques. Pour y remédier, ils se sont concentrés sur trois domaines clés : réduire le volume de données déplacées lors de chaque étape de décodage, générer plus de tokens par passe avant et minimiser la surcharge CPU et GPU associée à chaque étape.

La stratégie d'optimisation résultante combine des techniques établies comme la quantification avec des implémentations novatrices du décodage spéculatif et le développement de noyaux personnalisés. En profilant rigoureusement chaque étape du processus d'inférence, l'équipe a éliminé les transferts mémoire redondants et rationalisé la préparation des métadonnées. Ces changements permettent au système de gérer les requêtes concurrentes plus efficacement, accélérant ainsi les itérations sur les benchmarks et fournissant des résultats plus rapides aux utilisateurs finaux sans sacrifier la précision.

Comment cela fonctionne

Le fondement de cette accélération réside dans la quantification consciente de la charge de travail. Plutôt que d'appliquer une réduction uniforme de la précision, l'équipe a utilisé une précision mixte, conservant une haute précision pour les couches sensibles tout en compressant les autres en FP8 ou NVFP4. Ils ont employé une quantification consciente des activations, qui calibre les facteurs d'échelle à l'aide de données in-distribution pour garantir que la compression reste effectivement sans perte pour les tâches cibles. Fait notable, l'encodeur visuel a été laissé non quantifié pour préserver la précision visuelle, démontrant une approche ciblée pour équilibrer taille et fidélité.

Pour surmonter la limitation autorégressive consistant à générer un token par étape, l'équipe a mis en œuvre un décodage spéculatif utilisant un modèle brouillon appelé DFlash2. Ce modèle plus petit prédit plusieurs tokens futurs en parallèle, que le modèle cible plus grand vérifie ensuite lors d'une seule passe avant. DFlash2 utilise des convolutions locales et un sélecteur de candidats pour atteindre des taux d'acceptation supérieurs à 12 tokens, augmentant considérablement le débit. Le processus est mathématiquement garanti sans perte, car toute spéculation incorrecte est rejetée sans affecter la distribution de sortie finale.

Des gains supplémentaires proviennent des optimisations profondes du moteur. L'équipe a adapté ReplaySSM pour optimiser la gestion d'état dans les Gated Delta Networks, réduisant le coût de retour en arrière des tokens spéculatifs rejetés. Ils ont également fusionné des opérations consécutives, telles que l'activation SiLU et la quantification, pour maintenir les résultats intermédiaires dans les registres rapides du GPU au lieu de les écrire dans la mémoire à large bande passante. Côté CPU, ils ont refactorisé la construction des métadonnées pour éviter les calculs redondants entre différents groupes de cache, réduisant le nombre de lancements GPU requis par étape de décodage de plusieurs centaines à un seul.

Détails clés

• Gain de performance : La vitesse de génération de tokens est passée d'environ 285 tokens/s à environ 2800 tokens/s sur un seul GPU B200.
• Stratégie de quantification : Précision mixte utilisant BF16, FP8 et NVFP4, avec calibration consciente des activations pour maintenir la précision.
• Décodage spéculatif : Utilise DFlash2, un modèle brouillon à diffusion par blocs qui atteint des taux d'acceptation supérieurs à 12 tokens par passe.
• Optimisation des noyaux : Des noyaux CUDA personnalisés ont réduit le temps de vérification de ReplaySSM de 72–80 µs à 23–24 µs par couche.
• Efficacité CPU : La construction des métadonnées a été optimisée pour ne nécessiter qu'un seul lancement GPU par étape, contre environ 290 auparavant.
• Disponibilité : OCRia-1 est actuellement disponible sur LlamaParse, et les optimisations généralisables seront reversées à vLLM.

Pourquoi c'est important

Pour les ingénieurs développant des produits d'IA documentaire, cette étude de cas souligne que la puissance brute du matériel est souvent insuffisante sans une optimisation logicielle soignée. Le passage des opérations limitées par le calcul aux opérations limitées par la mémoire dans les grands modèles signifie que réduire les déplacements de données est fréquemment plus impactant que d'ajouter davantage d'unités de calcul. Comprendre le modèle roofline et identifier si une charge de travail est limitée par la bande passante ou le calcul permet aux équipes de prioriser les bonnes optimisations, telles que la quantification ou la fusion de noyaux, plutôt que de faire évoluer aveuglément les ressources.

L'utilisation du décodage spéculatif démontre une voie pratique pour briser le goulot d'étranglement séquentiel de la génération autorégressive. En vérifiant plusieurs tokens en parallèle, les développeurs peuvent améliorer significativement l'interactivité et le débit pour les utilisateurs traitant de longs documents. Cependant, cela nécessite une mise en œuvre minutieuse pour s'assurer que le modèle brouillon s'aligne bien avec le modèle cible, car de mauvais taux d'acceptation peuvent annuler les bénéfices. Le succès de l'équipe OCRia avec DFlash2 montre que investir dans des architectures de modèles brouillons spécialisés peut produire des retours substantiels.

Enfin, l'accent mis sur le profilage et l'amélioration incrémentale souligne l'importance de la mesure en ingénierie des systèmes. Beaucoup des gains provenaient de petites optimisations cumulatives comme la fusion d'opérations ou le raffinement de la gestion des métadonnées. Ces changements, bien que modestes individuellement, ont collectivement contribué à plus de 60 % du débit supplémentaire. Cette approche encourage les développeurs à regarder au-delà des paramètres de haut niveau des frameworks et à plonger dans les détails d'implémentation de leurs moteurs d'inférence pour libérer des performances cachées.

Ce que vous pouvez faire

• Profilez votre pipeline d'inférence pour déterminer s'il est limité par la bande passante mémoire ou par le calcul avant de choisir des stratégies d'optimisation.
• Implémentez une quantification à précision mixte avec calibration consciente des activations pour réduire l'utilisation mémoire sans perte significative de précision.
• Explorez le décodage spéculatif avec un modèle brouillon léger pour générer plusieurs tokens par passe avant.
• Fusionnez les opérations de noyau consécutives pour minimiser les écritures dans la mémoire à large bande passante et garder les données dans les registres.
• Optimisez la préparation des métadonnées côté CPU en mettant en cache les structures partagées et en réduisant les lancements GPU redondants.
• Reversez les optimisations généralisables vers des moteurs d'inférence open source comme vLLM pour bénéficier à la communauté plus large.