Cải thiện Hiệu năng CPU của Trình mã hóa Token Unigram

Chúng tôi đã viết lại trình mã hóa token Unigram của mình từ đầu như một dự án tập trung vào hiệu năng. Với độ dài đầu vào trong môi trường production, bộ mã hóa mới giảm độ trễ p50 khoảng 5× so với crate tokenizers của Hugging Face, ~2× so với SentencePiece (C++) và ~1.5× so với trình mã hóa của IREE (C), mà không hề phân bổ heap ở trạng thái ổn định. Trong môi trường production, điều này giúp giảm lượng sử dụng CPU trong ngăn xếp suy luận (inference stack) của chúng tôi đi 5-6× và cắt giảm được độ trễ hàng chục mili-giây của mô hình xếp hạng lại (reranker).

Chi phí suy luận LLM thường được coi là công việc của GPU: bộ đệm KV, các nhân chú ý (attention kernels), việc định tuyến chuyên gia (expert routing). Nhưng một phần đáng kể lưu lượng truy cập trong môi trường production chạy trên các mô hình nhỏ hơn nhiều, nhỏ hơn từ hai đến ba bậc độ lớn so với các mô hình transformer tiên tiến. Các mô hình nhúng (embedding models), bộ phân loại (classifiers) và mô hình xếp hạng lại vẫn chạy trên GPU, nhưng quá trình truyền xuôi (forward pass) đủ ngắn để mọi thứ xung quanh bắt đầu trở nên quan trọng.

Trường hợp rõ ràng nhất là một mô hình xếp hạng lại chấm điểm hàng trăm tài liệu ứng viên cho mỗi yêu cầu. Với một mô hình nhỏ, quá trình tính toán trên GPU thường hoàn thành trong vòng chưa đầy mười mili-giây, nhưng mọi đầu vào vẫn phải trải qua quá trình mã hóa token ở phía CPU trước, chuyển đổi văn bản thành các ID từ vựng. Khi lô (batch) lớn, điều đó có thể chiếm một phần đáng kể trong tổng độ trễ của yêu cầu.

Công việc này nhắm tới XLM-RoBERTa, một từ vựng Unigram 250K token được huấn luyện bằng SentencePiece. Các bộ mã hóa thuộc họ RoBERTa đã được tinh chỉnh là lựa chọn phổ biến cho các tác vụ xếp hạng, truy xuất và độ tương đồng. Chúng tôi báo cáo độ trễ đơn luồng trên mỗi lần mã hóa xuyên suốt; trong môi trường production, chúng tôi song song hóa trên các lõi bằng cách sử dụng rayon.

Xuyên suốt bài viết này, chúng tôi coi crate tokenizers của Hugging Face là tiêu chuẩn tham chiếu để đo lường điểm chuẩn, vì đây là trình mã hóa token Rust mặc định mà hầu hết các nhóm đều sử dụng. SentencePiece (C++) và IREE (C) được đưa vào để so sánh với tư cách là hai triển khai chính khác trong môi trường production.

Quá trình mã hóa token Unigram, được Kudo giới thiệu vào năm 2018 (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) và được tích hợp trong SentencePiece, coi việc phân đoạn là một bài toán tìm đường đi có xác suất cao nhất: mỗi token từ vựng có một xác suất logarit đã học, và bộ mã hóa chọn phân đoạn có tổng điểm các token đạt giá trị cao nhất. Bài toán được giải quyết bằng thuật toán Viterbi (1967), với các vị trí byte là các lớp của đồ thị và các token từ vựng là các cạnh trải dài trên một phạm vi byte liền kề. Công thức truy hồi là một quy hoạch động (DP) một chiều trên các vị trí byte:

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

Vòng lặp ngoài có tính tuyến tính theo độ dài đầu vào, bị giới hạn ở mỗi vị trí bởi token từ vựng dài nhất (48 byte UTF-8 đối với XLM-RoBERTa). Quá trình duyệt cây trie bên trong là đường dẫn nóng (hot path): trên đầu vào 16k token, nó chạy hàng trăm nghìn lần chuyển đổi và hoàn thành hàng chục triệu chỉ thị trên mỗi lần mã hóa.

tokenizers v0.22 của Hugging Face triển khai trực tiếp phép truy hồi, sử dụng bảng băm (hash-map) để tra cứu ngược và các nút con dạng bảng băm tại mọi nút 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
}

