Повышение производительности CPU Unigram-токенизатора

Мы с нуля переписали наш Unigram-токенизатор в рамках проекта, направленного на повышение производительности. При стандартных производственных длинах входных данных новый кодировщик сокращает p50-задержку примерно в 5 раз по сравнению с библиотекой Hugging Face tokenizers, примерно в 2 раза по сравнению с SentencePiece (C++) и примерно в 1,5 раза по сравнению с токенизатором IREE (C), без каких-либо выделений памяти в куче в установившемся режиме. В промышленной эксплуатации это сократило использование ресурсов процессора в нашем стеке логического вывода в 5–6 раз и уменьшило задержку ранжирования на десятки миллисекунд.

Стоимость логического вывода LLM обычно рассматривается как нагрузка на GPU: KV-кэши, механизмы внимания, маршрутизация экспертов. Однако значительная часть производственного трафика приходится на гораздо более компактные модели, на два-три порядка меньше, чем передовые трансформеры. Модели эмбеддингов, классификаторы и ранжировщики по-прежнему работают на GPU, но прямой проход настолько короток, что всё, что его окружает, начинает иметь значение.

Наиболее показательный пример — ранжировщик, который оценивает сотни документов-кандидатов на запрос. С небольшой моделью вычисления на GPU обычно завершаются за несколько миллисекунд, но каждый входной сигнал всё равно должен сначала пройти через токенизацию на CPU, преобразующую текст в идентификаторы словаря. Когда пакет большой, это может составлять значительную часть общей задержки запроса.

Эта работа ориентирована на XLM-RoBERTa, Unigram-словарь из 250 тыс. токенов, обученный с помощью 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 г.), где байтовые позиции выступают в роли слоев графа, а токены словаря — в роли ребер, охватывающих непрерывный байтовый диапазон. Рекуррентность представляет собой одномерное динамическое программирование по байтовым позициям:

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). Проход по внутреннему префиксному дереву — это «горячий» путь: на входе из 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 циклов на сброс при неверном догадке) и вызовы аллокатора. Эталонная реализация движка unigram сталкивается со всеми четырьмя:

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

Справочная информация о узких местах токенизатора, механизмах и измеренном влиянии.

Доминирующей статьей расходов является выделение памяти на каждое совпадение. При 1 тыс. токенов эталонная реализация завершает примерно 7 млн инструкций на кодирование при IPC 2,8, большая часть из которых приходится на вызовы аллокатора и преобразование строк. Промахи L2 имеют значение только при более чем ~1 тыс. токенов, когда совокупный объем памяти начинает переполнять кэш на ядро.

p50 и p99 задержка справочного токенизатора в зависимости от размера ввода.

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

Количество байтов, выделяемых на одно кодирование справочным токенизатором для разных размеров ввода.

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

Мы сравниваем три производственных токенизатора, работающих по алгоритму Unigram с той же моделью XLM-RoBERTa:

Движок	Язык	Префиксное дерево
Hugging Face	Rust	HashMap-children префиксное дерево
SentencePiece	C++	Darts double-array префиксное дерево
IREE токенизатор	C	Darts double-array префиксное дерево

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

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

Все движки собраны с полными оптимизациями (O3, нативная целевая архитектура, оптимизация времени компоновки). Измерения проводились с привязкой к одному свободному ядру на 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	7 295
SentencePiece	128 мкс	1,83 млн	3,39	1 559
IREE	112 мкс	2,28 млн	3,06	1

Задержка токенизатора, число инструкций, IPC и выделения памяти при 514 токенах.

SP и IREE достигают примерно трети задержки эталона. Они выполняют примерно в два раза меньше инструкций, работают при более высоком IPC и выделяют на 1–2 порядка меньше памяти.

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

Для каждого шага обхода дерева требуется только две загрузки из массива, одно сложение целых чисел и одно сравнение, без хеширования, поиска по указателям и выделения памяти в куче на узел. При построении каждый узел размещается так, чтобы его потомки попадали на свободные слоты в массиве проверки, позволяя нескольким родителям совместно использовать области памяти, если их потомки различаются по следующему за ними байту.

