Подобряване на производителността на CPU за Unigram токенизатор

Преработихме нашия Unigram токенизатор от нулата като целенасочен проект за производителност. При производствени дължини на входа новият енкодер намалява p50 латентността с приблизително 5 пъти спрямо библиотеката за токенизатори на Hugging Face, ~2 пъти спрямо SentencePiece (C++) и ~1.5 пъти спрямо токенизатора на IREE (C), без никакви заделяния на памет в стационарно състояние. В производството това намали използването на процесора в нашия инферентен стек с 5-6 пъти и отсече двуцифрени милисекунди от латентността на реранкера.

Цената за инференс на LLM обикновено се определя като работа на GPU: KV кешове, атеншън ядра, експертно насочване. Но значителна част от производствения трафик работи на много по-малки модели, с два до три порядъка по-малки от фронтиър трансформерите. Моделите за вграждане (embedding), класификаторите и реранкерите все още работят на GPU, но фазата на прав преход е достатъчно кратка, че всичко около нея започва да има значение.

Най-ясният пример е реранкер, който оценява стотици документи кандидати за заявка. При малък модел изчисленията на GPU обикновено приключват за едноцифрени милисекунди, но всеки вход все пак трябва първо да премине през токенизация от страната на CPU, превръщайки текста във идентификатори на речник. Когато партидата е голяма, това може да представлява значителна част от общата латентност на заявката.

Тази работа е насочена към XLM-RoBERTa, 250-хиляден Unigram речник, обучен със SentencePiece. Финно настроените енкодери от семейството на RoBERTa са често срещан избор за задачи за ранкиране, извличане и сходство. Тук докладваме латентност на еднонишков код за кодиране; в производството паралелизираме между ядрата с помощта на rayon.

В тази публикация третираме tokenizers crate на Hugging Face като референтната база, спрямо която правим бенчмарк, тъй като това е стандартният Rust токенизатор, към който повечето екипи посегат. SentencePiece (C++) и IREE (C) са включени за сравнение като другите две основни производствени реализации.

Токенизацията Unigram, въведена от Кудо през 2018 г. (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) и реализирана в SentencePiece, представя сегментирането като проблем с най-вероятен път: всеки токен от речника има научена лог-вероятност, а енкодерът избира сегментацията, чиито оценки на токените дават най-висока обща стойност. Проблемът се решава с алгоритъма на Витерби (1967 г.), като позициите на байтовете са слоеве на граф, а токените от речника са ръбове, обхващащи непрекъснат диапазон от байтове. Рекурентността е едномерно динамично програмиране (DP) върху позициите на байтовете:

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

Външният цикъл е линеен спрямо дължината на входа, ограничен на всяка позиция от най-дългия токен в речника (48 UTF-8 байта за XLM-RoBERTa). Вътрешното обхождане на дървото (trie) е горещият път: при вход от 16 000 токена то изпълнява стотици хиляди преходи и завършва десетки милиони инструкции за едно кодиране.

Tokenizers v0.22 на Hugging Face реализира рекурентността директно, с хеш-таблица за обратно търсене и деца-хеш-таблици във всеки възел на дървото:

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
}

А вътрешният цикъл (опростен) – всяко съвпадение при търсене на префикс заделя нов низ и след това проверява хеш-таблицата:

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
}

Кодът е четим Rust, но се извършва много работа, която е невидима на ниво изходен код. Следващия раздел разглобява къде отива времето.

Преди да стартира Витерби, входът преминава през уникодна нормализация (NFKC) и маркиране на границите на думите. Профайлингът показва предварителната обработка на по-малко от 5% от общото време за кодиране в нашите сценарии, така че за останалата част от публикацията се фокусираме върху правата стъпка на Витерби.

Има два начина да се ускорим CPU кодът: да се изпълняват по-малко инструкции или да се постигне по-висок IPC. Cachegrind измерва първото, perf stat измерва второто. Всяка оптимизация в тази публикация води до едното или другото.

Процесорите блокират поради четири неща: пропуски в кеша (L1d ~4 цикъла, пропуск в L3 до DRAM е стотици), вериги със занесимо зареждане, които механизмът за извънредно изпълнение не може да паралелизира, грешни предвиждания на клонове (~15-20 цикъла изчистване на грешно предположение) и извиквания на алокатора. Референтната реализация на unigram енкодера се сблъсква с всичките четири:

Ботълнек	Механизъм	Измерено въздействие
Алокация на съвпадение	String::from_utf8 върху съвпадналите байтове плюс AHashMap търсене за всяко съвпадение в дървото; ID-то на токена вече е в листото, но референцията не го съхранява там	7 295 алокации при 512 токена, 14 542 при 1k, 299 171 при 16k
Преследване на указатели на байт	AHashMap във всеки възел на дървото, хеширане на следващия байт и деференциране на два указателя на стъпка; структурата Node също така групира студени полета като is_leaf, изтегляйки ги в кеша при всяко посещение	Процентът на пропуски в L1d остава нисък (~2.5% при 1k), защото горещите страници на алокатора живеят в L1; цената се появява като латентност на зависимо зареждане
L2 thrashing при дълги входове	DP таблицата, префиксният буфер и списъкът с изходни токени се заделят наново при всяко извикване	Процентът на пропуски в L2 се качва от 12% при 128 токена до 50% при 16k, тъй като кумулативният отпечатък на алокациите надвишава 2 MB L2 на ядро

Референтни тесни места, механизми и измерени ефекти на токенизатора.

Доминиращата цена е алокацията на съвпадение. При 1k токена референтната реализация завършва (изпълнява) приблизително 7 млн. инструкции на кодиране при IPC 2.8, повечето от тях в извиквания на алокатора и низови преобразувания. Пропуските в L2 имат значение само след ~1k токена, след като кумулативният отпечатък на алокациите прелее кеша на ядрото.

p50 и p99 латентност на референтния токенизатор според размера на входа.

Ниво на пропуски в L2 кеша на референтния токенизатор според размера на входа.

Заделени байтове на encode от референтния токенизатор при различни размери на входа.

Алокацията на токен е постоянна: около 2 KB и ~18 алокации на токен, независимо от размера на входа. Нелинейността на латентността идва от L2: кумулативните байтове преминават 2 MB около 1k токена, процентът на пропуските се качва от 8% до ~50% и кривата на p99 се огъва, за да съответства.

Сравняваме три производствени токенизатора, всички изпълняващи един и същ алгоритъм Unigram срещу същия модел XLM-RoBERTa:

Енкодер	Език	Дърво (Trie)
Hugging Face	Rust	HashMap-children prefix trie
SentencePiece	C++	Darts double-array trie
IREE токенизатор	C	Darts double-array trie

Имплементации на токенизатора, езици и представяния на trie, използвани в сравнението.

Свързваме SentencePiece и IREE чрез директен C/C++ FFI, а не чрез техните Rust обвивки; обвивките добавят алокации на извикване и низови преобразувания, които не са част от цената на основната библиотека.

Всички двигатели са изградени с пълни оптимизации (O3, насочване към native, оптимизация по време на линкване). Измерванията се правят фиксирани към едно тихо ядро на Intel Xeon Platinum 8488C. За всеки двигател и размер на входа събираме латентност на реалното време (p50 и p99, 10 000 итерации след 1 000 загрявания), завършени инструкции чрез Valgrind cachegrind, IPC и проценти на пропуски в кеша от хардуерните броячи на perf stat, както и заделяния на купчината, преброени от обвивка на глобален алокатор около измерения прозорец.

При 514 токена (+2 над 512 е инжектиране на BOS/EOS):

Реализация	p50 латентност	Инструкции	IPC	Алокации
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

Латентност на токенизатора, брой инструкции, IPC и заделяния при 514 токена.

SP и IREE постигат приблизително една трета от латентността на референцията. Те завършват около два пъти по-малко инструкции, работят при по-висок IPC и заделят с един до два порядъка по-малко.

SentencePiece и IREE използват двойно масивно дърво, въведено от Аое през 1989 г. (Aoe, 1989; Karoonboonyanan). Дървото е кодирано в два плоски цели масива, base и check:

cpp

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

Търсенето на дете е:

cpp

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

Всяка стъпка от обхождането на дървото се нуждае само от две зареждания от масив, едно цяло число за събиране и едно сравнение, без хеширане, без преследване на указатели и без алокация на купа за възел. Конструкцията поставя всеки възел така, че неговите деца да попадат на свободни слотове в масива за валидиране, позволявайки на множество родители да споделят региони, стига техните деца да се различават по байта, който следват.

За речника от 250K на XLM-RoBERTa цялото дърво се побира в ~9 MB непрекъсната памет (~780K слота × 12 байта за отместване, валидиране и ID на токен на слот). Работният набор на едно кодиране е от порядъка на 100 KB, удобно в L2, а хардуерният префетчър се фиксира върху редовния достъп.

Референцията, за разлика от тях, поддържа хеш-таблица с деца за всеки възел, което струва приблизително четири зависими зареждания на стъпка от байт. Хешира се байтът, зарежда се кофата, следва се указател към записа в кофата, следва се друг указател към дъщерния възел на купата, след което се прави справка в собствената му хеш-таблица за следващия байт. Обектите в купата се разпръскват и повечето зареждания пропускат кеша.

SP и IREE съхраняват ID-то на токена директно във възела на дървото. Референтната реализация на дървото връща само съвпадналите байтове, така че енкодерът трябва да задели нов низ за всяко съвпадение, да го потърси в странична хеш-таблица, за да възстанови ID-то на токена, и да прочете оценката от отделен вектор. SP връща ID-то като съхранена стойност на дървото и чете оценката от плосък масив с плаваща запетая:

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 използва API с курсор с поддържане на състояние вместо да връща всички съвпадения в едно извикване, но моделът на разходите е същият: ID-тата на токените и оценките са достъпни като индекси на плосък масив, без алокация или хеширане при самото обхождане на дървото. IREE заделя само един изходен буфер на кодиране; SP допълнително заделя въззел от решетката (lattice node) за всяко съвпадение по време на DP, общо приблизително 1 559 алокации при 514 токена.

Преди да променим дървото, изолираме цената на алгоритъма от цената на реализацията: колко бързо работи референцията с премахната ненужна работа?

Базовата линия запазва дървото с хеш-таблица от деца на референцията, но прави две промени. Собствена буферна структура на извикващия притежава DP таблицата, предварително обработените байтове и изходните ID-та на токените, така че горещият път не прави никакви алокации след загряване. ID-то и оценката на токена се съхраняват директно във възела на дървото, премахвайки преобразуването на низове на съвпадение и търсенето в страничната хеш-таблица.

Вътрешен цикъл:

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

Това е същият алгоритъм, със същата форма на дървото и същия модел на достъп, само че без алокации и без излишни търсения. Резултати при 514 токена:

Енкодер	p50	Инструкции	IPC	Алокации
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
Базова линия	156 µs	1.53M	2.62	0

Базовата линия без заделяния в сравнение със съществуващите имплементации на токенизатора при 514 токена.

Премахването на материализацията на низове на съвпадение и второстепенното търсене в хеш-таблица намалява инструкциите 2.4 пъти спрямо референцията. IPC спада с тях (2.86 → 2.62), тъй като преследването на указатели в хеш-таблицата, преди скривано зад работата на алокатора, сега е доминиращата цена на байт.

p50 и p99 латентност на референтния токенизатор и базовата линия без заделяния според размера на входа.

Ниво на пропуски в L2 кеша на референтния и базовия токенизатор без заделяния според размера на входа.

Базовата линия намалява p50 латентността с 2-3 пъти спрямо референцията при всеки размер, а процентът на пропуски в L2 остава близо до 2-3% до 2k токена спрямо 8-48% при референцията. Оставащата цена е преследването на указатели в HashMap, с което се заема следващият раздел.

Дървото с хеш-таблица е преследване на указатели. SP и IREE решават това с двумасивно дърво (вижте предходния раздел) и ние заимстваме същата структура от данни. Разликата е, че SP и IREE са библиотеки с общо предназначение, поддържащи стрийминг на вход, изходни отмествания и множество режими на кодиране; тази гъвкавост им струва допълнително счетоводство около дървото (решетка в SP, осемстепенен конвейер в IREE). Ние не се нуждаем от нито едно от тях, затова вграждаме дървото директно в цикъла на Витерби и изхвърляме останалото.

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

Една стъпка по байт струва две четения от масив, събиране на цяло число и сравнение. Цялото дърво е около 9 MB, но всяко отделно кодиране обхожда само ~100 KB от него, така че горещите данни остават в L2.

Енкодер	p50	p99	Инструкции	L2 пропуски %	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%
Базова линия	155 µs	1 170 µs	1.53M	2.59%	16.66%
+ Darts	68 µs	80 µs	1.10M	6.45%	11.42%