Và vòng lặp bên trong (đã được đơn giản hóa) – mỗi kết quả khớp tìm kiếm tiền tố sẽ cấp phát một chuỗi mới và sau đó dò tìm trong bảng băm:

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
}

Mã nguồn là ngôn ngữ Rust dễ đọc, nhưng có rất nhiều công việc diễn ra mà không thể nhìn thấy ở cấp độ mã nguồn. Phần tiếp theo sẽ phân tích thời gian được tiêu tốn vào những đâu.

Trước khi Viterbi chạy, đầu vào phải trải qua quá trình chuẩn hóa Unicode (NFKC) và đánh dấu biên giới từ (word-boundary). Việc lập hồ sơ hiệu năng (profiling) cho thấy quá trình tiền xử lý chiếm chưa đến 5% tổng thời gian mã hóa trong các kịch bản của chúng tôi, vì vậy chúng tôi tập trung vào quá trình truyền xuôi Viterbi cho phần còn lại của bài viết.

Có hai cách để tăng tốc mã hóa CPU: thực hiện ít chỉ thị hơn, hoặc đạt được IPC cao hơn. Cachegrind đo lường cách thứ nhất, perf stat đo lường cách thứ hai. Mọi tối ưu hóa trong bài viết này đều hướng đến một trong hai cách đó.

CPU bị nghẽn ở bốn yếu tố: các lần truy cập bộ đệm không thành công (cache misses) (L1d khoảng 4 chu kỳ, một lần L3-miss-to-DRAM mất hàng trăm chu kỳ), chuỗi tải phụ thuộc mà công cụ thực thi ngoài thứ tự (out-of-order engine) không thể song song hóa, dự đoán nhánh sai (~15-20 chu kỳ bị hủy cho mỗi lần đoán sai), và các lệnh gọi bộ cấp phát (allocator calls). Triển khai công cụ unigram tham chiếu gặp phải cả bốn vấn đề sau:

Nút thắt cổ chai	Cơ chế	Tác động đã đo lường
Phân bổ cho mỗi kết quả khớp	String::from_utf8 trên các byte đã khớp cộng với một lượt tra cứu AHashMap cho mỗi kết quả khớp trie; ID token đã có sẵn trong nút lá (leaf), nhưng bản tham chiếu không lưu trữ nó ở đó	7.295 lần phân bổ ở 512 token, 14.542 ở 1k, 299.171 ở 16k
Truy tìm con trỏ trên mỗi byte	Một AHashMap ở mọi nút trie, băm byte tiếp theo và giải tham chiếu hai con trỏ trên mỗi bước; cấu trúc Node cũng gom nhóm các trường ít dùng (cold fields) như is_leaf, kéo chúng vào bộ đệm cache trong mỗi lần truy cập	Tỷ lệ trượt L1d ở mức thấp (~2.5% ở 1k) vì các trang nóng của bộ cấp phát nằm trong L1; chi phí thể hiện dưới dạng độ trễ tải phụ thuộc
Hiện tượng tranh chấp L2 (L2 thrashing) trên các đầu vào dài	Bảng DP, bộ nhớ đệm tạm tiền tố (prefix scratch) và danh sách token đầu ra đều được phân bổ mới trong mỗi lần gọi	Tỷ lệ trượt L2 tăng từ 12% ở 128 token lên 50% ở 16k khi dung lượng phân bổ tích lũy vượt quá 2 MB L2 trên mỗi lõi

Tham chiếu các điểm nghẽn, cơ chế và tác động đo được của tokenizer.

Chi phí lớn nhất là việc phân bổ cho mỗi kết quả khớp. Ở mức 1k token, triển khai tham chiếu hoàn thành khoảng 7 triệu chỉ thị cho mỗi lần mã hóa ở mức IPC là 2.8, phần lớn trong số đó là lệnh gọi bộ cấp phát và chuyển đổi chuỗi. Các lần trượt L2 chỉ quan trọng khi vượt quá ~1k token, một khi dung lượng phân bổ tích lũy làm tràn bộ đệm trên mỗi lõi.

Độ trễ p50 và p99 của tokenizer tham chiếu theo kích thước đầu vào.

