Améliorer la performance CPU du tokenizer Unigram

Nous avons réimplémenté notre tokenizer Unigram à partir de zéro dans le cadre d'un projet ciblé sur la performance. Pour les longueurs d'entrée de production, le nouvel encodeur réduit la latence p50 d'environ 5× par rapport au crate de tokenizers Hugging Face, d'environ 2× par rapport à SentencePiece (C++) et d'environ 1,5× par rapport au tokenizer d'IREE (C), sans aucune allocation de tas en état stationnaire. En production, cela a réduit l'utilisation du CPU dans notre pile d'inférence de 5 à 6× et a permis de gagner des dizaines de millisecondes sur la latence du reranker.

Le coût de l'inférence LLM est généralement présenté comme un travail GPU : caches KV, noyaux d'attention, routage d'experts. Mais une partie significative du trafic de production s'exécute sur des modèles beaucoup plus petits, deux à trois ordres de grandeur plus petits que les transformers de pointe. Les modèles d'embedding, les classificateurs et les rerankers fonctionnent 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 reranker notant des centaines de documents candidats par requête. Avec un petit modèle, le calcul GPU se termine généralement en quelques millisecondes, mais chaque entrée doit toujours passer par une tokenisation côté CPU, transformant 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 affinés sont un choix courant pour les tâches de classement, de recherche et de similarité. Nous rapportons la latence par encodage monothread tout au long ; en production, nous parallélisons sur les cœurs en utilisant rayon.

Tout au long de cet article, nous traitons le crate tokenizers de Hugging Face comme la référence par rapport à laquelle nous effectuons nos benchmarks, car il s'agit du tokenizer Rust par défaut que la plupart des équipes utilisent. SentencePiece (C++) et IREE (C) sont inclus pour comparaison en tant que deux autres implémentations de production majeures.

La tokenisation Unigram, introduite par Kudo en 2018 (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) et intégrée dans SentencePiece, définit 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 des jetons totalisent la valeur la plus élevée. Le problème est résolu avec l'algorithme de Viterbi (1967), avec les positions des octets comme couches de graphe et les jetons de vocabulaire comme arêtes couvrant une plage d'octets contiguë. La récurrence est une DP 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). La recherche dans le trie interne est le chemin critique : sur une entrée de 16k jetons, elle exécute des centaines de milliers de transitions et traite des dizaines de millions d'instructions par encodage.

Le tokenizers v0.22 de Hugging Face implémente directement la récurrence, avec une table de hachage pour la recherche inversée et des enfants de table de hachage à chaque nœud du 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 déroule de manière invisible au niveau de la source. La section suivante détaille où passe le temps.

Avant l'exécution de Viterbi, l'entrée passe par la normalisation Unicode (NFKC) et le marquage des limites de mots. Le profilage situe 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.

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

Les CPU stagnent sur quatre choses : les défauts de cache (L1d ~4 cycles, un défaut L3-vers-DRAM en représente des centaines), les chaînes de chargement dépendantes que le moteur de traitement hors ordre ne peut pas paralléliser, les prédictions de branchement erronées (~15-20 cycles vidés par mauvaise estimation) et les appels à l'allocateur. L'implémentation du moteur unigram de référence rencontre ces quatre problèmes :

Goulot d'étranglement	Mécanisme	Impact mesuré
Allocation par correspondance	String::from_utf8 sur les octets correspondants plus une recherche AHashMap par correspondance dans le trie ; l'ID du 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
Poursuite de pointeur par octet	Un AHashMap à chaque nœud du trie, hachant l'octet suivant et déréférençant deux pointeurs par étape ; le struct Node regroupe également des champs froids comme is_leaf, les faisant entrer dans le cache à chaque visite	Le taux de défaut L1d reste bas (~2,5 % à 1k) car les pages chaudes de l'allocateur vivent en L1 ; le coût apparaît comme la latence de chargement dépendant
Thrashing L2 sur les entrées longues	La table DP, le scratch de préfixe et la liste des jetons de sortie sont tous nouvellement alloués à chaque appel	Le taux de défaut L2 grimpe de 12 % à 128 jetons à 50 % à 16k car l'empreinte d'allocation cumulative dépasse le L2 de 2 Mo par cœur

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

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

