Skip to content

Group cold long-context BPE work by content - #50

Closed
michaelfeil wants to merge 54 commits into
crusoecloud:mainfrom
michaelfeil:mf/bpe-content-affinity
Closed

Group cold long-context BPE work by content#50
michaelfeil wants to merge 54 commits into
crusoecloud:mainfrom
michaelfeil:mf/bpe-content-affinity

Conversation

@michaelfeil

Copy link
Copy Markdown
Contributor

Stack

Depends on #49. Until #49 lands, GitHub shows the parent commits in this PR; after it lands, this diff reduces to the scheduler layer and its benchmark/blog documentation.

Summary

  • assign equal pretokens to the same BPE worker on the first eligible long encode
  • record per-split output boundaries and restore original token order in one linear pass
  • retain the existing contiguous scheduler outside exact Kimi/o200k patterns, inputs below 64 KiB, and subsequent warm encodes
  • add a reproducible Python comparison against tiktoken 0.13.0 with exact token-ID parity checks
  • add the draft long-context performance post in blog.md

Why

Contiguous chunks send repeated words and punctuation to multiple workers. With empty thread-local caches, each worker repeats the same BPE computation. Content affinity computes a repeated pretoken on one worker, then serves later occurrences from that worker cache.

The one-shot gate is important: unconditional affinity substantially regressed warm Kimi encodes. Restricting it to the first eligible encode preserves the existing warm path (0.492 ms before vs 0.478 ms after at 200k characters, within noise). Qwen and GLM retain the contiguous path after mixed-model A/B results showed no reliable affinity gain.

Performance

Intel Xeon Platinum 8480+, 8 pinned physical cores, median of 5 cold Rust rounds:

Kimi context parent this PR incremental vs main
208,710 tokens 11.83 ms 10.29 ms -13.0% -26.6%
1,003,641 tokens 36.79 ms 32.57 ms -11.5% -29.4%

Public Python API vs tiktoken 0.13.0 on identical gpt-oss/o200k output:

cores tokens fastokens tiktoken speedup
1 204,056 54.67 ms 71.63 ms 1.31x
1 1,022,785 149.44 ms 287.75 ms 1.93x
8 204,056 13.19 ms 64.09 ms 4.86x
8 1,022,785 37.65 ms 261.43 ms 6.94x

Validation

  • cargo test --release --lib (226 passed, 9 ignored)
  • cargo test --release content_affinity --lib
  • python -m py_compile examples/tiktoken_long_context.py
  • benchmark asserts every fastokens token ID equals tiktoken before timing
  • cargo fmt --check
  • git diff --check

michaelfeil and others added 30 commits July 10, 2026 19:48
[codex] release GIL during Python encode
…al-options

[codex] reject unsupported special token options
[codex] support split special token encoding
Rename fork crate package to fasttokens-b10
[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>
michaelfeil and others added 24 commits July 11, 2026 21:26
* 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants