Підвищення продуктивності CPU токенізатора Unigram

Ми переписали наш токенізатор Unigram з нуля в рамках цілеспрямованого проєкту з підвищення продуктивності. За довжин вхідних даних, що використовуються у продакшені, новий кодер знижує затримку p50 приблизно в 5 разів порівняно з крейтом токенізаторів Hugging Face, приблизно у 2 рази порівняно з SentencePiece (C++) та приблизно в 1,5 раза порівняно з токенізатором IREE (C), без жодних виділень пам'яті в купі у сталого стану. У продакшені це зменшило використання процесора в нашому інференсному стеку в 5–6 разів і скоротило затримку реранкера на двозначну кількість мілісекунд.

Витрати на інференс LLM зазвичай розглядаються через роботу графічного процесора: KV-кеші, ядра уваги, роутинг експертів. Але значна частина продакшен-трафіку працює на набагато менших моделях — на два-три порядки менших за передові трансформери. Моделі ембедингів, класифікатори та реранкери все ще виконуються на GPU, але прямий прохід (forward pass) достатньо короткий, і все, що навколо нього, починає відігравати важливу роль.

Найяскравішим прикладом є реранкер, який оцінює сотні кандидатів-документів на один запит. У разі невеликої моделі обчислення на GPU зазвичай завершуються за однозначну кількість мілісекунд, але кожен вхідний документ все одно повинен спочатку пройти токенізацію на CPU, перетворюючи текст на ідентифікатори словника. Коли батч великий, це може становити значну частку від загальної затримки запиту.

Ця робота зосереджена на XLM-RoBERTa — словнику Unigram із 250 тис. токенів, натренованому за допомогою SentencePiece. Доопрацьовані кодери сімейства RoBERTa є поширеним вибором для завдань ранжування, пошуку (retrieval) та визначення схожості. Далі ми наводимо однопотокову затримку на одне кодування; у продакшені ми розпаралелюємо обчислення за ядрами за допомогою rayon.

У цій статті ми розглядаємо крейт токенізаторів від 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) — це критичний шлях (hot path): на вхідних даних розміром 16 КБ він виконує сотні тисяч переходів і завершує десятки мільйонів інструкцій на одне кодування.

Токенізатори 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, але виконується багато прихованої від погляду на рівні вихідного коду роботи. У наступному розділі детально розбито, куди саме витрачається час.

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

Існує два способи прискорити код на CPU: виконувати менше інструкцій або досягти вищого показника IPC. Cachegrind вимірює перше, perf stat вимірює друге. Кожна оптимізація в цій статті впливає на одне або інше.

Процесори зупиняються через чотири речі: промахи кешу (L1d — приблизно 4 цикли, промах L3 із зверненням до DRAM — сотні), ланцюжки залежних завантажень, які механізм позачергового виконання не може розпаралелити, хибні передбачення розгалужень (приблизно 15–20 циклів скидання за кожну неправильну спробу) та виклики розподільника пам'яті (allocator). Еталонна реалізація рушія unigram стикається з усіма чотирма проблемами:

Вузьке місце	Механізм	Виміряний вплив
Виділення пам'яті на кожен збіг	String::from_utf8 для збіжних байтів плюс пошук у AHashMap для кожного збігу в дереві; ідентифікатор токена вже міститься в листі, але еталонна реалізація не зберігає його там	7295 виділень для 512 токенів, 14 542 для 1 КБ, 299 171 для 16 КБ
Пошук за вказівниками на кожен байт	AHashMap у кожному вузлі дерева, хешування наступного байта та розіменування двох вказівників на кожному кроці; структура Node також групує неактивні поля, такі як is_leaf, підтягуючи їх до кешу під час кожного відвідування	Частка промахів L1d залишається низькою (~2,5% при 1 КБ), оскільки гарячі сторінки розподільника знаходяться в L1; витрати проявляються у вигляді затримки залежних завантажень
Транінг L2 на довгих вхідних даних	Таблиця DP, буфер префіксів та список вихідних токенів щоразу виділяються заново під час кожного виклику	Частка промахів L2 зростає з 12% при 128 токенах до 50% при 16 КБ, оскільки сукупний обсяг виділеної пам'яті перевищує обсяг L2 (2 МБ на ядро)

Довідкові вузькі місця токенізатора, механізми та виміряний вплив.

Основною витратою є виділення пам'яті на кожен збіг. За 1 КБ токенів еталонна реалізація виконує (завершує) приблизно 7 млн інструкцій на кодування при IPC 2.8, більшість з яких припадає на виклики розподільника пам'яті та перетворення рядків. Промахи L2 стають важливими лише після ~1 КБ токенів, коли сукупний обсяг виділеної пам'яті переповнює кеш на ядро.

