GPUTOK：基于 GPU 的字节级 BPE 分词
计算与语言
2026-03-04 v1 人工智能
分布式、并行与集群计算
机器学习
摘要
随着大型语言模型向百万级上下文窗口发展，CPU 分词器成为主要瓶颈，因为它们逐步处理文本而强大的 GPU 闲置。我们构建了一个遵循 GPT-2 合并规则的 GPU 字节级 BPE 分词器。它包括基本的 BlockBPE 风格内核和一个更快、优化的版本，利用 cuCollections 静态映射、CUB 归约和 pybind11 的 Python 接口。在 WikiText103 序列最长达 131k 标记时，优化后的 GPU 分词器产生的标记与 CPU 版本相同，对最长输入，速度约为 tiktoken 的 1.7 倍，约为 HuggingFace GPT-2 分词器的 7.6 倍。Nsight 分析显示，70-80% 的 CUDA API 时间用于内存分配，因此添加内存池应提供最大的速度提升。针对使用 WikiText103 提示的生成任务的测试表明，我们的 GPU 分词器的输出在相似性和重叠指标上与 tiktoken 和 HuggingFace GPT-2 基本一致（相差约一percentage point），这意味着它在保持输出质量的同时，使长上下文推理更加实用。
引用
@article{arxiv.2603.02597,
title = {GPUTOK: GPU Accelerated Byte Level BPE Tokenization},
author = {Venu Gopal Kadamba and Kanishkha Jaisankar},
journal= {arXiv preprint arXiv:2603.02597},
year = {2026}
}