Résumé Technique : Inférence Efficace de LLM de Diffusion sur Appareils via NPU Mobile

1. Énoncé du Problème

Bien que les modèles de langage de diffusion (dLLM) offrent une alternative prometteuse aux modèles autorégressifs en permettant la génération parallèle de jetons (tokens), leur déploiement sur les appareils mobiles se heurte à des obstacles importants. Contrairement au décodage autorégressif, qui génère les jetons de manière séquentielle, les dLLM dénoisent de manière itérative des séquences masquées. Ce processus introduit trois défis spécifiques lors de l'utilisation des unités de traitement neuronal (NPU) mobiles :

1. Engagement de Jetons et Réduction de la Charge de Travail : Dans le décodage par blocs des dLLM, à mesure que le débruitage progresse, le nombre de jetons masqués restants dans un bloc diminue. Cela conduit à des « étapes de queue » (tail steps) où la charge de travail effective est trop faible pour exploiter pleinement les moteurs de tenseurs à haut débit des NPU, ce qui entraîne une mauvaise amortisation matérielle.
2. Révision de Jetons et Réutilisation du Cache KV : Les dLLM nécessitent un affinement itératif de l'état des jetons. Les jetons engagés tôt dans un bloc peuvent encore nécessiter une révision à mesure que le contexte devient disponible. Cela crée des mises à jour éparses et irrégulières du cache Clé-Valeur (KV) qui perturbent les modèles d'exécution denses et réguliers privilégiés par les NPU.
3. Accès à la Mémoire et Contraintes d'Espace d'Adressage : Les NPU mobiles disposent d'un espace d'adressage virtuel (VA) limité. L'inférence des dLLM nécessite un échange fréquent des poids du modèle, des caches KV et des activations entre la mémoire système et le NPU. Cela entraîne des surcoûts de remappage et de transfert de données qui peuvent annuler les bénéfices computationnels de l'accélération par NPU.

Les solutions existantes reposent souvent sur un déchargement (offloading) naïf vers le NPU, qui ne parvient pas à traiter la nature dynamique de l'inférence dLLM face au modèle d'exécution statique et basé sur des tuiles des NPU mobiles.

2. Méthodologie : llada.cpp

Les auteurs proposent llada.cpp, le premier cadre d'inférence conscient du NPU conçu pour aligner la nature de décodage parallèle des dLLM avec les caractéristiques d'exécution des NPU mobiles. Le framework intègre trois techniques de base :

A. Décodage Spéculatif Multi-Blocs

Pour résoudre le problème de la réduction de la charge de travail lors des phases tardives de décodage, llada.cpp étend dynamiquement la fenêtre d'exécution.