Tỷ lệ lỗi cache L2 của tokenizer tham chiếu theo kích thước đầu vào.

Số byte được cấp phát cho mỗi lần encode của tokenizer tham chiếu trên các kích thước đầu vào.

Phân bổ trên mỗi token là hằng số: khoảng 2 KB và ~18 lần phân bổ cho mỗi token, bất kể kích thước đầu vào. Độ phi tuyến tính của độ trễ xuất phát từ L2: số byte tích lũy vượt mốc 2 MB quanh 1k token, tỷ lệ trượt tăng từ 8% lên ~50%, và đường cong p99 uốn cong theo tương ứng.

Chúng tôi so sánh ba trình mã hóa token trong môi trường production, tất cả đều chạy cùng thuật toán Unigram trên cùng một mô hình XLM-RoBERTa:

Công cụ	Ngôn ngữ	Trie
Hugging Face	Rust	Trie tiền tố HashMap-children
SentencePiece	C++	Trie mảng kép Darts
Trình mã hóa token IREE	C	Trie mảng kép Darts

Các triển khai tokenizer, ngôn ngữ và biểu diễn trie được dùng trong so sánh.

Chúng tôi liên kết SentencePiece và IREE thông qua FFI C/C++ trực tiếp thay vì các trình bao bọc (wrappers) Rust của chúng; các trình bao bọc này bổ sung thêm các phân bổ trên mỗi lần gọi và chuyển đổi chuỗi vốn không phải là một phần chi phí của thư viện gốc.

Tất cả các công cụ đều được xây dựng với các tối ưu hóa đầy đủ (O3, nhắm mục tiêu gốc, tối ưu hóa thời gian liên kết). Các phép đo được chạy cố định trên một lõi ít hoạt động trên Intel Xeon Platinum 8488C. Đối với mỗi công cụ và kích thước đầu vào, chúng tôi thu thập độ trễ thời gian thực (p50 và p99, 10.000 vòng lặp sau 1.000 vòng khởi động), số chỉ thị hoàn thành thông qua Valgrind cachegrind, tỷ lệ IPC và trượt bộ đệm từ các bộ đếm phần cứng perf stat, và các phân bổ heap được đếm bởi một trình bao bọc bộ cấp phát toàn cục xung quanh cửa sổ được đo lường.

Ở mức 514 token (+2 so với 512 là việc chèn BOS/EOS):

Triển khai	Độ trễ p50	Chỉ thị	IPC	Phân bổ
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

Độ trễ tokenizer, số lượng lệnh, IPC và số lượt cấp phát tại 514 token.

SP và IREE đạt độ trễ xấp xỉ một phần ba so với bản tham chiếu. Chúng hoàn thành số chỉ thị ít hơn khoảng một nửa, chạy ở IPC cao hơn và phân bổ ít hơn từ một đến hai bậc độ lớn.

Cả SentencePiece và IREE đều sử dụng trie mảng kép do Aoe giới thiệu vào năm 1989 (Aoe, 1989; Karoonboonyanan). Trie được mã hóa thành hai mảng số nguyên phẳng, base và check:

cpp

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

Tra cứu nút con là:

cpp

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

Mỗi bước duyệt trie chỉ cần hai lần đọc mảng, một phép cộng số nguyên và một phép so sánh, không có băm, không truy tìm con trỏ và không phân bổ heap trên mỗi nút. Quá trình xây dựng đặt mỗi nút sao cho các nút con của nó rơi vào các ô trống trong mảng kiểm tra, cho phép nhiều nút cha chia sẻ các vùng miễn là các nút con của chúng khác nhau về byte theo sau.

Đối với từ vựng 250K của XLM-RoBERTa, toàn bộ trie nằm gọn trong khoảng 9 MB bộ nhớ liên tục (~780K ô × 12 byte cho độ lệch, kiểm tra và ID token cho mỗi ô). Tập hợp làm việc nóng (hot working set) của một lần mã hóa đơn lẻ có quy mô khoảng 100 KB, nằm thoải mái trong L2 và bộ nạp trước phần cứng (hardware prefetcher) khóa chặt vào mẫu truy cập đều đặn.

