Βελτίωση της απόδοσης CPU του Unigram Tokenizer

Υλοποιήσαμε εκ νέου τον Unigram tokenizer μας από την αρχή ως ένα στοχευμένο έργο βελτίωσης επιδόσεων. Σε τυπικά μήκη εισόδου παραγωγής, ο νέος κωδικοποιητής μειώνει την καθυστέρηση p50 κατά περίπου 5× σε σύγκριση με το crate tokenizers της Hugging Face, ~2× σε σύγκριση με το SentencePiece (C++) και ~1.5× σε σύγκριση με τον tokenizer του IREE (C), χωρίς καμία κατανομή σωρού (heap allocation) σε σταθερή κατάσταση. Στην παραγωγή, αυτό μείωσε τη χρήση CPU στο σύστημα εξαγωγής συμπερασμάτων (inference stack) κατά 5-6× και αφαίρεσε διψήφια χιλιοστά του δευτερολέπτου από την καθυστέρηση του reranker.

Το κόστος εξαγωγής συμπερασμάτων LLM συνήθως πλαισιώνεται ως εργασία GPU: κρυφές μνήμες KV, πυρήνες προσοχής (attention kernels), δρομολόγηση εμπειρογνωμόνων (expert routing). Ωστόσο, ένα σημαντικό μέρος της κίνησης παραγωγής εκτελείται σε πολύ μικρότερα μοντέλα, δύο έως τρεις τάξεις μεγέθους μικρότερα από τα transformers αιχμής. Τα μοντέλα ενσωμάτωσης (embedding models), οι ταξινομητές και οι rerankers εξακολουθούν να εκτελούνται σε GPUs, αλλά το πέρασμα προς τα εμπρός (forward pass) είναι αρκετά σύντομο ώστε όλα τα υπόλοιπα γύρω του να αρχίσουν να έχουν σημασία.

Η πιο ξεκάθαρη περίπτωση είναι ένας reranker που αξιολογεί εκατοντάδες υποψήφια έγγραφα ανά αίτημα. Με ένα μικρό μοντέλο, ο υπολογισμός GPU συνήθως ολοκληρώνεται σε μονές χιλιοστά του δευτερολέπτου, αλλά κάθε είσοδος πρέπει ακόμα να περάσει πρώτα από την τμηματοποίηση (tokenization) από την πλευρά της CPU, μετατρέποντας το κείμενο σε αναγνωριστικά λεξιλογίου (vocabulary IDs). Όταν η παρτίδα (batch) είναι μεγάλη, αυτό μπορεί να αντιπροσωπεύει ένα σημαντικό κλάσμα της συνολικής καθυστέρησης αιτήματος.

Αυτό το έργο στοχεύει στο XLM-RoBERTa, ένα λεξιλόγιο Unigram 250K token που εκπαιδεύτηκε με το SentencePiece. Οι κωδικοποιητές της οικογένειας RoBERTa με λεπτή ρύθμιση (fine-tuned) αποτελούν μια κοινή επιλογή για εργασίες κατάταξης, ανάκτησης και ομοιότητας. Αναφέρουμε την καθυστέρηση ανά κωδικοποίηση με ένα νήμα (single-threaded) παντού· στην παραγωγή κάνουμε παράλληλη επεξεργασία σε πυρήνες χρησιμοποιώντας το rayon.

Σε όλη αυτή την ανάρτηση αντιμετωπίζουμε το tokenizers crate της Hugging Face ως την αναφορά με την οποία κάνουμε benchmark, δεδομένου ότι είναι ο προεπιλεγμένος Rust tokenizer που επιλέγουν οι περισσότερες ομάδες. Το SentencePiece (C++) και το IREE (C) περιλαμβάνονται για σύγκριση ως οι άλλες δύο σημαντικές υλοποιήσεις παραγωγής.

Η τμηματοποίηση Unigram, που εισήχθη από τον Kudo το 2018 (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) και ενσωματώθηκε στο SentencePiece, πλαισιώνει την κατάτμηση ως πρόβλημα μονοπατιού με τη μεγαλύτερη πιθανότητα: κάθε token λεξιλογίου έχει μια εκπαιδευμένη λογαριθμική πιθανότητα (log-probability), και ο κωδικοποιητής επιλέγει την κατάτμηση της οποίας τα σκορ των token αθροίζονται στην υψηλότερη τιμή. Το πρόβλημα λύνεται με τον αλγόριθμο Viterbi (1967), με τις θέσεις byte ως επίπεδα γραφήματος και τα token λεξιλογίου ως ακμές που εκτείνονται σε ένα συνεχές εύρος byte. Η αναδρομή είναι ένας μονοδιάστατος δυναμικός προγραμματισμός (DP) πάνω στις θέσεις 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

Ο εξωτερικός βρόχος είναι γραμμικός ως προς το μήκος εισόδου, περιορισμένος σε κάθε θέση από το μεγαλύτερο token λεξιλογίου (48 byte UTF-8 για το XLM-RoBERTa). Η εσωτερική διαδρομή trie είναι το κρίσιμο μονοπάτι (hot path): σε μια είσοδο 16k token, εκτελεί εκατοντάδες χιλιάδες μεταβάσεις και ολοκληρώνει δεκάδες εκατομμύρια εντολές ανά κωδικοποίηση.

Το tokenizers v0.22 της Hugging Face υλοποιεί την αναδρομή απευθείας, με έναν πίνακα κατακερματισμού (hash-map) για την αντίστροφη αναζήτηση και hash-map-children σε κάθε κόμβο 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
}

Και ο εσωτερικός βρόχος (απλοποιημένος) – κάθε αντιστοιχία αναζήτησης προφίματος (prefix-search match) δεσμεύει (allocates) ένα νέο string και στη συνέχεια ελέγχει τον πίνακα κατακερματισμού:

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, αλλά συμβαίνει πολλή δουλειά που είναι αόρατη σε επίπεδο πηγαίου κώδικα. Η επόμενη ενότητα αναλύει πού πηγαίνει ο χρόνος.

Πριν εκτελεστεί το Viterbi, η είσοδος περνάει από κανονικοποίηση Unicode (NFKC) και σήμανση ορίων λέξεων. Η ανάλυση απόδοσης (profiling) τοποθετεί την προεπεξεργασία σε λιγότερο από 5% του συνολικού χρόνου κωδικοποίησης στα σενάριά μας, οπότε εστιαζόμαστε στο πέρασμα προς τα εμπρός (forward pass) του Viterbi για το υπόλοιπο της ανάρτησης.

Υπάρχουν δύο τρόποι για να επιταχυνθεί ο κώδικας CPU: εκτέλεση λιγότερων εντολών ή επίτευξη υψηλότερου IPC. Το Cachegrind μετρά το πρώτο, το perf stat μετρά το δεύτερο. Κάθε βελτιστοποίηση σε αυτή την ανάρτηση καταλήγει στο ένα ή στο άλλο.

Οι CPUs καθυστερούν σε τέσσερα πράγματα: αστοχίες κρυφής μνήμης (cache misses - L1d ~4 κύκλοι, μια αστοχία L3 προς DRAM είναι εκατοντάδες), αλυσίδες εξαρτώμενων φορτίων (dependent-load chains) που η μηχανή εκτός σειράς (out-of-order engine) δεν μπορεί να κάνει παράλληλες, λανθασμένες προβλέψεις διακλάδωσης (~15-20 κύκλοι εκκαθάρισης ανά λανθασμένη πρόβλεψη) και κλήσεις καταμεριστή (allocator calls). Η υλοποίηση της μηχανής unigram αναφοράς χτυπά και τα τέσσερα:

Σημείο συμφόρησης	Μηχανισμός	Μετρημένος αντίκτυπος
Κατανομή ανά αντιστοιχία	String::from_utf8 στα αντιστοιχισμένα byte συν μια αναζήτηση AHashMap ανά αντιστοιχία trie· το ID του token βρίσκεται ήδη στο φύλλο, αλλά η αναφορά δεν το αποθηκεύει εκεί	7.295 κατανομές σε 512 token, 14.542 σε 1k, 299.171 σε 16k
Αναζήτηση δεικτών (pointer chase) ανά byte	Ένας AHashMap σε κάθε κόμβο trie, που κατακερματίζει το επόμενο byte και αποεξαρτά (dereferences) δύο δείκτες ανά βήμα· η δομή Node ομαδοποιεί επίσης ψυχρά πεδία όπως το is_leaf, τραβώντας τα στην κρυφή μνήμη σε κάθε επίσκεψη	Το ποσοστό αστοχίας L1d παραμένει χαμηλό (~2.5% στο 1k) επειδή οι θερμές σελίδες του καταμεριστή βρίσκονται στην L1· το κόστος εμφανίζεται ως καθυστέρηση εξαρτώμενου φορτίου
Thrashing L2 σε μεγάλες εισόδους	Ο πίνακας DP, το πρόχειρο προφίματος (prefix scratch) και η λίστα token εξόδου δεσμεύονται εκ νέου σε κάθε κλήση	Το ποσοστό αστοχίας L2 ανεβαίνει από 12% σε 128 token σε 50% σε 16k καθώς το αθροιστικό αποτύπωμα κατανομής υπερβαίνει την L2 2 MB ανά πυρήνα

