DiffusionGemma : Google joue pour le titre de champion de l’inférence locale

Plus de 1000 tokens par seconde sur un seul GPU

Avec son expérimentation DiffusionGemma, Google veut aller plus loin sur ce dernier point. Exécuté sur un GPU Nvidia H100 (80 Go de VRAM), le LLM basé sur Gemma 4 A4B-26B (26 milliards de paramètres au total, mais moins de 4 actifs) produirait plus de 1100 tokens par seconde, contre 300 pour son modèle d’origine avec la prédiction multitoken enclenchée et environ 700 tokens/seconde sur un GPU « prosumer » Geforce RTX 5090 (32 Go de VRAM). Google a également travaillé avec Nvidia pour prendre en charge la carte RTX 4090, doté de 24 Go de VRAM. DiffusionGemma réclame 18 Go de VRAM quand il est compressé dans un format 4 bits. Nvidia l’a testé sur son DGX Spark et DGX Station. Résultat, le modèle génère respectivement 150 tokens/s et jusqu’à 2000 tokens/s.

Jusqu’alors, Mistral AI et OpenAI ont eu recours au service de Cerebras, un fabricant de puces spécialisées, pour obtenir des pointes de vitesse équivalentes. La méthode de Google n’implique pas la livraison d’une nouvelle puce ou d’un TPU. Explications.

Vers des modèles taillés pour l’inférence locale

Pour rappel, la plupart des grands modèles de langage sont dits « autorégressifs ». Pour faire simple, ils prédisent le mot suivant dans une phrase à partir des termes déjà présents en entrée. L’opération est répétée jusqu’à obtenir le résultat final. Cette technique, rappelle Google, est davantage limitée par la mémoire que le calcul.

C’est l’une des raisons qui ont poussé Nvidia à proposer des designs d’accélérateurs dotés de grandes quantités de VRAM HBM3e et GDDR7.

À gros trait, la génération se produit en deux phases. Il faut d’abord transférer les poids en mémoire vers les unités de traitement, puis que ces unités exécutent la tâche de génération. C’est la première phase qui prend le plus de temps. « Parce que les paramètres n’ont besoin d’être chargés qu’une seule fois par étape, peu importe la taille du lot en entrée, générer un token prend presque autant de temps qu’il y ait 1 ou 256 utilisateurs », expliquent les ingénieurs de Google.

Le reste du temps, les unités de calcul sont au repos. Ce n’est pas un problème majeur pour les fournisseurs de cloud et de LLM. L’architecture mixture of experts et les librairies logicielles associées ont permis de paralléliser les modèles et de saturer les GPU ou les TPU. La vitesse relativement lente des LLM est finalement un compromis acceptable au moment de servir des millions d’utilisateurs simultanément.

Or pour un usage local ou sur une instance dotée d’une seule carte accélératrice, ce comportement conduit à la sous-utilisation des ressources de calcul.

Pour exploiter la puissance de calcul, Google a fait en sorte que DiffusionGemma génère des blocs de 256 tokens en une seule fois.

« Le modèle initialise une séquence vierge de 256 tokens aléatoires – appelée “canevas” – puis évalue et affine simultanément l’ensemble du canevas de manière itérative », indique le fournisseur cloud. « Ainsi, le modèle n’est plus limité par la mémoire, mais par le calcul ».

Ce canevas découle de la mise en place d’un mécanisme d’attention bidirectionnelle. L’attention est le terme qui définit la manière dont le modèle prend en compte le « contexte » – les données en entrées ou déjà en mémoire – pour générer des tokens. Depuis l’avènement des LLM avec GPT, l’attention est unidirectionnelle, uniquement portée vers les séquences passées pour prédire le token suivant.

Le retour en grâce de l’attention bidirectionnelle

La bidirectionnalité n’est pas nouvelle. Elle était un attribut de l’architecture Transformer et, par conséquent, des modèles de type BERT. Ici, elle permet que les paramètres derrière chaque token prennent en compte tous les autres tokens du canevas.

De ce fait, DiffusionGemma profite des mêmes avantages que les modèles de complétion de texte et de code.

À ceci près que le mécanisme d’attention bidirectionnelle est utilisé pour un exercice d’autocorrection à l’échelle du canevas en plusieurs passes pour ne conserver que les meilleurs résultats.