Для словаря XLM-RoBERTa из 250 тыс. токенов всё дерево умещается в ~9 МБ непрерывной памяти (~780 тыс. слотов × 12 байт для смещения, проверки и ID токена на слот). Горячий рабочий набор одного кодирования составляет около 100 КБ, что комфортно помещается в 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 дополнительно выделяет узел решетки на совпадение во время DP, итого около 1 559 аллокаций при 514 токенах.

Прежде чем менять дерево, мы изолируем алгоритмическую стоимость от стоимости реализации: насколько быстро работает эталон с удаленной ненужной работой?

Базовая линия сохраняет префиксное дерево эталона с хеш-картой детей, но вносит два изменения. Структура scratch, принадлежащая вызывающей стороне, содержит таблицу 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 мкс	3,66 млн	2,86	7 295
SentencePiece	127 мкс	1,81 млн	3,55	1 559
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, который рассматривается в следующем разделе.

Дерево на базе 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 мкс	3 340 мкс	3,66 млн	8,17%	14,66%
SentencePiece	127 мкс	1 140 мкс	1,82 млн	4,74%	4,07%
IREE	111 мкс	132 мкс	1,61 млн	5,82%	3,87%
Базовая линия	155 мкс	1 170 мкс	1,53 млн	2,59%	16,66%
+ Darts	68 мкс	80 мкс	1,10 млн	6,45%	11,42%

Задержка токенизатора, инструкции и частота промахов кэша после внедрения двойного массивного trie.

Задержка падает в 2,3 раза от базовой и в 4,8 раза от эталонной. Количество инструкций падает ещё на 28% от базовой. Отсечение решетки и конвейера позволяет нам выполнять 1,10 млн инструкций на кодирование против 1,82 млн у SP и 1,61 млн у IREE при той же форме дерева; разрыв — это учет, который они ведут для API общего назначения, который нам не нужен.

Дерево Darts из Оптимизации 1 требует две зависимые загрузки на каждый байтовый шаг. Сначала базовое смещение родителя, затем запись проверки в слоте потенциального потомка, чтобы подтвердить, что родитель владеет им. Запись проверки только подтверждает владение. Тот же вопрос может быть решен битовой картой каждого узла, которая записывает, какие байты имеют потомков. По конструкции дерева на базе двойного массива бит для байта b в битовой карте родителя установлен именно тогда, когда потенциальный слот для этого байта принадлежит родителю. Битовая карта и массив проверки несут одну и ту же информацию, поэтому мы можем полностью удалить массив проверки из структуры времени выполнения. Битовая карта и массив проверки несут одну и ту же информацию, поэтому мы можем полностью удалить массив проверки из структуры.

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

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

Пока мы этим занимаемся, мы упаковываем четыре поля каждого узла (битовая карта, база, ID токена и значение) в одну 64-байтную строку кэша. Теперь одно посещение узла подтягивает всё: значение и ID токена, необходимые в терминальных узлах, плюс битовую карту и базу для следующего шага. Массив проверки используется только при построении и удаляется из макета времени выполнения. Индексация горячего пути не проверяется, так как структура дерева и цикла уже гарантируют границы. Компромисс — Пространство: при 64 байтах на узел × ~780 тыс. слотов дерево вырастает с ~9 МБ до ~50 МБ.

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

Движок	p50	p99	Инструкции	L2 доступы	L2 промах%	L3 промах%
Hugging Face	327 мкс	3 340 мкс	3,66 млн	15,0 тыс.	8,17%	14,66%
SentencePiece	127 мкс	1 140 мкс	1,82 млн	7,5 тыс.	4,74%	4,07%
IREE	111 мкс	132 мкс	1,61 млн	4,7 тыс.	5,82%	3,87%
Базовая линия	155 мкс	1 170 мкс	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 и встроенной упаковки.

Задержка систематически падает ещё на 4,5% при 514 токенах и растет с длиной ввода, достигая 6–7% при 4 тыс. и 16 тыс. Отказ от массива проверки превращает каждый побайтовый шаг дерева из двух загрузок строки кэша в одну, что в 2,5 раза меньше доступов L2 на кодирование (1,8 тыс. против 4,6 тыс. у простого Darts). Частота промахов прыгает с 6% до 20%, но большая часть этого — знаменатель: при менее чем в два раза меньшем количестве доступов те промахи, которые у нас есть, занимают большую долю.

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

