Amélioration des performances du processeur du stimulateur de jetons unigramme

Nous avons réimplémenté notre stimulateur de jetons Unigram à partir de zéro dans le cadre d'un projet ciblé sur les performances.

AuteursPerplexity Research

Nous avons réimplémenté notre stimulateur de jetons Unigram à partir de zéro dans le cadre d'un projet ciblé sur les performances. Pour des longueurs d'entrée de production, le nouvel encodeur réduit la latence p50 d'environ 5× par rapport au package de stimulateurs de jetons de Hugging Face, d'environ 2× par rapport à SentencePiece (C++) et d'environ 1,5× par rapport au stimulateur de jetons d'IREE (C), sans aucune allocation de tas à l'état stable. En production, cela a réduit l'utilisation du processeur dans notre pile d'inférence de 5 à 6× et a retranché deux chiffres de millisecondes de la latence du réclasseur.

Motivation

Le coût d'inférence des LLM est généralement présenté comme un travail de processeur graphique (GPU) : caches KV, noyaux d'attention, routage d'experts. Mais une part significative du trafic de production s'exécute sur des modèles beaucoup plus petits, de deux à trois ordres de grandeur plus petits que les transformateurs frontaliers. Les modèles d'intégration, les classificateurs et les réclasseurs s'exécutent toujours sur des GPU, mais la passe avant est suffisamment courte pour que tout ce qui l'entoure commence à compter.

Le cas le plus évident est celui d'un réclasseur évaluant des centaines de documents candidats par requête. Avec un petit modèle, le calcul du GPU se termine généralement en un seul chiffre de millisecondes, mais chaque entrée doit toujours passer d'abord par la tokenisation côté processeur (CPU), ce qui transforme le texte en ID de vocabulaire. Lorsque le lot est important, cela peut représenter une fraction significative de la latence totale de la requête.

Ce travail cible XLM-RoBERTa, un vocabulaire Unigram de 250 000 jetons entraîné avec SentencePiece. Les encodeurs de la famille RoBERTa finement ajustés sont un choix courant pour les tâches de classement, de recherche et de similarité. Nous rapportons la latence par encodage à un seul fil d'exécution partout ; en production, nous parallélisons sur les cœurs à l'aide de rayon.

Tout au long de cet article, nous traitons le package de stimulateurs de jetons de Hugging Face comme la référence par rapport à laquelle nous effectuons des benchmarks, car il s'agit du stimulateur de jetons Rust par défaut vers lequel la plupart des équipes se tournent. SentencePiece (C++) et IREE (C) sont inclus pour la comparaison en tant que deux autres principales implémentations de production.

L'algorithme de Viterbi

La tokenisation Unigram, introduite par Kudo en 2018 (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) et livrée dans SentencePiece, formule la segmentation comme un problème de chemin le plus probable : chaque jeton de vocabulaire a une log-probabilité apprise, et l'encodeur choisit la segmentation dont les scores de jetons s'additionnent pour donner la valeur la plus élevée. Le problème est résolu avec l'algorithme de Viterbi (1967), les positions des octets servant de couches de graphe et les jetons de vocabulaire d'arêtes couvrant une plage d'octets contiguë. La récurrence est une programmation dynamique unidimensionnelle sur les positions des octets :

python
best = [-inf] * (n + 1) best[0] = 0 back = [None] * (n + 1) for i in 0..n: if best[i] == -inf: continue if byte i is a UTF-8 continuation byte: continue walk the vocabulary trie from byte position i: for each matching token t ending at position j: cand = best[i] + score(t) if cand > best[j]: best[j] = cand back[j] = (i, token_id(t)) if no token matched, fall back to <unk> backtrack from position n using back[] to recover token IDs

La boucle externe est linéaire par rapport à la longueur de l'entrée, limitée à chaque position par le jeton de vocabulaire le plus long (48 octets UTF-8 pour XLM-RoBERTa). Le parcours de l'arbre trie interne est le chemin critique : sur une entrée de 16k jetons, il exécute des centaines de milliers de transitions et retire des dizaines de millions d'instructions par encodage.