Latence p50 et p99 du tokenizer de référence selon la taille d’entrée.

Taux de ratés du cache L2 du tokenizer de référence selon la taille d’entrée.

Octets alloués par encodage par le tokenizer de référence selon la taille d’entrée.

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

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

Moteur	Langage	Trie
Hugging Face	Rust	Trie de préfixes HashMap-children
SentencePiece	C++	Trie à double tableau Darts
Tokenizer IREE	C	Trie à double tableau Darts

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

Nous lions SentencePiece et IREE via FFI C/C++ direct plutôt que via 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 lors de l'édition de liens). Les mesures sont effectuées en épinglant un cœur calme sur un Intel Xeon Platinum 8488C. Pour chaque moteur et chaque taille d'entrée, nous collectons la latence horloge (p50 et p99, 10 000 itérations après 1 000 échauffements), les instructions traitées via Valgrind cachegrind, l'IPC et les taux de défaut de cache des compteurs matériels 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 (le +2 au-delà de 512 est 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 tokenizer, nombre d’instructions, IPC et allocations à 514 jetons.

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

SentencePiece et IREE utilisent tous deux le trie à double tableau introduit par Aoe en 1989 (Aoe, 1989 ; Karoonboonyanan). Le 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 de la traversée du trie ne nécessite que deux chargements de tableau, une addition d'entier et une comparaison, sans hachage, sans poursuite 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 de XLM-RoBERTa, le 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 en L2, et le pré-lecteur matériel se verrouille sur le motif d'accès régulier.

La référence, en revanche, conserve une table de hachage d'enfants par nœud, coûtant environ quatre chargements dépendants 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 chargements manquent le cache.

SP et IREE stockent tous deux l'ID du jeton directement dans le nœud du trie. L'implémentation du trie de référence renvoie uniquement les octets correspondants, donc l'encodeur doit allouer une nouvelle chaîne pour chaque correspondance, la rechercher dans une table de hachage latérale pour récupérer l'ID du jeton et lire le score à partir d'un vecteur séparé. SP renvoie l'ID comme valeur stockée du trie et lit le score à partir d'un tableau flottant 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 sous forme d'index de tableau plat, sans allocation ni hachage lors de la marche du trie elle-même. IREE alloue un seul tampon de sortie par encodage ; SP alloue en plus un nœud de réseau par correspondance pendant la DP, totalisant environ 1 559 allocations à 514 jetons.

Avant de changer le 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 avec le travail inutile supprimé ?

La référence conserve le trie hash-map-children de la référence mais apporte deux changements. Un struct scratch appartenant à l'appelant contient la table DP, les octets prétraités et les ID de jetons de sortie, de sorte que le chemin critique n'effectue aucune allocation après l'échauffement. L'ID du jeton et le score sont stockés directement dans le nœud du trie, supprimant la conversion de chaîne par correspondance et la recherche dans la table de hachage latérale.

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 de trie et le même motif d'accès, juste 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
Référence	156 µs	1,53M	2,62	0

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

La suppression de la matérialisation de chaîne 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 eux (2,86 → 2,62) car la poursuite de pointeur de la table de hachage, précédemment cachée derrière le travail de l'allocateur, est maintenant le coût dominant par octet.

Latence p50 et p99 du tokenizer de référence et de la référence sans allocation selon la taille d’entrée.

Taux de ratés du cache L2 du tokenizer de référence et de la 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éfaut L2 reste proche de 2-3 % jusqu'à 2k jetons contre 8-48 % pour la référence. Le coût restant est la poursuite de pointeur HashMap, que la section suivante aborde.

Le trie hash-map est une poursuite de pointeur. SP et IREE résolvent cela avec un 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 à usage général prenant en charge l'entrée en streaming, les décalages de sortie et plusieurs modes d'encodage ; cette flexibilité leur coûte une comptabilité supplémentaire autour du trie (un réseau dans SP, un pipeline à huit étages dans IREE). Nous n'avons besoin de rien de tout cela, donc nous inlinons le trie directement dans la boucle Viterbi et abandonnons 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'entier et une comparaison. Le trie entier fait environ 9 Mo, mais n'importe quel encodage ne parcourt que ~100 Ko de celui-ci, donc les données chaudes restent en L2.

Moteur	p50	p99	Instructions	L2 défaut %	L3 défaut %
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 %
Référence	155 µs	1 170 µs	1,53M	2,59 %	16,66 %
+ Darts	68 µs	80 µs	1,10M	6,45 %	11,42 %

Latence du tokenizer, instructions et taux de ratés du cache après l’adoption du trie double-array.

L'horloge tombe de 2,3× par rapport à la référence et de 4,8× par rapport à la référence. Les instructions chutent encore de 28 % par rapport à la référence. Couper le réseau et le pipeline nous permet de traiter 1,10M d'instructions par encodage contre 1,82M pour SP et 1,61M pour IREE à la même forme de trie ; l'écart est la comptabilité qu'ils conservent pour les API à usage général dont nous n'avons pas besoin.

Le trie Darts de l'Optimisation 1 nécessite deux chargements dépendants par étape d'octet. D'abord le décalage de base du parent, puis l'entrée de vérification à l'emplacement de l'enfant candidat pour confirmer que le parent le possède. L'entrée de vérification valide uniquement la propriété. La même question peut être résolue par un bitmap par nœud qui enregistre quels octets ont des enfants. Par la construction du 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 de vérification portent la même information, nous pouvons donc supprimer entièrement le tableau de vérification 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, qui s'exécute en environ un cycle. L'instruction POPCNT associée, qui compte tous les bits définis dans un registre, est ce que les autres conceptions de trie indexées par bitmap utilisent pour calculer un décalage d'enfant directement à partir du bitmap. Nous n'en avons pas besoin car la construction Darts stocke déjà ce décalage dans le champ de 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)
}

Pendant que nous y sommes, nous emballons les quatre champs par nœud (bitmap, base, ID de jeton et score) dans une ligne de cache de 64 octets. Une visite de nœud unique extrait maintenant tout : le score et l'ID de jeton dont nous avons besoin aux nœuds terminaux, plus le bitmap et la base pour l'étape suivante. Le tableau de vérification n'est utilisé que pendant la construction et est supprimé de la mise en page d'exécution. L'indexation du chemin critique n'est pas vérifiée car la construction du trie et la structure de la boucle garantissent déjà les limites. Le compromis est l'Espace : à 64 octets par nœud × ~780k emplacements, le trie passe de ~9 Mo à ~50 Mo.

Résultats à 514 jetons :

Optimisation 3 : Pages énormes pour le Trie	p50	p99	Instructions	Accès L2	L2 défaut %	L3 défaut %
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 %
Référence	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 emballage en ligne	65 µs	77 µs	1,06M	1,8K	20,53 %	15,01 %

Résultats du tokenizer à 514 jetons après l’ajout du bitmap et du packing en ligne.

L'horloge tombe 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 de vérification transforme l'étape de trie par octet de deux chargements de ligne de cache en un, ce qui représente 2,5× moins d'accès L2 par encodage (1,8K contre les 4,6K du Darts simple). Le taux de défaut bondit de 6 % à 20 %, mais la majeure partie de cela est le dénominateur : avec moins de la moitié des accès, les défauts que nous avons prennent une part plus importante.

Les données du trie font environ 50 Mo (780 000 nœuds à 64 octets chacun). Sur un processus Linux par défaut, cette mémoire est sauvegardée par des pages de 4 Ko, donc le trie complet s'étend sur environ 12 000 pages virtuelles distinctes, bien au-delà de ce que le tampon de recherche de traduction de données de premier niveau du CPU 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 brûlons environ 9 000 cycles dans les marches de table de page, environ 3 % du budget par encodage.

La solution est de sauvegarder le trie avec des pages énormes de 2 Mo au lieu de celles de 4 Ko. Les mêmes 50 Mo s'étendent maintenant sur 25 pages, tiennent entièrement dans le tampon de recherche de traduction, et les répertoires de table de page eux-mêmes cessent de concurrencer le trie pour le cache L2.