Решение — поддерживать дерево с помощью 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% в реальном времени в зависимости от размеров входных данных, что согласуется и пропорционально длине рабочего набора:

Токены	Без huge pages	С huge pages	Δ
130	15,7 мкс	15,2 мкс	−3,6%
514	65,4 мкс	63,1 мкс	−3,4%
1 026	129,5 мкс	125,3 мкс	−3,3%
4 098	773 мкс	679 мкс	−12,0%
16 386	3 808 мкс	3 554 мкс	−6,8%

Задержка токенизатора с большими страницами и без них для разных размеров ввода.

Самый большой выигрыш наблюдается на 4 тыс. токенов, где рабочий набор достаточно велик, чтобы трафик таблиц страниц активно конкурировал с данными дерева за полосу пропускания L2. За 4 тыс. прирост сокращается, потому что промахи L3 доминируют, а проходы по таблицам страниц становятся меньшей долей общего времени.

Объединяя три оптимизации, вот итоговый движок, измеренный против всех остальных токенизаторов RoBERTa, которые мы тестировали. Наряду с нативными бинарными файлами SentencePiece на C++ и IREE на C, мы также провели бенчмарк готовых оберток crate для Rust, которые использовал бы производственный пользователь, sentencepiece и iree-tokenizer. Crate sentencepiece показывает задержку в 1,6–1,9 раза выше, чем нативный бинарный файл C++ при любом размере входных данных, при этом накладные расходы возникают из-за того, что обертка выделяет новый список токенов, каждый из которых содержит собственную копию текста фрагмента, при каждом кодировании. Crate iree-tokenizer определяет размер своего внутреннего буфера для каждого ввода, что позволяет избежать постоянного выделения памяти на кодирование, которое оплачивает нативный бинарный файл IREE независимо от размера ввода, и это наиболее важно при коротких входных данных, где выделение памяти в противном случае доминировало бы над фактической работой по токенизации. При производственных длинах ввода постоянная стоимость амортизируется, и собственные накладные расходы обертки на вызов при копировании токенов обратно в Rust выталкивают задержку выше, чем у нативного бинарного файла.

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

Итоговый движок выполняет тот же алгоритм Витерби по тому же словарю из 250 000 токенов, ничего не выделяет на горячем пути и выдает токен-точный вывод по сравнению с эталоном. Каждый побайтовый шаг теперь является одной 64-байтной загрузкой строки кэша, одним тестом битовой карты и одним сложением целых чисел, вместо проверки HashMap, вызова аллокатора и нескольких ветвей специальных случаев в базовой реализации. Инструкции на кодирование упали с 3,66 млн до 1,04 млн, что в 3,5 раза меньше. Время в реальном времени соответствует этому соотношению при коротких входных данных и расширяется дальше при длинных, где выделения памяти эталона на токен начинают переполнять L2 и L3.

На практике эта работа сократила использование CPU в слое предобработки нашего стека логического вывода в 5–6 раз и уменьшила задержку ранжировщика на десятки миллисекунд.

Основная часть сокращения задержки пришла от легковесных изменений, которые удалили бесполезную работу: отказ от аллокаций на кодирование, замена дерева HashMap-children на двойной массив и удаление накладных расходов общего API. Использование существующей битовой карты узла в качестве проверки валидности и поддержка дерева 2 МБ огромными страницами в сумме составили 10–20% сверх основного перехода на двойной массив дерева, в зависимости от размера входных данных.

Готовые обертки Rust вокруг SentencePiece и IREE иллюстрируют ту же точку зрения в обратном порядке. Поверх нативных токенизаторов C и C++ граница FFI и аллокации на вызов, вносимые обертками, видны в цифрах задержки. Исходный код Rust для итогового кодировщика Витерби доступен по адресу 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}
}

Мотивация

Алгоритм Витерби

Эталонная реализация

Заметка о предварительной обработке

На что тратит время эталонная реализация

Как мы измеряли

Почему SentencePiece и IREE работают быстрее: двойной массив префиксного дерева

Базовая линия: Порт эталона с нулевым выделением памяти

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

Оптимизация 2: Упаковка битмапов и встроенная упаковка

Оптимизация 3: Огромные страницы для префиксного дерева

Итоговое сравнение

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

Ссылки

Цитирование