Poprawa wydajności procesora tokenizatora unigram

W ramach projektu skoncentrowanego na wydajności wdrożyliśmy nasz tokenizator Unigram od zera. Przy produkcyjnych długościach danych wejściowych nowy koder skraca opóźnienie p50 o około 5-krotnie w porównaniu z crate tokenizers od Hugging Face, około 2-krotnie w porównaniu z SentencePiece (C++) i około 1,5-krotnie w porównaniu z tokenizatorem IREE (C), przy zerowej liczbie alokacji sterty w stanie ustalonym. W produkcji pozwoliło to na 5-6-krotne zmniejszenie wykorzystania procesora w naszym stosie wnioskowania oraz skrócenie opóźnienia rerankera o wartości dwucyfrowe w milisekundach.

Koszt wnioskowania LLM jest zazwyczaj postrzegany jako praca GPU: pamięci podręczne KV, kernele uwagi, routing ekspercki. Jednak znacząca część ruchu produkcyjnego odbywa się na znacznie mniejszych modelach, od dwóch do trzech rzędów wielkości mniejszych niż transformatory typu frontier. Modele osadzania (embedding), klasyfikatory i rerankery nadal działają na GPU, ale przejście w przód (forward pass) jest na tyle krótkie, że wszystko, co je otacza, zaczyna mieć znaczenie.

Najbardziej oczywistym przypadkiem jest reranker oceniający setki dokumentów kandydatów na żądanie. W przypadku małego modelu obliczenia GPU zazwyczaj kończą się w ciągu kilku milisekund, ale każde dane wejściowe muszą najpierw przejść przez tokenizację po stronie procesora, zamieniając tekst na identyfikatory słownictwa. Gdy partia danych jest duża, może to stanowić znaczącą część całkowitego opóźnienia żądania.

Ta praca koncentruje się na XLM-RoBERTa, słownictwie Unigram o wielkości 250 tys. tokenów wytrenowanym przy użyciu SentencePiece. Dostrojone kodery z rodziny RoBERTa są częstym wyborem w zadaniach rankingu, wyszukiwania i podobieństwa. Raportujemy opóźnienie jednowątkowe na kodowanie; w produkcji stosujemy równoległość między rdzeniami przy użyciu rayon.

W całym poście traktujemy tokenizers crate od Hugging Face jako punkt odniesienia, z którym dokonujemy porównań, ponieważ jest to domyślny tokenizator Rust, po który sięgają zespoły. SentencePiece (C++) i IREE (C) zostały uwzględnione do porównania jako dwie inne główne implementacje produkcyjne.

Tokenizacja Unigram, wprowadzona przez Kudo w 2018 r. (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) i dostarczona w SentencePiece, definiuje segmentację jako problem ścieżki o największym prawdopodobieństwie: każdy token słownika ma wyuczone logarytmiczne prawdopodobieństwo, a koder wybiera segmentację, której sumaryczne wyniki tokenów są najwyższe. Problem rozwiązuje się za pomocą algorytmu Viterbiego (1967), z pozycjami bajtów jako warstwami grafu i tokenami słownika jako krawędziami obejmującymi ciągły zakres bajtów. Rekurencja jest jednowymiarowym programowaniem dynamicznym (DP) względem pozycji bajtów:

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

Pętla zewnętrzna jest liniowa względem długości wejścia, ograniczona w każdej pozycji przez najdłuższy token słownika (48 bajtów UTF-8 dla XLM-RoBERTa). Przejście przez wewnętrzne trie jest ścieżką krytyczną: przy danych wejściowych o wielkości 16 tys. tokenów wykonuje setki tysięcy przejść i realizuje dziesiątki milionów instrukcji na kodowanie.

tokenizers v0.22 od Hugging Face implementuje rekurencję bezpośrednio, używając mapy mieszającej (hash-map) do wyszukiwania wstecznego i mapy mieszającej dla dzieci (hash-map-children) w każdym węźle 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
}