Linux a deux façons de sauvegarder la mémoire avec des pages de 2 Mo. L'opportuniste est les pages énormes transparentes (THP) : votre programme demande au noyau de mettre à niveau une région mémoire de 4 Ko à 2 Mo, et un fil de noyau en arrière-plan essaie de le faire quand il a le temps et la mémoire physique contiguë disponible. L'autre alternative est d'allouer le trie directement à partir d'un pool de pages énormes 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 énormes sont réservées ; le 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 % de latence horloge selon les tailles d'entrée, cohérent et proportionnel à la longueur de l'ensemble de travail :

Jetons	Sans pages énormes	Avec pages énormes	Δ
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 tokenizer avec et sans grandes pages selon la taille d’entrée.

La plus grande victoire se situe à 4k jetons, où l'ensemble de travail est suffisamment grand pour que le trafic de table de page concurrençait activement les données du trie pour la bande passante L2. Au-delà de 4k, le gain diminue car les défauts L3 dominent et les marches de table de page deviennent une part plus petite du temps total.

En réunissant les trois optimisations, voici le moteur final mesuré par rapport à tous les autres tokenizers RoBERTa que nous avons testés. Parallèlement aux binaires natifs C++ SentencePiece et C IREE, nous avons également évalué les crates de wrapper Rust prêts à l'emploi qu'un utilisateur de production utiliserait, sentencepiece et iree-tokenizer. Le crate sentencepiece mesure une latence 1,6-1,9× plus élevée que le binaire natif C++ à chaque taille d'entrée, la surcharge provenant du wrapper allouant une nouvelle liste de morceaux de jetons, chacun détenant sa propre copie du texte du morceau, à chaque encodage. Le crate iree-tokenizer dimensionne son tampon interne à chaque entrée, ce qui évite une allocation constante par encodage que le binaire natif IREE paie indépendamment de la taille de l'entrée, et cela compte le plus aux 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 la surcharge par appel du wrapper lors de la copie des jetons vers Rust pousse sa latence au-dessus de celle du binaire natif.

Latence p50 et p99 par encodage selon les implémentations de tokenizer 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 exacte par rapport à la référence. Chaque étape d'octet est maintenant un seul chargement de ligne de cache de 64 octets, un test de bitmap et une addition d'entier, contre une sonde HashMap, un appel d'allocateur et plusieurs branches de cas spéciaux dans l'implémentation de référence. Les instructions par encodage sont passées de 3,66M à 1,04M, une baisse de 3,5×. L'horloge correspond à ce ratio aux entrées courtes et s'élargit davantage aux entrées longues, où les allocations par jeton de la référence commencent à déborder L2 et L3.

En pratique, ce travail a réduit l'utilisation du CPU dans la couche de prétraitement de notre pile d'inférence de 5 à 6× et a permis de gagner des dizaines de millisecondes sur la latence du reranker.

La majeure partie de la réduction de latence provient de changements légers qui ont supprimé le travail gaspillé : suppression des allocations par encodage, remplacement du trie HashMap-children par un double tableau et suppression de la surcharge de l'API générique. L'utilisation du bitmap par nœud existant comme vérification de validité et la sauvegarde du trie avec des pages énormes de 2 Mo combinées ont représenté 10-20 % en plus du passage plus large à un trie à double tableau, selon la taille de l'entrée.

Les crates de wrapper Rust prêts à l'emploi autour de SentencePiece et IREE illustrent le même point à l'inverse. Par-dessus les tokenizers natifs C et C++, la limite FFI et les allocations par appel que les wrappers introduisent sont visibles dans les chiffres de latence. La source Rust pour l'encodeur Viterbi final est disponible sur github.com/perplexityai/pplx-garden.

Hugging Face tokenizers — https://github.com/huggingface/tokenizers

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

IREE tokenizer — 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

Huge pages — https://docs.kernel.org/admin-guide/mm/hugetlbpage.html

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}
}

Motivation

L'algorithme de Viterbi

Implémentation de référence

Une note sur le prétraitement

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

Comment nous avons mesuré

Pourquoi SentencePiece et IREE sont plus rapides : le Trie à double tableau

Référence : un portage à zéro allocation de la référence

Optimisation 1 : Trie à double tableau sans abstractions supplémentaires

Optimisation 2 : bitmap et empaquetage en ligne

Optimisation 3 : pages de taille importante pour le trie

Comparaison finale

Réflexions finales

Références

Citation