यूनिग्राम टोकेनाइज़र CPU प्रदर्शन में सुधार

हमने अपने यूनिग्राम टोकेनाइज़र को एक केंद्रित प्रदर्शन परियोजना के रूप में स्क्रैच से फिर से लागू किया।

लेखकPerplexity Research

हमने अपने यूनिग्राम टोकेनाइज़र (Unigram tokenizer) को एक केंद्रित प्रदर्शन परियोजना के रूप में स्क्रैच से फिर से लागू किया। उत्पादन इनपुट लंबाई पर, नया एनकोडर p50 विलंबता को Hugging Face टोकेनाइज़र क्रेट की तुलना में लगभग 5×, SentencePiece (C++) की तुलना में ~2× और IREE के टोकेनाइज़र (C) की तुलना में ~1.5× कम करता है, जिसमें शून्य स्थिर-अवस्था हीप आवंटन (zero steady-state heap allocations) होता है। उत्पादन में इसने हमारे इंफ़रेंस स्टैक में CPU उपयोग को 5-6× कम कर दिया और रिरैंकर (reranker) विलंबता को दो अंकों की मिलीसेकंड तक कम किया।

प्रेरणा

LLM इंफ़रेंस लागत को आमतौर पर GPU कार्य के रूप में तैयार किया जाता है: KV कैश, अटेंशन कर्नेल, एक्सपर्ट रूटिंग। लेकिन उत्पादन ट्रैफ़िक का एक महत्वपूर्ण हिस्सा बहुत छोटे मॉडलों पर चलता है, जो फ्रंटियर ट्रांसफ़ॉर्मर्स की तुलना में दो से तीन परिमाण क्रम (orders of magnitude) छोटे होते हैं। एम्बेडिंग मॉडल, क्लासिफ़ायर और रिरैंकर अभी भी GPU पर चलते हैं, लेकिन फ़ॉरवर्ड पास इतना छोटा होता है कि इसके आस-पास की हर चीज़ मायने रखने लगती है।

सबसे स्पष्ट मामला एक रिरैंकर का है जो प्रति अनुरोध सैकड़ों उम्मीदवार दस्तावेज़ों को स्कोर करता है। एक छोटे मॉडल के साथ, GPU कंप्यूट आमतौर पर एकल-अंक मिलीसेकंड में समाप्त हो जाता है, लेकिन प्रत्येक इनपुट को पहले CPU-साइड टोकेनाइज़ेशन से गुज़रना पड़ता है, जो टेक्स्ट को शब्दावली IDs में बदल देता है। जब बैच बड़ा होता है, तो यह कुल अनुरोध विलंबता के एक महत्वपूर्ण हिस्से का कारण बन सकता है।

यह कार्य XLM-RoBERTa को लक्षित करता है, जो SentencePiece के साथ प्रशिक्षित 250K-टोकन यूनिग्राम शब्दावली है। फाइन-ट्यून किए गए RoBERTa-फ़ैमिली एनकोडर रैंकिंग, पुनर्प्राप्ति और समानता कार्यों के लिए एक सामान्य विकल्प हैं। हम पूरे एन्कोडिंग के दौरान एकल-थ्रेडेड प्रति-एनकोड विलंबता की रिपोर्ट करते हैं; उत्पादन में हम rayon का उपयोग करके कोर में समानांतर प्रसंस्करण (parallelize) करते हैं।

इस पोस्ट के दौरान हम Hugging Face के टोकेनाइज़र क्रेट को उस संदर्भ के रूप में मानते हैं जिसके विरुद्ध हम बेंचमार्क करते हैं, क्योंकि यह डिफ़ॉल्ट रस्ट (Rust) टोकेनाइज़र है जिसे अधिकांश टीमें चुनती हैं। SentencePiece (C++) और IREE (C) को तुलना के लिए अन्य दो प्रमुख उत्पादन कार्यान्वयन के रूप में शामिल किया गया है।

विटेर्बी एल्गोरिदम (The Viterbi Algorithm)

