본문으로 건너뛰기

GPT-2 토큰화를 Rust로 가속한 방법

GPT-2 tokenizer를 Rust와 직접 vocabulary lookup으로 최적화해 Moby-Dick 32 KiB 처리 시간을 7.01 ms에서 3.09 ms로 줄였습니다.

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

LLM은 원문이 아니라 token ID sequence를 입력받으므로 tokenization이 요청 처리 경로에 포함되며, context 압축 시점 판단이나 비용 추정, prompt routing에도 반복 사용될 수 있습니다. 글은 GPT-2의 Python reference encoder를 Rust로 옮긴 뒤, 인접 pair 전체를 반복 검색하는 baseline, priority queue와 linked list를 결합한 방식, 원본 byte 범위와 decoded vocabulary를 직접 조회하는 V3를 비교합니다. priority queue 방식은 작은 정규식 match에서 heap과 간접 참조 비용 때문에 32 KiB Moby-Dick 기준 9.25 ms로 baseline의 7.01 ms보다 느렸고, V3는 vocabulary-first lookup으로 3.09 ms를 기록했습니다. V3에서는 BPE merging 비중이 Moby-Dick에서 51%에서 9%로 줄었지만 regex matching이 55%로 가장 큰 단계가 되었으며, 작은 작업을 대량 수행하는 tokenization에서는 복잡한 전역 자료구조보다 연속 메모리 순회와 직접 조회가 유리하다는 결론으로 이어집니다.

섹션별 상세

