Speed up cold long-context BPE tokenization - #49
Closed
michaelfeil wants to merge 51 commits into
Closed
Conversation
[codex] release GIL during Python encode
release: Bump version to 0.2.0
…al-options [codex] reject unsupported special token options
[codex] support split special token encoding
Rename fork crate package to fasttokens-b10
This reverts commit f90d62b.
[codex] Revert structural token encoding API
* Add structural token encoding API * Return raw structural token encoding * Address structural token review comments * Preserve merges across structural placeholders * Test restored bare added structural tokens * Preserve merges for restored tag literals * Clarify structural matcher errors * Restore placeholder boundary parity
* Add separate Python CI workflow * Harden Python CI workflow permissions * Avoid caching Hugging Face token * Use dependency-based tokenizer cache key * Restrict Python support to 3.12 * Keep Python support metadata unchanged --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
* Add tiktoken conversion helper * Avoid network in tiktoken conversion tests * Address tiktoken conversion review feedback * Refine tiktoken conversion readability * Make tiktoken conversion optional * Clean up tiktoken conversion test import * Address tiktoken conversion review comments * Add tiktoken model conversion helper * Tighten tiktoken model conversion tests * Improve tiktoken model error handling * Clarify tiktoken model conversion logic * Polish tiktoken model conversion helpers * Fix tiktoken conversion test whitespace * Refine tiktoken conversion naming * Clarify tiktoken byte mapping constants * Run Python tests in CI with tiktoken * Vendor tokenizer assets * Add vendored Kimi K2.5 tiktoken gz conversion test * Fix docstring spacing in kimi conversion test * Add Kimi K2.5 structural tokenizer parity assertions * Refactor Kimi parity test helpers and finalize assertions --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Michael Feil <63565275+michaelfeil@users.noreply.github.com>
* Add numpy ID accessors * Clear encoding metadata after numpy move * Use generic encoding into_numpy API * Address numpy encoding review comments
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
* Add segmented encode support for tiktoken conversion * Clarify tiktoken-safe segment encoding * Default segmented encode to tiktoken parity * Bump segment encoding release to 0.2.5 * Delete vendored_tokenizers/.gitignore
* Add basic chat template rendering * Expose basic chat template API * Document chat template helpers * Improve native chat template rendering * Improve chat template parity and rendering performance * Make chat template rendering fail closed for tokenization * Tighten chat template render parity * Render chat templates from Python objects in a single hop Convert messages directly to minijinja values via depythonize instead of building an intermediate serde_json tree, keeping the JSON path only for continue_final_message mutation. Cuts the context-construction floor for a 2.1k-message conversation from ~5.5ms to ~3.4ms and the full Kimi 2.5 render from ~14.3ms to ~10.9ms with byte-for-byte transformers parity. Also reject kwargs that shadow reserved context names, rewrite generation markers with whitespace-control preserved, and expose set_special_tokens for persistent special-token configuration. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Reject chat template kwargs the template never reads Compute the template's undeclared variables once at compile time and validate apply_chat_template kwargs against them, so typos like enable_thinkng raise TypeError with a did-you-mean suggestion instead of silently rendering the wrong prompt. Opt out per call with fastokens_strict_template=False. Expose the variable set to Python as chat_template_variables for frontend-side validation. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Fix CI: rustfmt and missing jinja2 test dependency transformers' apply_chat_template needs jinja2 at call time, so the parity tests failed with ImportError despite transformers importing fine. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Address review: shim special-token persistence and indented tojson separators Persist set_special_tokens through shim copy, deepcopy, pickle, and shim-from-shim construction (pickle state grows to a 6-tuple, older 4/5-tuples still load). Honor tojson separators when indent is set by replacing serde_json's PrettyFormatter with an indent-aware HfJsonFormatter that matches json.dumps byte-for-byte, including verbatim separator whitespace before newlines and compact empty containers. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Gate native chat template rendering behind opt-in for patched transformers patch_transformers gains apply_chat_template=False: when enabled, render- only tokenize=False calls route to the fastokens renderer with the tokenizer's resolved template and special_tokens_map; positional args, tokenize=True, dict outputs, batched conversations, and any native error fall back to transformers' jinja2 (warned once). Strict kwarg validation stays off on this path to preserve transformers' ignore-unknown-kwargs contract. FASTOKENS_APPLY_CHAT_TEMPLATE overrides the argument in either direction as a fleet-wide kill switch, and patching logs which engine is active. Direct Tokenizer/shim calls are unaffected and always native. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Install jinja2 in the patch-transformers CI job The new opt-in routing tests exercise transformers' own renderer (as parity reference and via the tokenize=True fallthrough), which needs jinja2 at call time. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Add async_apply_chat_template Split apply_chat_template into a GIL-held preparation step returning a PreparedChatTemplateCall (owns no Python references) and a CPU-bound render, shared verbatim by both bindings: sync runs the render under allow_threads, async on the Tokio blocking pool via future_into_py. Awaiting keeps large renders off the event loop (max observed stall 1.6ms during a 10ms / 800k-char render vs blocking the full duration); input conversion still runs at call time since it reads live Python objects. Same error semantics, shim forwarder and stubs included. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Run async methods on pyo3-async-runtimes' managed runtime Drop the crate's own global Tokio runtime: async_apply_chat_template, async_encode_batch, and async_decode_batch now spawn_blocking on pyo3_async_runtimes::tokio::get_runtime(), the runtime future_into_py already schedules on. One thread pool instead of two nested ones, and the direct tokio dependency goes away. Investigated the reported hang: correct single-loop usage resolves in every probe (including the reporter's exact script and wheel); the forever-pending symptom reproduces only when an awaitable created in one event loop is awaited in another, where the completion callback targets the closed original loop. That misuse also aborts interpreter shutdown inside pyo3-async-runtimes' callback thread, before and after this change. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Reject return_dict without chat-template tokenization --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
* Add direct numpy encode materialization * Lazily materialize encoding masks * Add draining encoding to_numpy * Keep encoding numpy API focused * Keep numpy materialization on Encoding * Materialize only requested numpy fields
…Kimi (#25) * Add specialized split scanner tier (Qwen2/Qwen3 pattern) Recognize known pre-tokenization split regexes byte-for-byte at Split construction and scan them with a hand-written walker — first-byte dispatch, SWAR ASCII-letter runs, a packed 2-bit-per-codepoint Unicode class table — instead of the regex engines. Adapted from gigatoken (https://github.com/marcelroed/gigatoken, MIT). Tier order becomes scanner -> PCRE2 JIT -> fancy-regex; the scanner activates only on an exact recognized pattern with Isolated behavior and invert=false, so all other patterns are untouched. Large inputs are scanned across threads at token-safe chunk boundaries. On the serving-relevant workload — long-context, per-request, distinct inputs (cold split cache) — Qwen3 full encode is ~1.8-2.15x faster across 10k-200k char prompts, single- or multi-threaded. (On a single re-encoded string the PCRE2 incremental cache ties it; that regime does not represent independent requests.) Parity is enforced by differential tests vs fancy-regex on the exact pattern: 60+ edge cases + a 4000-round fuzz + an integration test covering the parallel chunked path. The GLM/dolma2 variant (Qwen2N3) is included and parity-tested but gated off pending an unexplained Split-integration slowdown; see docs/perf/specialized-split-scanners.md. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Fix formatting * Add Kimi (o200k-family) specialized split scanner Port gigatoken's scalar o200k-family walker (MIT) into fast_split_o200k.rs and wire the Kimi (moonshotai K2) scheme into the specialized-scanner tier. The Kimi pattern uses [\p{Han}]+ runs and &&[^\p{Han}] class intersections that PCRE2 cannot parse, so the PCRE2 tier runs a hand-rewritten pattern with a per-codepoint (?!\p{Han}) lookahead on every letter char — slow on both ASCII and CJK. The scanner classifies via one packed-table load (o200k 7-class + Han-split 10-class tables, built from ICU GeneralCategory + Script=Han). Full encode() on the vendored Kimi tokenizer, distinct inputs, A/B vs the PCRE2 tier: ~2.2-2.8x faster across 10k-50k char ASCII and CJK-mixed prompts. The walker is parameterized (CONTRACTIONS/DIGITS3/SLASH/HAN), so o200k (gpt-oss) and Nemotron are a pattern-const + dispatch line away as follow-ups. Parity: 30+ hand-picked edge cases (camelCase phase automaton, contractions, Han runs/numerals/symbols, combining marks, mixed CJK) + a 6000-round fuzz vs fancy-regex on the exact pattern, plus an integration test through Split::pre_tokenize including the parallel chunked path. 214 tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Enable Qwen2N3 (GLM/dolma2/OLMo) split scanner The scanner and parity tests were already present but dispatch was gated off after a microbenchmark showed the GLM (\p{N}{1,3}) scanner ~5x slower than the Qwen2 scanner. That measurement was a first-touch / page-warming artifact of re-encoding one fixed 4 MiB buffer in a tight loop: swapping the A/B order reversed which scheme measured slow, and on distinct inputs (real serving) the two scanners are within noise in both the sequential and parallel paths. Distinct-input A/B on the vendored GLM-5 tokenizer, full encode() vs the PCRE2 tier: ~2.0-2.3x faster across 10k-200k char prompts — the same win as Qwen2. The parallel-path integration test now covers both patterns with digit-heavy blocks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Add o200k (gpt-oss) split scanner and knext drift guard Wire the o200k_base scheme (GPT-4o, gpt-oss) into the specialized-scanner tier: the fast_split_o200k walker was already parameterized, so this is advance_pos::<CONTRACTIONS=true, DIGITS3=true, SLASH=true, HAN=false> plus the pattern const. Every vendored tokenizer is now scanner-backed (qwen3->Qwen2, glm-5.2->Qwen2N3, kimi-k2.5->Kimi, gpt-oss->O200k). Also add a guard for knext / Kimi K3: it ships as a tiktoken model whose pat_str is byte-identical to KIMI_PATTERN and reaches from_pattern verbatim through the tiktoken->tokenizer.json conversion, so it is caught by the Kimi scanner with no new code. recognizes_only_known_patterns now reassembles that pat_str exactly as tokenization_kimi.py builds it and asserts equality + Kimi routing, so an upstream regex tweak fails loudly instead of silently dropping to PCRE2. Parity: o200k gets 6000-round differential fuzz vs fancy-regex (shared with the Kimi fuzz) + edge cases + an integration test through Split::pre_tokenize including the parallel path. 216 tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Add large multilingual corpus differential test for split scanners Vendor a ~2 MB real-prose corpus (tests/fixtures/corpus_multilingual.txt): ~75% English (OpenWebText) + ~25% across eight non-Latin scripts (Han, kana, Hangul, Cyrillic, Arabic, Devanagari, Thai, Greek, Hebrew, from Wikipedia). Provenance and licenses documented in PROVENANCE.md. Each scheme now has a `*_matches_regex_on_corpus` test that runs the whole corpus through the scanner and fancy-regex in lockstep, asserting byte-identical token spans token-by-token (gigatoken-style, with byte offset + surrounding-token diagnostics on divergence). This closes the gap vs the short synthetic fuzz: real multilingual prose exercises the Han-run, mark, case-boundary, and mixed-script-transition paths — notably Kimi's `[\p{Han}]+` / `&&[^\p{Han}]` rules — at scale. Hermetic (no network, no ~/data provisioning), runs every CI. All four schemes match over 2.12 MB: qwen2 324601, qwen2n3 317780, kimi 312168, o200k 296620 tokens, zero divergence. Since the scanner only replaces the split stage, span-identity over a large real corpus is end-to-end token-id equivalence for the change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Fix O200k parallel safe-boundary: slash after newline is crossable find_safe_boundary treated any `\n` followed by a non-whitespace ASCII byte as an uncrossable chunk boundary. That holds for Qwen2/Qwen2N3/Kimi (punct tail `[\r\n]`), but O200k's tail is `[\r\n/]*`, so a punctuation token absorbs a `\n` and a following `/` as one token (e.g. `.\n/` -> one token). A boundary placed between the `\n` and the `/` is then crossed: on a multi-core host with input >= 16 KiB, the preceding chunk scans past its authority range (debug_assert in debug builds; overlapping ranges and a changed token sequence in release), and the next chunk re-scans the `/`. Exclude `/` as a boundary-follower for the O200k scheme only. Regression tests: a direct find_safe_boundary assertion (O200k skips the slash boundary, other schemes keep it) and a parallel-path Split::pre_tokenize test over >2*MIN_CHUNK_SIZE of mixed `\n/` and `\n<letter>` text asserting fast == generic. Both fail without the fix. Reported by review (codex). Thanks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Python 3.9 reached end-of-life in October 2025. Raising the abi3 floor to 3.10 lets PyO3 use PyUnicode_AsUTF8AndSize (limited API since 3.10) for &str/String extraction instead of the PyUnicode_AsUTF8String fallback, which allocates a temporary bytes object and copies twice per string crossing the boundary. Measured (same source, only the abi3 feature flag changed, Python 3.12, taskset-pinned, distinct inputs, 2 rounds): encode_batch+into_numpy ~10% faster, apply_chat_template ~7% faster (up to ~15% on 1MB chat histories), single encode ~12% faster. Still a single forward- compatible abi3 wheel; only the floor moves from 3.9 to 3.10. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Why
After the specialized split scanners in #25, BPE is the dominant stage on long contexts. Most pretokens are short, but the existing path still paid heap and general-priority-queue costs; repeated cold work also contended on allocating shared hash maps.
The small path is covered by a differential test against the general heap implementation. Inputs above 32 bytes retain the existing merge path, and the cache remains an optimization only.
Performance
Intel Xeon Platinum 8480+, 8 pinned physical cores,
RAYON_NUM_THREADS=8, median of 5 cold rounds:Each timed iteration constructs a fresh tokenizer before encoding; construction is excluded and reported separately.
Validation
cargo test --release --lib(223 passed, 9 ignored)cargo fmt --checkgit diff --check