Αναφορά στα σημεία συμφόρησης, στους μηχανισμούς και στις μετρημένες επιπτώσεις του tokenizer.

Το κυρίαρχο κόστος είναι η κατανομή ανά αντιστοιχία. Σε 1k token, η υλοποίηση αναφοράς ολοκληρώνει περίπου 7M εντολές ανά κωδικοποίηση με IPC 2.8, οι περισσότερες από τις οποίες σε κλήσεις καταμεριστή και μετατροπές string. Οι αστοχίες L2 έχουν σημασία μόνο μετά από ~1k token, μόλις το αθροιστικό αποτύπωμα κατανομής υπερχειλίσει την κρυφή μνήμη ανά πυρήνα.

Καθυστέρηση p50 και p99 του tokenizer αναφοράς ανά μέγεθος εισόδου.

Ποσοστό αστοχιών L2 cache του tokenizer αναφοράς ανά μέγεθος εισόδου.

Bytes που δεσμεύονται ανά encode από τον tokenizer αναφοράς σε διαφορετικά μεγέθη εισόδου.

Η κατανομή ανά token είναι σταθερή: περίπου 2 KB και ~18 κατανομές ανά token, ανεξάρτητα από το μέγεθος εισόδου. Η μη γραμμικότητα της καθυστέρησης προέρχεται από την L2: τα αθροιστικά byte ξεπερνούν τα 2 MB γύρω στα 1k token, το ποσοστό αστοχίας ανεβαίνει από 8% σε ~50%, και η καμπύλη p99 καμπυλώνει ανάλογα.

Συγκρίνουμε τρεις tokenizers παραγωγής, οι οποίοι εκτελούν όλοι τον ίδιο αλγόριθμο Unigram με το ίδιο μοντέλο XLM-RoBERTa:

Μηχανή	Γλώσσα	Trie
Hugging Face	Rust	Trie προφίματος με HashMap-children
SentencePiece	C++	Darts double-array trie
Tokenizer IREE	C	Darts double-array trie

Υλοποιήσεις tokenizer, γλώσσες και αναπαραστάσεις trie που χρησιμοποιούνται στη σύγκριση.

Συνδέουμε το SentencePiece και το IREE μέσω άμεσου C/C++ FFI αντί για τα wrappers τους σε Rust· τα wrappers προσθέτουν κατανομές ανά κλήση και μετατροπές string που δεν αποτελούν μέρος του κόστους της υποκείμενης βιβλιοθήκης.

Όλες οι μηχανές είναι κατασκευασμένες με πλήρεις βελτιστοποιήσεις (O3, native targeting, link-time optimization). Οι μετρήσεις εκτελούνται καρφιτσωμένες σε έναν ήσυχο πυρήνα σε έναν Intel Xeon Platinum 8488C. Για κάθε μηχανή και μέγεθος εισόδου συλλέγουμε λανθάνουσα καθυστέρηση wall-clock (p50 και p99, 10.000 επαναλήψεις μετά από 1.000 warmup), εντολές που ολοκληρώθηκαν μέσω του Valgrind cachegrind, IPC και ρυθμούς αστοχιών cache από counters υλικού του perf stat, καθώς και δεσμεύσεις heap που μετρώνται από έναν global allocator wrapper γύρω από το μετρούμενο παράθυρο.

Σε 514 token (τα +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

Καθυστέρηση tokenizer, αριθμός εντολών, IPC και δεσμεύσεις στα 514 tokens.

Το SP και το IREE καταλήγουν περίπου στο ένα τρίτο της καθυστέρησης της αναφοράς. Ολοκληρώνουν περίπου τις μισές εντολές, εκτελούνται σε υψηλότερο IPC και κατανέμουν μία έως δύο τάξεις μεγέθους λιγότερο.

Το SentencePiece και το IREE χρησιμοποιούν και τα δύο το double-array trie που εισήγαγε ο Aoe το 1989 (Aoe, 1989· Karoonboonyanan). Το trie είναι κωδικοποιημένο σε δύο επίπεδους ακέραιους πίνακες, τους 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 */ }

