Make BPE Pre-tokenizer Thread-Safe by Using Per-call Cursors

[Sayan Shaw (sayanshaw24)]

Make BPE pre-tokenizer thread-safe by using per-call cursors

Problem

PR #1068 optimized BPE tokenization by caching compiled PreTokenizerWithRegEx instances as member variables (cached_splitters_). This avoided re-compiling regex patterns on every Tokenize() call.

However, PreTokenizerWithRegEx has mutable cursor state (m_text, m_last_char, m_utf8_text) that gets overwritten on every Set() + GetNextToken() call. When multiple threads call Tokenize() on the same tokenizer instance concurrently, they race on this shared cursor, causing corrupted token output.

This was observed as flaky test failures in downstream consumers (microsoft/foundry-local#1008).

Root Cause

Mutable state	Mutated by	Thread-unsafe?
reg_splitter.m_text	Set(), GetNextToken()	Yes — cursor position shared
reg_splitter.m_last_char	TryMatch()	Yes — match state shared
reg_splitter.m_utf8_text	Set()	Yes — UTF-8 buffer shared
cached_splitters_ pointer	Lazy const_cast init	Yes — TOCTOU race

Fix

operators/tokenizer/bpe_utils.hpp

• Add CreateCursor() method that returns a lightweight PreTokenizerWithRegEx sharing compiled state but with independent cursor state.
• Change activated_matchers_ to shared_ptr<vector<RegexMatchFunc>> — avoids heap-allocating a copy of the matchers vector on every CreateCursor() call.
• Change fallback_patterns_ to shared_ptr<const std::regex> — enforces immutability of compiled patterns and enables sharing across cursors.

operators/tokenizer/bpe_kernels.cc

• In Tokenize() and SpmTokenize(), replace direct use of cached_splitters_->reg_splitter with auto splitter = cached_splitters_->reg_splitter.CreateCursor() — each call gets its own cursor on the stack.
• Same for the sequence pre-tokenizer path: each seq_splitter gets a local cursor.
• Replace const_cast lazy init with std::call_once to eliminate the TOCTOU race on cached_splitters_ initialization.
• Guard CompilePreTokenizer() with once_flag to prevent double-compile when called eagerly at load + lazily via call_once.

operators/tokenizer/bpe_kernels.h

• Add std::once_flag compile_pretokenizer_flag_ member for std::call_once.

Performance Impact

None. The compiled regex patterns and matcher function vectors (the expensive parts) are shared via shared_ptr — no heap allocations in CreateCursor(). Each cursor is ~40 bytes of stack state (a string_view, a char32_t, an empty string, and two shared_ptr copies).

Verified with test_performance.py — GPT-2 long_english: 16.8 MB/s (vs 15.6 MB/s before fix). No regression.

Testing

All existing tokenizer tests pass.

[Copilot]

Pull request overview

This PR addresses a concurrency bug in the BPE tokenizer pre-tokenization path by ensuring each Tokenize()/SpmTokenize() call uses its own PreTokenizerWithRegEx cursor state while still reusing cached compiled patterns.

Changes:

• Introduces PreTokenizerWithRegEx::CreateCursor() to create per-call cursor instances while sharing compiled regex/matchers.
• Switches pre-tokenizer compilation initialization to std::call_once to avoid racy lazy initialization of cached_splitters_.
• Updates the BPE kernel tokenization paths to use per-call cursor instances (including sequence pre-tokenizers).

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
operators/tokenizer/bpe_utils.hpp	Adds CreateCursor() and switches fallback regex ownership to shared_ptr for sharing across cursors.
operators/tokenizer/bpe_kernels.h	Adds <mutex> and a std::once_flag to support thread-safe pre-tokenizer compilation.
operators/tokenizer/bpe_kernels.cc	Uses std::call_once and per-call cursors to make regex pre-tokenization thread-safe.

Suppressed comments (1)

operators/tokenizer/bpe_kernels.cc:656

• Same issue as in Tokenize(): std::call_once will still run the lambda once even if cached_splitters_ was already compiled during model load, causing a redundant CompilePreTokenizer() on the first SpmTokenize() call. Guard the body so the once-call becomes a no-op when cached_splitters_ is already initialized.

  std::call_once(compile_pretokenizer_flag_, [this]() {
    const_cast<KernelBpeTokenizer*>(this)->CompilePreTokenizer();
  });

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.