GPU加速字节级BPE分词，长文本处理快7.6倍
GPUTOK: GPU Accelerated Byte Level BPE Tokenization
- 用GPU并行执行GPT-2字节级BPE分词，避开CPU串行瓶颈
- 在131k token输入下比HuggingFace版本快7.6倍，与CPU结果一致
- 适合长上下文推理场景，保持分词质量同时显著提速
随着大语言模型上下文窗口扩展至百万级别token，传统基于CPU的分词器成为性能瓶颈，因其实现方式为逐步处理而无法利用强大GPU。本文构建了一种遵循GPT-2合并规则的基于GPU的字节级BPE分词器，包含基础BlockBPE内核及优化版本，后者采用cuCollections静态映射、CUB归约和pybind11接口支持Python调用。在最长达131,000词元的WikiText103序列上，优化版GPU分词器输出与CPU版本完全一致，最长输入下相比tiktoken快1.7倍，比HuggingFace GPT-2分词器快7.6倍。Nsight分析显示70-80%的CUDA API时间消耗于内存分配，引入内存池有望带来最大性能提升。生成任务测试表明，该分词器在相似性与重叠度指标上与tiktoken和HuggingFace GPT-2相差不足1个百分点，保证了输出质量的同时使长上下文推理更高效可行。
原文摘要 · Abstract (English)
As large language models move toward million-token context windows, CPU tokenizers become a major slowdown because they process text one step at a time while powerful GPUs sit unused. We built a GPU-based byte-level BPE tokenizer that follows GPT-2's merge rules. It includes a basic BlockBPE-style kernel and a faster, optimized version that uses cuCollections static map, CUB reductions, and a pybind11 interface for Python. On WikiText103 sequences up to 131k tokens, the optimized GPU tokenizer produces the same tokens as a CPU version and, for the longest inputs, is about 1.7x faster than tiktoken and about 7.6x faster than the HuggingFace GPT-2 tokenizer. Nsight profiling shows that 70-80% of CUDA API time goes to memory allocation, so adding memory pooling should give the biggest speed boost next. Tests on generation tasks using WikiText103 prompts show that our GPU tokenizer's outputs stay within about one percentage point of tiktoken and HuggingFace GPT-2 on similarity and overlap metrics, meaning it keeps output quality while making long-context inference more practical.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。