The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Yu Lin; Yiming Wang; Runyuan Cai; Hanze Liu; Xiaodong Zeng

Articles · arXiv 2609.18063

L'autre moitié du mur de la mémoire : servir des MoE de 35B depuis un SSD grâce à la prédiction de routage entraînée

The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Yu Lin, Yiming Wang, Runyuan Cai, Hanze Liu, Xiaodong Zeng

Publié le 2026-09-16cs.AI arXiv PDF

Lire avec traduction et IA

Résumé en trois lignes

🎯 Les MoE de classe 35B nécessitent ~19,5 Go même en int4, dépassant la RAM grand public, tandis que le déchargement naïf sur SSD bloque sur la latence disque car le routage dépend de la couche précédente.
🔧 Edge0 diffuse les experts MoE depuis le SSD via un prerouter par couche prédisant le routage avec un token d'avance (« prédiction comme routage »), associé à un LoRA de récupération parallèle non fusionné distillé sur la trajectoire étudiante.
📈 Sert un MoE 35B à ~20,4 tok/s pour 2,9 Gio de mémoire pic sur un Mac mini M4 Pro (24 Go), à seulement 3,9 points d'écart moyen sur les benchmarks par rapport à l'enseignant fp16.

Résumé

L'inférence des modèles de type mélange d'experts (MoE) sur du matériel grand public est limitée par la mémoire des poids : un modèle de la classe 35B nécessite 19,5 Go en 4 bits, et la parcimonie réduit le calcul par token, mais pas les octets devant être conservés en mémoire. Le déchargement naïf sur SSD ne suffit pas à lui seul, car les experts de la couche N+1 doivent être sélectionnés avant que la sortie de la couche N n'existe, empêchant les lectures de débuter assez tôt pour être masquées par le calcul. Nous présentons Edge0, un moteur d'inférence MoE en streaming qui comble cet écart grâce à un prerouter : une tête dédiée par couche prédit le routage de la couche suivante avec un token d'avance, et cette prédiction est directement utilisée comme routage effectif, de sorte que l'ensemble des experts préchargés coïncide avec l'ensemble routé sans aucun abandon de token. Un LoRA de récupération non fusionné, entraîné sur la trajectoire de l'étudiant, compense la perte de qualité due à la quantification int4 et au remplacement du routage. Sur une unique machine de 24 Go, Edge0 sert un MoE de 35B à 20 tok/s avec une mémoire active maximale de 3 Gio, se situant en moyenne à quelques points de son enseignant fp16 sur cinq benchmarks publics. Une version 8B fonctionne sur le même framework, et le code, les points de contrôle ainsi que les adaptateurs sont open source.

Résumé original (anglais)

Mixture-of-experts (MoE) inference on consumer hardware is bounded by weight memory: a 35B-class model is 19.5GB at 4-bit, and sparsity shrinks the compute per token, not the bytes that must be held. Naive offloading to SSD does not help on its own, because layer N+1's experts must be chosen before layer N's output exists, so the reads cannot start early enough to hide behind compute. We present Edge0, a streaming MoE inference engine that closes the gap with a prerouter: a per-layer head predicts the next layer's routing one token ahead, and the prediction is consumed as the routing itself, so the staged expert set equals the routed set and nothing is dropped. An unmerged recovery LoRA, trained on the student path, pays back the quality lost to int4 quantization and routing replacement. On a single 24GB machine, Edge0 serves a 35B MoE at 20tok/s inside 3GiB of peak active memory, within a few points of its fp16 teacher on average across five public benchmarks. An 8B tier runs on the same framework, and the framework, checkpoints, and adapters are open source.

Guide de lecture

Problème et motivation

Les modèles de type mélange d'experts (MoE) augmentent le nombre total de paramètres tout en maintenant un faible nombre de paramètres actifs par token (par exemple, en n'activant qu'environ 3 milliards de paramètres par token dans un modèle de 35B). Cependant, l'inférence sur du matériel grand public reste limitée par la mémoire statique requise pour les poids : un MoE de 35B nécessite environ 19,5 Go, même quantifié en int4. Sur une machine de 24 Go, cette empreinte monopolise l'espace nécessaire au système d'exploitation et au cache KV, entraînant une pagination mémoire critique et une dégradation importante du débit de décodage (par exemple, 3,9 tok/s pour un modèle entièrement résident en mémoire sur un Mac mini M4 Pro).

Décharger les poids des experts sur un SSD réduit l'empreinte mémoire, mais un chargement à la demande naïf introduit un goulet d'étranglement de sérialisation strict : la sélection des experts pour la couche $N + 1$ dépend de la sortie de la couche $N$ . Les lectures SSD ne peuvent pas commencer suffisamment tôt pour être masquées par les calculs, bloquant l'exécution sur la latence du disque une fois par couche et par token.