Κάθε βήμα της διέλευσης trie χρειάζεται μόνο δύο αναγνώσεις πίνακα, μία προσθήκη ακεραίων και μία σύγκριση, χωρίς κατακερματισμό, χωρίς αναζήτηση δεικτών και χωρίς κατανομή σωρού ανά κόμβο. Η κατασκευή τοποθετεί κάθε κόμβο έτσι ώστε τα θυγατρικά του να πέφτουν σε ελεύθερες θέσεις στον πίνακα επικύρωσης, επιτρέποντας σε πολλαπλούς γονείς να μοιράζονται περιοχές εφόσον τα θυγατρικά τους διαφέρουν στο byte που ακολουθούν.

Για το λεξιλόγιο 250K του XLM-RoBERTa, ολόκληρο το trie χωράει σε ~9 MB συνεχούς μνήμης (~780K θέσεις × 12 byte για το offset, την επικύρωση και το ID token ανά θέση). Το σύνολο ενεργής εργασίας μιας μίας κωδικοποίησης είναι της τάξης των 100 KB, άνετα στην L2, και ο προεπιλογέας υλικού (hardware prefetcher) κλειδώνει στο κανονικό μοτίβο πρόσβασης.

Η αναφορά, αντίθετα, διατηρεί έναν πίνακα κατακερματισμού παιδιών ανά κόμβο, κοστίζοντας περίπου τέσσερα εξαρτώμενα φορτία ανά βήμα byte. Κατακερματισμός του byte, φόρτωση του κάδου, ακολουθία ενός δείκτη στην καταχώρηση του κάδου, ακολουθία άλλου δείκτη στον θυγατρικό κόμβο στον σωρό, και στη συνέχεια εξέταση του δικού του πίνακα κατακερματισμού για το επόμενο byte. Τα αντικείμενα σωρού διασκορπίζονται και οι περισσότερες φορτώσεις χάνουν την κρυφή μνήμη.

Τόσο το SP όσο και το IREE αποθηκεύουν το ID του token απευθείας στον κόμβο του trie. Η υλοποίηση αναφοράς του trie επιστρέφει μόνο τα bytes που ταιριάστηκαν, οπότε ο encoder πρέπει να δεσμεύει ένα νέο string για κάθε ταίριασμα, να το αναζητά σε ένα δευτερεύον hash-map για να ανακτήσει το ID του token και να διαβάζει το score από ένα ξεχωριστό vector. Το SP επιστρέφει το ID ως αποθηκευμένη τιμή του trie και διαβάζει το score από έναν επίπεδο float array:

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 χρησιμοποιεί ένα κρατούν κατάσταση (stateful) API δρομέα αντί να επιστρέφει όλα τα ταιριάσματα σε μία μόνο κλήση, αλλά το μοντέλο κόστους είναι το ίδιο: τα ID των token και τα σκορ είναι προσπελάσιμα ως δείκτες επίπεδου πίνακα (flat-array), χωρίς καμία κατανομή (allocation) ή κατακερματισμό (hashing) κατά την ίδια την περιήγηση στο trie. Το IREE δεσμεύει μόνο έναν buffer εξόδου ανά κωδικοποίηση· το SP επιπλέον δεσμεύει έναν κόμβο δικτύου ανά ταιριάσματος κατά τη διάρκεια του DP, συνολικά περίπου 1.559 δεσμεύσεις στα 514 token.

Πριν αλλάξουμε το trie, απομονώνουμε το κόστος του αλγορίθμου από το κόστος υλοποίησης: πόσο γρήγορα εκτελείται η αναφορά αφού αφαιρεθεί η περιττή εργασία;

Η βάση σύγκρισης διατηρεί το trie παιδιών-χάρτη κατακερματισμού (hash-map) της αναφοράς, αλλά κάνει δύο αλλαγές. Μια δομή scratch που ανήκει στον καλούντα διατηρεί τον πίνακα DP, τα προεπεξεργασμένα byte και τα ID των token εξόδου, οπότε η κύρια διαδρομή (hot path) δεν κάνει καμία δέσμευση μετά την προθέρμανση. Το ID του token και το σκορ αποθηκεύονται απευθείας στον κόμβο trie, αφαιρώντας τη μετατροπή συμβολοσειράς ανά ταίριασμα και την αναζήτηση στον παράλληλο χάρτη κατακερματισμού.