C’est justement cette phase d’autocorrection qui explique pourquoi Google n’a pas appelé son modèle « BidirectionnalGemma ». Ici, les chercheurs ont tenté d’appliquer la technique qui sous-tend l’existence des modèles de diffusion, le véritable nom des modèles capables de générer des images. À l’entraînement, un modèle de diffusion ajoute progressivement du bruit aléatoire dans des séquences de données (la plupart du temps des images, des vidéos, du son), jusqu’à ce qu’elles soient méconnaissables, puis apprend à reconstruire la distribution de données initiale.

La phase d’autocorrection est perçue comme une « réduction du bruit ». Avec BERT ou les modèles de complétion de code, la manière de représenter ce bruit dans un corpus consistait à masquer automatiquement certains mots à l’aide d’un token spécial. Ici, Google exploite une autre approche : la diffusion d’état uniforme. « Au lieu d’un token de masquage rigide, le bruit est introduit en remplaçant les mots d’origine par des tokens aléatoires présents dans le vocabulaire du modèle », expliquent les chercheurs de Google DeepMind.   

Lors de la phase de correction, le modèle évalue les tokens du canevas, leur attribue un score de confiance et les mets à jour si le résultat est faible.

« Si la probabilité d’un token tombe en dessous d’un seuil en raison de l’apparition d’un nouveau contexte lors des étapes suivantes, celui-ci est remplacé par un nouveau token aléatoire », ajoutent-ils. « Ce cycle permet une correction continue des erreurs et un affinement parallèle du canevas ».

À l’instar de BERT, DiffusionGemma est doté d’une architecture encodeur-décodeur où l’encodeur autorégressif est chargé de « traiter et mettre en cache le contexte de la requête » de manière incrémentale (le « prefill ») dans la mémoire du GPU (KV Cache), tandis que le décodeur porte le mécanisme d’attention bidirectionnelle pour « débruiter » le canevas. L’attention n’est pas seulement portée sur le canevas en cours de traitement, mais sur l’ensemble du cache afin de traiter l’entièreté du contexte. La fenêtre fait encore 256 000 tokens, tandis que la fenêtre d’attention glissante contient 1024 tokens à la fois.

De ce point de vue là, un modèle autorégressif est un peu plus simple : il n’a besoin que d’un décodeur pour produire du texte. Et les chercheurs ne détaillent pas les potentiels inconvénients de la bidirectionnalité à l’entraînement. Peut-on profiter de la loi empirique de la mise à l’échelle (« scaling laws ») ?

DiffusionGemma ne fait pas mieux que son aîné et est taillé pour les GPU Nvidia

DiffusionGemma profite toutefois des avantages obtenus ces quatre dernières années, à savoir l’architecture mixture of experts – qui n’activent qu’un petit nombre de paramètres à l’inférence – et la quantisation NVFP4 (la compression) qui permet de diminuer l’espace nécessaire à la mise en cache des poids et des tokens.

Malgré tout, DiffusionGemma n’est pas meilleur que Gemma 4 A4B-26B. Il perd 5 points de pourcentage sur l’évaluation MMLU Pro, 8 points sur le test LiveCodeBenchv6, 9 points sur le banc d’essai GPQA Diamond, 19 points sur l’examen de mathématiques AIME 2026, et 12 points sur le test agentique t2-bench. Artificial Analysis n’en a pas encore fait le test.

Un fine-tuning sur un domaine spécifique permettrait de tirer partie à la fois de la vitesse d’exécution et de meilleurs résultats, suggère Google. D’autant que le fournisseur n’a pas retiré du modèle l’encodeur capable d’interpréter les images et les vidéos.

L’autre point à prendre en compte, c’est que cette technique est principalement valide pour les GPU Nvidia. « Parce que cette accélération repose sur l’exploitation de l’intensité arithmétique élevée des accélérateurs, les architectures à mémoire unifiée comme celles des Mac Apple Silicon – qui sont souvent limitées par la bande passante mémoire plutôt que par la puissance de calcul lors de l’inférence – pourraient ne pas bénéficier de la même accélération par rapport aux modèles autorégressifs comme Gemma 4 », écrivent les chercheurs de Google en note de page. Nvidia a d’ailleurs félicité Google pour ce lancement. Habituellement, le géant du cloud se concentre sur ses TPU.

Sous licence Apache 2.0, DiffusionGemma est disponible depuis HuggingFace, Gemini Enterprise Agent Platform, Unsloth, NeMo et peut s’exécuter sur vLLM.