From d1ad772c90d7212c7ec59b59b6a856a5674d7ce5 Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 6 Aug 2026 16:59:53 +0900 Subject: [PATCH 1/3] 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/3] 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 ddce493b15156e9f5639e3a2ef3a10c1b34761ec Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 6 Aug 2026 21:28:39 +0900 Subject: [PATCH 3/3] perf(bpe): wire the word cache into the batched model path, fused #2306 gives the model a whole chunk of pre-tokens at a time by overriding `Model::tokenize_spans`. #2307 wires the word cache into `tokenize_pipeline`. Those are different functions: the pipeline calls the batched one, so landing both leaves the cache unreachable on the path that actually runs. This wires it into both. On the batched path the probe is fused into the emit. `WordCache::probe_emit` checks the home slot and, on an inline hit, stores all MAX_INLINE_IDS lanes straight at the caller's cursor, returning only the count. `lookup` hands back a `&[u32]` instead, which makes the caller re-read the slot to build a fat pointer and then copy a run whose length it learns at run time -- three trips over one 32-byte line that a single load already brought in. Anything that is not an inline home-slot hit falls back to the full window walk, reusing the placement so the word is never hashed twice. `lookup` and the walk are `#[inline]`. The output buffer is written through a cursor and closed with one `set_len`, so a hit costs no capacity check and no length store. It also reserves 2 ids per pre-token rather than 1: `spans.len()` is a *lower* bound (92% of English pre-tokens are one id, 98% at most two), so reserving it made the buffer grow -- and memcpy what it already held -- partway through most chunks. The fold stays in front of the cache: it answers a whole-vocabulary-entry word in one probe, which is cheaper than a cache probe, so folded words never enter the cache. They are already as cheap as a hit. Also drops `fold_by_flag`. Whether the config declared `ignore_merges` is a load -time question and does not belong in a per-pre-token branch: a declaring config asks every hit to fold, so every entry gets the bit; otherwise only the entries that prove they reduce to themselves earn it. `fold_id` is then one probe and one bit test with no policy left in it, and `ignore_merges` stops being a field. Fixes a sparse-id bug on the way: `prove_fold` bounded its walk by `vocab.len()`, the entry count, so any entry with an id above it was left unproven. Ids may be sparse, so the bound is the id space -- added as `BucketVocabStore::id_space`. No effect on gpt2, whose ids are dense, but the unified fold path above is wrong without it whenever a config leaves gaps. `PipelineToken` is asserted layout-identical to `u32`, since the fused probe writes ids through a pointer into the token buffer. 367 tests pass, 0 failures. --- .../tk-encode/src/models/bpe/legacy/model.rs | 2 +- tokenizers/tk-encode/src/models/bpe/model.rs | 167 ++++++++++++++---- tokenizers/tk-encode/src/utils/word_cache.rs | 74 +++++++- .../tk-encode/src/vocab/bucket_vocab_store.rs | 9 + 4 files changed, 214 insertions(+), 38 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/legacy/model.rs b/tokenizers/tk-encode/src/models/bpe/legacy/model.rs index 9e8b5adf0..20c0b9830 100644 --- a/tokenizers/tk-encode/src/models/bpe/legacy/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/legacy/model.rs @@ -302,7 +302,7 @@ pub struct BPE { /// Contains the mapping between Pairs and their (rank, new_id). pub merges: MergeMap, /// Contains the cache for optimizing the encoding step. - pub(super) cache: Option, + pub(crate) cache: Option, /// Dropout probability for merges. 0.0 = no dropout is the default. At 1.0, tokenization will /// perform no merges, so the result will just be characters. pub dropout: Option, diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 18b03f837..3512c635f 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -11,6 +11,7 @@ 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, MAX_INLINE_IDS, ProbeEmit, WordCache}; use crate::vocab::bucket_vocab_store::BucketVocabStore; const GATE_MULTI: u16 = 8; @@ -44,17 +45,20 @@ fn content_start(bytes: &[u8]) -> usize { } } +// `tokenize_spans` hands the word cache a `*mut u32` pointing into the `Vec` it is +// filling, so the probe can store ids straight at the cursor. That is only sound while a token is +// layout-identical to its id. +const _: () = assert!(size_of::() == size_of::()); +const _: () = assert!(align_of::() == align_of::()); + pub struct PipelineBPE { pub(super) atoms: Atoms, pub(super) tables: BpeTables, pub(super) affixes: Option, pub(super) vocab: BucketVocabStore, - ignore_merges: bool, - /// Whether the fold is decided per entry, by the bit each vocabulary entry carries in its id. - /// False when the config declares `ignore_merges`, which folds every hit unconditionally. - /// See [`PipelineBPE::prove_fold`]. - fold_by_flag: bool, byte_to_gate: [u16; 256], + /// Slots for the per-scratch word cache, from the config. `None` disables it. + cache_capacity: Option, } // A `PipelineBPE` holds exactly one `Atoms`, so `Chars`' 1 KB byte-fallback table costs nothing. @@ -83,8 +87,11 @@ impl PipelineBPE { fuse_unk, continuing_subword_prefix, end_of_word_suffix, + cache, .. } = model; + // A capacity of zero means "no cache"; anything else sizes the per-scratch table. + let cache_capacity = cache.map(|cache| cache.capacity).filter(|&c| c > 0); let prefix = continuing_subword_prefix.unwrap_or_default(); let suffix = end_of_word_suffix.unwrap_or_default(); if prefix.len() + 4 + suffix.len() > AFFIX_BUF { @@ -156,22 +163,26 @@ impl PipelineBPE { atoms, tables, affixes, - ignore_merges, - fold_by_flag: false, vocab, byte_to_gate: build_byte_to_gate(), + cache_capacity, }; - // A config that already declares `ignore_merges` folds every hit and needs no proof. - if !built.ignore_merges { - // Two phases because the proof runs the merge engine, which borrows `built`: work out - // the answers first, then set the bit on each entry that earned it. - let proven = built.prove_fold(); - for (id, foldable) in proven.iter().enumerate() { - if *foldable { - built.vocab.set_foldable(id as u32); - } + // Every entry carries a foldable bit, so the encode path is one probe and one bit test + // with no policy left in it. The policy is decided here, once: a config that declares + // `ignore_merges` asks for every hit to fold, so every entry gets the bit; otherwise only + // the entries that prove they reduce to themselves earn it. + // + // Two phases because the proof runs the merge engine, which borrows `built`: work out the + // answers first, then set the bit on each entry that earned it. + let proven = if ignore_merges { + vec![true; built.vocab.id_space()] + } else { + built.prove_fold() + }; + for (id, foldable) in proven.iter().enumerate() { + if *foldable { + built.vocab.set_foldable(id as u32); } - built.fold_by_flag = true; } Ok(built) } @@ -181,7 +192,9 @@ impl PipelineBPE { /// /// We replace the old "ignore_merges" with something that actually ignores whether or not the flag was set. fn prove_fold(&self) -> Vec { - let len = self.vocab.len(); + // The id space, not the entry count: ids may be sparse, and bounding the walk by + // `vocab.len()` would leave every entry above it unproven. + let len = self.vocab.id_space(); let mut proven = vec![false; len]; let mut symbols = Vec::with_capacity(64); let mut scratch = QueueScratch::default(); @@ -210,16 +223,10 @@ impl PipelineBPE { /// entry that may be folded. `None` sends the word to the merge engines. #[inline(always)] fn fold_id(&self, sequence: &str) -> Option { - if self.fold_by_flag { - // One probe; the proven bit is part of the id that probe already returned. - let (id, foldable) = self.vocab.get_bytes_foldable(sequence.as_bytes())?; - return foldable.then_some(id); - } - // The config declared `ignore_merges`: fold every hit, as it asks. - if self.ignore_merges { - return self.vocab.get_bytes(sequence.as_bytes()); - } - None + // 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())?; + foldable.then_some(id) } /// Converts a word to symbols and merges it. The gate, indexed by the word's first *content* @@ -261,6 +268,9 @@ pub struct BpeScratch { pub(crate) symbols: Vec, /// Entry arena and the two queue tiers. pub(crate) queue: QueueScratch, + /// Words already seen, so a repeat costs a probe instead of a merge. It lives in the scratch + /// so it outlives the encode call that fills it -- otherwise it would never see a word twice. + pub(crate) word_cache: Option, } impl pipeline::ModelScratch for BpeScratch {} @@ -278,18 +288,44 @@ impl pipeline::Model for PipelineBPE { return Ok(()); } + // The fold goes first: it answers a word that is itself a vocabulary entry in one probe, + // which is cheaper than a cache probe. Folded words therefore never enter the cache -- + // they are already as cheap as a hit. if let Some(id) = self.fold_id(sequence) { output.push(PipelineToken { id }); return Ok(()); } - let BpeScratch { symbols, queue } = scratch; + let BpeScratch { + symbols, + queue, + word_cache, + } = scratch; + // A word seen before costs a probe instead of a merge. + let insert_at = if let Some(cache) = word_cache.as_mut() { + match cache.lookup(sequence.as_bytes()) { + Lookup::Hit(ids) => { + output.extend(ids.iter().map(|&id| PipelineToken { id })); + return Ok(()); + } + Lookup::Miss(at) => Some(at), + } + } else { + None + }; + + 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 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) = insert_at + { + cache.insert(at, output[start..].iter().map(|token| token.id)); + } Ok(()) } @@ -307,10 +343,17 @@ impl pipeline::Model for PipelineBPE { 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()); + let BpeScratch { + symbols, + queue, + word_cache, + } = scratch; + // 92% of English pre-tokens are one id and 98% are at most two, so reserve for two apiece + // plus the probe's headroom. `spans.len()` alone is a *lower* bound, which would make the + // buffer grow -- and memcpy what it already holds -- partway through most chunks. + output.reserve(2 * 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 @@ -319,15 +362,70 @@ impl pipeline::Model for PipelineBPE { if sequence.is_empty() { continue; } + // One capacity check per word, covering both the fold's single write and the probe's + // `MAX_INLINE_IDS` lanes. After it, writing that many past `cursor` is in bounds. + 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 fold goes first -- see `tokenize_pipeline` for why. if let Some(id) = self.fold_id(sequence) { - output.push(PipelineToken { id }); + // 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; continue; } + let mut placement = None; + if let Some(cache) = word_cache.as_mut() { + // The probe writes the ids at the cursor itself, so a hit is one load of the slot + // and one unconditional store of its lanes -- the ids never become a slice and the + // line is never read twice. + // SAFETY: the check above leaves `MAX_INLINE_IDS` slots past `cursor`, and + // `PipelineToken` is a single `u32` (asserted below), so the cast is sound. + let found = unsafe { + cache.probe_emit( + sequence.as_bytes(), + output.as_mut_ptr().add(cursor).cast::(), + ) + }; + match found { + ProbeEmit::Wrote(n) => { + cursor += n; + continue; + } + // A hit the fast path could not serve: the probe already found the ids, so + // copy those rather than probing a second time. + 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; + } + ProbeEmit::Miss(at) => placement = Some(at), + } + } + // SAFETY: `cursor` counts what the fast paths wrote; the merge below uses `output` + // through its normal API, so its length has to be true again first. + unsafe { output.set_len(cursor) }; + let start = output.len(); self.merge_word(sequence, symbols, queue); 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 + { + cache.insert(at, output[start..].iter().map(|token| token.id)); + } + cursor = output.len(); + capacity = output.capacity(); } + // SAFETY: `cursor` counts every token written above. + unsafe { output.set_len(cursor) }; Ok(()) } @@ -335,6 +433,7 @@ impl pipeline::Model for PipelineBPE { Self::Scratch { symbols: Vec::with_capacity(64), queue: QueueScratch::default(), + word_cache: self.cache_capacity.map(WordCache::new), } } } diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index a92ac3079..a073bd699 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -75,6 +75,10 @@ static DISCRIMINANT_HASHER: RandomState = RandomState::with_seeds( 0x3f84_d5b5_b547_0917, ); +/// How many ids a [`WordCacheSlot`] holds inline before it has to spill. A probe writes this +/// many lanes unconditionally, so it is also the headroom [`WordCache::probe_emit`] needs. +pub const MAX_INLINE_IDS: usize = 3; + /// A table mapping words (`[u8]`) to the token ids they encode to (`[u32]`) pub struct WordCache { /// The cache slots as a contiguous table. See [`WordCacheSlot`] for more details. @@ -121,12 +125,64 @@ 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)) + } + + /// Probe and emit in one step: on an inline hit in the home slot the ids are written straight + /// to `dst` and the count returned, so nothing goes back to the slot and nothing becomes a + /// slice. + /// + /// This is the shape the hot path wants. [`Self::lookup`] hands back a `&[u32]`, which means + /// the caller re-reads the slot to build a fat pointer and then copies a run whose length it + /// only learns at run time -- three trips over one 32-byte line that a single load already + /// brought in. Here that line is read once, all [`MAX_INLINE_IDS`] lanes are stored + /// unconditionally, and the caller advances its cursor by the count: no branch on the length, + /// no second load, no slice. + /// + /// Falls back to the full window walk for anything else. The table is sized well above its + /// load, so a word's home slot is usually the one it was placed in and the walk is a few + /// percent of words. + /// + /// # Safety + /// `dst` must have room for [`MAX_INLINE_IDS`] `u32` writes. `word` must not be empty -- + /// an empty word keys to zero, which is also what an untouched slot holds. + #[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 placement = make_lookup_key(word, self.placement_mask); + // SAFETY: `index` is masked with `placement_mask` (`next_pow2 - 1`), and the table is + // `next_pow2 + WINDOW_SIZE` long, so the home slot is always in bounds. + let slot = unsafe { *self.cached_words.as_ptr().add(placement.index) }; + // An untouched slot holds `LookupKey(0)`, which no non-empty word can key to, so a key + // match here is a real hit -- the same 127-bit argument the window walk makes. + if slot.key == placement.key && !slot.is_spilled() { + // SAFETY: the caller guarantees room for `MAX_INLINE_IDS`. 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. + 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), + } + } + + /// The window walk, once a word has been keyed and placed. Split out of [`Self::lookup`] so + /// [`Self::probe_emit`] can fall back to it without hashing the word a second time. + #[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 +219,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 +320,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) { @@ -461,6 +517,18 @@ pub enum Lookup<'a> { Miss(InsertPlacement), } +/// What [`WordCache::probe_emit`] found. `Wrote` is the fast path: the ids are already at the +/// caller's cursor and only the count comes back. +pub enum ProbeEmit<'a> { + /// An inline hit in the home slot. [`MAX_INLINE_IDS`] lanes were written at `dst`; this many + /// of them are live. + Wrote(usize), + /// A hit the fast path could not serve -- a spilled entry, or one placed off its home slot. + /// The ids were found, so the caller copies these rather than probing again. + Hit(&'a [u32]), + Miss(InsertPlacement), +} + struct Window { window: [u8; WordCache::WINDOW_SIZE], offset: usize, diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index 788e4c641..5d9a099e4 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -266,6 +266,15 @@ impl BucketVocabStore { self.n } + /// One past the highest id this vocabulary can hold. + /// + /// Ids are not dense: a config may leave gaps, so [`Self::len`] counts entries and is *not* an + /// id bound. Anything that walks ids has to bound itself by this and skip the holes, which + /// [`Self::id_to_token_bytes`] reports as `None`. + pub fn id_space(&self) -> usize { + self.id_to_slot.len() + } + pub fn is_empty(&self) -> bool { self.n == 0 }