यूनिग्राम टोकेनाइज़ेशन, जिसे 2018 में कुडो द्वारा (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) पेश किया गया और SentencePiece में शिप किया गया, सेगमेंटेशन को सबसे-संभावित-पथ (most-probable-path) समस्या के रूप में तैयार करता है: प्रत्येक शब्दावली टोकन में एक सीखी गई लॉग-प्रायिकता होती है, और एनकोडर उस सेगमेंटेशन को चुनता है जिसके टोकन स्कोर का योग सबसे अधिक होता है। समस्या को विटेर्बी एल्गोरिदम (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

बाहरी लूप इनपुट लंबाई में रैखिक होता है, जो प्रत्येक पोजीशन पर सबसे लंबे शब्दावली टोकन (XLM-RoBERTa के लिए 48 UTF-8 बाइट्स) द्वारा सीमित होता है। आंतरिक ट्राई वॉक (trie walk) हॉट पाथ है: 16k-टोकन इनपुट पर, यह सैकड़ों हज़ारों ट्रांज़िशन चलाता है और प्रति एनकोड लाखों निर्देशों को रिटायर करता है।

संदर्भ कार्यान्वयन

Hugging Face के टोकेनाइज़र v0.22 सीधे पुनरावृत्ति को लागू करते हैं, जिसमें रिवर्स लुकअप के लिए हैश-मैप और प्रत्येक ट्राई नोड पर हैश-मैप-चिल्ड्रन होते हैं:

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 }

कोड पठनीय रस्ट है, लेकिन बहुत सारा काम होता है जो स्रोत स्तर पर अदृश्य है। अगला भाग बताता है कि समय कहाँ जाता है।

प्रीप्रोसेसिंग पर एक नोट

विटेर्बी के चलने से पहले, इनपुट यूनिकोड सामान्यीकरण (NFKC) और वर्ड-बाउंड्री मार्किंग से गुज़रता है। प्रोफ़ाइलिंग हमारे परिदृश्यों में प्रीप्रोसेसिंग को कुल एनकोड समय के 5% से कम पर रखती है, इसलिए हम पोस्ट के बाकी हिस्सों के लिए विटेर्बी फ़ॉरवर्ड पास पर ध्यान केंद्रित करते हैं।

जहाँ संदर्भ कार्यान्वयन अपना समय व्यतीत करता है

CPU कोड को गति देने के दो तरीके हैं: कम निर्देश निष्पादित करना, या उच्च IPC प्राप्त करना। Cachegrind पहले को मापता है, perf stat दूसरे को मापता है। इस पोस्ट का हर अनुकूलन एक या दूसरे पर टिकता है।

CPU चार चीजों पर रुकते हैं: कैश मिस (L1d ~4 चक्र, एक L3-मिस-टू-DRAM सैकड़ों है), आश्रित-लोड चेन जिन्हें आउट-ऑफ़-ऑर्डर इंजन समानांतर नहीं कर सकता, ब्रांच मिसप्रेडिक्ट्स (~15-20 चक्र प्रति गलत अनुमान फ़्लश किए गए), और एलोकेटर कॉल। संदर्भ यूनिग्राम इंजन कार्यान्वयन चारों को हिट करता है:

अवरोध (Bottleneck)

तंत्र

मापा गया प्रभाव

प्रति मैच आवंटन

मैच किए गए बाइट्स पर String::from_utf8 प्लस प्रति ट्राई मैच एक AHashMap लुकअप; टोकन ID पहले से ही लीफ़ में है, लेकिन संदर्भ इसे वहाँ संग्रहीत नहीं करता है

512 टोकन पर 7,295 आवंटन, 1k पर 14,542, 16k पर 299,171

प्रति बाइट पॉइंटर चेज़

प्रत्येक ट्राई नोड पर एक AHashMap, अगले बाइट को हैश करना और प्रति चरण दो पॉइंटर्स को डीरेफ़रेंस करना; नोड स्ट्रक्ट कोल्ड फ़ील्ड्स जैसे is_leaf को भी बंडल करता है, प्रत्येक यात्रा पर उन्हें कैश में खींचता है

L1d मिस दर कम बनी रहती है (~2.5% 1k पर) क्योंकि एलोकेटर हॉट पेज L1 में रहते हैं; लागत आश्रित-लोड विलंबता के रूप में दिखाई देती है

लंबी इनपुट पर L2 थ्रैशिंग

DP टेबल, प्रीफ़िक्स स्क्रैच, और आउटपुट टोकन सूची प्रत्येक कॉल पर ताज़ा आवंटित की जाती हैं