Εσωτερικός βρόχος:

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

Πρόκειται για τον ίδιο αλγόριθμο, με το ίδιο σχήμα trie και το ίδιο μοτίβο πρόσβασης, απλώς χωρίς δεσμεύσεις και χωρίς περιττές αναζητήσεις. Αποτελέσματα στα 514 token:

Μηχανή	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

Η μηδενικής δέσμευσης baseline σε σύγκριση με τις υπάρχουσες υλοποιήσεις tokenizer στα 514 tokens.

Η αφαίρεση της υλοποίησης string ανά ταίριασμα και του δευτερεύοντος lookup στο hash-map μειώνει τις εντολές κατά 2.4× σε σχέση με την υλοποίηση αναφοράς. Το IPC μειώνεται μαζί τους (2.86 → 2.62), επειδή το pointer chase στο hash-map, που προηγουμένως κρυβόταν πίσω από τη δουλειά του allocator, είναι πλέον το κυρίαρχο κόστος ανά byte.

Καθυστέρηση p50 και p99 του tokenizer αναφοράς και της μηδενικής δέσμευσης baseline ανά μέγεθος εισόδου.

Ποσοστό αστοχιών L2 cache του tokenizer αναφοράς και της μηδενικής δέσμευσης baseline ανά μέγεθος εισόδου.

Η βάση σύγκρισης μειώνει την καθυστέρηση p50 κατά 2-3× σε σύγκριση με την αναφορά σε κάθε μέγεθος, και το ποσοστό αστοχίας L2 παραμένει κοντά στο 2-3% έως τα 2k token σε αντίθεση με το 8-48% της αναφοράς. Το εναπομείναν κόστος είναι η αναζήτηση δεικτών HashMap, την οποία αντιμετωπίζει η επόμενη ενότητα.

Το trie χάρτη κατακερματισμού βασίζεται σε αναζήτηση δεικτών. Τα SP και IREE το λύνουν αυτό με ένα trie διπλού πίνακα (δείτε την προηγούμενη ενότητα) και δανειζόμαστε την ίδια δομή δεδομένων. Η διαφορά είναι ότι τα SP και IREE είναι βιβλιοθήκες γενικής χρήσης που υποστηρίζουν ροή εισόδου (streaming), offsets εξόδου και πολλαπλές λειτουργίες κωδικοποίησης· αυτή η ευελιξία τους κοστίζει επιπλέον λογιστική γύρω από το trie (ένα πλέγμα στο SP, ένας αγωγός οκτώ σταδίων στο IREE). Δεν χρειαζόμαστε τίποτα από αυτά, οπότε ενσωματώνουμε το trie απευθείας στον βρόχο Viterbi και παραλείπουμε τα υπόλοιπα.

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

Ένα βήμα byte κοστίζει δύο αναγνώσεις πίνακα, μια πρόσθεση ακεραίων και μια σύγκριση. Ολόκληρο το trie είναι γύρω στα 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%

Καθυστέρηση tokenizer, εντολές και ποσοστά αστοχιών cache μετά την υιοθέτηση του double-array trie.

Ο χρόνος εκτέλεσης μειώνεται κατά 2,3× από τη βάση σύγκρισης και 4,8× από την αναφορά. Οι εντολές μειώνονται κατά ένα επιπλέον 28% σε σχέση με τη βάση σύγκρισης. Η κατάργηση του πλέγματος και του αγωγού μάς επιτρέπει να εξοικονομούμε 1,10M εντολές ανά κωδικοποίηση σε σύγκριση με τα 1,82M του SP και 1,61M του IREE με το ίδιο σχήμα trie· η διαφορά οφείλεται στη λογιστική που διατηρούν για API γενικής χρήσης τα οποία δεν χρειαζόμαστε.