Oraz pętla wewnętrzna (uproszczona) – każde dopasowanie wyszukiwania prefiksów alokuje nowy ciąg znaków, a następnie sonduje mapę mieszającą:

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
}

Kod jest czytelnym Rustem, ale wykonuje się w nim wiele pracy, która jest niewidoczna na poziomie źródłowym. Następna sekcja wyjaśnia, na co zużywany jest czas.

Zanim uruchomiony zostanie algorytm Viterbiego, dane wejściowe przechodzą przez normalizację Unicode (NFKC) i oznaczanie granic wyrazów. Profilowanie wskazuje, że przetwarzanie wstępne zajmuje mniej niż 5% całkowitego czasu kodowania w naszych scenariuszach, dlatego w dalszej części postu skupiamy się na przejściu w przód algorytmu Viterbiego.

Istnieją dwa sposoby na przyspieszenie kodu procesora: wykonanie mniejszej liczby instrukcji lub uzyskanie wyższej wartości IPC (instrukcji na cykl). Cachegrind mierzy pierwszą, a perf stat mierzy drugą wartość. Każda optymalizacja w tym poście koncentruje się na jednej lub drugiej.

Procesory zwalniają z czterech powodów: chybienia pamięci podręcznej (L1d ~4 cykle, L3-miss-to-DRAM to setki cykli), łańcuchy zależnych odczytów, których silnik out-of-order nie może zrównoleglić, błędne przewidywania rozgałęzień (~15-20 cykli traconych przy błędnym odgadnięciu) oraz wywołania alokatora. Implementacja silnika referencyjnego unigram natrafia na wszystkie cztery problemy:

Wąskie gardło	Mechanizm	Zmierzony wpływ
Alokacja na dopasowanie	String::from_utf8 na dopasowanych bajtach plus odwołanie do AHashMap na dopasowanie trie; identyfikator tokena znajduje się już w liściu, ale implementacja referencyjna nie przechowuje go tam	7295 alokacji dla 512 tokenów, 14 542 dla 1 tys., 299 171 dla 16 tys.
Wyszukiwanie wskaźnika na bajt	AHashMap w każdym węźle trie, haszowanie następnego bajtu i dereferencja dwóch wskaźników na krok; struktura Node zawiera również zimne pola, takie jak is_leaf, wciągając je do pamięci podręcznej przy każdej wizycie	Wskaźnik chybień L1d pozostaje niski (~2,5% dla 1 tys.) ponieważ gorące strony alokatora znajdują się w L1; koszt objawia się jako opóźnienie zależnych odczytów
Trashing L2 przy długich danych wejściowych	Tabela DP, roboczy bufor prefiksów i lista tokenów wyjściowych są za każdym razem świeżo alokowane	Wskaźnik chybień L2 rośnie z 12% dla 128 tokenów do 50% dla 16 tys., gdy skumulowany ślad alokacji przekracza 2 MB pamięci podręcznej L2 na rdzeń

Odniesienie do wąskich gardeł tokenizera, mechanizmów i zmierzonych efektów.

Dominującym kosztem jest alokacja na dopasowanie. Przy 1 tys. tokenów implementacja referencyjna realizuje (kończy) około 7 mln instrukcji na kodowanie przy IPC wynoszącym 2,8, z których większość przypada na wywołania alokatora i konwersje ciągów znaków. Chybienia L2 mają znaczenie tylko powyżej ~1 tys. tokenów, gdy skumulowany ślad alokacji przepełnia pamięć podręczną na rdzeń.

p50 i p99 opóźnienia referencyjnego tokenizera według rozmiaru danych wejściowych.

Współczynnik nietrafień w pamięci podręcznej L2 referencyjnego tokenizera według rozmiaru danych wejściowych.

Liczba bajtów alokowanych na jedno kodowanie przez referencyjny tokenizator w różnych rozmiarach danych wejściowych.

Alokacja na token jest stała: około 2 KB i ~18 alokacji na token, niezależnie od rozmiaru danych wejściowych. Nieliniowość opóźnienia wynika z L2: skumulowane bajty przekraczają 2 MB w okolicach 1 tys. tokenów, wskaźnik chybień rośnie z 8% do ~50%, a krzywa p99 wygina się, aby to odzwierciedlić.