Ngược lại, bản tham chiếu duy trì một bảng băm các nút con cho mỗi nút, tốn khoảng bốn lần tải phụ thuộc cho mỗi bước byte. Băm byte, tải thùng (bucket), đi theo con trỏ đến mục nhập thùng, đi theo con trỏ khác đến nút con trên heap, sau đó tham khảo bảng băm của chính nó cho byte tiếp theo. Các đối tượng heap phân tán và hầu hết các lần tải đều trượt cache.

Cả SP và IREE đều lưu trữ ID token trực tiếp trong nút trie. Triển khai trie tham chiếu chỉ trả về các byte đã khớp, vì vậy bộ mã hóa phải phân bổ một chuỗi mới cho mỗi kết quả khớp, tra cứu nó trong một bảng băm phụ để khôi phục ID token, và đọc điểm số từ một véc-tơ riêng biệt. SP trả về ID dưới dạng giá trị được lưu trữ của trie và đọc điểm số từ một mảng số thực phẳng:

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 sử dụng API con trỏ có trạng thái (stateful cursor API) thay vì trả về tất cả các kết quả khớp trong một lệnh gọi duy nhất, nhưng mô hình chi phí là như nhau: ID token và điểm số có thể truy cập được dưới dạng các chỉ mục mảng phẳng, không có phân bổ hoặc băm trong chính quá trình duyệt trie. IREE chỉ phân bổ một bộ đệm đầu ra cho mỗi lần mã hóa; SP bổ sung phân bổ một nút mạng lưới (lattice node) cho mỗi kết quả khớp trong quá trình DP, tổng cộng khoảng 1.559 lần phân bổ ở 514 token.

Trước khi thay đổi trie, chúng tôi cô lập chi phí thuật toán khỏi chi phí triển khai: bản tham chiếu chạy nhanh đến mức nào khi loại bỏ các công việc không cần thiết?

Bản cơ sở giữ lại trie bảng băm các nút con của bản tham chiếu nhưng thực hiện hai thay đổi. Một cấu trúc bộ đệm tạm (scratch struct) do người gọi sở hữu chứa bảng DP, các byte đã được tiền xử lý và ID token đầu ra, do đó đường dẫn nóng không thực hiện phân bổ nào sau khi khởi động. ID token và điểm số được lưu trữ trực tiếp trong nút trie, loại bỏ việc chuyển đổi chuỗi trên mỗi kết quả khớp và việc tra cứu bảng băm phụ.

Vòng lặp bên trong:

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

Đây là cùng một thuật toán, với cùng hình dạng trie và cùng mẫu truy cập, chỉ là không có phân bổ và không có tra cứu dư thừa. Kết quả ở mức 514 token:

Công cụ	p50	Chỉ thị	IPC	Phân bổ
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
Bản cơ sở	156 µs	1.53M	2.62	0

Đường cơ sở không cấp phát so với các triển khai tokenizer hiện có tại 514 token.

Việc loại bỏ quá trình tạo chuỗi trên mỗi kết quả khớp và tra cứu bảng băm phụ giúp cắt giảm số chỉ thị xuống 2.4× so với bản tham chiếu. IPC giảm theo (2.86 → 2.62) bởi vì việc truy tìm con trỏ bảng băm, vốn trước đây bị ẩn sau công việc của bộ cấp phát, nay trở thành chi phí chiếm ưu thế trên mỗi byte.

Độ trễ p50 và p99 của tokenizer tham chiếu và đường cơ sở không cấp phát theo kích thước đầu vào.

Tỷ lệ lỗi cache L2 của tokenizer tham chiếu và đường cơ sở không cấp phát theo kích thước đầu vào.

Bản cơ sở cắt giảm độ trễ p50 đi 2-3× so với bản tham chiếu ở mọi kích thước, và tỷ lệ trượt L2 của nó duy trì ở mức gần 2-3% qua 2k token so với mức 8-48% của bản tham chiếu. Chi phí còn lại là việc truy tìm con trỏ HashMap, mà phần tiếp theo sẽ giải quyết.

Trie bảng băm là một quá trình truy tìm con trỏ. SP và IREE giải quyết vấn đề này bằng trie mảng kép (xem phần trước) và chúng tôi mượn cấu trúc dữ liệu tương tự. Sự khác biệt là SP và IREE là các thư viện đa năng hỗ trợ đầu vào dạng luồng, độ lệch đầu ra và nhiều chế độ mã hóa; sự linh hoạt đó đánh đổi bằng việc ghi sổ kế toán phụ xung quanh trie (một lưới trong SP, một đường ống tám giai đoạn trong IREE). Chúng ta không cần bất kỳ điều gì trong số đó, vì vậy chúng tôi nhúng thẳng trie vào vòng lặp Viterbi và loại bỏ phần còn lại.

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