01
GPT-2는 임의의 UTF-8을 처리하기 위해 256개 byte를 기본 alphabet으로 사용하고, 먼저 정규식으로 입력을 독립적인 영역으로 나눈 뒤 각 영역에 BPE를 적용합니다. BPE는 학습된 vocab.bpe의 merge rank를 기준으로 유효한 인접 pair 중 가장 낮은 rank를 반복 선택하므로, goldshire는 8번의 merge를 거쳐 Ġgold와 shire 두 token ID가 됩니다. 이 구조는 단순한 왼쪽에서 오른쪽 스캔이 아니라 merge가 새 후보를 만들 때마다 전체 인접 관계를 다시 고려해야 하는 이유를 설명합니다.
02
Rust baseline은 GPT-2의 regex, byte map, merge ranks, encoder.json을 같은 형태로 옮기고 각 정규식 match마다 인접 pair 전체를 반복 순회합니다. Python encoder.py가 Moby-Dick 32 KiB에서 42.43 ms, 0.7 MiB/s를 기록한 반면 Rust baseline은 7.01 ms, 6.0 MiB/s를 기록해 구현 언어와 실행 경로만으로도 차이가 났습니다. 다만 goldshire처럼 최종 token 수보다 많은 merge round가 필요한 입력에서는 이미 결합된 영역을 반복 방문하는 비용이 남아 있었습니다.
rust
pub fn encode(input: &str) -> Vec { tokenizer().encode(input)} fn encode(&self, input: &str) -> Vec { self.pattern.find_iter(input) // Each match is one independent BPE region. .flat_map(|item| self.merge(item.unwrap().as_str())) // Final BPE pieces map to GPT-2 token IDs. .map(|piece| self.token_ids[&self.symbols[piece]]) .collect()}

정규식 match를 독립적인 BPE 영역으로 처리한 뒤 최종 조각을 GPT-2 token ID로 변환해 Vec에 모읍니다.

rust
loop { let best = pieces.windows(2) // The earliest vocab.bpe line has the best rank. .filter_map(|pair| { self.merge_ranks .get(&(pair[0], pair[1])) .map(|&(rank, result)| { (pair[0], pair[1], rank, result) }) }) .min_by_key(|pair| pair.2); let Some((first, second, _, result)) = best else { return pieces; }; // Apply this merge wherever the pair occurs. pieces = merge_all(pieces, first, second, result);}

인접 pair 전체에서 가장 낮은 merge rank를 찾고, 해당 pair가 나타나는 위치를 결합하는 baseline merge loop입니다.

03
Version 2는 후보 pair를 priority queue에 넣고 linked-list sequence로 이웃을 유지해 merge 뒤 새로 영향을 받는 pair만 갱신하도록 구성했습니다. 큐에서 꺼낸 항목은 앞선 merge 때문에 낡았을 수 있어 다시 검사해야 하며, 이 과정에서 pair map, stale entry, node indirection, 정렬과 분기 비용이 추가됩니다. 실제 정규식 match 평균 길이는 Moby-Dick 4.29 byte, React 4.50 byte였고, heap 방식은 각각 9.25 ms와 5.4 MiB/s로 baseline의 7.01 ms와 6.0 MiB/s보다 불리했습니다.
rust
let bytes = piece.as_bytes(); // A learned token needs no BPE work at all.if let Some(&token) = ranks.get(bytes) { return vec![token];}

완전한 정규식 match가 이미 vocabulary에 있는지 먼저 확인해 BPE merge를 건너뛰고 token ID를 바로 반환합니다.

rust
struct Part { start: usize, rank: u32, // Candidate merged-token ID, or u32::MAX.} // Allocate once. The input bytes never change.let mut parts = initial_parts(bytes); while let Some((index, _)) = lowest_rank(&parts) { // Recheck the two candidates that will touch after removal. update_rank(bytes, &mut parts, index.saturating_sub(1)); update_rank(bytes, &mut parts, index); // Delete one boundary. Vec shifts entries but does not reallocate. parts.remove(index + 1);}

원본 byte를 복사하지 않고 살아남은 byte 경계만 Vec로 관리하면서 가장 낮은 rank의 경계를 제거해 merge를 반복합니다.

04
V3는 encoder.json 항목을 초기화 시점에 raw bytes로 decode하고, 원본 입력 bytes 안의 연속 범위를 그대로 vocabulary map에 조회합니다. 각 Part는 byte 시작 위치와 merge rank를 보유하며, merge가 일어나면 문자열이나 입력 자체를 복사하지 않고 경계 하나를 제거한 뒤 왼쪽과 오른쪽 후보만 갱신합니다. 이 방식은 Moby-Dick에서 3.09 ms, 9.4 MiB/s를 기록했으며, cache-first 변형은 3.23 ms, 8.8 MiB/s, tiktoken-rs r50k_base는 2.15 ms, 11.7 MiB/s를 기록했습니다.
05
V3에서 BPE merging 비중은 Moby-Dick 기준 baseline의 51%에서 9%로, React 기준 40%에서 8%로 낮아졌습니다. 그 결과 regex matching 비중은 두 입력 모두 55%가 되었고, lookup/output도 각각 17%와 20%로 커졌으며 byte mapping과 cache 비용도 상대적으로 드러났습니다. tokenization의 병목은 한 단계의 복잡성보다 아주 작은 텍스트 조각에 대해 regex, lookup, merge를 매우 많이 수행하는 작업량에 있다는 해석으로 이어집니다.
06
토큰 수만 대략 추정하면 되는 기능에는 실제 token ID를 계산하지 않고 input bytes를 N으로 나누는 고정 추정치가 더 저렴할 수 있습니다. 반면 입력을 여러 chunk로 받아 처리할 때는 chunk별 token count나 token sequence를 단순히 합칠 수 없으며, 뒤의 bytes가 앞 chunk의 regex 경계나 BPE merge를 바꿀 수 있습니다. 글은 tiktoken의 encode_with_unstable을 stable token과 가능한 completion sequence를 구분하는 관련 접근으로 언급하고, streaming tokenization에서는 안전하게 확정할 수 있는 tail을 판단해야 한다고 정리합니다.

용어 해설

Byte-Pair Encoding
Byte-Pair Encoding은 입력을 기본 byte 단위로 쪼갠 뒤 학습된 merge rank에 따라 자주 함께 나타나는 인접 조각을 반복 결합하는 토큰화 방식입니다. GPT-2에서는 UTF-8을 가역적으로 처리하고 최종 byte sequence를 token ID로 변환하는 데 사용됩니다.
사전 토큰화(Pre-tokenization)
사전 토큰화는 BPE를 적용하기 전에 정규식으로 입력을 단어 형태, 숫자, 축약형, 구두점, 공백 같은 독립 영역으로 나누는 단계입니다. 영역 경계를 두어 서로 다른 텍스트 유형이 하나의 토큰으로 합쳐지는 것을 막습니다.
Merge 순위(Merge Rank)
Merge 순위는 BPE가 인접한 두 조각을 결합할 때 적용하는 우선순위입니다. GPT-2는 유효한 인접 pair 가운데 순위 숫자가 가장 낮은 pair를 선택하며, 결합 결과가 새로운 후보 pair를 만들면 다음 반복에서 다시 평가합니다.
LRU 캐시(LRU Cache)
LRU 캐시는 최근에 사용한 항목을 제한된 용량으로 보관하고 새 항목이 들어올 때 오래 사용되지 않은 항목부터 제거하는 방식입니다. 이 글의 Rust tokenizer는 정규식 match를 BPE 결과로 바꾼 중간 결과를 최대 256개까지 캐시합니다.
우선순위 큐(Priority Queue)
우선순위 큐는 후보 pair를 우선순위 순서로 꺼내는 자료구조입니다. 글에서는 BPE merge 후보를 GPT-2의 rank 순서로 처리하고, merge 뒤 새로 인접해진 이웃만 갱신하는 데 사용했지만 작은 match에서는 자료구조 비용 때문에 baseline보다 느렸습니다.
FxHashMap
FxHashMap은 Rust의 범용 hash map 대신 사용할 수 있는 hash map 구현으로, 이 글의 V3 tokenizer가 byte sequence와 token ID 또는 merge rank을 빠르게 연결하는 데 사용합니다. 직접적인 vocabulary lookup과 함께 전체 merge 작업과 메모리 할당을 줄이는 역할을 합니다.

기술

  • GPT-2
  • BPE
  • Rust
  • regex
  • encoder.json
  • vocab.bpe
  • tiktoken
  • tiktoken-rs
  • FxHashMap
  • Rayon

활용 사례

  • LLM 요청 비용 추정
  • context 압축 시점 판단
  • prompt routing
  • 증분 입력 처리
  • streaming tokenization
AI 분석 전체 내용 보기

AI 요약 · 북마크 · 개인 피드 설정 — 무료

출처 · 인용 안내

원문 발행 2026. 09. 02.수집 2026. 09. 02.출처 타입 RSS

인용 시 "요약 출처: AI Trends (aitrends.kr)"를 표기하고, 사실 확인은 원문 보기 기준으로 진행해 주세요. 자세한 기준은 운영 정책을 참고해 주세요.