Architecture du wafer et contraintes de mémoire
Cerebras utilise le wafer WSE-3T, chaque puce intégrant 44 GB de SRAM sur‑silicon. Cette capacité fixe impose une limite stricte sur le nombre de paramètres pouvant être stockés localement. En quantifiant les poids à 4 bits, les 3 000 000 000 000 paramètres du modèle Kimi K3 occupent environ 16 GB, soit bien inférieur à la capacité du wafer. Cette marge permet de placer une couche complète sur un wafer sans fractionner les poids, ce qui élimine les transferts de données de poids entre puces.
Cartographie du modèle Kimi K3 sur les wafers
Kimi K3 comporte 69 couches à attention linéaire et 24 couches denses à attention complète. Les couches denses nécessitent, en plus des poids, un cache KV de 8 bits par jeton. En supposant une requête moyenne de 256 k tokens, le cache occupe une part importante de la SRAM restante. Ainsi, chaque wafer dédié à une couche dense réserve le reste de sa mémoire (≈28 GB) au cache KV, tandis que les wafers contenant des couches linéaires peuvent accueillir deux couches (2 × 16 GB = 32 GB) car elles ne stockent qu’un état fixe par requête. Le calcul donne 35 wafers pour les couches linéaires (69 ÷ 2 ≈ 35) et 24 wafers supplémentaires pour les couches denses, soit un total de 59 wafers nécessaires pour héberger l’ensemble du modèle.
Le compilateur Cerebras traduit cette répartition en une topologie de régions de wafers, chaque région correspondant à une ou plusieurs couches. Les poids restent immobiles sur chaque wafer, ce qui minimise les déplacements de données pendant l’inférence.
Pipeline d'exécution et gestion de la concurrence
L’exécution s’organise sous forme de pipeline : les activations, et non les poids, traversent les wafers via RDMA sur Ethernet, avec un délai d’environ 2 µs par saut. Chaque wafer agit comme une station d’assemblage, traitant les activations d’une couche avant de les transmettre à la suivante. Cette architecture permet de garder le pipeline constamment rempli. Au lieu d’attendre la fin d’une requête, le système injecte de nouvelles requêtes dès qu’une station devient disponible, créant ainsi un flux continu de micro‑lots.
Le nombre maximal de requêtes simultanées dépend de la capacité du cache KV sur les wafers denses. En soustrayant la taille d’une couche dense (≈16 GB) de la capacité totale du wafer (44 GB), il reste ≈28 GB pour le cache. En divisant cet espace par la consommation d’un cache de 256 k tokens (8 bits × 256 k ≈ 256 KB), on obtient plusieurs dizaines de requêtes concurrentes par wafer dense. Le nombre total de requêtes en vol est alors proportionnel au nombre de stages du pipeline multiplié par la taille du micro‑lot, offrant un débit de 1 000 à 2 000 tokens / s pour le modèle complet.
Cette approche se généralise aux modèles supérieurs à 3 T de paramètres : on calcule d’abord la mémoire requise pour les poids et le cache, on ajoute des wafers supplémentaires jusqu’à ce que chaque couche trouve sa place, puis on relie les wafers en pipeline. Chaque wafer additionnel augmente à la fois la capacité de calcul et la mémoire disponible, permettant d’étendre la taille du modèle au‑delà de la capacité d’un seul wafer.