Το Darts trie από τη Βελτιστοποίηση 1 απαιτεί δύο εξαρτημένες φορτώσεις ανά βήμα byte. Πρώτα τη βασική μετατόπιση (base offset) του γονέα, και στη συνέχεια την καταχώριση ελέγχου (check entry) στην υποψήφια θύρα παιδιού για να επιβεβαιωθεί ότι ο γονέας την κατέχει. Η καταχώριση ελέγχου απλώς επικυρώνει την κυριότητα. Το ίδιο ερώτημα μπορεί να απαντηθεί από έναν χάρτη bit (bitmap) ανά κόμβο που καταγράφει ποια byte έχουν παιδιά. Βάσει της κατασκευής του trie διπλού πίνακα, το bit για το byte b στον χάρτη bit του γονέα τίθεται σε ισχύ ακριβώς όταν η υποψήφια θύρα για αυτό το byte ανήκει στον γονέα. Ο χάρτης bit και ο πίνακας ελέγχου φέρουν τις ίδιες πληροφορίες, οπότε μπορούμε να αφαιρέσουμε εξ ολοκλήρου τον πίνακα ελέγχου από τη δομή χρόνου εκτέλεσης.

Η αναζήτηση στον χάρτη bit μεταγλωττίζεται σε έναν ενιαίο έλεγχο bit σε μία λέξη 64-bit, ο οποίος εκτελείται σε περίπου έναν κύκλο. Η σχετική εντολή POPCNT, η οποία μετράει όλα τα ενεργά bit σε έναν καταχωρητή, είναι αυτή που χρησιμοποιούν άλλοι σχεδιασμοί trie με δείκτες bitmap για να υπολογίσουν μια μετατόπιση παιδιού απευθείας από τον χάρτη bit. Δεν τη χρειαζόμαστε επειδή η κατασκευή 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)
}

Παράλληλα, συσκευάζουμε τα τέσσερα πεδία ανά κόμβο (χάρτης bit, βάση, ID token και σκορ) σε μία γραμμή κρυφής μνήμης (cache line) 64 byte. Μια επίσκεψη σε έναν ενιαίο κόμβο φέρνει πλέον τα πάντα: το σκορ και το ID του token που χρειαζόμαστε στους τερματικούς κόμβους, καθώς και τον χάρτη bit και τη βάση για το επόμενο βήμα. Ο πίνακας ελέγχου χρησιμοποιείται μόνο κατά την κατασκευή και παραλείπεται από τη διάταξη χρόνου εκτέλεσης. Η ευρετηρίαση στην κρίσιμη διαδρομή δεν ελέγχεται επειδή η κατασκευή του trie και η δομή του βρόχου εγγυώνται ήδη τα όρια. Το τίμημα είναι ο Χώρος: στα 64 byte ανά κόμβο × ~780K θύρες, το trie αυξάνεται από ~9 MB σε ~50 MB.

Αποτελέσματα στα 514 token:

Μηχανή	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%
+ Χάρτης bit και Ενσωματωμένη Συσκευασία	65 µs	77 µs	1,06M	1,8K	20,53%	15,01%

Αποτελέσματα tokenizer στα 514 tokens μετά την προσθήκη bitmap και inline packing.

Ο χρόνος εκτέλεσης μειώνεται συστηματικά κατά ένα επιπλέον 4,5% στα 514 token και αυξάνεται ανάλογα με το μήκος της εισόδου, φτάνοντας το 6-7% στα 4k και 16k. Η παράλειψη του πίνακα ελέγχου μετατρέπει το βήμα trie ανά byte από δύο φορτώσεις γραμμής κρυφής μνήμης σε μία, πράγμα που σημαίνει 2,5× λιγότερες προσβάσεις L2 ανά κωδικοποίηση (1,8K έναντι των 4,6K του απλού Darts). Το ποσοστό αστοχίας εκτοξεύεται από 6% σε 20%, αλλά το μεγαλύτερο μέρος αυτού οφείλεται στον παρονομαστή: με λιγότερες από τις μισές προσβάσεις, οι αστοχίες που όντως συμβαίνουν καταλαμβάνουν μεγαλύτερο μερίδιο.

