Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

GPUTOK: GPU Accelerated Byte Level BPE Tokenization

Venu Gopal Kadamba, Kanishkha Jaisankar

TL;DR

A GPU-based byte-level BPE tokenizer that follows GPT-2's merge rules that includes a basic BlockBPE-style kernel and a faster, optimized version that uses cuCollections static map, CUB reductions, and a pybind11 interface for Python.

Abstract

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.

GPUTOK: GPU Accelerated Byte Level BPE Tokenization

TL;DR

A GPU-based byte-level BPE tokenizer that follows GPT-2's merge rules that includes a basic BlockBPE-style kernel and a faster, optimized version that uses cuCollections static map, CUB reductions, and a pybind11 interface for Python.

Abstract

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.
Paper Structure (29 sections, 1 equation, 4 figures, 2 tables)

This paper contains 29 sections, 1 equation, 4 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Our approach
  3. Literature Review
  4. Algorithmic Variants of BPE.
  5. Systems and Implementation Strategies.
  6. GPU Hash Tables and Parallel Primitives.
  7. Pros and Cons of Existing Approaches.
  8. Why Existing Work Is Not Enough.
  9. Positioning relative to BlockBPE and prior GPU tokenizers
  10. Proposed Approach
  11. System architecture
  12. GPT2 byte encoder reconstruction
  13. Merge table representation
  14. Baseline kernel
  15. Optimized kernel
  16. ...and 14 more sections

Figures (4)

  • Figure 1: CUDA API breakdown from Nsight Systems on the long-document workload. Memory allocations dominate API time, indicating that device-side memory pooling could substantially reduce overhead.
  • Figure 2: Tokenization latency versus sequence length on WikiText103 in latency mode. The optimized GPU kernel overtakes tiktoken at around two thousand tokens and reaches a clear advantage at the longest context lengths.
  • Figure 3: Speedup of the optimized GPU tokenizer relative to CPU-based tokenizers. Values above one indicate that the GPU implementation is faster. The crossover point where the GPU becomes advantageous is near two thousand tokens.
  • Figure 4: Relative throughput comparison at the longest sequence length. The GPU tokenizer surpasses both tiktoken and HuggingFace GPT2.