Porównujemy trzy produkcyjne tokenizatory, wszystkie uruchamiające ten sam algorytm Unigram z tym samym modelem XLM-RoBERTa:

Silnik	Język	Trie
Hugging Face	Rust	Trie prefiksów HashMap-children
SentencePiece	C++	Darts double-array trie
Tokenizator IREE	C	Darts double-array trie

Implementacje tokenizera, języki i reprezentacje trie użyte w porównaniu.

Łączymy SentencePiece i IREE przez bezpośrednie C/C++ FFI, a nie przez ich wrappery w Rust; wrappery dodają alokacje na wywołanie i konwersje ciągów znaków, które nie są częścią kosztu bazowej biblioteki.

Wszystkie silniki są zbudowane z pełnymi optymalizacjami (O3, targetowanie natywne, optymalizacja czasu łączenia). Pomiary są przypięte do jednego nieobciążonego rdzenia na procesorze Intel Xeon Platinum 8488C. Dla każdego silnika i rozmiaru danych wejściowych zbieramy opóźnienie rzeczywiste (p50 i p99, 10 000 iteracji po 1 000 rozgrzewkowych), instrukcje zrealizowane przez Valgrind cachegrind, IPC oraz wskaźniki chybień pamięci podręcznej z liczników sprzętowych perf stat, a także alokacje sterty zliczane przez wrapper alokatora globalnego wokół mierzonego okna.

Dla 514 tokenów (+2 powyżej 512 to wstrzyknięcie BOS/EOS):

Implementacja	Opóźnienie p50	Instrukcje	IPC	Alokacje
Hugging Face	349 µs	3.60M	2.89	7295
SentencePiece	128 µs	1.83M	3.39	1559
IREE	112 µs	2.28M	3.06	1

Opóźnienie tokenizera, liczba instrukcji, IPC i alokacje przy 514 tokenach.

SP i IREE osiągają około jedną trzecią opóźnienia implementacji referencyjnej. Realizują około połowę mniej instrukcji, działają przy wyższym IPC i alokują od jednego do dwóch rzędów wielkości mniej.

SentencePiece i IREE używają podwójnej tablicy trie (double-array trie) wprowadzonej przez Aoe w 1989 r. (Aoe, 1989; Karoonboonyanan). Trie jest zakodowane w dwóch płaskich tablicach liczb całkowitych, base i check:

cpp

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

Wyszukiwanie potomka to:

cpp

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

Każdy krok przejścia przez trie potrzebuje tylko dwóch odczytów z tablicy, jednego dodawania liczb całkowitych i jednego porównania, bez haszowania, bez pościgu za wskaźnikami i bez alokacji sterty na węzeł. Konstrukcja umieszcza każdy węzeł tak, aby jego potomkowie trafiali na wolne miejsca w tablicy walidacyjnej, pozwalając wielu rodzicom dzielić regiony, o ile ich potomkowie różnią się bajtem, po którym następują.

Dla słownika XLM-RoBERTa o wielkości 250 tys. tokenów całe trie mieści się w ~9 MB ciągłej pamięci (~780 tys. miejsc × 12 bajtów na offset, walidację i ID tokena na miejsce). Gorący zestaw roboczy pojedynczego kodowania jest rzędu 100 KB, mieści się wygodnie w L2, a sprzętowy prefetcher blokuje się na regularnym wzorcu dostępu.

Implementacja referencyjna natomiast utrzymuje mapę mieszającą dzieci dla każdego węzła, co kosztuje około czterech zależnych odczytów na krok bajtowy. Zahaszuj bajt, wczytaj kubełek (bucket), przejdź wskaźnikiem do wpisu w kubełku, przejdź kolejnym wskaźnikiem do węzła potomnego na stercie, a następnie skonsultuj własną mapę mieszającą dla następnego bajtu. Obiekty na stercie są rozproszone i większość odczytów chybia pamięć podręczną.