L2 मिस दर 128 टोकन पर 12% से बढ़कर 16k पर 50% हो जाती है क्योंकि संचयी आवंटन फ़ुटप्रिंट प्रति-कोर 2 MB L2 से अधिक हो जाता है

संदर्भ टोकनाइज़र की बाधाएँ, तंत्र, और मापे गए प्रभाव।

प्रमुख लागत प्रति मैच आवंटन है। 1k टोकन पर, संदर्भ कार्यान्वयन 2.8 के IPC पर प्रति एनकोड लगभग 7M निर्देशों को रिटायर (पूरा) करता है, जिनमें से अधिकांश एलोकेटर कॉल और स्ट्रिंग रूपांतरणों में होते हैं। L2 मिस केवल ~1k टोकन के बाद मायने रखते हैं, एक बार जब संचयी आवंटन फ़ुटप्रिंट प्रति-कोर कैश से अधिक हो जाता है।

इनपुट आकार के अनुसार संदर्भ टोकनाइज़र p50 और p99 विलंबता।
इनपुट आकार के अनुसार संदर्भ टोकनाइज़र L2 कैश मिस दर।
इनपुट आकारों के पार संदर्भ टोकनाइज़र द्वारा प्रति encode आवंटित बाइट्स।

प्रति-टोकन आवंटन स्थिर है: इनपुट आकार की परवाह किए बिना, प्रति टोकन लगभग 2 KB और ~18 आवंटन। विलंबता गैर-रेखीयता L2 से आती है: संचयी बाइट्स 1k टोकन के आसपास 2 MB को पार करते हैं, मिस दर 8% से बढ़कर ~50% हो जाती है, और p99 वक्र मैच करने के लिए झुक जाता है।

हमने कैसे मापा

हम तीन उत्पादन टोकेनाइज़र की तुलना करते हैं, सभी एक ही XLM-RoBERTa मॉडल के विरुद्ध एक ही यूनिग्राम एल्गोरिदम चला रहे हैं:

इंजन

भाषा

ट्राई (Trie)

Hugging Face

रस्ट (Rust)

HashMap-चिल्ड्रन प्रीफ़िक्स ट्राई

SentencePiece

C++

डार्ट्स डबल-एरे ट्राई

IREE टोकेनाइज़र

C

डार्ट्स डबल-एरे ट्राई

तुलना में उपयोग किए गए टोकनाइज़र इम्प्लीमेंटेशन, भाषाएँ, और ट्राइ प्रतिनिधित्व।

हम SentencePiece और IREE को उनके रस्ट रैपर के बजाय सीधे C/C++ FFI के माध्यम से लिंक करते हैं; रैपर प्रति-कॉल आवंटन और स्ट्रिंग रूपांतरण जोड़ते हैं जो अंतर्निहित लाइब्रेरी की लागत का हिस्सा नहीं हैं।

सभी इंजन पूर्ण अनुकूलन (O3, नेटिव टारगेटिंग, लिंक-टाइम ऑप्टिमाइज़ेशन) के साथ निर्मित हैं। मापन एक इंटेल ज़ीओन प्लेटिनम 8488C पर एक शांत कोर पर पिन किए गए हैं। प्रत्येक इंजन और इनपुट आकार के लिए हम वॉल-क्लॉक विलंबता (p50 और p99, 1,000 वार्मअप के बाद 10,000 पुनरावृत्ति), Valgrind cachegrind के माध्यम से रिटायर किए गए निर्देश, perf stat हार्डवेयर काउंटर से IPC और कैश मिस दरें, और मापी गई विंडो के चारों ओर एक वैश्विक एलोकेटर रैपर द्वारा गिने गए हीप आवंटन एकत्र करते हैं।