p50 і p99 затримки довідкового токенізатора за розміром вхідних даних.

Частота промахів у кеші L2 довідкового токенізатора за розміром вхідних даних.

Кількість байтів, виділених на одне кодування довідковим токенізатором у різних розмірах вхідних даних.

Виділення пам'яті на один токен є константним: приблизно 2 КБ і ~18 виділень на токен, незалежно від розміру вхідних даних. Нелінійність затримки зумовлена L2: сукупні байти перетинають позначку 2 МБ приблизно на 1 КБ токенів, частка промахів зростає з 8% до ~50%, а крива p99 вигинається відповідно.

Ми порівнюємо три продакшен-токенізатори, які використовують той самий алгоритм Unigram для тієї самої моделі XLM-RoBERTa:

Рушій	Мова	Префіксне дерево (Trie)
Hugging Face	Rust	Префіксне дерево з дочірніми елементами у вигляді HashMap
SentencePiece	C++	Двомасивно-масивне дерево Darts (double-array trie)
Токенізатор IREE	C	Двомасивно-масивне дерево Darts

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

Ми підключаємо SentencePiece та IREE через прямий інтерфейс FFI C/C++, а не через їхні обгортки на Rust; обгортки додають виділення пам'яті на кожен виклик та перетворення рядків, які не є частиною витрат базової бібліотеки.

Усі рушії зібрані з повними оптимізаціями (O3, цільова оптимізація під нативний процесор, оптимізація під час компонування — LTO). Вимірювання виконувалися на одному закріпленому тихому ядрі процесора 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 мкс	3,60 млн	2,89	7295
SentencePiece	128 мкс	1,83 млн	3,39	1559
IREE	112 мкс	2,28 млн	3,06	1

Затримка токенізатора, кількість інструкцій, IPC та алокації для 514 токенів.

SP та IREE мають затримку приблизно втричі меншу, ніж у еталонної реалізації. Вони виконують приблизно вдвічі менше інструкцій, працюють з вищим IPC та виділяють на один-два порядки менше пам'яті.

SentencePiece та IREE використовують двомасивне префіксне дерево (double-array trie), представлене Аое у 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 */ }

Кожен крок обходу дерева потребує лише двох завантажень з масиву, одного додавання цілих чисел та одного порівняння, без хешування, без переслідування вказівників та без виділення пам'яті в купі для кожного вузла. Така побудова розміщує кожен вузол так, щоб його дочірні елементи потрапляли на вільні слоти в масиві перевірки, що дозволяє кільком батьківським вузлам спільно використовувати регіони, якщо їхні дочірні елементи відрізняються за байтом, за яким вони слідують.

Для словника XLM-RoBERTa розміром 250 КБ усе дерево поміщається приблизно у 9 МБ неперервної пам'яті (~780 тис. слотів × 12 байтів для зміщення, перевірки та ідентифікатора токена на слот). Активний робочий набір одного кодування становить порядку 100 КБ, що комфортно поміщається в L2, а апаратний механізм вибірки (prefetcher) налаштовується на регулярний шаблон доступу.

Натомість еталонна реалізація зберігає хеш-таблицю дочірніх елементів для кожного вузла, що коштує приблизно чотири залежні завантаження на байтовий крок. Захешувати байт, завантажити бакет, перейти за вказівником на елемент бакета, перейти за іншим вказівником на дочірній вузол у купі, а потім звернутися до його власної хеш-таблиці для наступного байта. Об'єкти в купі розкидані, і більшість завантажень призводять до промахів кешу.

SP та IREE зберігають ідентифікатор токена безпосередньо у вузлі дерева. Еталонна реалізація повертає лише знайдені байти, тому кодер змушений виділяти новий рядок для кожного збігу, шукати його в побічній хеш-таблиці для відновлення ідентифікатора токена та зчитувати оцінку з окремого вектору. SP повертає ідентифікатор як збережене значення дерева та зчитує оцінку з плоского масиву чисел із плавною комою:

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 з підтримкою стану (stateful cursor API) замість повернення всіх збігів за один виклик, але модель витрат така ж сама: ідентифікатори токенів та оцінки доступні як індекси плоского масиву, без виділення пам'яті чи хешування під час самого обходу дерева. IREE виділяє лише один вихідний буфер на кодування; SP додатково виділяє вузол решітки на кожен збіг під час динамічного програмування, що становить загалом близько 1559 виділень для 514 токенів.

Перш ніж змінювати дерево, ми ізолюємо вартість алгоритму від вартості реалізації: наскільки швидко працює еталонний код, якщо видалити непотрібну роботу?

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