Латентност на токенизатора, инструкции и нива на пропуски в кеша след въвеждане на double-array trie.

Реалното време спада 2.3 пъти спрямо базовата линия и 4.8 пъти спрямо референцията. Инструкциите спадат с още 28% спрямо базовата линия. Премахването на решетката и конвейера ни позволява да изпълняваме 1.10 млн. инструкции на кодиране срещу 1.82 млн. на SP и 1.61 млн. на IREE при същата форма на дървото; разликата е счетоводството, което те поддържат за общоцелеви API-та, от които не се нуждаем.

Дървото Darts от Оптимизация 1 се нуждае от две зависими зареждания на стъпка по байт. Първо базовото отместване на родителя, след това записа check в слота на кандидата за дете, за да се потвърди, че родителят го притежава. Записът check служи само за валидиране на собствеността. Същият въпрос може да бъде отговорен от битмап за всеки възел, който записва кои байтове имат деца. Чрез конструкцията на двойно масивното дърво, битът за байт b в битмапа на родителя е зададен точно тогава, когато кандидат-слотът за този байт е притежаван от родителя. Битмапът и масивът check носят една и съща информация, така че можем напълно да премахнем масива check от структурата по време на изпълнение.

Търсенето в битмапа се компилира до едно тестване на бит срещу една 64-битова дума, което се изпълнява за около цикъл. Свързаната инструкция POPCNT, която броя всички зададени битове в регистър, е това, което други дизайни на индексирани с битмап дървета използват за изчисляване на отместване на дете директно от битмапа. Не се нуждаем от нея, защото конструкцията на Darts вече съхранява това отместване в полето base на възела.

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

Между другото, опаковаме четирите полета на възел (битмап, base, ID на токен и оценка) в една 64-байтова кеш линия. Едно посещение на възел вече изтегля всичко: оценката и ID-то на токена, от които се нуждаем в крайните възли, плюс битмапа и base за следващата стъпка. Масивът check се използва само по време на изграждането и е отстранен от разположението по време на изпълнение. Индексирането в горещия път е без проверки (unchecked), защото изграждането на дървото и структурата на цикъла вече гарантират границите. Компромисът е в пространството: при 64 байта на възел × ~780K слота дървото нараства от ~9 MB до ~50 MB.

Резултати при 514 токена:

Енкодер	p50	p99	Инструкции	L2 достъпи	L2 пропуски %	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%
Базова линия	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%
+ Битмап и вградено опаковане	65 µs	77 µs	1.06M	1.8K	20.53%	15.01%

Резултати за токенизатора при 514 токена след добавяне на bitmap и вградена организация.

Реалното време систематично спада с още 4.5% при 514 токена и нараства с дължината на входа, достигайки 6-7% при 4k и 16k. Премахването на масива check превръща стъпката през дървото на байт от две зареждания на кеш линия в едно, което е 2.5 пъти по-малко L2 достъпи на кодиране (1.8K спрямо 4.6K при обикновения Darts). Процентът на пропуските скочи от 6% на 20%, но по-голямата част от това е знаменателят: с по-малко от половината достъпи, пропуските, които все пак имаме, заемат по-голям дял.

Данните на дървото са около 50 MB (780 000 въззела по 64 байта всеки). В стандартен Linux процес тази памет се поддържа от 4 KB страници, така че цялото дърво обхваща около 12 000 отделни виртуални страници, далеч над това, което буферът за транслация на адреси (TLB) от първо ниво на процесора може да побере (Sapphire Rapids има 96 записа). Всяка стъпка на Витерби докосва различен възел и при кодиране на 512 токена изгаряме приблизително 9 000 цикъла в обхождане на таблици със страници, около 3% от бюджета за кодиране.

Решението е да се поддържа дървото с 2 MB huge pages вместо 4 KB такива. Същите 50 MB сега обхващат 25 страници, побират се изцяло в Translation Lookaside Buffer, а самите директории на таблиците със страници спират да се конкурират с дървото за L2 кеш.

Linux има два начина за поддържане на паметта с 2 MB страници. Опортунистичният е прозрачни огромни страници (THP): вашата програма моли ядрото да надстрои региона на паметта от 4 KB на 2 MB страници, а фонова нишка на ядрото се опитва да го направи, когато има време и налично непрекъсната физическа памет. Другата алтернатива е да се задели дървото директно от предварително резервиран пул от огромни страници, който ядрото отделя при зареждане:

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