SP i IREE przechowują ID tokena bezpośrednio w węźle trie. Implementacja trie referencyjnego zwraca tylko dopasowane bajty, więc koder musi alokować świeży ciąg znaków dla każdego dopasowania, wyszukać go w bocznej mapie mieszającej, aby odzyskać ID tokena, i odczytać wynik z oddzielnego wektora. SP zwraca ID jako wartość zapisaną w trie i odczytuje wynik z płaskiej tablicy zmiennoprzecinkowej:

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 używa stanowego API kursora zamiast zwracać wszystkie dopasowania w jednym wywołaniu, ale model kosztów jest taki sam: ID tokenów i wyniki są dostępne jako indeksy płaskiej tablicy, bez alokacji i haszowania podczas samego przejścia przez trie. IREE alokuje tylko jeden bufor wyjściowy na kodowanie; SP dodatkowo alokuje węzeł sieci podczas DP, co daje łącznie około 1559 alokacji dla 514 tokenów.

Zanim zmienimy trie, izolujemy koszt algorytmu od kosztu implementacji: jak szybko działa implementacja referencyjna po usunięciu zbędnej pracy?

Punkt odniesienia utrzymuje trie z mapą mieszającą dla dzieci z implementacji referencyjnej, ale wprowadza dwie zmiany. Należąca do obiektu wywołującego struktura robocza (scratch struct) przechowuje tabelę DP, przetworzone bajty i identyfikatory tokenów wyjściowych, więc ścieżka krytyczna nie wykonuje żadnych alokacji po rozgrzewce. Identyfikator tokena i wynik są przechowywane bezpośrednio w węźle trie, co eliminuje konwersję ciągów na dopasowanie i wyszukiwanie w bocznej mapie mieszającej.

Pętla wewnętrzna:

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

To ten sam algorytm, z tym samym kształtem trie i tym samym wzorcem dostępu, tylko bez alokacji i zbędnych wyszukiwań. Wyniki dla 514 tokenów:

Silnik	p50	Instrukcje	IPC	Alokacje
Hugging Face	326 µs	3.66M	2.86	7295
SentencePiece	127 µs	1.81M	3.55	1559
IREE	111 µs	1.58M	3.26	1
Baseline	156 µs	1.53M	2.62	0

Wersja bazowa bez alokacji w porównaniu z istniejącymi implementacjami tokenizera przy 514 tokenach.

Usunięcie materializacji ciągów na dopasowanie i drugorzędnego wyszukiwania w mapie mieszającej zmniejsza liczbę instrukcji 2,4-krotnie w porównaniu z implementacją referencyjną. IPC spada wraz z nimi (2,86 → 2,62), ponieważ pościg za wskaźnikami mapy mieszającej, wcześniej ukryty za pracą alokatora, jest teraz dominującym kosztem na bajt.

p50 i p99 opóźnienia referencyjnego tokenizera i tokenizera bazowego bez alokacji według rozmiaru danych wejściowych.

Współczynnik nietrafień w pamięci podręcznej L2 referencyjnego tokenizera i tokenizera bazowego bez alokacji według rozmiaru danych wejściowych.

Baseline skraca opóźnienie p50 2-3-krotnie w porównaniu z implementacją referencyjną dla każdego rozmiaru, a wskaźnik chybień L2 pozostaje bliski 2-3% dla 2 tys. tokenów w porównaniu z 8-48% w implementacji referencyjnej. Pozostałym kosztem jest pościg za wskaźnikami HashMap, którym zajmuje się następna sekcja.

Trie mapy mieszającej to pościg za wskaźnikami. SP i IREE rozwiązują to za pomocą double-array trie (patrz poprzednia sekcja), a my zapożyczamy tę samą strukturę danych. Różnica polega na tym, że SP i IREE to biblioteki ogólnego przeznaczenia obsługujące przesyłanie strumieniowe, offsety wyjściowe i wiele trybów kodowania; ta elastyczność kosztuje je dodatkową księgowość wokół trie (sieć w SP, ośmioetapowy potok w IREE). My nie potrzebujemy niczego z tego, więc wklejamy trie bezpośrednio do pętli Viterbiego i porzucamy resztę.

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