Внутрішній цикл:

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 мкс	3,66 млн	2,86	7295
SentencePiece	127 мкс	1,81 млн	3,55	1559
IREE	111 мкс	1,58 млн	3,26	1
Базова лінія	156 мкс	1,53 млн	2,62	0

Базова модель без алокацій у порівнянні з наявними реалізаціями токенізатора для 514 токенів.

Усунення матеріалізації рядків для кожного збігу та вторинного пошуку в хеш-таблиці скорочує кількість інструкцій у 2,4 раза порівняно з еталоном. IPC при цьому падає (з 2,86 до 2,62), оскільки пошук за вказівниками в хеш-таблиці, який раніше приховувався за роботою розподільника пам'яті, тепер є основним джерелом витрат на байт.

p50 і p99 затримки довідкового токенізатора та базового токенізатора без алокацій за розміром вхідних даних.

Частота промахів у кеші L2 довідкового токенізатора та базового токенізатора без алокацій за розміром вхідних даних.

Базова лінія скорочує затримку p50 у 2–3 рази порівняно з еталоном за будь-якого розміру, а частка промахів L2 залишається на рівні 2–3% до 2 КБ токенів проти 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 МБ, але під час кодування обходиться лише близько 100 КБ з них, тому гарячі дані залишаються в L2.

Рушій	p50	p99	Інструкції	Промахи L2 %	Промахи L3 %
Hugging Face	327 мкс	3340 мкс	3,66 млн	8,17%	14,66%
SentencePiece	127 мкс	1140 мкс	1,82 млн	4,74%	4,07%
IREE	111 мкс	132 мкс	1,61 млн	5,82%	3,87%
Базова лінія	155 мкс	1170 мкс	1,53 млн	2,59%	16,66%
+ Darts	68 мкс	80 мкс	1,10 млн	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, ідентифікатор токена та оцінку) в один 64-байтний рядок кешу (cache line). Одне відвідування вузла тепер завантажує все: оцінку та ідентифікатор токена, які нам потрібні у кінцевих вузлах, а також бітову маску та base для наступного кроку. Масив check використовується лише під час побудови і виключається з макета виконання. Індексація на критичному шляху не перевіряється, оскільки побудова дерева та структура циклу вже гарантують дотримання меж. Компромісом є простір: при 64 байтах на вузол і приблизно 780 тис. слотів розмір дерева зростає з ~9 МБ до ~50 МБ.

Результати для 514 токенів:

Рушій	p50	p99	Інструкції	Звернення до L2	Промахи L2 %	Промахи L3 %
Hugging Face	327 мкс	3340 мкс	3,66 млн	15,0 К	8,17%	14,66%
SentencePiece	127 мкс	1140 мкс	1,82 млн	7,5 К	4,74%	4,07%
IREE	111 мкс	132 мкс	1,61 млн	4,7 К	5,82%	3,87%
Базова лінія	155 мкс	1170 мкс	1,53 млн	6,1 К	2,59%	16,66%
+ Darts	68 мкс	80 мкс	1,10 млн	4,6 К	6,45%	11,42%
+ Бітова маска та вбудоване пакування	65 мкс	77 мкс	1,06 млн	1,8 К	20,53%	15,01%

Результати токенізатора для 514 токенів після додавання bitmap та inline packing.

Час виконання за годинником систематично зменшується ще на 4,5% для 514 токенів і зростає зі збільшенням довжини вхідних даних, досягаючи 6–7% при 4 КБ і 16 КБ. Відмова від масиву check скорочує кількість завантажень рядків кешу на байтовий крок дерева з двох до одного, що забезпечує у 2,5 раза менше звернень до L2 на кодування (1,8 К проти 4,6 К у звичайному Darts). Відсоток промахів стрибає з 6% до 20%, але більша частина цього припадає на знаменник: маючи більш ніж удвічі менше звернень, наявні промахи починають становити більшу частку.

Дані дерева становлять близько 50 МБ (780 000 вузлів по 64 байти кожен). У стандартному процесі Linux ця пам'ять підтримується сторінками по 4 КБ, тому повне дерево охоплює близько 12 000 окремих віртуальних сторінок, що набагато перевищує місткість первинного буфера асоціативної трансляції даних процесора (TLB у Sapphire Rapids має 96 елементів). Кожен крок Вітербі торкається іншого вузла, і під час кодування 512 токенів ми витрачаємо приблизно 9000 циклів на обхід таблиць сторінок, що становить близько 3% від бюджету на одне кодування.

Вирішення полягає в тому, щоб підтримувати дерево за допомогою величезних сторінок (huge pages) розміром 2 МБ замість сторінок по 4 КБ. Ті самі 50 МБ тепер охоплюють 25 сторінок, повністю поміщаються в буфері TLB, а самі каталоги таблиць сторінок перестають конкурувати з деревом за кеш L2.