Một bước byte tốn hai lần đọc mảng, một phép cộng số nguyên và một phép so sánh. Toàn bộ trie có dung lượng khoảng 9 MB, nhưng bất kỳ lần mã hóa đơn lẻ nào chỉ duyệt qua ~100 KB trong số đó, vì vậy dữ liệu nóng được giữ lại trong L2.

Công cụ	p50	p99	Chỉ thị	L2 trượt%	L3 trượt%
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%
Bản cơ sở	155 µs	1.170 µs	1.53M	2.59%	16.66%
+ Darts	68 µs	80 µs	1.10M	6.45%	11.42%

Độ trễ tokenizer, số lệnh và tỷ lệ lỗi cache sau khi áp dụng double-array trie.

Thời gian thực giảm 2.3× so với bản cơ sở và 4.8× so với bản tham chiếu. Số chỉ thị giảm thêm 28% so với bản cơ sở. Việc loại bỏ mạng lưới và đường ống cho phép chúng tôi hoàn thành 1.10M chỉ thị trên mỗi lần mã hóa so với mức 1.82M của SP và 1.61M của IREE ở cùng hình dạng trie; khoảng cách này đến từ phần ghi sổ kế toán mà họ duy trì cho các API đa năng mà chúng ta không cần.

Trie Darts từ Tối ưu hóa 1 cần hai lần tải phụ thuộc cho mỗi bước byte. Đầu tiên là độ lệch base của nút cha, sau đó là mục check tại ô chứa nút con ứng viên để xác nhận nút cha sở hữu nó. Mục check chỉ xác nhận quyền sở hữu. Câu hỏi tương tự có thể được trả lời bằng một bitmap trên mỗi nút ghi lại các byte nào có nút con. Theo cấu trúc của trie mảng kép, bit cho byte b trong bitmap của nút cha được bật chính xác khi ô ứng viên cho byte đó thuộc quyền sở hữu của nút cha. Bitmap và mảng check mang cùng một thông tin, vì vậy chúng ta có thể loại bỏ hoàn toàn mảng check khỏi cấu trúc thời gian chạy.

Việc tra cứu bitmap biên dịch thành một kiểm tra bit đơn lẻ trên một từ 64-bit (64-bit word), chạy trong khoảng một chu kỳ. Chỉ thị POPCNT liên quan, đếm tất cả các bit được bật trong một thanh ghi, là thứ mà các thiết kế trie được lập chỉ mục bằng bitmap khác sử dụng để tính trực tiếp độ lệch nút con từ bitmap. Chúng ta không cần nó vì cấu trúc Darts đã lưu trữ độ lệch đó trong trường base của nút.

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

Trong lúc làm việc đó, chúng tôi đóng gói bốn trường trên mỗi nút (bitmap, base, ID token và điểm số) thành một dòng bộ đệm (cache line) 64 byte. Một lần truy cập nút duy nhất giờ đây kéo vào mọi thứ: điểm số và ID token chúng ta cần ở các nút cuối (terminal nodes), cộng với bitmap và base cho bước tiếp theo. Mảng check chỉ được sử dụng trong quá trình xây dựng và bị loại bỏ khỏi bố cục thời gian chạy. Việc lập chỉ mục đường dẫn nóng không được kiểm tra giới hạn (unchecked) vì việc xây dựng trie và cấu trúc vòng lặp đã đảm bảo các giới hạn đó. Sự đánh đổi là không gian (Space): ở mức 64 byte cho mỗi nút × ~780K ô, trie tăng trưởng từ ~9 MB lên ~50 MB.

Kết quả ở mức 514 token:

Công cụ	p50	p99	Chỉ thị	Truy cập L2	L2 trượt%	L3 trượt%
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%
Bản cơ sở	155 µs	1.170 µs	1.53M	6.1K	2.59%	16.66%
+ Darts	68 µs	80 µs	1.10M	4.6K	6.45%	11.42%
+ Bitmap và đóng gói nội tuyến	65 µs	77 µs	1.06M	1.8K	20.53%	15.01%

