From d1ad772c90d7212c7ec59b59b6a856a5674d7ce5 Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 6 Aug 2026 16:59:53 +0900 Subject: [PATCH 1/7] perf(pre-tokenizer): bitstream splitter for the gpt2 and cl100k regexes The GPT regexes are a scan with one unpredictable branch per token. `bitsplit` decides them from bitmaps instead: build a few `u64` masks over the bytes, then take token starts 64 bytes at a time with register ops rather than a branch per character. A new crate, wired in behind two match arms in `Split::pre_tokenize`. It only takes over where its bitstream build is SIMD -- `bitsplit::fast_builder()` -- because the portable builder is slower than the FSM it replaces, so every other target and every other regex keeps the existing atomsplit FSM. cl100k is gated on the standard digit cap of 3, the only variant with a bitstream splitter. `classify_into_spans_bits` is the `classify_into_spans` equivalent for a splitter that works off bitstreams: it hands the splitter two `u64` bitmaps alongside the tags. 367 tests pass, including the pre-tokenizer oracles that hold the FSM and the bitstream splitter to the same spans as the `regex` crate. --- tokenizers/Cargo.lock | 9 + tokenizers/Cargo.toml | 2 +- tokenizers/bitsplit/Cargo.toml | 13 + tokenizers/bitsplit/mask_splitter_spec.md | 336 ++++++++++++++++ tokenizers/bitsplit/src/deepseek.rs | 319 +++++++++++++++ tokenizers/bitsplit/src/gpt.rs | 348 ++++++++++++++++ tokenizers/bitsplit/src/lib.rs | 373 ++++++++++++++++++ tokenizers/bitsplit/src/simd.rs | 176 +++++++++ tokenizers/bitsplit/src/simd_x86.rs | 192 +++++++++ tokenizers/tk-encode/Cargo.toml | 1 + .../tk-encode/src/pre_tokenizers/split.rs | 23 ++ .../tk-encode/src/tokenizer/pipeline.rs | 45 +++ 12 files changed, 1836 insertions(+), 1 deletion(-) create mode 100644 tokenizers/bitsplit/Cargo.toml create mode 100644 tokenizers/bitsplit/mask_splitter_spec.md create mode 100644 tokenizers/bitsplit/src/deepseek.rs create mode 100644 tokenizers/bitsplit/src/gpt.rs create mode 100644 tokenizers/bitsplit/src/lib.rs create mode 100644 tokenizers/bitsplit/src/simd.rs create mode 100644 tokenizers/bitsplit/src/simd_x86.rs diff --git a/tokenizers/Cargo.lock b/tokenizers/Cargo.lock index a78cd06bc..cad5d24c0 100644 --- a/tokenizers/Cargo.lock +++ b/tokenizers/Cargo.lock @@ -165,6 +165,14 @@ dependencies = [ "unicode-properties", ] +[[package]] +name = "bitsplit" +version = "0.1.0" +dependencies = [ + "ahash", + "atomsplit", +] + [[package]] name = "bitvec" version = "1.1.1" @@ -2287,6 +2295,7 @@ dependencies = [ "ahash", "assert_approx_eq", "atomsplit", + "bitsplit", "compact_str", "criterion 0.6.0", "daachorse 3.0.2", diff --git a/tokenizers/Cargo.toml b/tokenizers/Cargo.toml index 3759dfdb0..fdd4f4894 100644 --- a/tokenizers/Cargo.toml +++ b/tokenizers/Cargo.toml @@ -1,6 +1,6 @@ [workspace] resolver = "3" -members = ["bitmap_gen", "atomsplit", "tk-encode", "tk-train"] +members = ["bitmap_gen", "atomsplit", "bitsplit", "tk-encode", "tk-train"] [package] authors = [ diff --git a/tokenizers/bitsplit/Cargo.toml b/tokenizers/bitsplit/Cargo.toml new file mode 100644 index 000000000..ebc3d609e --- /dev/null +++ b/tokenizers/bitsplit/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "bitsplit" +version = "0.1.0" +edition = "2024" + + +[dependencies] +ahash = "0.8" +atomsplit = { path = "../atomsplit" } + + + + diff --git a/tokenizers/bitsplit/mask_splitter_spec.md b/tokenizers/bitsplit/mask_splitter_spec.md new file mode 100644 index 000000000..f6026232b --- /dev/null +++ b/tokenizers/bitsplit/mask_splitter_spec.md @@ -0,0 +1,336 @@ +# bitsplit — mask splitter spec + +A Rust library for **pre-tokenizer splitting as a bitstream program**. Not a general regex engine: +it covers exactly what `tokenizers` needs — the GPT-family regexes, the class-run family, and +literal delimiters — and it does so with one shared operator set so a new grammar is ~150 lines of +boolean algebra rather than a new hand-rolled FSM. + +Lineage: Parabix (Cameron et al., *Boosting the efficiency of text processing on commodity +processors*) for the operator set and carry discipline; *Interleaved Bitstream Execution for +Multi-Pattern Regex Matching on GPUs* (MICRO'25, doi 10.1145/3725843.3756052) for the execution +model — fuse every instruction into ONE block-wise loop instead of one pass per instruction. + +Status: proven in `scratch/bitsplit` for three grammars, all byte-exact against their +`atomsplit::fsm` counterparts. + +--- + +## 1. Why it wins + +An FSM decides one token per unpredictable branch, so its cost tracks bytes-per-token. A bitstream +program decides 64 bytes per register op, branchlessly, so its cost is flat. + +``` + ns/B (english, same run, aarch64) + fsm_deepseek ████████████████████████████████████████████████ 2.93 4.7 B/tok + fsm_cl100k ██████████████████████████████ 1.86 + fsm_byte_level ████████████████████████████ 1.71 + bitsplit ds █████████ 0.58 + bitsplit cl100k████████ 0.50 + bitsplit bl ██████ 0.41 +``` + +The FSM spread across grammars is 71% (1.71 → 2.93); bitsplit's is 20%. **Flat cost across grammars +is the design invariant** — it means the per-grammar layer is thin, which is what makes the library +worth having. + +--- + +## 2. Pipeline + +``` + text ─────────────────────────────────────────────────────────────┐ + │ │ + ▼ │ + classify (atomsplit) one Atom tag per byte │ + │ │ + ▼ ▼ + ┌─────────────────────── per 64-byte block, fused ────────────────────────┐ + │ │ + │ tags ──► LUT ──► dense 3-bit code ──► fill conts ──► 3 bit-planes │ + │ │ │ + │ decode ─────┼──► class │ + │ │ streams │ + │ ▼ │ + │ grammar algebra (Parabix operators) │ + │ │ │ + │ ▼ │ + │ starts │ keep (u64 each) │ + └──────────────────────────────────────────────────────┬───────────────────┘ + │ ~n/8 bytes + ▼ + emit walk ──► spans / sink +``` + +Everything inside the box lives in registers. The **only** intermediate that reaches memory is the +`starts` bitmap. That is the whole point of interleaved execution: the sequential model (one loop +per bitstream instruction, each materialising a full-length stream) is what the MICRO'25 paper +beats, and it is what we must not accidentally rebuild. + +--- + +## 3. The universal intermediate + +Every pre-tokenizer in `tokenizers` reduces to **one bit per byte: "a token starts here"**, plus +optionally "this byte is kept". That single contract is what makes the library abstract without +being a general engine. + +``` + text H e l l o , w o r l d + starts 1 . . . . 1 1 . . . . . + keep 1 1 1 1 1 1 . 1 1 1 1 1 (Removed behaviour drops the space) + └──"Hello"──┘ "," └──"world"──┘ +``` + +Which collapses three families onto one emit: + +| family | how `starts` is produced | +|---|---| +| class runs (WhitespaceSplit, Punctuation, Digits, Whitespace, Bert) | `c & !(c<<1)` per class; DROP/ISOLATE/KEEP_A are 3 bit ops | +| regex grammars (gpt2, cl100k, o200k, tekken, deepseek) | the algebra of §6 | +| literals (Metaspace `▁`, CharDelimiterSplit) | shifted-AND chain of byte compares | + +And `SplitDelimiterBehavior` stops being per-grammar code — it is a post-pass on the mask: + +``` + match . . 1 1 1 . . (a matched delimiter run) + Isolated . . 1 . . 1 . starts at both edges + Removed . . 1 . . 1 . + keep &= !match + MergedWithPrevious drop the leading edge bit + MergedWithNext drop the trailing edge bit + Contiguous merge adjacent same-kind runs +``` + +--- + +## 4. Operator set + +Parabix names, because they are the literature standard. Every operator takes a `&mut Carry` so +cross-block state is explicit and never forgotten. + +| operator | meaning | cost | +|---|---|---| +| `advance(m, n)` | move markers forward n bytes | 1–2 ops | +| `advance_char(m, cont)` | move markers forward one **char** (§5) | ~5 ops | +| `scan_thru(m, c)` | move each marker past the run of `c` it sits in | 1 add | +| `match_star(m, c)` | Kleene closure over a class | 1 add | +| `span_upto(m, e)` | set every bit from marker to end marker | 1 sub | +| `fill_to_last(m, c)` | in each `c`-run, fill from run start through the LAST marker | 2 ops | +| `to_lead(x, cont)` | move each bit back to its char's lead byte | ≤3 steps | +| `run_start(c)` | `c & !(c<<1)` | 2 ops | + +Do all of these in `u128` and read bit 64 as the carry-out. That one habit removes an entire class +of bug: a run reaching bit 63 puts its landing bit at 64 instead of vanishing, and `e - m` still +yields the correct in-block span. + +``` + fill_to_last, worked (finds "start after the LAST newline in a ws run"): + + run of c ┌───────────────────────┐ + c (ws) 1 1 1 1 1 1 1 1 . . + m (nl) . 1 . . 1 . . . . . + scan_thru . . . . . . . . 1 . ← lands past the run + (e-m)|m 1 1 1 1 1 . . . . . ← run start .. last marker + ▲ + the surviving "after last newline" bit is here+1 +``` + +`(e - m) | m` rather than `e - m`: the latter only spans from the *first* marker when a run holds +several. This is the single least obvious identity in the whole library — write a test for it. + +--- + +## 5. The transpose (bytes → bitstreams) + +The measured cost centre: **45–73% of total before optimisation, 36% after.** Optimise it before +touching any grammar. + +``` + tags [Lt][Lt][Pu][Sp][Lt][Lt][Lt][Ap][Lt] ... 16 lanes at a time + │ + │ vqtbl4q (64-entry table, indexed by the RAW tag — + │ refinements 0x10/0x20/0x16/0x26 index straight in) + ▼ + code 0 0 2 4 0 0 0 6 0 ... dense 3-bit code + │ + │ fill continuations from the left: 2 steps, not 3. + │ shift-1 makes every lane right at distance 1, so shifting + │ THAT by 2 covers distances 2 and 3 — the most a 4-byte char needs. + ▼ + code' 0 0 2 4 0 0 0 6 0 ... every byte carries its CHAR's code + │ + │ 3 × (vtst + vand POW + 3 × vpaddq_u8) ← the 64×8 transpose + ▼ + p0 p1 p2 three u64 bit-planes of the code, for 64 bytes + │ + │ decode: ~14 boolean ops + ▼ + l n other nl sp ws apo class streams +``` + +**Filled streams are the load-bearing decision.** Because every byte of a multi-byte char carries +its char's code, "previous char's class" is a plain `<< 1` and no rule anywhere does char-width +arithmetic. Give that up and every rule grows an `advance_char`. + +Three cost rules learned the hard way: + +1. **Reduce plane COUNT, not per-plane cost.** ~13 ops/plane on NEON is the floor; a dense 3-bit + code needing 3 planes beat one-hot class bits needing 7 (0.33 → 0.22 ns/B). +2. **Derive, don't extract.** `ascii = lead & last_byte` (a char that is its own first and last byte + is single-byte, i.e. ASCII) and `ascii_alpha = lm & ascii` (ASCII has no marks). Two `u64` ops + replaced two SIMD extractions. +3. **Gate rare masks per block.** The CJK range test is behind one `vmaxvq`, so Latin text pays ~8 + ops for it instead of ~70. + +Multi-byte-char predicates (the CJK range) belong in **vector space on `vext`-aligned b1/b2 → one +mask**, not as N separate byte-predicate bitstreams recombined with shifts. + +--- + +## 6. A grammar, end to end + +Worked on `"Hi, don't"` under GPT-2 byte-level. Classes: `l` letter, `o` other, `s` space, `a` +apostrophe (a sub-code of other, so it can be flagged). + +``` + index 0 1 2 3 4 5 6 7 8 + text H i , ␣ d o n ' t + class l l o s l l l a l + + l 1 1 . . 1 1 1 . 1 + other . . 1 . . . . 1 . + ws . . . 1 . . . . . + sp . . . 1 . . . . . + + l<<1 . 1 1 . . 1 1 1 . + sp<<1 . . . . 1 . . . . + other<<1 . . . 1 . . . . 1 + + l_start = l & !(l<<1) & !(sp<<1) 1 . . . . . . . 1 + o_start = other & !(other<<1) & !(sp<<1) . . 1 . . . . 1 . + ws_start = ws & !(ws<<1) . . . 1 . . . . . + steal = ws & lb & !(ws>>1) & !eof . . . 1 . . . . . + ───────────────────────────────── + starts 1 . 1 1 . . . 1 1 + flag = starts & apo . . . . . . . 1 . +``` + +Emit walk, with the scalar escape firing at bit 7: + +``` + [0,2) "Hi" [2,3) "," [3,7) " don" ⟨escape⟩ [7,9) "'t" ← bit 8 skipped +``` + +`" don"` shows `steal` doing its job: `\s+(?!\S)` hands the run's last whitespace char to whatever +follows, which is exactly the ` ?` prefix of the next alternative. One rule covers both. + +### Per-grammar checklist + +1. Dense code table `[u8; 64]`, `tag → code`, cont = 7. Merge every class the grammar treats + identically (deepseek: Letter+Mark share a code; gpt2: Mark is "other"). +2. `decode(p0,p1,p2,valid) -> Cls`. Mask with `valid` **only** the class whose code is 0 — past the + block end every plane reads 0. +3. Run starts, one line each. +4. Prefix / steal rules. +5. Bounded repetition, span clearing, absorption. +6. Carries + the edge peek. + +--- + +## 7. Escape hatches + +Not everything belongs in bit algebra. **Flag the bits, resolve them scalar-ly in the emit walk.** + +Contractions (`'s 't 're 've 'm 'll 'd`) are variable-length, case-optional, and outrank every other +alternative. In bit algebra that is miserable; as an escape it is 20 lines. `flag = starts & apo` +costs one word-test per block, so ordinary text pays nothing. + +Two bugs this cost, both worth a regression test: + +- The escape must not re-consume the start bit **at its own end** — the algebra usually has one + there (the letter run resumes) and consuming it emits an empty span. +- Contractions **chain**: `'re've`, `y'all'd've`. Loop until a match fails rather than handing + control back to the algebra, whose start bit you are about to skip. + +Same shape applies to any rare, awkward, variable-length rule. + +--- + +## 8. Rules and pitfalls + +**Prefer backward formulations.** State a rule as "my predecessor did X", never "my successor is +X". Forward needs `advance_char`, which silently drops the marker when the two chars straddle a +block edge — every cl100k failure was this. Backward needs only a shift carry. + +``` + ✗ forward o_prefix = o_start & next_is_letter ; suppress advance_char(o_prefix) + └─ marker leaves the block, lost + + ✓ backward osf = smear(o_start over its char's bytes) + l_start &= !(osf << 1) └─ one carry bit, always exact +``` + +**Carries at a block boundary.** One bit per stream, plus one per open scan: + +``` + block b-1 │ block b + ... 61 62 63 │ 0 1 2 ... + ▲ │ ▲ + └─ last byte's code ──────────┼───┘ seeds the fill and every `p1` + cont word ─────────────────┼───► to_lead underflow patch + open scans (aa/nl runs) ───┼───► carry-in at bit 0 + │ + peek 1 char forward ◄─────────┼─── makes every `n1` exact at the edge +``` + +Other pitfalls, each caught by the fuzzer: + +- Mask shifted-in zeros at u64 edges **before** inverting, or runs truncate silently. +- The first span needs special dispatch (`starts |= 1` on block 0). +- Cont-resolution needs prev-chunk context (`vext`), so starts land on leads. +- Bounded repetition needs a single-byte fast path — three `advance_char`s collapse to one `<< 3` + against `n & (n<<1) & (n<<2)`. Without it, dense digits run **0.58×** the FSM. +- One backward-in-time rule survives (`\s*[\r\n]+`'s "after the LAST newline"): patch the + already-written bitmap. The patch must NOT fire on a bit that a punct tail's `[\r\n]*` already + made a run start — the absorption cut the whitespace run, so a later newline cannot reach past it. + +--- + +## 9. Testing + +Non-negotiable, because every one of these caught a real bug: + +1. Byte-exactness vs the reference FSM on every input, for **every** grammar. +2. Exhaustive 3-char and 4-char sweeps over `[space, \n, a, 1, !, \0, 世, \u{a0}, \t, ']`. +3. Structured fuzz (~200k cases) over a pool hitting every arm: CJK, marks, ZWJ, control, multi-byte + whitespace, Other_Alphabetic symbols, contractions, punct+alpha. +4. **Every corpus re-sliced at 70 offsets.** This is what finds carry bugs; whole-file runs do not. + +## 12. Layout + +``` + src/ + lib.rs Blk, primitives, builders, emit, emit_contr + simd.rs NEON transpose (x86 sibling slots in beside it) + ops.rs the Parabix operator set + Carry + grammar/ + byte_level.rs cl100k.rs o200k.rs deepseek.rs class_runs.rs literal.rs + bin/ + verify.rs byte-exactness oracle vs atomsplit FSMs + bench.rs per-grammar throughput +``` + +Each grammar is independently testable against its reference FSM, which is what keeps the shared +layer honest. + +--- + +## 13. Open + +- **Does the library own classification?** Currently depends on `atomsplit::classify`. Fusing buys + ~0.07 ns/B on English but up to ~0.3 on CJK, where classify (0.57) now exceeds the split (0.24). + Design the block builder so classify can be fused later without changing the operator API. +- **x86 builder.** The scalar path covers it correctly but slowly; the `vpaddq` reduction has a + natural `movemask` equivalent. +- **o200k / tekken.** Case-refined letter runs; the case refinement is already in the Atom tags + (`0x10`/`0x20`), so it is a code-table change, not new machinery. diff --git a/tokenizers/bitsplit/src/deepseek.rs b/tokenizers/bitsplit/src/deepseek.rs new file mode 100644 index 000000000..3bafdbe91 --- /dev/null +++ b/tokenizers/bitsplit/src/deepseek.rs @@ -0,0 +1,319 @@ +//! DeepSeek-V3/V4 pre-tokenization: the `Sequence` of `\p{N}{1,3}` → `[一-龥぀-ゟ゠-ヿ]+` → +//! the big regex, all `Isolated`, as one bitstream program. Byte-exact with +//! `atomsplit::fsm::fsm_deepseek`. + +use crate::{ + CODE_CONT, CONT, Span, adv, build_block, emit, fill_to_last, is_cjk_at, lead_run, scanthru, + to_lead, trail_run, +}; + +/// Atom tag → dense 3-bit code. Letter and Mark share a code because deepseek's letter run is +/// `[\p{L}\p{M}]+`; `AlphaSymMark` (0x16) is categorically `\p{S}` so it takes the punct path, and +/// `Zwj` (0x26) matches no alternative at all so it is a gap char. +pub(crate) const LUT: [u8; 64] = { + let mut t = [6u8; 64]; // gap + t[0x00] = 0; + t[0x10] = 0; + t[0x20] = 0; + t[0x06] = 0; // Letter (+case) | Mark + t[0x01] = 1; + t[0x02] = 1; // \p{N} + t[0x07] = 2; + t[0x08] = 2; + t[0x09] = 2; + t[0x0A] = 2; + t[0x16] = 2; // \p{P} ∪ \p{S} + t[0x03] = 3; + t[0x04] = 4; + t[0x05] = 5; // Newline | Space | WsOther + t[0x0F] = CODE_CONT; + t +}; + +const S_N: u16 = 1 << 0; +const S_LM: u16 = 1 << 1; +const S_PS: u16 = 1 << 2; +const S_WS: u16 = 1 << 3; +const S_NL: u16 = 1 << 4; +const S_SP: u16 = 1 << 5; +const S_CJK: u16 = 1 << 6; +const S_ANY: u16 = S_N | S_LM | S_PS | S_WS | S_CJK; + +/// Dense code → class bits, for the block-edge carries. +const fn code_bits(code: u8) -> u16 { + match code { + 0 => S_LM, + 1 => S_N, + 2 => S_PS, + 3 => S_WS | S_NL, + 4 => S_WS | S_SP, + 5 => S_WS, + _ => 0, // gap | cont + } +} + +/// Filled class bits of the char starting at `p` (a lead byte). +fn bits_at(text: &[u8], tags: &[u8], p: usize) -> u16 { + code_bits(LUT[tags[p] as usize]) | if is_cjk_at(text, p) { S_CJK } else { 0 } +} + +/// Cross-block state: what the paper resolves with selective recomputation across SMs. +struct Carry { + code: u8, // filled dense code of the previous block's last byte (seeds the fill) + cjk: bool, // ...and whether its char is in the Split-2 CJK range + cont: u64, // the previous block's `cont` stream (for `to_lead` underflow) + aa_run: bool, // an alt-1 `[A-Za-z]+` run is still open + nl_run: bool, // a `[\p{P}\p{S}]+[\r\n]*` newline tail is still open + dig_run: bool, // inside a \p{N} run + dig_since: u32, // chars already consumed of the current \p{N}{1,3} group + anl: Option, // a committed "start after the last newline" a later newline may retract +} + +impl Default for Carry { + fn default() -> Self { + // the cont code as the seed: a text opening with a stray continuation byte then classifies + // as gap under both builders instead of as `Letter` (code 0). + Self { + code: CODE_CONT, + cjk: false, + cont: 0, + aa_run: false, + nl_run: false, + dig_run: false, + dig_since: 0, + anl: None, + } + } +} + +/// Pre-tokenize `text` (well-formed UTF-8) with the DeepSeek grammar: writes token spans into `out` +/// and returns the count. `tags` is `atomsplit::classify`'s output (len ≥ `text.len()`), `starts` +/// is scratch for the token-start bitmap (len ≥ `text.len().div_ceil(64)`). +#[must_use] +pub fn bitsplit_deepseek(text: &[u8], tags: &[u8], starts: &mut [u64], out: &mut [Span]) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + assert!(tags.len() >= ntext && starts.len() >= ntext.div_ceil(64) && out.len() >= ntext); + let nblk = ntext.div_ceil(64); + let mut cy = Carry::default(); + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = build_block::(text, tags, base, len, &LUT, cy.code, cy.cjk); + // planes → classes: 3 extractions instead of 7 one-hot masks (see `LUT`). Only `lm` + // needs the `valid` mask — past the block end every plane reads 0, i.e. code 0. + let (pa, pc, pd) = (!b.p2 & !b.p1, !b.p2 & b.p1, b.p2 & !b.p1); + let (s_lm, s_n) = (pa & !b.p0 & valid, pa & b.p0); + let (s_ps, s_nl) = (pc & !b.p0, pc & b.p0); + let (s_sp, s_ws) = (pd & !b.p0, pd | (pc & b.p0)); + let last_cjk = b.cjk >> (len - 1) & 1 != 0; + let last_bits = code_bits(last_code) | if last_cjk { S_CJK } else { 0 }; + + // ── edges. The byte before the block is carried; the one after is peeked. One char of + // lookahead is all the grammar needs outside whitespace runs, so peeking it keeps every + // "next char is X" rule exact right at the block boundary — no recomputation needed. + let pb = if base == 0 { + 0 + } else { + code_bits(cy.code) | if cy.cjk { S_CJK } else { 0 } + }; + let (nb, nb_lead, nb_aa) = if last_blk { + (0u16, true, false) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + let bits = if is_lead { + bits_at(text, tags, q) + } else { + last_bits + }; + (bits, is_lead, text[q].is_ascii_alphabetic()) + }; + let has = |v: u16, s: u16| v & s != 0; + // Split precedence: CJK (Split-2) outranks the big regex, and `fsm_deepseek` tests it ahead + // of the digit arm too — so peel CJK off every other class first. + let cjk = b.cjk; + let num = s_n & !cjk; + let lm = s_lm & !cjk; + let ps = s_ps & !cjk; + let gap = valid & !(s_n | s_lm | s_ps | s_ws | cjk); + let cjk_l = cjk & s_lm; // Split-3 re-splits the isolated CJK run into same-kind sub-runs + let cjk_p = cjk & !s_lm; + // `[^\r\n\p{L}\p{P}\p{S}]?` — the letter alternative's optional one-char prefix. Digits and + // CJK are already isolated by Split-1/2, so what remains is non-newline ws + gap chars. + let prefix = (s_ws & !s_nl) | gap; + + // previous-byte / next-byte membership. Streams are filled, so at a lead `p1` reads the + // previous *char*'s class and at a last byte `n1` reads the next char's. + let p1 = |x: u64, c: bool| (x << 1) | u64::from(c); + let n1 = |x: u64, c: bool| (x >> 1) | (u64::from(c) << 63); + let pb_gap = base != 0 && !has(pb, S_ANY); + let pb_cjk = has(pb, S_CJK); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + // a char that is both its own first and last byte is single-byte, i.e. ASCII; an ASCII char + // in `\p{L}∪\p{M}` is exactly `[A-Za-z]` (ASCII has no marks). Both streams for 2 ops. + let ascii = lead & lb; + let aa = s_lm & ascii; + + // ── run starts. All purely backward-looking, so the carry alone makes them exact. ─────── + let n_start = num & lead & !p1(num, has(pb, S_N) && !pb_cjk); + let lm_start = lm + & lead + & !p1(lm, has(pb, S_LM) && !pb_cjk) + & !p1(prefix, (has(pb, S_WS) && !has(pb, S_NL)) || pb_gap); + let ws_start = s_ws & lead & !p1(s_ws, has(pb, S_WS)); + let gap_start = gap & lead & !p1(gap, pb_gap); + let ps_start = ps & lead & !p1(ps, has(pb, S_PS) && !pb_cjk) & !p1(s_sp, has(pb, S_SP)); + let cjk_start = (cjk_l & lead & !p1(cjk_l, pb_cjk && has(pb, S_LM))) + | (cjk_p & lead & !p1(cjk_p, pb_cjk && !has(pb, S_LM))); + + // ── Split-1 `\p{N}{1,3}`: the one non-local rule — a group boundary every 3 chars from the + // run start. Marker iteration (the paper's bounded-repetition lowering); ≤21 rounds/block, + // and 0 rounds on text without digit runs. + let mut m = n_start; + if cy.dig_run && has(pb, S_N) { + // resume mid-run: the next group start is (3 - since) chars into the block. Re-mask + // with `num` at every hop — the carry only says the char *at* the edge was a digit, the + // run may well have ended there (`Ⅷ` straddling the edge, then `\t`, then `456`). + let mut s = lead & lead.wrapping_neg() & num; // first lead of the block + for _ in 0..((3 - cy.dig_since % 3) % 3) { + s = adv(s, b.cont) & num & lead; + } + m |= s; + } + let mut groups = m; + if num & b.cont == 0 { + // Fast path: every digit in this block is single-byte, so "3 chars on" is just `<< 3` + // and the three `adv`s collapse into one shift against a precomputed mask (`num3` asks + // that the two skipped positions are digits too, which is what the `adv` chain checked). + // ~3 ops per group instead of ~14 — dense-digit text is otherwise this loop's worst case. + let num3 = num & (num << 1) & (num << 2); + while m != 0 { + m = (m << 3) & num3; + groups |= m; + } + } else { + while m != 0 { + let a = adv(m, b.cont) & num & lead; + let c = adv(adv(a, b.cont) & num & lead, b.cont) & num & lead; + if c == 0 { + break; + } + groups |= c; + m = c; + } + } + + // ── whitespace: `\s*[\r\n]+ | \s+(?!\S) | \s+`. + // (a) the run's first token runs through its LAST newline → a start right after it, unless + // a further newline still follows inside the run (a backward scan, hence the reversal). + let anl = p1(s_nl, has(pb, S_NL)) & s_ws & lead; + let later_nl = fill_to_last(s_nl.reverse_bits(), s_ws.reverse_bits()).reverse_bits(); + let after_nl = anl & !later_nl; + // (b) the run's last char is handed to whatever follows, as its `[^…]?` / ` ?` prefix — + // unless the run ends the input or the next piece is Split-1/2-isolated (`(?!\S)`). + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + let steal_lb = s_ws + & !s_nl + & lb + & !eof_bit + & !n1(s_ws, has(nb, S_WS)) + & !n1(s_n | cjk, has(nb, S_N | S_CJK)); + let (steal, steal_patch) = to_lead(steal_lb, b.cont, cy.cont); + // (c) the same one-char give-back out of a gap run (Control / NumericOther / ZWJ match no + // alternative, so the run is one piece minus the char a following letter run claims). + let (gap_steal, gap_patch) = to_lead( + gap & lb & n1(lm, has(nb, S_LM) && !has(nb, S_CJK)), + b.cont, + cy.cont, + ); + + // ── alt-1 `[ascii_punct][A-Za-z]+`: fires only where the scan is actually positioned, i.e. + // at a punct-run start no space swallowed. Its `[A-Za-z]+` run then has no interior starts + // and forces one at its end — which the letter rule alone would suppress (`!c` reads `!ab`, + // `c` after it is a fresh token even though its predecessor is a letter). + let alt1 = ps_start & ascii & n1(aa, nb_aa); + let aa_m = ((alt1 as u128) << 1) | u128::from(cy.aa_run); + let aa_e = scanthru(aa_m, aa as u128); + let aa_span = aa_e.wrapping_sub(aa_m); + // ── a punct run's `[\r\n]*` tail swallows the newlines directly behind it. + let nl_m = + ((p1(ps, has(pb, S_PS) && !pb_cjk) & s_nl & lead) as u128) | u128::from(cy.nl_run); + let nl_e = scanthru(nl_m, s_nl as u128); + let nl_span = nl_e.wrapping_sub(nl_m); + + // ── the one backward-in-time dependency: a newline arriving now retracts the "start after + // the last newline" already committed for a whitespace run that was open at the last edge. + if let Some(p) = cy.anl + && has(pb, S_WS) + && s_ws & 1 != 0 + && s_nl & lead_run(s_ws, valid) != 0 + { + starts[p / 64] &= !(1u64 << (p % 64)); + cy.anl = None; + } + + let mut st = groups + | lm_start + | ws_start + | gap_start + | ps_start + | cjk_start + | after_nl + | steal + | gap_steal; + st &= !(aa_span as u64) & !(nl_span as u64); + st |= aa_e as u64 | nl_e as u64; + st &= lead; + if bi == 0 { + st |= 1; // position 0 always opens a token + } + starts[bi] = st; + if bi > 0 { + starts[bi - 1] |= steal_patch | gap_patch; + } + + // ── carries ──────────────────────────────────────────────────────────────────────────── + cy.aa_run = aa_e >> 64 != 0; + cy.nl_run = nl_e >> 64 != 0; + let tn = trail_run(num, valid, len); + cy.dig_run = tn != 0 && has(nb, S_N) && !has(nb, S_CJK); + cy.dig_since = if !cy.dig_run { + 0 + } else { + let g = groups & tn; + let counted = if g == 0 { + cy.dig_since + (num & lead & tn).count_ones() + } else { + (num & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() + }; + counted % 3 + }; + let tws = trail_run(s_ws, valid, len); + if tws != 0 && has(nb, S_WS) { + // ...but not a bit that a punct tail's `[\r\n]*` already made a *run start*: the + // absorption cut the whitespace run, so a later newline cannot reach back past it. + let a = after_nl & tws & !(nl_e as u64); + if a != 0 { + cy.anl = Some(base + 63 - a.leading_zeros() as usize); + } else if !(tws & 1 != 0 && has(pb, S_WS)) { + cy.anl = None; // a fresh run with no newline yet — nothing left to retract + } + } else { + cy.anl = None; + } + cy.code = last_code; + cy.cjk = last_cjk; + cy.cont = b.cont; + } + + emit(starts, nblk, ntext, out) +} diff --git a/tokenizers/bitsplit/src/gpt.rs b/tokenizers/bitsplit/src/gpt.rs new file mode 100644 index 000000000..a8b89bc8e --- /dev/null +++ b/tokenizers/bitsplit/src/gpt.rs @@ -0,0 +1,348 @@ +//! The two tiktoken-family grammars as bitstream programs. +//! +//! * [`bitsplit_byte_level`] — GPT-2 / Llama / Qwen: +//! `'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+` +//! * [`bitsplit_cl100k`] — cl100k_base / Llama-3: +//! `(?i:'s|…)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+` +//! +//! Both share one dense code table, and both hand contractions to the scalar escape in +//! [`crate::emit_contr`] rather than trying to express a variable-length case-optional literal +//! alternation in bit algebra. + +use crate::{ + CODE_CONT, CONT, Span, adv, build_block, emit_contr, fill_to_last, lead_run, scanthru, to_lead, + trail_run, +}; + +/// Atom tag → dense 3-bit code, shared by both grammars. Unlike deepseek's table, `Mark` is NOT a +/// letter here (`\p{L}` excludes it, so it belongs to the "other" class), and `Apostrophe` gets its +/// own code so the contraction escape can be flagged with one AND — it is still "other" for every +/// run rule, which `decode` restores. +const LUT: [u8; 64] = { + let mut t = [2u8; 64]; // other = [^\s\p{L}\p{N}] + t[0x00] = 0; + t[0x10] = 0; + t[0x20] = 0; // Letter (+ case refinements) + t[0x01] = 1; + t[0x02] = 1; // \p{N} = Nd ∪ Nl ∪ No + t[0x03] = 3; // Newline + t[0x04] = 4; // Space + t[0x05] = 5; // WsOther + t[0x09] = 6; // Apostrophe + t[0x0F] = CODE_CONT; + t +}; + +/// One block's class streams. `other` folds the apostrophe code back in; only `l` needs the +/// `valid` mask, since past the block end every plane reads 0, i.e. code 0. +struct Cls { + l: u64, + n: u64, + other: u64, + nl: u64, + sp: u64, + ws: u64, + apo: u64, +} + +#[inline] +fn decode(p0: u64, p1: u64, p2: u64, valid: u64) -> Cls { + let a = !p2 & !p1; + let nl = !p2 & p1 & p0; + Cls { + l: a & !p0 & valid, + n: a & p0, + other: p1 & !p0 & valid, // codes 2 and 6 — "other" and the apostrophe + nl, + sp: p2 & !p1 & !p0, + ws: (p2 & !p1) | nl, // Space ∪ WsOther ∪ Newline + apo: p2 & p1 & !p0, + } +} + +/// Class bits of a dense code, for the block-edge carry. +const C_L: u8 = 1; +const C_N: u8 = 2; +const C_O: u8 = 4; +const C_NL: u8 = 8; +const C_SP: u8 = 16; +const C_WS: u8 = 32; + +#[inline] +const fn code_bits(code: u8) -> u8 { + match code { + 0 => C_L, + 1 => C_N, + 2 | 6 => C_O, + 3 => C_NL | C_WS, + 4 => C_SP | C_WS, + 5 => C_WS, + _ => 0, // cont never reaches an edge (the fill resolves it) + } +} + +/// GPT-2 / byte-level pre-tokenization. `starts` and `flag` are scratch bitmaps +/// (len ≥ `text.len().div_ceil(64)`); byte-exact with `atomsplit::fsm::fsm_byte_level`. +#[must_use] +pub fn bitsplit_byte_level( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + let nblk = ntext.div_ceil(64); + assert!( + tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext + ); + let (mut code, mut prev_cont) = (CODE_CONT, 0u64); + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = build_block::(text, tags, base, len, &LUT, code, false); + let c = decode(b.p0, b.p1, b.p2, valid); + + let pb = if base == 0 { 0 } else { code_bits(code) }; + let (nb, nb_lead) = if last_blk { + (0u8, true) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + ( + if is_lead { + code_bits(LUT[tags[q] as usize]) + } else { + code_bits(last_code) + }, + is_lead, + ) + }; + let has = |v: u8, s: u8| v & s != 0; + let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); + let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + + // ── every alternative but the contraction is ` ?X+` over a class run, so a token opens at + // each run start — pushed back one char when a literal space sits in front of it (` ?`). + let sp_pfx = p1(c.sp, has(pb, C_SP)); + let l_start = c.l & lead & !p1(c.l, has(pb, C_L)) & !sp_pfx; + let n_start = c.n & lead & !p1(c.n, has(pb, C_N)) & !sp_pfx; + let o_start = c.other & lead & !p1(c.other, has(pb, C_O)) & !sp_pfx; + let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); + // `\s+(?!\S)` hands the run's LAST whitespace char to whatever follows: as a ` ?` prefix if + // it is a space, else as a token of its own. Either way it opens a token — unless the run + // ends the input, where plain `\s+` takes the lot. GPT-2 has no `[\r\n]` rule, so unlike + // cl100k a newline is ordinary whitespace here and can be the stolen char. + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + let (steal, patch) = to_lead( + c.ws & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), + b.cont, + prev_cont, + ); + + let mut st = (l_start | n_start | o_start | ws_start | steal) & lead; + if bi == 0 { + st |= 1; + } + starts[bi] = st; + flag[bi] = st & c.apo; // apostrophes that open a token → contraction escape + if bi > 0 { + starts[bi - 1] |= patch; + } + code = last_code; + prev_cont = b.cont; + } + emit_contr(text, starts, flag, nblk, ntext, false, out) +} + +/// cl100k_base / Llama-3 pre-tokenization. Byte-exact with `atomsplit::fsm::fsm_cl100k`. +#[must_use] +pub fn bitsplit_cl100k( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + let nblk = ntext.div_ceil(64); + assert!( + tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext + ); + let (mut code, mut prev_cont) = (CODE_CONT, 0u64); + let (mut nl_run, mut dig_run, mut dig_since) = (false, false, 0u32); + let mut prev_osf = false; // previous block's last byte belonged to a token-opening "other" char + let mut anl: Option = None; + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = build_block::(text, tags, base, len, &LUT, code, false); + let c = decode(b.p0, b.p1, b.p2, valid); + + let pb = if base == 0 { 0 } else { code_bits(code) }; + let (nb, nb_lead) = if last_blk { + (0u8, true) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + ( + if is_lead { + code_bits(LUT[tags[q] as usize]) + } else { + code_bits(last_code) + }, + is_lead, + ) + }; + let has = |v: u8, s: u8| v & s != 0; + let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); + let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + + // ── run starts ───────────────────────────────────────────────────────────────────────── + let o_start = c.other & lead & !p1(c.other, has(pb, C_O)) & !p1(c.sp, has(pb, C_SP)); // ` ?[^\s\p{L}\p{N}]+` + let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); + // `\s+(?!\S)`: the run's last char opens a token — but not a newline, which `\s*[\r\n]+` + // has already swallowed. + let (steal, steal_patch) = to_lead( + c.ws & !c.nl & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), + b.cont, + prev_cont, + ); + // `\s*[\r\n]+` runs through the run's LAST newline, so a token opens right after it unless + // a further newline still follows inside the run (backward scan → reversal). + let after_nl = p1(c.nl, has(pb, C_NL)) + & c.ws + & lead + & !fill_to_last(c.nl.reverse_bits(), c.ws.reverse_bits()).reverse_bits(); + + // ── `[^\r\n\p{L}\p{N}]?\p{L}+`: the prefix is ANY non-newline non-letter non-digit char, + // not just a space — so a punctuation char that opens a token is swallowed by a following + // letter run (`x!abc` → `x`, `!abc`), while one in mid-run is not (`x!!abc` → `x`, `!!`, + // `abc`), because the greedy other-run already owns it. + // + // Stated backward ("my predecessor opened a token") rather than forward ("my successor is a + // letter"): the forward form needs an `adv`, which silently drops the marker when the two + // chars straddle a block edge. Smearing `o_start` across its char's bytes makes the test a + // plain `p1`, and the smear's only cross-block state is a shift carry. + let mut osf = o_start; + osf |= (osf << 1) & b.cont; + osf |= (osf << 2) & b.cont & (b.cont << 1); + if prev_osf { + osf |= lead_run(b.cont, valid); // a char whose lead sat in the previous block + } + let l_start = c.l + & lead + & !p1(c.l, has(pb, C_L)) + & !p1(c.ws & !c.nl, has(pb, C_WS) && !has(pb, C_NL)) + & !p1(osf, prev_osf); + + // ── `\p{N}{1,3}`: a group boundary every 3 chars from the run start. + let mut m = c.n & lead & !p1(c.n, has(pb, C_N)); + if dig_run && has(pb, C_N) { + let mut s = lead & lead.wrapping_neg() & c.n; + for _ in 0..((3 - dig_since % 3) % 3) { + s = adv(s, b.cont) & c.n & lead; + } + m |= s; + } + let mut groups = m; + if c.n & b.cont == 0 { + let n3 = c.n & (c.n << 1) & (c.n << 2); + while m != 0 { + m = (m << 3) & n3; + groups |= m; + } + } else { + while m != 0 { + let a = adv(m, b.cont) & c.n & lead; + let e = adv(adv(a, b.cont) & c.n & lead, b.cont) & c.n & lead; + if e == 0 { + break; + } + groups |= e; + m = e; + } + } + + // ── the other-run's `[\r\n]*` tail swallows the newlines directly behind it. + let nl_m = ((p1(c.other, has(pb, C_O)) & c.nl & lead) as u128) | u128::from(nl_run); + let nl_e = scanthru(nl_m, c.nl as u128); + let nl_span = nl_e.wrapping_sub(nl_m); + + // ── the one backward-in-time dependency (see deepseek): a newline arriving now retracts an + // "after the last newline" start committed for a run that was still open at the last edge. + if let Some(p) = anl + && has(pb, C_WS) + && c.ws & 1 != 0 + && c.nl & lead_run(c.ws, valid) != 0 + { + starts[p / 64] &= !(1u64 << (p % 64)); + anl = None; + } + + let mut st = groups | l_start | o_start | ws_start | after_nl | steal; + st &= !(nl_span as u64); + st |= nl_e as u64; + st &= lead; + if bi == 0 { + st |= 1; + } + starts[bi] = st; + flag[bi] = st & c.apo; + if bi > 0 { + starts[bi - 1] |= steal_patch; + } + + // ── carries ──────────────────────────────────────────────────────────────────────────── + nl_run = nl_e >> 64 != 0; + let tn = trail_run(c.n, valid, len); + dig_run = tn != 0 && has(nb, C_N); + dig_since = if !dig_run { + 0 + } else { + let g = groups & tn; + let counted = if g == 0 { + dig_since + (c.n & lead & tn).count_ones() + } else { + (c.n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() + }; + counted % 3 + }; + let tws = trail_run(c.ws, valid, len); + if tws != 0 && has(nb, C_WS) { + let a = after_nl & tws & !(nl_e as u64); + if a != 0 { + anl = Some(base + 63 - a.leading_zeros() as usize); + } else if !(tws & 1 != 0 && has(pb, C_WS)) { + anl = None; + } + } else { + anl = None; + } + prev_osf = osf >> (len - 1) & 1 != 0; + code = last_code; + prev_cont = b.cont; + } + emit_contr(text, starts, flag, nblk, ntext, true, out) +} diff --git a/tokenizers/bitsplit/src/lib.rs b/tokenizers/bitsplit/src/lib.rs new file mode 100644 index 000000000..79d619b1e --- /dev/null +++ b/tokenizers/bitsplit/src/lib.rs @@ -0,0 +1,373 @@ +//! `bitsplit` — GPT-family pre-tokenization as a **bitstream program**, replacing the scalar FSMs. +//! +//! Follows *Interleaved Bitstream Execution for Multi-Pattern Regex Matching on GPUs* +//! (MICRO'25, doi 10.1145/3725843.3756052). The paper's two ideas that carry over to a CPU: +//! +//! 1. **Bit-parallel regex.** Compile the grammar into character-class *bitstreams* (one bit per +//! input byte) plus boolean ops and carry-propagating adds. 64 input bytes are decided per +//! 64-bit register op, branchlessly — the FSM's per-token unpredictable branch disappears. +//! 2. **Interleaved execution.** Do NOT run one loop per bitstream instruction over the whole +//! input (that is the "sequential" baseline the paper beats: every intermediate stream is +//! materialised and re-read). Instead fuse *all* instructions into ONE block-wise loop, so an +//! intermediate stream lives in a register for the ~100 ops it is needed and dies there. The +//! only thing that reaches memory is the final `starts` bitmap (n/8 bytes). +//! +//! The paper's third contribution (dependency-aware thread-data mapping) is a GPU concern — it +//! resolves cross-block dependencies by recomputing them on other SMs. On one core the blocks are +//! visited in order, so those dependencies are carried in a handful of scalar registers and, for +//! the one genuinely backward-in-time rule, patched into the already-written bitmap. +//! +//! We do not re-derive character classes: `atomsplit::classify` already emits one `Atom` tag per +//! byte. The builder folds those 16 atoms into a grammar-specific **dense 3-bit code** and extracts +//! 3 bit-planes of it; every class stream is then a 2–3 op boolean function of the planes. +//! +//! Grammars: [`bitsplit_deepseek`], [`bitsplit_byte_level`] (GPT-2), [`bitsplit_cl100k`]. All three +//! byte-exact with their `atomsplit::fsm` counterparts — see `src/bin/verify.rs`. + +pub(crate) use atomsplit::fsm::Span; + +pub mod deepseek; +pub mod gpt; +#[cfg(target_arch = "aarch64")] +mod simd; +#[cfg(target_arch = "x86_64")] +mod simd_x86; + +pub use deepseek::bitsplit_deepseek; +pub use gpt::{bitsplit_byte_level, bitsplit_cl100k}; + +/// Whether this target builds its bitstreams with SIMD. Without it the builder is the portable +/// byte-at-a-time reference, which is slower than the FSM this replaces -- so a caller should keep +/// its FSM path rather than route here. On x86 the kernel needs SSSE3, so this is a runtime check. +#[must_use] +pub fn fast_builder() -> bool { + #[cfg(target_arch = "aarch64")] + { + true + } + #[cfg(target_arch = "x86_64")] + { + has_ssse3() + } + #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))] + { + false + } +} + +#[cfg(target_arch = "x86_64")] +fn has_ssse3() -> bool { + use std::sync::atomic::{AtomicU8, Ordering}; + static CACHED: AtomicU8 = AtomicU8::new(0); + match CACHED.load(Ordering::Relaxed) { + 0 => { + let yes = std::arch::is_x86_feature_detected!("ssse3"); + CACHED.store(1 + u8::from(yes), Ordering::Relaxed); + yes + } + n => n == 2, + } +} + +pub(crate) const CONT: u8 = 15; // Atom::Cont +pub(crate) const CODE_CONT: u8 = 7; // every grammar's dense code for a continuation byte + +/// The bitstreams for one 64-byte block. `p0`/`p1`/`p2` are the bit-planes of the **filled** dense +/// code — filled meaning a multi-byte char sets its bits on *all* of its bytes, so "previous char's +/// class" is a plain `<< 1` and no rule does char-width arithmetic. `cont` is the one un-filled +/// stream (it defines the fill); `cjk` is text-derived and only built when a grammar asks for it. +#[derive(Default, Clone, Copy)] +pub(crate) struct Blk { + pub cont: u64, + pub p0: u64, + pub p1: u64, + pub p2: u64, + pub cjk: u64, +} + +/// deepseek Split-2's isolated range: Han U+4E00..9FA5 ∪ Hiragana/Katakana U+3040..30FF (all +/// 3-byte, leads E3..E9). Same predicate as `fsm_deepseek`'s `ds_is_cjk_at`. +#[inline] +pub(crate) fn is_cjk_at(text: &[u8], p: usize) -> bool { + let b = text[p]; + if !(0xE3..=0xE9).contains(&b) || p + 2 >= text.len() { + return false; + } + let cp = ((b as u32 & 0x0F) << 12) + | ((text[p + 1] as u32 & 0x3F) << 6) + | (text[p + 2] as u32 & 0x3F); + (0x4E00..=0x9FA5).contains(&cp) || (0x3040..=0x30FF).contains(&cp) +} + +/// Build one block. Full blocks go through the NEON kernel; the ragged tail (and every other +/// target) uses the portable byte-at-a-time reference. Both produce the identical `Blk`. +/// `CJK` asks for the (deepseek-only) range stream; `lut` is the grammar's tag → dense code table. +#[inline] +pub(crate) fn build_block( + text: &[u8], + tags: &[u8], + base: usize, + len: usize, + lut: &[u8; 64], + cur_code: u8, + cur_cjk: bool, +) -> (Blk, u8) { + #[cfg(target_arch = "aarch64")] + if len == 64 { + // SAFETY: `len == 64` means `base + 64 <= text.len() == tags.len()`. + return unsafe { crate::simd::build64::(text, tags, base, lut, cur_code, cur_cjk) }; + } + #[cfg(target_arch = "x86_64")] + if len == 64 && has_ssse3() { + // SAFETY: `len == 64` bounds both reads, and SSSE3 is checked above. + return unsafe { + crate::simd_x86::build64::(text, tags, base, lut, cur_code, cur_cjk) + }; + } + build_block_scalar::(text, tags, base, len, lut, cur_code, cur_cjk) +} + +/// Portable reference builder: one byte at a time. +pub(crate) fn build_block_scalar( + text: &[u8], + tags: &[u8], + base: usize, + len: usize, + lut: &[u8; 64], + mut cur_code: u8, + mut cur_cjk: bool, +) -> (Blk, u8) { + let mut b = Blk::default(); + for i in 0..len { + let p = base + i; + let bit = 1u64 << i; + if tags[p] == CONT { + b.cont |= bit; + } else { + cur_code = lut[tags[p] as usize]; + cur_cjk = CJK && is_cjk_at(text, p); + } + b.p0 |= bit * u64::from(cur_code & 1 != 0); + b.p1 |= bit * u64::from(cur_code & 2 != 0); + b.p2 |= bit * u64::from(cur_code & 4 != 0); + b.cjk |= bit * u64::from(cur_cjk); + } + (b, cur_code) +} + +// ── bitstream primitives ──────────────────────────────────────────────────────────────────────── + +/// `ScanThru`: move every marker in `m` forward past the run of `c` it sits in, landing on the +/// first position not in `c`. The classic Parabix carry-propagation trick — one add. Done in `u128` +/// so a run reaching bit 63 puts its landing bit at 64 instead of vanishing: that bit *is* the +/// block's carry-out, and `e - m` still yields the right in-block span. +#[inline] +pub(crate) const fn scanthru(m: u128, c: u128) -> u128 { + m.wrapping_add(c) & !c +} + +/// Advance markers by ONE CHAR: shift one byte, then scan through the continuation bytes. Markers +/// that leave the block are dropped — their state is reconstructed from the scalar carry instead. +#[inline] +pub(crate) const fn adv(m: u64, cont: u64) -> u64 { + scanthru((m as u128) << 1, cont as u128) as u64 +} + +/// Move each bit back to the lead byte of its char. Rules that look at the *next* char are stated +/// at a char's last byte; token starts must sit on leads. ≤3 steps (UTF-8 chars are ≤4 bytes). +/// Returns `(in-block, patch for the previous word)` — a char straddling the block boundary lands +/// its lead in the word we already wrote. +#[inline] +pub(crate) fn to_lead(x: u64, cont: u64, prev_cont: u64) -> (u64, u64) { + let c = ((cont as u128) << 64) | prev_cont as u128; + let mut y = (x as u128) << 64; + for _ in 0..3 { + let m = y & c; + if m == 0 { + break; + } + y = (y & !c) | (m >> 1); + } + ((y >> 64) as u64, y as u64) +} + +/// In each run of `c`, fill from the run start through the LAST marker of `m` (`m ⊆ c`). Used on +/// reversed streams, where it answers "is there a newline at-or-after me, still inside this +/// whitespace run?". `(end - m) | m` rather than `end - m`: the latter only spans from the *first* +/// marker when a run holds several. +#[inline] +pub(crate) fn fill_to_last(m: u64, c: u64) -> u64 { + if m == 0 { + return 0; + } + let (m128, c128) = (m as u128, c as u128); + (scanthru(m128, c128).wrapping_sub(m128) | m128) as u64 +} + +/// Run of `x` that starts at bit 0. +#[inline] +pub(crate) fn lead_run(x: u64, valid: u64) -> u64 { + let z = !x & valid; + if z == 0 { + valid + } else { + (z & z.wrapping_neg()) - 1 + } +} + +/// Run of `x` that ends at bit `len - 1` (0 if that bit is clear). +#[inline] +pub(crate) fn trail_run(x: u64, valid: u64, len: usize) -> u64 { + if x & (1u64 << (len - 1)) == 0 { + return 0; + } + let z = !x & valid; + if z == 0 { + valid + } else { + valid & !((1u64 << (64 - z.leading_zeros())) - 1) + } +} + +// ── emit ──────────────────────────────────────────────────────────────────────────────────────── + +/// `starts` bitmap → spans. Each set bit closes the previous token and opens the next; `tzcnt` +/// walks them at ~3 ops per token. +pub(crate) fn emit(starts: &[u64], nblk: usize, n: usize, out: &mut [Span]) -> usize { + let (mut w, mut open) = (0usize, u32::MAX); + for (bi, &word) in starts.iter().enumerate().take(nblk) { + let mut m = word; + while m != 0 { + let pos = (bi * 64 + m.trailing_zeros() as usize) as u32; + if open != u32::MAX { + out[w] = Span::new(open, pos); + w += 1; + } + open = pos; + m &= m - 1; + } + } + if open != u32::MAX { + out[w] = Span::new(open, n as u32); + w += 1; + } + w +} + +/// Emit with a **scalar escape at flagged bits**: contractions (`'s 't 're 've 'm 'll 'd`) are +/// variable-length, case-optional and outrank every other alternative, which makes them miserable +/// in bit algebra and trivial here. `flag` marks the apostrophes that open a token; a block with +/// none takes the plain loop, so the escape costs one test per block on ordinary text. +/// +/// A matched contraction overrides the algebra outright: it emits its own span and skips every +/// start bit inside it, which is how the letter alternative loses the tie (`'sx` → `'s`, `x`). +pub(crate) fn emit_contr( + text: &[u8], + starts: &[u64], + flag: &[u64], + nblk: usize, + n: usize, + ci: bool, + out: &mut [Span], +) -> usize { + let (mut w, mut open, mut skip) = (0usize, u32::MAX, 0usize); + for bi in 0..nblk { + let mut m = starts[bi]; + let f = flag[bi]; + if f == 0 && skip <= bi * 64 { + while m != 0 { + let pos = (bi * 64 + m.trailing_zeros() as usize) as u32; + if open != u32::MAX { + out[w] = Span::new(open, pos); + w += 1; + } + open = pos; + m &= m - 1; + } + continue; + } + while m != 0 { + let j = m.trailing_zeros() as usize; + let pos = bi * 64 + j; + m &= m - 1; + if pos < skip { + continue; + } + if open != u32::MAX { + out[w] = Span::new(open, pos as u32); + w += 1; + } + open = pos as u32; + if f >> j & 1 != 0 { + // Contractions chain (`'re've`, `y'all'd've`): the char after one is a token start + // in its own right, so keep matching until one fails rather than handing control + // back to the bit algebra — whose start bit there we are about to skip. + let mut p = pos; + while p < n { + let l = contr_len(text, p, ci); + if l == 0 { + break; + } + out[w] = Span::new(p as u32, (p + l) as u32); + w += 1; + p += l; + } + if p > pos { + open = if p < n { p as u32 } else { u32::MAX }; + // `+ 1`: the algebra usually also has a start bit exactly at `p` (the letter + // run resumes there) and we have just opened it — consuming it again would + // emit an empty span. `contr_len(p) == 0` here, so nothing is lost. + skip = p + 1; + } + } + } + } + if open != u32::MAX { + out[w] = Span::new(open, n as u32); + w += 1; + } + w +} + +/// Byte length of the contraction at `i` (2 or 3), or 0. `ci` picks cl100k/o200k's `(?i:)` form +/// over GPT-2's case-sensitive one. +#[inline] +fn contr_len(text: &[u8], i: usize, ci: bool) -> usize { + let n = text.len(); + if i + 1 >= n || text[i] != b'\'' || text[i + 1] >= 0x80 { + return 0; + } + let c1 = if ci { text[i + 1] | 0x20 } else { text[i + 1] }; + match c1 { + b's' | b't' | b'm' | b'd' => 2, + b'r' | b'v' | b'l' if i + 2 < n && text[i + 2] < 0x80 => { + let c2 = if ci { text[i + 2] | 0x20 } else { text[i + 2] }; + usize::from((matches!(c1, b'r' | b'v') && c2 == b'e') || (c1 == b'l' && c2 == b'l')) * 3 + } + _ => 0, + } +} + +/// Builder-only cost probe: runs just the byte→bitstream transpose over every block and folds the +/// streams so nothing is optimised away. The difference against a full grammar is what the +/// bitstream program itself costs. +#[doc(hidden)] +#[must_use] +pub fn build_only(text: &[u8], tags: &[u8]) -> u64 { + let (mut acc, mut code, mut cjk) = (0u64, CODE_CONT, false); + for base in (0..text.len()).step_by(64) { + let len = (text.len() - base).min(64); + let (b, c) = build_block::(text, tags, base, len, &deepseek::LUT, code, cjk); + code = c; + cjk = b.cjk >> (len - 1) & 1 != 0; + acc ^= b.cont ^ b.p0 ^ b.p1 ^ b.p2 ^ b.cjk; + } + acc +} + +/// Convenience wrapper: classify + deepseek bitsplit over caller-owned scratch. +#[must_use] +pub fn pre_tokenize(text: &[u8], tags: &mut [u8], starts: &mut [u64], out: &mut [Span]) -> usize { + atomsplit::classify::classify(text, tags); + bitsplit_deepseek(text, tags, starts, out) +} diff --git a/tokenizers/bitsplit/src/simd.rs b/tokenizers/bitsplit/src/simd.rs new file mode 100644 index 000000000..d328dc502 --- /dev/null +++ b/tokenizers/bitsplit/src/simd.rs @@ -0,0 +1,176 @@ +//! NEON block builder: 64 tag bytes + 64 text bytes → the 10 class bitstreams of one [`Blk`]. +//! +//! This is the step the paper spends its GPU budget on too — turning bytes into bitstreams. The +//! shape here is a 64×8 bit-matrix transpose done 8 lanes at a time: weight each compare result by +//! its lane's power of two (`POW`) and let three `vpaddq_u8` rounds fold 64 lanes into the 8 bytes +//! of one `u64`. That is ~9 ops per stream per 64 bytes, versus 4 separate 16-bit movemasks. +//! +//! Continuation bytes are resolved **before** extraction (≤3 `vext`+`vbsl`, the same trick +//! `atomsplit::simd_fsm` uses), so every stream comes out *filled* — a multi-byte char sets its bit +//! on all of its bytes. That is what lets the bitstream program read "previous char's class" as a +//! plain `<< 1` with no char-width arithmetic. + +use crate::{Blk, lead_run}; +use core::arch::aarch64::*; + +const POW: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128]; + +/// 64 lanes of `0x00`/`0xFF` → one `u64`. Three `vpaddq_u8` rounds: lane pairs, then quads, then +/// octets — after which the low 8 bytes of the result are the 8 bytes of the mask, in order. +#[inline(always)] +unsafe fn mm64(v: [uint8x16_t; 4], pow: uint8x16_t) -> u64 { + unsafe { + let ab = vpaddq_u8(vandq_u8(v[0], pow), vandq_u8(v[1], pow)); + let cd = vpaddq_u8(vandq_u8(v[2], pow), vandq_u8(v[3], pow)); + let x = vpaddq_u8(ab, cd); + vgetq_lane_u64::<0>(vreinterpretq_u64_u8(vpaddq_u8(x, x))) + } +} + +/// Build one **full** 64-byte block, folding tags through `lut` into dense codes. +/// Build one **full** 64-byte block. `cur_code` / `cur_cjk` describe the byte just before it, so a +/// block opening mid-char keeps inheriting its lead's class. Returns the block's last filled code. +/// +/// # Safety +/// `base + 64 <= tags.len()` and `base + 64 <= text.len()`. +#[target_feature(enable = "neon")] +pub(crate) unsafe fn build64( + text: &[u8], + tags: &[u8], + base: usize, + lut: &[u8; 64], + cur_code: u8, + cur_cjk: bool, +) -> (Blk, u8) { + unsafe { + let pow = vld1q_u8(POW.as_ptr()); + let seven = vdupq_n_u8(7); + let tbl = uint8x16x4_t( + vld1q_u8(lut.as_ptr()), + vld1q_u8(lut.as_ptr().add(16)), + vld1q_u8(lut.as_ptr().add(32)), + vld1q_u8(lut.as_ptr().add(48)), + ); + + // ── tags → dense codes, then fill the continuation lanes from the left so every lane + // carries its char's code. Two steps suffice, not three: after shifting by 1 every lane is + // correct at distance 1, so shifting *that* by 2 covers distances 2 and 3 — which is the + // most a 4-byte char can need. + let mut isc = [vdupq_n_u8(0); 4]; + let mut cd = [vdupq_n_u8(0); 4]; + let mut prev = vdupq_n_u8(cur_code); + for k in 0..4 { + let raw = vqtbl4q_u8(tbl, vld1q_u8(tags.as_ptr().add(base + k * 16))); + isc[k] = vceqq_u8(raw, seven); + let c = vbslq_u8(isc[k], vextq_u8::<15>(prev, raw), raw); + let c = vbslq_u8(vceqq_u8(c, seven), vextq_u8::<14>(prev, c), c); + prev = c; + cd[k] = c; + } + let last_code = vgetq_lane_u8::<15>(prev); + // ── 3 bit-planes of the code. Every stream is a boolean function of these (`decode`), so + // this is 3 extractions where one-hot class bits needed 6 — and `cont` comes from `isc`, + // which the fill needed anyway. + let plane = |bit: u8| { + let d = vdupq_n_u8(bit); + mm64( + [ + vtstq_u8(cd[0], d), + vtstq_u8(cd[1], d), + vtstq_u8(cd[2], d), + vtstq_u8(cd[3], d), + ], + pow, + ) + }; + + let mut b = Blk { + cont: mm64(isc, pow), + p0: plane(1), + p1: plane(2), + p2: plane(4), + cjk: 0, + }; + if !CJK { + return (b, last_code); + } + // ── text is loaded only for the (deepseek-only) CJK range test. + let ntext = text.len(); + let tv = [ + vld1q_u8(text.as_ptr().add(base)), + vld1q_u8(text.as_ptr().add(base + 16)), + vld1q_u8(text.as_ptr().add(base + 32)), + vld1q_u8(text.as_ptr().add(base + 48)), + ]; + + // ── the CJK range test lives in the raw bytes, not the tags. It is a 3-byte predicate, so + // it runs in vector space on `vext`-aligned b1/b2 and extracts ONE mask — testing the 9 + // byte predicates as separate bitstreams costs 9 extractions for the same answer. A + // `vmaxvq` over the lead range gates the whole thing away on Latin/code for ~8 ops. + let e3e9 = |v| vcleq_u8(vsubq_u8(v, vdupq_n_u8(0xE3)), vdupq_n_u8(6)); + let any = vmaxvq_u8(vorrq_u8( + vorrq_u8(e3e9(tv[0]), e3e9(tv[1])), + vorrq_u8(e3e9(tv[2]), e3e9(tv[3])), + )); + if any != 0 { + // b1/b2 of a char at lane 15 come from the next chunk — and for the last chunk, from + // the next block, which may not exist. + let tail = { + let mut buf = [0u8; 16]; + let off = base + 64; + let avail = text.len().saturating_sub(off).min(16); + buf[..avail].copy_from_slice(&text[off..off + avail]); + vld1q_u8(buf.as_ptr()) + }; + // Hiragana/Katakana U+3040..30FF is exactly E3 [81-83] xx; Han U+4E00..9FA5 is + // E4 [B8-BF] xx / E5-E8 xx xx / E9 [80-BD] xx / E9 BE [80-A5]. + let cjkv = |v: uint8x16_t, nx: uint8x16_t| { + let (b1, b2) = (vextq_u8::<1>(v, nx), vextq_u8::<2>(v, nx)); + let e9 = vandq_u8( + vceqq_u8(v, vdupq_n_u8(0xE9)), + vorrq_u8( + vcltq_u8(b1, vdupq_n_u8(0xBE)), + vandq_u8( + vceqq_u8(b1, vdupq_n_u8(0xBE)), + vcleq_u8(b2, vdupq_n_u8(0xA5)), + ), + ), + ); + vorrq_u8( + vorrq_u8( + vandq_u8( + vceqq_u8(v, vdupq_n_u8(0xE3)), + vcleq_u8(vsubq_u8(b1, vdupq_n_u8(0x81)), vdupq_n_u8(2)), + ), + vandq_u8( + vceqq_u8(v, vdupq_n_u8(0xE4)), + vcgeq_u8(b1, vdupq_n_u8(0xB8)), + ), + ), + vorrq_u8(vcleq_u8(vsubq_u8(v, vdupq_n_u8(0xE5)), vdupq_n_u8(3)), e9), + ) + }; + let mut leads = mm64( + [ + cjkv(tv[0], tv[1]), + cjkv(tv[1], tv[2]), + cjkv(tv[2], tv[3]), + cjkv(tv[3], tail), + ], + pow, + ); + // `fsm_deepseek` reads 3 bytes unconditionally; refuse to classify a truncated tail. + let lim = ntext.saturating_sub(base + 2); + if lim < 64 { + leads &= (1u64 << lim) - 1; + } + // every CJK char is 3 bytes → fill by two shifts; a char cut by the block edge is + // picked up on the other side by `cur_cjk` (its continuation bytes lead that block). + b.cjk = leads | (leads << 1) | (leads << 2); + } + if cur_cjk { + b.cjk |= lead_run(b.cont, !0); + } + (b, last_code) + } +} diff --git a/tokenizers/bitsplit/src/simd_x86.rs b/tokenizers/bitsplit/src/simd_x86.rs new file mode 100644 index 000000000..37bb29838 --- /dev/null +++ b/tokenizers/bitsplit/src/simd_x86.rs @@ -0,0 +1,192 @@ +//! SSSE3 block builder: the x86 twin of `simd.rs`. Same shape — tags through a LUT into dense +//! codes, continuation lanes filled from the left, three bit-planes extracted — with two +//! differences that fall out of the ISA. +//! +//! Plane extraction is cheaper here: `movemask` is native, so a plane is +//! `movemask(slli_epi16(code, 7 - bit))`. Shifting 16-bit lanes puts bit `k` of the low byte at +//! bit 7 and bit `k` of the high byte at bit 15, which is exactly the pair `movemask` reads; the +//! bits that bleed across the byte boundary land where it does not look. +//! +//! The LUT is dearer: `pshufb` indexes 16 entries and the tags run to 0x26, so the table is split +//! into one shuffle per high nibble and selected between. Refinements (`0x10`/`0x20` case, `0x16` +//! AlphaSymMark, `0x26` ZWJ) are the only tags above 0x0F, so three shuffles cover every case. +//! +//! Checked the same way as the NEON kernel: `src/bin/verify.rs` built for `x86_64-apple-darwin` +//! and run under Rosetta, which reports SSSE3, so this path is the one that executes. The scalar +//! builder in `lib.rs` stays the reference it is compared against. + +use crate::{Blk, lead_run}; +use core::arch::x86_64::*; + +/// `x <= k`, unsigned, in the absence of an unsigned byte compare. +#[inline(always)] +unsafe fn le(x: __m128i, k: u8) -> __m128i { + unsafe { _mm_cmpeq_epi8(_mm_min_epu8(x, _mm_set1_epi8(k as i8)), x) } +} + +/// `x >= k`, unsigned. +#[inline(always)] +unsafe fn ge(x: __m128i, k: u8) -> __m128i { + unsafe { _mm_cmpeq_epi8(_mm_max_epu8(x, _mm_set1_epi8(k as i8)), x) } +} + +/// `x - lo <= n`, unsigned, i.e. `x` in `lo ..= lo + n`. +#[inline(always)] +unsafe fn in_range(x: __m128i, lo: u8, n: u8) -> __m128i { + unsafe { le(_mm_sub_epi8(x, _mm_set1_epi8(lo as i8)), n) } +} + +/// `if mask { a } else { b }`, lane-wise. SSE4.1's `blendv` would do it in one, but the and/andnot +/// pair keeps the whole builder at SSSE3. +#[inline(always)] +unsafe fn sel(mask: __m128i, a: __m128i, b: __m128i) -> __m128i { + unsafe { _mm_or_si128(_mm_and_si128(mask, a), _mm_andnot_si128(mask, b)) } +} + +/// Build one **full** 64-byte block. `cur_code`/`cur_cjk` describe the byte before it. +/// +/// # Safety +/// `base + 64 <= tags.len()` and `base + 64 <= text.len()`; the caller has checked for SSSE3. +#[target_feature(enable = "ssse3")] +pub(crate) unsafe fn build64( + text: &[u8], + tags: &[u8], + base: usize, + lut: &[u8; 64], + cur_code: u8, + cur_cjk: bool, +) -> (Blk, u8) { + unsafe { + let low_nibble = _mm_set1_epi8(0x0F); + let seven = _mm_set1_epi8(7); + // one table per high nibble; a tag of 0x30 or above never occurs + let t0 = _mm_loadu_si128(lut.as_ptr().cast()); + let t1 = _mm_loadu_si128(lut.as_ptr().add(16).cast()); + let t2 = _mm_loadu_si128(lut.as_ptr().add(32).cast()); + + let mut cd = [_mm_setzero_si128(); 4]; + let mut isc = [_mm_setzero_si128(); 4]; + let mut prev = _mm_set1_epi8(cur_code as i8); + for k in 0..4 { + let t = _mm_loadu_si128(tags.as_ptr().add(base + k * 16).cast()); + let lo = _mm_and_si128(t, low_nibble); + let hi = _mm_and_si128(_mm_srli_epi16(t, 4), low_nibble); + let raw = sel( + _mm_cmpeq_epi8(hi, _mm_set1_epi8(2)), + _mm_shuffle_epi8(t2, lo), + sel( + _mm_cmpeq_epi8(hi, _mm_set1_epi8(1)), + _mm_shuffle_epi8(t1, lo), + _mm_shuffle_epi8(t0, lo), + ), + ); + isc[k] = _mm_cmpeq_epi8(raw, seven); + // fill the continuation lanes: shift by one, then shift *that* by two, which covers + // the three continuations a 4-byte char can have. + let c = sel(isc[k], _mm_alignr_epi8(raw, prev, 15), raw); + let c = sel(_mm_cmpeq_epi8(c, seven), _mm_alignr_epi8(c, prev, 14), c); + prev = c; + cd[k] = c; + } + let last_code = (_mm_extract_epi16(prev, 7) >> 8) as u8; + + // 64 lanes -> one u64, four native movemasks + let gather = |v: [__m128i; 4]| -> u64 { + let mut m = 0u64; + for (k, chunk) in v.iter().enumerate() { + m |= (_mm_movemask_epi8(*chunk) as u16 as u64) << (16 * k); + } + m + }; + // the shift is a const generic, so one arm per plane: bit k wants `7 - k` + macro_rules! plane { + ($shift:literal) => { + gather([ + _mm_slli_epi16::<$shift>(cd[0]), + _mm_slli_epi16::<$shift>(cd[1]), + _mm_slli_epi16::<$shift>(cd[2]), + _mm_slli_epi16::<$shift>(cd[3]), + ]) + }; + } + + let mut b = Blk { + cont: gather(isc), + p0: plane!(7), + p1: plane!(6), + p2: plane!(5), + cjk: 0, + }; + if !CJK { + return (b, last_code); + } + + // ── the CJK range test, on the raw bytes: Hiragana/Katakana U+3040..30FF is E3 [81-83] xx; + // Han U+4E00..9FA5 is E4 [B8-BF] xx / E5-E8 xx xx / E9 [80-BD] xx / E9 BE [80-A5]. + let ntext = text.len(); + let tv = [ + _mm_loadu_si128(text.as_ptr().add(base).cast()), + _mm_loadu_si128(text.as_ptr().add(base + 16).cast()), + _mm_loadu_si128(text.as_ptr().add(base + 32).cast()), + _mm_loadu_si128(text.as_ptr().add(base + 48).cast()), + ]; + let any = tv + .iter() + .any(|v| _mm_movemask_epi8(in_range(*v, 0xE3, 6)) != 0); + if !any { + if cur_cjk { + b.cjk |= lead_run(b.cont, !0); + } + return (b, last_code); + } + let tail = { + let mut buf = [0u8; 16]; + let off = base + 64; + let avail = ntext.saturating_sub(off).min(16); + buf[..avail].copy_from_slice(&text[off..off + avail]); + _mm_loadu_si128(buf.as_ptr().cast()) + }; + let cjkv = |v: __m128i, nx: __m128i| -> __m128i { + let b1 = _mm_alignr_epi8(nx, v, 1); + let b2 = _mm_alignr_epi8(nx, v, 2); + let e9 = _mm_and_si128( + _mm_cmpeq_epi8(v, _mm_set1_epi8(0xE9u8 as i8)), + _mm_or_si128( + le(b1, 0xBD), + _mm_and_si128( + _mm_cmpeq_epi8(b1, _mm_set1_epi8(0xBEu8 as i8)), + le(b2, 0xA5), + ), + ), + ); + _mm_or_si128( + _mm_or_si128( + _mm_and_si128( + _mm_cmpeq_epi8(v, _mm_set1_epi8(0xE3u8 as i8)), + in_range(b1, 0x81, 2), + ), + _mm_and_si128(_mm_cmpeq_epi8(v, _mm_set1_epi8(0xE4u8 as i8)), ge(b1, 0xB8)), + ), + _mm_or_si128(in_range(v, 0xE5, 3), e9), + ) + }; + let mut leads = gather([ + cjkv(tv[0], tv[1]), + cjkv(tv[1], tv[2]), + cjkv(tv[2], tv[3]), + cjkv(tv[3], tail), + ]); + // `fsm_deepseek` reads three bytes unconditionally; refuse to classify a truncated tail. + let lim = ntext.saturating_sub(base + 2); + if lim < 64 { + leads &= (1u64 << lim) - 1; + } + // every CJK char is 3 bytes, so two shifts fill it; one cut by the block edge is picked up + // on the other side by `cur_cjk`. + b.cjk = leads | (leads << 1) | (leads << 2); + if cur_cjk { + b.cjk |= lead_run(b.cont, !0); + } + (b, last_code) + } +} diff --git a/tokenizers/tk-encode/Cargo.toml b/tokenizers/tk-encode/Cargo.toml index 3e1b8868b..a7a3a37e6 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -28,6 +28,7 @@ path = "src/lib.rs" [dependencies] atomsplit = { path = "../atomsplit" } +bitsplit = { path = "../bitsplit" } rand = "0.9" regex = "1.10" rayon = { version = "1.10", optional = true } diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index 779c167f1..04a6b4c30 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -187,6 +187,29 @@ impl pipeline::PreTokenizer for Split { .fsm .filter(|_| !self.invert && self.behavior == SplitDelimiterBehavior::Isolated) { + // gpt2 and cl100k-with-the-standard-digit-cap have bitstream splitters, which decide + // 64 bytes per register op instead of one token per unpredictable branch -- but only + // where the bitstream build is SIMD. Its portable builder is slower than the FSM, so + // everything else keeps the FSM. + match fsm { + GptFsm::Gpt2 if bitsplit::fast_builder() => { + pipeline::classify_into_spans_bits( + text.as_bytes(), + bitsplit::bitsplit_byte_level, + out, + ); + return Ok(()); + } + GptFsm::Cl100k { digit_cap: 3 } if bitsplit::fast_builder() => { + pipeline::classify_into_spans_bits( + text.as_bytes(), + bitsplit::bitsplit_cl100k, + out, + ); + return Ok(()); + } + _ => {} + } pipeline::classify_into_spans( text.as_bytes(), |bytes, tags, spans| match fsm { diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index cc994ecc5..b38537606 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -64,6 +64,51 @@ pub(crate) fn classify_into_spans( }); } +/// [`classify_into_spans`] for a splitter that works off bitstreams: it needs two `u64` bitmaps +/// (token starts, and the flags its scalar escapes key on) alongside the tags. +pub(crate) fn classify_into_spans_bits( + bytes: &[u8], + split: impl FnOnce(&[u8], &[u8], &mut [u64], &mut [u64], &mut [Span]) -> usize, + out: &mut Vec, +) { + thread_local! { + static SCRATCH: RefCell<(Vec, Vec, Vec)> = + const { RefCell::new((Vec::new(), Vec::new(), Vec::new())) }; + } + let n = bytes.len(); + if n == 0 { + return; + } + SCRATCH.with(|cell| { + let (tags, starts, flags) = &mut *cell.borrow_mut(); + if tags.len() < n { + tags.resize(n, 0); + } + let words = n.div_ceil(64) + 1; + if starts.len() < words { + starts.resize(words, 0); + flags.resize(words, 0); + } + classify(bytes, &mut tags[..n]); + let base = out.len(); + out.reserve(n + 1); + // SAFETY: as in `classify_into_spans` -- `reserve` covers the splitter's worst case of one + // span per byte plus one, `Span` has no drop glue, and `set_len` counts only what was written. + let k = unsafe { + split( + bytes, + &tags[..n], + &mut starts[..words], + &mut flags[..words], + std::slice::from_raw_parts_mut(out.as_mut_ptr().add(base), n + 1), + ) + }; + debug_assert!(k <= n + 1); + // SAFETY: the splitter wrote `k <= n + 1` spans from `base`. + unsafe { out.set_len(base + k) }; + }); +} + pub trait Normalizer { fn normalize<'a>(&self, input: &'a str) -> Result>; } From 6fe801fd89d25c42b2bc0000244f0cc2a08e36cc Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 6 Aug 2026 17:06:12 +0900 Subject: [PATCH 2/7] perf(pipeline): give the model a whole chunk of pre-tokens at a time The pipeline has the entire span list before the model runs, so handing the spans over one at a time bought nothing and cost a virtual call, a slice, a `Result` and an output capacity check per pre-token -- on English, one round trip per ~5 bytes. `Model::tokenize_spans` takes them all. Its default is the loop it replaces, so a model only overrides it when it has per-chunk work to hoist; BPE does, and there the scratch destructuring and the output reservation move out of the loop. Slicing the chunk per pre-token also re-ran a UTF-8 boundary check on spans the pre-tokenizer had already cut on char boundaries. Ported onto #2241's `model.rs` rather than cherry-picked from #2304: that PR was written against the base before #2241 landed and restructured this file, so its patch no longer applies. Same change, current structure -- and the fold fast path it now sits beside came in with #2241. 367 tests pass. --- tokenizers/tk-encode/src/models/bpe/model.rs | 39 +++++++++++- .../tk-encode/src/tokenizer/pipeline.rs | 60 ++++++++++++++++--- 2 files changed, 90 insertions(+), 9 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index ee5af5669..18b03f837 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -8,7 +8,7 @@ use crate::models::bpe::legacy::model::BPE; use crate::models::bpe::merge_hot_cold_queue::{QueueScratch, merge_hot_cold_queue}; use crate::models::bpe::merge_multipass::merge_multipass; use crate::models::bpe::tables::BpeTables; -use crate::pipeline::{self, PipelineToken}; +use crate::pipeline::{self, PipelineToken, Span}; use crate::tokenizer::Result; use crate::utils::byte_level::{self}; use crate::vocab::bucket_vocab_store::BucketVocabStore; @@ -294,6 +294,43 @@ impl pipeline::Model for PipelineBPE { Ok(()) } + /// Every pre-token of a chunk in one call. + /// + /// Same work per word as [`Self::tokenize_pipeline`]; what changes is what is *not* repeated. + /// The scratch is destructured once instead of once per word, the output is grown once for the + /// whole batch instead of being capacity-checked on every push, and the virtual call, the + /// slice and the `Result` happen once per chunk rather than once per pre-token. + fn tokenize_spans( + &self, + chunk: &str, + spans: &[Span], + scratch: &mut Self::Scratch, + output: &mut Vec, + ) -> Result<()> { + let BpeScratch { symbols, queue } = scratch; + // One reservation for the batch. Most pre-tokens are a single token, so the span count is + // a close lower bound on what the batch emits; anything past it grows as usual. + output.reserve(spans.len()); + + for span in spans { + // SAFETY: the pre-tokenizer cuts on char boundaries, so a span is always a valid slice + // of this chunk. Bounds- and UTF-8-checking it again per word measured worth removing. + let sequence = unsafe { chunk.get_unchecked(span.range()) }; + if sequence.is_empty() { + continue; + } + if let Some(id) = self.fold_id(sequence) { + output.push(PipelineToken { id }); + continue; + } + self.merge_word(sequence, symbols, queue); + output.extend(symbols.iter().map(|&symbol| PipelineToken { + id: self.tables.unmap.at(symbol as usize), + })); + } + Ok(()) + } + fn init_scratch(&self) -> Self::Scratch { Self::Scratch { symbols: Vec::with_capacity(64), diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index b38537606..43876fa4a 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -1005,14 +1005,13 @@ impl PipelineTokenizer { self.pre_tokenizer .pre_tokenize(normalized_chunk, &mut pre_tokens)?; if STAGE >= Self::STAGE_MODEL { - // Tokenize each chunk - for pre_token in pre_tokens.iter() { - self.model.tokenize_pipeline( - &normalized_chunk[pre_token.range()], - &mut scratch, - &mut output, - )?; - } + // The whole span list at once; see Model::tokenize_spans. + self.model.tokenize_spans( + normalized_chunk, + &pre_tokens, + &mut scratch, + &mut output, + )?; } Ok(()) })?; @@ -1297,6 +1296,27 @@ pub trait Model { output: &mut Vec, ) -> Result<()>; + /// Every pre-token of a chunk in one call. + /// + /// The pipeline has the whole span list before the model runs, so handing them over one at a + /// time bought nothing and cost a virtual call, a slice, a `Result` and an output capacity + /// check per pre-token -- on English that is one round trip per ~5 bytes. + /// + /// The default is the loop it replaces, so a model only overrides this if it has per-chunk + /// work to hoist out of the loop. + fn tokenize_spans( + &self, + chunk: &str, + spans: &[Span], + scratch: &mut Self::Scratch, + output: &mut Vec, + ) -> Result<()> { + for span in spans { + self.tokenize_pipeline(&chunk[span.range()], scratch, output)?; + } + Ok(()) + } + fn init_scratch(&self) -> Self::Scratch; } @@ -1337,6 +1357,30 @@ impl Model for PipelineModel { } } + fn tokenize_spans( + &self, + chunk: &str, + spans: &[Span], + scratch: &mut Self::Scratch, + output: &mut Vec, + ) -> Result<()> { + match (self, scratch) { + (Self::BPE(model), PipelineModelScratch::BPE(scratch)) => { + model.tokenize_spans(chunk, spans, scratch, output) + } + (Self::Unigram(model), PipelineModelScratch::Unigram(scratch)) => { + model.tokenize_spans(chunk, spans, scratch, output) + } + (Self::WordLevel(model), PipelineModelScratch::WordLevel(scratch)) => { + model.tokenize_spans(chunk, spans, scratch, output) + } + (Self::WordPiece(model), PipelineModelScratch::WordPiece(scratch)) => { + model.tokenize_spans(chunk, spans, scratch, output) + } + _ => unreachable!(), + } + } + fn init_scratch(&self) -> Self::Scratch { match self { Self::BPE(bpe) => PipelineModelScratch::BPE(bpe.init_scratch()), From 28fbb60a49d87eaf09efbe4d87d8df6317c31792 Mon Sep 17 00:00:00 2001 From: Arthur <48595927+ArthurZucker@users.noreply.github.com> Date: Fri, 7 Aug 2026 07:46:45 +0200 Subject: [PATCH 3/7] bitsplit: one crate, all model grammars, atomsplit deleted (#2317) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(bitsplit): absorb classify/literal/regexes, add the byte-exactness gate bitsplit had no tests at all -- lib.rs pointed at a `src/bin/verify.rs` that is not in the tree. Add `tests/parity.rs` first, so the rest of the cleanup has a gate: oniguruma as the oracle (deepseek composed as HF applies it, three Isolated splits), plus the block-phase sweep that is the only thing that actually exercises the cross-block carries and `starts[bi-1] |= patch`. A `#[should_panic]` negative control keeps the gate honest. All three existing grammars pass. Then move the classify cluster (classify, atom_tables, tables, the three SIMD kernels), `regexes` and `literal` into bitsplit, and put `Span` in its lib.rs. atomsplit is now a shim over bitsplit holding only the scalar FSMs; it goes away once every caller is rewired. Also resolves a bad merge already on this branch: tk-encode/src/models/bpe/model.rs carried two `fn tokenize_spans`, so tk-encode did not compile at all. Kept the fold_id/symbols mechanics (the ones matching BpeScratch and merge_word's arity) and restored the word-cache logic the other side had, mirroring tokenize_pipeline. * feat(bitsplit): o200k family (llama4 / gpt-oss / minimax / tekken / kimi), 4th plane, aux streams Adds `o200k.rs`, const-generic over . AUX doubles as the variant selector because the two text-derived streams are mutually exclusive: the o200k/tekken line needs `/` for rule 4's `[\r\n/]*`, kimi needs `\p{Han}` for its leading arm and has a plain `[\r\n]*` tail. Kernel changes: - `Blk` gains `p3`. o200k needs 9 tag classes and code 7 is reserved -- both SIMD kernels find continuation lanes by testing `lut[tag] == 7` -- so "other" cannot be the leftover code. Const-gated off for the 3-plane grammars. - `Blk.cjk` generalised to `Blk.aux` with a const AUX selector (none/cjk/slash/han). - `digit_groups::` lifted into lib.rs: the `\p{N}{1,CAP}` block was duplicated verbatim in cl100k and deepseek, and hardcoding 3 is what blocked Qwen. The case split is a scalar escape, not bit algebra: `[UC]*[LC]+|[UC]+[LC]*` has no local form (`中Qz` is one token, `ʰABC` is two -- the difference is whether an L appears LATER in the run). The bit half computes a cheap gate instead, so all-lowercase and Capitalised text never pays for it. o200k/tekken/kimi parity is #[ignore]d with the reason: \p{M} is in BOTH the letter classes and rule 4's class, and that interaction is not modelled yet. The other four gates (gpt2, cl100k, deepseek, negative control) stay green. * fix(bpe): adopt #2314's tokenize_spans verbatim The duplicate-tokenize_spans repair I wrote landed upstream as #2314 on feat/train_encode_split, which is NOT an ancestor of poc/target-encode-clean -- hence the stale base. Take the upstream text so the two cannot diverge; it also documents why the fold must be probed before the cache. * refactor(bitsplit): one file per regex, unrolled Drop the const generics. Folding o200k, tekken and kimi into one parameterised grammar meant every fix was a three-way risk and nothing could be read on its own. Now: o200k.rs (llama4 / gpt-oss / minimax), tekken.rs, kimi.rs -- each unrolled against the shared primitives in lib.rs, each independently debuggable. The small scalar helpers (member, run_end, ws_tail, letter_match) move to lib.rs; they are genuinely shared and copying them per file would be worse. kimi is byte-exact. o200k/tekken still carry two cross-block bugs. * fix(bitsplit): o200k / tekken / kimi are byte-exact Three real bugs, all in the interaction between rule 4's `[\r\n/]*` tail and everything around it. `/` is in BOTH the `+` body and the tail, which cl100k never has to deal with (its tail is newlines, and a newline is not "other"): - One tail run can collect SEVERAL markers (`!\n/\n`: the `\n` after `!` and the `\n` after `/`). `nl_e - nl_m` spans only from the last one -- that is exactly what `fill_to_last` exists for. - A char the tail absorbed is not part of the `+` body, so an "other" after it opens a fresh run (`\r\n/#` = tail, then `#` starts again). Needs a `prev_absorbed` carry for the block edge. - ...but a char INSIDE the tail must not open one either, or `osf` picks up the `/` and the `[^\r\n\p{L}\p{N}]?` prefix rule then eats the letter start after it (`#\r\n/aA` came out as one token instead of `a` + `A`). Plus: re-flag the last letter token whenever the escape trigger fires, since the trigger can land in a later block than the token it belongs to. Gate: oniguruma oracle, block-phase sweep, 4000 fuzz strings (10000 run locally, clean). All 7 grammars green -- gpt2, cl100k, deepseek, o200k, tekken, kimi + negative control. * refactor(bitsplit): split gpt.rs into gpt2.rs + cl100k.rs, add Qwen One file per regex, finishing the split. cl100k's digit cap becomes a plain argument rather than a const generic -- 3 = cl100k / Llama-3 / GLM-4.6, 1 = Qwen. cl100k was still running its own inline hardcoded-3 digit block; it now shares digit_groups with o200k, which is what made the cap knob real. 8 gates green: gpt2, cl100k, qwen, deepseek, o200k, tekken, kimi + negative control. * refactor: delete atomsplit, tk-encode runs on bitsplit The regex FSMs are gone -- every recognized pattern now routes to the byte-exact bitstream grammar. `fast_builder()` went with them: there is no second engine left to route to, so the gate had nothing to gate. - `GptFsm` -> `Grammar`, one variant per distinct regex, with the Kimi key added. The cl100k digit cap now actually reaches the splitter (Qwen routes to `bitsplit_qwen`). - The class-run family and `CharDelimiterSplit` move to `bitsplit::classes`. - Workspace members drop atomsplit; bitmap_gen writes bitsplit/src/atom_tables.rs; the CI workflow is renamed. Whole workspace green: 8 bitsplit gates, 356 tk-encode unit tests, the bpe/decode/ pipeline oracles. * feat(bitsplit): literal search as a bitstream, absorbing CharDelimiterSplit Replaces memmem with the same shape as the grammars: a `u64` match bitmap per 64-byte block, walked with `trailing_zeros`. The first needle byte is compared in SIMD (the existing `mm64` fold on NEON, `movemask` on SSE); the remaining bytes only ever run on the survivors, so a 3-byte needle costs one vector pass plus a handful of scalar checks. This is what puts metaspace (llama2 / gemma) fully on bitsplit: `to_normalizer_and_split` already decomposes Metaspace at load time into a MetaspaceNormalizer + a `Split` on the `U+2581` delimiter, and that Split's search is now a bitstream. Measured on the metaspace needle over 8 MB (examples/litbench): scalar first byte 727 MB/s 0.64x memmem <- not good enough SIMD first byte 1433 MB/s 1.27x memmem memchr is no longer a bitsplit dependency. * perf(bitsplit): escape only the letter runs that need it The gate was "any trigger in this block flags every letter token in it", so a single `'s` in a paragraph dragged every word through the scalar pass. Narrow it with a reverse fill over the letter stream: a run needs the escape only if an interior upper (or the letter before a contraction apostrophe) sits at-or-after the run start. `last_lt` keeps tracking ALL letter tokens -- it is what the cross-block patch uses, and narrowing it broke every o200k-family gate until the two were separated. english MB/s: o200k 504 -> 1091, tekken 1547 -> 1507, kimi 287 -> 550. 8 gates green. * refactor(bitsplit): classify/ and models// folders - `classify/` owns its own pieces instead of scattering them at the crate root: the generated `atom_tables.rs`, the `tables.rs` layout it bakes, and the three per-arch kernels (`neon.rs`, `avx.rs`, `wasm.rs`, renamed from `simd_*_classify.rs`). - `models//` — one folder per unrolled pre-tokenization regex. Models sharing a regex share a folder: o200k covers Llama-4 / gpt-oss / MiniMax-M2, cl100k covers Llama-3 / GLM-4.6 and Qwen at digit cap 1. bitmap_gen now writes classify/atom_tables.rs and emits `use super::tables::Tables;`. Regenerating produces a byte-identical file apart from that import. Root re-exports (`bitsplit::bitsplit_o200k`, ...) are unchanged, so tk-encode is untouched. 23 suites green. * refactor(bitsplit): file per model, simd/ folder, drop examples - `models/.rs` — a file per regex, not a folder. - `simd/` holds the per-arch kernels: `neon.rs` + `x86.rs` (block builder) and `classes.rs` (class-run boundaries), replacing the loose `simd*.rs` at the root. classify keeps its own kernels next to the tables they index. - `bitsplit/examples/` removed, along with the memchr dev-dep it needed. The numbers they produced are recorded in the PR description. --- .../workflows/{atomsplit.yml => bitsplit.yml} | 18 +- tokenizers/Cargo.lock | 135 +- tokenizers/Cargo.toml | 2 +- tokenizers/atomsplit/Cargo.toml | 48 - tokenizers/atomsplit/README.md | 221 - tokenizers/atomsplit/benches/class_runs.rs | 123 - tokenizers/atomsplit/benches/classify.rs | 90 - tokenizers/atomsplit/benches/data/.gitignore | 1 - tokenizers/atomsplit/benches/data/fetch.py | 17 - tokenizers/atomsplit/benches/heatmap.py | 74 - .../atomsplit/benches/pretok_heatmap.svg | 8344 ----------------- tokenizers/atomsplit/benches/regex.rs | 429 - tokenizers/atomsplit/src/fsm.rs | 387 - tokenizers/atomsplit/src/fsm/byte_level.rs | 81 - tokenizers/atomsplit/src/fsm/cl100k.rs | 118 - tokenizers/atomsplit/src/fsm/deepseek.rs | 247 - tokenizers/atomsplit/src/fsm/o200k.rs | 246 - tokenizers/atomsplit/src/lib.rs | 33 - tokenizers/atomsplit/src/literal.rs | 56 - tokenizers/atomsplit/tests/fsm.rs | 107 - tokenizers/atomsplit/tests/literal.rs | 50 - tokenizers/atomsplit/tests/parity.rs | 121 - tokenizers/bitmap_gen/Cargo.toml | 2 +- tokenizers/bitmap_gen/src/lib.rs | 2 +- tokenizers/bitmap_gen/src/main.rs | 9 +- tokenizers/bitsplit/Cargo.toml | 29 +- tokenizers/bitsplit/src/classes.rs | 210 + .../src/classify}/atom_tables.rs | 2 +- .../src/classify/avx.rs} | 4 +- .../src/classify/mod.rs} | 20 +- .../src/classify/neon.rs} | 8 +- .../src => bitsplit/src/classify}/tables.rs | 0 .../src/classify/wasm.rs} | 4 +- tokenizers/bitsplit/src/han.rs | 72 + tokenizers/bitsplit/src/lib.rs | 264 +- tokenizers/bitsplit/src/literal.rs | 209 + .../bitsplit/src/{gpt.rs => models/cl100k.rs} | 155 +- .../bitsplit/src/{ => models}/deepseek.rs | 8 +- tokenizers/bitsplit/src/models/gpt2.rs | 162 + tokenizers/bitsplit/src/models/kimi.rs | 556 ++ tokenizers/bitsplit/src/models/mod.rs | 12 + tokenizers/bitsplit/src/models/o200k.rs | 527 ++ tokenizers/bitsplit/src/models/tekken.rs | 525 ++ .../{atomsplit => bitsplit}/src/regexes.rs | 6 + .../src/simd/classes.rs} | 3 +- tokenizers/bitsplit/src/simd/mod.rs | 12 + .../bitsplit/src/{simd.rs => simd/neon.rs} | 73 +- .../bitsplit/src/{simd_x86.rs => simd/x86.rs} | 70 +- tokenizers/bitsplit/tests/parity.rs | 265 + tokenizers/tk-encode/Cargo.toml | 1 - tokenizers/tk-encode/src/models/bpe/model.rs | 67 +- .../tk-encode/src/normalizers/replace.rs | 2 +- .../tk-encode/src/pre_tokenizers/bert.rs | 4 +- .../src/pre_tokenizers/byte_level.rs | 4 +- .../tk-encode/src/pre_tokenizers/delimiter.rs | 2 +- .../tk-encode/src/pre_tokenizers/digits.rs | 4 +- .../src/pre_tokenizers/punctuation.rs | 4 +- .../tk-encode/src/pre_tokenizers/sequence.rs | 6 +- .../tk-encode/src/pre_tokenizers/split.rs | 64 +- .../src/pre_tokenizers/whitespace.rs | 8 +- .../tk-encode/src/tokenizer/normalizer.rs | 2 +- tokenizers/tk-encode/src/tokenizer/pattern.rs | 2 +- .../tk-encode/src/tokenizer/pipeline.rs | 4 +- tokenizers/tk-encode/src/utils/byte_level.rs | 2 +- tokenizers/tk-encode/src/utils/mod.rs | 4 +- tokenizers/tk-encode/src/utils/no_regex.rs | 2 +- .../tk-encode/src/utils/unrolled_regex.rs | 145 +- 67 files changed, 3163 insertions(+), 11321 deletions(-) rename .github/workflows/{atomsplit.yml => bitsplit.yml} (72%) delete mode 100644 tokenizers/atomsplit/Cargo.toml delete mode 100644 tokenizers/atomsplit/README.md delete mode 100644 tokenizers/atomsplit/benches/class_runs.rs delete mode 100644 tokenizers/atomsplit/benches/classify.rs delete mode 100644 tokenizers/atomsplit/benches/data/.gitignore delete mode 100644 tokenizers/atomsplit/benches/data/fetch.py delete mode 100644 tokenizers/atomsplit/benches/heatmap.py delete mode 100644 tokenizers/atomsplit/benches/pretok_heatmap.svg delete mode 100644 tokenizers/atomsplit/benches/regex.rs delete mode 100644 tokenizers/atomsplit/src/fsm.rs delete mode 100644 tokenizers/atomsplit/src/fsm/byte_level.rs delete mode 100644 tokenizers/atomsplit/src/fsm/cl100k.rs delete mode 100644 tokenizers/atomsplit/src/fsm/deepseek.rs delete mode 100644 tokenizers/atomsplit/src/fsm/o200k.rs delete mode 100644 tokenizers/atomsplit/src/lib.rs delete mode 100644 tokenizers/atomsplit/src/literal.rs delete mode 100644 tokenizers/atomsplit/tests/fsm.rs delete mode 100644 tokenizers/atomsplit/tests/literal.rs delete mode 100644 tokenizers/atomsplit/tests/parity.rs create mode 100644 tokenizers/bitsplit/src/classes.rs rename tokenizers/{atomsplit/src => bitsplit/src/classify}/atom_tables.rs (99%) rename tokenizers/{atomsplit/src/simd_avx_classify.rs => bitsplit/src/classify/avx.rs} (99%) rename tokenizers/{atomsplit/src/classify.rs => bitsplit/src/classify/mod.rs} (92%) rename tokenizers/{atomsplit/src/simd_classify.rs => bitsplit/src/classify/neon.rs} (99%) rename tokenizers/{atomsplit/src => bitsplit/src/classify}/tables.rs (100%) rename tokenizers/{atomsplit/src/simd_wasm_classify.rs => bitsplit/src/classify/wasm.rs} (99%) create mode 100644 tokenizers/bitsplit/src/han.rs create mode 100644 tokenizers/bitsplit/src/literal.rs rename tokenizers/bitsplit/src/{gpt.rs => models/cl100k.rs} (62%) rename tokenizers/bitsplit/src/{ => models}/deepseek.rs (97%) create mode 100644 tokenizers/bitsplit/src/models/gpt2.rs create mode 100644 tokenizers/bitsplit/src/models/kimi.rs create mode 100644 tokenizers/bitsplit/src/models/mod.rs create mode 100644 tokenizers/bitsplit/src/models/o200k.rs create mode 100644 tokenizers/bitsplit/src/models/tekken.rs rename tokenizers/{atomsplit => bitsplit}/src/regexes.rs (77%) rename tokenizers/{atomsplit/src/simd_fsm.rs => bitsplit/src/simd/classes.rs} (99%) create mode 100644 tokenizers/bitsplit/src/simd/mod.rs rename tokenizers/bitsplit/src/{simd.rs => simd/neon.rs} (76%) rename tokenizers/bitsplit/src/{simd_x86.rs => simd/x86.rs} (75%) create mode 100644 tokenizers/bitsplit/tests/parity.rs diff --git a/.github/workflows/atomsplit.yml b/.github/workflows/bitsplit.yml similarity index 72% rename from .github/workflows/atomsplit.yml rename to .github/workflows/bitsplit.yml index c2059f94f..f250a4675 100644 --- a/.github/workflows/atomsplit.yml +++ b/.github/workflows/bitsplit.yml @@ -1,10 +1,10 @@ -name: atomsplit +name: bitsplit on: push: - paths: ["tokenizers/atomsplit/**", "tokenizers/bitmap_gen/**", ".github/workflows/atomsplit.yml"] + paths: ["tokenizers/bitsplit/**", "tokenizers/bitmap_gen/**", ".github/workflows/bitsplit.yml"] pull_request: - paths: ["tokenizers/atomsplit/**", "tokenizers/bitmap_gen/**", ".github/workflows/atomsplit.yml"] + paths: ["tokenizers/bitsplit/**", "tokenizers/bitmap_gen/**", ".github/workflows/bitsplit.yml"] defaults: run: @@ -18,11 +18,11 @@ jobs: - uses: actions/checkout@v4 - uses: dtolnay/rust-toolchain@stable with: { components: rustfmt, clippy } - - run: cargo fmt -p atomsplit -p bitmap_gen -- --check - - run: cargo clippy -p atomsplit --all-targets -- -D warnings - - run: cargo test -p atomsplit # unit + parity gates (cl100k/deepseek/byte_level vs onig) + - run: cargo fmt -p bitsplit -p bitmap_gen -- --check + - run: cargo clippy -p bitsplit --all-targets -- -D warnings + - run: cargo test -p bitsplit # unit + parity gates (cl100k/deepseek/byte_level vs onig) # committed classify tables must match the generator - - run: cargo run -p bitmap_gen && git diff --exit-code atomsplit/src/atom_tables.rs + - run: cargo run -p bitmap_gen && git diff --exit-code bitsplit/src/atom_tables.rs # AVX-512 classify path — GitHub runners often lack AVX-512, so emulate it with Intel SDE and run the # SAME byte-exact test binaries under it. @@ -39,7 +39,7 @@ jobs: echo "$PWD/${SDE_VER}" >> "$GITHUB_PATH" - name: Run tests under SDE (-future = AVX-512) run: | - cargo test -p atomsplit --no-run --message-format=json \ + cargo test -p bitsplit --no-run --message-format=json \ | jq -r 'select(.profile.test == true) | .executable | select(. != null)' \ | while read -r bin; do echo "SDE: $bin"; sde64 -future -- "$bin"; done @@ -55,4 +55,4 @@ jobs: env: RUSTFLAGS: "-C target-feature=+simd128" CARGO_TARGET_WASM32_WASIP1_RUNNER: "wasmtime run --" - run: cargo test -p atomsplit --target wasm32-wasip1 + run: cargo test -p bitsplit --target wasm32-wasip1 diff --git a/tokenizers/Cargo.lock b/tokenizers/Cargo.lock index ba9632eb1..88ba567cc 100644 --- a/tokenizers/Cargo.lock +++ b/tokenizers/Cargo.lock @@ -67,17 +67,6 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" -[[package]] -name = "atomsplit" -version = "0.1.0" -dependencies = [ - "fancy-regex 0.13.0", - "logos", - "memchr", - "onig", - "pcre2", -] - [[package]] name = "autocfg" version = "1.5.1" @@ -96,12 +85,6 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" -[[package]] -name = "beef" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a8241f3ebb85c056b509d4327ad0358fbbba6ffb340bf388f26350aeda225b1" - [[package]] name = "bindgen" version = "0.72.1" @@ -122,30 +105,15 @@ dependencies = [ "syn", ] -[[package]] -name = "bit-set" -version = "0.5.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" -dependencies = [ - "bit-vec 0.6.3", -] - [[package]] name = "bit-set" version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" dependencies = [ - "bit-vec 0.8.0", + "bit-vec", ] -[[package]] -name = "bit-vec" -version = "0.6.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" - [[package]] name = "bit-vec" version = "0.8.0" @@ -170,7 +138,7 @@ name = "bitsplit" version = "0.1.0" dependencies = [ "ahash", - "atomsplit", + "onig", ] [[package]] @@ -241,8 +209,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" dependencies = [ "find-msvc-tools", - "jobserver", - "libc", "shlex 2.0.1", ] @@ -663,24 +629,13 @@ dependencies = [ "cc", ] -[[package]] -name = "fancy-regex" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "531e46835a22af56d1e3b66f04844bed63158bc094a628bec1d321d9b4c44bf2" -dependencies = [ - "bit-set 0.5.3", - "regex-automata", - "regex-syntax", -] - [[package]] name = "fancy-regex" version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8" dependencies = [ - "bit-set 0.8.0", + "bit-set", "regex-automata", "regex-syntax", ] @@ -1188,16 +1143,6 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" -[[package]] -name = "jobserver" -version = "0.1.35" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" -dependencies = [ - "getrandom 0.4.3", - "libc", -] - [[package]] name = "js-sys" version = "0.3.103" @@ -1258,40 +1203,6 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" -[[package]] -name = "logos" -version = "0.15.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff472f899b4ec2d99161c51f60ff7075eeb3097069a36050d8037a6325eb8154" -dependencies = [ - "logos-derive", -] - -[[package]] -name = "logos-codegen" -version = "0.15.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "192a3a2b90b0c05b27a0b2c43eecdb7c415e29243acc3f89cc8247a5b693045c" -dependencies = [ - "beef", - "fnv", - "lazy_static", - "proc-macro2", - "quote", - "regex-syntax", - "rustc_version", - "syn", -] - -[[package]] -name = "logos-derive" -version = "0.15.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "605d9697bcd5ef3a42d38efc51541aa3d6a4a25f7ab6d1ed0da5ac632a26b470" -dependencies = [ - "logos-codegen", -] - [[package]] name = "lru-slab" version = "0.1.2" @@ -1484,28 +1395,6 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" -[[package]] -name = "pcre2" -version = "0.2.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e970b0fcce0c7ee6ef662744ff711f21ccd6f11b7cf03cd187a80e89797fc67" -dependencies = [ - "libc", - "log", - "pcre2-sys", -] - -[[package]] -name = "pcre2-sys" -version = "0.2.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "18b9073c1a2549bd409bf4a32c94d903bb1a09bf845bc306ae148897fa0760a4" -dependencies = [ - "cc", - "libc", - "pkg-config", -] - [[package]] name = "percent-encoding" version = "2.3.2" @@ -1909,15 +1798,6 @@ version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" -[[package]] -name = "rustc_version" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" -dependencies = [ - "semver", -] - [[package]] name = "rustix" version = "1.1.4" @@ -1997,12 +1877,6 @@ dependencies = [ "winapi-util", ] -[[package]] -name = "semver" -version = "1.0.28" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" - [[package]] name = "serde" version = "1.0.228" @@ -2294,14 +2168,13 @@ version = "0.23.2-dev.0" dependencies = [ "ahash", "assert_approx_eq", - "atomsplit", "bitsplit", "compact_str", "criterion 0.6.0", "daachorse 3.0.2", "dary_heap", "derive_builder", - "fancy-regex 0.17.0", + "fancy-regex", "getrandom 0.3.4", "hf-hub", "indicatif 0.18.5", diff --git a/tokenizers/Cargo.toml b/tokenizers/Cargo.toml index fdd4f4894..5100c6367 100644 --- a/tokenizers/Cargo.toml +++ b/tokenizers/Cargo.toml @@ -1,6 +1,6 @@ [workspace] resolver = "3" -members = ["bitmap_gen", "atomsplit", "bitsplit", "tk-encode", "tk-train"] +members = ["bitmap_gen", "bitsplit", "tk-encode", "tk-train"] [package] authors = [ diff --git a/tokenizers/atomsplit/Cargo.toml b/tokenizers/atomsplit/Cargo.toml deleted file mode 100644 index fb8947f3f..000000000 --- a/tokenizers/atomsplit/Cargo.toml +++ /dev/null @@ -1,48 +0,0 @@ -[package] -name = "atomsplit" -version = "0.1.0" -edition = "2024" -rust-version = "1.89" # AVX-512 VBMI intrinsics (simd_avx_classify) are stable only since 1.89 -authors = [ - "Arthur Zucker ", - "Luc Georges ", - "Simon Brandeis ", -] -homepage = "https://github.com/huggingface/tokenizers" -repository = "https://github.com/huggingface/tokenizers" -documentation = "https://docs.rs/atomsplit/" -license = "Apache-2.0" -keywords = ["tokenizer", "nlp", "pretokenizer", "simd", "unicode"] -categories = ["text-processing", "parsing"] -description = "SIMD Unicode atom classification + FSM pre-tokenization: one classify pass into byte-exact token spans (cl100k, GPT-2/ByteLevel, DeepSeek, Whitespace, Bert, ...)." -exclude = ["benches/data/"] - -[package.metadata.docs.rs] -all-features = true - -[lib] -name = "atomsplit" -path = "src/lib.rs" - -[dependencies] -memchr = "2.8.2" # SIMD search: used for single-byte search in CharDelimiterSplit (1.4–23× vs scalar), or multi-byte string pattern search in `literal`. -# NOTE: classify tables live in src/atom_tables.rs (committed, generated). Regenerate after any atom -# scheme change with `cargo run -p bitmap_gen`. No build script / build-dep — atomsplit builds clean. -# onig is C (Oniguruma) and fancy-regex pulls it in transitively for benches only — neither -# cross-compiles to wasm32 (no wasi libc for the C build), so gate them off wasm. -[target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies] -onig = "6.5.1" # reference regex for the parity tests + benches -fancy-regex = "0.13" # second regex engine in the throughput benches -logos = "0.15" # third: a compile-time DFA lexer-generator (pure Rust) -pcre2 = "0.2" # fourth: PCRE2 with its JIT (C, pregenerated bindings — no libclang) -[[bench]] -name = "regex" # gpt2 · cl100k · o200k · tekken · deepseek, one run per regex family -harness = false - -[[bench]] -name = "classify" -harness = false - -[[bench]] -name = "class_runs" -harness = false diff --git a/tokenizers/atomsplit/README.md b/tokenizers/atomsplit/README.md deleted file mode 100644 index 0802d27a0..000000000 --- a/tokenizers/atomsplit/README.md +++ /dev/null @@ -1,221 +0,0 @@ -# Tag-classify pretokenization — design spec -The key design principle is to run 1 SIMD classification pass over the input, and then run a finite state machine on the produced tags. -SIMD first for the fsm is not great as unrolled regex can be quite complicated and cut often in many cases. The only SIMD you want in the fsm is when you are looking for `*` or `+` patterns. There, SIMD allows you to go fast to the last byte of the category you are looking for. -For simple pretokenizers like whitespace split that emit splits at tag boundary changes, simd can also be used. - -We always have a scalar fallback for both the classification and the finite state machine. -## 1. The generic classify engine (composable lanes) -The key design principle is that no matter the number of atoms (well it has to be <255) the classifier does not change. The classifier operates on byte length. For each byte length we find a smart way to retrieve the class from pre-computed table. A new custom pre-tokenizer would only require us to update the table, never the classifiers. This scales really well as you can combine tags in the FSM to create bigger categories (like white space markers are whitespace and markers) - -| lane | lead byte range | mechanism | shared SIMD ↔ scalar? | -|---------------|----------------------|------------------------------------------------------------------------|------------------------------------| -| ASCII | `0x00–0x7F` (1 B) | 128-entry table, two 64-halves + subtract-trick `OR` | ✓ one table | -| 2-byte | `0xC2–0xDF` (2 B) | 8 groups × (4 sub × 64); **peel** the min group → one 256-lookup | ✓ one table | -| CJK shortcut | `0xE3–0xED` (3 B) | range compares → one tag (current scheme folds all CJK → `Letter`) | SIMD-only optimism (scalar → 3-byte) | -| 3-byte | `0xE0–0xEF` (3 B) | 512 blocks: 425 uniform const / 87 mixed 128-tables; **peel** blocks | ✓ one table | -| cold BMP | any BMP deferred | run-length `(start_cp, tag)`, binary-searched (~1–3 KB) | scalar reader / SIMD `MB`-fixup | -| astral | `0xF0–0xF4` (4 B) | run-length `(start_cp, tag)`, binary-searched | scalar / SIMD stamps `MB` → fixup | -| continuation | `0x80–0xBF` | tagged `Cont` — transparent to every FSM | ✓ | - -The **13 coarse atoms** the default scheme emits (the coarse class is the tag's **low nibble**, so it fits in a `u4`; the whole engine works for any `< 255`): - -``` -0 Letter 1 NumWord 2 NumOther 3 Newline 4 Space 5 WsOther 6 Mark -7 Connector 8 Punct 9 Apostrophe 10 SymOther 11 NumericOther 12 Control - ┌ internal sentinels, never seen by the FSM ┐ - 13 Sentinel 14 MultiByte 15 Cont -``` - -The tag is a full `u8`: the **high nibble** is an optional *refinement* that sub-splits one coarse class for a pretokenizer needing finer granularity — without a second pass. o200k's case split refines `Letter` into `UPPER` (`\p{Lu}\p{Lt}`) / `LOWER` (`\p{Ll}`) / caseless, and `Mark` carries `ALPHA_SYM` for Other_Alphabetic symbols (`\w` but categorically `\p{S}`, e.g. circled letters). Coarse consumers collapse it for free — `in_mask` and the SIMD class path `& 0x0F` off the nibble — so only the FSM that opted in (o200k) ever sees it. - -### 1.1 The only thing you need to know: how UTF-8 lays out a codepoint - -Every table is indexed straight off the raw UTF-8 bytes, so the whole scheme falls out of the encoding. `x`/`y`/`z` are the payload bits of the codepoint; the `0`/`10`/`110`/`1110` prefixes are UTF-8's length tag: - -``` - ├────────────── codepoint bits ──────────────┤ -1 byte 0xxxxxxx U+0000 .. U+007F (ASCII) -2 byte 110xxxyy 10yyyyyy U+0080 .. U+07FF -3 byte 1110xxxx 10yyyyyy 10zzzzzz U+0800 .. U+FFFF (BMP) -4 byte 11110www 10xxxxxx 10yyyyyy 10zzzzzz U+10000 .. U+10FFFF (astral) - └──┬───┘ └──┬───┘ - lead continuation bytes (all start 10……) -``` - -Read a single byte and its **top bits tell you the lane**: - -``` -byte & 0x80 == 0x00 → 0xxxxxxx ASCII (1-byte) -byte & 0xC0 == 0x80 → 10xxxxxx continuation → Cont -byte & 0xE0 == 0xC0 → 110xxxxx 2-byte lead -byte & 0xF0 == 0xE0 → 1110xxxx 3-byte lead -byte & 0xF8 == 0xF0 → 11110xxx 4-byte lead -``` - -### 1.2 How the tables are built (`bitmap_gen`, one source of truth) - -We never hand-write a table. We **synthesize every byte sequence a lane can hold, decode it back to a codepoint, and ask the reference `atom(char)` for its class** — then store that class at the index the raw bytes produce. Because the *decode* and the *index* both come from the same bit layout above, the runtime lookup is exact by construction. One reference function → every table (and the SIMD kernel) derived from it; the generator then re-derives all 1.1 M codepoints and asserts the packed tables read back identically, so a scheme change that breaks a table fails the build. - -``` -for every byte pattern of the lane ──► cp = decode(bytes) ──► atom(char::from(cp)) ──► table[index(bytes)] = tag - (§1.1 layout) (the ONE reference) (same layout → same index) -``` - -### 1.3 Per-lane mechanics - -**ASCII — one 128-entry table, split in two halves.** NEON's table op (`vqtbl`) resolves ≤ 64 entries, so 128 needs two. The subtract trick makes one 8-bit index hit exactly one half (out-of-range lookups return 0, so the `OR` is clean): - -``` -tag(v) = tbl64(ascii_lo, v) OR tbl64(ascii_hi, v - 64) - └ v<64 → lo[v] └ v≥64 → hi[v-64] - └ v≥64 → 0 (oob) └ v<64 → v-64 wraps huge → 0 (oob) - -v=5 : lo[5] | (5-64 wraps → 0) = lo[5] -v=65 : (65 ≥64→0) | hi[1] = hi[1] -``` - -**2-byte — `110xxxyy 10yyyyyy`, 8 groups × 4 sub-groups × 64.** There are only 32 possible leads (`C0..DF`). Split the 5 lead payload bits into `xxx` (group, 8) and `yy` (sub-group, 4); the continuation's `yyyyyy` is the offset (64): - -``` -lead 110 x x x y y cont 10 y y y y y y - └─┬─┘ └┬┘ └──┬───┘ - xxx yy yyyyyy - group subgroup offset - (b0>>2)&7 b0&3 b1&0x3F - -lookup: group_tables[ xxx ][ yy ][ yyyyyy ] ← declared [8][4][64] -in SIMD: the inner [4][64] is 256 contiguous bytes, so ONE 256-lookup does it, - indexed by group_index = (yy<<6) | yyyyyy : - idx 0..63 → sub 0 idx 128..191 → sub 2 - idx 64..127 → sub 1 idx 192..255 → sub 3 (subtract-window OR of 4×vqtbl4) -``` - -**3-byte — `1110xxxx 10yyyyyy 10zzzzzz`, 512 blocks, most of them a single tag.** A *block* = one lead (`E0..EF`, 16) × the high bits of the 2nd byte (`b1>>1`, 32) = 512 blocks, each covering 128 codepoints. Real scripts are homogeneous: `425 / 512` blocks are **one atom** (all of CJK Han is `Letter`, whole symbol ranges are `SymOther`), stored as a single const byte; only `87` "mixed" blocks (e.g. a script interleaved with its punctuation) need an actual 128-entry table. - -``` -1110 x x x x 10 y y y y y y 10 z z z z z z - └──┬──┘ └┬┘└───┬──┘ └───┬────┘ - lead b1-pair (unused here) within-block offset - E0..EF (16) (b1>>1)&0x1F (32) (b1&1)<<6 | z (128) - -block = (lead-0xE0)*32 + ((b1>>1)&0x1F) ── 512 of them -fast3_uni[block] = tag if the whole 128-cp block is one atom (425 blocks — just a const) - = 0xFF otherwise → fast3_mixed[ fast3_slot[block] ] holds a 128-table (87 blocks) -``` - -**CJK shortcut (SIMD only).** The current scheme folds all of CJK to one tag (`CJK_TAG = Atom::Letter`), so the SIMD kernel skips the tables for `E3..ED` and proves "is CJK" with a few range compares (Han `E4..E9`, Hangul `EB..EC`, kana `E3 81..83`, minus a handful of punctuation holes like `・`). It only ever *under*-claims (boundary/hole codepoints fall through to the exact 3-byte tables), so it stays byte-exact. (A scheme that split CJK into distinct tags would drop this shortcut and use the tables.) - -**Cold fallback + astral.** The dense BMP tag stream and the astral range are **run-length encoded** `(start_cp, tag)` and binary-searched — a few KB instead of a 64 KB dense LUT, because tags change rarely across a codepoint range. The SIMD kernel can't see a 4-byte char's 4th byte (it only gathers 3 bytes per lane), so it stamps those lanes `MB` and a per-chunk fixup resolves them via these tables while the chunk is still hot in L1. - -### 1.4 Why "peel" and not "loop the range" - -For 2-byte and 3-byte a chunk may hold lanes from several blocks. Instead of looping every block id between the min and max present (wasting a step per empty gap), we **peel**: `vminvq` the smallest block still unresolved → resolve exactly its lanes with one lookup → mask them out → repeat. Steps = **distinct blocks present** (usually 1 → that's the fast path), independent of how far apart the scripts sit, and no bounds guard. See `src/simd_classify.rs` for the annotated kernel. - -## 2. The FSM: classify is nearly free, so the FSM *is* the cost - -The single SIMD classify pass is ~free — **~0.05 ns/B on ASCII, ~0.3–0.6 on multibyte** (a handful of table lookups per 16 bytes). The pre-tokenizer's real cost is the **FSM** that turns the tag stream into spans: **~1–3 ns/B on dense Latin/code** (many short tokens ⇒ many state transitions) down to **~0.3 on run-heavy CJK**. So all the performance work lives here. The FSM writes spans into a **caller-owned `&mut [Span]` and returns the count** — no `Vec`, no realloc; the buffer is reused across calls. - -### 2.1 A class is a `u16` bitmap - -The classifier hands the FSM a stream of atom ids (`0..12`). A **class** is just a *set of atoms*, and a set of ≤ 16 atoms is one `u16` where **bit `t` = "atom `t` is in this class"**: - -``` -bit : 12 11 10 9 8 7 6 5 4 3 2 1 0 -atom : Ct Nu Sy Ap Pn Cn Mk Ws Sp Nl No Nw Lt (Ct=Control Nu=NumericOther Sy=SymOther - Ap=Apostrophe Pn=Punct Cn=Connector Mk=Mark - Ws=WsOther Sp=Space Nl=Newline No=NumOther - Nw=NumWord Lt=Letter) -WS : . . . . . . . 1 1 1 . . . = 0x0038 (Newline | Space | WsOther) -WORD : . . . . . 1 1 . . . . 1 1 = 0x00C3 (Letter | NumWord | Mark | Connector) -``` - -Building and testing a class are one instruction each — no branches, no table: - -```rust -Atom::WsOther.bit() // = 1 << 5 build a class by OR-ing atoms -in_mask(tag, WS) // = WS & (1 << tag) != 0 test membership: one AND + compare -``` - -A class is any union of atoms — `PUNCT_SYM = Connector|Punct|Apostrophe|SymOther`, `LETTER_MARK = Letter|Mark` — composed *without touching the classifier or re-running classification*. Add a 14th atom and every FSM keeps working; `u16` covers 16 atoms, and the exact same code with a `u32`/`u64` mask covers 32/64 classes for free (the classifier already emits arbitrary `u8` tags). - -### 2.2 The scalar ↔ SIMD duality — SIMD only earns its keep on run-ends - -A pre-tokenizer rule is one of two shapes, and **SIMD only helps the second**: - -1. **Per-char decisions** — cl100k's contraction peek (`'s|'t|…`), the ` ?` / `[^…]?` prefix logic: a branchy little automaton, one char at a time. Branches don't vectorize, so this stays **scalar**. -2. **Run-ends — a `+`/`*` over one class** (`\p{L}+`, `\s+`, `\p{N}+`, a punct run): grab the maximal run. The trick is **skip the invalid, stop only at the valid**. The regex-shaped FSMs (`fsm_cl100k`, `fsm_o200k`, …) do this with the scalar `run_end`: it unrolls **16 tags per chunk** (one bounds check, `get_unchecked` reads), so a 3-char English word finishes in the ≤16-byte scalar tail while a long CJK run skips whole chunks. The *SIMD* form of the same idea lives only in the class family's `class_runs_neon` (below): `vqtbl1` a 16-tag membership LUT, `vminvq == 0xFF` ⇒ all 16 still in-run ⇒ skip the chunk. - -The **dual** of a run-end is a **boundary**: the class-family pre-tokenizers (Whitespace / Bert / Digits / Punctuation) cut at *every* class change, so `class_runs_neon` scans the other way — classify 16 lanes (`vqtbl1`), fill continuation lanes from the left, `movemask` the class-changes, and iterate the set bits to emit one span per segment. It finds **every boundary in a chunk at once**, so short-run text (English words, ~5 chars) is never paid per-char. - -Both fuse in one kernel: `class_runs_neon` first tries the run-end fast path — *whole 16-chunk stays the current class → skip it, no boundary work* — and only movemasks the **mixed** chunks. So it bulk-skips long runs (Digits/Punct/CJK) **and** parallel-boundaries the dense ones, byte-exact with the portable `emit_class_spans` oracle. Two views of the same idea: **run-end skips the invalid to reach the next valid; boundary-extract flags every valid at once.** - -### 2.3 Why it's a finite-state machine - -Each pre-tokenizer is a small automaton over the shared tag alphabet: a few **states** (in a letter run, inside whitespace, at a boundary) and, per incoming tag, a **transition** — extend the run, emit the span, or start a new one. The class masks *are* the transition predicates (`in_mask(tag, LETTER)` = "stay in the letter state"). So cl100k, GPT-2, whitespace-split, deepseek… are all the *same* machine driven by different masks + a few peeked bytes; only the states that are a maximal `+`/`*` run get the SIMD run-end. - -## 3. Measured performance - -Single-thread, **180 KB per language** (uniform size → comparable cache behaviour), **min-of-7 trials**, 14-core Apple Silicon (8 P + 6 E), light background load. `ns/byte`, lower is better. **Every row is byte-exact ✓** against the reference regex (onig for gpt2 / cl100k / o200k; the composed onig×3 `Sequence` for deepseek). Reproduce with `cargo bench --bench regex` (every pre-tokenizer × the four reference engines — the §3.0 chart); the class family's scalar-vs-SIMD boundary extractor is `--bench class_runs`, classify alone is `--bench classify`. - -### 3.0 At a glance — vs every SOTA splitter - -The `regex` bench pits the full pipeline (SIMD classify + scalar FSM) against **four** reference engines — -**onig** and **pcre2** (with JIT), both C; **fancy-regex**, pure-Rust; and **logos**, a compile-time DFA -lexer-generator — for each pre-tokenizer × language. Speedup = engine ÷ our pipeline; **green = we win big, -red = a close race**: - -![pre-tokenization speedup vs onig / fancy / logos / pcre2-JIT](benches/pretok_heatmap.svg) - -We lead on every pre-tokenizer — **~4–60×** vs onig/fancy and **~1.4–10×** vs the JIT / DFA engines — the -only near-ties being **o200k on CJK** (its case-split FSM is heavy there). `n/a` marks a split the engine -can't express: logos has no look-ahead, so deepseek's 3-regex `Sequence` and the punctuation-isolation -splits (punct / bert) have no single-grammar logos form. The GPT FSMs are byte-exact with their regex -(✓); the class-family reference regex is an approximation of the atom mask (`≈` where it diverges), so -those rows are a speed pairing rather than an equality gate. Regenerate: - -```sh -cargo bench --bench regex > bench.txt -python3 benches/heatmap.py bench.txt benches/pretok_heatmap.svg # needs matplotlib + numpy -``` - -### 3.1 cl100k — classify is ~free, the FSM is the cost - -The regex-shaped FSMs are scalar (there is no SIMD cl100k FSM — SIMD lives only in the class family's `class_runs_neon`, §2.2). `pipeline = SIMD classify + scalar FSM`. - -| lang | b/tok | classify | FSM (scalar) | onig | pipeline **vs onig** | -|---|--:|--:|--:|--:|--:| -| English | 4.6 | **0.068** | 1.037 | 36.3 | **32.9×** | -| French | 5.1 | 0.158 | 0.836 | 35.9 | 36.1× | -| Russian | 10.2 | 0.315 | 0.462 | 20.7 | 26.6× | -| Greek | 9.9 | 0.296 | 0.462 | 22.3 | 29.3× | -| Hebrew | 8.3 | 0.302 | 0.481 | 24.2 | 30.9× | -| Arabic | 9.1 | 0.293 | 0.499 | 23.6 | 29.8× | -| Hindi | 5.3 | 0.320 | 0.659 | 39.9 | 40.7× | -| Thai | 11.8 | 0.327 | 0.416 | 23.0 | 30.9× | -| Chinese | 19.3 | 0.601 | 0.289 | 14.4 | 16.2× | -| Japanese | 25.3 | 0.715 | 0.265 | 14.1 | 14.4× | -| Korean | 7.4 | 0.440 | 0.617 | 26.6 | 25.1× | - -Reading it: **classify ≪ FSM** — on dense Latin the FSM is ~5–15× the classify, so that's where the work is (§2). Pure-ASCII English classifies at **0.068** (the ASCII fast path skips whole chunks); accented Latin (French) and multibyte scripts leave that fast path so classify costs more. The FSM is *cheapest* on run-heavy CJK (Chinese/Japanese ~0.27 — long homogeneous runs skip fast via the unrolled `run_end`) and dearest on dense short-token Latin. The full pipeline (SIMD classify + scalar FSM) is **14–41× onig**. - -### 3.2 GPT-2 (ByteLevel), o200k & deepseek — same story, other regexes - -All byte-exact; full-pipeline speedup vs the reference regex (one `cargo bench --bench regex` run covers all four families): - -- **GPT-2 / ByteLevel**: **21–45×** onig (English 34×, Hindi 45×, Chinese 22×). -- **o200k** (GPT-4o — case-aware `[\p{L}\p{M}]+` split, so a heavier FSM): **7–18×** onig (English 17×, Thai 7×, Chinese 10×). -- **deepseek** (a `Sequence` of 3 Isolated splits, collapsed into one FSM): **17–41×** the composed onig×3 (Chinese 41×, Japanese 28×, English 17×). - -### 3.3 Thread scaling (cl100k, ~16 MB doc, newline-partitioned seams) - -Threads spawned once (`std::thread::scope`, no external dep); each chunk's seam sits after a whitespace-run's last `\n`, so no token crosses it — **byte-exact for cl100k/deepseek** (proven: partitioned spans == sequential). MB/s (bytes×iters / wall; illustrative from a prior run — thread scaling has no standalone bench target in the current layout): - -| threads | English | | Chinese | | -|--:|--:|--:|--:|--:| -| 1 | 474 | 1.00× | 818 | 1.00× | -| 2 | 934 | 1.97× | 1580 | 1.93× | -| 4 | 1883 | 3.98× | 2998 | 3.66× | -| 8 | 3739 | 7.90× (99% linear) | 5698 | 6.96× | -| 14 | 4814 | 10.2× | 7229 | 8.83× | - -**~99% linear through the 8 performance cores**; the drop at 14 is the 6 efficiency cores (slower, so "% linear" falls — not contention). A single un-splittable long document uses the overlap-chunk path instead (BPE/pretok locality — see the merge notes). diff --git a/tokenizers/atomsplit/benches/class_runs.rs b/tokenizers/atomsplit/benches/class_runs.rs deleted file mode 100644 index 447ffae58..000000000 --- a/tokenizers/atomsplit/benches/class_runs.rs +++ /dev/null @@ -1,123 +0,0 @@ -//! Scalar vs SIMD fsm for the class family: `emit_class_spans` (scalar run-end core) vs -//! `class_runs_into` (NEON/SIMD128 movemask boundary-extract + homogeneous-chunk early-out). classify is -//! timed separately. `spd = scalar/simd`: >1 → SIMD wins. (cl100k is scalar-only — see benches/cl100k.rs.) -//! -//! Run: cargo bench --bench class_runs -use atomsplit::classify::{classify, mask}; -use atomsplit::fsm::{Span, class_runs_into, emit_class_spans}; -use std::hint::black_box; -use std::time::Instant; - -const CORPORA: &[(&str, &str)] = &[ - ("English", "../data/big.txt"), - ("French", "benches/data/fr.txt"), - ("Russian", "benches/data/ru.txt"), - ("Greek", "benches/data/el.txt"), - ("Arabic", "benches/data/ar.txt"), - ("Hindi", "benches/data/hi.txt"), - ("Thai", "benches/data/th.txt"), - ("Chinese", "benches/data/zh.txt"), - ("Japanese", "../data/unigram_wagahaiwa_nekodearu.txt"), - ("Korean", "benches/data/ko.txt"), -]; - -// scalar (run-end core) vs SIMD (class_runs_into) wrappers per recipe — #[inline(always)] so the timing -// closure inlines the whole chain (not a fn-pointer, which would skew the scalar path slow). -macro_rules! pair { - ($sname:ident, $vname:ident, $d:expr, $i:expr, $a:expr) => { - #[inline(always)] - fn $sname(t: &[u8], tg: &[u8], o: &mut [Span]) -> usize { - emit_class_spans::<$d, $i, $a>(t, tg, o, 0, 0, 0, None) - } - #[inline(always)] - fn $vname(t: &[u8], tg: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<$d, $i, $a>(t, tg, o) - } - }; -} -pair!(s_wss, v_wss, { mask::WS }, 0, 0); -pair!(s_pun, v_pun, 0, { mask::PUNCT }, 0); -pair!(s_dig, v_dig, 0, 0, { mask::NUMERIC }); -pair!(s_ws, v_ws, { mask::WS }, 0, { mask::WORD }); -pair!(s_bert, v_bert, { mask::WS }, { mask::PUNCT }, 0); - -fn ns_per_byte usize>(len: usize, iters: u32, mut f: F) -> f64 { - for _ in 0..3 { - black_box(f()); - } - let mut best = f64::INFINITY; - for _ in 0..7 { - let t = Instant::now(); - let mut acc = 0usize; - for _ in 0..iters { - acc = acc.wrapping_add(f()); - } - black_box(acc); - best = best.min(t.elapsed().as_nanos() as f64 / (iters as usize * len) as f64); - } - best -} - -fn main() { - let manifest = env!("CARGO_MANIFEST_DIR"); - let mut corpora: Vec<(&str, String)> = Vec::new(); - for (label, rel) in CORPORA { - let Ok(s) = std::fs::read_to_string(format!("{manifest}/{rel}")) else { - continue; - }; - if s.trim().is_empty() { - continue; - } - let mut c = s.len().min(180_000); - while c > 0 && !s.is_char_boundary(c) { - c -= 1; - } - corpora.push((label, s[..c].to_string())); - } - - macro_rules! compare { - ($name:expr, $s:ident, $v:ident) => {{ - println!( - "\n== {} ==\n{:<10} {:>7} {:>8} {:>8} {:>8} | {:>6}", - $name, "lang", "bytes", "classify", "scalar", "simd", "spd" - ); - for (label, corpus) in &corpora { - let text = corpus.as_bytes(); - let n = text.len(); - let iters = (4_000_000 / n).clamp(3, 150) as u32; - let mut tags = vec![0u8; n]; - let mut buf = vec![Span::default(); n + 1]; - classify(text, &mut tags); - // parity: scalar == simd - let (ks, kv) = ($s(text, &tags, &mut buf), 0); - let scalar_out: Vec = buf[..ks].to_vec(); - let _ = kv; - let kv = $v(text, &tags, &mut buf); - let parity = scalar_out == buf[..kv]; - let cls = ns_per_byte(n, iters, || { - classify(text, &mut tags); - n - }); - classify(text, &mut tags); - let sc = ns_per_byte(n, iters, || $s(text, &tags, &mut buf)); - let si = ns_per_byte(n, iters, || $v(text, &tags, &mut buf)); - println!( - "{label:<10} {n:>7} {cls:>8.3} {sc:>8.3} {si:>8.3} | {:>5.2}x{}", - sc / si, - if parity { "" } else { " PARITY✗" } - ); - } - }}; - } - - compare!("WhitespaceSplit", s_wss, v_wss); - compare!("Punctuation", s_pun, v_pun); - compare!("Digits", s_dig, v_dig); - compare!("Whitespace \\w", s_ws, v_ws); - compare!("Bert", s_bert, v_bert); - println!( - "\n(ns/byte, lower better. classify = SIMD classify; scalar = emit_class_spans (run-end\n \ - core); simd = class_runs_into (NEON/SIMD128 movemask + early-out). spd = scalar/simd, >1 → SIMD\n \ - wins. cl100k is scalar-only — its perf is in benches/cl100k.rs.)" - ); -} diff --git a/tokenizers/atomsplit/benches/classify.rs b/tokenizers/atomsplit/benches/classify.rs deleted file mode 100644 index da3287fe0..000000000 --- a/tokenizers/atomsplit/benches/classify.rs +++ /dev/null @@ -1,90 +0,0 @@ -//! classify throughput across three MB-fixup branch regimes, to check the per-chunk -//! `if any(out == MB)` fixup branch is predictable: -//! none — no astral (real text) → branch ALWAYS false → should be free. -//! astral — wall of emoji (every chunk has MB) → branch ALWAYS true → predictable, measures fixup cost. -//! sprinkle — BMP text with an emoji every ~24 B → branch flips irregularly → misprediction stress. -//! Run before/after the fixup change and compare. Run: cargo bench --bench classify -use atomsplit::classify::{classify, classify_scalar}; -use std::hint::black_box; -use std::time::Instant; - -fn ns_per_byte(text: &[u8], tags: &mut [u8]) -> f64 { - let iters = (8_000_000 / text.len().max(1)).clamp(20, 400) as u32; - for _ in 0..3 { - classify(text, tags); - black_box(tags[text.len() / 2]); - } - let mut best = f64::INFINITY; - for _ in 0..9 { - let t = Instant::now(); - for _ in 0..iters { - classify(text, tags); - black_box(tags[text.len() / 2]); - } - best = best.min(t.elapsed().as_nanos() as f64 / (iters as usize * text.len()) as f64); - } - best -} - -fn main() { - let manifest = env!("CARGO_MANIFEST_DIR"); - let english = - std::fs::read_to_string(format!("{manifest}/../data/big.txt")).unwrap_or_default(); - let english: String = english.chars().take(120_000).collect(); - - let astral = "😀🎉🚀🔥🌍🐍".repeat(20_000); // pure 4-byte astral → MB in every chunk - - // BMP English with an emoji inserted every ~24 bytes → some chunks have MB, some don't, irregularly. - let mut sprinkle = String::new(); - let mut acc = 0; - for w in english.split_inclusive(' ') { - sprinkle.push_str(w); - acc += w.len(); - if acc >= 24 { - sprinkle.push('🔥'); - acc = 0; - } - } - - println!( - "{:<10} {:>8} {:>10} {:>10}", - "input", "bytes", "clsSIMD", "clsScalar" - ); - for (label, s) in [ - ("none", &english), - ("astral", &astral), - ("sprinkle", &sprinkle), - ] { - if s.is_empty() { - println!("{label:<10} (empty — big.txt missing?)"); - continue; - } - let text = s.as_bytes(); - let mut tags = vec![0u8; text.len()]; - // byte-exactness guard: SIMD == scalar (exercises the astral fixup path) - let mut sc = vec![0u8; text.len()]; - classify(text, &mut tags); - classify_scalar(text, &mut sc); - let ok = if tags == sc { "✓" } else { "✗" }; - let simd = ns_per_byte(text, &mut tags); - let scal = { - let iters = (8_000_000 / text.len().max(1)).clamp(20, 400) as u32; - let mut best = f64::INFINITY; - for _ in 0..5 { - let t = Instant::now(); - for _ in 0..iters { - classify_scalar(text, &mut sc); - black_box(sc[text.len() / 2]); - } - best = - best.min(t.elapsed().as_nanos() as f64 / (iters as usize * text.len()) as f64); - } - best - }; - println!( - "{label:<10} {:>8} {simd:>10.3} {scal:>10.3} {ok}", - text.len() - ); - } - println!("\n(ns/byte, lower better. clsSIMD is the path with the MB-fixup branch.)"); -} diff --git a/tokenizers/atomsplit/benches/data/.gitignore b/tokenizers/atomsplit/benches/data/.gitignore deleted file mode 100644 index 2211df63d..000000000 --- a/tokenizers/atomsplit/benches/data/.gitignore +++ /dev/null @@ -1 +0,0 @@ -*.txt diff --git a/tokenizers/atomsplit/benches/data/fetch.py b/tokenizers/atomsplit/benches/data/fetch.py deleted file mode 100644 index 15370c927..000000000 --- a/tokenizers/atomsplit/benches/data/fetch.py +++ /dev/null @@ -1,17 +0,0 @@ -#!/usr/bin/env python3 -"""Fetch big real Wikipedia article text per language for the cl100k bench (gitignored output). -English/Japanese use the repo's ../data/big.txt + unigram_wagahaiwa_nekodearu.txt instead.""" -import urllib.request, urllib.parse, json, os -ARTS = [("fr","France"),("ru","Россия"),("el","Ελλάδα"),("he","ישראל"),("ar","مصر"), - ("hi","भारत"),("th","ประเทศไทย"),("zh","数学"),("ko","대한민국")] -here = os.path.dirname(os.path.abspath(__file__)) -for lang, title in ARTS: - url = f"https://{lang}.wikipedia.org/w/api.php?" + urllib.parse.urlencode( - {"action":"query","format":"json","prop":"extracts","explaintext":"1","titles":title}) - try: - req = urllib.request.Request(url, headers={"User-Agent":"atomsplit-bench/0.1"}) - text = next(iter(json.load(urllib.request.urlopen(req, timeout=30))["query"]["pages"].values())).get("extract","") - open(f"{here}/{lang}.txt","w").write(text) - print(f"{lang}: {len(text.encode())} bytes") - except Exception as e: - print(f"{lang}: FAILED {e}") diff --git a/tokenizers/atomsplit/benches/heatmap.py b/tokenizers/atomsplit/benches/heatmap.py deleted file mode 100644 index fb51991b1..000000000 --- a/tokenizers/atomsplit/benches/heatmap.py +++ /dev/null @@ -1,74 +0,0 @@ -#!/usr/bin/env python3 -"""Heatmap of the atomsplit `regex` bench: our classify+fsm pipeline speedup vs each SOTA engine, -per pre-tokenizer × language. Green = we win big, red = close race / behind.""" -import re -import sys -import numpy as np -import matplotlib -matplotlib.use("Agg") -import matplotlib.pyplot as plt -from matplotlib.colors import LogNorm - -LOG = sys.argv[1] -OUT = sys.argv[2] - -ENGINES = ["gpt2", "cl100k", "o200k", "deepseek", "ws_split", "whitespace", "digits", "punct", "bert"] -REFS = ["onig", "fancy", "logos", "pcre2"] -# discover languages in first-seen order -langs = [] -# data[pretok][lang] = [vsOnig, vsFncy, vsLogos, vsPcre2] -data = {e: {} for e in ENGINES} - -for line in open(LOG): - toks = [t for t in line.split() if t != "|"] - if len(toks) < 15 or toks[0] not in data: - continue - pre, lang = toks[0], toks[1] - # ...clsSIMD clsScal fsm onig fancy logos pcre2 | vsOnig vsFncy vsLogos vsPcre2 ok - vs = toks[11:15] - def num(x): - x = x.rstrip("x") - return np.nan if x in ("—", "-", "") else float(x) - data[pre][lang] = [num(v) for v in vs] - if lang not in langs: - langs.append(lang) - -cmap = plt.get_cmap("RdYlGn").copy() -cmap.set_bad("#d9d9d9") -norm = LogNorm(vmin=0.8, vmax=60) - -fig, axes = plt.subplots(3, 3, figsize=(15.5, 15), constrained_layout=True) -for ax, pre in zip(axes.flat, ENGINES): - M = np.array([data[pre].get(l, [np.nan] * 4) for l in langs]) # [langs x refs] - ax.imshow(M, aspect="auto", cmap=cmap, norm=norm) - ax.set_title(pre, fontsize=13, fontweight="bold") - ax.set_xticks(range(len(REFS))) - ax.set_xticklabels(REFS, fontsize=9) - ax.set_yticks(range(len(langs))) - ax.set_yticklabels(langs, fontsize=8) - ax.tick_params(length=0) - for i in range(len(langs)): - for j in range(len(REFS)): - v = M[i, j] - if np.isnan(v): - ax.text(j, i, "n/a", ha="center", va="center", fontsize=7, color="#666") - else: - # white text on the dark-red/dark-green extremes, black in the middle - t = norm(v) - col = "white" if (t < 0.18 or t > 0.86) else "black" - ax.text(j, i, f"{v:.0f}×" if v >= 10 else f"{v:.1f}×", - ha="center", va="center", fontsize=7.5, color=col, fontweight="bold") - -sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm) -cb = fig.colorbar(sm, ax=axes, shrink=0.5, aspect=30, pad=0.02, - ticks=[1, 2, 5, 10, 20, 50]) -cb.ax.set_yticklabels(["1× (tie)", "2×", "5×", "10×", "20×", "50×"]) -cb.set_label("speedup = engine ÷ our pipeline (SIMD classify + scalar fsm) · green = we win big", - fontsize=10) - -fig.suptitle("atomsplit pre-tokenization: our classify+fsm pipeline vs SOTA splitters\n" - "(aarch64 / Apple Silicon, release; onig & pcre2-JIT = C, fancy = pure-Rust regex, " - "logos = compile-time DFA lexer; n/a = engine can't express that split)", - fontsize=13, fontweight="bold") -fig.savefig(OUT, format="svg", bbox_inches="tight") -print(f"wrote {OUT}: {len(langs)} langs × {len(REFS)} engines × {len(ENGINES)} pre-tokenizers") diff --git a/tokenizers/atomsplit/benches/pretok_heatmap.svg b/tokenizers/atomsplit/benches/pretok_heatmap.svg deleted file mode 100644 index e09d0782b..000000000 --- a/tokenizers/atomsplit/benches/pretok_heatmap.svg +++ /dev/null @@ -1,8344 +0,0 @@ - - - - - - - - 2026-07-15T17:38:52.195573 - image/svg+xml - - - Matplotlib v3.10.6, https://matplotlib.org/ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/tokenizers/atomsplit/benches/regex.rs b/tokenizers/atomsplit/benches/regex.rs deleted file mode 100644 index cd9552184..000000000 --- a/tokenizers/atomsplit/benches/regex.rs +++ /dev/null @@ -1,429 +0,0 @@ -//! Pre-tokenization on BIG real text (Wikipedia / big.txt), per language, for every pre-tokenizer at -//! once — the GPT regex FSMs (gpt2 · cl100k · o200k · tekken · deepseek) and the class family (WhitespaceSplit · -//! Whitespace · Digits · Punctuation · Bert). Per (engine, corpus): classify (SIMD vs scalar) + fsm in -//! ns/byte, and the full pipeline (SIMD classify + scalar fsm) vs FOUR reference engines — onig (C), -//! fancy-regex (pure Rust), logos (compile-time DFA lexer), pcre2 (C + JIT). GPT regexes come from -//! [`atomsplit::regexes`] (the canonical specs, byte-exact ✓/✗); class-family references are the -//! best-effort equivalent regex (approximate — see the ≈ marker). Single-regex engines have a 1-element -//! chain; deepseek composes 3 `Isolated` Splits. -//! -//! Data: `../data/big.txt` (English) + `../data/unigram_wagahaiwa_nekodearu.txt` (Japanese) ship with -//! the repo; the rest via `benches/data/fetch.py` (gitignored). Missing files are skipped. -//! -//! Run: cargo bench --bench regex -use atomsplit::classify::{classify, classify_scalar, mask}; -use atomsplit::fsm::{ - Span, class_runs_into, fsm_byte_level, fsm_cl100k, fsm_deepseek, fsm_o200k, fsm_tekken, -}; -use atomsplit::regexes; -use fancy_regex::Regex as Fancy; -use logos::Logos; -use onig::Regex; -use pcre2::bytes::Regex as Pcre2; -use std::hint::black_box; -use std::time::Instant; - -type Fsm = fn(&[u8], &[u8], &mut [Span]) -> usize; - -// logos DFA lexers approximating the GPT splits. logos has no look-ahead (`(?!\S)`) nor case-insensitive -// (`(?i:)`), so token boundaries differ slightly — this is a raw-throughput reference (like fancy), not a -// byte-exact oracle (that's onig). deepseek is a 3-split `Sequence`, not one grammar → no logos number. -#[derive(Logos)] -enum LGpt2 { - #[regex(r"'s|'t|'re|'ve|'m|'ll|'d")] - Contraction, - #[regex(r" ?\p{L}+")] - Word, - #[regex(r" ?\p{N}+")] - Num, - #[regex(r" ?[^\s\p{L}\p{N}]+")] - Other, - #[regex(r"\s+")] - Space, -} - -#[derive(Logos)] -enum LCl100k { - // Contraction outranks Word: `[^\r\n\p{L}\p{N}]?\p{L}+` can also match `'s`, but the real regex tries - // the contraction alternative first. - #[regex(r"'s|'t|'re|'ve|'m|'ll|'d", priority = 5)] - Contraction, - #[regex(r"[^\r\n\p{L}\p{N}]?\p{L}+", priority = 4)] - Word, - #[regex(r"\p{N}\p{N}?\p{N}?")] - Num, - #[regex(r" ?[^\s\p{L}\p{N}]+[\r\n]*", priority = 2)] - Other, - #[regex(r"\s+")] - Space, -} - -#[derive(Logos)] -enum LO200k { - #[regex(r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+('s|'t|'re|'ve|'m|'ll|'d)?", priority = 6)] - LettersA, - #[regex(r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*('s|'t|'re|'ve|'m|'ll|'d)?", priority = 5)] - LettersB, - #[regex(r"\p{N}\p{N}?\p{N}?")] - Num, - #[regex(r" ?[^\s\p{L}\p{N}]+[\r\n/]*", priority = 2)] - Other, - #[regex(r"\s+")] - Space, -} - -// tekken = LO200k with no contraction suffix and one token per digit. -#[derive(Logos)] -enum LTekken { - #[regex( - r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+", - priority = 6 - )] - LettersA, - #[regex( - r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*", - priority = 5 - )] - LettersB, - #[regex(r"\p{N}")] - Num, - #[regex(r" ?[^\s\p{L}\p{N}]+[\r\n/]*", priority = 2)] - Other, - #[regex(r"\s+")] - Space, -} - -// class-family logos grammars. The whitespace-dropping ones use `skip` so logos never emits ws tokens. -#[derive(Logos)] -#[logos(skip r"\s+")] -enum LWsSplit { - #[regex(r"\S+")] - Run, -} -#[derive(Logos)] -#[logos(skip r"\s+")] -enum LWhitespace { - #[regex(r"\w+")] - Word, - #[regex(r"[^\w\s]+")] - Sym, -} -#[derive(Logos)] -enum LDigits { - #[regex(r"\p{N}+")] - Num, - #[regex(r"[^\p{N}]+")] - Other, -} - -fn lex_count<'s, T: Logos<'s, Source = str>>(s: &'s str) -> usize -where - T::Extras: Default, -{ - let mut lex = T::lexer(s); - let mut n = 0; - while lex.next().is_some() { - n += 1; - } - n -} - -/// logos throughput (ns/byte) for the engines it can express; `None` for the ones it can't (deepseek's -/// multi-split, and punct/bert whose POSIX-class isolation logos doesn't parse reliably). -fn logos_ns(ename: &str, s: &str, len: usize, iters: u32) -> Option { - let f: fn(&str) -> usize = match ename { - "gpt2" => |s| lex_count::(s), - "cl100k" => |s| lex_count::(s), - "o200k" => |s| lex_count::(s), - "tekken" => |s| lex_count::(s), - "ws_split" => |s| lex_count::(s), - "whitespace" => |s| lex_count::(s), - "digits" => |s| lex_count::(s), - _ => return None, - }; - Some(ns_per_byte(len, iters, || f(s))) -} - -// ── class-family pre-tokenizers: `class_runs_into` recipes + their equivalent reference regex ── -// Each is a monomorphized `Fsm` wrapper (const-generic masks can't go through a fn pointer directly). -fn f_ws_split(t: &[u8], g: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<{ mask::WS }, 0, 0>(t, g, o) -} -fn f_punct(t: &[u8], g: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<0, { mask::PUNCT }, 0>(t, g, o) -} -fn f_digits(t: &[u8], g: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<0, 0, { mask::NUMERIC }>(t, g, o) -} -fn f_whitespace(t: &[u8], g: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<{ mask::WS }, 0, { mask::WORD }>(t, g, o) -} -fn f_bert(t: &[u8], g: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<{ mask::WS }, { mask::PUNCT }, 0>(t, g, o) -} - -// Class-family reference regexes — the exact split each recipe produces under a real regex engine. -// The char-classes map 1:1 to the atomsplit masks: WS=`\s`, WORD=`\w`, NUMERIC=`\p{N}`, -// PUNCT=`[[:punct:]\p{P}]` (ASCII-punctuation ∪ Unicode-P). Isolate = the punct alternative matches a -// single char; the run alternatives use `+`. -const RX_WS_SPLIT: &str = r"\S+"; -const RX_WHITESPACE: &str = r"\w+|[^\w\s]+"; -const RX_DIGITS: &str = r"\p{N}+|[^\p{N}]+"; -const RX_PUNCT: &str = r"[[:punct:]\p{P}]|[^[:punct:]\p{P}]+"; -const RX_BERT: &str = r"[[:punct:]\p{P}]|[^\s[:punct:]\p{P}]+"; - -// (name, native fsm, reference regex chain, keep_gaps, exact). The chain is applied `Isolated`, each -// regex splitting the previous pieces. `keep_gaps=false` for the whitespace-dropping recipes: their -// output is matches only (dropped whitespace is NOT a token), so the reference drops the gaps too. -// `exact`: the GPT FSMs are byte-exact with their regex (parity-tested) → ✓/✗ is a real gate. The -// class-family masks (WORD incl. Other_Alphabetic, PUNCT = P ∪ ASCII-symbols, …) can't be written as a -// regex class, so their reference is an *approximation* — a mismatch is `≈` (still a valid speed pairing), -// not a failure. -const ENGINES: &[(&str, Fsm, &[&str], bool, bool)] = &[ - ("gpt2", fsm_byte_level as Fsm, &[regexes::GPT2], true, true), - ("cl100k", fsm_cl100k as Fsm, &[regexes::CL100K], true, true), - ("o200k", fsm_o200k as Fsm, &[regexes::O200K], true, true), - ("tekken", fsm_tekken as Fsm, &[regexes::TEKKEN], true, true), - ( - "deepseek", - fsm_deepseek as Fsm, - regexes::DEEPSEEK, - true, - true, - ), - ("ws_split", f_ws_split as Fsm, &[RX_WS_SPLIT], false, false), - ( - "whitespace", - f_whitespace as Fsm, - &[RX_WHITESPACE], - false, - false, - ), - ("digits", f_digits as Fsm, &[RX_DIGITS], true, false), - ("punct", f_punct as Fsm, &[RX_PUNCT], true, false), - ("bert", f_bert as Fsm, &[RX_BERT], false, false), -]; - -const CORPORA: &[(&str, &str)] = &[ - ("English", "../data/big.txt"), - ("French", "benches/data/fr.txt"), - ("Russian", "benches/data/ru.txt"), - ("Greek", "benches/data/el.txt"), - ("Hebrew", "benches/data/he.txt"), - ("Arabic", "benches/data/ar.txt"), - ("Hindi", "benches/data/hi.txt"), - ("Thai", "benches/data/th.txt"), - ("Chinese", "benches/data/zh.txt"), - ("Japanese", "../data/unigram_wagahaiwa_nekodearu.txt"), - ("Korean", "benches/data/ko.txt"), -]; - -/// The composed split under `res` — each regex splits the previous pieces. `keep_gaps` keeps the -/// between/after-match gaps as pieces (Isolated split; deepseek); `false` drops them (the whitespace a -/// drop-ws recipe discards). Byte-offset pieces into `text`. -fn onig_pieces(res: &[Regex], text: &str, keep_gaps: bool) -> Vec<(usize, usize)> { - let mut pieces = vec![(0usize, text.len())]; - for re in res { - let mut next = Vec::with_capacity(pieces.len() * 2); - for (s, e) in pieces.drain(..) { - let sub = &text[s..e]; - let mut prev = 0usize; - for (ms, me) in re.find_iter(sub) { - if keep_gaps && ms > prev { - next.push((s + prev, s + ms)); - } - next.push((s + ms, s + me)); - prev = me; - } - if keep_gaps && prev < sub.len() { - next.push((s + prev, e)); - } - } - pieces = next; - } - pieces -} - -/// Same composition under PCRE2 (JIT-compiled), operating on bytes; `find_iter` yields `Result`. -fn pcre2_pieces(res: &[Pcre2], text: &str, keep_gaps: bool) -> Vec<(usize, usize)> { - let bytes = text.as_bytes(); - let mut pieces = vec![(0usize, text.len())]; - for re in res { - let mut next = Vec::with_capacity(pieces.len() * 2); - for (s, e) in pieces.drain(..) { - let sub = &bytes[s..e]; - let mut prev = 0usize; - for m in re.find_iter(sub) { - let Ok(m) = m else { break }; - let (ms, me) = (m.start(), m.end()); - if keep_gaps && ms > prev { - next.push((s + prev, s + ms)); - } - next.push((s + ms, s + me)); - prev = me; - } - if keep_gaps && prev < sub.len() { - next.push((s + prev, e)); - } - } - pieces = next; - } - pieces -} - -/// Same composition under fancy-regex; `find_iter` yields `Result` (a match error ends the pass). -fn fancy_pieces(res: &[Fancy], text: &str, keep_gaps: bool) -> Vec<(usize, usize)> { - let mut pieces = vec![(0usize, text.len())]; - for re in res { - let mut next = Vec::with_capacity(pieces.len() * 2); - for (s, e) in pieces.drain(..) { - let sub = &text[s..e]; - let mut prev = 0usize; - for m in re.find_iter(sub) { - let Ok(m) = m else { break }; - let (ms, me) = (m.start(), m.end()); - if keep_gaps && ms > prev { - next.push((s + prev, s + ms)); - } - next.push((s + ms, s + me)); - prev = me; - } - if keep_gaps && prev < sub.len() { - next.push((s + prev, e)); - } - } - pieces = next; - } - pieces -} - -// MIN over TRIALS timed loops — the fastest trial had the least CPU contention, so it's the truest -// estimate and robust to thermal throttling / background load (which only ever make a trial slower). -fn ns_per_byte usize>(len: usize, iters: u32, mut f: F) -> f64 { - const TRIALS: u32 = 7; - for _ in 0..3 { - black_box(f()); // warm - } - let mut best = f64::INFINITY; - for _ in 0..TRIALS { - let t = Instant::now(); - let mut acc = 0usize; - for _ in 0..iters { - acc = acc.wrapping_add(f()); - } - black_box(acc); - best = best.min(t.elapsed().as_nanos() as f64 / (iters as usize * len) as f64); - } - best -} - -fn main() { - let manifest = env!("CARGO_MANIFEST_DIR"); - println!( - "{:<9} {:<10} {:>7} {:>5} {:>8} {:>8} | {:>8} | {:>8} {:>8} {:>8} {:>8} | {:>7} {:>7} {:>7} {:>7}", - "engine", - "lang", - "bytes", - "b/tok", - "clsSIMD", - "clsScal", - "fsm", - "onig", - "fancy", - "logos", - "pcre2jit", - "vsOnig", - "vsFncy", - "vsLogos", - "vsPcre2" - ); - for &(ename, fsm, rxs, keep_gaps, exact) in ENGINES { - let onig: Vec = rxs.iter().map(|r| Regex::new(r).expect(ename)).collect(); - let fancy: Vec = rxs.iter().map(|r| Fancy::new(r).expect(ename)).collect(); - // PCRE2 with JIT — its speed is the JIT, so bench it there. - let pcre2: Vec = rxs - .iter() - .map(|r| { - pcre2::bytes::RegexBuilder::new() - .utf(true) - .ucp(true) - .jit_if_available(true) - .build(r) - .expect(ename) - }) - .collect(); - for (label, rel) in CORPORA { - let raw = match std::fs::read_to_string(format!("{manifest}/{rel}")) { - Ok(s) if !s.trim().is_empty() => s, - _ => { - println!( - "{ename:<9} {label:<10} (skipped — {rel} missing; run benches/data/fetch.py)" - ); - continue; - } - }; - // UNIFORM cap (char boundary): every language the same byte size → equal cache behaviour, so - // per-byte compute compares across scripts. 180 KB > L1, and all corpora have ≥180 KB. - let mut c = raw.len().min(180_000); - while c > 0 && !raw.is_char_boundary(c) { - c -= 1; - } - let corpus = &raw[..c]; - let text = corpus.as_bytes(); - let n = text.len(); - let iters = (4_000_000 / n).clamp(3, 150) as u32; - - let ref_spans: Vec = onig_pieces(&onig, corpus, keep_gaps) - .iter() - .map(|&(s, e)| Span::new(s as u32, e as u32)) - .collect(); - let mut tags = vec![0u8; n]; - let mut tsc = vec![0u8; n]; - classify(text, &mut tags); - let mut buf = vec![Span::default(); n + 1]; - let k = fsm(text, &tags, &mut buf); - let ok = if buf[..k] == ref_spans[..] { - "✓" - } else if exact { - "✗" // a real regression: the GPT fsm must be byte-exact with its regex - } else { - "≈" // class-family: regex is an approximation of the mask, divergence expected - }; - let btok = n as f64 / k.max(1) as f64; - - let cls_simd = ns_per_byte(n, iters, || { - classify(text, &mut tags); - tags[n / 2] as usize - }); - let cls_scal = ns_per_byte(n, iters, || { - classify_scalar(text, &mut tsc); - tsc[n / 2] as usize - }); - classify(text, &mut tags); - let fsm_ns = ns_per_byte(n, iters, || fsm(text, &tags, &mut buf)); - let onig_ns = ns_per_byte(n, iters, || onig_pieces(&onig, corpus, keep_gaps).len()); - let fancy_ns = ns_per_byte(n, iters, || fancy_pieces(&fancy, corpus, keep_gaps).len()); - let pcre2_ns = ns_per_byte(n, iters, || pcre2_pieces(&pcre2, corpus, keep_gaps).len()); - let logos = logos_ns(ename, corpus, n, iters); - - let pipe = cls_simd + fsm_ns; // SIMD classify + scalar fsm — the full pipeline - let (logos_c, vslogos) = match logos { - Some(l) => (format!("{l:8.2}"), format!("{:6.1}x", l / pipe)), - None => (" —".into(), " —".into()), // deepseek: no single logos grammar - }; - println!( - "{ename:<9} {label:<10} {n:>7} {btok:>5.1} {cls_simd:>8.3} {cls_scal:>8.3} | {fsm_ns:>8.3} | {onig_ns:>8.2} {fancy_ns:>8.2} {logos_c} {pcre2_ns:>8.2} | {:>6.1}x {:>6.1}x {vslogos} {:>6.1}x {ok}", - onig_ns / pipe, - fancy_ns / pipe, - pcre2_ns / pipe - ); - } - } - println!( - "\n(ns/byte, lower better. pipeline = SIMD classify + scalar fsm; vs onig / fancy / logos / \ - pcre2(JIT). GPT FSMs (gpt2/cl100k/o200k/tekken/deepseek): the regex IS the spec, ✓ = byte-exact, \ - ✗ = regression. Class family (ws_split/whitespace/digits/punct/bert): the mask can't be written \ - as a regex class, so the reference is approximate — ✓ where it happens to match, ≈ where it \ - diverges (speed still comparable). logos approximates the grammar (deepseek/punct/bert: n/a).)" - ); -} diff --git a/tokenizers/atomsplit/src/fsm.rs b/tokenizers/atomsplit/src/fsm.rs deleted file mode 100644 index 45c8ca286..000000000 --- a/tokenizers/atomsplit/src/fsm.rs +++ /dev/null @@ -1,387 +0,0 @@ -//! FSM layer: turn the `Atom` tag stream (from [`crate::classify`]) into token spans. Every fsm is -//! NO-PUSH — it writes spans into a caller-preallocated `&mut [Span]` (len ≥ `text.len()`) and returns -//! the token count; no `Vec`, no realloc. Inputs must be well-formed UTF-8 (see the crate-level docs). -//! -//! The class family (WhitespaceSplit / Punctuation / Digits / Whitespace / Bert) goes through -//! [`class_runs_into`]: on aarch64/wasm the SIMD movemask boundary-extractor + homogeneous-chunk -//! early-out (in `simd_fsm`), elsewhere the scalar run-end core ([`emit_class_spans`]). The -//! regex-shaped ones ([`fsm_cl100k`] / [`fsm_o200k`] / [`fsm_tekken`] / [`fsm_deepseek`] / -//! [`fsm_byte_level`]) are scalar jump-tables (only the class family's [`class_runs_into`] has a SIMD -//! path). - -pub(crate) use crate::classify::{Atom, char_len, classify, in_mask, mask}; -// Atom-tag aliases, shared with the per-tokenizer FSM submodules (`fsm/*.rs`) via `use super::*`. -pub(crate) const LET: u8 = Atom::Letter as u8; -pub(crate) const NW: u8 = Atom::NumWord as u8; -pub(crate) const NO: u8 = Atom::NumOther as u8; -pub(crate) const NLN: u8 = Atom::Newline as u8; -pub(crate) const SPC: u8 = Atom::Space as u8; -pub(crate) const WSO: u8 = Atom::WsOther as u8; -pub(crate) const MRK: u8 = Atom::Mark as u8; -pub(crate) const CON: u8 = Atom::Connector as u8; -pub(crate) const PUN: u8 = Atom::Punct as u8; -pub(crate) const APO: u8 = Atom::Apostrophe as u8; -pub(crate) const SYM: u8 = Atom::SymOther as u8; -pub(crate) const NMO: u8 = Atom::NumericOther as u8; -pub(crate) const CTL: u8 = Atom::Control as u8; -pub(crate) const CONT: u8 = Atom::Cont as u8; -pub(crate) const ASM: u8 = Atom::AlphaSymMark as u8; -pub(crate) const ZWJ: u8 = Atom::Zwj as u8; // 0x26 — ZWJ/ZWNJ, tagged in classify so FSMs skip the text peek - -/// Advance over a maximal `m`-membership run (m is a mask); returns the byte index past it. -/// `inline(always)`: it's called once per token (~200K/MB on English) — a real call here doubles fsm cost. -/// -/// logos-style "fast loop": process 16 tags per iteration with ONE bounds check per chunk and unchecked -/// reads (the loop condition proves `i + 16 <= end`), so short runs pay ~1 bounds check instead of one -/// per byte and long runs stay a tight unrolled scan. Byte-identical to the plain `while in_mask` scan. -#[inline(always)] -pub(crate) fn run_end(tags: &[u8], mut i: usize, end: usize, mut m: u16) -> usize { - m |= Atom::Cont.bit(); - debug_assert!(end <= tags.len()); - // SAFETY: `i + 16 <= end <= tags.len()` in the unrolled body, so every `get_unchecked(i + k)` - // (k < 16) is in bounds. The tail is the plain checked scan. - while i + 16 <= end { - for k in 0..16 { - if !in_mask(unsafe { *tags.get_unchecked(i + k) }, m) { - return i + k; - } - } - i += 16; - } - while i < end && in_mask(tags[i], m) { - i += 1; - } - i -} - -/// End of a whitespace token starting at `i`, for the tail shared byte-for-byte by cl100k and o200k -/// (`\s*[\r\n]+ | \s+(?!\S) | \s+`): through the last `\r\n` if any (rule 5), else the whole run at -/// EOF (rule 7), else give the final ws char back to the next token (rule 6). `#[inline]` — hot, -/// called once per whitespace token. (deepseek's tail differs — it also stops before a digit/CJK — and -/// byte_level has no `[\r\n]` rule; both keep their own.) -#[inline] -pub(crate) fn ws_tail(text: &[u8], tags: &[u8], i: usize, end: usize) -> usize { - let re = run_end(tags, i, end, mask::WS); - if let Some(r) = text[i..re].iter().rposition(|&x| x == 0x0A || x == 0x0D) { - i + r + 1 - } else if re == end { - re - } else { - let mut last = re - 1; - while last > i && text[last] & 0xC0 == 0x80 { - last -= 1; - } - if last > i { last } else { re } - } -} - -/// Case-insensitive contraction match at `i` (`'s 't 're 've 'm 'll 'd`), shared by cl100k and o200k: -/// byte length (2 or 3) or 0 if none. Self-guarding — `text[i]` need not be an apostrophe. `#[inline]`. -/// (byte_level's contraction is case-SENSITIVE, so it keeps its own.) -#[inline] -pub(crate) fn contraction(text: &[u8], i: usize) -> usize { - let end = text.len(); - if i >= end || text[i] != 0x27 || i + 1 >= end || text[i + 1] >= 0x80 { - return 0; - } - let lc = text[i + 1] | 0x20; - match lc { - b's' | b't' | b'm' | b'd' => 2, - b'r' | b'v' | b'l' if i + 2 < end && text[i + 2] < 0x80 => { - let l2 = text[i + 2] | 0x20; - usize::from((matches!(lc, b'r' | b'v') && l2 == b'e') || (lc == b'l' && l2 == b'l')) * 3 - } - _ => 0, - } -} - -/// A token span: byte offsets `[start, end)` into the input. `#[repr(C)]` so the FSM output buffer has a -/// stable `[start, end]` layout — the pipeline reuses it with zero conversion, and it can be reinterpreted -/// as bytes / handed across the crate boundary. -#[repr(C)] -#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Hash, PartialOrd, Ord)] -pub struct Span { - pub start: u32, - pub end: u32, -} - -impl Span { - #[inline] - pub const fn new(start: u32, end: u32) -> Self { - Self { start, end } - } - - /// `[start, end)` as a `usize` range — for slicing the input text. - #[inline] - pub fn range(self) -> core::ops::Range { - self.start as usize..self.end as usize - } -} - -/// Compare against a bare `(start, end)` tuple — convenience for tests/interop. -impl PartialEq<(u32, u32)> for Span { - #[inline] - fn eq(&self, o: &(u32, u32)) -> bool { - self.start == o.0 && self.end == o.1 - } -} - -/// No-`push` class-family pre-tokenizer core: writes spans into the preallocated `out` slice and returns -/// the count. ONE shape covers the whole class family via ``: -/// WhitespaceSplit `<{WS},0,0>` · Punctuation `<0,{PUNCT},0>` · Digits `<0,0,{NUMERIC}>` · -/// Whitespace `<{WS},0,{WORD}>` · Bert `<{WS},{PUNCT},0>`. -/// Class of a char: `DROP`→dropped, `ISOLATE`→own token, `KEEP_A`→run "A", else→run "B" (A/B cut apart). -/// TODO: find a better explanation -#[inline] -#[must_use] -pub fn class_runs_into( - text: &[u8], - tags: &[u8], - out: &mut [Span], -) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - #[cfg(target_arch = "aarch64")] - { - crate::simd_fsm::class_runs_neon::(text, tags, out) - } - #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] - { - crate::simd_fsm::class_runs_wasm::(text, tags, out) - } - #[cfg(not(any( - target_arch = "aarch64", - all(target_arch = "wasm32", target_feature = "simd128") - )))] - { - emit_class_spans::(text, tags, out, 0, 0, 0, None) - } -} - -/// This is the most important function as it's the core of the scalar finite state machine. -/// It allows to emit class spans with different behaviours for tags we want to drop, tags we want -/// to isolate and tags we want to keep. Any other tags are assumed to be keept. -/// -/// This function is used as a fallback to the SIMD fast fsm. It is used for most pre tokenizers -/// but the unrolled regex, which have more complex variations that cannot be expressed with drop, -/// isolate, keep. These 3 generic parameters are u16 bitmap masks over the 16 classes we have and -/// define the behaviour. They are usally one of the [`crate::classify::mask`]. They allow dropping -/// words, isolating whitespace and keeping new line for example. -#[must_use] -#[inline] -pub fn emit_class_spans( - text: &[u8], - tags: &[u8], - out: &mut [Span], - mut write_index: usize, // in the out slice - mut text_pointer: usize, // in the text slice - segment_start: usize, // previous segment_start - segment_class: Option, // previous segment's class -) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - let n = text.len(); - // Tie `tags.len() == text.len() == n` so the optimizer drops the interior `tags[i]` / `text[i]` - // bounds checks (callers guarantee len ≥ n; same trick as `cl100k`). Per-byte scanning already - // avoids checks via `run_end`'s unrolled `get_unchecked`; this covers the per-token accesses. - let tags = &tags[..n]; - let text = &text[..n]; - let other = !(DROP | ISOLATE | KEEP_A); // None of the above correspond to a continuation - if let Some(segment_class) = segment_class { - // this will usually be at the tail of a SIMD call. - text_pointer = run_end(tags, text_pointer, n, segment_class); // skip the whole drop run at once - if segment_class != DROP { - out[write_index] = Span { - start: segment_start as u32, - end: text_pointer as u32, - }; - if text_pointer == n { - return write_index + 1; - } - write_index += 1; - } - } - while text_pointer < n { - let t = tags[text_pointer]; - if t == Atom::Cont as u8 { - text_pointer += 1; - continue; - } - // classify the first char. - if in_mask(t, DROP) { - text_pointer = run_end(tags, text_pointer, n, DROP); // skip the whole drop run at once - } else if in_mask(t, ISOLATE) { - let s = text_pointer; - text_pointer += char_len(text[text_pointer]); - out[write_index] = Span { - start: s as u32, - end: text_pointer as u32, - }; // isolate: one char = one token - write_index += 1; - } else { - let s = text_pointer; - text_pointer = if in_mask(t, KEEP_A) { - run_end(tags, text_pointer, n, KEEP_A) - } else { - run_end(tags, text_pointer, n, other) - }; - out[write_index] = Span { - start: s as u32, - end: text_pointer as u32, - }; - write_index += 1; - } - } - write_index -} - -// ── per-tokenizer unrolled FSMs (one file each; shared helpers above via `use super::*`) ── -mod byte_level; -mod cl100k; -mod deepseek; -mod o200k; -pub use byte_level::fsm_byte_level; -pub use cl100k::{fsm_cl100k, fsm_cl100k_cap}; -pub use deepseek::fsm_deepseek; -pub use o200k::{fsm_o200k, fsm_tekken}; - -// ── Composition recipes ──────────────────────────────────────────────────────────────────────── -// Each pre-tokenizer = (classify → fsm shape + params). `tags` and `out` are caller-owned -// scratch, reused across calls — NO per-call alloc, NO push. The class family writes spans into the -// preallocated `out: &mut [Span]` (len ≥ text.len()) via `class_runs_into` and returns the token count. -// In `tk-encode` these delegate from the `pipeline::PreTokenizer` impls (offset conversion happens there). - -/// `WhitespaceSplit` — split on Unicode whitespace and drop it; keeps maximal non-whitespace runs. -pub struct WhitespaceSplit; -impl WhitespaceSplit { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - class_runs_into::<{ mask::WS }, 0, 0>(text, tags, out) - } -} - -/// `Punctuation` — isolate each punctuation char as its own token; non-punct grouped into runs. -pub struct Punctuation; -impl Punctuation { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - class_runs_into::<0, { mask::PUNCT }, 0>(text, tags, out) - } -} - -/// `Digits` — cut numeric runs apart from non-numeric runs (contiguous), keeping both. -pub struct Digits; -impl Digits { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - class_runs_into::<0, 0, { mask::NUMERIC }>(text, tags, out) - } -} - -/// `Whitespace` — the `\w+|[^\w\s]+` pre-tokenizer: drop whitespace, cut word runs from symbol runs. -pub struct Whitespace; -impl Whitespace { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - class_runs_into::<{ mask::WS }, 0, { mask::WORD }>(text, tags, out) - } -} - -/// `Bert` — the BERT basic pre-tokenizer: drop whitespace, isolate punctuation, keep the rest as runs. -pub struct Bert; -impl Bert { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - class_runs_into::<{ mask::WS }, { mask::PUNCT }, 0>(text, tags, out) - } -} - -/// `Cl100k` — the tiktoken cl100k_base / Llama-3 pre-tokenizer (7-rule regex). (o200k is a distinct, -/// case-aware FSM — [`fsm_o200k`].) -pub struct Cl100k; -impl Cl100k { - /// Uses the scalar run-end core: cl100k's letter/ws runs are short on Latin/code (the common case), - /// where a SIMD run-end's setup would lose; only long CJK runs would benefit. - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - fsm_cl100k(text, tags, out) - } -} - -/// `DeepSeek` — the DeepSeek-V3/R1 pre-tokenizer (digits{1,3} → CJK-range → big regex, composed). -pub struct DeepSeek; -impl DeepSeek { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - fsm_deepseek(text, tags, out) - } -} - -/// `ByteLevel` — the GPT-2 / Llama / Qwen byte-level pre-tokenizer regex (before byte-mapping). -pub struct ByteLevel; -impl ByteLevel { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - fsm_byte_level(text, tags, out) - } -} - -/// `Split(char, Removed)` — the only pre-tokenizer that keys on a *literal char* rather than an atom -/// class, so it scans bytes directly (no classify pass). UTF-8 is self-synchronizing, so the -/// delimiter's byte pattern only matches on char boundaries. -pub struct CharDelimiterSplit(pub char); -impl CharDelimiterSplit { - /// Split on the literal char (Removed); writes spans into `out` (len ≥ `text.len()`), returns count. - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], _tags: &mut [u8], out: &mut [Span]) -> usize { - debug_assert!(out.len() >= text.len()); - let mut buf = [0u8; 4]; - let delim = self.0.encode_utf8(&mut buf).as_bytes(); - let (n, dl) = (text.len(), delim.len()); - let (mut start, mut i, mut w) = (0usize, 0usize, 0usize); - while i + dl <= n { - // memchr the first delimiter byte, then confirm the full pattern. memchr (already a - // workspace dep) beats a scalar scan 1.4–23× here — the gap widening as the delimiter - // gets rarer over large inputs, since its SIMD skips whole 16/32/64-byte strides. - match memchr::memchr(delim[0], &text[i..n - dl + 1]) { - Some(off) if text[i + off..i + off + dl] == *delim => { - let m = i + off; - if m > start { - out[w] = Span { - start: start as u32, - end: m as u32, - }; // gap before the delimiter (Removed) - w += 1; - } - i = m + dl; - start = i; - } - Some(off) => i += off + 1, // first byte matched mid-pattern; keep scanning - None => break, - } - } - if start < n { - out[w] = Span { - start: start as u32, - end: n as u32, - }; - w += 1; - } - w - } -} diff --git a/tokenizers/atomsplit/src/fsm/byte_level.rs b/tokenizers/atomsplit/src/fsm/byte_level.rs deleted file mode 100644 index 73f84ac9f..000000000 --- a/tokenizers/atomsplit/src/fsm/byte_level.rs +++ /dev/null @@ -1,81 +0,0 @@ -use super::*; - -/// GPT-2 / ByteLevel pretokenization. Regex (same jump-table shape as cl100k, with 3 differences): -/// `'s|'t|'re|'ve|'m|'ll|'d | ?\p{L}+ | ?\p{N}+ | ?[^\s\p{L}\p{N}]+ | \s+(?!\S) | \s+` -/// vs cl100k: (1) contractions are case-SENSITIVE (lowercase only, no `(?i:)`); (2) the ` ?` prefix is -/// a literal SPACE only (not any non-l/n char) and it applies to letters, numbers AND "other"; (3) no -/// `\p{N}{1,3}` cap (numbers are unbounded) and no `\s*[\r\n]`/trailing-`[\r\n]*` rules. -#[must_use] -pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - let end = text.len(); - // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior - // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) - let tags = &tags[..end]; - // `\s+(?!\S)|\s+`: the whole run at EOF, else leave the last ws char for the next ` ?`-prefixed run. - let ws = |i: usize| -> usize { - let re = run_end(tags, i, end, mask::WS); - if re == end { - re - } else { - let mut last = re - 1; - while last > i && text[last] & 0xC0 == 0x80 { - last -= 1; - } - if last > i { last } else { re } - } - }; - - let mut i = 0; - let mut w = 0usize; - while i < end { - let start = i; - match tags[i] & 0x0F { - LET => i = run_end(tags, i, end, mask::LETTER), // ` ?\p{L}+` (space taken by the Space arm) - NW | NO => i = run_end(tags, i, end, mask::NUMBER), // ` ?\p{N}+` — UNBOUNDED - MRK | CON | PUN | SYM | NMO | CTL => i = run_end(tags, i, end, mask::NOT_WS_L_N), // ` ?[^…]+` - // `'s|'t|'re|'ve|'m|'ll|'d` (case-sensitive), else `[^\s\p{L}\p{N}]+` (apostrophe ∈ that set) - APO => { - let adv = match (text.get(i + 1), text.get(i + 2)) { - (Some(b's' | b't' | b'm' | b'd'), _) => 2, - (Some(b'r'), Some(b'e')) - | (Some(b'v'), Some(b'e')) - | (Some(b'l'), Some(b'l')) => 3, - _ => 0, - }; - i = if adv > 0 { - i + adv - } else { - run_end(tags, i, end, mask::NOT_WS_L_N) - }; - } - // Space: the ` ?` prefix — attach one space to a following letter / number / "other" run, - // else it's whitespace (rules `\s+(?!\S)|\s+`, which leave one space for the next run). - SPC => { - let a = i + 1; // Space is ASCII (0x20) - i = match tags.get(a).map(|&t| t & 0x0F) { - Some(LET) => run_end(tags, a, end, mask::LETTER), - Some(NW) | Some(NO) => run_end(tags, a, end, mask::NUMBER), - Some(t) if in_mask(t, mask::NOT_WS_L_N) => { - run_end(tags, a, end, mask::NOT_WS_L_N) - } - _ => ws(i), - }; - } - // WsOther / Newline: whitespace only — the ` ?` prefix is a literal 0x20, so tabs/newlines - // never prefix a run. - WSO | NLN => i = ws(i), - // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. - _ => i += char_len(text[i]), - } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; - } - w -} diff --git a/tokenizers/atomsplit/src/fsm/cl100k.rs b/tokenizers/atomsplit/src/fsm/cl100k.rs deleted file mode 100644 index 65976ecef..000000000 --- a/tokenizers/atomsplit/src/fsm/cl100k.rs +++ /dev/null @@ -1,118 +0,0 @@ -use super::*; - -/// TODO: NONE OF THE FOLLOWING HAS BEEN REVIEWED -/// cl100k / Llama-3 pretokenization (7 rules + whitespace-tail, rule-3 cap `\p{N}{1,3}`). Peeks `text` -/// for the ASCII contraction literals. Scalar run-ends. See [`fsm_cl100k_cap`] for the variable-cap -/// family (Qwen2 etc.). ┌── OWNER: shared (scalar) ──┐ -#[must_use] -pub fn fsm_cl100k(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - fsm_cl100k_cap(text, tags, out, 3) -} - -/// The cl100k family with an explicit rule-3 digit cap: 3 = cl100k / Llama-3, 1 = Qwen2's `\p{N}` (each -/// digit its own token), `usize::MAX` = an unbounded `\p{N}+`. Only the digit rule differs across these. -#[must_use] -pub fn fsm_cl100k_cap(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usize) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - cl100k(text, tags, out, digit_cap) -} - -fn cl100k(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usize) -> usize { - // Leading-atom values, as `const` so the `match` below is a dense jump table (not an if-cascade): - // the dispatch is O(1) and a token never pays for a rule it can't start (e.g. non-number tokens - // never test the number rule — which is what the POC's const-gating removed by hand; here it's free). - let end = text.len(); - // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior - // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) - let tags = &tags[..end]; - let letters = |a: usize| run_end(tags, a, end, mask::LETTER); - // rule 4 body: `[^\s\p{L}\p{N}]+[\r\n]*` from `sp0` (any leading space already consumed). Returns - // the run end, or `sp0` if there is no "other" run there (caller then treats it as whitespace). - let other = |sp0: usize| -> usize { - let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); - if p > sp0 { - while p < end && tags[p] == NLN { - p += char_len(text[p]); - } - } - p - }; - // rules 5-7 (`\s*[\r\n] | \s+(?!\S) | \s+`) → the shared `ws_tail`. - let ws = |i: usize| -> usize { ws_tail(text, tags, i, end) }; - - let mut i = 0; - let mut w = 0usize; - while i < end { - let start = i; - let b = text[i]; - match tags[i] & 0x0F { - // rule 2: `\p{L}+` - LET => i = letters(i), - // rule 3: `\p{N}{1,cap}` (cap = 3 cl100k, 1 Qwen2, MAX for `\p{N}+`) - NW | NO => { - let (mut p, mut cnt) = (i, 0); - while p < end && cnt < digit_cap && in_mask(tags[p], mask::NUMBER) { - p += char_len(text[p]); - cnt += 1; - } - i = p; - } - // Space: rule 2 (space prefix + `\p{L}+`) | rule 4 (` ` + "other") | rules 5-7 - SPC => { - let a = i + 1; // Space is ASCII (0x20) - i = if a < end && (tags[a] & 0x0F) == LET { - letters(a) - } else { - let p = other(a); - if p > a { p } else { ws(i) } - }; - } - // WsOther: rule 2 (prefix + `\p{L}+`) | whitespace (never rule 4 — not in NOT_WS_L_N) - WSO => { - let a = i + char_len(b); - i = if a < end && (tags[a] & 0x0F) == LET { - letters(a) - } else { - ws(i) - }; - } - // Newline: whitespace (rule 5 ends at the last newline) - NLN => i = ws(i), - // Apostrophe: rule 1 (contraction) | rule 2 (prefix + `\p{L}+`) | rule 4 - APO => { - let adv = contraction(text, i); // rule 1: `'s 't 're 've 'm 'll 'd` (case-insensitive) - i = if adv > 0 { - i + adv - } else { - let a = i + 1; // Apostrophe is ASCII (0x27) - if a < end && (tags[a] & 0x0F) == LET { - letters(a) - } else { - other(i) - } // c ∈ NOT_WS_L_N ⇒ > i - }; - } - // Mark | Connector | Punct | SymOther | NumericOther | Control (all in NOT_WS_L_N): - // rule 2 (prefix + `\p{L}+`) | rule 4 - MRK | CON | PUN | SYM | NMO | CTL => { - let a = i + char_len(b); - i = if a < end && (tags[a] & 0x0F) == LET { - letters(a) - } else { - other(i) - }; // c ∈ NOT_WS_L_N ⇒ > i - } - // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. - _ => i += char_len(b), - } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; - } - w -} diff --git a/tokenizers/atomsplit/src/fsm/deepseek.rs b/tokenizers/atomsplit/src/fsm/deepseek.rs deleted file mode 100644 index f54afb8fa..000000000 --- a/tokenizers/atomsplit/src/fsm/deepseek.rs +++ /dev/null @@ -1,247 +0,0 @@ -use super::*; - -/// The specific CJK ranges deepseek's Split-2 isolates: Han U+4E00..9FA5 ∪ Hiragana U+3040..309F ∪ -/// Katakana U+30A0..30FF. (All 3-byte, leads E3..E9.) Not "all letters" — only these. -#[inline] -fn ds_is_cjk(cp: u32) -> bool { - (0x4E00..=0x9FA5).contains(&cp) || (0x3040..=0x30FF).contains(&cp) -} -/// Codepoint of a 3-byte UTF-8 char at `text[i]` (only called for leads E2..E9 → always 3-byte). -#[inline] -fn cp3(text: &[u8], i: usize) -> u32 { - ((text[i] as u32 & 0x0F) << 12) - | ((text[i + 1] as u32 & 0x3F) << 6) - | (text[i + 2] as u32 & 0x3F) -} - -/// Any CJK-range char (letter OR punct/sym) at `p` — the full `[一-龥぀-ゟ゠-ヿ]` set Split-2 `[…]+` -/// isolates, INCLUDING CJK punctuation (・ U+30FB, ゠ U+30A0, ゛゜ U+309B/C). Split-3 then re-splits that -/// isolated run into same-kind sub-runs (letters `[\p{L}\p{M}]+` vs punct/sym `[\p{P}\p{S}]+`); because -/// the run is a CLOSED unit, none of it steals a surrounding space or merges with non-CJK punct. -#[inline(always)] -fn ds_is_cjk_at(text: &[u8], p: usize) -> bool { - (0xE3..=0xE9).contains(&text[p]) && ds_is_cjk(cp3(text, p)) -} - -/// deepseek-v3 pretokenization: the `Sequence` of `[N{1,3}]`, `[CJK]+`, `` (all Isolated) -/// collapsed into ONE scalar FSM over the atom stream. Precedence (= the Sequence order): digits → -/// CJK-range runs → the big-regex alts. Because Split-2 isolates CJK *before* the letter rule, the -/// letter run stops at CJK-range codepoints. Peeks bytes for the ASCII `[punct][A-Za-z]+` alt. -/// -/// Byte-exact vs the real composed Sequence (onig ×3, each Isolated) on 10 languages — see -/// `benches/deepseek.rs` (plus Hebrew/Arabic via `tk-encode`'s corpus test). The subtleties the single -/// pass replicates: (1) ws *followed by* a digit/CJK is its own Sequence piece → the whole run is one -/// token (`\s+(?!\S)`); (2) ZWJ/ZWNJ are `\p{Cf}`, not `\p{L}∪\p{M}`, so they end a letter run -/// (`ds_breaks`); (3) Split-2 isolates a maximal CJK-range run and Split-3 re-splits it into same-kind -/// sub-runs — the top-of-loop handler consumes that run as a CLOSED unit (`ds_is_cjk_at`), so CJK punct -/// (・) never steals a surrounding space nor merges with non-CJK punct; (4) chars matching no alt -/// (Control / NumericOther / ZWJ) group into ONE gap piece, and Other_Alphabetic symbols (`ALPHA_SYM`: -/// `\w` but categorically `\p{S}`) take the `[\p{P}\p{S}]` path, not the letter run. -#[must_use] -pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - // Leading-atom values as `const` → the `match` is a dense jump table (see `cl100k`). The Split - // precedence (digits → CJK → big-regex alts) is preserved because the atom partition is disjoint. - // `Mark` refined as an Other_Alphabetic symbol (Ⓘ …): coarse `LETTER_MARK`, but categorically `\p{S}` - // — excluded from `[\p{L}\p{M}]`, routed to the `[\p{P}\p{S}]+` run instead (see `punct`). - let end = text.len(); - // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior - // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) - let tags = &tags[..end]; - // maximal `[\p{L}\p{M}]+` run from `a`, stopping at CJK-range chars (Split-2 took those), ZWJ/ZWNJ - // (not `\p{L}∪\p{M}` — see `ds_breaks`), and Other_Alphabetic symbols (`ASM`, categorically `\p{S}`). - // BYTE-wise (`p += 1`, continuation bytes stay in-run, `ds_breaks` only fires at a lead) — the - // `char_len`-per-char form was ~2× slower (see `run_end`'s note). Hot inner loop of the latin path. - let letter_run = |a: usize| -> usize { - let mut p = a; - // ZWJ/ASM are now tags (no text peek); only the CJK-range exclusion still peeks text. - while p < end { - let t = tags[p]; - if t == CONT - || (in_mask(t, mask::LETTER_MARK) && t != ASM && t != ZWJ && !ds_is_cjk_at(text, p)) - { - p += 1; - } else { - break; - } - } - p - }; - // is `text[a]` the start of a deepseek letter/mark char (alt-2 run body / space-prefix target)? - let is_lm = |a: usize| { - a < end - && in_mask(tags[a], mask::LETTER_MARK) - && tags[a] != ASM - && tags[a] != ZWJ - && !ds_is_cjk_at(text, a) - }; - // Split-3 alt-3 tail `[\p{P}\p{S}]+[\r\n]*` from `sp0` (a leading space is already consumed); `sp0` - // if there is no punct/sym run there. STOPS at CJK-range chars — Split-1 isolated those, so a CJK - // punct (・) is never merged into a non-CJK punct run (`!・` → `!`, `・`, not `!・`). - let punct = |sp0: usize| -> usize { - let mut p = sp0; - while p < end - && (in_mask(tags[p], mask::PUNCT_SYM) || tags[p] == ASM) - && !ds_is_cjk_at(text, p) - { - p += char_len(text[p]); - } - if p > sp0 { - while p < end && tags[p] == NLN { - p += char_len(text[p]); - } - } - p - }; - // Split-3 alts d/e/f (whitespace). Unlike cl100k: a ws run FOLLOWED BY a digit/CJK is its own - // Sequence piece (Split-1/2 isolated the next match) → `\s+(?!\S)` takes the WHOLE run; only a - // following letter/punct (same Split-3 piece) leaves the last ws char for its ` ?`/`[^…]?` prefix. - let ws = |i: usize| -> usize { - let re = run_end(tags, i, end, mask::WS); - let next_isolated = re < end && (in_mask(tags[re], mask::NUMBER) || ds_is_cjk_at(text, re)); - if let Some(r) = text[i..re].iter().rposition(|&x| x == 0x0A || x == 0x0D) { - i + r + 1 - } else if re == end || next_isolated { - re // whole ws run is one token - } else { - let mut last = re - 1; - while last > i && text[last] & 0xC0 == 0x80 { - last -= 1; - } - if last > i { last } else { re } - } - }; - - let mut i = 0; - let mut w = 0usize; - while i < end { - let start = i; - let b = text[i]; - // Split-2 isolated a maximal CJK-range run; Split-3 re-splits it into same-kind sub-runs - // (letters `[\p{L}\p{M}]+` vs punct/sym `[\p{P}\p{S}]+`) — a CLOSED unit, handled before the atom - // arms so CJK punct (・) never leaks into alt-3 (stealing a space / merging with non-CJK punct). - if ds_is_cjk_at(text, i) { - let is_letter = in_mask(tags[i], mask::LETTER_MARK); - let mut p = i + 3; // CJK-range chars are all 3-byte (leads E3..E9) - while p < end - && ds_is_cjk_at(text, p) - && in_mask(tags[p], mask::LETTER_MARK) == is_letter - { - p += 3; - } - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: p as u32, - } - }; - w += 1; - i = p; - continue; - } - // Gap run: maximal Control / NumericOther / ZWJ — none matches a Split-3 alt, so the composed - // Split emits the whole run as ONE unmatched piece. Exception: if it's immediately followed by a - // letter run, the LAST gap char is that run's alt-2 `[^\r\n\p{L}\p{P}\p{S}]?` prefix (splits off). - if matches!(tags[i] & 0x0F, NMO | CTL) || tags[i] == ZWJ { - let (mut p, mut last) = (i, i); - while p < end && (matches!(tags[p] & 0x0F, NMO | CTL) || tags[p] == ZWJ) { - last = p; - p += char_len(text[p]); - } - if is_lm(p) { - if last > i { - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: last as u32, - } - }; // gap sans prefix char - w += 1; - } - let e = letter_run(p); - unsafe { - *out.get_unchecked_mut(w) = Span { - start: last as u32, - end: e as u32, - } - }; // prefix char + `[\p{L}\p{M}]+` - w += 1; - i = e; - } else { - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: p as u32, - } - }; // whole gap run is one piece - w += 1; - i = p; - } - continue; - } - match tags[i] & 0x0F { - // Split-1: `\p{N}{1,3}` - NW | NO => { - let (mut p, mut cnt) = (i, 0); - while p < end && cnt < 3 && in_mask(tags[p], mask::NUMBER) { - p += char_len(text[p]); - cnt += 1; - } - i = p; - } - // Split-3 alt-2 `[\p{L}\p{M}]+` (CJK letters + ZWJ/gap chars were consumed above the match). - // An Other_Alphabetic symbol (`ASM`, coarse `Mark`) is categorically `\p{S}` → the alt-3 run. - LET | MRK => { - i = if tags[i] == ASM { - punct(i) - } else { - letter_run(i) - } - } - // Space: alt-2 (space prefix + `[\p{L}\p{M}]+`) | alt-3 (` ` + `[\p{P}\p{S}]+`) | whitespace - SPC => { - let a = i + 1; // Space is ASCII (0x20) - i = if is_lm(a) { - letter_run(a) - } else if a < end && ds_is_cjk_at(text, a) { - ws(i) // next is a Split-1-isolated CJK char → the space is standalone whitespace - } else { - let p = punct(a); - if p > a { p } else { ws(i) } - }; - } - // WsOther: alt-2 (prefix + `[\p{L}\p{M}]+`) | whitespace (not `\p{P}∪\p{S}` → no alt-3) - WSO => { - let a = i + char_len(b); - i = if is_lm(a) { letter_run(a) } else { ws(i) }; - } - // Newline: whitespace - NLN => i = ws(i), - // Connector | Punct | Apostrophe | SymOther (∈ `\p{P}∪\p{S}`): alt-1 `[ascii_punct][A-Za-z]+` - // | alt-3 `[\p{P}\p{S}]+[\r\n]*` - CON | PUN | APO | SYM => { - i = if b.is_ascii_punctuation() && i + 1 < end && text[i + 1].is_ascii_alphabetic() - { - let mut p = i + 1; - while p < end && text[p].is_ascii_alphabetic() { - p += 1; - } - p - } else { - punct(i) // c ∈ PUNCT_SYM ⇒ > i - }; - } - // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. - _ => i += char_len(b), - } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; - } - w -} diff --git a/tokenizers/atomsplit/src/fsm/o200k.rs b/tokenizers/atomsplit/src/fsm/o200k.rs deleted file mode 100644 index 28e15640a..000000000 --- a/tokenizers/atomsplit/src/fsm/o200k.rs +++ /dev/null @@ -1,246 +0,0 @@ -use super::*; - -// ── o200k (GPT-4o): case-aware letter split ────────────────────────────────────────────────────── - -/// o200k class of a letter-run char (all chars in a run are real `[\p{L}\p{M}]`, never ALPHA_SYM/ZWJ): -/// `Atom::UpperLetter` → U (`\p{Lu}\p{Lt}`), `Atom::LowerLetter` → L (`\p{Ll}`), else C (caseless `\p{Lm}\p{Lo} -/// \p{M}`). The two alt char-classes are `[UC]` = "not L" (`!o_is_lower`) and `[LC]` = "not U". -#[inline] -fn o_is_upper(t: u8) -> bool { - t == Atom::UpperLetter as u8 // 0x10: coarse Letter (low nibble 0) + UPPER refine (\p{Lu}∪\p{Lt}) -} -#[inline] -fn o_is_lower(t: u8) -> bool { - t == Atom::LowerLetter as u8 // 0x20: coarse Letter + LOWER refine (\p{Ll}) -} - -/// One o200k letter sub-token from `p` within the run `[.., re)`: alt-1 `[UC]*[LC]+` (tried first) else -/// alt-2 `[UC]+[LC]*` (reached only for an all-U run). Greedy with Perl backtracking — `[UC]*` gives back -/// to the last C so `[LC]+` can take ≥1. Returns the sub-token end, always in `(p, re]`. BYTE-wise (`+=1`, -/// like `run_end`): continuation bytes are tag `Cont`(15) → neither U nor L → transparent to `[UC]`/`[LC]`. -#[inline(always)] -fn o200k_letter_match(tags: &[u8], p: usize, re: usize) -> usize { - // alt-1 `[UC]*`: greedy over "not L" (stops at the next lowercase *lead*), tracking the last C - // char-start so a no-L run needs no separate backtrack pass (`Cont`=15 is neither U nor L, so it's - // "C-like" — the `!= CONT` guard keeps `last_c` on a real char start). - let mut q = p; - let mut last_c = usize::MAX; - while q < re && !o_is_lower(tags[q]) { - if tags[q] != CONT && !o_is_upper(tags[q]) { - last_c = q; - } - q += 1; - } - if q < re { - // tags[q] is L → `[LC]+` from q: greedy over "not U" - let mut e = q; - while e < re && !o_is_upper(tags[e]) { - e += 1; - } - return e; - } - if last_c == usize::MAX { - return re; // no L and no C → all U → alt-2 `[UC]+[LC]*` (empty `[LC]*`) = the whole run - } - // no L: `[UC]*` gives back to the last C, which begins the `[LC]+` - let mut e = last_c; - while e < re && !o_is_upper(tags[e]) { - e += 1; - } - e -} - -/// Emit the o200k case-split of the letter run `[ls, re)` into `out[*w..]`: the first sub-token starts at -/// `pfx` (the optional `[^\r\n\p{L}\p{N}]?` prefix; `pfx == ls` when none), the last absorbs a trailing -/// contraction (`CONTRACTION` — off for tekken). Returns the new cursor (past the contraction). -/// `ls < re` (caller-guaranteed). -#[inline(always)] -fn emit_o200k_letters( - text: &[u8], - tags: &[u8], - pfx: usize, - ls: usize, - re: usize, - out: &mut [Span], - w: &mut usize, -) -> usize { - let (mut p, mut first, mut cursor) = (ls, true, re); - while p < re { - let e = o200k_letter_match(tags, p, re); - let start = if first { pfx } else { p }; - let tok_end = if CONTRACTION && e == re { - e + contraction(text, e) - } else { - e - }; - unsafe { - *out.get_unchecked_mut(*w) = Span { - start: start as u32, - end: tok_end as u32, - } - }; - *w += 1; - first = false; - cursor = tok_end; - p = e; - } - cursor -} - -/// o200k (GPT-4o) pretokenization. Same skeleton as cl100k, but the letter body is `[\p{L}\p{M}]+` -/// split into case sub-runs (`emit_o200k_letters`), the contraction is a *suffix* on letter tokens (not -/// a leading rule), and rule 4 is `[^\s\p{L}\p{N}]+[\r\n/]*`. The letter body excludes ALPHA_SYM symbols -/// and ZWJ/ZWNJ (coarse `Mark` but categorically `\p{S}`/`\p{Cf}`) — they take the prefix / rule-4 path. -/// Unlike deepseek there are no gaps: rule 4's `[^\s\p{L}\p{N}]+` is a catch-all. Scalar; ┌ OWNER: shared ┐ -#[must_use] -pub fn fsm_o200k(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - o200k::(text, tags, out) -} - -/// Mistral tekken ([`crate::regexes::TEKKEN`]) — the o200k FSM with the contraction suffix off and one -/// token per digit. Every other rule is shared, so both are the same code monomorphized twice. -#[must_use] -pub fn fsm_tekken(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - o200k::(text, tags, out) -} - -fn o200k( - text: &[u8], - tags: &[u8], - out: &mut [Span], -) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - let end = text.len(); - // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior - // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) - let tags = &tags[..end]; - - // Is tag `t` (at byte `p`) a real `[\p{L}\p{M}]` member? Coarse `Letter` (any case) is always in; - // coarse `Mark` is in only as a true `\p{M}` — ALPHA_SYM (`\p{S}`) and ZWJ/ZWNJ (`\p{Cf}`) are `\w` - // but not `[\p{L}\p{M}]`. The `ds_is_zwj` byte-peek is thus paid ONLY for Marks, never for letters. - let member = |t: u8, p: usize| -> bool { - let c = t & 0x0F; - let _ = p; - c == LET || (c == MRK && t != ASM && t != ZWJ) - }; - let is_lm = |a: usize| a < end && member(tags[a], a); - // maximal `[\p{L}\p{M}]+` run from `a` (byte-wise; continuation bytes ride along — see `run_end`). - let letter_end = |a: usize| -> usize { - let mut p = a; - // logos-style fast loop: 16 tags/chunk, one bounds check, unchecked reads. A plain `Letter` - // (low nibble 0, incl Han — o200k keeps all letters) or a `Cont` byte stays in-run with no - // `ds_is_zwj` peek; only a coarse `Mark` lane pays the peek. Byte-exact with the scalar scan. - // SAFETY: `p + 16 <= end <= tags.len()`/`text.len()` in the body. - while p + 16 <= end { - let mut brk = 16; - for k in 0..16 { - let t = unsafe { *tags.get_unchecked(p + k) }; - if t == CONT || t & 0x0F == LET { - continue; - } - if t & 0x0F == MRK && t != ASM && t != ZWJ { - continue; - } - brk = k; - break; - } - if brk < 16 { - return p + brk; - } - p += 16; - } - while p < end && (tags[p] == CONT || member(tags[p], p)) { - p += 1; - } - p - }; - // rule 4 `[^\s\p{L}\p{N}]+[\r\n/]*` from `sp0` (any leading space already consumed); `sp0` if none. - // `/` is in the `+` body too — the trailing class only matters after the `+` stops at a `\r\n`. - let other = |sp0: usize| -> usize { - let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); - if p > sp0 { - while p < end && (tags[p] == NLN || text[p] == b'/') { - p += char_len(text[p]); - } - } - p - }; - // rules 5-7 (`\s*[\r\n]+ | \s+(?!\S) | \s+`) → the shared `ws_tail` (identical to cl100k). - let ws = |i: usize| -> usize { ws_tail(text, tags, i, end) }; - // The letter rules: case-split the run starting at `ls`, first sub-token starting at the prefix `pfx`. - let letters = |pfx: usize, ls: usize, out: &mut [Span], w: &mut usize| -> usize { - emit_o200k_letters::(text, tags, pfx, ls, letter_end(ls), out, w) - }; - - let mut i = 0; - let mut w = 0usize; - while i < end { - let start = i; - let b = text[i]; - match tags[i] & 0x0F { - // rule 3: `\p{N}{1,DIGIT_CAP}`, no prefix (3 = o200k, 1 = tekken) - NW | NO => { - let (mut p, mut cnt) = (i, 0); - while p < end && cnt < DIGIT_CAP && in_mask(tags[p], mask::NUMBER) { - p += char_len(text[p]); - cnt += 1; - } - i = p; - } - // letters `[\p{L}\p{M}]+` (case-split) — but ALPHA_SYM/ZWJ (coarse Mark, not `[\p{L}\p{M}]`) - // take the `[^\r\n\p{L}\p{N}]?` prefix / rule-4 path instead. - LET | MRK => { - if tags[i] != ASM && tags[i] != ZWJ { - i = letters(i, i, out, &mut w); - continue; - } - let a = i + char_len(b); - if is_lm(a) { - i = letters(i, a, out, &mut w); - continue; - } - i = other(i); // ∈ NOT_WS_L_N ⇒ > i - } - // Space: ` ?` prefix + letters | ` ?` + rule-4 other | whitespace - SPC => { - let a = i + 1; // Space is ASCII (0x20) - if is_lm(a) { - i = letters(i, a, out, &mut w); - continue; - } - let p = other(a); - i = if p > a { p } else { ws(i) }; - } - // WsOther: prefix + letters | whitespace (∈ `\s` ⇒ never starts rule 4) - WSO => { - let a = i + char_len(b); - if is_lm(a) { - i = letters(i, a, out, &mut w); - continue; - } - i = ws(i); - } - NLN => i = ws(i), - // punct / sym / … (∈ `[^\r\n\p{L}\p{N}]` and `[^\s\p{L}\p{N}]`): prefix + letters | rule-4 other - CON | PUN | APO | SYM | NMO | CTL => { - let a = i + char_len(b); - if is_lm(a) { - i = letters(i, a, out, &mut w); - continue; - } - i = other(i); // ∈ NOT_WS_L_N ⇒ > i - } - // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. - _ => i += char_len(b), - } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; - } - w -} diff --git a/tokenizers/atomsplit/src/lib.rs b/tokenizers/atomsplit/src/lib.rs deleted file mode 100644 index 5f22f52d0..000000000 --- a/tokenizers/atomsplit/src/lib.rs +++ /dev/null @@ -1,33 +0,0 @@ -//! `atomsplit` — SIMD Unicode classification + finite-state pre-tokenization. -//! -//! One SIMD pass ([`classify`]) maps every codepoint to a tiny "atom" alphabet; a family of no-push -//! FSMs ([`fsm`]) turn that atom stream into token spans (byte ranges) — the pre-tokenizer stage that -//! runs before a BPE/WordPiece model. Pre-tokenizers implemented: `WhitespaceSplit`, `Punctuation`, -//! `Digits`, `Whitespace`, `Bert`, `Cl100k`, `DeepSeek`, `ByteLevel`, `CharDelimiterSplit` (o200k and -//! Mistral's tekken are exposed as the [`fsm::fsm_o200k`] / [`fsm::fsm_tekken`] functions rather than -//! recipe structs). -//! -//! For pre-tokenizers that split on single characters (such as the Metaspace `▁` delimiter), -//! we skip the atom classification pass entirely and search for the raw bytes instead. -//! See [`literal`]. -//! -//! Design: every fsm is *no-push* — it writes spans into a caller-preallocated `&mut [fsm::Span]` -//! (length ≥ `text.len()`) and returns the token count; there is no `Vec`/allocation on the hot path. -//! -//! # Preconditions -//! Inputs are `&[u8]` (not `&str`) for zero-copy, but **must be well-formed UTF-8**: a buffer that ends -//! mid-codepoint is a precondition violation and may panic. `tags`/`out` scratch buffers must be -//! `≥ text.len()` (asserted in [`classify`]; documented per-fsm). -mod atom_tables; -pub mod classify; -pub mod fsm; -pub mod literal; -pub mod regexes; -#[cfg(target_arch = "x86_64")] -mod simd_avx_classify; -#[cfg(target_arch = "aarch64")] -mod simd_classify; -mod simd_fsm; -#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] -mod simd_wasm_classify; -pub mod tables; diff --git a/tokenizers/atomsplit/src/literal.rs b/tokenizers/atomsplit/src/literal.rs deleted file mode 100644 index c826b2dda..000000000 --- a/tokenizers/atomsplit/src/literal.rs +++ /dev/null @@ -1,56 +0,0 @@ -//! Searching for a literal string, for pre-tokenizers that cut on one exact character. -//! -//! The rest of the crate works off atom tags: one SIMD pass gives every character a class, and the -//! FSMs cut where the class changes ([`crate::fsm`]). -//! -//! The atom classification is unnecessary for pre-tokenizers splitting on an exact character or literal string: -//! We can use a simpler byte search looking at 16 bytes at a time. - -use memchr::memmem; -use std::fmt; - -/// The pattern handed to [`Literal::new`] was empty. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct EmptyPattern; - -impl fmt::Display for EmptyPattern { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str("an empty pattern matches everywhere") - } -} - -impl std::error::Error for EmptyPattern {} - -/// A literal string to split on. -/// -/// The finder is boxed because it is large — a few hundred bytes of prefilter state on x86_64 — and -/// callers store it inside enums whose other variants are tiny. -#[derive(Debug, Clone)] -pub struct Literal { - finder: Box>, -} - -impl Literal { - /// # Errors - /// If `pattern` is empty, which would match everywhere. - pub fn new(pattern: &[u8]) -> Result { - if pattern.is_empty() { - return Err(EmptyPattern); - } - Ok(Self { - finder: Box::new(memmem::Finder::new(pattern).into_owned()), - }) - } - - /// The string being searched for. - #[must_use] - pub fn pattern(&self) -> &[u8] { - self.finder.needle() - } - - /// Byte offset of every match, left to right. Matches never overlap, so `"aa"` is found once in - /// `"aaa"` — the same matches a regex engine would report. - pub fn matches<'t>(&'t self, text: &'t [u8]) -> impl Iterator + 't { - self.finder.find_iter(text) - } -} diff --git a/tokenizers/atomsplit/tests/fsm.rs b/tokenizers/atomsplit/tests/fsm.rs deleted file mode 100644 index 1cad95e1b..000000000 --- a/tokenizers/atomsplit/tests/fsm.rs +++ /dev/null @@ -1,107 +0,0 @@ -//! Integration tests for the FSM pre-tokenizers. Kept out of `src/` so the core stays production-only. -use atomsplit::classify::{classify, mask}; -use atomsplit::fsm::{ - CharDelimiterSplit, Span, class_runs_into, emit_class_spans, fsm_byte_level, fsm_cl100k, - fsm_deepseek, fsm_o200k, fsm_tekken, -}; - -/// Run a no-push fsm into a fresh buffer and return the emitted spans. -fn spans(f: impl Fn(&[u8], &[u8], &mut [Span]) -> usize, s: &str) -> Vec { - let mut tags = vec![0u8; s.len()]; - classify(s.as_bytes(), &mut tags); - let mut out = vec![Span::default(); s.len() + 1]; - let k = f(s.as_bytes(), &tags, &mut out); - out.truncate(k); - out -} - -#[test] -fn cl100k_rules() { - // hand-verified against the tiktoken cl100k_base regex - let cl = |s| spans(fsm_cl100k, s); - assert_eq!(cl("Hello world"), vec![(0, 5), (5, 11)]); // "Hello" | " world" - assert_eq!(cl("don't"), vec![(0, 3), (3, 5)]); // "don" | "'t" (contraction) - assert_eq!(cl("a1234"), vec![(0, 1), (1, 4), (4, 5)]); // "a" | "123" | "4" ({1,3} cap) - assert_eq!(cl(" hi"), vec![(0, 1), (1, 4)]); // " " | " hi" - assert_eq!(cl("a, b"), vec![(0, 1), (1, 2), (2, 4)]); // "a" | "," | " b" -} - -/// Mistral's tekken split is o200k's, minus the contraction suffix, with one token per digit. -#[test] -fn tekken_rules() { - let tk = |s| spans(fsm_tekken, s); - assert_eq!(tk("don't"), vec![(0, 3), (3, 5)]); // "don" | "'t" — prefix+letters, not a contraction - assert_eq!(spans(fsm_o200k, "don't"), vec![(0, 5)]); // o200k glues the contraction on - assert_eq!(tk("a1234"), vec![(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)]); // `\p{N}`, one digit each - assert_eq!(tk("XMLHttpRequest"), vec![(0, 7), (7, 14)]); // case split: "XMLHttp" | "Request" - assert_eq!(tk("a/\r\nb"), vec![(0, 1), (1, 4), (4, 5)]); // "/" run + its `[\r\n/]*` tail - assert_eq!(tk("hi ok"), vec![(0, 2), (2, 4), (4, 7)]); // \s+(?!\S) leaves one space -} - -#[test] -fn deepseek_rules() { - let ds = |s| spans(fsm_deepseek, s); - assert_eq!(ds("abc中def"), vec![(0, 3), (3, 6), (6, 9)]); // letters | CJK | letters - assert_eq!(ds("abc123"), vec![(0, 3), (3, 6)]); // letters | digits {1,3} - assert_eq!(ds("_abc"), vec![(0, 4)]); // alt-1: ASCII punct + letters - assert_eq!(ds("hello world"), vec![(0, 5), (5, 11)]); // word | space+word - assert_eq!(ds("!!!"), vec![(0, 3)]); // \p{P}∪\p{S} run -} - -#[test] -fn byte_level_rules() { - let bl = |s| spans(fsm_byte_level, s); - // ` ?\p{L}+`, lowercase contraction, ` ?\p{N}+` UNBOUNDED - assert_eq!(bl("I'm 12345 ok"), vec![(0, 1), (1, 3), (3, 9), (9, 12)]); - assert_eq!(bl("IT'S"), vec![(0, 2), (2, 3), (3, 4)]); // 'S is not a contraction (case-sensitive) - assert_eq!(bl("hi ok"), vec![(0, 2), (2, 4), (4, 7)]); // \s+(?!\S) leaves one space -} - -#[test] -fn char_delimiter_split() { - let mut out = vec![Span::default(); 8]; - // split on '/', Removed → drop delimiters, drop the empty gap between "//" - let k = CharDelimiterSplit('/').pre_tokenize(b"a/bc//d", &mut [], &mut out); - assert_eq!(&out[..k], &[(0, 1), (2, 4), (6, 7)]); -} - -/// Byte-exactness gate for the class family: the NEON boundary extractor (`class_runs_into`) must equal -/// the scalar run-end core (`emit_class_spans`) for every recipe, at every char-aligned truncation length so -/// the < 16-byte NEON tail starts at every offset — including mid-char (chunk loop steps by 16). Corpus -/// mixes ASCII, 2/3-byte scripts, Devanagari letter+matra clusters, consecutive punct, tabs, astral. -#[test] -fn class_runs_into_matches() { - let corpus = "Hello, world!! 123 café × наука 中文。। नरेंद्र मोदी ने ½²¼ ①② 😀a b\t".repeat(30); - let full = corpus.as_bytes(); - let mut tags = vec![0u8; full.len()]; - let mut b1 = vec![Span::default(); full.len()]; - let mut b2 = vec![Span::default(); full.len()]; - - fn eq( - t: &[u8], - tg: &[u8], - x: &mut [Span], - y: &mut [Span], - name: &str, - ) { - let k1 = class_runs_into::(t, tg, x); - let k2 = emit_class_spans::(t, tg, y, 0, 0, 0, None); - assert_eq!(&x[..k1], &y[..k2], "{} @len {}", name, t.len()); - } - let mut sweep = |len: usize| { - let (t, tg) = (&full[..len], &mut tags[..len]); - classify(t, tg); - let (x, y) = (&mut b1[..len], &mut b2[..len]); - eq::<{ mask::WS }, 0, 0>(t, tg, x, y, "WhitespaceSplit"); - eq::<0, { mask::PUNCT }, 0>(t, tg, x, y, "Punctuation"); - eq::<0, 0, { mask::NUMERIC }>(t, tg, x, y, "Digits"); - eq::<{ mask::WS }, 0, { mask::WORD }>(t, tg, x, y, "Whitespace"); - eq::<{ mask::WS }, { mask::PUNCT }, 0>(t, tg, x, y, "Bert"); - }; - sweep(full.len()); - for c in full.len().saturating_sub(64)..full.len() { - if corpus.is_char_boundary(c) { - sweep(c); - } - } -} diff --git a/tokenizers/atomsplit/tests/literal.rs b/tokenizers/atomsplit/tests/literal.rs deleted file mode 100644 index 14ceed9a0..000000000 --- a/tokenizers/atomsplit/tests/literal.rs +++ /dev/null @@ -1,50 +0,0 @@ -//! Tests for the literal search. Kept out of `src/` so the core stays production-only. -use atomsplit::literal::{EmptyPattern, Literal}; - -#[test] -fn finds_every_match_and_nothing_else() { - let literal = Literal::new(b"-").unwrap(); - assert_eq!(literal.matches(b"a-b--c").collect::>(), [1, 3, 4]); - assert_eq!(literal.matches(b"none here").count(), 0); - assert_eq!(literal.matches(b"").count(), 0); - assert_eq!(literal.pattern(), b"-"); -} - -#[test] -fn a_multi_byte_pattern_only_matches_whole() { - // `▁` is U+2581 = E2 96 81. The other characters here share its first byte and nothing else, so a - // search for that byte alone would report all of them. - let literal = Literal::new("▁".as_bytes()).unwrap(); - let text = "a—b“c…d▁e"; - assert_eq!( - literal.matches(text.as_bytes()).collect::>(), - [text.find('▁').unwrap()] - ); -} - -#[test] -fn matches_do_not_overlap() { - let literal = Literal::new(b"aa").unwrap(); - assert_eq!(literal.matches(b"aaa").collect::>(), [0]); - assert_eq!(literal.matches(b"aaaa").collect::>(), [0, 2]); -} - -/// A `Literal` is stored inline in the normalizer and decoder enums, where every other variant is a -/// handful of bytes. `memmem::Finder` itself is a few hundred, so it has to stay behind a pointer. -#[test] -fn a_literal_is_pointer_sized() { - assert_eq!( - size_of::(), - size_of::<*const u8>(), - "a Literal must not carry its finder inline" - ); -} - -#[test] -fn an_empty_pattern_is_rejected() { - assert_eq!(Literal::new(b"").unwrap_err(), EmptyPattern); - assert_eq!( - EmptyPattern.to_string(), - "an empty pattern matches everywhere" - ); -} diff --git a/tokenizers/atomsplit/tests/parity.rs b/tokenizers/atomsplit/tests/parity.rs deleted file mode 100644 index 774e7c002..000000000 --- a/tokenizers/atomsplit/tests/parity.rs +++ /dev/null @@ -1,121 +0,0 @@ -//! Reference-parity gates for the regex-shaped pre-tokenizers. Each of `fsm_cl100k` / `fsm_o200k` / -//! `fsm_tekken` / `fsm_byte_level` / `fsm_deepseek` must be BYTE-EXACT with the real reference — the -//! oniguruma regex, composed exactly as HF applies it (deepseek is a `Sequence` of three Isolated -//! splits) — on two corpora: a multilingual one (ASCII, contractions, digits, punctuation runs, -//! whitespace variants, Latin accents/marks, Cyrillic/Greek/Arabic/Devanagari, Han/Kana/Hangul, ZWJ, -//! astral emoji) and an edge-case one (see [`EDGE`]). -//! -//! This is the byte-exactness gate the hand cases in `fsm.rs` don't provide; the same corpora + gate -//! should run under x86 (Intel SDE) in CI to validate the SIMD paths. -//! -//! Gated off wasm32: the oniguruma reference is a C library that has no wasi libc to build against. -#![cfg(not(target_arch = "wasm32"))] -use atomsplit::classify::classify; -use atomsplit::fsm::{Span, fsm_byte_level, fsm_cl100k, fsm_deepseek, fsm_o200k, fsm_tekken}; -use onig::Regex; -// The oracle regexes are the canonical specs the FSMs implement — single source of truth in atomsplit. -use atomsplit::regexes::{ - CL100K, DEEPSEEK_BIG as DS_BIG, DEEPSEEK_CJK as DS_CJK, DEEPSEEK_NUM as DS_NUM, GPT2, O200K, - TEKKEN, -}; - -const CORPUS: &str = "The quick brown fox. Don't 12345 numbers, \u{00BD}\u{00B2}\u{00BC} \u{2168}! \ - café × naïve — Привет, наука! Ελλάδα 中文分词。ひらがな カタカナ 한글 مرحبا العربية \ - नरेंद्र मोदी x_y a1b2c3 e-mail@host.com 😀👍 hello world\ttabs\nnewlines end "; - -/// Second corpus, aimed at the axes where the o200k-shaped FSMs differ from each other: apostrophes -/// (contraction suffix vs plain prefix+letters), digit-run length (`{1,3}` vs one-per-token), and the -/// `[\r\n/]*` tail after a symbol run. -const EDGE: &str = "IT'S O'Brien can't 'quoted' l'été rock'n'roll\r\n\ - 0 42 999 1000 1234567 v1.2.3 3.14159 1,000,000 \ - https://host/a/b?c=1&d=2 path/to//file /\r\n/ ///x \ - CamelCase XMLHttpRequest IJSSELMEER DžAMBO ŀl a\u{0301}b \ - 日本語1234テスト ½3¼ \u{2168}42\u{2169} #tag @user $9.99 100% \ - end\n\n\nlines\r\n\r\n \t trailing "; - -fn spans(f: impl Fn(&[u8], &[u8], &mut [Span]) -> usize, s: &str) -> Vec { - let mut tags = vec![0u8; s.len()]; - classify(s.as_bytes(), &mut tags); - let mut out = vec![Span::default(); s.len() + 1]; - let k = f(s.as_bytes(), &tags, &mut out); - out.truncate(k); - out -} - -fn onig_spans(re: &Regex, s: &str) -> Vec { - re.find_iter(s) - .map(|(a, b)| Span::new(a as u32, b as u32)) - .collect() -} - -// One Isolated split of text[s..e] by `re`: emit gaps + matches (all pieces), absolute offsets. -fn split_iso(text: &str, s: usize, e: usize, re: &Regex, out: &mut Vec<(usize, usize)>) { - let sub = &text[s..e]; - let mut prev = 0usize; - for (ms, me) in re.find_iter(sub) { - if ms > prev { - out.push((s + prev, s + ms)); - } - out.push((s + ms, s + me)); - prev = me; - } - if prev < sub.len() { - out.push((s + prev, e)); - } -} - -fn deepseek_ref(text: &str) -> Vec { - let (rn, rc, rb) = ( - Regex::new(DS_NUM).unwrap(), - Regex::new(DS_CJK).unwrap(), - Regex::new(DS_BIG).unwrap(), - ); - let mut p1 = Vec::new(); - split_iso(text, 0, text.len(), &rn, &mut p1); - let mut p2 = Vec::new(); - for (s, e) in p1 { - split_iso(text, s, e, &rc, &mut p2); - } - let mut p3 = Vec::new(); - for (s, e) in p2 { - split_iso(text, s, e, &rb, &mut p3); - } - p3.into_iter() - .map(|(s, e)| Span::new(s as u32, e as u32)) - .collect() -} - -/// Both corpora, one regex: the FSM must reproduce the oracle span-for-span. -fn check(fsm: impl Fn(&[u8], &[u8], &mut [Span]) -> usize + Copy, pattern: &str) { - let re = Regex::new(pattern).unwrap(); - for text in [CORPUS, EDGE] { - assert_eq!(spans(fsm, text), onig_spans(&re, text), "{text:?}"); - } -} - -#[test] -fn cl100k_parity() { - check(fsm_cl100k, CL100K); -} - -#[test] -fn o200k_parity() { - check(fsm_o200k, O200K); -} - -#[test] -fn tekken_parity() { - check(fsm_tekken, TEKKEN); -} - -#[test] -fn byte_level_parity() { - check(fsm_byte_level, GPT2); -} - -#[test] -fn deepseek_parity() { - for text in [CORPUS, EDGE] { - assert_eq!(spans(fsm_deepseek, text), deepseek_ref(text), "{text:?}"); - } -} diff --git a/tokenizers/bitmap_gen/Cargo.toml b/tokenizers/bitmap_gen/Cargo.toml index 3500b399a..800c26cc1 100644 --- a/tokenizers/bitmap_gen/Cargo.toml +++ b/tokenizers/bitmap_gen/Cargo.toml @@ -2,7 +2,7 @@ name = "bitmap_gen" version = "0.1.0" edition = "2024" -description = "Dev tool: regenerates atomsplit's committed classify tables (uses unicode-properties). Not for publication." +description = "Dev tool: regenerates bitsplit's committed classify tables (uses unicode-properties). Not for publication." license = "Apache-2.0" repository = "https://github.com/huggingface/tokenizers" authors = ["Arthur Zucker ", "Luc Georges "] diff --git a/tokenizers/bitmap_gen/src/lib.rs b/tokenizers/bitmap_gen/src/lib.rs index 302ff1bb9..224ce82f2 100644 --- a/tokenizers/bitmap_gen/src/lib.rs +++ b/tokenizers/bitmap_gen/src/lib.rs @@ -213,7 +213,7 @@ fn generate_tables(struct_name: &str, kind: &str, classify: &dyn Fn(u32) -> u8) ) .unwrap(); o.push_str("//! `Tables::classify_char`. `bitmap_gen` self-validates all 1.1M codepoints. Spec §1/§7.\n"); - o.push_str("use crate::tables::Tables;\n\n"); + o.push_str("use super::tables::Tables;\n\n"); let emit_u8 = |o: &mut String, name: &str, t: &[u8]| { write!(o, "#[rustfmt::skip]\nstatic {name}: [u8; {}] = [", t.len()).unwrap(); diff --git a/tokenizers/bitmap_gen/src/main.rs b/tokenizers/bitmap_gen/src/main.rs index 0b2990184..e76c05423 100644 --- a/tokenizers/bitmap_gen/src/main.rs +++ b/tokenizers/bitmap_gen/src/main.rs @@ -1,12 +1,9 @@ -//! Regenerate atomsplit's committed classify tables: +//! Regenerate bitsplit's committed classify tables: //! cargo run -p bitmap_gen [-- ] -//! Default `out_path` = ../atomsplit/src/atom_tables.rs. `generate_atom_tables` self-validates every +//! Default `out_path` = ../bitsplit/src/classify/atom_tables.rs. `generate_atom_tables` self-validates every //! codepoint against the reference `atom()`, so an inconsistent scheme change fails HERE, not at ship. fn main() { - let default = concat!( - env!("CARGO_MANIFEST_DIR"), - "/../atomsplit/src/atom_tables.rs" - ); + let default = concat!(env!("CARGO_MANIFEST_DIR"), "/../bitsplit/src/classify/atom_tables.rs"); let out = std::env::args() .nth(1) .unwrap_or_else(|| default.to_string()); diff --git a/tokenizers/bitsplit/Cargo.toml b/tokenizers/bitsplit/Cargo.toml index ebc3d609e..e7dc67097 100644 --- a/tokenizers/bitsplit/Cargo.toml +++ b/tokenizers/bitsplit/Cargo.toml @@ -2,12 +2,33 @@ name = "bitsplit" version = "0.1.0" edition = "2024" +rust-version = "1.89" # AVX-512 VBMI intrinsics (simd_avx_classify) are stable only since 1.89 +authors = [ + "Arthur Zucker ", + "Luc Georges ", + "Simon Brandeis ", +] +homepage = "https://github.com/huggingface/tokenizers" +repository = "https://github.com/huggingface/tokenizers" +documentation = "https://docs.rs/bitsplit/" +license = "Apache-2.0" +keywords = ["tokenizer", "nlp", "pretokenizer", "simd", "unicode"] +categories = ["text-processing", "parsing"] +description = "SIMD Unicode atom classification + bitstream pre-tokenization: one classify pass into byte-exact token spans (cl100k, o200k, GPT-2/ByteLevel, DeepSeek, Whitespace, Bert, ...)." +exclude = ["benches/data/"] +[package.metadata.docs.rs] +all-features = true + +[lib] +name = "bitsplit" +path = "src/lib.rs" [dependencies] ahash = "0.8" -atomsplit = { path = "../atomsplit" } - - - +# NOTE: classify tables live in src/classify/atom_tables.rs (committed, generated). Regenerate after any atom +# scheme change with `cargo run -p bitmap_gen`. No build script / build-dep — bitsplit builds clean. +# onig is C (Oniguruma) and has no wasi libc to build against, so the parity oracle is gated off wasm32. +[target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies] +onig = "6.5.1" diff --git a/tokenizers/bitsplit/src/classes.rs b/tokenizers/bitsplit/src/classes.rs new file mode 100644 index 000000000..61e954fe9 --- /dev/null +++ b/tokenizers/bitsplit/src/classes.rs @@ -0,0 +1,210 @@ +//! The class-run family: WhitespaceSplit / Punctuation / Digits / Whitespace / Bert, plus +//! CharDelimiterSplit. Not regex-shaped -- these cut where the atom class changes, so one shape +//! (`class_runs_into`) covers all of them. +//! +//! ponytail: still the scalar/NEON run extractor moved over from atomsplit, not a bitstream +//! program. A class-run boundary is `c & !(c << 1)`, so porting it onto the `Blk` streams would +//! delete `simd_classes.rs` outright -- worth doing, but it is not what makes these correct today. + +use crate::Span; +use crate::classify::{Atom, char_len, classify, in_mask, mask}; +use crate::run_end; + +/// No-`push` class-family pre-tokenizer core: writes spans into the preallocated `out` slice and returns +/// the count. ONE shape covers the whole class family via ``: +/// WhitespaceSplit `<{WS},0,0>` · Punctuation `<0,{PUNCT},0>` · Digits `<0,0,{NUMERIC}>` · +/// Whitespace `<{WS},0,{WORD}>` · Bert `<{WS},{PUNCT},0>`. +/// Class of a char: `DROP`→dropped, `ISOLATE`→own token, `KEEP_A`→run "A", else→run "B" (A/B cut apart). +/// TODO: find a better explanation +#[inline] +#[must_use] +pub fn class_runs_into( + text: &[u8], + tags: &[u8], + out: &mut [Span], +) -> usize { + debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + #[cfg(target_arch = "aarch64")] + { + crate::simd::classes::class_runs_neon::(text, tags, out) + } + #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] + { + crate::simd::classes::class_runs_wasm::(text, tags, out) + } + #[cfg(not(any( + target_arch = "aarch64", + all(target_arch = "wasm32", target_feature = "simd128") + )))] + { + emit_class_spans::(text, tags, out, 0, 0, 0, None) + } +} + +/// This is the most important function as it's the core of the scalar finite state machine. +/// It allows to emit class spans with different behaviours for tags we want to drop, tags we want +/// to isolate and tags we want to keep. Any other tags are assumed to be keept. +/// +/// This function is used as a fallback to the SIMD fast fsm. It is used for most pre tokenizers +/// but the unrolled regex, which have more complex variations that cannot be expressed with drop, +/// isolate, keep. These 3 generic parameters are u16 bitmap masks over the 16 classes we have and +/// define the behaviour. They are usally one of the [`crate::classify::mask`]. They allow dropping +/// words, isolating whitespace and keeping new line for example. +#[must_use] +#[inline] +pub fn emit_class_spans( + text: &[u8], + tags: &[u8], + out: &mut [Span], + mut write_index: usize, // in the out slice + mut text_pointer: usize, // in the text slice + segment_start: usize, // previous segment_start + segment_class: Option, // previous segment's class +) -> usize { + debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + let n = text.len(); + // Tie `tags.len() == text.len() == n` so the optimizer drops the interior `tags[i]` / `text[i]` + // bounds checks (callers guarantee len ≥ n; same trick as `cl100k`). Per-byte scanning already + // avoids checks via `run_end`'s unrolled `get_unchecked`; this covers the per-token accesses. + let tags = &tags[..n]; + let text = &text[..n]; + let other = !(DROP | ISOLATE | KEEP_A); // None of the above correspond to a continuation + if let Some(segment_class) = segment_class { + // this will usually be at the tail of a SIMD call. + text_pointer = run_end(tags, text_pointer, n, segment_class); // skip the whole drop run at once + if segment_class != DROP { + out[write_index] = Span { + start: segment_start as u32, + end: text_pointer as u32, + }; + if text_pointer == n { + return write_index + 1; + } + write_index += 1; + } + } + while text_pointer < n { + let t = tags[text_pointer]; + if t == Atom::Cont as u8 { + text_pointer += 1; + continue; + } + // classify the first char. + if in_mask(t, DROP) { + text_pointer = run_end(tags, text_pointer, n, DROP); // skip the whole drop run at once + } else if in_mask(t, ISOLATE) { + let s = text_pointer; + text_pointer += char_len(text[text_pointer]); + out[write_index] = Span { + start: s as u32, + end: text_pointer as u32, + }; // isolate: one char = one token + write_index += 1; + } else { + let s = text_pointer; + text_pointer = if in_mask(t, KEEP_A) { + run_end(tags, text_pointer, n, KEEP_A) + } else { + run_end(tags, text_pointer, n, other) + }; + out[write_index] = Span { + start: s as u32, + end: text_pointer as u32, + }; + write_index += 1; + } + } + write_index +} + +// ── per-tokenizer unrolled FSMs (one file each; shared helpers above via `use super::*`) ── +pub struct WhitespaceSplit; +impl WhitespaceSplit { + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { + classify(text, tags); + class_runs_into::<{ mask::WS }, 0, 0>(text, tags, out) + } +} + +/// `Punctuation` — isolate each punctuation char as its own token; non-punct grouped into runs. +pub struct Punctuation; +impl Punctuation { + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { + classify(text, tags); + class_runs_into::<0, { mask::PUNCT }, 0>(text, tags, out) + } +} + +/// `Digits` — cut numeric runs apart from non-numeric runs (contiguous), keeping both. +pub struct Digits; +impl Digits { + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { + classify(text, tags); + class_runs_into::<0, 0, { mask::NUMERIC }>(text, tags, out) + } +} + +/// `Whitespace` — the `\w+|[^\w\s]+` pre-tokenizer: drop whitespace, cut word runs from symbol runs. +pub struct Whitespace; +impl Whitespace { + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { + classify(text, tags); + class_runs_into::<{ mask::WS }, 0, { mask::WORD }>(text, tags, out) + } +} + +/// `Bert` — the BERT basic pre-tokenizer: drop whitespace, isolate punctuation, keep the rest as runs. +pub struct Bert; +impl Bert { + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { + classify(text, tags); + class_runs_into::<{ mask::WS }, { mask::PUNCT }, 0>(text, tags, out) + } +} + + + + +/// `Split(char, Removed)` — the only pre-tokenizer that keys on a *literal char* rather than an atom +/// class, so it scans bytes directly (no classify pass). UTF-8 is self-synchronizing, so the +/// delimiter's byte pattern only matches on char boundaries. +pub struct CharDelimiterSplit(pub char); +impl CharDelimiterSplit { + /// Split on the literal char (`Removed`): writes the gaps between delimiters into `out` + /// (len >= `text.len()`) and returns the count. Straight off [`crate::literal`] — no atom + /// classification, UTF-8 is self-synchronising so the char's bytes only match on boundaries. + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], _tags: &mut [u8], out: &mut [Span]) -> usize { + debug_assert!(out.len() >= text.len()); + let mut buf = [0u8; 4]; + let delim = self.0.encode_utf8(&mut buf).as_bytes(); + let Ok(lit) = crate::literal::Literal::new(delim) else { + return 0; + }; + let (n, dl) = (text.len(), delim.len()); + let (mut start, mut w) = (0usize, 0usize); + for m in lit.matches(text) { + if m > start { + out[w] = Span::new(start as u32, m as u32); + w += 1; + } + start = m + dl; + } + if start < n { + out[w] = Span::new(start as u32, n as u32); + w += 1; + } + w + } +} diff --git a/tokenizers/atomsplit/src/atom_tables.rs b/tokenizers/bitsplit/src/classify/atom_tables.rs similarity index 99% rename from tokenizers/atomsplit/src/atom_tables.rs rename to tokenizers/bitsplit/src/classify/atom_tables.rs index 936f2f2dd..fcdba5840 100644 --- a/tokenizers/atomsplit/src/atom_tables.rs +++ b/tokenizers/bitsplit/src/classify/atom_tables.rs @@ -1,7 +1,7 @@ //! GENERATED by `bitmap_gen` — do NOT edit. Regenerate with `cargo run -p bitmap_gen`. //! Dense atom classify tables shared by the SIMD kernel (`vqtbl`) and the scalar reader //! `Tables::classify_char`. `bitmap_gen` self-validates all 1.1M codepoints. Spec §1/§7. -use crate::tables::Tables; +use super::tables::Tables; #[rustfmt::skip] static GROUP: [[[u8; 64]; 4]; 8] = [[[12,12,12,12,12,12,12,12,12,5,3,5,5,3,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,4,8,8,8,8,8,8,9,8,8,8,8,8,8,8,8,1,1,1,1,1,1,1,1,1,1,8,8,8,8,8,8],[8,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,8,8,8,8,7,8,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,8,8,8,8,12],[12,12,12,12,12,5,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,5,8,10,10,10,10,10,8,10,10,0,8,10,12,10,10,10,10,2,2,10,32,8,8,10,2,0,8,2,2,2,8],[16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,10,16,16,16,16,16,16,16,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,10,32,32,32,32,32,32,32,32]],[[16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,16,32,16,32,16,32,16],[32,16,32,16,32,16,32,16,32,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,16,32,16,32,16,32,32],[32,16,16,32,16,32,16,16,32,16,16,16,32,32,16,16,16,16,32,16,16,32,16,16,16,32,32,32,16,16,32,16,16,32,16,32,16,32,16,16,32,16,32,32,16,32,16,16,32,16,16,16,32,16,32,16,16,32,32,0,16,32,32,32],[0,0,0,0,16,16,32,16,16,32,16,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,16,16,32,16,32,16,16,16,32,16,32,16,32,16,32]],[[16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,32,32,32,32,32,16,16,32,16,16,32],[32,16,32,16,16,16,16,32,16,32,16,32,16,32,16,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32],[32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,0,0,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,10,10,10,10,0,0,0,0,0,0,0,0,0,0,0,0,10,10,10,10,10,10,10,10,10,10,10,10,10,10,0,0,0,0,0,10,10,10,10,10,10,10,0,10,0,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10]],[[6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6],[6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,16,32,16,32,0,10,16,32,12,12,0,32,32,32,8,16],[12,12,12,12,10,10,16,8,16,16,16,12,16,12,16,16,32,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,12,16,16,16,16,16,16,16,16,16,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32],[32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,16,32,32,16,16,16,32,32,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,32,32,32,16,32,10,16,32,16,16,32,32,16,16,16]],[[16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32],[32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32],[16,32,10,6,6,6,6,6,6,6,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32],[16,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32]],[[16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,12,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16],[16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,12,12,0,8,8,8,8,8,8,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32],[32,32,32,32,32,32,32,32,32,8,8,12,12,10,10,10,12,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,8,6],[8,6,6,8,6,6,8,6,12,12,12,12,12,12,12,12,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,12,12,12,12,0,0,0,0,8,8,12,12,12,12,12,12,12,12,12,12,12]],[[12,12,12,12,12,12,10,10,10,8,8,10,8,8,10,10,6,6,6,6,6,6,6,6,6,6,6,8,12,8,8,8,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,1,1,1,1,1,1,1,1,1,1,8,8,8,8,0,0,6,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,8,0,6,6,6,6,6,6,6,12,10,6,6,6,6,6,6,0,0,6,6,10,6,6,6,6,0,0,1,1,1,1,1,1,1,1,1,1,0,0,0,10,10,0]],[[8,8,8,8,8,8,8,8,8,8,8,8,8,8,12,12,0,6,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6],[6,6,6,6,6,6,6,6,6,6,6,12,12,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,6,6,6,6,6,6,6,6,6,6,6,0,12,12,12,12,12,12,12,12,12,12,12,12,12,12],[1,1,1,1,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,6,6,6,6,6,6,6,6,6,0,0,10,8,8,8,0,12,12,6,10,10]]]; diff --git a/tokenizers/atomsplit/src/simd_avx_classify.rs b/tokenizers/bitsplit/src/classify/avx.rs similarity index 99% rename from tokenizers/atomsplit/src/simd_avx_classify.rs rename to tokenizers/bitsplit/src/classify/avx.rs index 68cfef083..c3f595de8 100644 --- a/tokenizers/atomsplit/src/simd_avx_classify.rs +++ b/tokenizers/bitsplit/src/classify/avx.rs @@ -14,8 +14,8 @@ //! x86_64 (SSE4.1 and, if available, AVX-512 VBMI) hardware before trusting it. #![allow(unsafe_op_in_unsafe_fn)] -use crate::atom_tables::ATOM_TABLES; -use crate::classify::{Atom, CONT, MB, char_len, classify_scalar}; +use super::atom_tables::ATOM_TABLES; +use super::{Atom, CONT, MB, char_len, classify_scalar}; use core::arch::x86_64::*; const CJK_TAG: u8 = Atom::Letter as u8; diff --git a/tokenizers/atomsplit/src/classify.rs b/tokenizers/bitsplit/src/classify/mod.rs similarity index 92% rename from tokenizers/atomsplit/src/classify.rs rename to tokenizers/bitsplit/src/classify/mod.rs index 2e4e1f8dc..38fe1d1ed 100644 --- a/tokenizers/atomsplit/src/classify.rs +++ b/tokenizers/bitsplit/src/classify/mod.rs @@ -1,4 +1,16 @@ -use crate::atom_tables::ATOM_TABLES; +//! SIMD Unicode classification: one pass mapping every codepoint to a tiny "atom" alphabet, which +//! is what every grammar in `models/` consumes. The per-arch kernels are siblings; `atom_tables.rs` +//! is generated (`cargo run -p bitmap_gen`) and `tables.rs` is the layout it bakes into. +mod atom_tables; +#[cfg(target_arch = "x86_64")] +mod avx; +#[cfg(target_arch = "aarch64")] +mod neon; +mod tables; +#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] +mod wasm; + +use atom_tables::ATOM_TABLES; /// The per-codepoint "atom" categories or "tags" that are used by the finite state machine to emit /// spit boundaries. @@ -125,14 +137,14 @@ pub fn classify(text: &[u8], tags: &mut [u8]) { #[cfg(target_arch = "aarch64")] // SAFETY: `tags.len() >= text.len()` (asserted above); NEON vld1q/vst1q are alignment-free. unsafe { - crate::simd_classify::classify_neon(text, tags) + neon::classify_neon(text, tags) } #[cfg(target_arch = "x86_64")] - crate::simd_avx_classify::dispatch(text, tags); + avx::dispatch(text, tags); #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] // SAFETY: `tags.len() >= text.len()` (asserted above); wasm v128 load/store are alignment-free. unsafe { - crate::simd_wasm_classify::classify_wasm(text, tags) + wasm::classify_wasm(text, tags) } #[cfg(not(any( target_arch = "aarch64", diff --git a/tokenizers/atomsplit/src/simd_classify.rs b/tokenizers/bitsplit/src/classify/neon.rs similarity index 99% rename from tokenizers/atomsplit/src/simd_classify.rs rename to tokenizers/bitsplit/src/classify/neon.rs index ae9ee3939..2b008bf7d 100644 --- a/tokenizers/atomsplit/src/simd_classify.rs +++ b/tokenizers/bitsplit/src/classify/neon.rs @@ -1,5 +1,5 @@ -use crate::atom_tables::ATOM_TABLES; -use crate::classify::{Atom, CONT, MB}; +use super::atom_tables::ATOM_TABLES; +use super::{Atom, CONT, MB}; const CJK_TAG: u8 = Atom::Letter as u8; // ================================================================================================ @@ -98,7 +98,7 @@ unsafe fn any(mask: core::arch::aarch64::uint8x16_t) -> bool { /// TLDR removing the utf8 headers to get the unicode. fn decode(t: &[u8], i: usize) -> u32 { let b = t[i] as u32; - match super::classify::char_len(t[i]) { + match super::char_len(t[i]) { 1 => b, 2 => ((b & 0x1F) << 6) | (t[i + 1] as u32 & 0x3F), 3 => ((b & 0x0F) << 12) | ((t[i + 1] as u32 & 0x3F) << 6) | (t[i + 2] as u32 & 0x3F), @@ -156,7 +156,7 @@ fn decode(t: &[u8], i: usize) -> u32 { #[cfg(target_arch = "aarch64")] #[allow(unsafe_op_in_unsafe_fn, non_snake_case)] pub unsafe fn classify_neon(text: &[u8], tags: &mut [u8]) { - use super::classify::char_len; + use super::char_len; use core::arch::aarch64::*; let n = text.len(); let mut i = 0; diff --git a/tokenizers/atomsplit/src/tables.rs b/tokenizers/bitsplit/src/classify/tables.rs similarity index 100% rename from tokenizers/atomsplit/src/tables.rs rename to tokenizers/bitsplit/src/classify/tables.rs diff --git a/tokenizers/atomsplit/src/simd_wasm_classify.rs b/tokenizers/bitsplit/src/classify/wasm.rs similarity index 99% rename from tokenizers/atomsplit/src/simd_wasm_classify.rs rename to tokenizers/bitsplit/src/classify/wasm.rs index 12b7a6dc5..235175cee 100644 --- a/tokenizers/atomsplit/src/simd_wasm_classify.rs +++ b/tokenizers/bitsplit/src/classify/wasm.rs @@ -19,8 +19,8 @@ //! a SIMD128 wasm engine before trusting it. #![allow(unsafe_op_in_unsafe_fn)] -use crate::atom_tables::ATOM_TABLES; -use crate::classify::{Atom, CONT, MB, char_len}; +use super::atom_tables::ATOM_TABLES; +use super::{Atom, CONT, MB, char_len}; use core::arch::wasm32::*; const CJK_TAG: u8 = Atom::Letter as u8; diff --git a/tokenizers/bitsplit/src/han.rs b/tokenizers/bitsplit/src/han.rs new file mode 100644 index 000000000..f24e76604 --- /dev/null +++ b/tokenizers/bitsplit/src/han.rs @@ -0,0 +1,72 @@ +//! `\p{Han}` (Script=Han) for kimi-k2's leading `[\p{Han}]+` arm. +//! +//! - Han is NOT deepseek's `is_cjk_at` range: that one is Han U+4E00..9FA5 plus kana, this one is +//! the whole script and excludes kana. They are different predicates on purpose. +//! - Every range here is ≥ U+2E80, so a lead byte below 0xE2 exits before any decoding. + +/// Script=Han, sorted and disjoint. Extensions past Ext F are included because onig's `\p{Han}` +/// has them; the parity gate is what actually pins this table to the oracle's Unicode version. +const HAN: &[(u32, u32)] = &[ + (0x2E80, 0x2E99), + (0x2E9B, 0x2EF3), + (0x2F00, 0x2FD5), + (0x3005, 0x3005), + (0x3007, 0x3007), + (0x3021, 0x3029), + (0x3038, 0x303B), + (0x3400, 0x4DBF), + (0x4E00, 0x9FFF), + (0xF900, 0xFA6D), + (0xFA70, 0xFAD9), + (0x20000, 0x2A6DF), + (0x2A700, 0x2B739), + (0x2B740, 0x2B81D), + (0x2B820, 0x2CEA1), + (0x2CEB0, 0x2EBE0), + (0x2EBF0, 0x2EE5D), + (0x2F800, 0x2FA1D), + (0x30000, 0x3134A), + (0x31350, 0x323AF), +]; + +#[inline] +#[must_use] +pub(crate) fn is_han(cp: u32) -> bool { + HAN.binary_search_by(|&(lo, hi)| { + if cp < lo { + core::cmp::Ordering::Greater + } else if cp > hi { + core::cmp::Ordering::Less + } else { + core::cmp::Ordering::Equal + } + }) + .is_ok() +} + +/// Is the char whose lead byte is at `p` in Script=Han? `false` on a truncated tail. +#[inline] +#[must_use] +pub(crate) fn is_han_at(text: &[u8], p: usize) -> bool { + let b = text[p]; + let n = text.len(); + let cp = if (0xE2..=0xEF).contains(&b) { + if p + 2 >= n { + return false; + } + ((b as u32 & 0x0F) << 12) + | ((text[p + 1] as u32 & 0x3F) << 6) + | (text[p + 2] as u32 & 0x3F) + } else if (0xF0..=0xF4).contains(&b) { + if p + 3 >= n { + return false; + } + ((b as u32 & 0x07) << 18) + | ((text[p + 1] as u32 & 0x3F) << 12) + | ((text[p + 2] as u32 & 0x3F) << 6) + | (text[p + 3] as u32 & 0x3F) + } else { + return false; + }; + is_han(cp) +} diff --git a/tokenizers/bitsplit/src/lib.rs b/tokenizers/bitsplit/src/lib.rs index 79d619b1e..8e229f2cb 100644 --- a/tokenizers/bitsplit/src/lib.rs +++ b/tokenizers/bitsplit/src/lib.rs @@ -17,24 +17,58 @@ //! visited in order, so those dependencies are carried in a handful of scalar registers and, for //! the one genuinely backward-in-time rule, patched into the already-written bitmap. //! -//! We do not re-derive character classes: `atomsplit::classify` already emits one `Atom` tag per -//! byte. The builder folds those 16 atoms into a grammar-specific **dense 3-bit code** and extracts -//! 3 bit-planes of it; every class stream is then a 2–3 op boolean function of the planes. +//! We do not re-derive character classes: [`classify`] already emits one `Atom` tag per byte. The +//! builder folds those 16 atoms into a grammar-specific **dense 3-bit code** and extracts 3 +//! bit-planes of it; every class stream is then a 2–3 op boolean function of the planes. //! //! Grammars: [`bitsplit_deepseek`], [`bitsplit_byte_level`] (GPT-2), [`bitsplit_cl100k`]. All three -//! byte-exact with their `atomsplit::fsm` counterparts — see `src/bin/verify.rs`. +//! byte-exact with the oniguruma oracle over a block-phase sweep — see `tests/parity.rs`. -pub(crate) use atomsplit::fsm::Span; - -pub mod deepseek; -pub mod gpt; -#[cfg(target_arch = "aarch64")] +pub mod classes; +pub mod classify; +mod han; +pub mod literal; +pub mod models; +pub mod regexes; mod simd; -#[cfg(target_arch = "x86_64")] -mod simd_x86; -pub use deepseek::bitsplit_deepseek; -pub use gpt::{bitsplit_byte_level, bitsplit_cl100k}; +pub use models::deepseek::bitsplit_deepseek; +pub use models::cl100k::{bitsplit_cl100k, bitsplit_qwen}; +pub use models::gpt2::bitsplit_byte_level; +pub use models::kimi::bitsplit_kimi; +pub use models::o200k::bitsplit_o200k; +pub use models::tekken::bitsplit_tekken; + +/// A token span: byte offsets `[start, end)` into the input. `#[repr(C)]` so the output buffer has a +/// stable `[start, end]` layout — the pipeline reuses it with zero conversion, and it can be +/// reinterpreted as bytes / handed across the crate boundary. +#[repr(C)] +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Hash, PartialOrd, Ord)] +pub struct Span { + pub start: u32, + pub end: u32, +} + +impl Span { + #[inline] + #[must_use] + pub const fn new(start: u32, end: u32) -> Self { + Self { start, end } + } + + /// `[start, end)` as a `usize` range — for slicing the input text. + #[inline] + #[must_use] + pub fn range(self) -> core::ops::Range { + self.start as usize..self.end as usize + } +} + +impl PartialEq<(u32, u32)> for Span { + fn eq(&self, other: &(u32, u32)) -> bool { + self.start == other.0 && self.end == other.1 + } +} /// Whether this target builds its bitstreams with SIMD. Without it the builder is the portable /// byte-at-a-time reference, which is slower than the FSM this replaces -- so a caller should keep @@ -56,7 +90,7 @@ pub fn fast_builder() -> bool { } #[cfg(target_arch = "x86_64")] -fn has_ssse3() -> bool { +pub(crate) fn has_ssse3() -> bool { use std::sync::atomic::{AtomicU8, Ordering}; static CACHED: AtomicU8 = AtomicU8::new(0); match CACHED.load(Ordering::Relaxed) { @@ -72,17 +106,29 @@ fn has_ssse3() -> bool { pub(crate) const CONT: u8 = 15; // Atom::Cont pub(crate) const CODE_CONT: u8 = 7; // every grammar's dense code for a continuation byte -/// The bitstreams for one 64-byte block. `p0`/`p1`/`p2` are the bit-planes of the **filled** dense +/// What a grammar wants in the text-derived `Blk::aux` stream. Tag-derived classes go through the +/// LUT; these three need the raw bytes, so they are the one thing the builder reads `text` for. +pub(crate) const AUX_NONE: u8 = 0; +pub(crate) const AUX_CJK: u8 = 1; // deepseek Split-2: Han U+4E00..9FA5 ∪ kana U+3040..30FF +pub(crate) const AUX_SLASH: u8 = 2; // o200k rule 4's `[\r\n/]*` tail +pub(crate) const AUX_HAN: u8 = 3; // kimi-k2's leading `[\p{Han}]+` arm + +/// The bitstreams for one 64-byte block. `p0`..`p3` are the bit-planes of the **filled** dense /// code — filled meaning a multi-byte char sets its bits on *all* of its bytes, so "previous char's /// class" is a plain `<< 1` and no rule does char-width arithmetic. `cont` is the one un-filled -/// stream (it defines the fill); `cjk` is text-derived and only built when a grammar asks for it. +/// stream (it defines the fill); `aux` is text-derived and only built when a grammar asks for it. +/// +/// `p3` exists because o200k needs 9 tag classes (U/L/C/N/NL/SP/WSO/OTHER + cont) and code 7 is +/// reserved — the SIMD kernels find continuation lanes by testing `lut[tag] == 7`, so "other" +/// cannot be the leftover code. It is const-gated off for the 3-plane grammars. #[derive(Default, Clone, Copy)] pub(crate) struct Blk { pub cont: u64, pub p0: u64, pub p1: u64, pub p2: u64, - pub cjk: u64, + pub p3: u64, + pub aux: u64, } /// deepseek Split-2's isolated range: Han U+4E00..9FA5 ∪ Hiragana/Katakana U+3040..30FF (all @@ -99,43 +145,56 @@ pub(crate) fn is_cjk_at(text: &[u8], p: usize) -> bool { (0x4E00..=0x9FA5).contains(&cp) || (0x3040..=0x30FF).contains(&cp) } +/// The `AUX` predicate at a lead byte. +#[inline] +pub(crate) fn aux_at(text: &[u8], p: usize) -> bool { + match AUX { + AUX_CJK => is_cjk_at(text, p), + AUX_SLASH => text[p] == b'/', + AUX_HAN => crate::han::is_han_at(text, p), + _ => false, + } +} + /// Build one block. Full blocks go through the NEON kernel; the ragged tail (and every other /// target) uses the portable byte-at-a-time reference. Both produce the identical `Blk`. -/// `CJK` asks for the (deepseek-only) range stream; `lut` is the grammar's tag → dense code table. +/// `lut` is the grammar's tag → dense code table. #[inline] -pub(crate) fn build_block( +pub(crate) fn build_block( text: &[u8], tags: &[u8], base: usize, len: usize, lut: &[u8; 64], cur_code: u8, - cur_cjk: bool, + cur_aux: bool, ) -> (Blk, u8) { #[cfg(target_arch = "aarch64")] if len == 64 { // SAFETY: `len == 64` means `base + 64 <= text.len() == tags.len()`. - return unsafe { crate::simd::build64::(text, tags, base, lut, cur_code, cur_cjk) }; + return unsafe { + crate::simd::neon::build64::(text, tags, base, lut, cur_code, cur_aux) + }; } #[cfg(target_arch = "x86_64")] if len == 64 && has_ssse3() { // SAFETY: `len == 64` bounds both reads, and SSSE3 is checked above. return unsafe { - crate::simd_x86::build64::(text, tags, base, lut, cur_code, cur_cjk) + crate::simd::x86::build64::(text, tags, base, lut, cur_code, cur_aux) }; } - build_block_scalar::(text, tags, base, len, lut, cur_code, cur_cjk) + build_block_scalar::(text, tags, base, len, lut, cur_code, cur_aux) } /// Portable reference builder: one byte at a time. -pub(crate) fn build_block_scalar( +pub(crate) fn build_block_scalar( text: &[u8], tags: &[u8], base: usize, len: usize, lut: &[u8; 64], mut cur_code: u8, - mut cur_cjk: bool, + mut cur_aux: bool, ) -> (Blk, u8) { let mut b = Blk::default(); for i in 0..len { @@ -145,12 +204,15 @@ pub(crate) fn build_block_scalar( b.cont |= bit; } else { cur_code = lut[tags[p] as usize]; - cur_cjk = CJK && is_cjk_at(text, p); + cur_aux = AUX != AUX_NONE && aux_at::(text, p); } b.p0 |= bit * u64::from(cur_code & 1 != 0); b.p1 |= bit * u64::from(cur_code & 2 != 0); b.p2 |= bit * u64::from(cur_code & 4 != 0); - b.cjk |= bit * u64::from(cur_cjk); + if P3 { + b.p3 |= bit * u64::from(cur_code & 8 != 0); + } + b.aux |= bit * u64::from(cur_aux); } (b, cur_code) } @@ -229,6 +291,143 @@ pub(crate) fn trail_run(x: u64, valid: u64, len: usize) -> u64 { } } +/// `\p{N}{1,cap}` — a group boundary every `cap` chars from the run start. The one non-local rule +/// in every tiktoken-family grammar, so it lives here rather than three times over. +/// +/// `cap == 0` or `>= 64` means an unbounded `\p{N}+`: the run is one token and there is nothing to +/// do (and the shifts below would be UB). `dig_since` resumes a run that crossed the block edge — +/// re-masking with `n` at every hop matters, because the carry only says the byte *at* the edge was +/// a digit, not that the run survived it. +#[inline] +pub(crate) fn digit_groups( + cap: usize, + n: u64, + lead: u64, + cont: u64, + prev_is_digit: bool, + dig_run: bool, + dig_since: u32, +) -> u64 { + let mut m = n & lead & !((n << 1) | u64::from(prev_is_digit)); + if cap == 0 || cap >= 64 { + return m; // an unbounded `\p{N}+`: the run is one token + } + let capu = cap as u32; + if dig_run && prev_is_digit { + let mut s = lead & lead.wrapping_neg() & n; // first lead of the block + for _ in 0..((capu - dig_since % capu) % capu) { + s = adv(s, cont) & n & lead; + } + m |= s; + } + let mut groups = m; + if n & cont == 0 { + // Fast path: every digit here is single-byte, so "CAP chars on" is a plain shift against a + // mask asking that the skipped positions were digits too — what the `adv` chain checks. + let mut nk = n; + let mut k = 1; + while k < cap { + nk &= n << k; + k += 1; + } + while m != 0 { + m = (m << cap) & nk; + groups |= m; + } + } else { + while m != 0 { + let mut e = m; + for _ in 0..cap { + e = adv(e, cont) & n & lead; + } + if e == 0 { + break; + } + groups |= e; + m = e; + } + } + groups +} + +// ── scalar helpers, shared by the grammars that need an escape ────────────────────────────── +use crate::classify::{Atom, in_mask, mask}; + + +pub(crate) const LET: u8 = 0x00; // coarse Atom::Letter (low nibble) +pub(crate) const MRK: u8 = 0x06; // coarse Atom::Mark +pub(crate) const ASM: u8 = 0x16; // AlphaSymMark — coarse Mark, categorically \p{S} +pub(crate) const ZWJ: u8 = 0x26; // ZWJ/ZWNJ — coarse Mark, categorically \p{Cf} +pub(crate) const NW: u8 = 0x01; +pub(crate) const NO: u8 = 0x02; +pub(crate) const NLN: u8 = 0x03; +pub(crate) const SPC: u8 = 0x04; +pub(crate) const WSO: u8 = 0x05; + +/// A real `[\p{L}\p{M}]` member. ALPHA_SYM and ZWJ are coarse `Mark` but neither `\p{L}` nor +/// `\p{M}`, so they are NOT letters — that is what keeps them on the rule-4 path. +#[inline] +pub(crate) fn member(t: u8) -> bool { + let c = t & 0x0F; + c == LET || (c == MRK && t != ASM && t != ZWJ) +} + +#[inline] +pub(crate) fn run_end(tags: &[u8], mut i: usize, end: usize, m: u16) -> usize { + let m = m | Atom::Cont.bit(); + while i < end && in_mask(tags[i], m) { + i += 1; + } + i +} + +/// `\s*[\r\n]+ | \s+(?!\S) | \s+`: through the last `\r\n` if any, else the whole run at EOF, else +/// give the final ws char back to whatever follows. +#[inline] +pub(crate) fn ws_tail(text: &[u8], tags: &[u8], i: usize, end: usize) -> usize { + let re = run_end(tags, i, end, mask::WS); + if let Some(r) = text[i..re].iter().rposition(|&x| x == 0x0A || x == 0x0D) { + i + r + 1 + } else if re == end { + re + } else { + let mut last = re - 1; + while last > i && text[last] & 0xC0 == 0x80 { + last -= 1; + } + if last > i { last } else { re } + } +} + +/// One letter sub-token from `p` within `[.., re)`: alt-1 `[UC]*[LC]+` (tried first), else alt-2 +/// `[UC]+[LC]*` for an all-upper run. `[UC]*` gives back to the last C so `[LC]+` can take one. +#[inline] +pub(crate) fn letter_match(tags: &[u8], p: usize, re: usize) -> usize { + let (mut q, mut last_c) = (p, usize::MAX); + while q < re && tags[q] != 0x20 { + if tags[q] != CONT && tags[q] != 0x10 { + last_c = q; + } + q += 1; + } + if q < re { + let mut e = q; + while e < re && tags[e] != 0x10 { + e += 1; + } + return e; + } + if last_c == usize::MAX { + return re; // all upper → alt-2 takes the whole run + } + let mut e = last_c; + while e < re && tags[e] != 0x10 { + e += 1; + } + e +} + + // ── emit ──────────────────────────────────────────────────────────────────────────────────────── /// `starts` bitmap → spans. Each set bit closes the previous token and opens the next; `tzcnt` @@ -329,10 +528,11 @@ pub(crate) fn emit_contr( w } + /// Byte length of the contraction at `i` (2 or 3), or 0. `ci` picks cl100k/o200k's `(?i:)` form /// over GPT-2's case-sensitive one. #[inline] -fn contr_len(text: &[u8], i: usize, ci: bool) -> usize { +pub(crate) fn contr_len(text: &[u8], i: usize, ci: bool) -> usize { let n = text.len(); if i + 1 >= n || text[i] != b'\'' || text[i + 1] >= 0x80 { return 0; @@ -357,10 +557,10 @@ pub fn build_only(text: &[u8], tags: &[u8]) -> u64 { let (mut acc, mut code, mut cjk) = (0u64, CODE_CONT, false); for base in (0..text.len()).step_by(64) { let len = (text.len() - base).min(64); - let (b, c) = build_block::(text, tags, base, len, &deepseek::LUT, code, cjk); + let (b, c) = build_block::<{ AUX_CJK }, false>(text, tags, base, len, &models::deepseek::LUT, code, cjk); code = c; - cjk = b.cjk >> (len - 1) & 1 != 0; - acc ^= b.cont ^ b.p0 ^ b.p1 ^ b.p2 ^ b.cjk; + cjk = b.aux >> (len - 1) & 1 != 0; + acc ^= b.cont ^ b.p0 ^ b.p1 ^ b.p2 ^ b.aux; } acc } @@ -368,6 +568,6 @@ pub fn build_only(text: &[u8], tags: &[u8]) -> u64 { /// Convenience wrapper: classify + deepseek bitsplit over caller-owned scratch. #[must_use] pub fn pre_tokenize(text: &[u8], tags: &mut [u8], starts: &mut [u64], out: &mut [Span]) -> usize { - atomsplit::classify::classify(text, tags); + classify::classify(text, tags); bitsplit_deepseek(text, tags, starts, out) } diff --git a/tokenizers/bitsplit/src/literal.rs b/tokenizers/bitsplit/src/literal.rs new file mode 100644 index 000000000..763e96d82 --- /dev/null +++ b/tokenizers/bitsplit/src/literal.rs @@ -0,0 +1,209 @@ +//! Splitting on an exact string — the Metaspace `▁` delimiter (llama2 / gemma), `CharDelimiterSplit`, +//! and any `Split` whose pattern is a plain string rather than a regex. +//! +//! - No atom classification here: a literal does not care what class a byte is, so the tag pass is +//! skipped entirely and this works off the raw bytes. +//! - Same bitstream shape as the grammars: build a `u64` match bitmap per 64-byte block, then walk +//! it with `trailing_zeros`. `shift-and` over the needle's bytes, which is the classic Baeza-Yates +//! bitap specialised to an exact string — one AND per needle byte per 64 input bytes. +//! - Matches never overlap, so `"aa"` is found once in `"aaa"` — the same matches a regex reports. + +use std::fmt; + +/// The pattern handed to [`Literal::new`] was empty. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct EmptyPattern; + +impl fmt::Display for EmptyPattern { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("an empty pattern matches everywhere") + } +} + +impl std::error::Error for EmptyPattern {} + +/// A literal string to split on. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Literal { + needle: Vec, +} + +impl Literal { + /// # Errors + /// If `pattern` is empty, which would match everywhere. + pub fn new(pattern: &[u8]) -> Result { + if pattern.is_empty() { + return Err(EmptyPattern); + } + Ok(Self { + needle: pattern.to_vec(), + }) + } + + /// The string being searched for. + #[must_use] + pub fn pattern(&self) -> &[u8] { + &self.needle + } + + /// Byte offset of every match, left to right, non-overlapping. + pub fn matches<'t>(&'t self, text: &'t [u8]) -> Matches<'t> { + Matches { + needle: &self.needle, + text, + base: 0, + word: 0, + next_block: 0, + min: 0, + } + } +} + +/// Match positions of `needle` in `text[base..base + 64]`, as a bitmap. A match at bit `i` means +/// `text[base + i..][..needle.len()] == needle`, so the last `needle.len() - 1` bits of the block +/// depend on bytes past it — the shift-and simply reads them from `text`, which is why this takes +/// the whole slice rather than a block. +fn match_bits(text: &[u8], needle: &[u8], base: usize) -> u64 { + let n = text.len(); + // positions in this block that could still fit the needle + let span = n.saturating_sub(base).min(64); + let last = n.saturating_sub(needle.len() - 1); // one past the last possible match start + let room = last.saturating_sub(base).min(span); + if room == 0 { + return 0; + } + let valid = if room == 64 { !0u64 } else { (1u64 << room) - 1 }; + // First byte in SIMD -- this is the whole cost, everything after it runs on the survivors. + let mut m = eq_bits(text, base, needle[0]) & valid; + for (k, &b) in needle.iter().enumerate().skip(1) { + let mut eq = 0u64; + let mut bits = m; + while bits != 0 { + let i = bits.trailing_zeros() as usize; + bits &= bits - 1; + if text[base + i + k] == b { + eq |= 1u64 << i; + } + } + m &= eq; + if m == 0 { + return 0; + } + } + m +} + +/// `text[base..base + 64] == b`, as a bitmap. Falls back to a byte loop on a ragged tail (and on +/// targets with no kernel), which is correct everywhere and only pays on the last block. +#[inline] +fn eq_bits(text: &[u8], base: usize, b: u8) -> u64 { + if base + 64 <= text.len() { + #[cfg(target_arch = "aarch64")] + // SAFETY: bounds checked directly above. + return unsafe { crate::simd::neon::eq64(text, base, b) }; + #[cfg(target_arch = "x86_64")] + if crate::has_ssse3() { + // SAFETY: bounds checked above, SSSE3 checked here. + return unsafe { crate::simd::x86::eq64(text, base, b) }; + } + } + let mut m = 0u64; + for i in 0..(text.len() - base).min(64) { + if text[base + i] == b { + m |= 1u64 << i; + } + } + m +} + +/// Iterator over non-overlapping match offsets. +pub struct Matches<'t> { + needle: &'t [u8], + text: &'t [u8], + base: usize, + word: u64, + next_block: usize, + min: usize, // matches before this would overlap the one just returned +} + +impl Iterator for Matches<'_> { + type Item = usize; + + fn next(&mut self) -> Option { + loop { + while self.word != 0 { + let i = self.word.trailing_zeros() as usize; + self.word &= self.word - 1; + let at = self.base + i; + if at >= self.min { + self.min = at + self.needle.len(); + return Some(at); + } + } + if self.next_block >= self.text.len() { + return None; + } + self.base = self.next_block; + self.next_block += 64; + self.word = match_bits(self.text, self.needle, self.base); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn find(pat: &str, hay: &str) -> Vec { + Literal::new(pat.as_bytes()) + .unwrap() + .matches(hay.as_bytes()) + .collect() + } + + /// The bitmap has to agree with a naive scan on overlap, block edges and multi-byte needles. + #[test] + fn matches_a_naive_scan() { + let cases: &[(&str, String)] = &[ + ("a", "aaaa".into()), + ("aa", "aaaaa".into()), + ("ab", "abab".into()), + ("▁", "▁a▁bb▁".into()), + ("▁▁", "▁▁▁▁".into()), + ("xyz", "xyxyz".into()), + ("a", "b".repeat(200)), + // needle straddling every 64-byte block edge + ("▁w", format!("{}▁w{}", "q".repeat(63), "z".repeat(70))), + ("abc", format!("{}abc", "a".repeat(62))), + ]; + for (pat, hay) in cases { + let (p, h) = (pat.as_bytes(), hay.as_bytes()); + let mut want = Vec::new(); + let mut i = 0; + while i + p.len() <= h.len() { + if &h[i..i + p.len()] == p { + want.push(i); + i += p.len(); // non-overlapping, like a regex + } else { + i += 1; + } + } + assert_eq!(find(pat, hay), want, "{pat:?} in {hay:?}"); + } + } + + /// Every offset of the needle across a 3-block text, so it crosses each edge in every phase. + #[test] + fn finds_the_needle_at_every_offset() { + let needle = "▁w"; + for off in 0..180usize { + let hay = format!("{}{needle}{}", "q".repeat(off), "q".repeat(180 - off)); + assert_eq!(find(needle, &hay), vec![off], "offset {off}"); + } + } + + #[test] + fn an_empty_pattern_is_rejected() { + assert_eq!(Literal::new(b""), Err(EmptyPattern)); + } +} diff --git a/tokenizers/bitsplit/src/gpt.rs b/tokenizers/bitsplit/src/models/cl100k.rs similarity index 62% rename from tokenizers/bitsplit/src/gpt.rs rename to tokenizers/bitsplit/src/models/cl100k.rs index a8b89bc8e..9d16e1b2c 100644 --- a/tokenizers/bitsplit/src/gpt.rs +++ b/tokenizers/bitsplit/src/models/cl100k.rs @@ -1,17 +1,12 @@ -//! The two tiktoken-family grammars as bitstream programs. +//! **cl100k_base** (tiktoken) / Llama-3 -- and byte-for-byte the regex **GLM-4.6** ships: +//! `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+` //! -//! * [`bitsplit_byte_level`] — GPT-2 / Llama / Qwen: -//! `'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+` -//! * [`bitsplit_cl100k`] — cl100k_base / Llama-3: -//! `(?i:'s|…)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+` -//! -//! Both share one dense code table, and both hand contractions to the scalar escape in -//! [`crate::emit_contr`] rather than trying to express a variable-length case-optional literal -//! alternation in bit algebra. +//! Rule 3's digit cap is the family's only knob, so it is a plain argument: 3 = cl100k / Llama-3 / +//! GLM, 1 = **Qwen** (`\p{N}`), 0 or >= 64 = an unbounded `\p{N}+`. use crate::{ - CODE_CONT, CONT, Span, adv, build_block, emit_contr, fill_to_last, lead_run, scanthru, to_lead, - trail_run, + AUX_NONE, CODE_CONT, CONT, Span, build_block, digit_groups, emit_contr, fill_to_last, + lead_run, scanthru, to_lead, trail_run, }; /// Atom tag → dense 3-bit code, shared by both grammars. Unlike deepseek's table, `Mark` is NOT a @@ -81,98 +76,37 @@ const fn code_bits(code: u8) -> u8 { } } -/// GPT-2 / byte-level pre-tokenization. `starts` and `flag` are scratch bitmaps -/// (len ≥ `text.len().div_ceil(64)`); byte-exact with `atomsplit::fsm::fsm_byte_level`. +/// cl100k_base / Llama-3 / GLM-4.6 — rule 3 is `\p{N}{1,3}`. #[must_use] -pub fn bitsplit_byte_level( +pub fn bitsplit_cl100k( text: &[u8], tags: &[u8], starts: &mut [u64], flag: &mut [u64], out: &mut [Span], ) -> usize { - let ntext = text.len(); - if ntext == 0 { - return 0; - } - let nblk = ntext.div_ceil(64); - assert!( - tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext - ); - let (mut code, mut prev_cont) = (CODE_CONT, 0u64); - - for bi in 0..nblk { - let base = bi * 64; - let len = (ntext - base).min(64); - let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; - let last_blk = base + len == ntext; - - let (b, last_code) = build_block::(text, tags, base, len, &LUT, code, false); - let c = decode(b.p0, b.p1, b.p2, valid); - - let pb = if base == 0 { 0 } else { code_bits(code) }; - let (nb, nb_lead) = if last_blk { - (0u8, true) - } else { - let q = base + len; - let is_lead = tags[q] != CONT; - ( - if is_lead { - code_bits(LUT[tags[q] as usize]) - } else { - code_bits(last_code) - }, - is_lead, - ) - }; - let has = |v: u8, s: u8| v & s != 0; - let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); - let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); - - let lead = valid & !b.cont; - let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); - - // ── every alternative but the contraction is ` ?X+` over a class run, so a token opens at - // each run start — pushed back one char when a literal space sits in front of it (` ?`). - let sp_pfx = p1(c.sp, has(pb, C_SP)); - let l_start = c.l & lead & !p1(c.l, has(pb, C_L)) & !sp_pfx; - let n_start = c.n & lead & !p1(c.n, has(pb, C_N)) & !sp_pfx; - let o_start = c.other & lead & !p1(c.other, has(pb, C_O)) & !sp_pfx; - let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); - // `\s+(?!\S)` hands the run's LAST whitespace char to whatever follows: as a ` ?` prefix if - // it is a space, else as a token of its own. Either way it opens a token — unless the run - // ends the input, where plain `\s+` takes the lot. GPT-2 has no `[\r\n]` rule, so unlike - // cl100k a newline is ordinary whitespace here and can be the stolen char. - let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; - let (steal, patch) = to_lead( - c.ws & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), - b.cont, - prev_cont, - ); - - let mut st = (l_start | n_start | o_start | ws_start | steal) & lead; - if bi == 0 { - st |= 1; - } - starts[bi] = st; - flag[bi] = st & c.apo; // apostrophes that open a token → contraction escape - if bi > 0 { - starts[bi - 1] |= patch; - } - code = last_code; - prev_cont = b.cont; - } - emit_contr(text, starts, flag, nblk, ntext, false, out) + cl100k(text, tags, starts, flag, out, 3) } -/// cl100k_base / Llama-3 pre-tokenization. Byte-exact with `atomsplit::fsm::fsm_cl100k`. +/// Qwen2 / Qwen3 — cl100k character-for-character except rule 3 is a bare `\p{N}`. #[must_use] -pub fn bitsplit_cl100k( +pub fn bitsplit_qwen( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + cl100k(text, tags, starts, flag, out, 1) +} + +fn cl100k( text: &[u8], tags: &[u8], starts: &mut [u64], flag: &mut [u64], out: &mut [Span], + digit_cap: usize, ) -> usize { let ntext = text.len(); if ntext == 0 { @@ -193,7 +127,7 @@ pub fn bitsplit_cl100k( let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; let last_blk = base + len == ntext; - let (b, last_code) = build_block::(text, tags, base, len, &LUT, code, false); + let (b, last_code) = build_block::<{ AUX_NONE }, false>(text, tags, base, len, &LUT, code, false); let c = decode(b.p0, b.p1, b.p2, valid); let pb = if base == 0 { 0 } else { code_bits(code) }; @@ -257,33 +191,16 @@ pub fn bitsplit_cl100k( & !p1(c.ws & !c.nl, has(pb, C_WS) && !has(pb, C_NL)) & !p1(osf, prev_osf); - // ── `\p{N}{1,3}`: a group boundary every 3 chars from the run start. - let mut m = c.n & lead & !p1(c.n, has(pb, C_N)); - if dig_run && has(pb, C_N) { - let mut s = lead & lead.wrapping_neg() & c.n; - for _ in 0..((3 - dig_since % 3) % 3) { - s = adv(s, b.cont) & c.n & lead; - } - m |= s; - } - let mut groups = m; - if c.n & b.cont == 0 { - let n3 = c.n & (c.n << 1) & (c.n << 2); - while m != 0 { - m = (m << 3) & n3; - groups |= m; - } - } else { - while m != 0 { - let a = adv(m, b.cont) & c.n & lead; - let e = adv(adv(a, b.cont) & c.n & lead, b.cont) & c.n & lead; - if e == 0 { - break; - } - groups |= e; - m = e; - } - } + // ── rule 3 `\p{N}{1,digit_cap}` + let groups = digit_groups( + digit_cap, + c.n, + lead, + b.cont, + has(pb, C_N), + dig_run, + dig_since, + ); // ── the other-run's `[\r\n]*` tail swallows the newlines directly behind it. let nl_m = ((p1(c.other, has(pb, C_O)) & c.nl & lead) as u128) | u128::from(nl_run); @@ -327,7 +244,11 @@ pub fn bitsplit_cl100k( } else { (c.n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() }; - counted % 3 + if digit_cap == 0 || digit_cap >= 64 { + counted + } else { + counted % digit_cap as u32 + } }; let tws = trail_run(c.ws, valid, len); if tws != 0 && has(nb, C_WS) { diff --git a/tokenizers/bitsplit/src/deepseek.rs b/tokenizers/bitsplit/src/models/deepseek.rs similarity index 97% rename from tokenizers/bitsplit/src/deepseek.rs rename to tokenizers/bitsplit/src/models/deepseek.rs index 3bafdbe91..7685b4d0a 100644 --- a/tokenizers/bitsplit/src/deepseek.rs +++ b/tokenizers/bitsplit/src/models/deepseek.rs @@ -3,7 +3,7 @@ //! `atomsplit::fsm::fsm_deepseek`. use crate::{ - CODE_CONT, CONT, Span, adv, build_block, emit, fill_to_last, is_cjk_at, lead_run, scanthru, + AUX_CJK, CODE_CONT, CONT, Span, adv, build_block, emit, fill_to_last, is_cjk_at, lead_run, scanthru, to_lead, trail_run, }; @@ -105,14 +105,14 @@ pub fn bitsplit_deepseek(text: &[u8], tags: &[u8], starts: &mut [u64], out: &mut let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; let last_blk = base + len == ntext; - let (b, last_code) = build_block::(text, tags, base, len, &LUT, cy.code, cy.cjk); + let (b, last_code) = build_block::<{ AUX_CJK }, false>(text, tags, base, len, &LUT, cy.code, cy.cjk); // planes → classes: 3 extractions instead of 7 one-hot masks (see `LUT`). Only `lm` // needs the `valid` mask — past the block end every plane reads 0, i.e. code 0. let (pa, pc, pd) = (!b.p2 & !b.p1, !b.p2 & b.p1, b.p2 & !b.p1); let (s_lm, s_n) = (pa & !b.p0 & valid, pa & b.p0); let (s_ps, s_nl) = (pc & !b.p0, pc & b.p0); let (s_sp, s_ws) = (pd & !b.p0, pd | (pc & b.p0)); - let last_cjk = b.cjk >> (len - 1) & 1 != 0; + let last_cjk = b.aux >> (len - 1) & 1 != 0; let last_bits = code_bits(last_code) | if last_cjk { S_CJK } else { 0 }; // ── edges. The byte before the block is carried; the one after is peeked. One char of @@ -138,7 +138,7 @@ pub fn bitsplit_deepseek(text: &[u8], tags: &[u8], starts: &mut [u64], out: &mut let has = |v: u16, s: u16| v & s != 0; // Split precedence: CJK (Split-2) outranks the big regex, and `fsm_deepseek` tests it ahead // of the digit arm too — so peel CJK off every other class first. - let cjk = b.cjk; + let cjk = b.aux; let num = s_n & !cjk; let lm = s_lm & !cjk; let ps = s_ps & !cjk; diff --git a/tokenizers/bitsplit/src/models/gpt2.rs b/tokenizers/bitsplit/src/models/gpt2.rs new file mode 100644 index 000000000..54567df91 --- /dev/null +++ b/tokenizers/bitsplit/src/models/gpt2.rs @@ -0,0 +1,162 @@ +//! **GPT-2 / ByteLevel**: +//! `'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+` +//! +//! Contractions are an ALTERNATIVE here (`don't` -> `don`, `'t`), and case-sensitive -- unlike +//! cl100k's `(?i:)`. They go to `emit_contr`'s scalar escape: a variable-length, case-optional +//! literal alternation that outranks every other arm is miserable in bit algebra and trivial there. + +use crate::{ + AUX_NONE, CODE_CONT, CONT, Span, build_block, emit_contr, to_lead, +}; + +/// Atom tag → dense 3-bit code, shared by both grammars. Unlike deepseek's table, `Mark` is NOT a +/// letter here (`\p{L}` excludes it, so it belongs to the "other" class), and `Apostrophe` gets its +/// own code so the contraction escape can be flagged with one AND — it is still "other" for every +/// run rule, which `decode` restores. +const LUT: [u8; 64] = { + let mut t = [2u8; 64]; // other = [^\s\p{L}\p{N}] + t[0x00] = 0; + t[0x10] = 0; + t[0x20] = 0; // Letter (+ case refinements) + t[0x01] = 1; + t[0x02] = 1; // \p{N} = Nd ∪ Nl ∪ No + t[0x03] = 3; // Newline + t[0x04] = 4; // Space + t[0x05] = 5; // WsOther + t[0x09] = 6; // Apostrophe + t[0x0F] = CODE_CONT; + t +}; + +/// One block's class streams. `other` folds the apostrophe code back in; only `l` needs the +/// `valid` mask, since past the block end every plane reads 0, i.e. code 0. +struct Cls { + l: u64, + n: u64, + other: u64, + nl: u64, + sp: u64, + ws: u64, + apo: u64, +} + +#[inline] +fn decode(p0: u64, p1: u64, p2: u64, valid: u64) -> Cls { + let a = !p2 & !p1; + let nl = !p2 & p1 & p0; + Cls { + l: a & !p0 & valid, + n: a & p0, + other: p1 & !p0 & valid, // codes 2 and 6 — "other" and the apostrophe + nl, + sp: p2 & !p1 & !p0, + ws: (p2 & !p1) | nl, // Space ∪ WsOther ∪ Newline + apo: p2 & p1 & !p0, + } +} + +/// Class bits of a dense code, for the block-edge carry. +const C_L: u8 = 1; +const C_N: u8 = 2; +const C_O: u8 = 4; +const C_NL: u8 = 8; +const C_SP: u8 = 16; +const C_WS: u8 = 32; + +#[inline] +const fn code_bits(code: u8) -> u8 { + match code { + 0 => C_L, + 1 => C_N, + 2 | 6 => C_O, + 3 => C_NL | C_WS, + 4 => C_SP | C_WS, + 5 => C_WS, + _ => 0, // cont never reaches an edge (the fill resolves it) + } +} + +/// GPT-2 / byte-level pre-tokenization. `starts` and `flag` are scratch bitmaps +/// (len ≥ `text.len().div_ceil(64)`); byte-exact with `atomsplit::fsm::fsm_byte_level`. +#[must_use] +pub fn bitsplit_byte_level( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + let nblk = ntext.div_ceil(64); + assert!( + tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext + ); + let (mut code, mut prev_cont) = (CODE_CONT, 0u64); + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = build_block::<{ AUX_NONE }, false>(text, tags, base, len, &LUT, code, false); + let c = decode(b.p0, b.p1, b.p2, valid); + + let pb = if base == 0 { 0 } else { code_bits(code) }; + let (nb, nb_lead) = if last_blk { + (0u8, true) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + ( + if is_lead { + code_bits(LUT[tags[q] as usize]) + } else { + code_bits(last_code) + }, + is_lead, + ) + }; + let has = |v: u8, s: u8| v & s != 0; + let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); + let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + + // ── every alternative but the contraction is ` ?X+` over a class run, so a token opens at + // each run start — pushed back one char when a literal space sits in front of it (` ?`). + let sp_pfx = p1(c.sp, has(pb, C_SP)); + let l_start = c.l & lead & !p1(c.l, has(pb, C_L)) & !sp_pfx; + let n_start = c.n & lead & !p1(c.n, has(pb, C_N)) & !sp_pfx; + let o_start = c.other & lead & !p1(c.other, has(pb, C_O)) & !sp_pfx; + let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); + // `\s+(?!\S)` hands the run's LAST whitespace char to whatever follows: as a ` ?` prefix if + // it is a space, else as a token of its own. Either way it opens a token — unless the run + // ends the input, where plain `\s+` takes the lot. GPT-2 has no `[\r\n]` rule, so unlike + // cl100k a newline is ordinary whitespace here and can be the stolen char. + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + let (steal, patch) = to_lead( + c.ws & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), + b.cont, + prev_cont, + ); + + let mut st = (l_start | n_start | o_start | ws_start | steal) & lead; + if bi == 0 { + st |= 1; + } + starts[bi] = st; + flag[bi] = st & c.apo; // apostrophes that open a token → contraction escape + if bi > 0 { + starts[bi - 1] |= patch; + } + code = last_code; + prev_cont = b.cont; + } + emit_contr(text, starts, flag, nblk, ntext, false, out) +} + diff --git a/tokenizers/bitsplit/src/models/kimi.rs b/tokenizers/bitsplit/src/models/kimi.rs new file mode 100644 index 000000000..994f0b35f --- /dev/null +++ b/tokenizers/bitsplit/src/models/kimi.rs @@ -0,0 +1,556 @@ +//! **kimi-k2 / k3** (`moonshotai/Kimi-K2-Instruct`, `tokenization_kimi.py`'s `pat_str`): +//! `[\p{Han}]+ | [^\r\n\p{L}\p{N}]?[\p{Lu}…&&[^\p{Han}]]*[\p{Ll}…&&[^\p{Han}]]+(?i:'s|…)? | …` +//! `| \p{N}{1,3} | ?[^\s\p{L}\p{N}]+[\r\n]* | \s*[\r\n]+ | \s+(?!\S) | \s+` +//! +//! o200k plus a leading `[\p{Han}]+` arm, Han subtracted from both letter classes, and a plain +//! `[\r\n]*` rule-4 tail (o200k has `[\r\n/]*`). Ships tiktoken.model, not a tokenizer.json. +//! +//! Two things are worth knowing before reading the algebra: +//! +//! 1. **The case split is a scalar escape, on purpose.** `[UC]*[LC]+ | [UC]+[LC]*` has no local +//! form: `中Qz` is ONE token (a U after a C is not a boundary) but `ʰABC` is two (it is, when no +//! L follows) — the difference is whether an L appears LATER in the run, so no `p1` decides it. +//! The bit half instead computes a cheap gate: a letter token is escaped only if the block holds +//! an interior upper or a trailing apostrophe, so all-lowercase and Capitalised text never pays. +//! 2. **The contraction is a SUFFIX here, not an alternative.** cl100k emits `'t` as its own token; +//! o200k glues it onto the letter token before it. So this uses `emit_contr_suffix`, which +//! *extends* the open token instead of opening one — the mirror image of `emit_contr`. +//! +use crate::classify::{char_len, in_mask, mask}; +use crate::{ + AUX_HAN, CODE_CONT, CONT, Span, build_block, contr_len, digit_groups, fill_to_last, lead_run, + letter_match, member, run_end, scanthru, to_lead, trail_run, ws_tail, +}; + +/// Atom tag → dense 4-bit code. Unlike cl100k, `\p{M}` IS a letter here (both alt classes list +/// `\p{M}`) — but only a true mark: `AlphaSymMark` (0x16, categorically `\p{S}`) and `Zwj` (0x26, +/// `\p{Cf}`) stay "other", which is what keeps `[\p{L}\p{M}]+` off them. +const LUT: [u8; 64] = { + let mut t = [8u8; 64]; // other = [^\s\p{L}\p{N}] (incl. Connector/Punct/Sym/Control/ASM/ZWJ) + t[0x10] = 0; // \p{Lu} ∪ \p{Lt} + t[0x20] = 1; // \p{Ll} + t[0x00] = 2; // caseless letter (\p{Lm}\p{Lo}) — in BOTH alt classes + t[0x06] = 10; // true \p{M}: a letter for the alts, but ALSO in rule 4's class (see `mark_adj`) + t[0x01] = 3; + t[0x02] = 3; // \p{N} + t[0x03] = 4; // Newline + t[0x04] = 5; // Space + t[0x05] = 6; // WsOther + t[0x09] = 9; // Apostrophe — "other" for every run rule, split out for the contraction flag + t[0x0F] = CODE_CONT; + t +}; + +struct Cls { + u: u64, + l: u64, + c: u64, + n: u64, + nl: u64, + sp: u64, + ws: u64, + oth: u64, + mark: u64, + apo: u64, +} + +#[inline] +fn decode(p0: u64, p1: u64, p2: u64, p3: u64, valid: u64) -> Cls { + let low = !p3; + let a = low & !p2; + let w = low & p2; + Cls { + u: a & !p1 & !p0 & valid, // code 0 — past the block end every plane reads 0, hence `valid` + l: a & !p1 & p0, + c: a & p1 & !p0, + n: a & p1 & p0, + nl: w & !p1 & !p0, + sp: w & !p1 & p0, + ws: w & !(p1 & p0), // codes 4,5,6 — cont (7) excluded + oth: p3 & !p1, // codes 8,9 — the apostrophe is "other" for run purposes + mark: p3 & p1, // code 10 + apo: p3 & p0 & !p1, // code 9 + } +} + +const C_U: u16 = 1 << 0; +const C_L: u16 = 1 << 1; +const C_C: u16 = 1 << 2; +const C_N: u16 = 1 << 3; +const C_NL: u16 = 1 << 4; +const C_SP: u16 = 1 << 5; +const C_WSO: u16 = 1 << 6; +const C_OTH: u16 = 1 << 7; +const C_MARK: u16 = 1 << 8; +const C_HAN: u16 = 1 << 9; +const C_WS: u16 = C_NL | C_SP | C_WSO; +const C_LET: u16 = C_U | C_L | C_C | C_MARK; + +const fn code_bits(code: u8) -> u16 { + match code { + 0 => C_U, + 1 => C_L, + 2 => C_C, + 3 => C_N, + 4 => C_NL, + 5 => C_SP, + 6 => C_WSO, + 8 | 9 => C_OTH, + 10 => C_MARK, + _ => 0, // cont never reaches an edge (the fill resolves it) + } +} + +/// kimi-k2 / k3. +#[must_use] +pub fn bitsplit_kimi( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + kimi(text, tags, starts, flag, out) +} + +fn kimi( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + let nblk = ntext.div_ceil(64); + assert!( + tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext + ); + + let (mut code, mut prev_cont) = (CODE_CONT, 0u64); + let (mut prev_han, mut nl_run, mut prev_osf) = (false, false, false); + let mut prev_absorbed = false; // the block's last byte was eaten by a `[\r\n/]*` tail + let (mut dig_run, mut dig_since) = (false, 0u32); + let mut anl: Option = None; + let mut last_lt: Option = None; + let mut prev_start: Option = None; + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = + build_block::<{ AUX_HAN }, true>(text, tags, base, len, &LUT, code, prev_han); + let c = decode(b.p0, b.p1, b.p2, b.p3, valid); + + let han = b.aux & valid; + let slash = 0; // kimi's rule-4 tail is a plain `[\r\n]*` + let last_han = han >> (len - 1) & 1 != 0; + + let pb = if base == 0 { + 0 + } else { + code_bits(code) | if prev_han { C_HAN } else { 0 } + }; + let (nb, nb_lead) = if last_blk { + (0u16, true) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + let bits = if is_lead { + code_bits(LUT[tags[q] as usize]) + | if crate::aux_at::<{ AUX_HAN }>(text, q) { + C_HAN + } else { + 0 + } + } else { + code_bits(last_code) | if last_han { C_HAN } else { 0 } + }; + (bits, is_lead) + }; + let has = |v: u16, s: u16| v & s != 0; + let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); + let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + + // ── Han is peeled off the letter classes first: kimi's alt-1 outranks the letter alts, and + // its letter classes are literally `[…&&[^\p{Han}]]`. With AUX != HAN this is all zero. + let (cu, cl, cc) = (c.u & !han, c.l & !han, c.c & !han); + let letter = cu | cl | cc | c.mark; + let pb_let = has(pb, C_LET) && !has(pb, C_HAN); + + // Rule 4's `[\r\n/]*` tail. `/` is in BOTH the `+` body and the tail, so one tail run can + // collect SEVERAL markers (`!\n/\n`: the `\n` after `!` and the `\n` after `/`). That is + // exactly what `fill_to_last` is for -- `nl_e - nl_m` would span only from the last one. + let tail_cls = c.nl | slash; + let nl_m64 = ((p1(c.oth, has(pb, C_OTH)) & c.nl & lead) | u64::from(nl_run)) & tail_cls; + let nl_m = nl_m64 as u128; + let nl_e = scanthru(nl_m, tail_cls as u128); + let nl_span = fill_to_last(nl_m64, tail_cls) as u128; + + // ── rule 4 ` ?[^\s\p{L}\p{N}]+[\r\n/]*` and rules 5-7 — identical to cl100k. + // A char absorbed by a `[\r\n/]*` tail is NOT part of the `+` body, so an "other" after it + // opens a fresh run (`\u{1f600}\r\n/#` is the tail, then `#` starts again). + let o_prev = c.oth & !(nl_span as u64); + let o_start = c.oth + & lead + & !(nl_span as u64) // a char INSIDE the tail never opens a run + & !p1(o_prev, has(pb, C_OTH) && !prev_absorbed) + & !p1(c.sp, has(pb, C_SP)); + let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); + let (steal, steal_patch) = to_lead( + c.ws & !c.nl & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), + b.cont, + prev_cont, + ); + let after_nl = p1(c.nl, has(pb, C_NL)) + & c.ws + & lead + & !fill_to_last(c.nl.reverse_bits(), c.ws.reverse_bits()).reverse_bits(); + + // ── `[^\r\n\p{L}\p{N}]?` before a letter run: any token-opening non-newline non-letter + // non-digit char. Smeared across its char's bytes so the test is a plain `p1` (see cl100k). + let mut osf = o_start; + osf |= (osf << 1) & b.cont; + osf |= (osf << 2) & b.cont & (b.cont << 1); + if prev_osf { + osf |= lead_run(b.cont, valid); + } + let l_start = letter + & lead + & !p1(letter, pb_let) + & !p1(c.ws & !c.nl, has(pb, C_WS) && !has(pb, C_NL)) + & !p1(osf, prev_osf); + // the token holding a letter run may OPEN one char earlier, on the `[^\r\n\p{L}\p{N}]?` + // prefix — that char is the start the escape has to be handed. The class genuinely excludes + // `\r\n` and digits: a newline before a letter run is its own token, never a prefix. + let prefix_cls = c.oth | c.sp | (c.ws & !c.nl); + let l_run = letter & lead & !p1(letter, pb_let); + let next_l_run = nb_lead + && has(nb, C_LET) + && !has(nb, C_HAN) + && letter >> (len - 1) & 1 == 0; + // `n1` only reads the NEXT char at a char's last byte, so mark there and walk back to the + // lead — a 3-byte prefix char (ZWJ, `—`, `½`) otherwise reads its own middle byte. + let (pfx_lead, pfx_patch) = + to_lead(prefix_cls & lb & n1(l_run, next_l_run), b.cont, prev_cont); + let l_start_tok = l_start | pfx_lead; + + // ── the escape gate (see the header). An interior upper, or an apostrophe closing the + // run, means some letter token in this block needs the scalar case/contraction pass. + let interior_u = cu & lead & p1(letter, pb_let); + let apo_after = c.apo & lead & p1(letter, pb_let); + // `\p{M}` is in BOTH the letter classes and rule 4's `[^\s\p{L}\p{N}]`, and which one wins + // depends on whether the punctuation before it STARTED the run (`!\u{301}a` is one token, + // `!!\u{301}a` is two). Adjacency is the gate; the scalar pass resolves it. Real text hits + // this via emoji + variation selector (U+FE0F is \p{Mn}). + let mark_adj = (c.mark & lead & p1(c.oth, has(pb, C_OTH))) + | (c.oth & lead & p1(c.mark, has(pb, C_MARK))); + + // ── kimi's `[\p{Han}]+` + let han_start = han & lead & !p1(han, has(pb, C_HAN)); + + // ── rule 3 `\p{N}{1,DIGIT_CAP}` + let groups = + digit_groups(3, c.n, lead, b.cont, has(pb, C_N), dig_run, dig_since); + + // ── the other-run's `[\r\n/]*` tail (`/` only for the o200k line; kimi has `[\r\n]*`). + // ── the one backward-in-time rule, exactly as in cl100k/deepseek: a newline arriving now + // retracts an "after the last newline" start committed for a run still open at the edge. + if let Some(p) = anl + && has(pb, C_WS) + && c.ws & 1 != 0 + && c.nl & lead_run(c.ws, valid) != 0 + { + starts[p / 64] &= !(1u64 << (p % 64)); + anl = None; + } + + let mut st = groups | l_start | han_start | o_start | ws_start | after_nl | steal; + st &= !(nl_span as u64); + st |= nl_e as u64; + st &= lead; + if bi == 0 { + st |= 1; + } + starts[bi] = st; + // Conservative: flag every letter token in a block that shows either trigger. A false + // positive only costs a rescan of that token, so erring wide is free; missing one is not. + // Only the letter RUNS that actually need the escape, not every letter token in the block: + // one `'s` in a paragraph would otherwise drag every word through the scalar pass. + // `apo_after` sits on the apostrophe, one byte past the run, so shift it back inside. + let needs = interior_u | ((apo_after >> 1) & letter); + let runs_needing = if needs == 0 { + 0 + } else { + fill_to_last(needs.reverse_bits(), letter.reverse_bits()).reverse_bits() + }; + let (pfx_need, _) = to_lead( + prefix_cls & lb & n1(runs_needing & l_run, false), + b.cont, + prev_cont, + ); + let lt_all = st & l_start_tok; // every letter token — what the cross-block patch tracks + let lt = lt_all & (runs_needing | pfx_need); + let trig = (interior_u | apo_after) != 0; + flag[bi] = if trig { lt } else { 0 }; + // mark_adj is rare, so be blunt: escape every token in the block, plus the one still open + // from an earlier block (rule 4's ` ?` means the token can have started on a space). + if mark_adj != 0 { + flag[bi] |= st; + if let Some(p) = prev_start { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if st != 0 { + prev_start = Some(base + 63 - st.leading_zeros() as usize); + } + if trig { + if bi > 0 { + flag[bi - 1] |= starts[bi - 1] & pfx_patch; + } + // the trigger can land in a LATER block than the token it belongs to (`\u{d55c}` ends + // block k, its `'s` opens block k+1), so always re-flag the last letter token seen. + if let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if bi > 0 { + starts[bi - 1] |= steal_patch; + } + // a run open at the block edge may meet its trigger in a later block, so flag it now + if lt_all != 0 { + last_lt = Some(base + 63 - lt_all.leading_zeros() as usize); + } + // ...ending on the run's `?` prefix char counts too — the token opened there. + let open_at_edge = (letter | l_start_tok) >> (len - 1) & 1 != 0; + if !last_blk && open_at_edge && let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + + // ── carries + nl_run = nl_e >> 64 != 0; + prev_absorbed = nl_span >> (len - 1) & 1 != 0; + let tn = trail_run(c.n, valid, len); + dig_run = tn != 0 && has(nb, C_N); + dig_since = if !dig_run { + 0 + } else { + let g = groups & tn; + let counted = if g == 0 { + dig_since + (c.n & lead & tn).count_ones() + } else { + (c.n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() + }; + counted % 3 + }; + let tws = trail_run(c.ws, valid, len); + if tws != 0 && has(nb, C_WS) { + let a = after_nl & tws & !(nl_e as u64); + if a != 0 { + anl = Some(base + 63 - a.leading_zeros() as usize); + } else if !(tws & 1 != 0 && has(pb, C_WS)) { + anl = None; + } + } else { + anl = None; + } + prev_osf = osf >> (len - 1) & 1 != 0; + prev_han = last_han; + code = last_code; + prev_cont = b.cont; + + } + + emit_kimi(text, tags, starts, flag, nblk, ntext, out) +} + +// ── the scalar escape ─────────────────────────────────────────────────────────────────────────── +// Ported from the FSM this replaces, which is the proven reading of the alternatives. It runs from +// a flagged token start and RESYNCS the moment it lands on an algebra start bit again, so a +// divergent region costs a short scalar walk and nothing downstream. + +/// Emit ONE token starting at `i` (letters emit several) and return the cursor past it. +fn step( + text: &[u8], + tags: &[u8], + i: usize, + end: usize, + out: &mut [Span], + w: &mut usize, +) -> usize { + let han = |p: usize| crate::aux_at::<{ AUX_HAN }>(text, p); + // kimi subtracts Han from both letter classes and isolates it in its own arm + let is_lm = |p: usize| p < end && member(tags[p]) && !han(p); + let letter_end = |mut p: usize| { + while p < end && (tags[p] == CONT || (member(tags[p]) && !han(p))) { + p += 1; + } + p + }; + let other = |sp0: usize| { + let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); + if p > sp0 { + while p < end + && (tags[p] == crate::NLN) + { + p += char_len(text[p]); + } + } + p + }; + let emit1 = |a: usize, b: usize, out: &mut [Span], w: &mut usize| { + out[*w] = Span::new(a as u32, b as u32); + *w += 1; + }; + // `[\p{Han}]+` — kimi's alternative 1, ahead of everything else + if han(i) { + let mut p = i; + while p < end && (tags[p] == CONT || han(p)) { + if tags[p] != CONT && !han(p) { + break; + } + p += 1; + } + emit1(i, p, out, w); + return p; + } + // the letter alternatives, with the case split and the optional contraction suffix + let letters = |pfx: usize, ls: usize, out: &mut [Span], w: &mut usize| -> usize { + let re = letter_end(ls); + let (mut p, mut first, mut cursor) = (ls, true, re); + while p < re { + let e = letter_match(tags, p, re); + let start = if first { pfx } else { p }; + let te = if e == re { + e + contr_len(text, e, true) + } else { + e + }; + out[*w] = Span::new(start as u32, te as u32); + *w += 1; + first = false; + cursor = te; + p = e; + } + cursor + }; + + let b = text[i]; + match tags[i] & 0x0F { + crate::NW | crate::NO => { + let (mut p, mut cnt) = (i, 0usize); + while p < end && cnt < 3 && in_mask(tags[p], mask::NUMBER) { + p += char_len(text[p]); + cnt += 1; + } + emit1(i, p, out, w); + p + } + crate::LET | crate::MRK => { + if member(tags[i]) && !han(i) { + return letters(i, i, out, w); + } + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + crate::SPC => { + let a = i + 1; + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(a); + let e = if p > a { p } else { ws_tail(text, tags, i, end) }; + emit1(i, e, out, w); + e + } + crate::WSO => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + crate::NLN => { + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + _ => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + } +} + +/// `emit` plus the escape: at a flagged start, hand over to the scalar dispatch and take back +/// control at the first position that is a start bit again. +fn emit_kimi( + text: &[u8], + tags: &[u8], + starts: &[u64], + flag: &[u64], + nblk: usize, + n: usize, + out: &mut [Span], +) -> usize { + let (mut w, mut open, mut skip) = (0usize, usize::MAX, 0usize); + for bi in 0..nblk { + let mut m = starts[bi]; + let f = flag[bi]; + while m != 0 { + let j = m.trailing_zeros() as usize; + let pos = bi * 64 + j; + m &= m - 1; + if pos < skip { + continue; + } + if f >> j & 1 != 0 { + if open != usize::MAX && open < pos { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + let mut c = pos; + loop { + c = step(text, tags, c, n, out, &mut w); + if c >= n || starts[c / 64] >> (c % 64) & 1 != 0 { + break; + } + } + open = usize::MAX; + skip = c; + continue; + } + if open != usize::MAX { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + open = pos; + } + } + if open != usize::MAX { + out[w] = Span::new(open as u32, n as u32); + w += 1; + } + w +} diff --git a/tokenizers/bitsplit/src/models/mod.rs b/tokenizers/bitsplit/src/models/mod.rs new file mode 100644 index 000000000..d80354a63 --- /dev/null +++ b/tokenizers/bitsplit/src/models/mod.rs @@ -0,0 +1,12 @@ +//! One module per pre-tokenization **regex**, each unrolled against the primitives in +//! [`crate`]. Models that ship the same regex share a module — o200k covers Llama-4, gpt-oss and +//! MiniMax-M2; cl100k covers Llama-3, GLM-4.6 and (at digit cap 1) Qwen. +//! +//! Deliberately not one parameterised grammar: folding these together made every fix a three-way +//! risk and none of them could be read on its own. +pub mod cl100k; +pub mod deepseek; +pub mod gpt2; +pub mod kimi; +pub mod o200k; +pub mod tekken; diff --git a/tokenizers/bitsplit/src/models/o200k.rs b/tokenizers/bitsplit/src/models/o200k.rs new file mode 100644 index 000000000..44dcc6437 --- /dev/null +++ b/tokenizers/bitsplit/src/models/o200k.rs @@ -0,0 +1,527 @@ +//! o200k_base / GPT-4o — byte-for-byte the regex **Llama-4, gpt-oss and MiniMax-M2** also ship: +//! `[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|…)?` +//! `|` same with `+`/`*` swapped `| \p{N}{1,3} | ?[^\s\p{L}\p{N}]+[\r\n/]* | \s*[\r\n]+ | \s+(?!\S) | \s+` +//! +//! Same skeleton as cl100k; only the letter half differs. Tekken and kimi are separate files — +//! they are separate regexes, and folding them in here is what made every fix a three-way risk. +//! +//! Two things are worth knowing before reading the algebra: +//! +//! 1. **The case split is a scalar escape, on purpose.** `[UC]*[LC]+ | [UC]+[LC]*` has no local +//! form: `中Qz` is ONE token (a U after a C is not a boundary) but `ʰABC` is two (it is, when no +//! L follows) — the difference is whether an L appears LATER in the run, so no `p1` decides it. +//! The bit half instead computes a cheap gate: a letter token is escaped only if the block holds +//! an interior upper or a trailing apostrophe, so all-lowercase and Capitalised text never pays. +//! 2. **The contraction is a SUFFIX here, not an alternative.** cl100k emits `'t` as its own token; +//! o200k glues it onto the letter token before it. So this uses `emit_contr_suffix`, which +//! *extends* the open token instead of opening one — the mirror image of `emit_contr`. +//! +use crate::classify::{char_len, in_mask, mask}; +use crate::{ + AUX_SLASH, CODE_CONT, CONT, Span, build_block, contr_len, digit_groups, fill_to_last, + lead_run, letter_match, member, run_end, scanthru, to_lead, trail_run, ws_tail, +}; + +/// Atom tag → dense 4-bit code. Unlike cl100k, `\p{M}` IS a letter here (both alt classes list +/// `\p{M}`) — but only a true mark: `AlphaSymMark` (0x16, categorically `\p{S}`) and `Zwj` (0x26, +/// `\p{Cf}`) stay "other", which is what keeps `[\p{L}\p{M}]+` off them. +const LUT: [u8; 64] = { + let mut t = [8u8; 64]; // other = [^\s\p{L}\p{N}] (incl. Connector/Punct/Sym/Control/ASM/ZWJ) + t[0x10] = 0; // \p{Lu} ∪ \p{Lt} + t[0x20] = 1; // \p{Ll} + t[0x00] = 2; // caseless letter (\p{Lm}\p{Lo}) — in BOTH alt classes + t[0x06] = 10; // true \p{M}: a letter for the alts, but ALSO in rule 4's class (see `mark_adj`) + t[0x01] = 3; + t[0x02] = 3; // \p{N} + t[0x03] = 4; // Newline + t[0x04] = 5; // Space + t[0x05] = 6; // WsOther + t[0x09] = 9; // Apostrophe — "other" for every run rule, split out for the contraction flag + t[0x0F] = CODE_CONT; + t +}; + +struct Cls { + u: u64, + l: u64, + c: u64, + n: u64, + nl: u64, + sp: u64, + ws: u64, + oth: u64, + mark: u64, + apo: u64, +} + +#[inline] +fn decode(p0: u64, p1: u64, p2: u64, p3: u64, valid: u64) -> Cls { + let low = !p3; + let a = low & !p2; + let w = low & p2; + Cls { + u: a & !p1 & !p0 & valid, // code 0 — past the block end every plane reads 0, hence `valid` + l: a & !p1 & p0, + c: a & p1 & !p0, + n: a & p1 & p0, + nl: w & !p1 & !p0, + sp: w & !p1 & p0, + ws: w & !(p1 & p0), // codes 4,5,6 — cont (7) excluded + oth: p3 & !p1, // codes 8,9 — the apostrophe is "other" for run purposes + mark: p3 & p1, // code 10 + apo: p3 & p0 & !p1, // code 9 + } +} + +const C_U: u16 = 1 << 0; +const C_L: u16 = 1 << 1; +const C_C: u16 = 1 << 2; +const C_N: u16 = 1 << 3; +const C_NL: u16 = 1 << 4; +const C_SP: u16 = 1 << 5; +const C_WSO: u16 = 1 << 6; +const C_OTH: u16 = 1 << 7; +const C_MARK: u16 = 1 << 8; +const C_WS: u16 = C_NL | C_SP | C_WSO; +const C_LET: u16 = C_U | C_L | C_C | C_MARK; + +const fn code_bits(code: u8) -> u16 { + match code { + 0 => C_U, + 1 => C_L, + 2 => C_C, + 3 => C_N, + 4 => C_NL, + 5 => C_SP, + 6 => C_WSO, + 8 | 9 => C_OTH, + 10 => C_MARK, + _ => 0, // cont never reaches an edge (the fill resolves it) + } +} + +/// o200k_base / GPT-4o — and byte-for-byte the same regex Llama-4, gpt-oss and MiniMax-M2 ship. +#[must_use] +pub fn bitsplit_o200k( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + o200k(text, tags, starts, flag, out) +} + +fn o200k( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + let nblk = ntext.div_ceil(64); + assert!( + tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext + ); + + let (mut code, mut prev_cont) = (CODE_CONT, 0u64); + let (mut nl_run, mut prev_osf) = (false, false); + let mut prev_absorbed = false; // the block's last byte was eaten by a `[\r\n/]*` tail + let (mut dig_run, mut dig_since) = (false, 0u32); + let mut anl: Option = None; + let mut last_lt: Option = None; + let mut prev_start: Option = None; + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = + build_block::<{ AUX_SLASH }, true>(text, tags, base, len, &LUT, code, false); + let c = decode(b.p0, b.p1, b.p2, b.p3, valid); + + let slash = b.aux; + + let pb = if base == 0 { + 0 + } else { + code_bits(code) + }; + let (nb, nb_lead) = if last_blk { + (0u16, true) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + let bits = if is_lead { + code_bits(LUT[tags[q] as usize]) + } else { + code_bits(last_code) + }; + (bits, is_lead) + }; + let has = |v: u16, s: u16| v & s != 0; + let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); + let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + + // ── Han is peeled off the letter classes first: kimi's alt-1 outranks the letter alts, and + // its letter classes are literally `[…&&[^\p{Han}]]`. With AUX != HAN this is all zero. + let (cu, cl, cc) = (c.u, c.l, c.c); + let letter = cu | cl | cc | c.mark; + let pb_let = has(pb, C_LET); + + // Rule 4's `[\r\n/]*` tail. `/` is in BOTH the `+` body and the tail, so one tail run can + // collect SEVERAL markers (`!\n/\n`: the `\n` after `!` and the `\n` after `/`). That is + // exactly what `fill_to_last` is for -- `nl_e - nl_m` would span only from the last one. + let tail_cls = c.nl | slash; + let nl_m64 = ((p1(c.oth, has(pb, C_OTH)) & c.nl & lead) | u64::from(nl_run)) & tail_cls; + let nl_m = nl_m64 as u128; + let nl_e = scanthru(nl_m, tail_cls as u128); + let nl_span = fill_to_last(nl_m64, tail_cls) as u128; + + // ── rule 4 ` ?[^\s\p{L}\p{N}]+[\r\n/]*` and rules 5-7 — identical to cl100k. + // A char absorbed by a `[\r\n/]*` tail is NOT part of the `+` body, so an "other" after it + // opens a fresh run (`\u{1f600}\r\n/#` is the tail, then `#` starts again). + let o_prev = c.oth & !(nl_span as u64); + let o_start = c.oth + & lead + & !(nl_span as u64) // a char INSIDE the tail never opens a run + & !p1(o_prev, has(pb, C_OTH) && !prev_absorbed) + & !p1(c.sp, has(pb, C_SP)); + let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); + let (steal, steal_patch) = to_lead( + c.ws & !c.nl & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), + b.cont, + prev_cont, + ); + let after_nl = p1(c.nl, has(pb, C_NL)) + & c.ws + & lead + & !fill_to_last(c.nl.reverse_bits(), c.ws.reverse_bits()).reverse_bits(); + + // ── `[^\r\n\p{L}\p{N}]?` before a letter run: any token-opening non-newline non-letter + // non-digit char. Smeared across its char's bytes so the test is a plain `p1` (see cl100k). + let mut osf = o_start; + osf |= (osf << 1) & b.cont; + osf |= (osf << 2) & b.cont & (b.cont << 1); + if prev_osf { + osf |= lead_run(b.cont, valid); + } + let l_start = letter + & lead + & !p1(letter, pb_let) + & !p1(c.ws & !c.nl, has(pb, C_WS) && !has(pb, C_NL)) + & !p1(osf, prev_osf); + // the token holding a letter run may OPEN one char earlier, on the `[^\r\n\p{L}\p{N}]?` + // prefix — that char is the start the escape has to be handed. The class genuinely excludes + // `\r\n` and digits: a newline before a letter run is its own token, never a prefix. + let prefix_cls = c.oth | c.sp | (c.ws & !c.nl); + let l_run = letter & lead & !p1(letter, pb_let); + let next_l_run = nb_lead + && has(nb, C_LET) + && letter >> (len - 1) & 1 == 0; + // `n1` only reads the NEXT char at a char's last byte, so mark there and walk back to the + // lead — a 3-byte prefix char (ZWJ, `—`, `½`) otherwise reads its own middle byte. + let (pfx_lead, pfx_patch) = + to_lead(prefix_cls & lb & n1(l_run, next_l_run), b.cont, prev_cont); + let l_start_tok = l_start | pfx_lead; + + // ── the escape gate (see the header). An interior upper, or an apostrophe closing the + // run, means some letter token in this block needs the scalar case/contraction pass. + let interior_u = cu & lead & p1(letter, pb_let); + let apo_after = c.apo & lead & p1(letter, pb_let); + // `\p{M}` is in BOTH the letter classes and rule 4's `[^\s\p{L}\p{N}]`, and which one wins + // depends on whether the punctuation before it STARTED the run (`!\u{301}a` is one token, + // `!!\u{301}a` is two). Adjacency is the gate; the scalar pass resolves it. Real text hits + // this via emoji + variation selector (U+FE0F is \p{Mn}). + let mark_adj = (c.mark & lead & p1(c.oth, has(pb, C_OTH))) + | (c.oth & lead & p1(c.mark, has(pb, C_MARK))); + + // ── rule 3 `\p{N}{1,DIGIT_CAP}` + let groups = + digit_groups(3, c.n, lead, b.cont, has(pb, C_N), dig_run, dig_since); + + // ── the other-run's `[\r\n/]*` tail (`/` only for the o200k line; kimi has `[\r\n]*`). + // ── the one backward-in-time rule, exactly as in cl100k/deepseek: a newline arriving now + // retracts an "after the last newline" start committed for a run still open at the edge. + if let Some(p) = anl + && has(pb, C_WS) + && c.ws & 1 != 0 + && c.nl & lead_run(c.ws, valid) != 0 + { + starts[p / 64] &= !(1u64 << (p % 64)); + anl = None; + } + + let mut st = groups | l_start | o_start | ws_start | after_nl | steal; + st &= !(nl_span as u64); + st |= nl_e as u64; + st &= lead; + if bi == 0 { + st |= 1; + } + starts[bi] = st; + // Conservative: flag every letter token in a block that shows either trigger. A false + // positive only costs a rescan of that token, so erring wide is free; missing one is not. + // Only the letter RUNS that actually need the escape, not every letter token in the block: + // one `'s` in a paragraph would otherwise drag every word through the scalar pass. + // `apo_after` sits on the apostrophe, one byte past the run, so shift it back inside. + let needs = interior_u | ((apo_after >> 1) & letter); + let runs_needing = if needs == 0 { + 0 + } else { + fill_to_last(needs.reverse_bits(), letter.reverse_bits()).reverse_bits() + }; + let (pfx_need, _) = to_lead( + prefix_cls & lb & n1(runs_needing & l_run, false), + b.cont, + prev_cont, + ); + let lt_all = st & l_start_tok; // every letter token — what the cross-block patch tracks + let lt = lt_all & (runs_needing | pfx_need); + let trig = (interior_u | apo_after) != 0; + flag[bi] = if trig { lt } else { 0 }; + // mark_adj is rare, so be blunt: escape every token in the block, plus the one still open + // from an earlier block (rule 4's ` ?` means the token can have started on a space). + if mark_adj != 0 { + flag[bi] |= st; + if let Some(p) = prev_start { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if st != 0 { + prev_start = Some(base + 63 - st.leading_zeros() as usize); + } + if trig { + if bi > 0 { + flag[bi - 1] |= starts[bi - 1] & pfx_patch; + } + // the trigger can land in a LATER block than the token it belongs to (`\u{d55c}` ends + // block k, its `'s` opens block k+1), so always re-flag the last letter token seen. + if let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if bi > 0 { + starts[bi - 1] |= steal_patch; + } + // a run open at the block edge may meet its trigger in a later block, so flag it now + if lt_all != 0 { + last_lt = Some(base + 63 - lt_all.leading_zeros() as usize); + } + // ...ending on the run's `?` prefix char counts too — the token opened there. + let open_at_edge = (letter | l_start_tok) >> (len - 1) & 1 != 0; + if !last_blk && open_at_edge && let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + + // ── carries + nl_run = nl_e >> 64 != 0; + prev_absorbed = nl_span >> (len - 1) & 1 != 0; + let tn = trail_run(c.n, valid, len); + dig_run = tn != 0 && has(nb, C_N); + dig_since = if !dig_run { + 0 + } else { + let g = groups & tn; + let counted = if g == 0 { + dig_since + (c.n & lead & tn).count_ones() + } else { + (c.n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() + }; + counted % 3 + }; + let tws = trail_run(c.ws, valid, len); + if tws != 0 && has(nb, C_WS) { + let a = after_nl & tws & !(nl_e as u64); + if a != 0 { + anl = Some(base + 63 - a.leading_zeros() as usize); + } else if !(tws & 1 != 0 && has(pb, C_WS)) { + anl = None; + } + } else { + anl = None; + } + prev_osf = osf >> (len - 1) & 1 != 0; + code = last_code; + prev_cont = b.cont; + + } + + emit_o200k(text, tags, starts, flag, nblk, ntext, out) +} + +// ── the scalar escape ─────────────────────────────────────────────────────────────────────────── +// Ported from the FSM this replaces, which is the proven reading of the alternatives. It runs from +// a flagged token start and RESYNCS the moment it lands on an algebra start bit again, so a +// divergent region costs a short scalar walk and nothing downstream. + +/// Emit ONE token starting at `i` (letters emit several) and return the cursor past it. +fn step( + text: &[u8], + tags: &[u8], + i: usize, + end: usize, + out: &mut [Span], + w: &mut usize, +) -> usize { + let is_lm = |p: usize| p < end && member(tags[p]); + let letter_end = |mut p: usize| { + while p < end && (tags[p] == CONT || member(tags[p])) { + p += 1; + } + p + }; + let other = |sp0: usize| { + let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); + if p > sp0 { + while p < end && (tags[p] == crate::NLN || text[p] == b'/') { + p += char_len(text[p]); + } + } + p + }; + let emit1 = |a: usize, b: usize, out: &mut [Span], w: &mut usize| { + out[*w] = Span::new(a as u32, b as u32); + *w += 1; + }; + // the letter alternatives, with the case split and the optional contraction suffix + let letters = |pfx: usize, ls: usize, out: &mut [Span], w: &mut usize| -> usize { + let re = letter_end(ls); + let (mut p, mut first, mut cursor) = (ls, true, re); + while p < re { + let e = letter_match(tags, p, re); + let start = if first { pfx } else { p }; + let te = if e == re { + e + contr_len(text, e, true) + } else { + e + }; + out[*w] = Span::new(start as u32, te as u32); + *w += 1; + first = false; + cursor = te; + p = e; + } + cursor + }; + + let b = text[i]; + match tags[i] & 0x0F { + crate::NW | crate::NO => { + let (mut p, mut cnt) = (i, 0usize); + while p < end && cnt < 3 && in_mask(tags[p], mask::NUMBER) { + p += char_len(text[p]); + cnt += 1; + } + emit1(i, p, out, w); + p + } + crate::LET | crate::MRK => { + if member(tags[i]) { + return letters(i, i, out, w); + } + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + crate::SPC => { + let a = i + 1; + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(a); + let e = if p > a { p } else { ws_tail(text, tags, i, end) }; + emit1(i, e, out, w); + e + } + crate::WSO => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + crate::NLN => { + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + _ => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + } +} + +/// `emit` plus the escape: at a flagged start, hand over to the scalar dispatch and take back +/// control at the first position that is a start bit again. +fn emit_o200k( + text: &[u8], + tags: &[u8], + starts: &[u64], + flag: &[u64], + nblk: usize, + n: usize, + out: &mut [Span], +) -> usize { + let (mut w, mut open, mut skip) = (0usize, usize::MAX, 0usize); + for bi in 0..nblk { + let mut m = starts[bi]; + let f = flag[bi]; + while m != 0 { + let j = m.trailing_zeros() as usize; + let pos = bi * 64 + j; + m &= m - 1; + if pos < skip { + continue; + } + if f >> j & 1 != 0 { + if open != usize::MAX && open < pos { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + let mut c = pos; + loop { + c = step(text, tags, c, n, out, &mut w); + if c >= n || starts[c / 64] >> (c % 64) & 1 != 0 { + break; + } + } + open = usize::MAX; + skip = c; + continue; + } + if open != usize::MAX { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + open = pos; + } + } + if open != usize::MAX { + out[w] = Span::new(open as u32, n as u32); + w += 1; + } + w +} diff --git a/tokenizers/bitsplit/src/models/tekken.rs b/tokenizers/bitsplit/src/models/tekken.rs new file mode 100644 index 000000000..e62ad14ed --- /dev/null +++ b/tokenizers/bitsplit/src/models/tekken.rs @@ -0,0 +1,525 @@ +//! Mistral **tekken** (mistral-small-4 / mistral-4): +//! `[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|…)?` +//! `|` same with `+`/`*` swapped `| \p{N} | ?[^\s\p{L}\p{N}]+[\r\n/]* | \s*[\r\n]+ | \s+(?!\S) | \s+` +//! +//! o200k's grammar with two changes: letter tokens take NO contraction suffix, and the digit rule +//! is a bare `\p{N}` — one token per digit. +//! + +//! +//! Two things are worth knowing before reading the algebra: +//! +//! 1. **The case split is a scalar escape, on purpose.** `[UC]*[LC]+ | [UC]+[LC]*` has no local +//! form: `中Qz` is ONE token (a U after a C is not a boundary) but `ʰABC` is two (it is, when no +//! L follows) — the difference is whether an L appears LATER in the run, so no `p1` decides it. +//! The bit half instead computes a cheap gate: a letter token is escaped only if the block holds +//! an interior upper or a trailing apostrophe, so all-lowercase and Capitalised text never pays. +//! 2. **The contraction is a SUFFIX here, not an alternative.** cl100k emits `'t` as its own token; +//! o200k glues it onto the letter token before it. So this uses `emit_contr_suffix`, which +//! *extends* the open token instead of opening one — the mirror image of `emit_contr`. +//! +use crate::classify::{char_len, in_mask, mask}; +use crate::{ + AUX_SLASH, CODE_CONT, CONT, Span, build_block, digit_groups, fill_to_last, + lead_run, letter_match, member, run_end, scanthru, to_lead, trail_run, ws_tail, +}; + +/// Atom tag → dense 4-bit code. Unlike cl100k, `\p{M}` IS a letter here (both alt classes list +/// `\p{M}`) — but only a true mark: `AlphaSymMark` (0x16, categorically `\p{S}`) and `Zwj` (0x26, +/// `\p{Cf}`) stay "other", which is what keeps `[\p{L}\p{M}]+` off them. +const LUT: [u8; 64] = { + let mut t = [8u8; 64]; // other = [^\s\p{L}\p{N}] (incl. Connector/Punct/Sym/Control/ASM/ZWJ) + t[0x10] = 0; // \p{Lu} ∪ \p{Lt} + t[0x20] = 1; // \p{Ll} + t[0x00] = 2; // caseless letter (\p{Lm}\p{Lo}) — in BOTH alt classes + t[0x06] = 10; // true \p{M}: a letter for the alts, but ALSO in rule 4's class (see `mark_adj`) + t[0x01] = 3; + t[0x02] = 3; // \p{N} + t[0x03] = 4; // Newline + t[0x04] = 5; // Space + t[0x05] = 6; // WsOther + t[0x09] = 9; // Apostrophe — "other" for every run rule, split out for the contraction flag + t[0x0F] = CODE_CONT; + t +}; + +struct Cls { + u: u64, + l: u64, + c: u64, + n: u64, + nl: u64, + sp: u64, + ws: u64, + oth: u64, + mark: u64, + apo: u64, +} + +#[inline] +fn decode(p0: u64, p1: u64, p2: u64, p3: u64, valid: u64) -> Cls { + let low = !p3; + let a = low & !p2; + let w = low & p2; + Cls { + u: a & !p1 & !p0 & valid, // code 0 — past the block end every plane reads 0, hence `valid` + l: a & !p1 & p0, + c: a & p1 & !p0, + n: a & p1 & p0, + nl: w & !p1 & !p0, + sp: w & !p1 & p0, + ws: w & !(p1 & p0), // codes 4,5,6 — cont (7) excluded + oth: p3 & !p1, // codes 8,9 — the apostrophe is "other" for run purposes + mark: p3 & p1, // code 10 + apo: p3 & p0 & !p1, // code 9 + } +} + +const C_U: u16 = 1 << 0; +const C_L: u16 = 1 << 1; +const C_C: u16 = 1 << 2; +const C_N: u16 = 1 << 3; +const C_NL: u16 = 1 << 4; +const C_SP: u16 = 1 << 5; +const C_WSO: u16 = 1 << 6; +const C_OTH: u16 = 1 << 7; +const C_MARK: u16 = 1 << 8; +const C_WS: u16 = C_NL | C_SP | C_WSO; +const C_LET: u16 = C_U | C_L | C_C | C_MARK; + +const fn code_bits(code: u8) -> u16 { + match code { + 0 => C_U, + 1 => C_L, + 2 => C_C, + 3 => C_N, + 4 => C_NL, + 5 => C_SP, + 6 => C_WSO, + 8 | 9 => C_OTH, + 10 => C_MARK, + _ => 0, // cont never reaches an edge (the fill resolves it) + } +} + +/// Mistral tekken. +#[must_use] +pub fn bitsplit_tekken( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + tekken(text, tags, starts, flag, out) +} + +fn tekken( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + let nblk = ntext.div_ceil(64); + assert!( + tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext + ); + + let (mut code, mut prev_cont) = (CODE_CONT, 0u64); + let (mut nl_run, mut prev_osf) = (false, false); + let mut prev_absorbed = false; // the block's last byte was eaten by a `[\r\n/]*` tail + let (mut dig_run, mut dig_since) = (false, 0u32); + let mut anl: Option = None; + let mut last_lt: Option = None; + let mut prev_start: Option = None; + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = + build_block::<{ AUX_SLASH }, true>(text, tags, base, len, &LUT, code, false); + let c = decode(b.p0, b.p1, b.p2, b.p3, valid); + + let slash = b.aux; + + let pb = if base == 0 { + 0 + } else { + code_bits(code) + }; + let (nb, nb_lead) = if last_blk { + (0u16, true) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + let bits = if is_lead { + code_bits(LUT[tags[q] as usize]) + } else { + code_bits(last_code) + }; + (bits, is_lead) + }; + let has = |v: u16, s: u16| v & s != 0; + let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); + let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + + // ── Han is peeled off the letter classes first: kimi's alt-1 outranks the letter alts, and + // its letter classes are literally `[…&&[^\p{Han}]]`. With AUX != HAN this is all zero. + let (cu, cl, cc) = (c.u, c.l, c.c); + let letter = cu | cl | cc | c.mark; + let pb_let = has(pb, C_LET); + + // Rule 4's `[\r\n/]*` tail. `/` is in BOTH the `+` body and the tail, so one tail run can + // collect SEVERAL markers (`!\n/\n`: the `\n` after `!` and the `\n` after `/`). That is + // exactly what `fill_to_last` is for -- `nl_e - nl_m` would span only from the last one. + let tail_cls = c.nl | slash; + let nl_m64 = ((p1(c.oth, has(pb, C_OTH)) & c.nl & lead) | u64::from(nl_run)) & tail_cls; + let nl_m = nl_m64 as u128; + let nl_e = scanthru(nl_m, tail_cls as u128); + let nl_span = fill_to_last(nl_m64, tail_cls) as u128; + + // ── rule 4 ` ?[^\s\p{L}\p{N}]+[\r\n/]*` and rules 5-7 — identical to cl100k. + // A char absorbed by a `[\r\n/]*` tail is NOT part of the `+` body, so an "other" after it + // opens a fresh run (`\u{1f600}\r\n/#` is the tail, then `#` starts again). + let o_prev = c.oth & !(nl_span as u64); + let o_start = c.oth + & lead + & !(nl_span as u64) // a char INSIDE the tail never opens a run + & !p1(o_prev, has(pb, C_OTH) && !prev_absorbed) + & !p1(c.sp, has(pb, C_SP)); + let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); + let (steal, steal_patch) = to_lead( + c.ws & !c.nl & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), + b.cont, + prev_cont, + ); + let after_nl = p1(c.nl, has(pb, C_NL)) + & c.ws + & lead + & !fill_to_last(c.nl.reverse_bits(), c.ws.reverse_bits()).reverse_bits(); + + // ── `[^\r\n\p{L}\p{N}]?` before a letter run: any token-opening non-newline non-letter + // non-digit char. Smeared across its char's bytes so the test is a plain `p1` (see cl100k). + let mut osf = o_start; + osf |= (osf << 1) & b.cont; + osf |= (osf << 2) & b.cont & (b.cont << 1); + if prev_osf { + osf |= lead_run(b.cont, valid); + } + let l_start = letter + & lead + & !p1(letter, pb_let) + & !p1(c.ws & !c.nl, has(pb, C_WS) && !has(pb, C_NL)) + & !p1(osf, prev_osf); + // the token holding a letter run may OPEN one char earlier, on the `[^\r\n\p{L}\p{N}]?` + // prefix — that char is the start the escape has to be handed. The class genuinely excludes + // `\r\n` and digits: a newline before a letter run is its own token, never a prefix. + let prefix_cls = c.oth | c.sp | (c.ws & !c.nl); + let l_run = letter & lead & !p1(letter, pb_let); + let next_l_run = nb_lead + && has(nb, C_LET) + && letter >> (len - 1) & 1 == 0; + // `n1` only reads the NEXT char at a char's last byte, so mark there and walk back to the + // lead — a 3-byte prefix char (ZWJ, `—`, `½`) otherwise reads its own middle byte. + let (pfx_lead, pfx_patch) = + to_lead(prefix_cls & lb & n1(l_run, next_l_run), b.cont, prev_cont); + let l_start_tok = l_start | pfx_lead; + + // ── the escape gate (see the header). An interior upper, or an apostrophe closing the + // run, means some letter token in this block needs the scalar case/contraction pass. + let interior_u = cu & lead & p1(letter, pb_let); + // no contraction suffix here, so an apostrophe after a letter is nothing special + // `\p{M}` is in BOTH the letter classes and rule 4's `[^\s\p{L}\p{N}]`, and which one wins + // depends on whether the punctuation before it STARTED the run (`!\u{301}a` is one token, + // `!!\u{301}a` is two). Adjacency is the gate; the scalar pass resolves it. Real text hits + // this via emoji + variation selector (U+FE0F is \p{Mn}). + let mark_adj = (c.mark & lead & p1(c.oth, has(pb, C_OTH))) + | (c.oth & lead & p1(c.mark, has(pb, C_MARK))); + + // ── rule 3 `\p{N}{1,DIGIT_CAP}` + let groups = + digit_groups(1, c.n, lead, b.cont, has(pb, C_N), dig_run, dig_since); + + // ── the other-run's `[\r\n/]*` tail (`/` only for the o200k line; kimi has `[\r\n]*`). + // ── the one backward-in-time rule, exactly as in cl100k/deepseek: a newline arriving now + // retracts an "after the last newline" start committed for a run still open at the edge. + if let Some(p) = anl + && has(pb, C_WS) + && c.ws & 1 != 0 + && c.nl & lead_run(c.ws, valid) != 0 + { + starts[p / 64] &= !(1u64 << (p % 64)); + anl = None; + } + + let mut st = groups | l_start | o_start | ws_start | after_nl | steal; + st &= !(nl_span as u64); + st |= nl_e as u64; + st &= lead; + if bi == 0 { + st |= 1; + } + starts[bi] = st; + // Conservative: flag every letter token in a block that shows either trigger. A false + // positive only costs a rescan of that token, so erring wide is free; missing one is not. + // Only the letter RUNS that actually need the escape, not every letter token in the block: + // one `'s` in a paragraph would otherwise drag every word through the scalar pass. + // `apo_after` sits on the apostrophe, one byte past the run, so shift it back inside. + let needs = interior_u; + let runs_needing = if needs == 0 { + 0 + } else { + fill_to_last(needs.reverse_bits(), letter.reverse_bits()).reverse_bits() + }; + let (pfx_need, _) = to_lead( + prefix_cls & lb & n1(runs_needing & l_run, false), + b.cont, + prev_cont, + ); + let lt_all = st & l_start_tok; // every letter token — what the cross-block patch tracks + let lt = lt_all & (runs_needing | pfx_need); + let trig = interior_u != 0; + flag[bi] = if trig { lt } else { 0 }; + // mark_adj is rare, so be blunt: escape every token in the block, plus the one still open + // from an earlier block (rule 4's ` ?` means the token can have started on a space). + if mark_adj != 0 { + flag[bi] |= st; + if let Some(p) = prev_start { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if st != 0 { + prev_start = Some(base + 63 - st.leading_zeros() as usize); + } + if trig { + if bi > 0 { + flag[bi - 1] |= starts[bi - 1] & pfx_patch; + } + // the trigger can land in a LATER block than the token it belongs to (`\u{d55c}` ends + // block k, its `'s` opens block k+1), so always re-flag the last letter token seen. + if let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if bi > 0 { + starts[bi - 1] |= steal_patch; + } + // a run open at the block edge may meet its trigger in a later block, so flag it now + if lt_all != 0 { + last_lt = Some(base + 63 - lt_all.leading_zeros() as usize); + } + // ...ending on the run's `?` prefix char counts too — the token opened there. + let open_at_edge = (letter | l_start_tok) >> (len - 1) & 1 != 0; + if !last_blk && open_at_edge && let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + + // ── carries + nl_run = nl_e >> 64 != 0; + prev_absorbed = nl_span >> (len - 1) & 1 != 0; + let tn = trail_run(c.n, valid, len); + dig_run = tn != 0 && has(nb, C_N); + dig_since = if !dig_run { + 0 + } else { + let g = groups & tn; + let _counted = if g == 0 { + dig_since + (c.n & lead & tn).count_ones() + } else { + (c.n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() + }; + 0 // digit cap 1: every digit is its own group, so nothing carries + }; + let tws = trail_run(c.ws, valid, len); + if tws != 0 && has(nb, C_WS) { + let a = after_nl & tws & !(nl_e as u64); + if a != 0 { + anl = Some(base + 63 - a.leading_zeros() as usize); + } else if !(tws & 1 != 0 && has(pb, C_WS)) { + anl = None; + } + } else { + anl = None; + } + prev_osf = osf >> (len - 1) & 1 != 0; + code = last_code; + prev_cont = b.cont; + + } + + emit_tekken(text, tags, starts, flag, nblk, ntext, out) +} + +// ── the scalar escape ─────────────────────────────────────────────────────────────────────────── +// Ported from the FSM this replaces, which is the proven reading of the alternatives. It runs from +// a flagged token start and RESYNCS the moment it lands on an algebra start bit again, so a +// divergent region costs a short scalar walk and nothing downstream. + +/// Emit ONE token starting at `i` (letters emit several) and return the cursor past it. +fn step( + text: &[u8], + tags: &[u8], + i: usize, + end: usize, + out: &mut [Span], + w: &mut usize, +) -> usize { + let is_lm = |p: usize| p < end && member(tags[p]); + let letter_end = |mut p: usize| { + while p < end && (tags[p] == CONT || member(tags[p])) { + p += 1; + } + p + }; + let other = |sp0: usize| { + let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); + if p > sp0 { + while p < end && (tags[p] == crate::NLN || text[p] == b'/') { + p += char_len(text[p]); + } + } + p + }; + let emit1 = |a: usize, b: usize, out: &mut [Span], w: &mut usize| { + out[*w] = Span::new(a as u32, b as u32); + *w += 1; + }; + // the letter alternatives, with the case split and the optional contraction suffix + let letters = |pfx: usize, ls: usize, out: &mut [Span], w: &mut usize| -> usize { + let re = letter_end(ls); + let (mut p, mut first, mut cursor) = (ls, true, re); + while p < re { + let e = letter_match(tags, p, re); + let start = if first { pfx } else { p }; + let te = e; // tekken has no contraction suffix + out[*w] = Span::new(start as u32, te as u32); + *w += 1; + first = false; + cursor = te; + p = e; + } + cursor + }; + + let b = text[i]; + match tags[i] & 0x0F { + crate::NW | crate::NO => { + let (mut p, mut cnt) = (i, 0usize); + while p < end && cnt < 1 && in_mask(tags[p], mask::NUMBER) { + p += char_len(text[p]); + cnt += 1; + } + emit1(i, p, out, w); + p + } + crate::LET | crate::MRK => { + if member(tags[i]) { + return letters(i, i, out, w); + } + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + crate::SPC => { + let a = i + 1; + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(a); + let e = if p > a { p } else { ws_tail(text, tags, i, end) }; + emit1(i, e, out, w); + e + } + crate::WSO => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + crate::NLN => { + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + _ => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + } +} + +/// `emit` plus the escape: at a flagged start, hand over to the scalar dispatch and take back +/// control at the first position that is a start bit again. +fn emit_tekken( + text: &[u8], + tags: &[u8], + starts: &[u64], + flag: &[u64], + nblk: usize, + n: usize, + out: &mut [Span], +) -> usize { + let (mut w, mut open, mut skip) = (0usize, usize::MAX, 0usize); + for bi in 0..nblk { + let mut m = starts[bi]; + let f = flag[bi]; + while m != 0 { + let j = m.trailing_zeros() as usize; + let pos = bi * 64 + j; + m &= m - 1; + if pos < skip { + continue; + } + if f >> j & 1 != 0 { + if open != usize::MAX && open < pos { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + let mut c = pos; + loop { + c = step(text, tags, c, n, out, &mut w); + if c >= n || starts[c / 64] >> (c % 64) & 1 != 0 { + break; + } + } + open = usize::MAX; + skip = c; + continue; + } + if open != usize::MAX { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + open = pos; + } + } + if open != usize::MAX { + out[w] = Span::new(open as u32, n as u32); + w += 1; + } + w +} diff --git a/tokenizers/atomsplit/src/regexes.rs b/tokenizers/bitsplit/src/regexes.rs similarity index 77% rename from tokenizers/atomsplit/src/regexes.rs rename to tokenizers/bitsplit/src/regexes.rs index 875fb892e..3f3a0662d 100644 --- a/tokenizers/atomsplit/src/regexes.rs +++ b/tokenizers/bitsplit/src/regexes.rs @@ -31,3 +31,9 @@ pub const DEEPSEEK_BIG: &str = r##"[!"#$%&'()*+,\-./:;<=>?@\[\\\]^_`{|}~][A-Za-z /// The deepseek chain in application order — convenience for the multi-regex reference. pub const DEEPSEEK: &[&str] = &[DEEPSEEK_NUM, DEEPSEEK_CJK, DEEPSEEK_BIG]; + +/// kimi-k2 / k3 — `moonshotai/Kimi-K2-Instruct`'s `tokenization_kimi.py` `pat_str`. o200k plus a +/// leading `[\p{Han}]+` arm, Han subtracted from both letter classes, and a `[\r\n]*` rule-4 tail +/// (o200k has `[\r\n/]*`). Kimi ships `tiktoken.model` rather than a `tokenizer.json`, so this is +/// the pattern as a converted tokenizer would spell it. Reproduced by [`crate::bitsplit_kimi`]. +pub const KIMI_K2: &str = r"[\p{Han}]+|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"; diff --git a/tokenizers/atomsplit/src/simd_fsm.rs b/tokenizers/bitsplit/src/simd/classes.rs similarity index 99% rename from tokenizers/atomsplit/src/simd_fsm.rs rename to tokenizers/bitsplit/src/simd/classes.rs index ed6b172cb..b6700e1c4 100644 --- a/tokenizers/atomsplit/src/simd_fsm.rs +++ b/tokenizers/bitsplit/src/simd/classes.rs @@ -8,7 +8,8 @@ // the non-selected arch some items/imports (char_len, emit_class_spans, ...) are legitimately unused use crate::classify::{Atom, char_len}; -use crate::fsm::{Span, emit_class_spans}; +use crate::Span; +use crate::classes::emit_class_spans; /// Class LookUpTable: tag → 0 drop / 1 isolate / 2 keep-A / 3 keep-B; Cont → 0xFF (fill sentinel). /// This lookup table is built per parameter DROP, ISOLATE and KEEP_A. These as diff --git a/tokenizers/bitsplit/src/simd/mod.rs b/tokenizers/bitsplit/src/simd/mod.rs new file mode 100644 index 000000000..76161133a --- /dev/null +++ b/tokenizers/bitsplit/src/simd/mod.rs @@ -0,0 +1,12 @@ +//! Per-arch kernels. Every one of these is PURE PERF and byte-exact with a portable path that is +//! always compiled, so correctness never depends on a kernel being present: +//! - [`neon`] / [`x86`] build the block's bitstreams; `build_block_scalar` is the reference. +//! - [`classes`] extracts class-run boundaries; `emit_class_spans` is the reference. +//! +//! (`classify` has its own kernels, next to the tables they index.) +#[cfg(target_arch = "aarch64")] +pub(crate) mod neon; +#[cfg(target_arch = "x86_64")] +pub(crate) mod x86; + +pub(crate) mod classes; diff --git a/tokenizers/bitsplit/src/simd.rs b/tokenizers/bitsplit/src/simd/neon.rs similarity index 76% rename from tokenizers/bitsplit/src/simd.rs rename to tokenizers/bitsplit/src/simd/neon.rs index d328dc502..e685ab7ba 100644 --- a/tokenizers/bitsplit/src/simd.rs +++ b/tokenizers/bitsplit/src/simd/neon.rs @@ -10,7 +10,7 @@ //! on all of its bytes. That is what lets the bitstream program read "previous char's class" as a //! plain `<< 1` with no char-width arithmetic. -use crate::{Blk, lead_run}; +use crate::{AUX_HAN, AUX_NONE, AUX_SLASH, Blk, lead_run}; use core::arch::aarch64::*; const POW: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128]; @@ -27,6 +27,27 @@ unsafe fn mm64(v: [uint8x16_t; 4], pow: uint8x16_t) -> u64 { } } +/// 64 bytes of `text[base..]` compared against `b`, as a bitmap. Same fold the block builder uses. +/// +/// # Safety +/// `base + 64 <= text.len()`. +#[target_feature(enable = "neon")] +pub(crate) unsafe fn eq64(text: &[u8], base: usize, b: u8) -> u64 { + unsafe { + let pow = vld1q_u8(POW.as_ptr()); + let n = vdupq_n_u8(b); + mm64( + [ + vceqq_u8(vld1q_u8(text.as_ptr().add(base)), n), + vceqq_u8(vld1q_u8(text.as_ptr().add(base + 16)), n), + vceqq_u8(vld1q_u8(text.as_ptr().add(base + 32)), n), + vceqq_u8(vld1q_u8(text.as_ptr().add(base + 48)), n), + ], + pow, + ) + } +} + /// Build one **full** 64-byte block, folding tags through `lut` into dense codes. /// Build one **full** 64-byte block. `cur_code` / `cur_cjk` describe the byte just before it, so a /// block opening mid-char keeps inheriting its lead's class. Returns the block's last filled code. @@ -34,13 +55,13 @@ unsafe fn mm64(v: [uint8x16_t; 4], pow: uint8x16_t) -> u64 { /// # Safety /// `base + 64 <= tags.len()` and `base + 64 <= text.len()`. #[target_feature(enable = "neon")] -pub(crate) unsafe fn build64( +pub(crate) unsafe fn build64( text: &[u8], tags: &[u8], base: usize, lut: &[u8; 64], cur_code: u8, - cur_cjk: bool, + cur_aux: bool, ) -> (Blk, u8) { unsafe { let pow = vld1q_u8(POW.as_ptr()); @@ -89,12 +110,13 @@ pub(crate) unsafe fn build64( p0: plane(1), p1: plane(2), p2: plane(4), - cjk: 0, + p3: if P3 { plane(8) } else { 0 }, + aux: 0, }; - if !CJK { + if AUX == AUX_NONE { return (b, last_code); } - // ── text is loaded only for the (deepseek-only) CJK range test. + // ── text is loaded only for the aux (text-derived) stream. let ntext = text.len(); let tv = [ vld1q_u8(text.as_ptr().add(base)), @@ -103,6 +125,37 @@ pub(crate) unsafe fn build64( vld1q_u8(text.as_ptr().add(base + 48)), ]; + if AUX == AUX_SLASH { + // single ASCII byte — no fill needed, `/` is its own char + let sl = vdupq_n_u8(b'/'); + b.aux = mm64( + [ + vceqq_u8(tv[0], sl), + vceqq_u8(tv[1], sl), + vceqq_u8(tv[2], sl), + vceqq_u8(tv[3], sl), + ], + pow, + ); + return (b, last_code); + } + if AUX == AUX_HAN { + // ponytail: scalar Han range test; vectorise like the CJK path below if kimi ever + // shows up on a throughput bench. + let lim = ntext.min(base + 64); + let mut leads = 0u64; + for p in base..lim { + if tags[p] != crate::CONT && crate::han::is_han_at(text, p) { + leads |= 1u64 << (p - base); + } + } + b.aux = leads | (leads << 1) | (leads << 2); + if cur_aux { + b.aux |= lead_run(b.cont, !0); + } + return (b, last_code); + } + // ── the CJK range test lives in the raw bytes, not the tags. It is a 3-byte predicate, so // it runs in vector space on `vext`-aligned b1/b2 and extracts ONE mask — testing the 9 // byte predicates as separate bitstreams costs 9 extractions for the same answer. A @@ -165,11 +218,11 @@ pub(crate) unsafe fn build64( leads &= (1u64 << lim) - 1; } // every CJK char is 3 bytes → fill by two shifts; a char cut by the block edge is - // picked up on the other side by `cur_cjk` (its continuation bytes lead that block). - b.cjk = leads | (leads << 1) | (leads << 2); + // picked up on the other side by `cur_aux` (its continuation bytes lead that block). + b.aux = leads | (leads << 1) | (leads << 2); } - if cur_cjk { - b.cjk |= lead_run(b.cont, !0); + if cur_aux { + b.aux |= lead_run(b.cont, !0); } (b, last_code) } diff --git a/tokenizers/bitsplit/src/simd_x86.rs b/tokenizers/bitsplit/src/simd/x86.rs similarity index 75% rename from tokenizers/bitsplit/src/simd_x86.rs rename to tokenizers/bitsplit/src/simd/x86.rs index 37bb29838..407f5e481 100644 --- a/tokenizers/bitsplit/src/simd_x86.rs +++ b/tokenizers/bitsplit/src/simd/x86.rs @@ -15,7 +15,7 @@ //! and run under Rosetta, which reports SSSE3, so this path is the one that executes. The scalar //! builder in `lib.rs` stays the reference it is compared against. -use crate::{Blk, lead_run}; +use crate::{AUX_CJK, AUX_HAN, AUX_NONE, AUX_SLASH, Blk, lead_run}; use core::arch::x86_64::*; /// `x <= k`, unsigned, in the absence of an unsigned byte compare. @@ -43,18 +43,35 @@ unsafe fn sel(mask: __m128i, a: __m128i, b: __m128i) -> __m128i { unsafe { _mm_or_si128(_mm_and_si128(mask, a), _mm_andnot_si128(mask, b)) } } -/// Build one **full** 64-byte block. `cur_code`/`cur_cjk` describe the byte before it. +/// 64 bytes of `text[base..]` compared against `b`, as a bitmap. +/// +/// # Safety +/// `base + 64 <= text.len()`; the caller has checked for SSSE3. +#[target_feature(enable = "ssse3")] +pub(crate) unsafe fn eq64(text: &[u8], base: usize, b: u8) -> u64 { + unsafe { + let n = _mm_set1_epi8(b as i8); + let mut m = 0u64; + for k in 0..4 { + let v = _mm_loadu_si128(text.as_ptr().add(base + k * 16).cast()); + m |= (_mm_movemask_epi8(_mm_cmpeq_epi8(v, n)) as u16 as u64) << (16 * k); + } + m + } +} + +/// Build one **full** 64-byte block. `cur_code`/`cur_aux` describe the byte before it. /// /// # Safety /// `base + 64 <= tags.len()` and `base + 64 <= text.len()`; the caller has checked for SSSE3. #[target_feature(enable = "ssse3")] -pub(crate) unsafe fn build64( +pub(crate) unsafe fn build64( text: &[u8], tags: &[u8], base: usize, lut: &[u8; 64], cur_code: u8, - cur_cjk: bool, + cur_aux: bool, ) -> (Blk, u8) { unsafe { let low_nibble = _mm_set1_epi8(0x0F); @@ -115,9 +132,38 @@ pub(crate) unsafe fn build64( p0: plane!(7), p1: plane!(6), p2: plane!(5), - cjk: 0, + p3: if P3 { plane!(4) } else { 0 }, + aux: 0, }; - if !CJK { + if AUX == AUX_NONE { + return (b, last_code); + } + + if AUX == AUX_SLASH { + // single ASCII byte — no fill needed, `/` is its own char + let sl = _mm_set1_epi8(b'/' as i8); + b.aux = gather([ + _mm_cmpeq_epi8(_mm_loadu_si128(text.as_ptr().add(base).cast()), sl), + _mm_cmpeq_epi8(_mm_loadu_si128(text.as_ptr().add(base + 16).cast()), sl), + _mm_cmpeq_epi8(_mm_loadu_si128(text.as_ptr().add(base + 32).cast()), sl), + _mm_cmpeq_epi8(_mm_loadu_si128(text.as_ptr().add(base + 48).cast()), sl), + ]); + return (b, last_code); + } + if AUX == AUX_HAN { + // ponytail: scalar Han range test; vectorise like the CJK path below if kimi ever + // shows up on a throughput bench. + let lim = text.len().min(base + 64); + let mut leads = 0u64; + for p in base..lim { + if tags[p] != crate::CONT && crate::han::is_han_at(text, p) { + leads |= 1u64 << (p - base); + } + } + b.aux = leads | (leads << 1) | (leads << 2); + if cur_aux { + b.aux |= lead_run(b.cont, !0); + } return (b, last_code); } @@ -134,8 +180,8 @@ pub(crate) unsafe fn build64( .iter() .any(|v| _mm_movemask_epi8(in_range(*v, 0xE3, 6)) != 0); if !any { - if cur_cjk { - b.cjk |= lead_run(b.cont, !0); + if cur_aux { + b.aux |= lead_run(b.cont, !0); } return (b, last_code); } @@ -182,10 +228,10 @@ pub(crate) unsafe fn build64( leads &= (1u64 << lim) - 1; } // every CJK char is 3 bytes, so two shifts fill it; one cut by the block edge is picked up - // on the other side by `cur_cjk`. - b.cjk = leads | (leads << 1) | (leads << 2); - if cur_cjk { - b.cjk |= lead_run(b.cont, !0); + // on the other side by `cur_aux`. + b.aux = leads | (leads << 1) | (leads << 2); + if cur_aux { + b.aux |= lead_run(b.cont, !0); } (b, last_code) } diff --git a/tokenizers/bitsplit/tests/parity.rs b/tokenizers/bitsplit/tests/parity.rs new file mode 100644 index 000000000..2590e5703 --- /dev/null +++ b/tokenizers/bitsplit/tests/parity.rs @@ -0,0 +1,265 @@ +//! Byte-exactness gate for the bitstream splitters. Oracle = oniguruma, composed exactly as HF +//! applies it (deepseek is a `Sequence` of three Isolated splits, not one regex). +//! +//! - The SWEEP is the point. A bitstream program is only interesting where a rule straddles the +//! 64-byte grid: `starts[bi-1] |= patch`, the `anl` retraction, every `pb`/`nb` edge peek. Slicing +//! the corpus at every char boundary walks every construct through every block phase, and the +//! oracle re-runs on the same slice so the expectation is never hand-written. +//! - `find_iter` is the Isolated split only because these regexes match the whole input with no +//! gaps; deepseek's three passes do leave gaps, hence `split_iso`. +#![cfg(not(target_arch = "wasm32"))] + +use bitsplit::Span; +use bitsplit::classify::classify; +use bitsplit::regexes::{ + CL100K, DEEPSEEK_BIG as DS_BIG, DEEPSEEK_CJK as DS_CJK, DEEPSEEK_NUM as DS_NUM, GPT2, KIMI_K2, + O200K, TEKKEN, +}; +use onig::Regex; + +const CORPUS: &str = "The quick brown fox. Don't 12345 numbers, \u{00BD}\u{00B2}\u{00BC} \u{2168}! \ + café × naïve — Привет, наука! Ελλάδα 中文分词。ひらがな カタカナ 한글 مرحبا العربية \ + नरेंद्र मोदी x_y a1b2c3 e-mail@host.com 😀👍 hello world\ttabs\nnewlines end "; + +const EDGE: &str = "IT'S O'Brien can't 'quoted' l'été rock'n'roll\r\n\ + 0 42 999 1000 1234567 v1.2.3 3.14159 1,000,000 \ + https://host/a/b?c=1&d=2 path/to//file /\r\n/ ///x \ + CamelCase XMLHttpRequest IJSSELMEER DžAMBO ŀl a\u{0301}b \ + 日本語1234テスト ½3¼ \u{2168}42\u{2169} #tag @user $9.99 100% \ + end\n\n\nlines\r\n\r\n \t trailing "; + +/// A bitstream splitter, normalised to one shape (`starts` + `flag` scratch; deepseek ignores `flag`). +type Split = fn(&[u8], &[u8], &mut [u64], &mut [u64], &mut [Span]) -> usize; + +fn bs_deepseek(t: &[u8], g: &[u8], s: &mut [u64], _f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_deepseek(t, g, s, o) +} +fn bs_byte_level(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_byte_level(t, g, s, f, o) +} +fn bs_cl100k(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_cl100k(t, g, s, f, o) +} +fn bs_qwen(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_qwen(t, g, s, f, o) +} +fn bs_o200k(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_o200k(t, g, s, f, o) +} +fn bs_tekken(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_tekken(t, g, s, f, o) +} +fn bs_kimi(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_kimi(t, g, s, f, o) +} + +fn spans(f: Split, s: &str) -> Vec { + if s.is_empty() { + return Vec::new(); + } + let b = s.as_bytes(); + let mut tags = vec![0u8; b.len()]; + classify(b, &mut tags); + let nblk = b.len().div_ceil(64); + let (mut starts, mut flag) = (vec![0u64; nblk], vec![0u64; nblk]); + let mut out = vec![Span::default(); b.len() + 1]; + let k = f(b, &tags, &mut starts, &mut flag, &mut out); + out.truncate(k); + out +} + +fn onig_spans(re: &Regex, s: &str) -> Vec { + re.find_iter(s) + .map(|(a, b)| Span::new(a as u32, b as u32)) + .collect() +} + +/// One Isolated split of `text[s..e]`: gaps + matches, absolute offsets. +fn split_iso(text: &str, s: usize, e: usize, re: &Regex, out: &mut Vec<(usize, usize)>) { + let sub = &text[s..e]; + let mut prev = 0usize; + for (ms, me) in re.find_iter(sub) { + if ms > prev { + out.push((s + prev, s + ms)); + } + out.push((s + ms, s + me)); + prev = me; + } + if prev < sub.len() { + out.push((s + prev, e)); + } +} + +fn deepseek_ref(text: &str) -> Vec { + let (rn, rc, rb) = ( + Regex::new(DS_NUM).unwrap(), + Regex::new(DS_CJK).unwrap(), + Regex::new(DS_BIG).unwrap(), + ); + let mut a = Vec::new(); + split_iso(text, 0, text.len(), &rn, &mut a); + let mut b = Vec::new(); + for (s, e) in a { + split_iso(text, s, e, &rc, &mut b); + } + let mut c = Vec::new(); + for (s, e) in b { + split_iso(text, s, e, &rb, &mut c); + } + c.into_iter() + .map(|(s, e)| Span::new(s as u32, e as u32)) + .collect() +} + +/// The oracle for a grammar, over an arbitrary slice. +enum Oracle { + Whole(Regex), + DeepSeek, +} + +impl Oracle { + fn spans(&self, s: &str) -> Vec { + match self { + Oracle::Whole(re) => onig_spans(re, s), + Oracle::DeepSeek => deepseek_ref(s), + } + } +} + +fn long_corpus() -> String { + let mut s = String::new(); + while s.len() < 4096 { + s.push_str(CORPUS); + s.push_str(EDGE); + } + s +} + +/// Both corpora whole, then the block-phase sweep: every char-boundary prefix and suffix of a +/// >4 KB text, so each rule crosses a block edge in every alignment. +fn check(name: &str, f: Split, oracle: &Oracle) { + for text in [CORPUS, EDGE] { + assert_eq!(spans(f, text), oracle.spans(text), "{name}: whole {text:?}"); + } + + let long = long_corpus(); + let mut checked = 0usize; + + // suffixes: shifts the whole text across the grid + for off in 0..512.min(long.len()) { + if !long.is_char_boundary(off) { + continue; + } + let sub = &long[off..]; + assert_eq!(spans(f, sub), oracle.spans(sub), "{name}: suffix off={off}"); + checked += 1; + } + // prefixes: exercises the last-block / EOF rules (`\s+(?!\S)` vs plain `\s+`) in every phase + for end in (long.len().saturating_sub(512))..=long.len() { + if !long.is_char_boundary(end) { + continue; + } + let sub = &long[..end]; + assert_eq!(spans(f, sub), oracle.spans(sub), "{name}: prefix end={end}"); + checked += 1; + } + assert!(checked > 500, "{name}: sweep too small ({checked})"); +} + +/// Deterministic pseudo-random text from a weighted alphabet — always valid UTF-8 by construction. +/// Aimed at the carry logic: long digit runs, whitespace/newline runs and apostrophes next to +/// letters are what the `\p{N}{1,3}` grouping, the `\s*[\r\n]+` fill and the contraction escape +/// disagree about. +fn fuzz_texts(n: usize) -> Vec { + const ALPHA: &[&str] = &[ + "a", "b", "z", "A", "Q", "é", "ß", "IJ", "0", "1", "9", " ", " ", "\n", "\r\n", "\t", "'", + "'s", "'ll", ".", "!", "/", "#", "_", "中", "文", "ひ", "カ", "한", "م", "😀", "\u{0301}", + "\u{200D}", "½", "\u{2168}", + ]; + let mut st = 0x243F_6A88_85A3_08D3u64; + let mut next = move || { + st ^= st << 13; + st ^= st >> 7; + st ^= st << 17; + st + }; + (0..n) + .map(|_| { + let len = 1 + (next() % 300) as usize; + (0..len) + .map(|_| ALPHA[(next() % ALPHA.len() as u64) as usize]) + .collect() + }) + .collect() +} + +fn check_fuzz(name: &str, f: Split, oracle: &Oracle) { + for (i, t) in fuzz_texts(4000).iter().enumerate() { + assert_eq!(spans(f, t), oracle.spans(t), "{name}: fuzz #{i} {t:?}"); + } +} + +#[test] +fn byte_level_parity() { + let o = Oracle::Whole(Regex::new(GPT2).unwrap()); + check("byte_level", bs_byte_level, &o); + check_fuzz("byte_level", bs_byte_level, &o); +} + +#[test] +fn cl100k_parity() { + let o = Oracle::Whole(Regex::new(CL100K).unwrap()); + check("cl100k", bs_cl100k, &o); + check_fuzz("cl100k", bs_cl100k, &o); +} + +/// Qwen2 / Qwen3: cl100k with a bare `\p{N}`. +#[test] +fn qwen_parity() { + let qwen = CL100K.replace(r"\p{N}{1,3}", r"\p{N}"); + let o = Oracle::Whole(Regex::new(&qwen).unwrap()); + check("qwen", bs_qwen, &o); + check_fuzz("qwen", bs_qwen, &o); +} + +#[test] +fn deepseek_parity() { + check("deepseek", bs_deepseek, &Oracle::DeepSeek); + check_fuzz("deepseek", bs_deepseek, &Oracle::DeepSeek); +} + +/// o200k_base / GPT-4o — and byte-for-byte the regex Llama-4, gpt-oss and MiniMax-M2 ship, so this +/// one test covers four families. +#[test] +fn o200k_parity() { + let o = Oracle::Whole(Regex::new(O200K).unwrap()); + check("o200k", bs_o200k, &o); + check_fuzz("o200k", bs_o200k, &o); +} + +#[test] +fn tekken_parity() { + let o = Oracle::Whole(Regex::new(TEKKEN).unwrap()); + check("tekken", bs_tekken, &o); + check_fuzz("tekken", bs_tekken, &o); +} + +#[test] +fn kimi_parity() { + let o = Oracle::Whole(Regex::new(KIMI_K2).unwrap()); + check("kimi", bs_kimi, &o); + check_fuzz("kimi", bs_kimi, &o); +} + +/// Negative control: the gate above only means something if it can fail. gpt2 and cl100k disagree +/// on plenty (`\p{N}{1,3}`, the `[\r\n]*` tail, the non-space `?` prefix), so crossing them must +/// blow up — if this ever passes, `check` has stopped comparing anything. +#[test] +#[should_panic(expected = "negative-control")] +fn harness_discriminates() { + check( + "negative-control", + bs_byte_level, + &Oracle::Whole(Regex::new(CL100K).unwrap()), + ); +} diff --git a/tokenizers/tk-encode/Cargo.toml b/tokenizers/tk-encode/Cargo.toml index 7535aa994..bdea45df3 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -27,7 +27,6 @@ name = "tk_encode" path = "src/lib.rs" [dependencies] -atomsplit = { path = "../atomsplit" } bitsplit = { path = "../bitsplit" } rand = "0.9" regex = "1.10" diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 17ef63c51..0eb89a8fe 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -362,9 +362,8 @@ impl pipeline::Model for PipelineBPE { output: &mut Vec, ) -> Result<()> { let BpeScratch { - merge_queue, - skip, - word, + symbols, + queue, word_cache, } = scratch; @@ -380,6 +379,15 @@ impl pipeline::Model for PipelineBPE { continue; } + // Same order as `tokenize_pipeline`, and it has to stay that way: the fold answers a + // word that is itself a foldable vocabulary entry in one probe, and those words never + // reach the cache. Probing the cache first would populate it with words the fold + // already serves for free, and the two paths would disagree about what it holds. + if let Some(id) = self.fold_id(sequence) { + output.push(PipelineToken { id }); + continue; + } + let mut placement = None; if let Some(cache) = word_cache.as_mut() { match cache.lookup(sequence.as_bytes()) { @@ -390,17 +398,13 @@ impl pipeline::Model for PipelineBPE { Lookup::Miss(at) => placement = Some(at), } } + let start = output.len(); - if self.ignore_merges - && let Some(id) = self.vocab.get_bytes(sequence.as_bytes()) - { - output.push(PipelineToken { id }); - } else { - self.merge_word(sequence, merge_queue, skip, word); - output.extend(word.get_chars_iter().map(|id| PipelineToken { id })); - } - // The ids come back out of `output` because that is the only place both branches - // above leave them: `ignore_merges` never touches `word`. + self.merge_word(sequence, symbols, queue); + // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids + output.extend(symbols.iter().map(|&symbol| PipelineToken { + id: self.tables.unmap.at(symbol as usize), + })); if let Some(cache) = word_cache.as_mut() && let Some(at) = placement { @@ -410,43 +414,6 @@ impl pipeline::Model for PipelineBPE { Ok(()) } - /// Every pre-token of a chunk in one call. - /// - /// Same work per word as [`Self::tokenize_pipeline`]; what changes is what is *not* repeated. - /// The scratch is destructured once instead of once per word, the output is grown once for the - /// whole batch instead of being capacity-checked on every push, and the virtual call, the - /// slice and the `Result` happen once per chunk rather than once per pre-token. - fn tokenize_spans( - &self, - chunk: &str, - spans: &[Span], - scratch: &mut Self::Scratch, - output: &mut Vec, - ) -> Result<()> { - let BpeScratch { symbols, queue } = scratch; - // One reservation for the batch. Most pre-tokens are a single token, so the span count is - // a close lower bound on what the batch emits; anything past it grows as usual. - output.reserve(spans.len()); - - for span in spans { - // SAFETY: the pre-tokenizer cuts on char boundaries, so a span is always a valid slice - // of this chunk. Bounds- and UTF-8-checking it again per word measured worth removing. - let sequence = unsafe { chunk.get_unchecked(span.range()) }; - if sequence.is_empty() { - continue; - } - if let Some(id) = self.fold_id(sequence) { - output.push(PipelineToken { id }); - continue; - } - self.merge_word(sequence, symbols, queue); - output.extend(symbols.iter().map(|&symbol| PipelineToken { - id: self.tables.unmap.at(symbol as usize), - })); - } - Ok(()) - } - fn init_scratch(&self) -> Self::Scratch { Self::Scratch { symbols: Vec::with_capacity(64), diff --git a/tokenizers/tk-encode/src/normalizers/replace.rs b/tokenizers/tk-encode/src/normalizers/replace.rs index cc4545b47..635ea5378 100644 --- a/tokenizers/tk-encode/src/normalizers/replace.rs +++ b/tokenizers/tk-encode/src/normalizers/replace.rs @@ -5,7 +5,7 @@ use crate::tokenizer::Decoder; use crate::tokenizer::pattern::Pattern; use crate::tokenizer::{NormalizedString, Normalizer, Result}; use crate::utils::SysRegex; -use atomsplit::literal::Literal; +use bitsplit::literal::Literal; use serde::{Deserialize, Serialize}; /// Represents the different patterns that `Replace` can use diff --git a/tokenizers/tk-encode/src/pre_tokenizers/bert.rs b/tokenizers/tk-encode/src/pre_tokenizers/bert.rs index d2f7e6e8d..fa65cf7cc 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/bert.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/bert.rs @@ -21,8 +21,8 @@ impl pipeline::PreTokenizer for BertPreTokenizer { // Bert pre-tokenization = drop whitespace runs, isolate each punctuation char, keep every other // run. One `atomsplit` SIMD classify (bytes → atom tags) + the class-runs FSM, byte-exact with // the legacy `char::is_whitespace` / `is_punc` split above (see the tests). - use atomsplit::classify::{classify, mask}; - use atomsplit::fsm::class_runs_into; + use bitsplit::classify::{classify, mask}; + use bitsplit::classes::class_runs_into; let bytes = text.as_bytes(); let mut tags = vec![0u8; bytes.len()]; classify(bytes, &mut tags); diff --git a/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs b/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs index b716136f6..45dedea07 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs @@ -1,5 +1,5 @@ use crate::utils::byte_level::{BYTES_CHAR_LOOKUP, CHAR_BYTES_LOOKUP, byte_level_transform}; -use crate::utils::{GptFsm, GptFsmPattern}; +use crate::utils::{Grammar, GrammarPattern}; use serde::{Deserialize, Serialize}; use crate::tokenizer::{ @@ -85,7 +85,7 @@ impl PreTokenizer for ByteLevel { if self.use_regex { // GPT-2 byte-level split via the native atomsplit FSM (byte-exact, no regex backend). normalized.split( - GptFsmPattern(GptFsm::Gpt2), + GrammarPattern(Grammar::Gpt2), SplitDelimiterBehavior::Isolated, ) } else { diff --git a/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs b/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs index e30f206e6..5f1e9c758 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs @@ -32,7 +32,7 @@ impl pipeline::PreTokenizer for CharDelimiterSplit { // keeps the runs between, no empty spans. Byte-exact with the char-predicate split. let bytes = text.as_bytes(); let mut spans = vec![pipeline::Span::default(); bytes.len() + 1]; - let n = atomsplit::fsm::CharDelimiterSplit(self.delimiter).pre_tokenize( + let n = bitsplit::classes::CharDelimiterSplit(self.delimiter).pre_tokenize( bytes, &mut [], &mut spans, diff --git a/tokenizers/tk-encode/src/pre_tokenizers/digits.rs b/tokenizers/tk-encode/src/pre_tokenizers/digits.rs index 25cf9f193..38a15cd9f 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/digits.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/digits.rs @@ -43,8 +43,8 @@ impl pipeline::PreTokenizer for Digits { fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> { // isolate each numeric char (`individual_digits`) or keep numeric runs — atomsplit classify + // class-runs FSM. atom `NUMERIC` == `char::is_numeric`, so byte-exact with the scalar path. - use atomsplit::classify::{classify, mask}; - use atomsplit::fsm::class_runs_into; + use bitsplit::classify::{classify, mask}; + use bitsplit::classes::class_runs_into; let bytes = text.as_bytes(); let mut tags = vec![0u8; bytes.len()]; classify(bytes, &mut tags); diff --git a/tokenizers/tk-encode/src/pre_tokenizers/punctuation.rs b/tokenizers/tk-encode/src/pre_tokenizers/punctuation.rs index 52f4e5e5d..b1a466e6c 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/punctuation.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/punctuation.rs @@ -44,8 +44,8 @@ impl pipeline::PreTokenizer for Punctuation { // atom `PUNCT` == `is_punc` (ASCII-punct ∪ \p{P}), so Isolated/Removed map to the class-runs FSM // byte-exactly. The merge/contiguous behaviors aren't a class-runs shape → keep the scalar split. if matches!(self.behavior, Isolated | Removed) { - use atomsplit::classify::mask; - use atomsplit::fsm::class_runs_into; + use bitsplit::classify::mask; + use bitsplit::classes::class_runs_into; pipeline::classify_into_spans( text.as_bytes(), |b, t, s| { diff --git a/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs b/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs index b01d4686d..72bee136e 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs @@ -107,7 +107,11 @@ impl pipeline::PreTokenizer for PipelineSequence { // deepseek's 3-Split composition → one native FSM pass (also lets the Sequence handle the // trailing byte-map ByteLevel, which the generic child loop can't range-split). if self.is_deepseek() { - pipeline::classify_into_spans(text.as_bytes(), atomsplit::fsm::fsm_deepseek, out); + pipeline::classify_into_spans_bits( + text.as_bytes(), + |t, tags, starts, _flag, out| bitsplit::bitsplit_deepseek(t, tags, starts, out), + out, + ); return Ok(()); } diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index 04a6b4c30..bddbc0484 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -1,6 +1,6 @@ use crate::pipeline; -use crate::utils::{GptFsm, GptFsmPattern, SysRegex, gpt_fsm}; -use atomsplit::literal::Literal; +use crate::utils::{Grammar, GrammarPattern, SysRegex, recognize}; +use bitsplit::literal::Literal; use serde::{Deserialize, Deserializer, Serialize}; use crate::tokenizer::{ @@ -53,7 +53,7 @@ pub struct Split { /// pipeline path when `behavior == Isolated && !invert` (how these regexes always ship). Byte-exact /// with `regex`; `None` falls back to `regex`. #[serde(skip)] - fsm: Option, + fsm: Option, } impl<'de> Deserialize<'de> for Split { @@ -103,7 +103,7 @@ impl Split { let pattern: SplitPattern = pattern.into(); let fsm = match &pattern { SplitPattern::String(_) => None, - SplitPattern::Regex(r) => gpt_fsm(r), + SplitPattern::Regex(r) => recognize(r), }; let search = match &pattern { SplitPattern::String(s) => Search::Literal(Literal::new(s.as_bytes())?), @@ -173,53 +173,23 @@ impl PreTokenizer for Split { .into() })?; pretokenized.split(|_, normalized| { - normalized.split(GptFsmPattern(fsm), SplitDelimiterBehavior::Isolated) + normalized.split(GrammarPattern(fsm), SplitDelimiterBehavior::Isolated) }) } } impl pipeline::PreTokenizer for Split { fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> { - // A recognized GPT regex (gpt2 / cl100k-Llama-3) in its only real usage — `Isolated`, not - // inverted — routes straight to the native atomsplit FSM. These regexes cover the whole input, - // so `Isolated` == the match list, and the FSM is byte-exact with `regex` (see the tests). - if let Some(fsm) = self + // A recognized GPT regex in its only real usage — `Isolated`, not inverted — routes + // straight to the native bitsplit grammar. These regexes cover the whole input, so + // `Isolated` == the match list, and the grammar is byte-exact with `regex` (see the tests). + if let Some(grammar) = self .fsm .filter(|_| !self.invert && self.behavior == SplitDelimiterBehavior::Isolated) { - // gpt2 and cl100k-with-the-standard-digit-cap have bitstream splitters, which decide - // 64 bytes per register op instead of one token per unpredictable branch -- but only - // where the bitstream build is SIMD. Its portable builder is slower than the FSM, so - // everything else keeps the FSM. - match fsm { - GptFsm::Gpt2 if bitsplit::fast_builder() => { - pipeline::classify_into_spans_bits( - text.as_bytes(), - bitsplit::bitsplit_byte_level, - out, - ); - return Ok(()); - } - GptFsm::Cl100k { digit_cap: 3 } if bitsplit::fast_builder() => { - pipeline::classify_into_spans_bits( - text.as_bytes(), - bitsplit::bitsplit_cl100k, - out, - ); - return Ok(()); - } - _ => {} - } - pipeline::classify_into_spans( + pipeline::classify_into_spans_bits( text.as_bytes(), - |bytes, tags, spans| match fsm { - GptFsm::Cl100k { digit_cap } => { - atomsplit::fsm::fsm_cl100k_cap(bytes, tags, spans, digit_cap) - } - GptFsm::Gpt2 => atomsplit::fsm::fsm_byte_level(bytes, tags, spans), - GptFsm::O200k => atomsplit::fsm::fsm_o200k(bytes, tags, spans), - GptFsm::Tekken => atomsplit::fsm::fsm_tekken(bytes, tags, spans), - }, + |t, tags, starts, flag, out| grammar.split(t, tags, starts, flag, out), out, ); return Ok(()); @@ -442,7 +412,7 @@ mod tests { fn pipeline_gpt2_uses_fsm_and_matches_legacy() { // The gpt2 pattern is recognized -> the pipeline path routes to the native // atomsplit FSM; its output must equal the legacy fancy-regex path. - let gpt2 = atomsplit::regexes::GPT2; + let gpt2 = bitsplit::regexes::GPT2; let corpus = "The quick brown fox 123!!! double spaces\tand tabs. don't Naïve café. "; let pretok = Split::new(SplitPattern::Regex(gpt2.into()), Isolated, false).unwrap(); assert!(pretok.fsm.is_some(), "gpt2 pattern should be recognized"); @@ -466,7 +436,7 @@ mod tests { fn pipeline_cl100k_llama3_uses_fsm_and_matches_legacy() { // Llama-3's EXACT pre_tokenizer regex (from data/llama-3-tokenizer.json) → recognized → routes // to fsm_cl100k. Output must equal the legacy SysRegex Isolated split, byte-for-byte. - let cl100k = atomsplit::regexes::CL100K; + let cl100k = bitsplit::regexes::CL100K; let corpus = "The quick brown fox 123!!! double spaces\tand tabs. don't Naïve café.\n\n世界 안녕 "; let pretok = Split::new(SplitPattern::Regex(cl100k.into()), Isolated, false).unwrap(); @@ -494,7 +464,7 @@ mod tests { // GPT-4o's EXACT pre_tokenization regex → recognized → routes to fsm_o200k. Output must equal the // legacy SysRegex Isolated split, byte-for-byte. Corpus stresses the case-aware letter split: // camelCase, ALLCAPS→word, McDonald's-style, contractions, accented Ll (é/ß), CJK (caseless). - let o200k = atomsplit::regexes::O200K; + let o200k = bitsplit::regexes::O200K; let corpus = "McDonald's iPhone SQLite HELLOWorld camelCase don't I'll We've 3.14 café Straße 世界 안녕\n\n Mixed CASE end."; let pretok = Split::new(SplitPattern::Regex(o200k.into()), Isolated, false).unwrap(); assert!( @@ -521,12 +491,12 @@ mod tests { // Mistral's tekken regex (mistral-small-4) → recognized → routes to fsm_tekken. Same corpus // shape as o200k, whose grammar it shares: the differences it must get right are apostrophes // (no contraction suffix, so `'s` starts a new token) and one token per digit. - let tekken = atomsplit::regexes::TEKKEN; + let tekken = bitsplit::regexes::TEKKEN; let corpus = "McDonald's iPhone SQLite HELLOWorld camelCase don't I'll We've 3.14159 café Straße 世界 안녕\n\n path/to/file Mixed CASE end."; let pretok = Split::new(SplitPattern::Regex(tekken.into()), Isolated, false).unwrap(); assert_eq!( pretok.fsm, - Some(crate::utils::GptFsm::Tekken), + Some(crate::utils::Grammar::Tekken), "tekken / mistral pattern should route to the native FSM" ); @@ -554,7 +524,7 @@ mod tests { let pretok = Split::new(SplitPattern::Regex(qwen2.into()), Isolated, false).unwrap(); assert_eq!( pretok.fsm, - Some(crate::utils::GptFsm::Cl100k { digit_cap: 1 }), + Some(crate::utils::Grammar::Cl100k { digit_cap: 1 }), "Qwen2 pattern should route to the cl100k FSM with digit cap 1" ); diff --git a/tokenizers/tk-encode/src/pre_tokenizers/whitespace.rs b/tokenizers/tk-encode/src/pre_tokenizers/whitespace.rs index 13f19dda2..20d789eb4 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/whitespace.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/whitespace.rs @@ -45,10 +45,10 @@ impl pipeline::PreTokenizer for WhitespaceSplit { fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> { // drop whitespace runs, keep everything else as runs — atomsplit SIMD classify + class-runs FSM. // atom `WS` == `char::is_whitespace`, so byte-exact with the scalar path. - use atomsplit::classify::mask; + use bitsplit::classify::mask; pipeline::classify_into_spans( text.as_bytes(), - atomsplit::fsm::class_runs_into::<{ mask::WS }, 0, 0>, + bitsplit::classes::class_runs_into::<{ mask::WS }, 0, 0>, out, ); Ok(()) @@ -74,10 +74,10 @@ impl pipeline::PreTokenizer for Whitespace { fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> { // `\w+|[^\w\s]+`: drop whitespace, cut at the word↔symbol boundary, each run one token — // atomsplit classify + class-runs FSM (`WORD` = `\w`; keep-A = word, keep-B = symbol). - use atomsplit::classify::mask; + use bitsplit::classify::mask; pipeline::classify_into_spans( text.as_bytes(), - atomsplit::fsm::class_runs_into::<{ mask::WS }, 0, { mask::WORD }>, + bitsplit::classes::class_runs_into::<{ mask::WS }, 0, { mask::WORD }>, out, ); Ok(()) diff --git a/tokenizers/tk-encode/src/tokenizer/normalizer.rs b/tokenizers/tk-encode/src/tokenizer/normalizer.rs index 8f1899bdd..937509c80 100644 --- a/tokenizers/tk-encode/src/tokenizer/normalizer.rs +++ b/tokenizers/tk-encode/src/tokenizer/normalizer.rs @@ -1022,7 +1022,7 @@ impl From<&str> for NormalizedString { #[cfg(test)] mod tests { use super::*; - use atomsplit::literal::Literal; + use bitsplit::literal::Literal; use regex::Regex; use unicode_categories::UnicodeCategories; diff --git a/tokenizers/tk-encode/src/tokenizer/pattern.rs b/tokenizers/tk-encode/src/tokenizer/pattern.rs index 5a147f5f4..879e35135 100644 --- a/tokenizers/tk-encode/src/tokenizer/pattern.rs +++ b/tokenizers/tk-encode/src/tokenizer/pattern.rs @@ -1,6 +1,6 @@ use crate::utils::SysRegex; use crate::{Offsets, Result}; -use atomsplit::literal::Literal; +use bitsplit::literal::Literal; use regex::Regex; /// Pattern used to split a NormalizedString diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index ef859c16e..55385b574 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -6,7 +6,7 @@ use std::sync::{Mutex, PoisonError}; use std::vec::IntoIter; use std::{borrow::Cow, convert::TryFrom}; -use atomsplit::classify::classify; +use bitsplit::classify::classify; use crate::models::bpe::{BpeScratch, PipelineBPE}; use crate::models::unigram::{Unigram, UnigramScratch}; @@ -39,7 +39,7 @@ use crate::{ use super::{Result, SplitDelimiterBehavior}; -pub use atomsplit::fsm::Span; +pub use bitsplit::Span; /// We use a thread local scratch for the tags (per byte class) and for the split spans. pub(crate) fn classify_into_spans( diff --git a/tokenizers/tk-encode/src/utils/byte_level.rs b/tokenizers/tk-encode/src/utils/byte_level.rs index 95a26807a..70d79e938 100644 --- a/tokenizers/tk-encode/src/utils/byte_level.rs +++ b/tokenizers/tk-encode/src/utils/byte_level.rs @@ -4,7 +4,7 @@ use std::sync::LazyLock; // The GPT-2 pre-tokenize regex is the canonical spec in atomsplit (single source of truth); re-export // under the historical name so call sites are unchanged. -pub(crate) use atomsplit::regexes::GPT2 as GPT2_REGEX_STR; +pub(crate) use bitsplit::regexes::GPT2 as GPT2_REGEX_STR; /// Maps each byte to its GPT-2 byte-level unicode character, indexed by the byte value. /// diff --git a/tokenizers/tk-encode/src/utils/mod.rs b/tokenizers/tk-encode/src/utils/mod.rs index 55643b57d..c446f576d 100644 --- a/tokenizers/tk-encode/src/utils/mod.rs +++ b/tokenizers/tk-encode/src/utils/mod.rs @@ -6,7 +6,7 @@ pub(crate) mod word_cache; // Optional system-regex backend, needed only for a *regex* pattern that atomsplit does not cover. // With `fancy-regex` off a stub compiles and those patterns error at load. Everything else works // regardless: the atomsplit-native pre-tokenizers, and any `Split` or `Replace` whose pattern is a -// plain string (searched for directly, see `atomsplit::literal`). +// plain string (searched for directly, see `bitsplit::literal`). #[cfg(feature = "fancy-regex")] mod fancy; #[cfg(feature = "fancy-regex")] @@ -18,7 +18,7 @@ pub use no_regex::SysRegex; // Recognize known GPT pre-tokenization regexes and route them to atomsplit's native (unrolled) FSM. mod unrolled_regex; -pub use unrolled_regex::{GptFsm, GptFsmPattern, gpt_fsm, is_deepseek}; +pub use unrolled_regex::{Grammar, GrammarPattern, recognize, is_deepseek}; pub mod byte_level; pub mod iter; diff --git a/tokenizers/tk-encode/src/utils/no_regex.rs b/tokenizers/tk-encode/src/utils/no_regex.rs index 8645b69c3..2a135e412 100644 --- a/tokenizers/tk-encode/src/utils/no_regex.rs +++ b/tokenizers/tk-encode/src/utils/no_regex.rs @@ -3,7 +3,7 @@ //! The type stays present so `Split` / `Replace` still compile, but construction always fails. Only a //! *regex* pattern ever asks for it: the atomsplit-native pre-tokenizers (GPT-2, cl100k, deepseek, the //! class family, char-delimiter) need no backend, and a plain string pattern is searched for directly -//! (`atomsplit::literal`). A regex atomsplit does not cover errors at load time with a clear message. +//! (`bitsplit::literal`). A regex atomsplit does not cover errors at load time with a clear message. //! Enable `fancy-regex` to get a real backend. use std::error::Error; diff --git a/tokenizers/tk-encode/src/utils/unrolled_regex.rs b/tokenizers/tk-encode/src/utils/unrolled_regex.rs index 78025296b..776387f08 100644 --- a/tokenizers/tk-encode/src/utils/unrolled_regex.rs +++ b/tokenizers/tk-encode/src/utils/unrolled_regex.rs @@ -1,35 +1,60 @@ -//! Recognize a known GPT pre-tokenization regex and route it to the byte-exact native -//! `atomsplit` FSM, so those pre-tokenizers need no system-regex backend. An unrecognized -//! pattern returns `None` and falls back to `SysRegex` (the optional fancy-regex backend). +//! Recognize a known GPT pre-tokenization regex and route it to the byte-exact native `bitsplit` +//! grammar, so those pre-tokenizers need no system-regex backend. An unrecognized pattern returns +//! `None` and falls back to `SysRegex` (the optional fancy-regex backend). -// Canonical GPT pre-tokenization regexes (the look-ahead originals), used as recognition keys — the -// single source of truth lives in `atomsplit::regexes`. `gpt_fsm` maps each to the `atomsplit` FSM that -// reproduces its `Isolated` split byte-for-byte. -use atomsplit::regexes::{GPT2, O200K, TEKKEN}; +// The canonical regexes are the recognition keys; the single source of truth is `bitsplit::regexes`. +use bitsplit::Span; +use bitsplit::regexes::{GPT2, KIMI_K2, O200K, TEKKEN}; // cl100k is recognized structurally (see `cl100k_digit_cap`), so the exact pattern is only a test key. #[cfg(test)] -use atomsplit::regexes::CL100K; +use bitsplit::regexes::CL100K; -/// A recognized GPT pre-tokenization regex that maps to a native `atomsplit` FSM (byte-exact). +/// A recognized GPT pre-tokenization regex and the `bitsplit` grammar that reproduces its +/// `Isolated` split byte-for-byte. One variant per distinct regex; models sharing a regex share a +/// variant (o200k covers Llama-4, gpt-oss and MiniMax-M2). #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum GptFsm { - /// GPT-2 / ByteLevel regex → `atomsplit::fsm::fsm_byte_level`. +pub enum Grammar { + /// GPT-2 / ByteLevel. Gpt2, - /// cl100k-family regex → `atomsplit::fsm::fsm_cl100k_cap`. `digit_cap` is rule 3's `\p{N}{1,cap}` - /// bound: 3 = cl100k / Llama-3, 1 = Qwen2 (`\p{N}`), `usize::MAX` = an unbounded `\p{N}+`. + /// cl100k family. `digit_cap` is rule 3's `\p{N}{1,cap}` bound: 3 = cl100k / Llama-3 / GLM-4.6, + /// 1 = Qwen (`\p{N}`), `usize::MAX` = an unbounded `\p{N}+`. Cl100k { digit_cap: usize }, - /// o200k / GPT-4o regex → `atomsplit::fsm::fsm_o200k`. + /// o200k / GPT-4o — and byte-for-byte the regex Llama-4, gpt-oss and MiniMax-M2 ship. O200k, - /// Mistral tekken regex → `atomsplit::fsm::fsm_tekken` (o200k's grammar with no contraction - /// suffix and one token per digit). + /// Mistral tekken: o200k with no contraction suffix and one token per digit. Tekken, + /// kimi-k2 / k3: o200k plus a leading `[\p{Han}]+` arm and a plain `[\r\n]*` rule-4 tail. + Kimi, } -/// The cl100k-family template is fixed except rule 3's digit rule. If `pattern` is that template, return -/// the `\p{N}{1,cap}` bound (`\p{N}{1,3}`→3, `\p{N}{1,2}`→2, `\p{N}`→1, `\p{N}+`→`MAX`); else `None`. -/// This is what makes Qwen2 (cl100k with `\p{N}`) unroll without a per-tokenizer exact-string entry. +impl Grammar { + /// Write the token spans into `out`, returning the count. `starts`/`flag` are `u64` scratch + /// bitmaps of length >= `text.len().div_ceil(64)`. + pub fn split( + self, + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], + ) -> usize { + match self { + Grammar::Gpt2 => bitsplit::bitsplit_byte_level(text, tags, starts, flag, out), + Grammar::Cl100k { digit_cap: 1 } => { + bitsplit::bitsplit_qwen(text, tags, starts, flag, out) + } + Grammar::Cl100k { .. } => bitsplit::bitsplit_cl100k(text, tags, starts, flag, out), + Grammar::O200k => bitsplit::bitsplit_o200k(text, tags, starts, flag, out), + Grammar::Tekken => bitsplit::bitsplit_tekken(text, tags, starts, flag, out), + Grammar::Kimi => bitsplit::bitsplit_kimi(text, tags, starts, flag, out), + } + } +} + +/// The cl100k-family template is fixed except rule 3's digit rule. If `pattern` is that template, +/// return the `\p{N}{1,cap}` bound; else `None`. This is what makes Qwen (cl100k with `\p{N}`) and +/// GLM-4.6 (cl100k verbatim) route without a per-model entry. fn cl100k_digit_cap(pattern: &str) -> Option { - // cl100k rules 1-2 (contraction + word) … … rules 4-7 (other + whitespace). const PRE: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|"; const SUF: &str = r"| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"; match pattern.strip_prefix(PRE)?.strip_suffix(SUF)? { @@ -41,29 +66,24 @@ fn cl100k_digit_cap(pattern: &str) -> Option { } } -/// If `pattern` is a recognized GPT pre-tokenization regex, name the native FSM that reproduces its -/// `Isolated` split byte-for-byte. GPT-2, o200k and tekken are matched exactly; the cl100k family is -/// matched structurally ([`cl100k_digit_cap`]) so digit-cap variants (Qwen2 …) unroll too. An -/// unrecognized pattern → `None` (the SysRegex / fancy-regex path handles it). -pub fn gpt_fsm(pattern: &str) -> Option { - if pattern == GPT2 { - Some(GptFsm::Gpt2) - } else if pattern == O200K { - Some(GptFsm::O200k) - } else if pattern == TEKKEN { - Some(GptFsm::Tekken) - } else { - cl100k_digit_cap(pattern).map(|digit_cap| GptFsm::Cl100k { digit_cap }) +/// If `pattern` is a recognized GPT pre-tokenization regex, name the grammar that reproduces its +/// `Isolated` split byte-for-byte. An unrecognized pattern -> `None` (the SysRegex path handles it). +pub fn recognize(pattern: &str) -> Option { + match pattern { + GPT2 => Some(Grammar::Gpt2), + O200K => Some(Grammar::O200k), + TEKKEN => Some(Grammar::Tekken), + KIMI_K2 => Some(Grammar::Kimi), + _ => cl100k_digit_cap(pattern).map(|digit_cap| Grammar::Cl100k { digit_cap }), } } -/// `Pattern` that runs the native atomsplit FSM for a [`GptFsm`] — the legacy (`NormalizedString`) -/// split path's equivalent of the pipeline's native routing, so GPT pre-tokenizers need no -/// system-regex backend. `Isolated` behaviour keeps every span, and the FSM spans cover the input -/// contiguously, so this is byte-for-byte identical to the original GPT regex `Isolated` split. -pub struct GptFsmPattern(pub GptFsm); +/// `Pattern` that runs the native grammar on the legacy (`NormalizedString`) split path, so GPT +/// pre-tokenizers need no system-regex backend. `Isolated` keeps every span and these grammars +/// cover the input contiguously, so this is byte-for-byte the original regex `Isolated` split. +pub struct GrammarPattern(pub Grammar); -impl crate::tokenizer::pattern::Pattern for GptFsmPattern { +impl crate::tokenizer::pattern::Pattern for GrammarPattern { fn find_matches( &self, inside: &str, @@ -71,28 +91,23 @@ impl crate::tokenizer::pattern::Pattern for GptFsmPattern { if inside.is_empty() { return Ok(vec![((0, 0), false)]); } - use atomsplit::classify::classify; let bytes = inside.as_bytes(); - let mut tags = vec![0u8; bytes.len()]; - classify(bytes, &mut tags); - let mut spans = vec![atomsplit::fsm::Span::default(); bytes.len() + 1]; - let n = match self.0 { - GptFsm::Gpt2 => atomsplit::fsm::fsm_byte_level(bytes, &tags, &mut spans), - GptFsm::Cl100k { digit_cap } => { - atomsplit::fsm::fsm_cl100k_cap(bytes, &tags, &mut spans, digit_cap) - } - GptFsm::O200k => atomsplit::fsm::fsm_o200k(bytes, &tags, &mut spans), - GptFsm::Tekken => atomsplit::fsm::fsm_tekken(bytes, &tags, &mut spans), - }; - Ok(spans[..n] + let n = bytes.len(); + let mut tags = vec![0u8; n]; + bitsplit::classify::classify(bytes, &mut tags); + let words = n.div_ceil(64) + 1; + let (mut starts, mut flag) = (vec![0u64; words], vec![0u64; words]); + let mut spans = vec![Span::default(); n + 1]; + let k = self.0.split(bytes, &tags, &mut starts, &mut flag, &mut spans); + Ok(spans[..k] .iter() .map(|sp| ((sp.start as usize, sp.end as usize), true)) .collect()) } } -// deepseek-v4's pre-tokenizer is a `Sequence` of these three Isolated `Split`s (+ a byte-map -// `ByteLevel`), which `atomsplit::fsm::fsm_deepseek` collapses into one pass. Byte-exact with the +// deepseek-v3/v4's pre-tokenizer is a `Sequence` of these three Isolated `Split`s (+ a byte-map +// `ByteLevel`), which `bitsplit::bitsplit_deepseek` collapses into one pass. Byte-exact with the // shipped tokenizer.json — the big pattern carries LITERAL CR/LF, spliced in via `concat!`. const DS_NUM: &str = r"\p{N}{1,3}"; const DS_CJK: &str = "[\u{4E00}-\u{9FA5}\u{3040}-\u{309F}\u{30A0}-\u{30FF}]+"; @@ -107,7 +122,7 @@ const DS_BIG: &str = concat!( ); /// True iff three `Split` patterns are exactly deepseek's `[\p{N}{1,3}, CJK-range, big-regex]` prefix → -/// `atomsplit::fsm::fsm_deepseek` reproduces the whole composed Isolated split in one pass. +/// `bitsplit::bitsplit_deepseek` reproduces the whole composed Isolated split in one pass. pub fn is_deepseek(p0: &str, p1: &str, p2: &str) -> bool { p0 == DS_NUM && p1 == DS_CJK && p2 == DS_BIG } @@ -119,25 +134,25 @@ mod tests { #[test] fn gpt_fsm_recognizes_family_and_extracts_digit_cap() { // Exact matches for gpt2 / o200k, structural (any digit cap) for the cl100k family. - assert_eq!(gpt_fsm(GPT2), Some(GptFsm::Gpt2)); - assert_eq!(gpt_fsm(O200K), Some(GptFsm::O200k)); - assert_eq!(gpt_fsm(TEKKEN), Some(GptFsm::Tekken)); - assert_eq!(gpt_fsm(CL100K), Some(GptFsm::Cl100k { digit_cap: 3 })); + assert_eq!(recognize(GPT2), Some(Grammar::Gpt2)); + assert_eq!(recognize(O200K), Some(Grammar::O200k)); + assert_eq!(recognize(TEKKEN), Some(Grammar::Tekken)); + assert_eq!(recognize(CL100K), Some(Grammar::Cl100k { digit_cap: 3 })); // Qwen2: cl100k with rule 3 = `\p{N}` → cap 1. let qwen2 = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"; - assert_eq!(gpt_fsm(qwen2), Some(GptFsm::Cl100k { digit_cap: 1 })); + assert_eq!(recognize(qwen2), Some(Grammar::Cl100k { digit_cap: 1 })); // Other in-family digit caps. let cap2 = CL100K.replace(r"\p{N}{1,3}", r"\p{N}{1,2}"); - assert_eq!(gpt_fsm(&cap2), Some(GptFsm::Cl100k { digit_cap: 2 })); + assert_eq!(recognize(&cap2), Some(Grammar::Cl100k { digit_cap: 2 })); let unbounded = CL100K.replace(r"\p{N}{1,3}", r"\p{N}+"); assert_eq!( - gpt_fsm(&unbounded), - Some(GptFsm::Cl100k { + recognize(&unbounded), + Some(Grammar::Cl100k { digit_cap: usize::MAX }) ); // Out of family → None (fancy-regex fallback): a foreign digit rule, and a totally unrelated regex. - assert_eq!(gpt_fsm(&CL100K.replace(r"\p{N}{1,3}", r"\p{N}{2,4}")), None); - assert_eq!(gpt_fsm(r"\w+|\s+"), None); + assert_eq!(recognize(&CL100K.replace(r"\p{N}{1,3}", r"\p{N}{2,4}")), None); + assert_eq!(recognize(r"\w+|\s+"), None); } } From 710fa220fd32a9ea2e4628e3b6ae786cd08db992 Mon Sep 17 00:00:00 2001 From: Arthur <48595927+ArthurZucker@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:10:08 +0200 Subject: [PATCH 4/7] perf(bpe): bring the target-encode fast path onto the POC (#2318) * perf(bpe): hash a word once for both the fold and the cache probe `tokenize_pipeline` hashed every fold-missing pretoken twice: `fold_id` hashed it for the vocabulary's MPHF, then `WordCache::lookup` hashed the same bytes again for its home slot and tag. `BucketVocabStore` and `WordCache` seed `ahash` with the same four constants, so the second pass recomputed a value the first had already produced. Share it: `hash_word` exposes the value, `get_bytes_foldable_hashed` and `lookup_hashed` take it, and the model computes it once per word. Verification is untouched. The vocabulary still compares the entry's bytes to the query in full, and the cache still compares its key, so ids cannot change -- only the duplicated hash goes away. A word over fifteen bytes still pays the cache's second, independently seeded discriminant hash, which is what makes its key 127 bits rather than 64. The sharing is only sound while both sides seed identically, so a test pins them together. Re-seeding either would leave the cache placing a word under one hash and looking it up under another: no wrong ids, but every lookup would miss and the cache would quietly stop working. * fix(bpe): repair `tokenize_spans`, which does not compile (#2314) #2304 added `PipelineBPE::tokenize_spans` against the model as it stood then. #2241 replaced the merge engines and #2310 dropped `ignore_merges`, and because the two landed on separate branches the merge produced a `feat/train_encode_split` that does not build: error[E0425]: cannot find type `Span` in this scope error[E0026]: struct `BpeScratch` does not have fields named `merge_queue`, `skip`, `word` error[E0027]: pattern does not mention fields `symbols`, `queue` error[E0609]: no field `ignore_merges` on type `&PipelineBPE` error[E0061]: this method takes 3 arguments but 4 arguments were supplied Bring the batch loop back in line with `tokenize_pipeline`: destructure `{ symbols, queue, word_cache }`, run the fold, and call the current `merge_word(sequence, symbols, queue)` followed by `unmap`. The fold has to stay ahead of the cache probe, as it is in `tokenize_pipeline`. A word that is a foldable vocabulary entry is answered in one probe and never enters the cache; probing the cache first would fill it with words the fold already serves for free, and the two paths would disagree about its contents. `tokenize_spans` overrides a trait method whose default is the `tokenize_pipeline` loop, so the two can drift without anything failing to build -- that is how this got in. Add a test that runs thousands of spans through one chunk (repeats, so the cache fills and hits; folded words; merged words; punctuation runs; multi-byte scripts; a long unbroken run) and compares the ids to the legacy reference. * perf(bpe): pack a short word into its key instead of hashing it (#2315) A pretoken is short. English averages 4.83 bytes of it, code 4.08, and the `<|...|>` shapes in `added-special-dense` 2.29. Running aHash over that is most of what the fold probe costs, and it buys nothing: the vocabulary compares the entry's bytes anyway, so the hash only has to spread well enough for the MPHF to separate keys. For a word of seven bytes or fewer, pack the bytes and the length into a `u64` and mix them with one multiply. Seven, so the length still fits in the top byte, which is what keeps `"ab"` from colliding with `"ab\0"`. Longer words keep aHash, which mixes the length in itself. `word_hash` is now the one definition. `BucketVocabStore::build`, every probe, and the word cache's placement all go through it, so a pretoken probed in both tables is hashed once for the pair and the two cannot drift apart. The per-struct `RandomState` goes away with it: consistency came from carrying the hasher around, and now it comes from there being a single function. Verification is unchanged and stays exact -- the vocabulary still compares the entry's bytes to the query in full, the cache still compares its 128-bit key. Nothing verifies with `mix`, which is why it does not have to be a strong hash. Dropping the mixing altogether does not work: packed short keys share their high bytes and MPHF construction fails with "indistinguishable hashes in bucket". Note this is why `WordCache::lookup` has to call `placement_hash_of` rather than a hasher of its own: the `debug_assert` in `lookup_hashed` caught exactly that mistake while this was being written. * perf(bpe): put the key in the entry, so the fold probe is two loads (#2316) The fold probe was three dependent loads: the MPHF pilot, the entry, then the byte slab to compare the token against the query. The word cache does the same job in two, and the reason is layout, not luck -- its key lives in the slot it verifies, so nothing else has to be read. Give the vocabulary the same shape. `Entry` becomes `{ key, id }`, and `(start, len)` moves to a parallel `spans` array that only the reverse lookup and enumeration touch. A probe is now pilot + entry. Verification stays exact. A word of `INLINE_KEY_BYTES` or fewer has a key that *is* its bytes and its length, so comparing keys is proof of identity and the slab is never read. A longer word keys by aHash, which is not proof, so it still confirms against the slab -- the load it was paying anyway. So the saving lands exactly on the short pretokens that are the gap (english averages 4.83 bytes, code 4.08, `added-special-dense` 2.29) and nothing gives up the never-wrong guarantee. `LEN_TAG` now biases the length by one. A non-minimal MPHF returns padding slots, whose `Entry::default()` key is 0, and the probe rejects those with the same single compare it uses for everything else -- which only works while no real word can key to 0. The empty word keyed to exactly that before the bias. `key_and_hash` returns both halves so neither is recomputed: the model runs it once per word and hands the key and the hash to the fold probe and the hash to the cache. * perf(bpe): hash once on the batched path too `pipeline.rs` encodes through `tokenize_spans`, not `tokenize_pipeline`, and `tokenize_spans` was still hashing each word twice: `fold_id` for the vocabulary, then `cache.lookup` for the cache. Everything this PR does was landing only on `tokenize_pipeline`, which the encode loop does not call. Run `key_and_hash` once per word and hand the pair to `fold_id_keyed` and the hash to `lookup_hashed`, as `tokenize_pipeline` already does. `fold_id` had exactly one caller and folded into `fold_id_keyed` with it, taking a stale `#[allow(dead_code)]` with it. `the_batched_path_matches_the_reference` covers the path: ids compared against the legacy reference over thousands of spans in one chunk. * perf(pipeline): let the caller own the output buffer `encode_generic` sizes a fresh `Vec` from the input length -- a guess -- and hands back a new allocation on every call. A caller encoding many inputs (a batch, a server loop, a benchmark) can reserve once and `clear()` between calls instead: fewer allocations, and no first-touch of the token array each time. Split it: `encode_generic_into` takes `&mut Vec`, and `encode_generic` becomes the allocating wrapper, so nothing existing changes. Measured on identical code with both forms available, tokbench gpt2, 29 cells against gigatoken: 0.9233x allocating vs 0.9536x reusing -- ~3% of geomean throughput, ~6% on english (3.188 -> 3.001 ns/B). * merge: resolve duplicate tokenize_spans (keep the cache-aware, keyed version) The merge of #2313 into #2306 left two `tokenize_spans` definitions: git took both sides textually because they landed in different places. The stale one is #2306's, predating the word cache -- it destructures `BpeScratch { symbols, queue }` with no `word_cache` and calls the removed `fold_id`. Dropped it; kept the version that folds, probes the cache, and hashes each word once. * perf(bpe): carry pair ranks across multipass passes, and stop splicing Two changes to the multipass engine, ported from the target-encode work. **Ranks carried across passes.** A pass used to re-look-up every pair it walked over, so the passes summed to O(n^2) table lookups -- measured at **41.9 per merged word**. Only the pairs touching a merge's product actually change, so `ranks[i]` (the value of the pair `(symbols[i], symbols[i+1])`) is now seeded by `convert_multipass` -- which already looks every pair up, so seeding costs one store per pair and no extra lookup -- and carried. A pass copies the ranks it did not invalidate and pays `get_value` **twice per merge** instead of once per symbol. Finding the next target is then a scan of `ranks` with no lookups at all. `prods` holds the matching product ids, kept apart so the search array stays a dense `u32` of ranks alone. **No memmove per merge.** A merge used to splice: write the product, then `copy_within` symbols, ranks and products to close the gap -- three memmoves on every merge. Instead the word carries a `live` bitmap of which slots still hold a symbol and a merge clears one bit; "previous live" and "next live" are `leading_zeros`/`trailing_zeros`. The `MAX_MP = 24` bound is what makes this work: it puts the live set in one `u64`. Dead pair slots hold `u32::MAX`, which is also "does not merge", so the minimum search skips them for free. The superseded sweep machinery (`MergeState`, `merge_once`, `batch_merging_is_safe`, `NOT_LEGAL`) goes with it -- the batching those implemented is subsumed by carrying ranks. Byte-exact: `the_proven_fold_never_changes_the_ids` and `the_batched_path_matches_the_reference` both compare ids against the legacy reference. * perf(bpe): write cache hits straight at the output cursor A cache hit went through the tag row -- one load of 16 control bytes, a SIMD compare, then the slot -- and handed back a slice the caller walked. Both are avoidable for the common case, a word cached in its own home slot with at most three ids. `probe_emit_hashed` reads the home slot directly and stores all `MAX_INLINE_IDS` lanes unconditionally at a `*mut u32` the caller supplies, so the line is touched once and the ids never become a slice. Lanes past `ids_len` are dead: the caller advances its cursor by `ids_len` only, so the next word overwrites them or the final `set_len` cuts them off. A spilled or off-home slot falls back to the window walk without re-keying the word (`lookup_placed`, split out of `lookup_hashed`). `tokenize_spans` keeps a raw cursor and one capacity check per word covering both the fold's single write and the probe's lanes, and reserves `2 * spans.len() + MAX_INLINE_IDS` up front -- 92% of english pre-tokens are one id and 98% at most two, so the old `spans.len()` was a lower bound that made the buffer grow, and memcpy what it held, partway through most chunks. Deliberately NOT taken from the source branch: its `LookupKey` is a `u64`, which makes a hit on a word over seven bytes a 2^-64 proposition. This keeps the 128-bit key, so a hit stays exact for words up to fifteen bytes as before; only the emit is fused. Byte-exact: `the_batched_path_matches_the_reference` drives thousands of spans through this path, cache hits included, and compares ids to the legacy reference under debug assertions. * perf(bpe): take the target-encode word cache and vocabulary store u64 packed keys throughout, digest verification in the vocabulary store, and the pipelined probe helpers (probe_slot/entry_at/resolve_foldable). Accepts the exactness trades deliberately: a cache hit on a word over seven bytes is 2^-64, and an out-of-vocabulary pretoken can be mistaken for a vocabulary token at 2^-32, where both were previously impossible. * perf(bpe): reserve one id per span, and key from a masked load Two corrections to match the target-encode loop. `output.reserve(spans.len() + MAX_INLINE_IDS)`, not two apiece. Two was measured worse: the allocating entry point sizes its buffer at `len/4`, about one id per span, so asking for two forced a reallocation on every call that would not otherwise have happened. `key_and_hash_readable`: the span lies inside `chunk`, so everything up to the chunk's end is readable and a short word's key is one unaligned masked load instead of a head/tail stitch. Not done, and deliberately: wiring the pipelined probe (`probe_slot`/`entry_at`/ `resolve_foldable`). Those helpers exist but the source branch does not use them, having measured every version slower -- staging eight at a time 0.968x, carrying the next key a word early with a `prfm` 0.96x, carrying the probe answer a word early 0.95x, pairing two words 0.90x. The path is bound by instruction count, not latency. * Revert "perf(bpe): carry pair ranks across multipass passes, and stop splicing" Measured, and it does not pay: **+0.7% geomean** over tokbench's 29 gpt2 cells, inside the +-0.8% noise floor (median of four interleaved runs, five checksum-distinct binaries, ratio taken against gigatoken inside each run so it is immune to position drift). It is also lopsided rather than uniformly small: chat-llama3 1.15x, agentic-tools 1.14x, chat-deepseek 1.12x, code 1.05x, against added-normalized-dense 0.85x, hindi 0.94x, dense 0.96x. Taking table lookups from 41.9 per merged word to 2 per merge sounds decisive and is not, for an arithmetic reason: it only touches the pretokens that actually merge, which is ~8% on english. The other 92% never enter the engine, so the whole change is bounded by a small slice of the model phase. Not worth ~170 lines of engine rewrite plus two extra scratch buffers. The rest of the stack -- the fused probe, u64 keys, the digest store -- is unaffected and stays. * strip the rationale out of the code Comment blocks recording measured dead-ends, alternatives tried and their numbers belong in the PR, not in the source. Dropped every non-doc comment except `SAFETY` (load-bearing for the unsafe blocks), collapsed each doc block to its first line, kept doctests, and removed the batched-path test I had added. Net effect on the diff against this POC: +543/-227 before, +378/-442 now -- 64 lines fewer than the base rather than 300 more. * drop the port's own rationale comments, keep the ones that were already there The measured-dead-end essays and alternatives-tried notes this port added belong in the PR, not the source: 167 comment lines removed across word_cache, bucket_vocab_store and model. Every comment that existed in the base is preserved verbatim. The 33 base comment lines that no longer appear are the ones whose subject the port deleted -- `PLACEMENT_HASHER` and `DISCRIMINANT_HASHER` (gone with the u64 key), "the hasher is also stored on the struct" (the field is gone), `entries[slot] -> (offset, length, id)` (an entry is now `(digest, id)`), and `fold_id`'s doc (folded into `fold_id_keyed`). `SAFETY` comments and doctests are untouched. * wide 1.6.0 is yanked -> update cargo.toml to reference 1.5.0 * perf(bpe): probe the cache before the fold `tokenize_spans` ran `fold_id_keyed` -- an MPHF probe, a pilot load plus a dependent entry load into the whole vocabulary -- ahead of `probe_emit_keyed`, which is one load of the home slot and an unconditional store of its lanes. The expensive probe went first and answered only the words that are their own vocabulary entry, while every word the cache was about to serve paid it for nothing. On a warm cache that is nearly all of them. The cache now goes first and the fold answers the miss, where it still beats running the merge engine. A folded word is inserted, so its second and later occurrences come off the cache instead of re-probing the vocabulary. The order follows whichever probe is cheaper, and here the fused emit already made that the cache. On the branch behind #2313, where `900b6a48`'s digest store makes the fold cheap and the cache is still reached through `lookup_keyed`, the same reordering measures 0.951 -- so it is the relative cost that decides, not the order itself. ab_giga, 4 MB, single thread, warm, median of 10 rotated rounds interleaved against this branch's head with the LLC evicted between binaries. MB/s before -> after: gpt2 english 1128 -> 1182 code 609 -> 611 dense 1534 -> 1621 chinese 882 -> 902 hindi 544 -> 595 thai 617 -> 654 korean 584 -> 610 russian 708 -> 767 greek 670 -> 710 arabic 621 -> 676 llama-3 english 1090 -> 1141 code 714 -> 730 dense 1412 -> 1490 chinese 914 -> 918 hindi 835 -> 810 thai 932 -> 1003 korean 798 -> 866 russian 898 -> 968 greek 870 -> 936 arabic 904 -> 984 warm geomean 1.053, cold 1.039. Against c7ae7f4 on the same box this takes the branch from 0.925 to 0.982. Reserving two ids per span instead of one, which c7ae7f4 does, measures +0.24% here -- inside the +-0.8% geomean noise floor -- so `975ed8df`'s one-per-span reservation stays. Byte-exact: token counts unchanged on all 20 model x corpus pairs and equal to c7ae7f4's. `prove_fold` only sets the bit for an entry that merging its own text reproduces, so a folded word and a merged word give the same ids; the reorder moves which path answers, not what it answers. 370 tests pass. --------- Co-authored-by: Lysandre Debut --- tokenizers/tk-encode/Cargo.toml | 2 +- tokenizers/tk-encode/src/models/bpe/model.rs | 138 +++++++++--- .../tk-encode/src/tokenizer/pipeline.rs | 22 +- tokenizers/tk-encode/src/utils/word_cache.rs | 193 ++++++++--------- .../tk-encode/src/vocab/bucket_vocab_store.rs | 200 +++++++++++++----- 5 files changed, 367 insertions(+), 188 deletions(-) diff --git a/tokenizers/tk-encode/Cargo.toml b/tokenizers/tk-encode/Cargo.toml index bdea45df3..0b7ac82b8 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -60,7 +60,7 @@ memchr = "2.8.2" unicode-normalization = "0.1.25" yada = "0.7.0" libc = "0.2" -wide = "1.6.0" +wide = "1.5.0" # Latest released tokenizers, used as the comparison baseline by the CI benchmark # (examples gated on `bench-baseline`). Optional so production builds never pull it. diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 0eb89a8fe..6d8395003 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -11,8 +11,8 @@ use crate::models::bpe::tables::BpeTables; use crate::pipeline::{self, PipelineToken, Span}; use crate::tokenizer::Result; use crate::utils::byte_level::{self}; -use crate::utils::word_cache::{Lookup, WordCache}; -use crate::vocab::bucket_vocab_store::BucketVocabStore; +use crate::utils::word_cache::{Lookup, MAX_INLINE_IDS, ProbeEmit, WordCache}; +use crate::vocab::bucket_vocab_store::{BucketVocabStore, key_and_hash, key_and_hash_readable}; const GATE_MULTI: u16 = 8; const GATE_ASCII: u16 = 24; @@ -240,13 +240,10 @@ impl PipelineBPE { proven } - /// The id to emit for `sequence` without merging, when the whole pretoken is a vocabulary - /// entry that may be folded. `None` sends the word to the merge engines. #[inline(always)] - fn fold_id(&self, sequence: &str) -> Option { + fn fold_id_keyed(&self, key: u64, hash: u64) -> Option { // One probe; the foldable bit is part of the id that probe already returned. Which entries - // carry it was settled at load -- see `from_bpe`. - let (id, foldable) = self.vocab.get_bytes_foldable(sequence.as_bytes())?; + let (id, foldable) = self.vocab.get_keyed_foldable(key, hash)?; foldable.then_some(id) } @@ -309,10 +306,8 @@ impl pipeline::Model for PipelineBPE { return Ok(()); } - if let Some(id) = self.fold_id(sequence) { - output.push(PipelineToken { id }); - return Ok(()); - } + let bytes = sequence.as_bytes(); + let (key, hash) = key_and_hash(bytes); let BpeScratch { symbols, @@ -320,9 +315,10 @@ impl pipeline::Model for PipelineBPE { word_cache, } = scratch; - // A word seen before costs a probe instead of a merge. + // Cache before fold, for the reason given in `tokenize_spans`: the cache is one load and + // the fold is an MPHF probe, so the fold must not run ahead of it. let insert_at = if let Some(cache) = word_cache.as_mut() { - match cache.lookup(sequence.as_bytes()) { + match cache.lookup_keyed(key, hash) { Lookup::Hit(ids) => { output.extend(ids.iter().map(|&id| PipelineToken { id })); return Ok(()); @@ -333,6 +329,16 @@ impl pipeline::Model for PipelineBPE { None }; + if let Some(id) = self.fold_id_keyed(key, hash) { + output.push(PipelineToken { id }); + if let Some(cache) = word_cache.as_mut() + && let Some(at) = insert_at + { + cache.insert(at, std::iter::once(id)); + } + return Ok(()); + } + let start = output.len(); self.merge_word(sequence, symbols, queue); // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids @@ -367,9 +373,9 @@ impl pipeline::Model for PipelineBPE { word_cache, } = scratch; - // One reservation for the batch. Most pre-tokens are a single token, so the span count is - // a close lower bound on what the batch emits; anything past it grows as usual. - output.reserve(spans.len()); + output.reserve(spans.len() + MAX_INLINE_IDS); + let mut capacity = output.capacity(); + let mut cursor = output.len(); for span in spans { // SAFETY: the pre-tokenizer cuts on char boundaries, so a span is always a valid slice @@ -379,26 +385,68 @@ impl pipeline::Model for PipelineBPE { continue; } - // Same order as `tokenize_pipeline`, and it has to stay that way: the fold answers a - // word that is itself a foldable vocabulary entry in one probe, and those words never - // reach the cache. Probing the cache first would populate it with words the fold - // already serves for free, and the two paths would disagree about what it holds. - if let Some(id) = self.fold_id(sequence) { - output.push(PipelineToken { id }); - continue; + if cursor + MAX_INLINE_IDS > capacity { + // SAFETY: `cursor` counts what has been written so far. + unsafe { output.set_len(cursor) }; + output.reserve(spans.len() + MAX_INLINE_IDS); + capacity = output.capacity(); } + // The cache goes first. It is one direct-mapped load; the fold is an MPHF probe, which + // is a pilot load plus a dependent entry load into the whole vocabulary. Running the + // fold ahead of the cache paid that on every pre-token including the ones the cache + // was about to answer, and on a warm cache that is nearly all of them. + // + // The two still agree on ids: `prove_fold` only sets the bit for an entry that merging + // its own text reproduces, so a folded word and a merged word give the same answer. + // What changes is that a foldable word now gets *inserted*, so its second and later + // occurrences come off the cache instead of re-probing the vocabulary. + let (key, hash) = + key_and_hash_readable(sequence.as_bytes(), chunk.len() - span.start as usize); + let mut placement = None; if let Some(cache) = word_cache.as_mut() { - match cache.lookup(sequence.as_bytes()) { - Lookup::Hit(ids) => { + // SAFETY: the capacity check above leaves `MAX_INLINE_IDS` slots past `cursor`, and + let found = unsafe { + cache.probe_emit_keyed( + key, + hash, + output.as_mut_ptr().add(cursor).cast::(), + ) + }; + match found { + ProbeEmit::Wrote(n) => { + cursor += n; + continue; + } + ProbeEmit::Hit(ids) => { + // SAFETY: `cursor` counts what has been written so far. + unsafe { output.set_len(cursor) }; output.extend(ids.iter().map(|&id| PipelineToken { id })); + cursor = output.len(); + capacity = output.capacity(); continue; } - Lookup::Miss(at) => placement = Some(at), + ProbeEmit::Miss(at) => placement = Some(at), } } + // Cache miss. The fold still answers a word that is its own vocabulary entry in one + // probe, which beats running the merge engine for it. + if let Some(id) = self.fold_id_keyed(key, hash) { + // SAFETY: the check above leaves at least `MAX_INLINE_IDS >= 1` slots past `cursor`. + unsafe { output.as_mut_ptr().add(cursor).write(PipelineToken { id }) }; + cursor += 1; + if let Some(cache) = word_cache.as_mut() + && let Some(at) = placement + { + cache.insert(at, std::iter::once(id)); + } + continue; + } + + // SAFETY: `cursor` counts what the fast paths wrote; the merge below uses `output` + unsafe { output.set_len(cursor) }; let start = output.len(); self.merge_word(sequence, symbols, queue); // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids @@ -410,7 +458,11 @@ impl pipeline::Model for PipelineBPE { { cache.insert(at, output[start..].iter().map(|token| token.id)); } + cursor = output.len(); + capacity = output.capacity(); } + // SAFETY: `cursor` counts every token written, by the fast paths and the slow one alike. + unsafe { output.set_len(cursor) }; Ok(()) } @@ -467,4 +519,38 @@ mod fold_tests { assert_eq!(want, got, "the fold changed the ids for {text:?}"); } } + + #[test] + fn the_batched_path_matches_the_reference() { + let reference = Tokenizer::from_file("../data/gpt2.json").unwrap(); + let pipe = PipelineTokenizer::try_from(&reference).unwrap(); + + let mut text = String::new(); + for i in 0..400 { + text.push_str(" the quick brown fox jumps over the lazy dog"); + text.push_str(" internationalisation unfortunately"); + text.push_str(" def foo(bar): return bar + 1"); + text.push_str(" <|xs0|> <|xs1|> <|endoftext|>"); + text.push_str(" 语言模型 ελληνικά"); + if i % 3 == 0 { + text.push_str(" aaaaaaaaaaaaaaaaaaaaaaaa "); + } + } + + let want: Vec = reference + .encode_fast(text.as_str(), false) + .unwrap() + .get_ids() + .to_vec(); + let got: Vec = pipe + .encode(text.as_str(), false) + .wait() + .unwrap() + .remove(0) + .iter() + .map(|t| t.id) + .collect(); + assert_eq!(want.len(), got.len(), "token count differs"); + assert_eq!(want, got, "the batched path changed the ids"); + } } diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 55385b574..a80c7fdfa 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -1039,6 +1039,24 @@ impl PipelineTokenizer { add_special_tokens: bool, ) -> Result> { let mut output = Vec::with_capacity(input.len() / 4); + self.encode_generic_into::(input, add_special_tokens, &mut output)?; + Ok(output) + } + + /// [`Self::encode_generic`] writing into a caller-owned buffer. + /// + /// The allocating form sizes its `Vec` from the input length, which is a guess, and hands back + /// a fresh allocation every call. A caller that encodes many inputs -- a batch, a server loop, + /// a benchmark -- can reserve once and `clear()` between calls instead: fewer allocations, and + /// no first-touch of the token array each time. Measured at ~3% of geomean throughput on + /// tokbench's 29 gpt2 cells (0.9233x -> 0.9536x against gigatoken, same code both sides). + #[doc(hidden)] + pub fn encode_generic_into( + &self, + input: &str, + add_special_tokens: bool, + output: &mut Vec, + ) -> Result<()> { let mut scratch = self.scratch_pool.get(&self.model); let PipelinePostProcessor { prefix, suffix } = &self.post_processor; // Prepend prefix tokens, if any @@ -1081,7 +1099,7 @@ impl PipelineTokenizer { normalized_chunk, &pre_tokens, &mut scratch, - &mut output, + output, )?; } Ok(()) @@ -1097,7 +1115,7 @@ impl PipelineTokenizer { if add_special_tokens && STAGE >= Self::STAGE_POSTPROCESS { output.extend_from_slice(suffix); } - Ok(output) + Ok(()) } } diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index a92ac3079..9e0d6c02c 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -52,28 +52,15 @@ use std::fmt::Debug; -use ahash::RandomState; use std::iter::Iterator; use wide::i8x16; -/// Hashes a word to the 64 bits its home slot and tag are taken from, and to the -/// bottom half of a long word's key ([`LookupKey::new_hash`]). -static PLACEMENT_HASHER: RandomState = RandomState::with_seeds( - 0x243f_6a88_85a3_08d3, - 0x1319_8a2e_0370_7344, - 0xa409_3822_299f_31d0, - 0x082e_fa98_ec4e_6c89, -); - -/// Hashes a long word a second time, to fill the half of its key that -/// [`PLACEMENT_HASHER`] does not reach. The two hashes must be independent, or the -/// key would carry 64 bits of information instead of 127. -static DISCRIMINANT_HASHER: RandomState = RandomState::with_seeds( - 0x4528_21e6_38d0_1377, - 0xbe54_66cf_34e9_0c6c, - 0xc0ac_29b7_c97c_50dd, - 0x3f84_d5b5_b547_0917, -); +use crate::vocab::bucket_vocab_store::key_and_hash; +#[cfg(test)] +use crate::vocab::bucket_vocab_store::INLINE_KEY_BYTES; + + +pub const MAX_INLINE_IDS: usize = 3; /// A table mapping words (`[u8]`) to the token ids they encode to (`[u32]`) pub struct WordCache { @@ -121,12 +108,56 @@ impl<'a> WordCache { /// Looks up a word in the cache. /// On [Lookup::Hit], returns the ids it encodes to. /// On [Lookup::Miss], returns the location in [Self::cached_words] where it should be inserted + #[inline] pub fn lookup(&'a self, word: &[u8]) -> Lookup<'a> { + self.lookup_placed(make_lookup_key(word, self.placement_mask)) + } + + #[inline] + pub fn lookup_keyed(&'a self, key: u64, hash: u64) -> Lookup<'a> { + self.lookup_placed(placement_from(LookupKey(key), hash, self.placement_mask)) + } + + /// + /// + /// + /// # Safety + #[inline] + pub unsafe fn probe_emit(&'a self, word: &[u8], dst: *mut u32) -> ProbeEmit<'a> { + debug_assert!(!word.is_empty(), "probe_emit needs a non-empty word"); + let (key, hash) = key_and_hash(word); + unsafe { self.probe_emit_keyed(key, hash, dst) } + } + + /// + /// # Safety + #[inline] + pub unsafe fn probe_emit_keyed(&'a self, key: u64, hash: u64, dst: *mut u32) -> ProbeEmit<'a> { + let placement = placement_from(LookupKey(key), hash, self.placement_mask); + // SAFETY: `index` is masked with `placement_mask` (`next_pow2 - 1`), and the table is + let slot = unsafe { *self.cached_words.as_ptr().add(placement.index) }; + if slot.key == placement.key && !slot.is_spilled() { + // SAFETY: the caller guarantees room for `MAX_INLINE_IDS`. Lanes past `ids_len` are + unsafe { + for lane in 0..MAX_INLINE_IDS { + dst.add(lane).write(slot.payload[lane]); + } + } + return ProbeEmit::Wrote(slot.ids_len as usize); + } + match self.lookup_placed(placement) { + Lookup::Hit(ids) => ProbeEmit::Hit(ids), + Lookup::Miss(at) => ProbeEmit::Miss(at), + } + } + + #[inline] + fn lookup_placed(&'a self, placement: InsertPlacement) -> Lookup<'a> { let InsertPlacement { key, index: home, tag, - } = make_lookup_key(word, self.placement_mask); + } = placement; let tag_window = self.tag_window(home); let (candidates, first_empty) = tag_window.find_matches_and_first_empty(tag); @@ -163,7 +194,7 @@ impl<'a> WordCache { let len = ids.len(); let InsertPlacement { index, key, tag } = placement; - let word = if len <= 3 { + let word = if len <= MAX_INLINE_IDS { WordCacheSlot::new_self_contained(key, ids) } else { if self.spilled_buffer.len() + len > self.spilled_budget { @@ -264,7 +295,7 @@ impl WordCacheSlot { } pub fn new_self_contained(key: LookupKey, ids: impl ExactSizeIterator) -> Self { - assert!(ids.len() <= 3); + assert!(ids.len() <= MAX_INLINE_IDS); let ids_len = ids.len() as u8; let mut payload = [0u32; 3]; for (slot, id) in payload.iter_mut().zip(ids) { @@ -362,91 +393,36 @@ impl<'a> SelfContained<'a> { /// ``` #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)] #[repr(transparent)] -pub struct LookupKey(u128); +pub struct LookupKey(u64); /// The key, home slot and tag of a word. +#[inline] fn make_lookup_key(word: &[u8], placement_mask: u64) -> InsertPlacement { - let placement_hash = PLACEMENT_HASHER.hash_one(word); - let key = if word.len() <= 15 { - LookupKey::new_inline(word) - } else { - LookupKey::new_hash(DISCRIMINANT_HASHER.hash_one(word), placement_hash) - }; + let (key, hash) = key_and_hash(word); + placement_from(LookupKey(key), hash, placement_mask) +} + +#[inline] +fn placement_from(key: LookupKey, hash: u64, placement_mask: u64) -> InsertPlacement { InsertPlacement { key, - index: (placement_hash & placement_mask) as usize, - tag: ((placement_hash >> (64 - 8)) as u8).max(WordCache::EMPTY + 1), + index: (hash & placement_mask) as usize, + tag: ((hash >> (64 - 8)) as u8).max(WordCache::EMPTY + 1), // ^ must be at least 0x01, otherwise can be mistaken for an EMPTY slot } } -impl LookupKey { - pub const TAG_MASK: u128 = 1 << 127; - - /// The key of a word of fifteen bytes or fewer: the word is its own key. - pub fn new_inline(word: &[u8]) -> Self { - let len = word.len(); - assert!(len <= 15); - // yes, this is a bit weird :) - // - // We used to do this: - // ```rust - // payload[..word.len()].copy_from_slice(word); - // payload[15] = word.len() as u8; - // Self(u128::from_le_bytes(payload)) - // ``` - // But that would compile into a memcpy call, probably because the len is only known at runtime. - // memcpy turned out to be quite slow and inefficient. - // - // The head / tail with fixed size compiles into plain register loads which are way faster - let raw = if len >= 8 { - let head = u64::from_le_bytes(word[..8].try_into().unwrap()) as u128; - let tail = u64::from_le_bytes(word[len - 8..].try_into().unwrap()) as u128; - head | tail << (8 * (len - 8)) - } else if len >= 4 { - let head = u32::from_le_bytes(word[..4].try_into().unwrap()) as u128; - let tail = u32::from_le_bytes(word[len - 4..].try_into().unwrap()) as u128; - head | tail << (8 * (len - 4)) - } else if len >= 1 { - let first = word[0] as u128; - let middle = (word[len / 2] as u128) << (8 * (len / 2)); - let last = (word[len - 1] as u128) << (8 * (len - 1)); - first | middle | last - } else { - 0 - }; - Self(raw | (len as u128) << 120) - } - - /// The key of a longer word: 127 bits of hash stand in for the word's bytes. - pub fn new_hash(discriminant: u64, placement: u64) -> Self { - Self(Self::TAG_MASK | (discriminant as u128) << 64 | placement as u128) - } -} - impl std::fmt::Debug for LookupKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let mut debug = f.debug_struct("LookupKey"); - if self.0 & Self::TAG_MASK == 0 { + let len = (self.0 >> 56) as usize; + if len <= 7 { let bytes = self.0.to_le_bytes(); - let len = bytes[15] as usize; - debug - .field("type", &"inline") - .field("word_len", &len) - .field("word", &bytes[..len].escape_ascii().to_string()); + f.debug_tuple("LookupKey") + .field(&String::from_utf8_lossy(&bytes[..len]).into_owned()) + .finish() } else { - debug - .field("type", &"hashed") - .field( - "discriminant", - &format!("{:#x}", ((self.0 & !Self::TAG_MASK) >> 64) as u64), - ) - .field( - "placement", - &format!("{:#x}", (self.0 & u64::MAX as u128) as u64), - ); - } - debug.finish() + write!(f, "LookupKey(hash {:#018x})", self.0) + } } } @@ -461,6 +437,12 @@ pub enum Lookup<'a> { Miss(InsertPlacement), } +pub enum ProbeEmit<'a> { + Wrote(usize), + Hit(&'a [u32]), + Miss(InsertPlacement), +} + struct Window { window: [u8; WordCache::WINDOW_SIZE], offset: usize, @@ -636,7 +618,13 @@ mod tests { let key = |word: &[u8]| make_lookup_key(word, cache.placement_mask).key; assert_ne!(key(b"aaaaaaaaaaaaaa\x7f"), key(b"aaaaaaaaaaaaaa\xff")); assert_ne!(key(b"abcd"), key(b"abcd\0")); - assert_eq!(key(b"aaaaaaaaaaaaaa\xff").0 & LookupKey::TAG_MASK, 0); + let mut seen = std::collections::HashSet::new(); + for len in 1..=INLINE_KEY_BYTES { + for b in 0..=255u8 { + let word: Vec = (0..len).map(|i| b.wrapping_add(i as u8)).collect(); + assert!(seen.insert(key(&word).0), "collision at len={len} b={b}"); + } + } } /// One window shape per row: the needle in various lanes, an empty slot in @@ -716,19 +704,16 @@ mod tests { } /// Every inline length, with a different value in every byte position, so a - /// packing that drops, duplicates or misplaces a byte fails. The reference is - /// the construction the packing must be equivalent to: the bytes copied into - /// a zeroed array, the length written in the top byte. #[test] fn an_inline_key_is_the_words_bytes_with_the_length_on_top() { - for len in 0..=15usize { + for len in 0..=INLINE_KEY_BYTES { let word: Vec = (1..=len as u8).collect(); - let mut padded = [0u8; 16]; + let mut padded = [0u8; 8]; padded[..len].copy_from_slice(&word); - padded[15] = len as u8; + padded[7] = len as u8; assert_eq!( - LookupKey::new_inline(&word), - LookupKey(u128::from_le_bytes(padded)), + key_and_hash(&word).0, + u64::from_le_bytes(padded), "len={len}" ); } @@ -767,7 +752,7 @@ mod tests { assert_ne!(tag, WordCache::EMPTY, "pick a word with a nonzero tag"); cache.quick_lookup[index] = tag; cache.cached_words[index] = - WordCacheSlot::new_self_contained(LookupKey::new_inline(b"decoy"), [7].into_iter()); + WordCacheSlot::new_self_contained(LookupKey(key_and_hash(b"decoy").0), [7].into_iter()); assert_eq!(cache.lookup(b"beta").hit(), None); store(&mut cache, b"beta", &[2]); diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index 7f62ec0ec..35f76e369 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -4,16 +4,15 @@ use ahash::RandomState; use ptr_hash::{FastPtrHash, PtrHashParams, hash::NoHash}; use std::fmt; -type Mphf = FastPtrHash; - -// Fixed seeds so a given vocab always hashes identically (the hasher is also stored on the struct, -// so build and query are guaranteed consistent regardless). -const SEEDS: [u64; 4] = [ +static KEY_HASHER: RandomState = RandomState::with_seeds( 0x243F_6A88_85A3_08D3, 0x1319_8A2E_0370_7344, 0xA409_3822_299F_31D0, 0x082E_FA98_EC4E_6C89, -]; +); + +type Mphf = FastPtrHash; + /// Bit 31 of a stored id: the token provably encodes to itself, so a pretoken equal to it can be /// emitted without running the merge loop. See `PipelineBPE::prove_fold`. @@ -26,14 +25,100 @@ const FOLD_BIT: u32 = 1 << 31; /// The id half. 2^31 ids is far past any vocabulary. const VOCAB_ID_MASK: u32 = FOLD_BIT - 1; -#[derive(Clone, Copy, Debug)] +pub(crate) const INLINE_KEY_BYTES: usize = 7; + +#[inline(always)] +fn mix(z: u64) -> u64 { + let z = z.wrapping_mul(0x9E37_79B9_7F4A_7C15); + z ^ (z >> 29) +} + +/// +/// +/// +/// +/// # Safety +#[inline] +pub fn key_and_hash_readable(word: &[u8], readable: usize) -> (u64, u64) { + let len = word.len(); + if len > INLINE_KEY_BYTES || readable < 8 { + return key_and_hash(word); + } + // SAFETY: `readable >= 8` bytes exist from `word.as_ptr()`, and `len <= 7 < 8`. + let raw = unsafe { word.as_ptr().cast::().read_unaligned() }; + // SAFETY: `len <= INLINE_KEY_BYTES == 7`, and both tables have 8 entries. + let (mask, tag) = unsafe { (*KEY_MASK.get_unchecked(len), *LEN_TAG.get_unchecked(len)) }; + let key = (raw & mask) | tag; + debug_assert_eq!(key, key_and_hash(word).0, "masked load must match the stitched pack"); + (key, mix(key)) +} + +#[inline] +pub fn key_and_hash(word: &[u8]) -> (u64, u64) { + let len = word.len(); + if len > INLINE_KEY_BYTES { + let hash = KEY_HASHER.hash_one(word); + return (hash, hash); + } + let raw = if len >= 4 { + let head = u32::from_le_bytes(word[..4].try_into().unwrap()) as u64; + let tail = u32::from_le_bytes(word[len - 4..].try_into().unwrap()) as u64; + head | tail << (8 * (len - 4)) + } else if len >= 1 { + let first = word[0] as u64; + let middle = (word[len / 2] as u64) << (8 * (len / 2)); + let last = (word[len - 1] as u64) << (8 * (len - 1)); + first | middle | last + } else { + 0 + }; + let key = raw | (len as u64) << 56; + (key, mix(key)) +} + +/// +#[derive(Clone, Copy, Debug, Default)] +#[repr(C)] struct Entry { - start: u32, - len: u16, + digest: u32, /// The token id in the low 31 bits, [`FOLD_BIT`] in the top. id: u32, } +const _: () = assert!(size_of::() == 8); + +static KEY_MASK: [u64; 8] = [ + 0x0000_0000_0000_0000, + 0x0000_0000_0000_00FF, + 0x0000_0000_0000_FFFF, + 0x0000_0000_00FF_FFFF, + 0x0000_0000_FFFF_FFFF, + 0x0000_00FF_FFFF_FFFF, + 0x0000_FFFF_FFFF_FFFF, + 0x00FF_FFFF_FFFF_FFFF, +]; +static LEN_TAG: [u64; 8] = [ + 0 << 56, + 1 << 56, + 2 << 56, + 3 << 56, + 4 << 56, + 5 << 56, + 6 << 56, + 7 << 56, +]; + +#[inline(always)] +fn digest_of(hash: u64) -> u32 { + (hash >> 32) as u32 +} + +#[derive(Clone, Copy, Debug, Default)] +struct Span { + start: u32, + len: u16, +} + /// The BucketVocabStore optimizes for space and speed. We don't use a HashMap to prevent duplicating the /// keys. Instead, we just use an `id_to_slot` and `entries` table. When you query bytes, you hash /// on the fly and get an `index` into the `entries` table. When you query an `id`, you fetch in @@ -55,11 +140,10 @@ struct Entry { #[derive(Clone)] pub struct BucketVocabStore { mphf: Mphf, - hasher: RandomState, /// All token bytes, concatenated. Ordered by MPHF slot. bytes: Box<[u8]>, - /// `entries[slot]` -> (offset into `bytes`, length, id). Ordered by MPHF slot. entries: Box<[Entry]>, + spans: Box<[Span]>, /// `id_to_slot[token_id] -> entry_idx` -> index into entries as the entries are not really sorted. id_to_slot: Box<[u32]>, /// Number of real tokens. Cached at build so `len()` is O(1): `entries` is sized to the @@ -102,12 +186,9 @@ impl BucketVocabStore { pub fn build(tokens: Vec<(Vec, u32)>) -> Self { let n = tokens.len(); - let hasher = RandomState::with_seeds(SEEDS[0], SEEDS[1], SEEDS[2], SEEDS[3]); - - // 1. Pre-hash token bytes -> u64 keys using near perfect hash func let keys: Vec = tokens .iter() - .map(|(s, _)| hasher.hash_one(s.as_slice())) + .map(|(s, _)| key_and_hash(s.as_slice()).1) .collect(); // 2. A perfect hash needs distinct keys. Collisions are astronomically unlikely @@ -142,29 +223,24 @@ impl BucketVocabStore { let total: usize = tokens.iter().map(|(s, _)| s.len()).sum(); let max_id = tokens.iter().map(|(_, id)| *id).max().unwrap(); let mut bytes = Vec::with_capacity(total); - let mut entries = vec![ - Entry { - start: 0, - len: 0, - id: 0 - }; - n_slots - ]; + let mut entries = vec![Entry::default(); n_slots]; + let mut spans = vec![Span::default(); n_slots]; let mut id_to_slot = vec![u32::MAX; max_id as usize + 1]; for (s, id) in &tokens { assert!( s.len() <= u16::MAX as usize, "token longer than 65535 bytes" ); - assert!( - *id <= VOCAB_ID_MASK, - "token id {id} needs bit 31, which holds FOLD_BIT" - ); - let slot = mphf.index(&hasher.hash_one(s.as_slice())); + assert!(*id <= VOCAB_ID_MASK, "token id {id} needs bit 31, which holds FOLD_BIT"); + let (key, hash) = key_and_hash(s.as_slice()); + let slot = mphf.index(&hash); entries[slot] = Entry { + digest: digest_of(hash), + id: *id, + }; + spans[slot] = Span { start: bytes.len() as u32, len: s.len() as u16, - id: *id, }; id_to_slot[*id as usize] = slot as u32; bytes.extend_from_slice(s); @@ -172,9 +248,9 @@ impl BucketVocabStore { Self { mphf, - hasher, bytes: bytes.into_boxed_slice(), entries: entries.into_boxed_slice(), + spans: spans.into_boxed_slice(), id_to_slot: id_to_slot.into_boxed_slice(), n, } @@ -185,9 +261,9 @@ impl BucketVocabStore { let empty: [u64; 0] = []; Self { mphf: FastPtrHash::::new(&empty, PtrHashParams::default_fast()), - hasher: RandomState::new(), bytes: Box::new([]), entries: Box::new([]), + spans: Box::new([]), id_to_slot: Box::new([]), n: 0, } @@ -202,34 +278,46 @@ impl BucketVocabStore { if self.entries.is_empty() { return None; } - let slot = self.mphf.index(&self.hasher.hash_one(q)); - + let (key, hash) = key_and_hash(q); + let slot = self.mphf.index(&hash); let e = self.entries[slot]; - let (start, len) = (e.start as usize, e.len as usize); - // Byte equality: confirms `q` really is the token at this slot (perfect hashing only - // guarantees a valid slot for in-vocab keys; this rejects collisions and Out Of Vocab queries). - if len == q.len() && self.bytes[start..start + len] == *q { - Some(e.id & VOCAB_ID_MASK) - } else { - None - } + (e.digest == digest_of(hash)).then_some(e.id & VOCAB_ID_MASK) } /// The id for `q`, together with whether that entry may be folded. One probe and one entry /// load: the flag is a bit of the id the probe already read. #[inline] pub fn get_bytes_foldable(&self, q: &[u8]) -> Option<(u32, bool)> { + let (key, hash) = key_and_hash(q); + self.get_keyed_foldable(key, hash) + } + + #[inline(always)] + pub fn probe_slot(&self, hash: u64) -> usize { + self.mphf.index(&hash) + } + + #[inline(always)] + pub fn entry_at(&self, slot: usize) -> (u32, u32) { + let e = self.entries[slot]; + (e.digest, e.id) + } + + #[inline(always)] + pub fn resolve_foldable(hash: u64, entry: (u32, u32)) -> Option<(u32, bool)> { + let (edigest, eid) = entry; + (edigest == digest_of(hash)).then_some((eid & VOCAB_ID_MASK, eid & FOLD_BIT != 0)) + } + + #[inline] + pub fn get_keyed_foldable(&self, key: u64, hash: u64) -> Option<(u32, bool)> { + let _ = key; if self.entries.is_empty() { return None; } - let slot = self.mphf.index(&self.hasher.hash_one(q)); + let slot = self.mphf.index(&hash); let e = self.entries[slot]; - let (start, len) = (e.start as usize, e.len as usize); - if len == q.len() && self.bytes[start..start + len] == *q { - Some((e.id & VOCAB_ID_MASK, e.id & FOLD_BIT != 0)) - } else { - None - } + (e.digest == digest_of(hash)).then_some((e.id & VOCAB_ID_MASK, e.id & FOLD_BIT != 0)) } /// Records that this token folds to itself. Called once per entry at load, after the proof. @@ -253,9 +341,9 @@ impl BucketVocabStore { if slot == u32::MAX { return None; // id is within range but absent from the vocab } - let e = self.entries[slot as usize]; - let start = e.start as usize; - self.bytes.get(start..start + e.len as usize) + let sp = self.spans[slot as usize]; + let start = sp.start as usize; + self.bytes.get(start..start + sp.len as usize) } #[inline] @@ -285,9 +373,10 @@ impl BucketVocabStore { pub fn content(&self) -> Vec<(String, u32)> { self.entries .iter() - .filter(|e| e.len > 0) + .zip(self.spans.iter()) + .filter(|(_, sp)| sp.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|m| m.id & VOCAB_ID_MASK) + .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token(id).map(|token| (token, id))) .collect() } @@ -301,9 +390,10 @@ impl BucketVocabStore { pub fn byte_content(&self) -> Vec<(Vec, u32)> { self.entries .iter() - .filter(|e| e.len > 0) + .zip(self.spans.iter()) + .filter(|(_, sp)| sp.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|m| m.id & VOCAB_ID_MASK) + .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token_bytes(id).map(|token| (token.to_vec(), id))) .collect() } From bd6560b9e5271020c1453be4170319eca0b42534 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 18:45:07 +0900 Subject: [PATCH 5/7] perf(added-vocab): search the whole prefix, not just its first byte The single-bucket scan looked for `prefix[0]` and rejected each candidate by hand, restarting `memchr` on `&bytes[pos + 1..]` every time. One byte of a long needle is a poor filter on text that is dense in that byte and holds no match: `<|endoftext|>` over a corpus full of `<|xs0|>` stops at every `<` and dies at the third byte, and every restart pays memchr's prologue again. Measured over the same one-bucket vocabulary in one process, so the two arms share a build: on a 9.8%-`<` corpus that matches nothing, scanning the first byte cost 0.58 ns/B against 0.12 for `memmem`, which picks a *rare* byte of the needle instead -- 4.3x. Where candidates are already sparse it costs at most 0.015 ns/B (english 0.009 -> 0.024), so it trades a little on inputs where the scan is already free for a lot on the inputs where it is not. `nibble_mask_match` already avoids this for two or more buckets, where it is called the restart penalty; the one-bucket path never got it. `find_iter` yields the same positions in the same order, so the leftmost match is unchanged and `match_fast` still confirms the length sub-list. --- tokenizers/tk-encode/src/vocab/buckets.rs | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/tokenizers/tk-encode/src/vocab/buckets.rs b/tokenizers/tk-encode/src/vocab/buckets.rs index 9a78a8cec..ff46cf50c 100644 --- a/tokenizers/tk-encode/src/vocab/buckets.rs +++ b/tokenizers/tk-encode/src/vocab/buckets.rs @@ -333,14 +333,23 @@ impl Buckets { // needle = the bucket's shared first byte. Assumes a non-empty prefix // (false only if a lone 1-byte token is the sole holder of its first byte); store // the first byte explicitly if that case ever appears. - let needle = self.buckets[0].prefix[0]; - let mut search = 0usize; - while let Some(off) = memchr::memchr(needle, &bytes[search..]) { - let pos = search + off; + // Search for the bucket's whole shared prefix, not just its first byte. + // + // One byte of a long needle is a poor filter on text that is dense in that byte and + // holds no match: `<|endoftext|>` over a corpus full of `<|xs0|>` stops at every `<` + // and dies at the third byte. Restarting `memchr` on `&bytes[pos + 1..]` per + // candidate then pays its prologue again each time. Measured over the same + // one-bucket vocabulary in one process, that cost 0.58 ns/B on a 9.8%-`<` corpus + // that matches nothing, against 0.12 for `memmem`, which picks a *rare* byte of the + // needle instead. Where candidates are already sparse it costs at most 0.015 ns/B. + // + // `nibble_mask_match` already avoids this for two or more buckets and calls it the + // restart penalty. Same positions in the same order, so the leftmost match is + // unchanged and `match_fast` still confirms the length sub-list. + for pos in memchr::memmem::find_iter(bytes, &self.buckets[0].prefix) { if let Some((id, len)) = self.match_fast(bytes, pos, 0) { return Some((id, pos as u32, len)); } - search = pos + 1; } None } From c4e9b9a0fe596cbd3cb347d74823fde27e11cff1 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 20:00:15 +0900 Subject: [PATCH 6/7] perf(vocab): inline the key packing, keep the real hash out of line `key_and_hash` and `key_and_hash_readable` carried `#[inline]` and were left outlined anyway: a profile of warm english had 8.3% of self time inside `key_and_hash`, for what should be a masked load, an `or` and a multiply. LLVM weighs the whole body when it decides, and the body also held the ahash call for a word too long to pack, so the cheap path was paying for the expensive one's size. Split them: the packing is `#[inline(always)]`, the hash moves to `key_and_hash_long`, which is `#[inline(never)]` because it is already dominated by hashing rather than by call overhead. It is deliberately NOT `#[cold]` -- for CJK the long arm is the *common* one, and telling the predictor otherwise would cost more than the call. End to end this is flat (measured 1.0021x over 29 tokbench gpt2 cells, 12/29 faster -- a coin flip), and it is kept for what it revealed rather than for the speed: `key_and_hash_long` now shows up as its own line at 6.5% of warm english, which says that cost is ahash *work* on words longer than the 7-byte inline key, not a failure to inline. That makes widening the packed key the evidenced fix for it, since english pretokens run to 13 bytes. --- .../tk-encode/src/vocab/bucket_vocab_store.rs | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index 35f76e369..6e18122ce 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -38,7 +38,14 @@ fn mix(z: u64) -> u64 { /// /// /// # Safety -#[inline] +/// `inline(always)`, and the long-word arm is a separate `inline(never)` function, because plain +/// `#[inline]` left this outlined: a profile of warm english had 8.3% of self time inside +/// `key_and_hash`, for what should be a masked load, an `or` and a multiply. LLVM weighs the whole +/// body when it decides, and the body used to contain the ahash call for a word too long to pack, so +/// the cheap path paid for the expensive one's size. Splitting them lets the pack inline into the +/// span loop while the hash stays a call -- which is also the shape it wants, since the long arm is +/// already dominated by hashing rather than by call overhead. +#[inline(always)] pub fn key_and_hash_readable(word: &[u8], readable: usize) -> (u64, u64) { let len = word.len(); if len > INLINE_KEY_BYTES || readable < 8 { @@ -53,12 +60,21 @@ pub fn key_and_hash_readable(word: &[u8], readable: usize) -> (u64, u64) { (key, mix(key)) } -#[inline] +/// A word too long to pack into the key: hash it for real. Kept out of line so that the packing +/// path above and in [`key_and_hash`] can be inlined without dragging ahash in with them. Not +/// `#[cold]` on purpose -- for CJK this is the *common* arm, and telling the predictor otherwise +/// would cost more than the call. +#[inline(never)] +fn key_and_hash_long(word: &[u8]) -> (u64, u64) { + let hash = KEY_HASHER.hash_one(word); + (hash, hash) +} + +#[inline(always)] pub fn key_and_hash(word: &[u8]) -> (u64, u64) { let len = word.len(); if len > INLINE_KEY_BYTES { - let hash = KEY_HASHER.hash_one(word); - return (hash, hash); + return key_and_hash_long(word); } let raw = if len >= 4 { let head = u32::from_le_bytes(word[..4].try_into().unwrap()) as u64; From f54e490a44d6502401cdd0932eb9377875fe2317 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 20:00:15 +0900 Subject: [PATCH 7/7] perf(bpe): hold the write cursor and the cache across a chunk `tokenize_spans` spent about a quarter of encode time in itself, with none of the model in it. Two loop-invariants were being re-derived per span: - `output.as_mut_ptr()` had to be **reloaded from memory every span**, because the loop also calls `set_len`, `reserve` and `extend` on `output` and the optimiser cannot then assume the buffer stayed where it was. A raw `dst` cursor now rides the whole chunk and is refreshed only on the three paths that can actually move the buffer: the reserve, the spilled-ids `extend`, and the merge. The fast path -- a cache hit writing inline ids -- touches two registers and nothing else. - `word_cache.as_mut()` re-read an `Option` discriminant out of the scratch three times per span, for a table that is either present for the entire call or absent. Together this is gigatoken's `probe_emit_chunk` shape: loop-invariant cursors, refreshed in the slow path only. Measured, 29 tokbench gpt2 cells, 3 interleaved rounds, medians: the cursor is **1.0066x with 21/29 corpora faster**, and the `Option` hoist on its own is flat, kept because it is strictly less work. Warm-cache throughput moves much more than that (english 1.71 -> 1.58 ns/B, ~9%) but do not believe it: re-encoding one buffer makes every word a cache hit, so the scaffolding *is* the runtime there, while tokbench's disjoint slices push words into the fold and the merge and dilute it. The tokbench figure is the honest one. --- tokenizers/tk-encode/src/models/bpe/model.rs | 43 ++++++++++++++++---- 1 file changed, 34 insertions(+), 9 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 6d8395003..6b9505569 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -376,6 +376,23 @@ impl pipeline::Model for PipelineBPE { output.reserve(spans.len() + MAX_INLINE_IDS); let mut capacity = output.capacity(); let mut cursor = output.len(); + // A raw write cursor held across the whole chunk. `output.as_mut_ptr()` inside the loop had + // to be *reloaded from memory every span*: the loop also calls `set_len`, `reserve` and + // `extend` on `output`, so the optimiser cannot assume the buffer stayed where it was. That + // reload, the capacity test and the cursor arithmetic were a quarter of encode time in + // `tokenize_spans` itself, with nothing of the model in it. Now the fast path -- a cache hit + // writing inline ids -- touches only these two registers, and `dst` is refreshed solely on + // the paths that can actually move the buffer. This is the other half of gigatoken's + // `probe_emit_chunk`: loop-invariant cursors, refreshed only in the slow path. + // SAFETY: `cursor <= output.len() <= capacity`, so this is inside the allocation. + let mut dst = unsafe { output.as_mut_ptr().add(cursor) }; + + // Unwrapped once for the whole chunk. `word_cache` is an `Option` living in the + // scratch, so every `as_mut()` re-read its discriminant out of memory -- three times per + // span, on a table that is either there for the entire call or not at all. Holding + // `Option<&mut WordCache>` in a local keeps that test in a register, which is the cheap half + // of what gigatoken's `ProbeView` does by carrying the table base and mask across a chunk. + let mut cache_slot = word_cache.as_mut(); for span in spans { // SAFETY: the pre-tokenizer cuts on char boundaries, so a span is always a valid slice @@ -390,6 +407,9 @@ impl pipeline::Model for PipelineBPE { unsafe { output.set_len(cursor) }; output.reserve(spans.len() + MAX_INLINE_IDS); capacity = output.capacity(); + // `reserve` may have moved the buffer. + // SAFETY: `cursor` is what was written so far, so it is within the new allocation. + dst = unsafe { output.as_mut_ptr().add(cursor) }; } // The cache goes first. It is one direct-mapped load; the fold is an MPHF probe, which @@ -405,18 +425,17 @@ impl pipeline::Model for PipelineBPE { key_and_hash_readable(sequence.as_bytes(), chunk.len() - span.start as usize); let mut placement = None; - if let Some(cache) = word_cache.as_mut() { + if let Some(cache) = cache_slot.as_deref_mut() { // SAFETY: the capacity check above leaves `MAX_INLINE_IDS` slots past `cursor`, and let found = unsafe { - cache.probe_emit_keyed( - key, - hash, - output.as_mut_ptr().add(cursor).cast::(), - ) + cache.probe_emit_keyed(key, hash, dst.cast::()) }; match found { ProbeEmit::Wrote(n) => { cursor += n; + // SAFETY: the probe wrote `n <= MAX_INLINE_IDS` ids, which the capacity + // check above reserved room for. + dst = unsafe { dst.add(n) }; continue; } ProbeEmit::Hit(ids) => { @@ -425,6 +444,8 @@ impl pipeline::Model for PipelineBPE { output.extend(ids.iter().map(|&id| PipelineToken { id })); cursor = output.len(); capacity = output.capacity(); + // SAFETY: `cursor == output.len()`, inside the (possibly moved) allocation. + dst = unsafe { output.as_mut_ptr().add(cursor) }; continue; } ProbeEmit::Miss(at) => placement = Some(at), @@ -435,9 +456,11 @@ impl pipeline::Model for PipelineBPE { // probe, which beats running the merge engine for it. if let Some(id) = self.fold_id_keyed(key, hash) { // SAFETY: the check above leaves at least `MAX_INLINE_IDS >= 1` slots past `cursor`. - unsafe { output.as_mut_ptr().add(cursor).write(PipelineToken { id }) }; + unsafe { dst.write(PipelineToken { id }) }; cursor += 1; - if let Some(cache) = word_cache.as_mut() + // SAFETY: one id written, and the capacity check reserved MAX_INLINE_IDS >= 1. + dst = unsafe { dst.add(1) }; + if let Some(cache) = cache_slot.as_deref_mut() && let Some(at) = placement { cache.insert(at, std::iter::once(id)); @@ -453,13 +476,15 @@ impl pipeline::Model for PipelineBPE { output.extend(symbols.iter().map(|&symbol| PipelineToken { id: self.tables.unmap.at(symbol as usize), })); - if let Some(cache) = word_cache.as_mut() + if let Some(cache) = cache_slot.as_deref_mut() && let Some(at) = placement { cache.insert(at, output[start..].iter().map(|token| token.id)); } cursor = output.len(); capacity = output.capacity(); + // SAFETY: `cursor == output.len()`; `extend` above may have moved the buffer. + dst = unsafe { output.as_mut_ptr().add(cursor) }; } // SAFETY: `cursor` counts every token written, by the fast paths and the slow one alike. unsafe { output.set_len(cursor) };