У Linux є два способи підтримки пам'яті за допомогою сторінок по 2 МБ. Опортуністичним методом є прозорі величезні сторінки (THP): ваша програма просить ядро підвищити регіон пам'яті зі сторінок 4 КБ до 2 МБ, і фоновий тред ядра намагається зробити це, коли має час і доступну неперервну фізичну пам'ять. Іншою альтернативою є виділення пам'яті для дерева безпосередньо з попередньо зарезервованого пулу величезних сторінок, який ядро виділяє під час завантаження:

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 мкс	15,2 мкс	−3,6%
514	65,4 мкс	63,1 мкс	−3,4%
1026	129,5 мкс	125,3 мкс	−3,3%
4098	773 мкс	679 мкс	−12,0%
16 386	3808 мкс	3554 мкс	−6,8%

Затримка токенізатора з великими сторінками та без них у різних розмірах вхідних даних.

Найбільший виграш спостерігається на 4 КБ токенів, де робочий набір достатньо великий, щоб трафік таблиць сторінок активно конкурував із даними дерева за пропускну здатність L2. Понад 4 КБ виграш зменшується, оскільки домінують промахи L3, а обхід таблиць сторінок стає меншою часткою загального часу.

Об'єднавши ці три оптимізації, нижче наведено фінальний рушій у порівнянні з усіма іншими токенізаторами RoBERTa, які ми тестували. Поряд з нативними бінарними файлами SentencePiece на C++ та IREE на C, ми також провели бенчмаркінг готових крейтів-обгорток на Rust, якими скористався б продакшен-розробник: sentencepiece та iree-tokenizer. Крейт sentencepiece демонструє затримку на 1.6–1.9 раза вищу, ніж нативний бінарний файл C++ за будь-якого розміру вхідних даних; накладні витрати виникають через те, що обгортка виділяє новий список частин токенів на кожне кодування, кожна з яких містить власну копію тексту частини. Крейт iree-tokenizer підганяє розмір свого внутрішнього буфера під кожен вхідний файл, що дозволяє уникнути постійного виділення пам'яті на кожне кодування, яке нативний бінарний файл IREE сплачує незалежно від розміру вхідних даних, і це має найбільше значення на коротких вхідних даних, де виділення пам'яті інакше домінувало б над фактичною роботою з токенізації. За довжин вхідних даних, що використовуються у продакшені, постійні витрати амортизуються, а власні накладні витрати обгортки на виклик від копіювання токенів назад у Rust піднімають її затримку вище показників нативного бінарного файлу.

p50 і p99 затримки на одне кодування для реалізацій токенізатора та розмірів вхідних даних.

Фінальний рушій виконує той самий алгоритм Вітербі для того ж словника на 250 000 токенів, не виділяє пам'ять на критичному шляху та створює побайтово точний вихідний результат порівняно з еталоном. Кожен байтовий крок тепер є одним завантаженням 64-байтного рядка кешу, однією перевіркою бітової маски та одним додаванням цілих чисел (порівняно з опитуванням HashMap, викликом розподільника та кількома гілками для особливих випадків у базовій реалізації). Кількість інструкцій на одне кодування впала з 3,66 млн до 1,04 млн — у 3,5 раза. Час виконання за годинником відповідає цьому співвідношенню на коротких вхідних даних і ще більше зростає на довгих, де виділення пам'яті на токен у еталона починає переповнювати L2 та L3.

На практиці ця робота зменшила використання процесора в шарі попередньої обробки нашого інференсного стеку в 5–6 разів і скоротила затримку реранкера на двозначну кількість мілісекунд.

Основна частина зниження затримки була досягнута завдяки легким змінам, які усунули марну роботу: відмові від виділення пам'яті на кожне кодування, заміні дерева з дочірніми елементами у вигляді HashMap на двомасивне дерево та усуненню накладних витрат універсального API. Використання існуючої бітової маски вузла як перевірки валідності та підтримка дерева за допомогою величезних сторінок 2 МБ у поєднанні забезпечили 10–20% прискорення на додачу до ширшого переходу на двомасивне префіксне дерево, залежно від розміру вхідних даних.

Готові крейти-обгортки на Rust навколо SentencePiece та IREE ілюструють те саме зворотно. Окрім нативних токенізаторів на C і C++, межа FFI та виділення пам'яті на кожен виклик, які вводять обгортки, помітні в показниках затримки. Вихідний код на Rust для фінального кодера Вітербі доступний за адресою github.com/perplexityai/pplx-garden.

Токенізатори Hugging Face — 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: Величезні сторінки для префіксного дерева

Фінальне порівняння

Підсумкові думки

Посилання

Цитування