Τα δεδομένα του trie είναι περίπου 50 MB (780.000 κόμβοι στα 64 byte ο καθένας). Σε μια προεπιλεγμένη διεργασία Linux, αυτή η μνήμη υποστηρίζεται από σελίδες 4 KB, οπότε ολόκληρο το trie εκτείνεται σε περίπου 12.000 διακριτές εικονικές σελίδες, πολύ πέρα από όσα μπορεί να χωρέσει ο buffer μετάφρασης διευθύνσεων πρώτου επιπέδου της CPU (Translation Lookaside Buffer – το Sapphire Rapids διαθέτει 96 καταχωρήσεις). Κάθε βήμα Viterbi αγγίζει έναν διαφορετικό κόμβο, και κατά τη διάρκεια μιας κωδικοποίησης 512 token σπαταλάμε περίπου 9.000 κύκλους σε περιηγήσεις στον πίνακα σελίδων (page-table walks), περίπου 3% του προϋπολογισμού ανά κωδικοποίηση.

Η λύση είναι η υποστήριξη του trie με τεράστιες σελίδες (huge pages) 2 MB αντί για σελίδες 4 KB. Τα ίδια 50 MB εκτείνονται πλέον σε 25 σελίδες, χωρούν εξ ολοκλήρου στον Translation Lookaside Buffer, και οι ίδιοι οι κατάλογοι πινάκων σελίδων παύουν να ανταγωνίζονται το trie για την κρυφή μνήμη L2.

Το Linux έχει δύο τρόπους να υποστηρίξει τη μνήμη με σελίδες 2 MB. Ο ευκαιριακός είναι τα transparent huge pages (THP): το πρόγραμμά σας ζητά από τον kernel να αναβαθμίσει μια περιοχή μνήμης από σελίδες 4 KB σε 2 MB, και ένα νήμα του kernel στο παρασκήνιο προσπαθεί να το κάνει όταν έχει χρόνο και υπάρχει διαθέσιμη συνεχόμενη φυσική μνήμη. Η άλλη εναλλακτική είναι να δεσμεύσετε το trie απευθείας από ένα προκρατημένο pool huge-page που ο kernel παραχωρεί κατά το boot:

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 τεράστιες σελίδες είναι δεσμευμένες· το trie καταναλώνει 24 από αυτές. Μετά το mmap, το /proc/<pid>/smaps επιβεβαιώνει την περιοχή: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Το κέρδος είναι 3–12% στον χρόνο εκτέλεσης σε διάφορα μεγέθη εισόδου, συνεπές και ανάλογο με το μήκος του ενεργού συνόλου δεδομένων:

Token	Χωρίς huge pages	Με huge pages	Δ
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%

Καθυστέρηση tokenizer με και χωρίς huge pages σε διαφορετικά μεγέθη εισόδου.

Το μεγαλύτερο κέρδος σημειώνεται στα 4k token, όπου το ενεργό σύνολο δεδομένων είναι αρκετά μεγάλο ώστε η κίνηση του πίνακα σελίδων να ανταγωνίζεται ενεργά τα δεδομένα trie για το εύρος ζώνης της L2. Πέρα από τα 4k το κέρδος συρρικνώνεται επειδή οι αστοχίες L3 κυριαρχούν και οι περιηγήσεις στον πίνακα σελίδων αποτελούν μικρότερο μερίδιο του συνολικού χρόνου.

Συνδυάζοντας τις τρεις βελτιστοποιήσεις, ακολουθεί η μέτρηση της τελικής μηχανής σε σύγκριση με κάθε άλλο tokenizer RoBERTa που δοκιμάσαμε. Παράλληλα με τα εγγενή δυαρχικά αρχεία SentencePiece σε C++ και IREE σε C, κάναμε επίσης benchmark τα έτοιμα πακέτα wrapper Rust που θα χρησιμοποιούσε ένας χρήστης σε περιβάλλον παραγωγής, τα sentencepiece και iree-tokenizer. Το πακέτο sentencepiece εμφανίζει 1.6-1.9× υψηλότερη καθυστέρηση από το εγγενές δυαρχικό C++ σε κάθε μέγεθος εισόδου, με το πρόσθετο κόστος να προέρχεται από το γεγονός ότι ο wrapper δεσμεύει μια νέα λίστα κομματιών token σε κάθε κωδικοποίηση, καθένα από τα οποία διατηρεί το δικό του αντίγραφο του κειμένου του κομματιού. Το πακέτο iree-tokenizer προσαρμόζει το μέγεθος του εσωτερικού του buffer σε κάθε είσοδο, γεγονός που αποφεύγει μια σταθερή δέσμευση ανά κωδικοποίηση την οποία το εγγενές δυαρχικό IREE πληρώνει ανεξάρτητα από το μέγεθος της εισόδου· αυτό έχει τη μεγαλύτερη σημασία σε μικρές εισόδους όπου η δέσμευση θα κυριαρχούσε διαφορετικά στην πραγματική διαδικασία δημιουργίας token. Σε μήκη εισόδου παραγωγής, το σταθερό κόστος αποσβένεται και το συνολικό πρόσθετο κόστος ανά κλήση του wrapper λόγω της αντιγραφής των token πίσω στη Rust ωθεί την καθυστέρησή του πάνω από εκείνη του εγγενούς δυαρχικού.