Kết quả tokenizer tại 514 token sau khi thêm bitmap và đóng gói nội tuyến.

Thời gian thực tiếp tục giảm có hệ thống thêm 4.5% ở 514 token và tăng theo độ dài đầu vào, đạt 6-7% ở 4k và 16k. Việc loại bỏ mảng check biến bước trie trên mỗi byte từ hai lần tải dòng bộ đệm thành một, tức là ít hơn 2.5× số lần truy cập L2 trên mỗi lần mã hóa (1.8K so với 4.6K của Darts thông thường). Tỷ lệ trượt tăng vọt từ 6% lên 20%, nhưng phần lớn trong số đó là do mẫu số: với số lần truy cập chưa bằng một nửa, những lần trượt mà chúng ta gặp phải chiếm tỷ trọng lớn hơn.

Dữ liệu trie có dung lượng khoảng 50 MB (780.000 nút với mỗi nút 64 byte). Trên một tiến trình Linux mặc định, bộ nhớ đó được hỗ trợ bởi các trang 4 KB, vì vậy toàn bộ trie trảiσ rộng trên khoảng 12.000 trang ảo khác nhau, vượt xa khả năng lưu trữ của Bộ đệm dịch địa chỉ dữ liệu cấp độ đầu tiên của CPU (Translation Lookaside Buffer - Sapphire Rapids có 96 mục). Mỗi bước Viterbi chạm vào một nút khác nhau, và qua một lần mã hóa 512 token, chúng ta tiêu tốn khoảng 9.000 chu kỳ cho việc duyệt bảng trang (page-table walks), chiếm khoảng 3% ngân sách cho mỗi lần mã hóa.

Giải pháp là hỗ trợ trie bằng các trang lớn 2 MB thay vì các trang 4 KB. 50 MB dữ liệu giờ đây trải rộng trên 25 trang, nằm trọn vẹn trong Translation Lookaside Buffer, và bản thân các thư mục bảng trang (page-table directories) không còn cạnh tranh với trie về băng thông bộ đệm L2.

Linux có hai cách để hỗ trợ bộ nhớ bằng các trang 2 MB. Cách linh hoạt là trang lớn trong suốt (transparent huge pages - THP): chương trình của bạn yêu cầu nhân hệ điều hành nâng cấp một vùng nhớ từ trang 4 KB lên 2 MB, và một luồng nhân nền cố gắng thực hiện điều đó khi có thời gian và bộ nhớ vật lý liên tục khả dụng. Giải pháp thay thế khác là phân bổ trie trực tiếp từ một nhóm trang lớn được đặt trước mà nhân hệ điều hành dành riêng khi khởi động:

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

Nhóm này cần tồn tại trên máy chủ (vm.nr_hugepages được thiết lập lúc khởi động). Trên các thể hiện (instances) production của chúng tôi, 10.561 trang lớn được dự trữ; trie tiêu thụ 24 trang trong số đó. Sau lệnh mmap, /proc/<pid>/smaps xác nhận vùng nhớ: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Mức tăng đạt 3–12% thời gian thực trên các kích thước đầu vào, nhất quán và tỷ lệ thuận với độ dài tập hợp làm việc:

Token	Không có trang lớn	Có trang lớn	Δ
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%

Độ trễ tokenizer có và không có huge pages trên các kích thước đầu vào.

Mức cải thiện lớn nhất là ở 4k token, nơi tập hợp làm việc đủ lớn để lưu lượng bảng trang đang cạnh tranh trực tiếp với dữ liệu trie về băng thông L2. Vượt quá 4k, mức tăng thu hẹp lại vì các lần trượt L3 chiếm ưu thế và việc duyệt bảng trang trở thành một phần nhỏ hơn trong tổng thời gian.

