GPUTOK: GPU Accelerated Byte Level BPE Tokenization

📄 arXiv: 2603.02597v1 📥 PDF

作者: Venu Gopal Kadamba, Kanishkha Jaisankar

分类: cs.CL, cs.AI, cs.DC, cs.LG

发布日期: 2026-03-03

💡 一句话要点

GPUTOK：利用GPU加速字节级BPE分词，提升长文本处理效率

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: GPU加速 BPE分词 大语言模型 长文本处理 CUDA cuCollections CUB规约

📋 核心要点

  1. 现有CPU分词器在处理百万token上下文的大语言模型时成为瓶颈，无法充分利用GPU算力。
  2. 提出一种基于GPU的字节级BPE分词器，优化了内存访问和计算，加速分词过程。
  3. 实验表明，该GPU分词器在长文本分词速度上显著优于tiktoken和HuggingFace GPT-2分词器，且保持了输出质量。

📝 摘要（中文）

本文提出了一种基于GPU的字节级BPE分词器，遵循GPT-2的合并规则。该分词器包含一个基础的BlockBPE风格的内核，以及一个更快速的优化版本，后者利用了cuCollections静态映射、CUB规约以及用于Python的pybind11接口。在WikiText103序列上，对于最长可达131k tokens的输入，优化的GPU分词器能够产生与CPU版本相同的tokens。实验结果表明，对于最长的输入，该GPU分词器比tiktoken快约1.7倍，比HuggingFace GPT-2分词器快约7.6倍。Nsight性能分析显示，70-80%的CUDA API时间花费在内存分配上，因此增加内存池将带来最大的速度提升。在WikiText103提示的生成任务测试中，该GPU分词器的输出在相似性和重叠度指标上与tiktoken和HuggingFace GPT-2的差距在一个百分点以内，这意味着它在保持输出质量的同时，使长上下文推理更具实用性。

🔬 方法详解

问题定义：现有的大语言模型正朝着百万token级别的上下文窗口发展，这使得CPU分词器成为主要的性能瓶颈。传统的CPU分词器以串行方式处理文本，无法充分利用GPU的并行计算能力，导致整体处理速度缓慢。因此，如何利用GPU加速分词过程，成为提升大语言模型性能的关键问题。

核心思路：本文的核心思路是将字节级BPE分词过程迁移到GPU上进行加速。通过设计高效的GPU内核，充分利用GPU的并行计算能力，从而显著提升分词速度。此外，论文还通过优化内存访问和数据结构，进一步提升了GPU分词器的性能。

技术框架：该GPU分词器主要包含以下几个模块：1) 基础的BlockBPE风格的GPU内核；2) 优化的GPU内核，利用cuCollections静态映射和CUB规约；3) pybind11接口，用于Python集成。整体流程为：首先将文本数据加载到GPU内存中，然后使用GPU内核进行BPE分词，最后将分词结果返回给CPU。

关键创新：该论文的关键创新在于设计并实现了一个高效的GPU字节级BPE分词器。与传统的CPU分词器相比，该GPU分词器能够充分利用GPU的并行计算能力，从而显著提升分词速度。此外，论文还通过优化内存访问和数据结构，进一步提升了GPU分词器的性能。

关键设计：优化的GPU内核使用了cuCollections静态映射来加速词汇查找，并使用CUB规约来高效地进行数据聚合。此外，论文还使用了pybind11接口，使得该GPU分词器可以方便地与Python集成，从而方便用户使用。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在WikiText103序列上，对于最长可达131k tokens的输入，优化的GPU分词器能够产生与CPU版本相同的tokens。对于最长的输入，该GPU分词器比tiktoken快约1.7倍，比HuggingFace GPT-2分词器快约7.6倍。在WikiText103提示的生成任务测试中，该GPU分词器的输出在相似性和重叠度指标上与tiktoken和HuggingFace GPT-2的差距在一个百分点以内。

🎯 应用场景

该研究成果可广泛应用于需要处理长文本的大语言模型应用中，例如长文档摘要、机器翻译、对话系统等。通过加速分词过程，可以显著提升这些应用的性能和用户体验，使得长上下文推理更具实用性。未来，该技术有望进一步推广到其他自然语言处理任务中。

📄 摘要（原文）

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.