Idées clés

Le SSD comme niveau de stockage des poids : Les poids des experts MoE restent sur SSD au format int4 et sont mappés en mémoire (mmap) dans un pool de streaming compact. La mémoire de crête est déterminée par le working set actif plutôt que par la taille totale des paramètres du modèle.
Prédiction comme routage (Prediction as Routing) : Au lieu d'attendre la fin des calculs de la couche $N$ , un prerouter par couche prédit le routage de la couche $N + 1$ avec un token d'avance complet. Le modèle adopte directement cette prédiction en tant que routage au moment du décodage. Cela élimine par construction tout abandon de token ou rechargement d'urgence, car l'ensemble préchargé correspond strictement à l'ensemble routé.
LoRA de récupération non fusionné : La quantification 4 bits et le routage approximé entraînent une perte de qualité. Edge0 entraîne un adaptateur de bas rang sur la trajectoire de routage de l'étudiant via distillation. Comme la fusion du LoRA dans la base int4 suivie d'une re-quantification efface la quasi-totalité du delta (seuls 18 % de rétention au niveau des logits subsistent), Edge0 sert l'adaptateur sous la forme d'un delta parallèle non fusionné avec un surcoût négligeable.

Méthode

Architecture de la tête de pré-routage (Prerouter Head)

Chaque tête de prerouter appartenant à la couche $N$ s'exécute au token $t$ sur la sortie de la normalisation de couche post-attention et prédit le routage pour la couche $N + 1$ au token $t + 1$ . Les caractéristiques d'entrée concatènent l'état caché avec deux vecteurs de routage one-hot top- $k$ (token actuel et précédent) :

$dim = d_{model} + 2 E$

Pour la classe 35B, cela donne $2048 + 2 \times 256 = 2560$ . La tête calcule :

$head (x) = fc_{2} (erf - gelu (fc_{1} (x))) + ℓ (x)$

où $ℓ (x)$ est une branche résiduelle linéaire initialisée à chaud à partir des poids du routeur de la couche suivante.

Exécution de l'adaptateur non fusionné

La couche MoE servie préserve les octets de base quantifiés et exécute en parallèle le delta LoRA :

$y = W_{int4} (x) + \frac{α}{r} B A x$

Entraînement en trois phases