Καθυστέρηση p50 και p99 ανά encode σε διαφορετικές υλοποιήσεις tokenizer και μεγέθη εισόδου.

Η τελική μηχανή εκτελεί τον ίδιο αλγόριθμο Viterbi πάνω στο ίδιο λεξιλόγιο 250.000 token, δεν πραγματοποιεί καμία δέσμευση μνήμης στην κρίσιμη διαδρομή και παράγει έξοδο ταυτόσημη σε επίπεδο token σε σχέση με την αναφορά. Κάθε βήμα byte είναι πλέον μια ενιαία φόρτωση γραμμής κρυφής μνήμης 64 byte, ένας έλεγχος χάρτη bit και μία πρόσθεση ακεραίων, σε αντίθεση με μια ανίχνευση HashMap, μια κλήση δεσμευτή και διάφορους κλάδους ειδικών περιπτώσεων στην υλοποίηση βάσης σύγκρισης. Οι εντολές ανά κωδικοποίηση μειώθηκαν από 3,66M σε 1,04M, μια μείωση κατά 3,5×. Ο χρόνος εκτέλεσης αντιστοιχεί σε αυτή την αναλογία στις μικρές εισόδους και διευρύνεται περαιτέρω στις μεγάλες εισόδους, όπου οι δεσμεύσεις ανά token της αναφοράς αρχίζουν να υπερχειλίζουν την L2 και την L3.

Στην πράξη, αυτή η εργασία μείωσε τη χρήση της CPU στο στρώμα προεπεξεργασίας της στοίβας εξαγωγής συμπερασμάτων (inference) κατά 5-6× και αφαίρεσε διψήφια χιλιοστά του δευτερολέπτου από την καθυστέρηση του reranker.

Το μεγαλύτερο μέρος της μείωσης της καθυστέρησης προήλθε από ελαφριές αλλαγές που αφαίρεσαν άσκοπη δουλειά: την κατάργηση των allocations ανά encode, την αντικατάσταση του trie με HashMap-children με ένα double-array και την αφαίρεση του overhead του γενικού API. Η χρήση του υπάρχοντος bitmap ανά κόμβο ως ελέγχου εγκυρότητας και η υποστήριξη του trie με huge pages 2 MB συνεισέφεραν μαζί 10-20% επιπλέον πάνω από τη μεγαλύτερη μετάβαση σε double-array trie, ανάλογα με το μέγεθος εισόδου.

Τα έτοιμα πακέτα wrapper Rust γύρω από τα SentencePiece και IREE απεικονίζουν το ίδιο σημείο αντίστροφα. Επιπλέον των εγγενών tokenizer σε C και C++, τα όρια FFI και οι δεσμεύσεις ανά κλήση που εισάγουν οι wrappers είναι ορατά στους αριθμούς καθυστέρησης. Ο πηγαίος κώδικας Rust για τον τελικό κωδικοποιητή Viterbi διατίθεται στη διεύθυνση github.com/perplexityai/pplx-garden.

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

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

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

Κίνητρο

Ο αλγόριθμος Viterbi

Υλοποίηση αναφοράς

Μια σημείωση για την προεπεξεργασία

Πού η υλοποίηση αναφοράς ξοδεύει τον χρόνο της

Πώς μετρήσαμε

Γιατί το SentencePiece και το IREE είναι ταχύτερα: Το Double-Array Trie

Βάση σύγκρισης (Baseline): Μια υλοποίηση αναφοράς χωρίς καμία δέσμευση μνήμης (Zero-Allocation)

Βελτιστοποίηση 1: Δομή Trie Διπλού Πίνακα Χωρίς Επιπλέον Αφαιρέσεις

Βελτιστοποίηση 2: Συσκευασία bitmap και inline

Βελτιστοποίηση 3: Τεράστιες σελίδες για το trie

Τελική σύγκριση

Συμπέρασμα

Βιβλιογραφία

Παραπομπή