514 टोकन पर (512 से अधिक +2 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

514 टोकनों पर टोकनाइज़र विलंबता, निर्देश संख्या, IPC, और आवंटन।

SP और IREE संदर्भ की विलंबता के लगभग एक तिहाई पर उतरते हैं। वे लगभग आधे निर्देश रिटायर करते हैं, उच्च IPC पर चलते हैं, और एक से दो परिमाण क्रम कम आवंटित करते हैं।

SentencePiece और IREE तेज़ क्यों हैं: डबल-एरे ट्राई

SentencePiece और IREE दोनों 1989 में Aoe द्वारा पेश किए गए डबल-एरे ट्राई का उपयोग करते हैं (Aoe, 1989; Karoonboonyanan)। ट्राई दो सपाट पूर्णांक एरे, बेस और चेक में एन्कोड किया गया है:

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 की 250K शब्दावली के लिए, पूरा ट्राई ~9 MB निरंतर मेमोरी में फिट बैठता है (~780K स्लॉट × प्रति स्लॉट ऑफ़सेट, सत्यापन और टोकन ID के लिए 12 बाइट्स)। एक एकल एनकोड का हॉट वर्किंग सेट 100 KB के क्रम का है, जो L2 में आराम से रहता है, और हार्डवेयर प्रीफ़ेचर नियमित एक्सेस पैटर्न पर लॉक हो जाता है।

इसके विपरीत, संदर्भ प्रति नोड बच्चों का एक हैश-मैप रखता है, जिसकी लागत प्रति बाइट चरण लगभग चार आश्रित लोड होती है। बाइट को हैश करें, बकेट लोड करें, बकेट प्रविष्टि के लिए एक पॉइंटर का पालन करें, हीप पर चाइल्ड नोड के लिए दूसरे पॉइंटर का पालन करें, फिर अगले बाइट के लिए अपने स्वयं के हैश-मैप से परामर्श करें। हीप ऑब्जेक्ट्स बिखर जाते हैं और अधिकांश लोड कैश मिस कर देते हैं।

SP और IREE दोनों टोकन ID को सीधे ट्राई नोड में संग्रहीत करते हैं। संदर्भ ट्राई कार्यान्वयन केवल मैच किए गए बाइट्स लौटाता है, इसलिए एनकोडर को प्रत्येक मैच के लिए एक ताज़ा स्ट्रिंग आवंटित करनी पड़ती है, टोकन ID को पुनर्प्राप्त करने के लिए इसे साइड हैश-मैप में देखना पड़ता है, और अलग वेक्टर से स्कोर पढ़ना पड़ता है। 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 का उपयोग करता है, लेकिन लागत मॉडल समान है: टोकन IDs और स्कोर सपाट-एरे इंडेक्स के रूप में पहुंच योग्य हैं, ट्राई वॉक पर कोई आवंटन या हैशिंग नहीं है। IREE प्रति एनकोड केवल एक आउटपुट बफ़र आवंटित करता है; SP DP के दौरान प्रति मैच अतिरिक्त रूप से एक लैटिस नोड आवंटित करता है, जो 514 टोकन पर कुल लगभग 1,559 आवंटन करता है।

बेसलाइन: संदर्भ का एक शून्य-आवंटन पोर्ट

ट्राई को बदलने से पहले, हम कार्यान्वयन लागत से एल्गोरिदम लागत को अलग करते हैं: अनावश्यक काम हटा दिए जाने पर संदर्भ कितनी तेज़ी से चलता है?

बेसलाइन संदर्भ के हैश-मैप-चिल्ड्रन ट्राई को रखता है लेकिन दो बदलाव करता है। एक कॉलर-स्वामित्व वाला स्क्रैच स्ट्रक्ट DP टेबल, प्रीप्रोसेस्ड बाइट्स और आउटपुट टोकन IDs रखता है, इसलिए हॉट पाथ वार्मअप के बाद कोई आवंटन नहीं करता है। टोकन ID और स्कोर सीधे ट्राई नोड में संग्रहीत किए जाते हैं, जिससे प्रति-मैच स्ट्रिंग रूपांतरण और साइड हैश-मैप लुकअप हट जाता है।

आंतरिक लूप:

rust
for pos in 0..n { if scratch.prep[pos] & 0xC0 == 0x80 { continue; } let s0 = scratch.best_score[pos]; if s0 == f32::NEG_INFINITY { continue; } let mut node = 0u32; for end in pos..n { let byte = scratch.prep[end]; match self.nodes[node as usize].children.get(&byte) { Some(&child) => { node = child; let n = &self.nodes[node as usize]; if n.is_terminal { let total = s0 + n.score; if total > scratch.best_score[end + 1] { scratch.best_score[end + 1] = total; scratch.best_start[end + 1] = pos as u32; scratch.best_id[end + 1] = n.token_id; } } } None => break, } } }

यह वही एल्गोरिदम है, जिसमें वही ट्राई आकार और वही एक्सेस पैटर्न है, बस कोई आवंटन नहीं और कोई अनावश्यक लुकअप नहीं है। 514 टोकन पर परिणाम:

इंजन

p50

निर्देश

IPC

आवंटन

Hugging Face

326 µs

3.66M

2.86

7,295

SentencePiece

127 µs

1.81M

3.55

1,559

IREE

111 µs

1.58M

3.26

1

बेसलाइन

156 µs

1.53M

2.62

0

514 टोकनों पर मौजूदा टोकनाइज़र इम्प्लीमेंटेशनों की तुलना में शून्य-आवंटन बेसलाइन।

प्रति-मैच स्ट्रिंग मटेरियलाइज़ेशन और सेकेंडरी हैश-मैप लुकअप को हटाने से संदर्भ की तुलना में 2.4× निर्देश कम हो जाते हैं। IPC उनके साथ गिर जाता है (2.86 → 2.62) क्योंकि हैश-मैप पॉइंटर चेज़, जो पहले एलोकेटर कार्य के पीछे छिपा था, अब प्रमुख प्रति-बाइट लागत है।

इनपुट आकार के अनुसार संदर्भ और शून्य-आवंटन बेसलाइन टोकनाइज़र p50 और p99 विलंबता।
इनपुट आकार के अनुसार संदर्भ और शून्य-आवंटन बेसलाइन टोकनाइज़र L2 कैश मिस दर।

बेसलाइन प्रत्येक आकार पर संदर्भ की तुलना में p50 विलंबता को 2-3× कम करता है, और इसकी L2 मिस दर 2k टोकन के माध्यम से 2-3% के पास रहती है जबकि संदर्भ की 8-48% होती है। शेष लागत HashMap पॉइंटर चेज़ है, जिसे अगला अनुभाग संबोधित करता है।

अनुकूलन 1: अतिरिक्त अमूर्तता के बिना डबल-एरे ट्राई

हैश-मैप ट्राई एक पॉइंटर चेज़ है। SP और IREE इसे डबल-एरे ट्राई (पिछला अनुभाग देखें) के साथ हल करते हैं, और हम उसी डेटा संरचना को उधार लेते हैं। अंतर यह है कि SP और IREE स्ट्रीमिंग इनपुट, आउटपुट ऑफ़सेट और कई एन्कोडिंग मोड का समर्थन करने वाली सामान्य-उद्देश्य लाइब्रेरी हैं; वह लचीलापन उन्हें ट्राई के चारों ओर अतिरिक्त बहीखाता लागत (SP में एक लैटिस, IREE में आठ-चरण पाइपलाइन) देता है। हमें इसकी कोई आवश्यकता नहीं है, इसलिए हम विटेर्बी लूप में ट्राई को सीधे इनलाइन करते हैं और बाकी को हटा देते हैं।

rust
fn child(&self, node: u32, byte: u8) -> Option<u32> { let base = self.base[node as usize]; let next = base + byte as i32; if self.check[next as usize] == node as i32 { Some(next as u32) } else { None } }

एक बाइट चरण में दो एरे रीड, एक पूर्णांक जोड़ और एक तुलना की लागत होती है। पूरा ट्राई ~9 MB है, लेकिन कोई भी एकल एनकोड केवल ~100 KB ही चलता है, इसलिए हॉट डेटा L2 में रहता है।

इंजन

p50

p99

निर्देश

L2 मिस%

L3 मिस%

Hugging Face

327 µs

3,340 µs

3.66M

8.17%

14.66%

SentencePiece

127 µs

1,140 µs

1.82M

4.74%

4.07%

IREE

111 µs

132 µs

1.61M

5.82%

3.87%

बेसलाइन

155 µs

1,170 µs

1.53M

2.59%

16.66%

+ डार्ट्स

68 µs

80 µs

1.10M

6.45%

11.42%

डबल-एरे ट्राइ अपनाने के बाद टोकनाइज़र विलंबता, निर्देश, और कैश मिस दरें।

वॉल-क्लॉक बेसलाइन से 2.3× और संदर्भ से 4.8× गिरता है। निर्देश बेसलाइन से 28% और गिरते हैं। लैटिस और पाइपलाइन को काटने से हम वही ट्राई आकार रखने पर SP के 1.82M और IREE के 1.61M के विरुद्ध प्रति एनकोड 1.10M निर्देश रिटायर कर पाते हैं; अंतर वह बहीखाता है जो वे सामान्य-उद्देश्य APIs के लिए रखते हैं जिनकी हमें आवश्यकता नहीं है।

ऑप्टिमाइज़ेशन 2: बिटमैप और इनलाइन पैकिंग

अनुकूलन 1 के डार्ट्स ट्राई को प्रति बाइट चरण में दो आश्रित लोड की आवश्यकता होती है। पहले पैरेंट का बेस ऑफ़सेट, फिर यह पुष्टि करने के लिए कैंडिडेट चाइल्ड स्लॉट पर चेक प्रविष्टि कि पैरेंट इसका स्वामी है। चेक प्रविष्टि केवल स्वामित्व को मान्य करती है। वही प्रश्न एक प्रति-नोड बिटमैप द्वारा उत्तर दिया जा सकता है जो रिकॉर्ड करता है कि किस बाइट के बच्चे हैं। डबल-एरे ट्राई के निर्माण से, पैरेंट के बिटमैप में बाइट b के लिए बिट ठीक तब सेट होता है जब उस बाइट के लिए कैंडिडेट स्लॉट पैरेंट का होता है। बिटमैप और चेक एरे समान जानकारी रखते हैं, इसलिए हम रनटाइम संरचना से चेक एरे को पूरी तरह से हटा सकते हैं।

बिटमैप लुकअप एक 64-बिट वर्ड के विरुद्ध एक एकल बिट टेस्ट में संकलित होता है, जो लगभग एक चक्र में चलता है। संबंधित POPCNT निर्देश, जो एक रजिस्टर में सभी सेट बिट्स को गिनता है, वह है जिसका उपयोग अन्य बिटमैप-इंडेक्स ट्राई डिज़ाइन बिटमैप से सीधे चाइल्ड ऑफ़सेट की गणना करने के लिए करते हैं। हमें इसकी आवश्यकता नहीं है क्योंकि डार्ट्स निर्माण पहले से ही उस ऑफ़सेट को नोड के बेस फ़ील्ड में संग्रहीत करता है।

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 बाइट्स × ~780K स्लॉट पर ट्राई ~9 MB से बढ़कर ~50 MB हो जाता है।

514 टोकन पर परिणाम:

इंजन

p50

p99

निर्देश

L2 एक्सेस

L2 मिस%

L3 मिस%

Hugging Face

327 µs

3,340 µs

3.66M

15.0K

8.17%

14.66%

SentencePiece

127 µs

1,140 µs

1.82M

7.5K

4.74%

4.07%

IREE

111 µs

132 µs

1.61M

4.7K

5.82%

3.87%

बेसलाइन

155 µs

1,170 µs

1.53M

6.1K

2.59%

16.66%

+ डार्ट्स

68 µs

80 µs

1.10M

4.6K

6.45%

11.42%

+ बिटमैप और इनलाइन पैकिंग

65 µs

77 µs

1.06M

1.8K

20.53%

15.01%

बिटमैप और इनलाइन पैकिंग जोड़ने के बाद 514 टोकनों पर टोकनाइज़र परिणाम।

वॉल-क्लॉक व्यवस्थित रूप से 514 टोकन पर 4.5% गिरता है और इनपुट लंबाई के साथ बढ़ता है, 4k और 16k पर 6-7% तक पहुँच जाता है। चेक एरे को छोड़ने से प्रति-बाइट ट्राई चरण दो कैश-लाइन लोड से एक में बदल जाता है, जो प्रति एनकोड 2.5× कम L2 एक्सेस है (साधारण डार्ट्स के 4.6K के मुकाबले 1.8K)। मिस दर 6% से बढ़कर 20% हो जाती है, लेकिन इसका अधिकांश हिस्सा हर में है: आधे से कम एक्सेस के साथ, जो मिस हमारे पास होते हैं, उनका हिस्सा बड़ा हो जाता है।

ऑप्टिमाइज़ेशन 3: ट्राई के लिए ह्यूज पेज

ट्राई डेटा लगभग 50 MB है (64 बाइट्स पर 780,000 नोड्स)। एक डिफ़ॉल्ट लिनक्स प्रक्रिया पर, वह मेमोरी 4 KB पेजों द्वारा समर्थित है, इसलिए पूरा ट्राई लगभग 12,000 अलग वर्चुअल पेजों तक फैला है, जो CPU के पहले-स्तर के डेटा ट्रांसलेशन लुकअसाइड बफ़र (Translation Lookaside Buffer) की क्षमता से कहीं अधिक है (सैफायर रैपिड्स में 96 प्रविष्टियां हैं)। प्रत्येक विटेर्बी चरण एक अलग नोड को छूता है, और 512-टोकन एनकोड पर हम पेज-टेबल वॉक में लगभग 9,000 चक्र जलाते हैं, जो प्रति-एनकोड बजट का लगभग 3% है।

इसका समाधान ट्राई को 4 KB के बजाय 2 MB ह्यूज पेजों के साथ बैक करना है। वही 50 MB अब 25 पेजों तक फैला है, ट्रांसलेशन लुकअसाइड बफ़र में पूरी तरह फिट बैठता है, और पेज-टेबल निर्देशिकाएँ स्वयं L2 कैश के लिए ट्राई के साथ प्रतिस्पर्धा करना बंद कर देती हैं।

लिनक्स के पास मेमोरी को 2 MB पेजों के साथ बैक करने के दो तरीके हैं। अवसरवादी पारदर्शी ह्यूज पेज (THP) है: आपका प्रोग्राम कर्नेल से 4 KB से 2 MB पेजों में मेमोरी क्षेत्र को अपग्रेड करने के लिए कहता है, और एक बैकग्राउंड कर्नेल थ्रेड इसे तब करने की कोशिश करता है जब उसके पास समय और निरंतर भौतिक मेमोरी उपलब्ध हो। दूसरा विकल्प ट्राई को सीधे एक प्री-रिजर्व्ड ह्यूज-पेज पूल से आवंटित करना है जिसे कर्नेल बूट पर अलग रखता है:

rust
let raw = unsafe { libc::mmap( std::ptr::null_mut(), alloc_len, libc::PROT_READ | libc::PROT_WRITE, libc::MAP_PRIVATE | libc::MAP_ANONYMOUS | libc::MAP_HUGETLB, -1, 0, ) };

पूल का होस्ट पर होना आवश्यक है (vm.nr_hugepages बूट पर सेट)। हमारे उत्पादन उदाहरणों पर, 10,561 ह्यूज पेज आरक्षित हैं; ट्राई उनमें से 24 का उपभोग करता है। mmap के बाद, /proc/<pid>/smaps क्षेत्र की पुष्टि करता है: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB।

लाभ इनपुट आकारों के पार 3-12% वॉल-क्लॉक है, जो वर्किंग-सेट लंबाई के लिए सुसंगत और आनुपातिक है:

टोकन

ह्यूज पेजों के बिना

ह्यूज पेजों के साथ

Δ

130

15.7 µs

15.2 µs

−3.6%

514

65.4 µs

63.1 µs

−3.4%

1,026

129.5 µs

125.3 µs

−3.3%

4,098

773 µs

679 µs

−12.0%

16,386

3,808 µs

3,554 µs

−6.8%

इनपुट आकारों के पार विशाल पृष्ठों के साथ और बिना टोकनाइज़र विलंबता।

सबसे बड़ी जीत 4k टोकन पर है, जहाँ वर्किंग सेट इतना बड़ा है कि पेज-टेबल ट्रैफ़िक L2 बैंडविड्थ के लिए ट्राई डेटा के साथ सक्रिय रूप से प्रतिस्पर्धा कर रहा था। 4k के बाद लाभ कम हो जाता है क्योंकि L3 मिस हावी हो जाते हैं और पेज-टेबल वॉक कुल समय का एक छोटा हिस्सा बन जाते हैं।

अंतिम तुलना

तीनों अनुकूलनों को एक साथ लाते हुए, यहाँ अंतिम इंजन है जिसे हमने परीक्षण किए गए हर दूसरे RoBERTa टोकेनाइज़र के विरुद्ध मापा। नेटिव C++ SentencePiece और C IREE बाइनरी के साथ, हमने उन ऑफ़-द-शेल्फ़ रस्ट रैपर क्रेट्स का भी बेंचमार्क किया जिन्हें एक उत्पादन उपयोगकर्ता चुनेगा, sentencepiece और iree-tokenizer। Sentencepiece क्रेट प्रत्येक इनपुट आकार पर नेटिव C++ बाइनरी की तुलना में 1.6-1.9× उच्च विलंबता मापता है, जिसमें ओवरहेड रैपर द्वारा प्रत्येक एनकोड पर टोकन टुकड़ों की एक ताज़ा सूची आवंटित करने से आता है, जिनमें से प्रत्येक टुकड़े के टेक्स्ट की अपनी प्रति रखता है। Iree-tokenizer क्रेट प्रत्येक इनपुट के लिए अपने आंतरिक बफ़र को आकार देता है, जो एक निरंतर प्रति-एनकोड आवंटन से बचाता है जो नेटिव IREE बाइनरी इनपुट आकार की परवाह किए बिना भुगतान करता है, और यह छोटे इनपुट पर सबसे अधिक मायने रखता है जहाँ आवंटन अन्यथा वास्तविक टोकेनाइज़ेशन कार्य पर हावी हो जाएगा। उत्पादन इनपुट लंबाई पर निरंतर लागत परिशोधित हो जाती है और टोकन को रस्ट में वापस कॉपी करने से रैपर का अपना प्रति-कॉल ओवरहेड इसकी विलंबता को नेटिव बाइनरी से ऊपर धकेलता है।

टोकनाइज़र इम्प्लीमेंटेशनों और इनपुट आकारों के पार प्रति encode p50 और p99 विलंबता।

अंतिम इंजन एक ही 250,000-टोकन शब्दावली पर वही विटेर्बी चलाता है, हॉट पाथ पर कुछ भी आवंटित नहीं करता है, और संदर्भ के विरुद्ध टोकन-सटीक आउटपुट उत्पन्न करता है। प्रत्येक बाइट चरण अब एक एकल 64-बाइट कैश-लाइन लोड, एक बिटमैप टेस्ट, और एक पूर्णांक जोड़ है, जो बेसलाइन कार्यान्वयन में HashMap जांच, एलोकेटर कॉल और कई विशेष-मामले शाखाओं से कम है। प्रति एनकोड निर्देश 3.66M से 1.04M तक गिर गए, जो 3.5× की गिरावट है। वॉल-क्लॉक छोटे इनपुट पर उस अनुपात से मेल खाता है और लंबे इनपुट पर और अधिक बढ़ जाता है, जहाँ संदर्भ का प्रति-टोकन आवंटन L2 और L3 को ओवरफ़्लो करना शुरू कर देता है।

समापन विचार

व्यवहार में इस कार्य ने हमारे इंफ़रेंस स्टैक की प्रीप्रोसेसिंग लेयर में CPU उपयोग को 5-6× कम कर दिया और रिरैंकर विलंबता को दो अंकों की मिलीसेकंड तक कम कर दिया।

विलंबता में कमी का अधिकांश हिस्सा हल्के परिवर्तनों से आया जिन्होंने बर्बाद काम को हटा दिया: प्रति-एनकोड आवंटन छोड़ना, HashMap-चिल्ड्रन ट्राई को डबल-एरे से बदलना, और जेनेरिक-API ओवरहेड को हटाना। वैधता जांच के रूप में मौजूदा प्रति-नोड बिटमैप का उपयोग करना और ट्राई को 2 MB ह्यूज पेजों के साथ बैक करना इनपुट आकार के आधार पर डबल-एरे ट्राई पर बड़े स्विच के अलावा 10-20% के लिए जिम्मेदार है।

SentencePiece और IREE के चारों ओर ऑफ़-द-शेल्फ़ रस्ट रैपर क्रेट्स उसी बिंदु को उल्टा स्पष्ट करते हैं। नेटिव C और C++ टोकेनाइज़र के ऊपर, FFI बाउंड्री और रैपर द्वारा पेश किए गए प्रति-कॉल आवंटन विलंबता संख्याओं में दिखाई देते हैं। अंतिम विटेर्बी एनकोडर के लिए रस्ट स्रोत github.com/perplexityai/pplx-garden पर उपलब्ध है।

संदर्भ (References)

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

ह्यूज पेज — 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}
}