Distillation de la tête : Les têtes de prerouter sont entraînées à imiter les véritables décisions de routage de la couche suivante de l'enseignant, tandis que le reste du réseau reste gelé.
SFT sur la trajectoire étudiante : Les adaptateurs LoRA de récupération (attachés aux projections d'attention, d'attention linéaire et d'experts partagés, mais pas aux experts en streaming) sont entraînés à l'aide de textes générés par l'enseignant tout en forçant le routage via le prerouter.
Distillation on-policy : L'étudiant génère du texte sous exécution pré-routée, l'enseignant de base en fp16 évalue les générations, et les paramètres du LoRA sont optimisés via la divergence de Kullback–Leibler inverse avec un objectif de probabilité top- $k$ sensible aux queues de distribution.

Résultats

Débit et mémoire (Mac mini M4 Pro, 24 Go) :
- edge0-35b ( $K = 4$ ) : Atteint 20,4 tok/s en décodage en utilisant seulement 2,9 Gio de mémoire active maximale à partir d'un checkpoint int4 de 19,5 Go (contre 3,9 tok/s et 18,2 Gio en mode entièrement résident).
- edge0-8b ( $K = 8$ ) : Atteint 28,0 tok/s en décodage avec 1,5 Gio de mémoire de crête.
Accélération du préchargement (MacBook M2, 16 Go) :
- Par rapport à une diffusion purement réactive à la demande, le préchargement par prerouter à 1 token augmente la vitesse de décodage de +80 % à +84 % selon les largeurs de routage ( $K = 2$ : de 4,8 à 8,6 tok/s ; $K = 4$ : de 3,5 à 6,4 tok/s ; $K = 8$ : de 1,8 à 3,3 tok/s).
- Le temps de chargement bloquant sur le thread principal chute de manière spectaculaire (par exemple, de 244,0 ms à 101,9 ms par étape pour $K = 4$ ).
Qualité du modèle (OpenCompass) :
- Sur cinq benchmarks (AIME 2026, HumanEval, GPQA-Diamond, MMLU-Pro, IFBench), edge0-35b obtient une moyenne de 79,2 contre 83,2 pour le modèle de base fp16 Qwen3.6 (écart de 3,9 points).
- edge0-8b atteint 69,9 contre 72,7 pour le modèle de base Ling 3.0 tiny fp16 (écart de 2,8 points).

Limites et questions ouvertes

Concurrence de requêtes : Edge0 fonctionne de manière sérialisée en FIFO pour des requêtes uniques ; le traitement par lots modifie le working set des experts et la dynamique du cache d'une façon non prise en compte par les profils à flux unique.
Surcoût de construction du graphe sur CPU : La vitesse de décodage est limitée par le surcoût du framework côté CPU (~44 ms/étape dans la construction de graphe MLX sur 40 couches), ce que les optimisations de SSD ne peuvent pas corriger.
Baisse de qualité sur le raisonnement : La perte de performance est concentrée sur les benchmarks de raisonnement complexe (par exemple, une baisse de 6,1 points sur AIME pour le 35B et de 10,0 points pour le 8B).
Dépendance matérielle : Les gains relatifs de débit apportés par le prerouter dépendent de la latence du disque et s'amenuisent lorsque les caches de poids sont chauds ou sur des supports de stockage ultra-rapides.

Termes clés

Edge0: Un moteur d'inférence MoE en streaming qui sert des modèles depuis un SSD en superposant les lectures et les calculs via la prédiction de routage et compense la perte de qualité par un LoRA de récupération non fusionné.
prerouter: Une tête neuronale légère par couche qui prédit les décisions de routage de la couche N+1 avec un token d'avance à partir de l'état caché de la couche N et de l'historique récent de routage.
prediction-as-routing: Le principe de conception selon lequel l'ensemble d'experts prédit par le prerouter est directement exécuté comme routage effectif de la couche lors du décodage, éliminant les abandons de tokens ou les pénalités de prédiction manquée à l'exécution.
recovery LoRA: Un adaptateur de bas rang entraîné par distillation sur la trajectoire pré-routée de l'étudiant et exécuté comme branche parallèle non fusionnée pour compenser les erreurs dues à la quantification int4 et à l'approximation du routage.
staged decode: Un mode d'exécution des experts utilisant un double tampon à emplacements fixes où les indices des experts sont mappés via une table de slots sur le GPU pour éviter les allocations mémoire à chaque étape.
incremental sticky-slot stack (incr_stack): Un mécanisme qui met à jour sur place les tenseurs d'experts empilés et persistants, en ne réécrivant que les emplacements d'experts modifiés à chaque étape plutôt qu'en réassemblant tous les tenseurs.
unmerged delta: L'exécution d'un adaptateur LoRA en parallèle avec les poids de base quantifiés plutôt que leur fusion dans ces derniers, empêchant la perte de précision de l'adaptation lors d'une re-quantification.
memory wall: L'écart croissant entre le débit des processeurs et la bande passante ou la capacité de la DRAM, qui, dans le service de LLM locaux, se répartit entre l'état dynamique du cache KV et le stockage statique des poids du modèle.

Questions

Pourquoi Edge0 utilise-t-il la prédiction comme routage au lieu de traiter les prédictions comme de simples indications ?

Traiter les prédictions comme de simples indications impose des mécanismes de secours à l'exécution ou l'abandon de tokens en cas d'erreur. Utiliser directement la prédiction comme routage garantit que l'ensemble des experts préparés correspond systématiquement à l'ensemble routé, reportant le coût de l'approximation sur l'entraînement hors ligne où la distillation restaure la qualité.

Pourquoi le LoRA de récupération est-il servi de manière non fusionnée plutôt qu'intégré dans les poids 4 bits ?

Les deltas de poids LoRA ont une valeur RMS d'environ 10^-3, inférieure au pas de quantification 4 bits. La fusion suivie d'une re-quantification détruit la majeure partie de l'effet de l'adaptateur, ne conservant que 18 % de rétention au niveau des logits, tandis qu'un delta parallèle non fusionné préserve une fidélité de 100 %.

Pourquoi Edge0 prédit-il avec un token d'avance plutôt qu'au sein du même token ?

Prédire au sein du même token nécessite des vidages de pipeline GPU (flushes) par couche qui prennent entre 30 et 100 ms par étape, soit plus que le temps de chargement qu'ils cherchent à masquer. Anticiper d'un token complet permet de prédire toutes les couches préparées en un seul flush groupé, offrant tout le temps de calcul d'une étape de décodage pour masquer les E/S du SSD.

Comment Edge0 décode-t-il plus rapidement qu'une base de référence entièrement chargée en mémoire sur une machine de 24 Go ?

Un modèle de 19,5 Go entièrement résident consomme 18,2 Gio de RAM, ne laissant pratiquement aucun espace pour le système d'exploitation et le cache KV, ce qui provoque une pagination mémoire massive et fait chuter le décodage à 3,9 tok/s. Edge0 diffuse les experts depuis le SSD avec ~2,9 Gio de mémoire active maximale, évitant la saturation de la RAM et atteignant 20,4 tok/s.

La présentation de cette page a été rédigée par une IA à partir de l'article et peut contenir des erreurs. Consultez l'article pour les détails.

Lire avec traduction et IA