Krok bajtowy kosztuje dwa odczyty tablicy, dodawanie liczb całkowitych i porównanie. Całe trie ma około 9 MB, ale każde pojedyncze kodowanie przechodzi tylko przez ~100 KB z niego, więc gorące dane pozostają w L2.

Silnik	p50	p99	Instrukcje	L2 miss%	L3 miss%
Hugging Face	327 µs	3340 µs	3.66M	8.17%	14.66%
SentencePiece	127 µs	1140 µs	1.82M	4.74%	4.07%
IREE	111 µs	132 µs	1.61M	5.82%	3.87%
Baseline	155 µs	1170 µs	1.53M	2.59%	16.66%
+ Darts	68 µs	80 µs	1.10M	6.45%	11.42%

Opóźnienie tokenizera, liczba instrukcji i współczynniki nietrafień w pamięci podręcznej po zastosowaniu trie z podwójną tablicą.

Czas zegarowy spada 2,3-krotnie w stosunku do baseline'u i 4,8-krotnie w stosunku do implementacji referencyjnej. Liczba instrukcji spada o kolejne 28% w stosunku do baseline'u. Usunięcie sieci i potoku pozwala nam realizować 1,10 mln instrukcji na kodowanie w porównaniu z 1,82 mln dla SP i 1,61 mln dla IREE przy tym samym kształcie trie; różnicą jest księgowość, którą one utrzymują dla API ogólnego przeznaczenia, których nie potrzebujemy.

Trie Darts z Optymalizacji 1 potrzebuje dwóch zależnych odczytów na krok bajtowy. Najpierw bazowy offset rodzica, następnie wpis kontrolny (check entry) w miejscu potencjalnego potomka, aby potwierdzić, że rodzic go posiada. Wpis kontrolny tylko waliduje własność. To samo pytanie można rozwiązać za pomocą bitmapy na węzeł, która rejestruje, które bajty mają potomków. Dzięki konstrukcji double-array trie, bit dla bajtu b w bitmapie rodzica jest ustawiony dokładnie wtedy, gdy miejsce docelowe dla tego bajtu jest własnością rodzica. Bitmapa i tablica kontrolna niosą te same informacje, więc możemy całkowicie porzucić tablicę kontrolną ze struktury czasu wykonania.

Wyszukiwanie bitmapy kompiluje się do testu pojedynczego bitu względem jednego 64-bitowego słowa, co trwa około cyklu. Powiązany z tym instrukcja POPCNT, która zlicza wszystkie ustawione bity w rejestrze, jest tym, co inne projekty trie indeksowanych bitmapami używają do bezpośredniego obliczania offsetu potomka z bitmapy. Nie potrzebujemy jej, ponieważ konstrukcja Darts już przechowuje ten offset w polu base węzła.

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

Przy okazji pakujemy cztery pola na węzeł (bitmapa, base, ID tokena i wynik) w jedną 64-bajtową linię pamięci podręcznej. Pojedyncza wizyta w węźle wciąga teraz wszystko: wynik i ID tokena, których potrzebujemy w węzłach końcowych, plus bitmapę i base dla następnego kroku. Tablica kontrolna jest używana tylko podczas konstrukcji i jest porzucana w układzie czasu wykonania. Indeksowanie ścieżki krytycznej jest niesprawdzane (unchecked), ponieważ konstrukcja trie i struktura pętli gwarantują już granice. Kompromisem jest przestrzeń: przy 64 bajtach na węzeł × ~780 tys. miejsc trie rośnie z ~9 MB do ~50 MB.

Wyniki dla 514 tokenów:

Silnik	p50	p99	Instrukcje	Dostępów do L2	L2 miss%	L3 miss%
Hugging Face	327 µs	3340 µs	3.66M	15.0K	8.17%	14.66%
SentencePiece	127 µs	1140 µs	1.82M	7.5K	4.74%	4.07%
IREE	111 µs	132 µs	1.61M	4.7K	5.82%	3.87%
Baseline	155 µs	1170 µs	1.53M	6.1K	2.59%	16.66%
+ Darts	68 µs	80 µs	1.10M	4.6K	6.45%	11.42%
+ Bitmap i pakowanie wbudowane	65 µs	77 µs	1.06M	1.8K	20.53%	15.01%

Wyniki tokenizera przy 514 tokenach po dodaniu bitmapy i pakowania inline.

Czas zegarowy systematycznie spada o kolejne 4,5% dla 514 tokenów i rośnie wraz z długością danych wejściowych, osiągając 6-7% dla 4 tys. i 16 tys. Porzucenie tablicy kontrolnej zmienia krok trie na bajt z dwóch odczytów linii pamięci podręcznej na jeden, co oznacza 2,5-krotnie mniej dostępów do L2 na kodowanie (1,8 tys. w porównaniu z 4,6 tys. w czystym Darts). Wskaźnik chybień skacze z 6% do 20%, ale większość z tego to mianownik: przy mniej niż połowie dostępów, te chybienia, które mamy, stanowią większy udział.

Dane trie to około 50 MB (780 000 węzłów po 64 bajty każdy). W domyślnym procesie Linuxa ta pamięć jest wspierana przez strony 4 KB, więc całe trie obejmuje około 12 000 różnych stron wirtualnych, znacznie więcej niż mieści Translation Lookaside Buffer pierwszego poziomu procesora (Sapphire Rapids ma 96 wpisów). Każdy krok Viterbiego dotyka innego węzła, a podczas kodowania 512 tokenów zużywamy około 9 000 cykli na przejścia przez tablice stron, około 3% budżetu na kodowanie.

Rozwiązaniem jest wsparcie trie za pomocą huge pages 2 MB zamiast stron 4 KB. Te same 50 MB obejmują teraz 25 stron, mieszczą się całkowicie w Translation Lookaside Buffer, a same katalogi tablic stron przestają konkurować z trie o pamięć podręczną L2.

Linux ma dwa sposoby na wspieranie pamięci za pomocą stron 2 MB. Oportunistycznym są przezroczyste huge pages (THP): program prosi jądro o ulepszenie regionu pamięci ze stron 4 KB na 2 MB, a wątek jądra w tle próbuje to zrobić, gdy ma czas i dostępną ciągłą pamięć fizyczną. Drugą alternatywą jest alokacja trie bezpośrednio z wstępnie zarezerwowanej puli huge pages, którą jądro odkłada przy starcie systemu:

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,
    )
};

Pula musi istnieć na hoście (vm.nr_hugepages ustawione przy starcie). Na naszych instancjach produkcyjnych zarezerwowanych jest 10 561 huge pages; trie zużywa 24 z nich. Po mmap, /proc/<pid>/smaps potwierdza region: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Zysk wynosi 3–12% czasu zegarowego dla różnych rozmiarów danych wejściowych, jest spójny i proporcjonalny do długości zestawu roboczego:

Tokeny	Bez huge pages	Z huge pages	Δ
130	15,7 µs	15,2 µs	−3,6%
514	65,4 µs	63,1 µs	−3,4%
1026	129,5 µs	125,3 µs	−3,3%
4098	773 µs	679 µs	−12,0%
16 386	3808 µs	3554 µs	−6,8%

Opóźnienie tokenizera z i bez ogromnych stron pamięci dla różnych rozmiarów danych wejściowych.

Największy zysk jest przy 4 tys. tokenów, gdzie zestaw roboczy jest na tyle duży, że ruch w tablicy stron aktywnie konkurował z danymi trie o pasmo L2. Powyżej 4 tys. zysk maleje, ponieważ dominują chybienia L3, a przejścia przez tablicę stron stają się mniejszym udziałem w całkowitym czasie.