Kết hợp ba tối ưu hóa lại với nhau, đây là công cụ cuối cùng được đo lường so với mọi trình mã hóa token RoBERTa khác mà chúng tôi đã thử nghiệm. Bên cạnh các tệp nhị phân C++ SentencePiece và C IREE gốc, chúng tôi cũng tiến hành đo điểm chuẩn các crate bao bọc Rust có sẵn mà người dùng trong môi trường production thường tìm đến, cụ thể là sentencepiece và iree-tokenizer. Crate sentencepiece đo được độ trễ cao hơn 1.6-1.9× so với tệp nhị phân C++ gốc ở mọi kích thước đầu vào, với phần chi phí phát sinh đến từ việc trình bao bọc phân bổ một danh sách các mảnh token mới, mỗi mảnh giữ một bản sao văn bản mảnh của riêng nó, trong mỗi lần mã hóa. Crate iree-tokenizer định kích thước bộ đệm nội bộ của nó theo từng đầu vào, giúp tránh việc phân bổ liên tục trên mỗi lần mã hóa mà tệp nhị phân IREE gốc phải trả bất kể kích thước đầu vào, và điều này quan trọng nhất ở các đầu vào ngắn nơi việc phân bổ sẽ chiếm phần lớn công việc mã hóa thực tế. Ở độ dài đầu vào trong môi trường production, chi phí hằng số này được phân bổ dần và chi phí phát sinh trên mỗi lần gọi của trình bao bọc từ việc sao chép các token trở lại Rust đẩy độ trễ của nó vượt lên trên tệp nhị phân gốc.

Độ trễ p50 và p99 trên mỗi lần encode across các triển khai tokenizer và kích thước đầu vào.

Công cụ cuối cùng chạy thuật toán Viterbi tương tự trên cùng từ vựng 250.000 token, không phân bổ bất cứ thứ gì trên đường dẫn nóng và tạo ra đầu ra chính xác từng token so với bản tham chiếu. Mỗi bước byte giờ đây là một lần tải dòng bộ đệm 64 byte đơn lẻ, một phép kiểm tra bitmap và một phép cộng số nguyên, giảm xuống từ một lượt dò tìm HashMap, một lệnh gọi bộ cấp phát và một số nhánh trường hợp đặc biệt trong triển khai cơ sở. Số chỉ thị trên mỗi lần mã hóa giảm từ 3.66M xuống 1.04M, một mức giảm 3.5×. Thời gian thực khớp với tỷ lệ đó ở các đầu vào ngắn và nới rộng hơn nữa ở các đầu vào dài, nơi các phân bổ trên mỗi token của bản tham chiếu bắt đầu làm tràn bộ đệm L2 và L3.

Trong thực tế, công việc này đã giảm lượng sử dụng CPU trong lớp tiền xử lý của ngăn xếp suy luận của chúng tôi đi 5-6× và cắt giảm độ trễ của mô hình xếp hạng lại đi hàng chục mili-giây.

Phần lớn việc giảm độ trễ đến từ các thay đổi nhẹ nhàng giúp loại bỏ công việc lãng phí: loại bỏ các phân bổ trên mỗi lần mã hóa, thay thế trie bảng băm các nút con bằng một mảng kép và loại bỏ chi phí API chung. Việc sử dụng bitmap trên mỗi nút hiện có làm cơ sở kiểm tra tính hợp lệ và hỗ trợ trie bằng các trang lớn 2 MB kết hợp lại chiếm 10-20% hiệu quả vượt trội bên cạnh việc chuyển đổi lớn sang trie mảng kép, tùy thuộc vào kích thước đầu vào.

Các crate bao bọc Rust có sẵn xung quanh SentencePiece và IREE minh họa điểm tương tự theo chiều ngược lại. Trên các trình mã hóa token C và C++ gốc, ranh giới FFI và các phân bổ trên mỗi lần gọi mà các trình bao bọc giới thiệu đều hiển thị rõ trong các con số độ trễ. Mã nguồn Rust cho bộ mã hóa Viterbi cuối cùng có sẵn tại 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}
}

Động lực

Thuật toán Viterbi

Triển khai tham chiếu

Lưu ý về tiền xử lý

Triển khai tham chiếu tiêu tốn thời gian ở đâu

Cách chúng tôi đo lường

Tại sao SentencePiece và IREE nhanh hơn: Trie mảng kép

Cơ sở: Bản chuyển đổi không phân bổ của bản tham chiếu

Tối ưu hóa 1: Trie mảng kép không có cácừu tượng hóa thừa

Tối ưu hóa 2: Đóng gói bitmap và inline

Tối ưu hóa 3: Trang kích thước lớn cho trie

So sánh cuối cùng

Lời kết

Tài liệu tham khảo

Trích dẫn