Implémentation de référence

Le stimulateur de jetons v0.22 de Hugging Face implémente la récurrence directement, avec une table de hachage pour la recherche inverse et des enfants de table de hachage à chaque nœud de l'arbre trie :

rust
vocab: Vec<(String, f64)> // indexed by token ID token_to_ids: AHashMap<String, u32> // reverse lookup trie: Trie<u8> // byte-level prefix trie struct Node<Label> { is_leaf: bool, children: AHashMap<Label, Node<Label>>, // heap-allocated map per node }

Et la boucle interne (simplifiée) – chaque correspondance de recherche de préfixe alloue une nouvelle chaîne, puis sonde la table de hachage :

rust
fn encode_optimized(&self, sentence: &str) -> Result<Vec<String>> { // alloc let mut best = vec![BestPathNode::default(); sentence.len() + 1]; let mut starts_at = 0; while starts_at < sentence.len() { for tok_bytes in self.trie.common_prefix_search(sentence.bytes().skip(starts_at)) { // alloc per match let token: String = String::from_utf8(tok_bytes).unwrap(); // hash lookup let id = self.token_to_ids.get(&token).unwrap(); let score = self.vocab[*id as usize].1; // DP update at best[starts_at + tok_bytes.len()] // ... } starts_at += mblen; } // backtrace into Vec<String>, one allocation per output token }

Le code est du Rust lisible, mais beaucoup de travail se produit qui est invisible au niveau de la source. La section suivante détaille où va le temps.

Une note sur le prétraitement

Avant que Viterbi ne s'exécute, l'entrée passe par la normalisation Unicode (NFKC) et le marquage des limites de mots. Le profilage place le prétraitement à moins de 5 % du temps total d'encodage dans nos scénarios, nous nous concentrons donc sur la passe avant de Viterbi pour le reste de l'article.

Où l'implémentation de référence passe son temps

Il y a deux façons d'accélérer le code du processeur : exécuter moins d'instructions ou obtenir un IPC plus élevé. Cachegrind mesure le premier, perf stat mesure le second. Chaque optimisation de cet article se rapporte à l'un ou à l'autre.

Les processeurs subissent des temps d'arrêt pour quatre raisons : les défauts de cache (L1d ~4 cycles, un défaut de cache L3 vers la DRAM se compte en centaines), les chaînes de chargement dépendantes que le moteur d'exécution désordonnée ne peut pas paralléliser, les mauvaises prédictions de branchement (~15-20 cycles vidés par mauvaise estimation) et les appels d'allocateur. L'implémentation du moteur unigramme de référence présente les quatre :

Goulot d'étranglement

Mécanisme

Impact mesuré

Allocation par correspondance

String::from_utf8 sur les octets correspondants plus une recherche AHashMap par correspondance dans l'arbre trie ; l'ID de jeton est déjà dans la feuille, mais la référence ne l'y stocke pas

7 295 allocations à 512 jetons, 14 542 à 1k, 299 171 à 16k

Recherche de pointeur par octet

Une AHashMap à chaque nœud de l'arbre trie, hachant l'octet suivant et déréférençant deux pointeurs par étape ; la structure Node regroupe également des champs froids comme is_leaf, les ramenant dans le cache à chaque visite

Le taux de défauts L1d reste bas (~2,5 % à 1k) parce que les pages chaudes de l'allocateur résident dans L1 ; le coût apparaît sous forme de latence de chargement dépendant

Excès de pagination (thrashing) L2 sur les entrées longues

La table de programmation dynamique, le tampon de travail des préfixes et la liste des jetons de sortie sont tous nouvellement alloués à chaque appel

Le taux de défauts L2 grimpe de 12 % à 128 jetons à 50 % à 16k à mesure que l'empreinte d'allocation cumulative dépasse le L2 de 2 Mo par cœur

Référence des goulots d'étranglement, des mécanismes et des impacts mesurés du tokeniseur.

Le coût dominant est l'allocation par correspondance. À 1k jetons, l'implémentation de référence retire (achève) environ 7M instructions par encodage à un IPC de 2,8, la plupart dans des appels d'allocateur et des conversions de chaînes. Les défauts L2 ne comptent qu'après ~1k jetons, une fois que l'empreinte d'allocation cumulative dépasse le cache par cœur.

Latence p50 et p99 du tokeniseur de référence selon la taille d’entrée.
Taux de ratés du cache L2 du tokeniseur de référence selon la taille d’entrée.
Octets alloués par encodage par le tokeniseur de référence selon la taille d’entrée.

L'allocation par jeton est constante : environ 2 Ko et ~18 allocations par jeton, quelle que soit la taille de l'entrée. La non-linéarité de la latence provient du L2 : les octets cumulés dépassent 2 Mo autour de 1k jetons, le taux de défauts grimpe de 8 % à ~50 %, et la courbe p99 fléchit pour correspondre.

Comment nous avons mesuré

Nous comparons trois stimulateurs de jetons de production, exécutant tous le même algorithme Unigram sur le même modèle XLM-RoBERTa :

Moteur

Langage

Arbre trie

Hugging Face

Rust

Arbre trie de préfixes à enfants HashMap

SentencePiece

C++

Arbre trie à double tableau Darts

Stimulateur de jetons IREE

C

Arbre trie à double tableau Darts

Implémentations de tokeniseur, langages et représentations de trie utilisés dans la comparaison.

Nous lions SentencePiece et IREE par le biais d'une FFI C/C++ directe plutôt que par leurs wrappers Rust ; les wrappers ajoutent des allocations par appel et des conversions de chaînes qui ne font pas partie du coût de la bibliothèque sous-jacente.

Tous les moteurs sont construits avec des optimisations complètes (O3, ciblage natif, optimisation au moment de l'édition de liens). Les mesures sont exécutées en étant épinglées à un cœur silencieux sur un Intel Xeon Platinum 8488C. Pour chaque moteur et taille d'entrée, nous collectons la latence horloge (p50 et p99, 10 000 itérations après 1 000 réchauffements), les instructions retirées via cachegrind de Valgrind, l'IPC et les taux de défauts de cache provenant des compteurs matériels de perf stat, ainsi que les allocations de tas comptées par un wrapper d'allocateur global autour de la fenêtre mesurée.

À 514 jetons (les +2 par rapport à 512 correspondent à l'injection BOS/EOS) :

Implémentation

Latence p50

Instructions

IPC

Allocations

Hugging Face

349 µs

3,60M

2,89

7 295

SentencePiece

128 µs

1,83M

3,39

1 559

IREE

112 µs

2,28M

3,06

1

Latence du tokeniseur, nombre d'instructions, IPC et allocations à 514 jetons.

SP et IREE atteignent environ un tiers de la latence de la référence. Ils retirent environ deux fois moins d'instructions, s'exécutent à un IPC plus élevé et allouent un à deux ordres de grandeur de moins.

Pourquoi SentencePiece et IREE sont plus rapides : L'arbre trie à double tableau

SentencePiece et IREE utilisent tous deux l'arbre trie à double tableau introduit par Aoe en 1989 (Aoe, 1989 ; Karoonboonyanan). L'arbre trie est encodé dans deux tableaux d'entiers plats, base et check :

cpp
int base[N]; // base[node] = offset for this node's children int check[N]; // check[child] == node validates the parent-child link

Une recherche d'enfant est :

cpp
int next = base[node] + byte; if (check[next] == node) { /* valid transition; child node = next */ }

Chaque étape du parcours de l'arbre trie ne nécessite que deux lectures de tableau, une addition d'entiers et une comparaison, sans hachage, sans recherche de pointeur et sans allocation de tas par nœud. La construction place chaque nœud de sorte que ses enfants tombent sur des emplacements libres dans le tableau de validation, permettant à plusieurs parents de partager des régions tant que leurs enfants diffèrent sur l'octet qu'ils suivent.

Pour le vocabulaire de 250K d'XLM-RoBERTa, l'arbre trie entier tient dans ~9 Mo de mémoire contiguë (~780K emplacements × 12 octets pour le décalage, la validation et l'ID de jeton par emplacement). L'ensemble de travail chaud d'un seul encodage est de l'ordre de 100 Ko, confortablement dans le L2, et le pré-récupérateur matériel se verrouille sur le modèle d'accès régulier.

La référence, en revanche, conserve une table de hachage des enfants par nœud, ce qui coûte environ quatre charges dépendantes par étape d'octet. Hacher l'octet, charger le seau, suivre un pointeur vers l'entrée du seau, suivre un autre pointeur vers le nœud enfant sur le tas, puis consulter sa propre table de hachage pour l'octet suivant. Les objets du tas se dispersent et la plupart des charges ratent le cache.

SP et IREE stockent tous deux l'ID du jeton directement dans le nœud de l'arbre trie. L'implémentation de l'arbre trie de référence renvoie uniquement les octets correspondants, de sorte que l'encodeur doit allouer une nouvelle chaîne pour chaque correspondance, la rechercher dans une table de hachage secondaire pour récupérer l'ID du jeton et lire le score à partir d'un vecteur séparé. SP renvoie l'ID en tant que valeur stockée de l'arbre trie et lit le score à partir d'un tableau de nombres à virgule flottante plat :

cpp
// unigram_model.cc:567 const size_t num_nodes = trie_->commonPrefixSearch( begin, trie_results.data(), trie_results.size(), static_cast<int>(end - begin)); for (size_t k = 0; k < num_nodes; ++k) { const int id = trie_results[k].value; // already in the trie const float score = GetScoreInlined(id); // flat float[] indexed by id Lattice::Node *node = lattice->Insert(begin_pos, length); node->id = id; node->score = score; }

IREE utilise une API de curseur avec état au lieu de renvoyer toutes les correspondances en un seul appel, mais le modèle de coût est le même : les ID de jetons et les scores sont accessibles en tant qu'indices de tableaux plats, sans allocation ni hachage lors du parcours de l'arbre trie lui-même. IREE n'alloue qu'un seul tampon de sortie par encodage ; SP alloue en outre un nœud de réseau (lattice) par correspondance pendant la programmation dynamique, totalisant environ 1 559 allocations à 514 jetons.

Référence : Un portage sans allocation de la référence

Avant de modifier l'arbre trie, nous isolons le coût de l'algorithme du coût de l'implémentation : à quelle vitesse la référence s'exécute-t-elle une fois le travail inutile supprimé ?

La référence conserve l'arbre trie à enfants de table de hachage de la référence, mais apporte deux modifications. Une structure de travail (scratch) appartenant à l'appelant contient la table de programmation dynamique, les octets prétraités et les ID de jetons de sortie, de sorte que le chemin critique ne fait aucune allocation après le réchauffement. L'ID de jeton et le score sont stockés directement dans le nœud de l'arbre trie, ce qui supprime la conversion de chaîne par correspondance et la recherche dans la table de hachage secondaire.

Boucle interne :

rust
for pos in 0..n { if scratch.prep[pos] & 0xC0 == 0x80 { continue; } let s0 = scratch.best_score[pos]; if s0 == f32::NEG_INFINITY { continue; } let mut node = 0u32; for end in pos..n { let byte = scratch.prep[end]; match self.nodes[node as usize].children.get(&byte) { Some(&child) => { node = child; let n = &self.nodes[node as usize]; if n.is_terminal { let total = s0 + n.score; if total > scratch.best_score[end + 1] { scratch.best_score[end + 1] = total; scratch.best_start[end + 1] = pos as u32; scratch.best_id[end + 1] = n.token_id; } } } None => break, } } }

Il s'agit du même algorithme, avec la même forme d'arbre trie et le même modèle d'accès, mais simplement sans allocations et sans recherches redondantes. Résultats à 514 jetons :

Moteur

p50

Instructions

IPC

Allocations

Hugging Face

326 µs

3,66M

2,86

7 295

SentencePiece

127 µs

1,81M

3,55

1 559

IREE

111 µs

1,58M

3,26

1

De base

156 µs

1,53M

2,62

0

La base de référence sans allocation comparée aux implémentations de tokeniseur existantes à 514 jetons.

La suppression de la matérialisation de chaînes par correspondance et de la recherche dans la table de hachage secondaire réduit les instructions de 2,4× par rapport à la référence. L'IPC chute avec elles (2,86 → 2,62) car la recherche de pointeur dans la table de hachage, auparavant masquée par le travail de l'allocateur, est désormais le coût dominant par octet.

Latence p50 et p99 du tokeniseur de référence et de la base de référence sans allocation selon la taille d’entrée.
Taux de ratés du cache L2 du tokeniseur de référence et de la base de référence sans allocation selon la taille d’entrée.

La référence réduit la latence p50 de 2 à 3× par rapport à la référence à chaque taille, et son taux de défauts L2 reste proche de 2 à 3 % jusqu'à 2k jetons, contre 8 à 48 % pour la référence. Le coût restant est la recherche de pointeur dans la HashMap, que la section suivante aborde.

Optimisation 1 : Arbre trie à double tableau sans abstractions supplémentaires

L'arbre trie à table de hachage est une recherche de pointeur. SP et IREE résolvent ce problème avec un arbre trie à double tableau (voir la section précédente), et nous empruntons la même structure de données. La différence est que SP et IREE sont des bibliothèques polyvalentes prenant en charge l'entrée en flux continus, les décalages de sortie et de multiples modes d'encodage ; cette flexibilité leur coûte de la comptabilité supplémentaire autour de l'arbre trie (un réseau dans SP, un pipeline à huit étapes dans IREE). Nous n'avons besoin de rien de tout cela, nous inlinons donc l'arbre trie directement dans la boucle de Viterbi et éliminons le reste.

rust
fn child(&self, node: u32, byte: u8) -> Option<u32> { let base = self.base[node as usize]; let next = base + byte as i32; if self.check[next as usize] == node as i32 { Some(next as u32) } else { None } }

Une étape d'octet coûte deux lectures de tableau, une addition d'entiers et une comparaison. L'arbre trie entier fait environ 9 Mo, mais chaque encodage individuel n'en parcourt que ~100 Ko, de sorte que les données chaudes restent dans le L2.

Moteur

p50

p99

Instructions

Défauts L2 %

Défauts L3 %

Hugging Face

327 µs

3 340 µs

3,66M

8,17 %

14,66 %

SentencePiece

127 µs

1 140 µs

1,82M

4,74 %

4,07 %

IREE

111 µs

132 µs

1,61M

5,82 %

3,87 %

De base

155 µs

1 170 µs

1,53M

2,59 %

16,66 %

+ Darts

68 µs

80 µs

1,10M

6,45 %

11,42 %

Latence du tokeniseur, instructions et taux de ratés du cache après l’adoption du trie à tableau double.

Le temps d'horloge chute de 2,3× par rapport à la référence et de 4,8× par rapport à la référence initiale. Les instructions chutent de 28 % supplémentaires par rapport à la référence. La suppression du réseau et du pipeline nous permet de retirer 1,10M instructions par encodage contre 1,82M pour SP et 1,61M pour IREE avec la même forme d'arbre trie ; l'écart correspond à la comptabilité qu'ils conservent pour les API polyvalentes dont nous n'avons pas besoin.

Optimisation 2 : bitmap et empaquetage en ligne

L'arbre trie Darts de l'optimisation 1 nécessite deux charges dépendantes par étape d'octet. D'abord le décalage de base du parent, puis l'entrée check au niveau de l'emplacement de l'enfant candidat pour confirmer que le parent le possède. L'entrée check valide uniquement la propriété. La même question peut recevoir une réponse par un bitmap par nœud qui enregistre quels octets ont des enfants. Par la construction de l'arbre trie à double tableau, le bit pour l'octet b dans le bitmap du parent est défini exactement lorsque l'emplacement candidat pour cet octet est possédé par le parent. Le bitmap et le tableau check transportent la même information, nous pouvons donc supprimer entièrement le tableau check de la structure d'exécution.

La recherche dans le bitmap se compile en un seul test de bit contre un mot de 64 bits, ce qui s'exécute en un cycle environ. L'instruction POPCNT associée, qui compte tous les bits définis dans un registre, est ce que d'autres conceptions d'arbres trie indexés par bitmap utilisent pour calculer directement un décalage d'enfant à partir du bitmap. Nous n'en avons pas besoin car la construction Darts stocke déjà ce décalage dans le champ base du nœud.

rust
#[repr(C, align(64))] struct NodeHot { bitmap: [u64; 4], // 32 B — "which bytes have children" base: i32, // 4 B — child slot offset token_id: i32, // 4 B — vocab id (-1 if not terminal) score: f64, // 8 B — Viterbi log-prob _pad: [u8; 16], // 16 B — pad to one cache line } // Lookup: one cache-line load, one bit test, one add. fn child(&self, parent: u32, byte: u8) -> Option<u32> { let h = &self.node_hot[parent as usize]; if !h.bitmap.test(byte) { return None; } Some(h.base.wrapping_add(byte as i32) as u32) }

Tant qu'on y est, nous empaquetons les quatre champs par nœud (bitmap, base, ID de jeton et score) dans une seule ligne de cache de 64 octets. Une seule visite de nœud extrait désormais tout : le score et l'ID de jeton dont nous avons besoin aux nœuds terminaux, ainsi que le bitmap et la base pour l'étape suivante. Le tableau check n'est utilisé que pendant la construction et est supprimé de la disposition d'exécution. L'indexation du chemin critique n'est pas vérifiée car la construction de l'arbre trie et la structure de la boucle garantissent déjà les limites. Le compromis est l'Espace : à 64 octets par nœud × ~780K emplacements, l'arbre trie passe de ~9 Mo à ~50 Mo.

Résultats à 514 jetons :

Moteur

p50

p99

Instructions

Accès L2

Défauts L2 %

Défauts L3 %

Hugging Face

327 µs

3 340 µs

3,66M

15,0K

8,17 %

14,66 %

SentencePiece

127 µs

1 140 µs

1,82M

7,5K

4,74 %

4,07 %

IREE

111 µs

132 µs

1,61M

4,7K

5,82 %

3,87 %

De base

155 µs

1 170 µs

1,53M

6,1K

2,59 %

16,66 %

+ Darts

68 µs

80 µs

1,10M

4,6K

6,45 %

11,42 %

+ Bitmap et empaquetage en ligne

65 µs

77 µs

1,06M

1,8K

20,53 %

15,01 %

Résultats du tokeniseur à 514 jetons après l’ajout du bitmap et de l’empaquetage en ligne.

Le temps d'horloge chute systématiquement de 4,5 % supplémentaires à 514 jetons et augmente avec la longueur de l'entrée, atteignant 6 à 7 % à 4k et 16k. La suppression du tableau check transforme l'étape de l'arbre trie par octet de deux charges de ligne de cache à une seule, ce qui représente 2,5× de moins d'accès L2 par encodage (1,8K contre 4,6K pour les Darts ordinaires). Le taux de défauts bondit de 6 % à 20 %, mais la majeure partie de cela est due au dénominateur : avec deux fois moins d'accès, les défauts que nous avons représentent une plus grande part.

Optimisation 3 : pages de grande taille pour le trie

Les données de l'arbre trie font environ 50 Mo (780 000 nœuds de 64 octets chacun). Sur un processus Linux par défaut, cette mémoire est prise en charge par des pages de 4 Ko, de sorte que l'arbre trie complet s'étend sur environ 12 000 pages virtuelles distinctes, bien au-delà de ce que le tampon de traduction anticipée (TLB) des données de premier niveau du processeur peut contenir (Sapphire Rapids a 96 entrées). Chaque étape de Viterbi touche un nœud différent, et sur un encodage de 512 jetons, nous consumons environ 9 000 cycles dans les parcours de table de pages, soit environ 3 % du budget par encodage.

La solution consiste à prendre en charge l'arbre trie avec des pages géantes de 2 Mo au lieu de pages de 4 Ko. Les mêmes 50 Mo s'étendent désormais sur 25 pages, tiennent entièrement dans le tampon de traduction anticipée, et les répertoires de tables de pages eux-mêmes cessent de concurrencer l'arbre trie pour le cache L2.

Linux offre deux façons de prendre en charge la mémoire avec des pages de 2 Mo. La méthode opportuniste est celle des pages géantes transparentes (THP) : votre programme demande au noyau de faire passer une région de mémoire de pages de 4 Ko à des pages de 2 Mo, et un thread d'arrière-plan du noyau essaie de le faire lorsqu'il a du temps et de la mémoire physique contiguë disponible. L'autre solution consiste à allouer l'arbre trie directement à partir d'un pool de pages géantes pré-réservé que le noyau met de côté au démarrage :

rust
let raw = unsafe { libc::mmap( std::ptr::null_mut(), alloc_len, libc::PROT_READ | libc::PROT_WRITE, libc::MAP_PRIVATE | libc::MAP_ANONYMOUS | libc::MAP_HUGETLB, -1, 0, ) };

Le pool doit exister sur l'hôte (vm.nr_hugepages défini au démarrage). Sur nos instances de production, 10 561 pages géantes sont réservées ; l'arbre trie en consomme 24. Après le mmap, /proc/<pid>/smaps confirme la région : KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Le gain est de 3 à 12 % sur le temps d'horloge à travers les différentes tailles d'entrée, de manière cohérente et proportionnelle à la longueur de l'ensemble de travail :

Jetons

Sans pages géantes

Avec pages géantes

Δ

130

15,7 µs

15,2 µs

−3,6 %

514

65,4 µs

63,1 µs

−3,4 %

1 026

129,5 µs

125,3 µs

−3,3 %

4 098

773 µs

679 µs

−12,0 %

16 386

3 808 µs

3 554 µs

−6,8 %

Latence du tokeniseur avec et sans grandes pages selon la taille d’entrée.

Le plus grand gain se situe à 4k jetons, où l'ensemble de travail est suffisamment grand pour que le trafic de la table de pages soit en concurrence active avec les données de l'arbre trie pour la bande passante L2. Au-delà de 4k, le gain diminue parce que les défauts L3 dominent et que les parcours de tables de pages représentent une part plus faible du temps total.

Comparaison finale

En rassemblant les trois optimisations, voici le moteur final mesuré par rapport à tous les autres stimulateurs de jetons RoBERTa que nous avons testés. Parallèlement aux binaires natifs C++ SentencePiece et C IREE, nous avons également évalué les packages de wrappers Rust prêts à l'emploi vers lesquels un utilisateur de production se tournerait, sentencepiece et iree-tokenizer. Le package sentencepiece mesure une latence supérieure de 1,6 à 1,9× par rapport au binaire C++ natif à chaque taille d'entrée, le surdébit provenant du wrapper qui alloue une nouvelle liste de morceaux de jetons, chacun détenant sa propre copie du texte du morceau, à chaque encodage. Le package iree-tokenizer dimensionne son tampon interne à chaque entrée, ce qui évite une allocation constante par encodage que le binaire IREE natif paie quelle que soit la taille de l'entrée, et cela compte le plus pour les entrées courtes où l'allocation dominerait autrement le travail de tokenisation réel. Aux longueurs d'entrée de production, le coût constant s'amortit et le surdébit par appel du wrapper lié à la copie des jetons dans Rust pousse sa latence au-dessus de celle du binaire natif.

Latence p50 et p99 par encodage selon les implémentations de tokeniseur et les tailles d’entrée.

Le moteur final exécute le même Viterbi sur le même vocabulaire de 250 000 jetons, n'alloue rien sur le chemin critique et produit une sortie identique en termes de jetons par rapport à la référence. Chaque étape d'octet est désormais un chargement unique de ligne de cache de 64 octets, un test de bitmap et une addition d'entiers, en baisse par rapport à une sonde de HashMap, un appel d'allocateur et plusieurs branches de cas particuliers dans l'implémentation de référence. Les instructions par encodage sont tombées de 3,66M à 1,04M, soit une baisse de 3,5×. Le temps d'horloge correspond à ce ratio pour les entrées courtes et s'élargit encore pour les entrées longues, où les allocations par jeton de la référence commencent à saturer L2 et L3.

Réflexions finales

En pratique, ce travail a réduit l'utilisation du processeur dans la couche de prétraitement de notre pile d'inférence de 5 à 6× et a retranché deux chiffres de millisecondes de la latence du réclasseur.

La plus grande partie de la réduction de la latence provient de modifications légères qui ont éliminé le travail gaspillé : suppression des allocations par encodage, remplacement de l'arbre trie à enfants HashMap par un double tableau, et suppression du surdébit de l'API générique. L'utilisation du bitmap par nœud existant comme vérification de validité et la prise en charge de l'arbre trie par des pages géantes de 2 Mo combinées ont représenté 10 à 20 % en plus du passage plus large à un arbre trie à double tableau, selon la taille de l'entrée.

Les packages de wrappers Rust prêts à l'emploi autour de SentencePiece et IREE illustrent le même point à l'envers. En plus des stimulateurs de jetons natifs C et C++, la limite FFI et les allocations par appel que les wrappers introduisent sont visibles dans les chiffres de latence. Le code source Rust pour l'encodeur Viterbi final est disponible sur github.com/perplexityai/pplx-garden.

Références

Stimulateurs de jetons Hugging Face — https://github.com/huggingface/tokenizers

SentencePiece — https://github.com/google/sentencepiece

Stimulateur de jetons IREE — https://github.com/iree-org/iree

Kudo, T. (2018). Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates. https://arxiv.org/abs/1804.10959

Viterbi, A. J. (1967). Error bounds for convolutional codes and an asymptotically optimum decoding algorithm. IEEE Transactions on Information Theory. https://en.wikipedia.org/wiki/Viterbi_algorithm

Aoe, J. (1989). An efficient digital search algorithm by using a double-array structure. IEEE Transactions on Software Engineering. https://ieeexplore.ieee.org/document/31365/

Karoonboonyanan, T. An Implementation of Double-Array Trie. https://linux.thai.net/~thep/datrie/datrie.html

POPCNT (x86) — https://www.felixcloutier.com/x86/popcnt

Pages géantes — https://docs.kernel.org/admin-guide/mm/hugetlbpage.html

Citation

BibTeX
@online{perplexity2026ImprovingUnigramTokenizerCpuPerformance,
  author = {{Perplexity} and Vladimir Zaytsev and Bex Brooker and Daisy Lee and Roshan Shankar},
  title = {{Improving unigram tokenizer CPU performance}},
  date = {2026-05-20},
  url = {https://www.perplexity.ai/hub/blog/improving-unigram-tokenizer-cpu-performance}
}