Mamba-3 : SSM optimise inference, benchmarks et architecture

Du training-first a l'inference-first

Mamba-2, sorti mi-2024, avait fait un choix clair : simplifier le mecanisme SSM (State Space Model) pour accelerer l'entrainement, quitte a sacrifier l'expressivite du modele. La matrice de transition diagonale avait ete reduite a un scalaire multiplie par l'identite. Resultat : 2 a 8x de gain en entrainement, mais un decodage token par token quasi entierement memory-bound.

Mamba-3, developpe par Together AI, Carnegie Mellon, Princeton et Cartesia AI, inverse la logique. Le postulat : en 2026, l'inference est le goulot d'etranglement. Les workflows agentiques (Codex, Claude Code), le RLVR (Reinforcement Learning with Verifiable Rewards) et le post-training consomment des volumes massifs de tokens generes. Optimiser pour l'inference n'est plus optionnel.

Trois axes d'amelioration

Discretisation exponentielle-trapezoidale

Mamba-3 introduit une nouvelle formule de recurrence derivee d'un schema de discretisation exponentiel-trapezoidal. Contrairement a la discretisation simplifiee de Mamba-2, cette approche permet au modele d'appliquer implicitement une convolution sur l'entree vers l'etat cache. Consequence directe : la convolution causale courte (short conv) qui etait un composant obligatoire de Mamba-1 et Mamba-2 devient superflue. Les tests montrent que sa suppression n'impacte pas les performances sur les taches de retrieval reelles (NIAH), meme si les taches synthetiques comme MQAR deviennent legerement plus difficiles.

Etats complexes via RoPE

La deuxieme amelioration porte sur le suivi d'etat. Mamba-3 utilise des nombres complexes pour modeliser le systeme SSM, interpretes comme des rotations via RoPE (Rotary Position Embedding). Cette approche elargit l'espace de representation sans necessiter de reimplementer les noyaux de calcul depuis zero. Les transitions complexes permettent de capturer des dependances temporelles plus fines qu'avec des valeurs reelles seules.

Mode MIMO (Multi-Input, Multi-Output)

La troisieme innovation est le passage d'un SSM SISO (Single-Input, Single-Output) a un SSM MIMO. Concretement, plusieurs SSM tournent en parallele avec des matrices B et C etendues. Le point crucial : le mode MIMO augmente le temps d'entrainement mais pas la latence de decodage. Cette asymetrie s'explique par la nature respectivement compute-bound (entrainement) et memory-bound (inference) des deux phases. Le mode MIMO ajoute du calcul (qui reste gratuit pendant l'inference car le GPU est deja bloque sur les acces memoire) sans ajouter de transferts memoire.

Architecture revisitee

Au-dela du coeur SSM, Mamba-3 modernise son architecture globale. L'ajout de QKNorm (ou BCNorm en terminologie SSM) stabilise l'entrainement et aligne Mamba-3 avec les pratiques standard des Transformers et de Gated DeltaNet. Les couches MLP entrelacees suivent la convention Transformer classique.

La suppression de la short conv merite attention. Historiquement introduite par H3 (inspiree des induction heads d'Anthropic) puis popularisee par Mamba-1, la convolution causale courte etait consideree comme indispensable pour les capacites de retrieval. Mamba-3 demontre que la combinaison des biais sur B et C apres BCNorm et de la nouvelle recurrence reproduit empiriquement le meme effet.

Benchmarks

A l'echelle 1,5B parametres, Mamba-3 SISO surpasse Mamba-2, Gated DeltaNet et Llama-3.2-1B (Transformer) en latence totale prefill+decode sur toutes les longueurs de sequence testees.

Sur les evaluations de language modeling downstream (taches de comprehension, raisonnement), Mamba-3 SISO surpasse Mamba-2 a architecture equivalente. La variante MIMO ajoute plus d'un point de pourcentage de precision a l'echelle 1B, avec un cout d'entrainement accru mais sans impact sur la latence d'inference.

L'ecart avec les Transformers se resserre significativement, en particulier sur les taches ou la longueur du contexte est un facteur (ou les Transformers souffrent de la croissance quadratique du KV cache).

Implementation et noyaux

L'equipe a publie l'ensemble des noyaux de calcul en open source. L'implementation combine trois frameworks : Triton pour le prototypage rapide, TileLang pour les operations de tiling et CuTe DSL pour l'optimisation hardware bas niveau. Cette approche hybride permet d'atteindre une utilisation optimale des tensor cores tout en gardant le code maintenable.

Implications pour l'ecosysteme

Mamba-3 ne pretend pas remplacer les Transformers sur toutes les taches. Son avantage se manifeste principalement sur les scenarios inference-heavy : agents de code, chatbots a haute volumetrie, generation de rollouts pour le RLVR.

La combinaison d'un etat fixe (pas de KV cache croissant), d'une complexite lineaire et d'un decodage compute-dense (grace au MIMO) positionne les SSM comme une alternative credible pour les cas d'usage ou le cout par token genere est le facteur limitant.

La publication en open source complete, incluant poids, code d'entrainement et noyaux optimises, permet a la communaute d'evaluer et d'integrer Mamba-3 dans les pipelines de production existants.