Пулът трябва да съществува на хоста (vm.nr_hugepages, зададен при зареждане). В нашите производствени екземпляри са резервирани 10 561 огромни страници; дървото консумира 24 от тях. След mmap, /proc/<pid>/smaps потвърждава региона: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Печалбата е 3–12% реално време в различните дължини на входа, последователна и пропорционална на дължината на работния набор:

Токени	Без огромни страници	С огромни страници	Δ
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%

Латентност на токенизатора с и без huge pages при различни размери на входа.

Най-голямата печалба е при 4k токена, където работният набор е достатъчно голям, че трафикът към таблиците със страници активно се е конкурирал с данните на дървото за L2 честотна лента. Отвъд 4k печалбата се свива, тъй като пропуските в L3 доминират и обхожданията на таблиците със страници стават по-малък дял от общото време.

Обединявайки трите оптимизации, ето крайната производителност на енкодера, измерена спрямо всеки друг токенизатор на RoBERTa, който тествахме. Наред с native C++ SentencePiece и C IREE бинарните файлове, направихме бенчмарк и на готови Rust крейтове-обвивки, към които би посегнал производствен потребител – sentencepiece и iree-tokenizer. Крейтът sentencepiece отчита 1.6-1.9 пъти по-висока латентност от native C++ бинарния файл при всеки размер на входа, като овърхедът идва от това, че обвивката заделя нов списък с частици от токени, всяка от които държи собствено копие на текста на частицата, при всяко кодиране. Крейтът iree-tokenizer оразмерява вътрешния си буфер спрямо всеки вход, което избягва постоянна алокация на кодиране, която native IREE бинарният файл заплаща независимо от размера на входа, и това е най-важно при къси входове, където алокацията иначе би доминирала над действителната работа по токенизация. При производствени дължини на входа постоянната цена се амортизира и собственият овърхед на обвивката от копирането на токените обратно в Rust избутва латентността ѝ над тази на native бинарния файл.

p50 и p99 латентност на encode за всяка имплементация на токенизатора и размер на входа.

Крайният енкодер изпълнява същия алгоритъм Витерби върху същия речник от 250 000 токена, не заделя нищо в горещия път и произвежда токено-точен изход спрямо референцията. Всяка стъпка по байт вече е едно зареждане на 64-байтова кеш линия, една проверка на битмап и едно събиране на цяло число, спрямо проверка в HashMap, извикване на алокатор и няколко клона за специални случаи в базовата реализация. Инструкциите за кодиране спаднаха от 3.66 млн. до 1.04 млн., което е 3.5 пъти спад. Реалното време съответства на това съотношение при къси входове и се разширява допълнително при дълги входове, където алокациите на токен на референцията започват да преливат L2 и L3.

На практика тази работа намали използването на процесора в слоя за предварителна обработка на нашия инферентен стек с 5-6 пъти и отсече двуцифрени милисекунди от латентността на реранкера.

Лъвският пай от намалението на латентността дойде от леки промени, които премахнаха излишната работа: премахване на алокациите на кодиране, замяна на дървото с хеш-таблица на децата с двумасивно дърво и премахване на овърхеда от общия API. Използването на съществуващия битмап на възел като проверка за валидност и поддържането на дървото с 2 MB huge pages комбинирано допринесоха с 10-20% върху по-големия преход към двумасивно дърво, в зависимост от размера на входа.

Готовите Rust крейтове-обвивки около SentencePiece и I/O демонстрират същата точка в обратна посока. Върху native C и C++ токенизаторите, FFI границата и алокациите на извикване, които обвивките въвеждат, са видими в числата за латентност. Изходният код на Rust за крайния енкодер Витерби е наличен на github.com/perplexityai/pplx-garden.

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

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

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

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

Мотивация

Алгоритъмът на Витерби

Референтна реализация

Бележка относно предварителната обработка

Къде референтната реализация харчи своето време

Как измерихме

Защо SentencePiece и IREE са по-бързи: Двумасивно дърво (Double-Array Trie)

Базова линия: Порт без алокации на референцията

Оптимизация 1: Двумасивно дърво без излишни абстракции

Оптимизация 2: Битова карта и вградено опаковане

Оптимизация 3: Огромни страници за трие структури

Крайно сравнение

Заключителни мисли

Препратки

Цитиране