యూనిగ్రామ్ టోకనైజర్ CPU పనితీరును మెరుగుపరచడం
మేము పనితీరు మెరుగుదల ప్రాజెక్ట్గా మా యూనిగ్రామ్ టోకనైజర్ను మొదటి నుండి పునఃరూపకల్పన చేశాము.
మేము పనితీరు మెరుగుదల ప్రాజెక్ట్గా మా యూనిగ్రామ్ టోకనైజర్ను మొదటి నుండి పునఃరూపకల్పన చేశాము. ఉత్పత్తి ఇన్పుట్ పొడవుల వద్ద, కొత్త ఎన్కోడర్ Hugging Face టోకనైజర్స్ క్రేట్తో పోలిస్తే p50 లేటెన్సీని సుమారు 5×, సెంటెన్స్పీస్ (C++)తో పోలిస్తే ~2×, మరియు IREE టోకనైజర్ (C)తో పోలిస్తే ~1.5× తగ్గిస్తుంది, దీనిలో స్థిరమైన హీప్ అలొకేషన్స్ సున్నా. ఉత్పత్తిలో ఇది మా ఇన్ఫరెన్స్ స్టాక్లో CPU వినియోగాన్ని 5-6× తగ్గించింది మరియు రీరాంకర్ లేటెన్సీ నుండి డబుల్-అంకె మిలిసెకండ్లను తగ్గించింది.
ప్రేరణ
LLM ఇన్ఫరెన్స్ ఖర్చు సాధారణంగా GPU పనిగా పరిగణించబడుతుంది: KV కాష్లు, అటెన్షన్ కర్నెల్స్, ఎక్స్పర్ట్ రూటింగ్. కానీ ప్రొడక్షన్ ట్రాఫిక్లో గణనీయమైన భాగం చాలా చిన్న మోడళ్లపై నడుస్తుంది, ఇవి ఫ్రాంటియర్ ట్రాన్స్ఫార్మర్ల కంటే రెండు నుండి మూడు ఆర్డర్ల పరిమాణం చిన్నవిగా ఉంటాయి. ఎంబెడింగ్ మోడల్స్, క్లాసిఫైయర్లు మరియు రీరాంకర్లు ఇప్పటికీ GPUలపైనే నడుస్తాయి, కానీ ఫార్వర్డ్ పాస్ చాలా చిన్నదిగా ఉండటం వలన దాని చుట్టూ ఉన్న ప్రతిదీ ప్రాముఖ్యత సంతరించుకుంటుంది.
అత్యంత స్పష్టమైన ఉదాహరణ ఏమిటంటే, ప్రతి అభ్యర్థనకు వందలాది అభ్యర్థి పత్రాలను స్కోర్ చేసే రీరాంకర్. చిన్న మోడల్తో, GPU గణన సాధారణంగా సింగిల్-డిజిట్ మిలిసెకండ్లలో ముగుస్తుంది, కానీ ప్రతి ఇన్పుట్ మొదట CPU-సైడ్ టోకనైజేషన్ గుండా వెళ్ళాలి, ఇది టెక్స్ట్ను వొకాబ్యులరీ IDలుగా మారుస్తుంది. బ్యాచ్ పెద్దదిగా ఉన్నప్పుడు, ఇది మొత్తం అభ్యర్థన లేటెన్సీలో గణనీయమైన భాగాన్ని ఆక్రమించగలదు.
ఈ పని XLM-RoBERTaను లక్ష్యంగా చేసుకుంది, ఇది SentencePieceతో శిక్షణ పొందిన 250K-టోకెన్ యూనిగ్రామ్ వొకాబ్యులరీ. ఫైన్-ట్యూన్ చేయబడిన RoBERTa-కుటుంబ ఎన్కోడర్లు రేటింగ్, రిట్రీవల్ మరియు సిమిలారిటీ టాస్క్ల కోసం సాధారణ ఎంపిక. మేము ప్రతి ఎన్కోడ్కు సింగిల్-థ్రెడెడ్ లేటెన్సీని నివేదిస్తాము; ఉత్పత్తిలో మేము rayon ఉపయోగించి కోర్లలో సమాంతరీకరణ చేస్తాము.
ఈ పోస్టు అంతటా, బెంచ్మార్క్గా మనం సరిపోల్చుకునే సూచికగా Hugging Face యొక్క tokenizers crateను తీసుకుంటున్నాము, ఎందుకంటే ఇది ఎక్కువ టీమ్లు ఎంచుకునే డిఫాల్ట్ Rust tokenizer. SentencePiece (C++) మరియు IREE (C)లను ఇతర రెండు ప్రధాన ప్రొడక్షన్ అమలులుగా పోలిక కోసం చేర్చాము.
విటర్బి అల్గారిథమ్
2018లో కుడో చేత పరిచయం చేయబడిన (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) మరియు SentencePieceలో అందించబడిన యూనిగ్రామ్ టోకనైజేషన్, సెగ్మెంటేషన్ను అత్యంత-సంభావ్యత-మార్గం సమస్యగా రూపొందిస్తుంది: ప్రతి వొకాబ్యులరీ టోకెన్ ఒక నేర్చుకున్న లాగ్-సంభావ్యతను కలిగి ఉంటుంది, మరియు ఎన్కోడర్ టోకెన్ స్కోర్ల మొత్తం అత్యధిక విలువను ఇచ్చే సెగ్మెంటేషన్ను ఎంచుకుంటుంది. ఈ సమస్య Viterbi algorithm (1967) తో పరిష్కరించబడుతుంది, బైట్ స్థానాలు గ్రాఫ్ పొరలుగా మరియు వొకాబ్యులరీ టోకెన్లు నిరంతర బైట్ పరిధిని విస్తరించే అంచులుగా ఉంటాయి. పునరావృతం అనేది బైట్ స్థానాలపై ఒక డైమెన్షనల్ DP:
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 బైట్లు). లోపలి ట్రీ వాక్ అనేది హాట్ పాత్: 16k-టోకెన్ ఇన్పుట్పై, ఇది లక్షలాది పరివర్తనాలను నిర్వహిస్తుంది మరియు ప్రతి ఎన్కోడ్కు కోట్లాది సూచనలను రిటైర్ చేస్తుంది.
రిఫరెన్స్ ఇంప్లిమెంటేషన్
Hugging Face యొక్క tokenizers v0.22 పునరావృతాన్ని నేరుగా అమలు చేస్తుంది, రివర్స్ లుకప్ కోసం హ్యాష్-మ్యాప్ మరియు ప్రతి ట్రీ నోడ్ వద్ద హ్యాష్-మ్యాప్-చిల్డ్రన్తో:
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
}మరియు లోపలి లూప్ (సరళీకృతమైనది) – ప్రతి ప్రిఫిక్స్-సెర్చ్ మ్యాచ్ తాజా స్ట్రింగ్ను కేటాయిస్తుంది మరియు ఆపై హ్యాష్-మ్యాప్ను పరిశీలిస్తుంది:
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
}కోడ్ చదవదగిన రస్ట్, కానీ మూల స్థాయిలో కనిపించని చాలా పని జరుగుతుంది. సమయం ఎక్కడికి వెళుతుందో తదుపరి విభాగం వివరిస్తుంది.
ప్రీప్రాసెసింగ్పై ఒక గమనిక
Viterbi రన్ అయ్యే ముందు, ఇన్పుట్ యూనికోడ్ సాధారణీకరణ (NFKC) మరియు వర్డ్-బౌండరీ మార్కింగ్ గుండా వెళుతుంది. మా దృశ్యాలలో మొత్తం ఎన్కోడ్ సమయంలో ప్రీప్రాసెసింగ్ 5% కంటే తక్కువగా ఉంటుందని ప్రొఫైలింగ్ చూపుతుంది, కాబట్టి మేము పోస్ట్ మిగిలిన భాగం కోసం Viterbi ఫార్వర్డ్ పాస్పై దృష్టి పెడతాము.
రిఫరెన్స్ ఇంప్లిమెంటేషన్ తన సమయాన్ని ఎక్కడ ఖర్చు చేస్తుంది
CPU కోడ్ను వేగవంత చేయడానికి రెండు మార్గాలు ఉన్నాయి: తక్కువ సూచనలను అమలు చేయడం, లేదా అధిక IPCని పొందడం. Cachegrind మొదటిదాన్ని కొలుస్తుంది, perf stat రెండవదాన్ని కొలుస్తుంది. ఈ పోస్ట్లోని ప్రతి ఆప్టిమైజేషన్ ఏదో ఒకదానిపై వర్తిస్తుంది.
CPUలు నాలుగు విషయాలపై ఆగిపోతాయి: కాష్ మిస్సులు (L1d ~4 సైకిల్స్, L3-మిస్-టు-DRAM వందల సంఖ్యలో ఉంటుంది), అవుట్-ఆఫ్-ఆర్డర్ ఇంజిన్ సమాంతరీకరించలేని డిపెండెంట్-లోడ్ చైన్స్, బ్రాంచ్ మిస్ప్రెడిక్ట్స్ (~తప్పు అంచనాకు 15-20 సైకిల్స్ ఫ్లష్ చేయబడతాయి), మరియు అలోకేటర్ కాల్స్. రిఫరెన్స్ యూనిగ్రామ్ ఇంజిన్ అమలు నాలుండింటినీ తాకుతుంది:
బాటిల్నెక్ | మెకానిజం | కొలవబడిన ప్రభావం |
ప్రతి మ్యాచ్కు అలోకేషన్ | సరిపోలిన బైట్లపై String::from_utf8 మరియు ప్రతి ట్రీ మ్యాచ్కు ఒక AHashMap లుకప్; టోకెన్ ID ఇప్పటికే లీఫ్లో ఉంది, కానీ రిఫరెన్స్ దానిని అక్కడ నిల్వ చేయదు | 512 టోకెన్ల వద్ద 7,295 అలోకేషన్స్, 1k వద్ద 14,542, 16k వద్ద 299,171 |
ప్రతి బైట్కు పాయింటర్ ఛేజ్ | ప్రతి ట్రీ నోడ్ వద్ద ఒక AHashMap, తదుపరి బైట్ను హ్యాష్ చేయడం మరియు ప్రతి దశకు రెండు పాయింటర్లను డీరెఫరెన్స్ చేయడం; Node స్ట్రక్ట్ is_leaf వంటి కోల్డ్ ఫీల్డ్లను కూడా బండిల్ చేస్తుంది, ప్రతి సందర్శనలో వాటిని కాష్లోకి లాగుతుంది | L1d మిస్ రేట్ తక్కువగా ఉంటుంది (1k వద్ద ~2.5%) ఎందుకంటే అలోకేటర్ హాట్ పేజీలు L1లో నివసిస్తాయి; ఖర్చు డిపెండెంట్-లోడ్ లేటెన్సీగా కనిపిస్తుంది |
పొడవైన ఇన్పుట్లపై L2 థ్రాషింగ్ | DP టేబుల్, ప్రిఫిక్స్ స్క్రాచ్, మరియు అవుట్పుట్ టోకెన్ జాబితా అన్నీ ప్రతి కాల్కి కొత్తగా కేటాయించబడతాయి | సంచిత అలోకేషన్ ఫుట్ప్రింట్ ప్రతి-కోర్ 2 MB L2ను మించినందున L2 మిస్ రేట్ 128 టోకెన్ల వద్ద 12% నుండి 16k వద్ద 50%కి పెరుగుతుంది |
సూచనాత్మక టోకెనైజర్ bottlenecks, mechanisms, మరియు కొలిచిన ప్రభావాలు.
ప్రధాన ఖర్చు ప్రతి మ్యాచ్కు జరిగే కేటాయింపు. 1k టోకెన్ల వద్ద, రిఫరెన్స్ ఇంప్లిమెంటేషన్ 2.8 IPC వద్ద ప్రతి ఎన్కోడ్కు సుమారు 7M సూచనలను రిటైర్ చేస్తుంది (పూర్తి చేస్తుంది), వాటిలో ఎక్కువ భాగం అలోకేటర్ కాల్స్ మరియు స్ట్రింగ్ మార్పిడిలలో ఉంటాయి. సంచిత అలోకేషన్ ఫుట్ప్రింట్ ప్రతి-కోర్ కాష్ను మించిపోయిన తర్వాత, L2 మిస్సులు ~1k టోకెన్ల దాటిన తర్వాత మాత్రమే ముఖ్యమైనవి.
ప్రతి టోకెన్కు allocation స్థిరంగానే ఉంటుంది: ఇన్పుట్ పరిమాణం ఎలా ఉన్నా, సుమారు 2 KB మరియు టోకెన్కు ~18 allocations. లేటెన్సీ రేఖీయేతరత L2 వల్ల వస్తుంది: కూడబడ్డ bytes 1k tokens వద్ద సుమారు 2 MB దాటుతాయి, miss rate 8% నుండి ~50% వరకు పెరుగుతుంది, మరియు p99 curve దానికనుగుణంగా వంగుతుంది.
మేము ఎలా కొలిచాము
మేము మూడు ప్రొడక్షన్ టోకనైజర్లను పోల్చాము, అన్నీ అదే XLM-RoBERTa మోడల్కు వ్యతిరేకంగా అదే యూనిగ్రామ్ అల్గారిథమ్ను నడుపుతున్నాయి:
ఇంజిన్ | భాష | ట్రీ |
Hugging Face | రస్ట్ | HashMap-children ప్రిఫిక్స్ ట్రీ |
సెంటెన్స్పీస్ | C++ | Darts డబుల్-అరే ట్రీ |
IREE టోకనైజర్ | C | Darts డబుల్-అరే ట్రీ |
పోలికలో ఉపయోగించిన టోకెనైజర్ అమలులు, భాషలు, మరియు trie ప్రతినిధీకరణలు.
మేము SentencePiece మరియు IREEలను వాటి రస్ట్ రేప్పర్ల కంటే ప్రత్యక్ష C/C++ FFI ద్వారా లింక్ చేస్తాము; రేప్పర్లు అంతర్లీన లైబ్రరీ ఖర్చులో భాగం కాని ప్రతి-కాల్ కేటాయింపులు మరియు స్ట్రింగ్ మార్పిడిలను జోడిస్తాయి.
అన్ని ఇంజిన్లు పూర్తి ఆప్టిమైజేషన్లతో నిర్మించబడ్డాయి (O3, నేటివ్ టార్గెటింగ్, లింక్-టైమ్ ఆప్టిమైజేషన్). కొలతలు Intel Xeon Platinum 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 |
సెంటెన్స్పీస్ | 128 µs | 1.83M | 3.39 | 1,559 |
IREE | 112 µs | 2.28M | 3.06 | 1 |
514 టోకెన్ల వద్ద టోకెనైజర్ latency, instruction count, IPC, మరియు allocations.
SP మరియు IREE రిఫరెన్స్ లేటెన్సీలో సుమారు మూడవ వంతు వద్ద ఉంటాయి. అవి దాదాపు సగం సూచనలను రిటైర్ చేస్తాయి, అధిక IPC వద్ద నడుస్తాయి, మరియు ఒకటి నుండి రెండు ఆర్డర్ల పరిమాణం తక్కువగా కేటాయిస్తాయి.
SentencePiece మరియు IREE ఎందుకు వేగంగా ఉన్నాయి: డబుల్-అరే ట్రీ
SentencePiece మరియు IREE రెండూ 1989లో Aoe చేత పరిచయం చేయబడిన డబుల్-అరే ట్రీని ఉపయోగిస్తాయి (Aoe, 1989; Karoonboonyanan). ట్రీ రెండు ఫ్లాట్ పూర్ణాంక శ్రేణులలో ఎన్కోడ్ చేయబడింది, base మరియు check:
int base[N]; // base[node] = offset for this node's children
int check[N]; // check[child] == node validates the parent-child linkచైల్డ్ లుకప్ అనేది:
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 ట్రీ నిల్వ చేసిన విలువగా IDని అందిస్తుంది మరియు ఫ్లాట్ ఫ్లోట్ అరే నుండి స్కోర్ను చదువుతుంది:
// 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 అన్ని matchesను ఒకే callలో తిరిగి ఇవ్వకుండా stateful cursor APIని ఉపయోగిస్తుంది, కానీ cost model అదే: token IDs మరియు scores flat-array indicesగా అందుబాటులో ఉంటాయి, trie walk స్వయంగా ఎలాంటి allocation లేదా hashing చేయదు. IREE ప్రతి encodeకు కేవలం ఒక output bufferను మాత్రమే allocate చేస్తుంది; SP అదనంగా DP సమయంలో ప్రతి matchకు ఒక lattice nodeను allocate చేస్తుంది, మొత్తం సుమారు 1,559 allocations 514 tokens వద్ద.
బేస్లైన్: రిఫరెన్స్ యొక్క జీరో-అలొకేషన్ పోర్ట్
ట్రీని మార్చడానికి ముందు, మేము అల్గారిథమ్ ఖర్చును ఇంప్లిమెంటేషన్ ఖర్చు నుండి వేరు చేస్తాము: అనవసరమైన పని తొలగించబడితే రిఫరెన్స్ ఎంత వేగంగా నడుస్తుంది?
బేస్లైన్ రిఫరెన్స్ యొక్క హ్యాష్-మ్యాప్-చిల్డ్రన్ ట్రీని ఉంచుతుంది కానీ రెండు మార్పులు చేస్తుంది. కాలర్-యాన్డ్ స్క్రాచ్ స్ట్రక్ట్ DP టేబుల్, ప్రీప్రాసెస్ చేయబడిన బైట్లు, మరియు అవుట్పుట్ టోకెన్ IDలను కలిగి ఉంటుంది, కాబట్టి వార్మప్ తర్వాత హాట్ పాత్ ఎలాంటి కేటాయింపులను చేయదు. టోకెన్ ID మరియు స్కోర్ నేరుగా ట్రీ నోడ్లో నిల్వ చేయబడతాయి, ప్రతి-మ్యా మ్యాచ్ స్ట్రింగ్ మార్పిడి మరియు సెకండరీ హ్యాష్-మ్యాప్ లుకప్ను తొలగిస్తాయి.
ఇన్నర్ లూప్:
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 |
సెంటెన్స్పీస్ | 127 µs | 1.81M | 3.55 | 1,559 |
IREE | 111 µs | 1.58M | 3.26 | 1 |
బేస్లైన్ | 156 µs | 1.53M | 2.62 | 0 |
514 టోకెన్ల వద్ద ఉన్న ప్రస్తుత టోకెనైజర్ అమలులతో పోల్చిన zero-allocation baseline.
ప్రతి-మ్యాచ్ స్ట్రింగ్ మెటీరియలైజేషన్ మరియు సెకండరీ హ్యాష్-మ్యాప్ లుకప్ను తొలగించడం రిఫరెన్స్తో పోలిస్తే సూచనలను 2.4× తగ్గిస్తుంది. వాటితో పాటు IPC తగ్గుతుంది (2.86 → 2.62) ఎందుకంటే హ్యాష్-మ్యాప్ పాయింటర్ ఛేజ్, గతంలో అలోకేటర్ పని వెనుక దాచబడి ఉండేది, ఇప్పుడు ప్రతి-బైట్ ప్రధాన ఖర్చుగా మారుతుంది.
బేస్లైన్ ప్రతి పరిమాణంలో రిఫరెన్స్తో పోలిస్తే p50 లేటెన్సీని 2-3× తగ్గిస్తుంది, మరియు దాని L2 మిస్ రేట్ రిఫరెన్స్ యొక్క 8-48% తో పోలిస్తే 2k టోకెన్ల వరకు 2-3% వద్ద ఉంటుంది. మిగిలిన ఖర్చు HashMap పాయింటర్ ఛేజ్, దీనిని తదుపరి విభాగం పరిష్కరిస్తుంది.
ఆప్టిమైజేషన్ 1: అదనపు అబ్స్ట్రాక్షన్స్ లేని డబుల్-అరే ట్రీ
హ్యాష్-మ్యాప్ ట్రీ అనేది పాయింటర్ ఛేజ్. SP మరియు IREE దీనిని డబుల్-అరే ట్రీతో పరిష్కరిస్తాయి (మునుపటి విభాగాన్ని చూడండి), మరియు మేము అదే డేటా స్ట్రక్చర్ను అరువు తీసుకుంటాము. వ్యత్యాసం ఏమిటంటే SP మరియు IREE స్ట్రీమింగ్ ఇన్పుట్, అవుట్పుట్ ఆఫ్సెట్లు, మరియు బహుళ ఎన్కోడింగ్ మోడ్లకు మద్దతు ఇచ్చే జనరల్-పర్పస్ లైబ్రరీలు; ఆ సౌలభ్యం వాటికి ట్రీ చుట్టూ అదనపు బుక్కీపింగ్ను ఖర్చు చేస్తుంది (SPలో లాటిస్, IREEలో ఎనిమిది-దశల పైప్లైన్). మాకు వీటిలో ఏదీ అవసరం లేదు, కాబట్టి మేము ట్రీని నేరుగా Viterbi లూప్లో ఇన్లైన్ చేస్తాము మరియు మిగిలిన వాటిని వదిలివేస్తాము.
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% |
సెంటెన్స్పీస్ | 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% |
double-array trieను స్వీకరించిన తర్వాత టోకెనైజర్ latency, instructions, మరియు cache miss rates.
వాల్-క్లాక్ బేస్లైన్ నుండి 2.3× మరియు రిఫరెన్స్ నుండి 4.8× పడిపోతుంది. సూచనలు బేస్లైన్ నుండి మరో 28% పడిపోతాయి. లాటిస్ మరియు పైప్లైన్ను తొలగించడం ద్వారా అదే ట్రీ ఆకారంలో SP యొక్క 1.82M మరియు IREE యొక్క 1.61M తో పోలిస్తే ప్రతి ఎన్కోడ్కు 1.10M సూచనలను రిటైర్ చేయడానికి మాకు వీలవుతుంది; మాకు అవసరం లేని జనరల్-పర్పస్ APIల కోసం అవి ఉంచే బుక్కీపింగ్ కారణంగా ఈ వ్యత్యాసం వస్తుంది.
ఆప్టిమైజేషన్ 2: బిట్మ్యాప్ మరియు ఇన్లైన్ ప్యాకింగ్
ఆప్టిమైజేషన్ 1 నుండి Darts ట్రీకి ప్రతి బైట్ దశకు రెండు డిపెండెంట్ లోడ్లు అవసరం. మొదట పేరెంట్ యొక్క బేస్ ఆఫ్సెట్, ఆపై పేరెంట్ దానిని కలిగి ఉందని నిర్ధారించడానికి అభ్యర్థి చైల్డ్ స్లాట్ వద్ద చెక్ ఎంట్రీ. చెక్ ఎంట్రీ కేవలం యాజమాన్యాన్ని ధృవీకరిస్తుంది. ఏ బైట్లకు పిల్లలు ఉన్నాయో రికార్డ్ చేసే ప్రతి-నోడ్ బిట్మ్యాప్ ద్వారా అదే ప్రశ్నకు సమాధానం ఇవ్వవచ్చు. డబుల్-అరే ట్రీ నిర్మాణం ద్వారా, పేరెంట్ బిట్మ్యాప్లోని బైట్ b కోసం బిట్ సరిగ్గా ఆ బైట్ కోసం అభ్యర్థి స్లాట్ పేరెంట్ చేత స్వంతం చేసుకోబడినప్పుడు సెట్ చేయబడుతుంది. బిట్మ్యాప్ మరియు చెక్ అరే ఒకే సమాచారాన్ని కలిగి ఉంటాయి, కాబట్టి మేము రన్-టైమ్ స్ట్రక్చర్ నుండి చెక్ అరేను పూర్తిగా తొలగించవచ్చు.
బిట్మ్యాప్ లుకప్ ఒకే 64-బిట్ పదానికి వ్యతిరేకంగా ఒకే బిట్ పరీక్షగా సంకలనం చేయబడుతుంది, ఇది దాదాపు ఒక సైకిల్లో నడుస్తుంది. సంబంధిత POPCNT సూచన, ఇది రిజిస్టర్లోని అన్ని సెట్ బిట్లను లెక్కిస్తుంది, బిట్మ్యాప్ నుండి నేరుగా చైల్డ్ ఆఫ్సెట్ను లెక్కించడానికి ఇతర బిట్మ్యాప్-ఇండెక్స్డ్ ట్రీ డిజైన్లు ఉపయోగించేది. Darts నిర్మాణం ఇప్పటికే ఆ ఆఫ్సెట్ను నోడ్ యొక్క base ఫీల్డ్లో నిల్వ చేస్తుంది కాబట్టి మనకు ఇది అవసరం లేదు.
#[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)
}అలాగే, నాలుగు per-node fields (bitmap, base, token ID, and score)ను ఒకే 64-byte cache lineలో pack చేస్తాము. ఇప్పుడు ఒకే node visitతో అన్నీ లోనికి వస్తాయి: terminal nodes వద్ద మనకు అవసరమైన score మరియు token ID, అలాగే తదుపరి step కోసం bitmap మరియు base. check arrayని కేవలం construction సమయంలోనే ఉపయోగించి runtime layout నుండి తొలగిస్తాము. Hot-path indexing uncheckedగానే ఉంటుంది, ఎందుకంటే trie construction మరియు loop structure boundsను ఇప్పటికే guarantee చేస్తాయి. Trade-off space: 64 bytes per node × ~780K slots వద్ద trie ~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% |
సెంటెన్స్పీస్ | 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% |
+ బిట్మ్యాప్ మరియు ఇన్లైన్ ప్యాకింగ్ | 65 µs | 77 µs | 1.06M | 1.8K | 20.53% | 15.01% |
bitmap మరియు inline packing జోడించిన తర్వాత 514 టోకెన్ల వద్ద టోకెనైజర్ ఫలితాలు.
Wall-clock 514 tokens వద్ద systematically మరో 4.5% తగ్గి, input lengthతో పాటు పెరుగుతుంది; 4k మరియు 16k వద్ద 6-7%కి చేరుతుంది. check arrayని తొలగించడం వల్ల per-byte trie step రెండు cache-line loads నుంచి ఒకదానికి తగ్గుతుంది, అంటే ప్రతి encodeకు L2 accesses 2.5× తక్కువగా ఉంటాయి (plain Darts యొక్క 4.6Kతో పోలిస్తే 1.8K). miss rate 6% నుండి 20%కి పెరుగుతుంది, కానీ అందులో ఎక్కువ భాగం denominator వల్లనే: accesses సగానికి కూడా తక్కువగా ఉన్నప్పుడు, మనకు ఉన్న misses పెద్ద వాటాను ఆక్రమిస్తాయి.
ఆప్టిమైజేషన్ 3: ట్రై కోసం హ్యూజ్ పేజీలు
ట్రీ డేటా దాదాపు 50 MB (ప్రతి ఒక్కటి 64 బైట్ల వద్ద 780,000 నోడ్లు). డిఫాల్ట్ లైనక్స్ ప్రాసెస్లో, ఆ మెమరీ 4 KB పేజీల ద్వారా బ్యాక్ చేయబడుతుంది, కాబట్టి పూర్తి ట్రీ దాదాపు 12,000 వేర్వేరు వర్చువల్ పేజీలను విస్తరిస్తుంది, CPU యొక్క ఫస్ట్-లవెల్ డేటా ట్రాన్స్లేషన్ లుకసైడ్ బఫర్ పట్టుకోగల దానికంటే చాలా ఎక్కువ (Sapphire Rapids 96 ఎంట్రీలను కలిగి ఉంది). ప్రతి Viterbi దశ వేర్వేరు నోడ్ను తాకుతుంది, మరియు 512-టోకెన్ ఎన్కోడ్ ద్వారా మనం పేజీ-టేబుల్ వాక్స్లో దాదాపు 9,000 సైకిళ్లను కాలిపోతాము, ఇది ప్రతి-ఎన్కోడ్ బడ్జెట్లో దాదాపు 3%.
పరిష్కారం trieను 4 KB pages బదులుగా 2 MB huge pagesతో back చేయడం. అదే 50 MB ఇప్పుడు 25 pagesను మాత్రమే ఆక్రమిస్తుంది, మొత్తం Translation Lookaside Bufferలో సరిపోతుంది, మరియు page-table directories తామే L2 cache కోసం trieతో పోటీ పడడం ఆగిపోతుంది.
Linuxకు 2 MB pagesతో memoryని back చేయడానికి రెండు మార్గాలు ఉన్నాయి. అవకాశవంతమైనది transparent huge pages (THP): మీ program ఒక memory regionను 4 KB నుంచి 2 MB pagesకు upgrade చేయమని kernelను కోరుతుంది, మరియు background kernel threadకు సమయం ఉన్నప్పుడు, అలాగే contiguous physical memory అందుబాటులో ఉన్నప్పుడు, అది చేయడానికి ప్రయత్నిస్తుంది. మరో alternative ఏమిటంటే, boot సమయంలో kernel పక్కన పెట్టే ముందే reserved huge-page pool నుండి trieను నేరుగా allocate చేయడం:
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% |
input sizes across huge pages తో మరియు లేకుండా టోకెనైజర్ latency.
అత్యపెద్ద విజయం 4k టోకెన్ల వద్ద ఉంది, ఇక్కడ వర్కింగ్ సెట్ L2 బ్యాండ్విడ్త్ కోసం ట్రీ డేటాతో పేజీ-టేబుల్ ట్రాఫిక్ చురుకుగా పోటీ పడేంత పెద్దదిగా ఉంటుంది. 4k దాటిన తర్వాత లాభం తగ్గిపోతుంది ఎందుకంటే L3 మిస్సులు ఆధిపత్యం చెలాయిస్తాయి మరియు పేజీ-టేబుల్ వాక్స్ మొత్తం సమయంలో చిన్న వాటాతో ఉంటాయి.
తుది పోలిక
మూడు ఆప్టిమైజేషన్లను కలిపి తీసుకుంటే, మేము పరీక్షించిన ప్రతి ఇతర RoBERTa టోకనైజర్కు వ్యతిరేకంగా కొలవబడిన తుది ఇంజిన్ ఇక్కడ ఉంది. స్థానిక C++ SentencePiece మరియు C IREE బైనరీలతో పాటు, ఉత్పత్తి వినియోగదారు చేరుకునే ఆఫ్ది-షెల్ఫ్ రస్ట్ రేప్పర్ క్రేట్లు, sentencepiece మరియు iree-tokenizerలను కూడా మేము బెంచ్మార్క్ చేసాము. sentencepiece క్రేట్ ప్రతి ఇన్పుట్ పరిమాణంలో స్థానిక C++ బైనరీ కంటే 1.6-1.9× ఎక్కువ లేటెన్సీని కొలుస్తుంది, రేప్పర్ ప్రతి ఎన్కోడ్లో టోకెన్ పీస్ల తాజా జాబితాను కేటాయించడం నుండి ఓవర్హెడ్ వస్తుంది, ప్రతి ఒక్కటి పీస్ టెక్స్ట్ యొక్క స్వంత కాపీని కలిగి ఉంటుంది. iree-tokenizer క్రేట్ దాని అంతర్గత బఫర్ను ప్రతి ఇన్పుట్కు సరిపోయేలా పరిమాణం చేస్తుంది, ఇది ఇన్పుట్ పరిమాణంతో సంబంధం లేకుండా స్థానిక IREE బైనరీ చెల్లించే స్థిరమైన ప్రతి-ఎన్కోడ్ కేటాయింపును నివారిస్తుంది, మరియు కేటాయింపు లేకపోతే వాస్తవ టోకనైజేషన్ పనిలో ఆధిపత్యం చెలాయించే చిన్న ఇన్పుట్ల వద్ద ఇది చాలా ముఖ్యం. ఉత్పత్తి ఇన్పుట్ పొడవుల వద్ద స్థిరమైన ఖర్చు అమార్టైజ్ అవుతుంది మరియు రస్ట్లోకి టోకెన్లను వెనక్కి కాపీ చేయడం నుండి రేప్పర్ యొక్క స్వంత ప్రతి-కాల్ ఓవర్హెడ్ దాని లేటెన్సీని స్థానిక బైనరీ కంటే ఎక్కువగా నెడుతుంది.
తుది ఇంజిన్ అదే 250,000-టోకెన్ వొకాబ్యులరీపై అదే Viterbiని నడుపుతుంది, హాట్ పాత్లో ఏమీ కేటాయించదు, మరియు రిఫరెన్స్కు వ్యతిరేకంగా టోకెన్-ఖచ్చితమైన అవుట్పుట్ను ఉత్పత్తి చేస్తుంది. ప్రతి బైట్ దశ ఇప్పుడు ఒకే 64-బైట్ కాష్-లైన్ లోడ్, ఒక బిట్మ్యాప్ పరీక్ష, మరియు ఒక పూర్ణాంక జోడింపు, బేస్లైన్ అమలులోని HashMap ప్రోబ్, అలోకేటర్ కాల్, మరియు అనేక ప్రత్యేక-కేసు బ్రాంచ్ల నుండి తగ్గింది. ప్రతి ఎన్కోడ్కు సూచనలు 3.66M నుండి 1.04Mకి పడిపోయాయి, ఇది 3.5× తగ్గుదల. చిన్న ఇన్పుట్ల వద్ద వాల్-క్లాక్ ఆ నిష్పత్తికి సరిపోతుంది మరియు పొడవైన ఇన్పుట్ల వద్ద మరింత విస్తరిస్తుంది, ఇక్కడ రిఫరెన్స్ యొక్క ప్రతి-టోకెన్ కేటాయింపులు L2 మరియు L3ని ఓవర్ఫ్లో చేయడం ప్రారంభిస్తాయి.
ముగింపు ఆలోచనలు
ఆచరణలో ఈ పని మా ఇన్ఫరెన్స్ స్టాక్ యొక్క ప్రీప్రాసెసింగ్ పొరలో CPU వినియోగాన్ని 5-6× తగ్గించింది మరియు రీరాంకర్ లేటెన్సీ నుండి డబుల్-అంకె మిలిసెకండ్లను తగ్గించింది.
లేటెన్సీ తగ్గింపులో ఎక్కువ భాగం వృధా అయిన పనిని తొలగించిన తేలికపాటి మార్పుల నుండి వచ్చింది: ప్రతి-ఎన్కోడ్ కేటాయింపులను తొలగించడం, HashMap-children ట్రీని డబుల్-అరేతో భర్తీ చేయడం, మరియు జనరిక్-API ఓవర్హెడ్ను స్ట్రిప్ చేయడం. డబుల్-అరే ట్రీకి పెద్ద మార్పుతో పాటు, చెల్లుబాటు తనిఖీగా ఇప్పటికే ఉన్న ప్రతి-నోడ్ బిట్మ్యాప్ను ఉపయోగించడం మరియు 2 MB హ్యూజ్ పేజీలతో ట్రీని బ్యాక్ చేయడం ఇన్పుట్ పరిమాణాన్ని బట్టి 10-20% తోడ్పడ్డాయి.
SentencePiece మరియు IREE చుట్టూ ఉన్న ఆఫ్ది-షెల్ఫ్ రస్ట్ రేప్పర్ క్రేట్లు వ్యతిరేక దిశలో అదే విషయాన్ని వివరిస్తాయి. స్థానిక C మరియు C++ టోకనైజర్లపై, FFI బౌండరీ మరియు రేప్పర్లు పరిచయం చేసే ప్రతి-కాల్ కేటాయింపులు లేటెన్సీ సంఖ్యలలో కనిపిస్తాయి. తుది Viterbi ఎన్కోడర్ కోసం రస్ట్ సోర్స్ 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
ఉద్ధృతి
@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}
}