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/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 ee5af5669..3512c635f 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -8,9 +8,10 @@ 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::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,19 +288,144 @@ 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(()) + } + /// 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, + 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 + // 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; + } + // 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) { + // 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(()) } @@ -298,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/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..43876fa4a 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>; } @@ -960,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(()) })?; @@ -1252,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; } @@ -1292,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()), 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 }