• Mécanisme : Au lieu de limiter la passe avant (forward pass) du NPU au bloc actuel, le système incorpore de manière spéculative des jetons provenant de blocs futurs dans la fenêtre de décodage actuelle.
• Stratégie : Il emploie un Glissement Continu (déplacement de la fenêtre vers la droite à mesure que les jetons s'engagent) et un Regard Anticipé Adaptatif (extension proactive de la fenêtre pour inclure des blocs futurs lorsque le nombre de jetons masqués du bloc actuel descend sous un certain seuil).
• Bénéfice : Cela maintient une densité élevée de jetons actifs pour le NPU, garantissant que les frais fixes (lancements de noyaux, transferts de données) sont amortis sur une charge de travail plus large et productive, sans violer la sémantique de décodage originale.

B. Révision Progressive à Double Voie

Pour gérer l'irrégularité des révisions de jetons sans bloquer le NPU, le framework sépare le calcul dense des mises à jour éparses.

• Stabilisation par Étapes des Jetons : Les jetons passent par trois états : Invisible (masqué/faible confiance), Visible (prédiction concrète mais révisable) et Stable (fiable pour une réutilisation à long terme).
• Exécution à Double Voie :
  • Voie NPU : Dédiée au débruitage dense du bloc actuel.
  • Voie CPU : Gère les mises à jour éparses pour les jetons « Visibles » qui nécessitent une révision. Le CPU recompute de manière asynchrone la confiance, rafraîchit les états KV et met à jour les identités des jetons en utilisant des noyaux spécialisés qui interprètent les formats de tuiles spécifiques au NPU (HMX).
• Saut de Logits Sélectif : Pour réduire davantage la charge du CPU, le système saute la projection de sortie et l'estimation de la confiance pour les jetons ayant déjà atteint l'état « Stable ».
• Fusion Différée : Les entrées KV rafraîchies par le CPU sont fusionnées dans le cache de préfixe uniquement aux limites des étapes (step boundaries), évitant ainsi une synchronisation trop fine.

C. Runtime de Mémoire Optimisé pour l'Échange (Swap)

Pour atténuer l'espace d'adressage limité du NPU et les surcoûts de transfert, le framework optimise la gestion de la mémoire.

• Cartographie de Tampon Guidée par Graphe : Le runtime analyse le graphe de calcul pour classer les tenseurs comme « à longue durée de vie » (poids, cache KV) ou « à courte durée de vie » (activations). Il utilise des métriques profilées pour déterminer quels objets restent résidents dans l'espace d'adressage virtuel visible par le NPU et lesquels sont mappés temporairement, minimisant ainsi la fragmentation.
• Transfert de Données par Pipeline : Le système utilise un double tampon de transit (double-buffered staging). Pendant que le NPU exécute les opérateurs, l'hôte prépare les descripteurs de tampons et transfère les données pour l'opérateur suivant dans un emplacement de transit séparé. Cela superpose les coûts de mouvement de données et de mappage avec le calcul, masquant ainsi la latence.

3. Contributions Clés

Le papier revendique les contributions suivantes :

• Identification du Décalage : Il caractérise le désalignement fondamental entre les modèles d'exécution statiques des NPU mobiles et l'inférence par blocs dynamiques des dLLM, soulignant spécifiquement les défis liés à l'engagement des jetons, à la révision et à l'accès à la mémoire.
• Framework llada.cpp : Il introduit le premier cadre d'inférence conscient du NPU pour les dLLM sur smartphones, combinant des innovations algorithmiques (décodage spéculatif, révision progressive) avec des optimisations au niveau système (runtime de mémoire).
• Évaluation Complète : Il valide le framework sur diverses plateformes matérielles (Qualcomm SM8650, SM8750, SM8850) et charges de travail (LLaDA-8B, Dream-7B), démontrant une traduction efficace du parallélisme des dLLM en une réduction de la latence.

4. Résultats

L'évaluation a été menée sur des smartphones incluant le OnePlus Ace5 Pro (SM8750). Les conclusions clés incluent :

• Réduction de la Latence : llada.cpp réduit la latence de génération de bout en bout de LLaDA-8B de 17× à 42× par rapport à une base CPU avec réutilisation du cache KV de préfixe.
• Comparaison avec les Modèles Autorégressifs : Sur les SoCs récents (SM8750, SM8850), llada.cpp atteint jusqu'à 3,9× d'accélération par rapport à un modèle autorégressif de taille égale (Llama-3-8B) tournant sur le même matériel, soulignant l'avantage du parallélisme de diffusion sur les NPU mobiles.
• Préservation de la Qualité : Le framework maintient une qualité de génération comparable aux bases CPU. Les études d'ablation montrent que si le décodage spéculatif multi-blocs seul peut légèrement réduire la précision en raison d'une utilisation prématurée des jetons futurs, le mécanisme de révision progressive à double voie restaure la qualité en gardant les jetons précoces révisables.
• Efficacité Énergétique : En fournissant un travail parallèle suffisant au NPU, llada.cpp réduit à la fois la puissance moyenne et l'énergie par requête. En revanche, l'exécution de modèles autorégressifs sur les NPU montrait une consommation d'énergie plus élevée en raison de la surcharge liée à l'ordonnancement de charges de travail séquentielles de petite taille.
• Scalabilité : Le système gère avec succès des exécutions de longue durée (jusqu'à 512 jetons) et différentes architectures de dLLM (Dream-7B), prouvant la robustesse de la gestion de la mémoire et des stratégies de décodage.

5. Signification et Revendications

Le papier positionne llada.cpp comme une étape critique pour rendre les dLLM pratiques pour les applications mobiles sensibles à la latence. Il soutient que le « calcul parallèle répété » des dLLM, souvent considéré comme un handicap, peut être converti en avantage en exploitant les capacités de calcul tensoriel des NPU mobiles modernes.

Les auteurs soulignent que le simple déchargement des dLLM vers les NPU est insuffisant ; le succès nécessite une approche de co-conception qui respecte les contraintes du matériel (formes de tuiles statiques, espace d'adressage limité) tout en exploitant les forces de l'algorithme (débruitage parallèle). Ce travail démontre qu'avec les bonnes optimisations système, les dLLM peuvent surpasser les modèles autorégressifs traditionnels sur les appareils mobiles, offrant une voie vers une IA à faible latence, respectueuse de la vie privée et accessible sur les smartphones.