Łącząc trzy optymalizacje, oto ostateczny silnik mierzony w stosunku do każdego innego tokenizatora RoBERTa, który testowaliśmy. Obok natywnych plików binarnych SentencePiece w C++ i IREE w C, przetestowaliśmy również gotowe wrappery w Rust, po które sięgnąłby użytkownik produkcyjny, czyli sentencepiece i iree-tokenizer. Crate sentencepiece mierzy opóźnienie wyższe o 1,6-1,9× niż natywny plik binarny C++ dla każdego rozmiaru danych wejściowych, a narzut pochodzi z wrappera alokującego świeżą listę części tokenów, z których każda przechowuje własną kopię tekstu, przy każdym kodowaniu. Crate iree-tokenizer dostosowuje swój wewnętrzny bufor do każdych danych wejściowych, co pozwala uniknąć stałej alokacji na kodowanie, którą natywny plik binarny IREE płaci niezależnie od rozmiaru danych wejściowych, co ma największe znaczenie przy krótkich danych wejściowych, gdzie alokacja w przeciwnym razie zdominowałaby właściwą pracę tokenizacji. Przy produkcyjnych długościach danych wejściowych koszt stały się amortyzuje, a własny narzut wrappera z kopiowania tokenów z powrotem do Rusta przesuwa opóźnienie powyżej natywnego pliku binarnego.

p50 i p99 opóźnienia na jedno kodowanie w różnych implementacjach tokenizera i rozmiarach danych wejściowych.

Ostateczny silnik uruchamia ten sam algorytm Viterbiego dla tego samego słownika o wielkości 250 000 tokenów, nie alokuje niczego na ścieżce krytycznej i produkuje wyjście identyczne z implementacją referencyjną. Każdy krok bajtowy jest teraz pojedynczym odczytem 64-bajtowej linii pamięci podręcznej, jednym testem bitmapy i jednym dodawaniem liczb całkowitych, zamiast sondowania HashMap, wywołania alokatora i kilku specjalnych rozgałęzień w implementacji bazowej. Liczba instrukcji na kodowanie spadła z 3,66 mln do 1,04 mln, co stanowi 3,5-krotny spadek. Czas zegarowy pasuje do tego stosunku przy krótkich danych wejściowych i rozszerza się dalej przy długich danych wejściowych, gdzie alokacje na token w implementacji referencyjnej zaczynają przepełniać L2 i L3.

W praktyce praca ta zredukowała wykorzystanie procesora w warstwie przetwarzania wstępnego naszego stosu wnioskowania o 5-6 razy i skróciła opóźnienie rerankera o dwucyfrowe wartości milisekundowe.

Większość redukcji opóźnienia pochodzi z lekkich zmian, które usunęły marnotrawstwo pracy: rezygnacja z alokacji na kodowanie, zastąpienie trie z mapą mieszającą dla dzieci strukturą double-array i usunięcie narzutu ogólnego API. Wykorzystanie istniejącej bitmapy na węzeł jako kontroli poprawności i wsparcie trie przez huge pages 2 MB złożyło się na dodatkowe 10-20% zysku na dodatek do głównego przejścia na double-array trie, w zależności od rozmiaru danych wejściowych.

Gotowe wrappery Rust wokół SentencePiece i IREE ilustrują ten sam punkt w odwrotny sposób. Ponad natywnymi tokenizatorami C i C++, granica FFI i alokacje na wywołanie, które wprowadzają wrappery, są widoczne w liczbach opóźnień. Kod źródłowy w Rusta dla ostatecznego kodera Viterbiego jest dostępny na 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}
}

Motywacja

Algorytm Viterbiego

Implementacja referencyjna

Uwaga na temat przetwarzania wstępnego

Gdzie implementacja referencyjna zużywa swój czas

Jak mierzyliśmy

Dlaczego SentencePiece i IREE są szybsze: Double-Array Trie

Punkt odniesienia: Port implementacji referencyjnej bez alokacji

Optymalizacja 1: Double-Array Trie bez dodatkowych abstrakcji

Optymalizacja 2: Bitmapa i pakowanie w linii

Optymalizacja 3: Ogromne strony dla drzewa trie

Ostateczne porównanie

Końcowe przemyślenia

Referencje

Cytat