From 36874bd28e0c57b0fdbddc53e7333fc7d3927ab1 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 01:23:10 +0000 Subject: [PATCH 01/18] bench: measure encoded BPE cache misses by symbol count --- examples/encoded_merge_bench.rs | 116 ++++++++++++++++++++++++++++++++ 1 file changed, 116 insertions(+) create mode 100644 examples/encoded_merge_bench.rs diff --git a/examples/encoded_merge_bench.rs b/examples/encoded_merge_bench.rs new file mode 100644 index 0000000..9a2bb31 --- /dev/null +++ b/examples/encoded_merge_bench.rs @@ -0,0 +1,116 @@ +//! Benchmark the generic encoded-BPE cache-miss merger. +//! +//! Each input is a fresh ASCII pretoken-shaped string, so it is already in the +//! representation consumed by `Bpe::tokenize` and cannot hit either BPE cache. +//! The leading `c` also keeps the whole input from matching the `ab` vocabulary +//! token, forcing the encoded merge path. Use `--symbols` to select the exact +//! initial-symbol bucket, including the 32/33 crossover. + +use std::{collections::HashSet, env, hint::black_box, time::Instant}; + +use fastokens::models::bpe::Bpe; +use serde_json::json; + +const DEFAULT_ITERATIONS: usize = 8_192; +const WARMUP: usize = 1_024; +const ALPHABET: &[u8] = b"abde"; + +fn fixture() -> Bpe { + serde_json::from_value(json!({ + "vocab": { + "a": 0, + "b": 1, + "c": 2, + "d": 3, + "e": 4, + "ab": 5 + }, + "merges": ["a b"] + })) + .expect("benchmark BPE fixture must deserialize") +} + +fn inputs(symbols: usize, count: usize, state: &mut u64) -> Vec { + assert!( + symbols >= 2, + "a measured input must not be a one-token match" + ); + let mut seen = HashSet::with_capacity(count); + let mut result = Vec::with_capacity(count); + while result.len() < count { + let mut value = *state; + let mut input = String::with_capacity(symbols); + input.push('c'); + for _ in 1..symbols { + // A deterministic stream gives every invocation the same workload, + // while the set makes each call a cache miss in the BPE caches. + value ^= value << 13; + value ^= value >> 7; + value ^= value << 17; + input.push(ALPHABET[value as usize & (ALPHABET.len() - 1)] as char); + } + *state = value.wrapping_add(0x9e37_79b9_7f4a_7c15); + if seen.insert(input.clone()) { + result.push(input); + } + } + result +} + +fn main() { + let mut symbols = None; + let mut iterations = DEFAULT_ITERATIONS; + let mut args = env::args().skip(1); + while let Some(arg) = args.next() { + match arg.as_str() { + "--symbols" => { + symbols = Some( + args.next() + .expect("--symbols needs a value") + .parse() + .expect("invalid symbol count"), + ) + } + "--iterations" => { + iterations = args + .next() + .expect("--iterations needs a value") + .parse() + .expect("invalid iteration count") + } + other => panic!("unknown argument {other:?}"), + } + } + let symbols = symbols.expect("usage: encoded_merge_bench --symbols N [--iterations N]"); + assert!( + symbols <= 64, + "the benchmark only covers the short branch and its guard" + ); + assert!(iterations > 0, "iterations must be positive"); + + let bpe = fixture(); + let mut state = 0x243f_6a88_85a3_08d3u64 ^ symbols as u64; + let warmup = inputs(symbols, WARMUP, &mut state); + let measured = inputs(symbols, iterations, &mut state); + + for input in &warmup { + black_box(bpe.tokenize(input).expect("benchmark input must tokenize")); + } + + let start = Instant::now(); + let mut checksum = 0u64; + for input in &measured { + let ids = bpe.tokenize(input).expect("benchmark input must tokenize"); + for id in ids { + checksum = checksum.rotate_left(7) ^ u64::from(id); + } + } + let elapsed = start.elapsed(); + black_box(checksum); + + let ns_per_op = elapsed.as_secs_f64() * 1e9 / iterations as f64; + eprintln!( + "encoded generic cache misses: symbols={symbols}, iterations={iterations}, checksum={checksum}" + ); + println!(r#"{{"metric":"ns/op","value":{ns_per_op:.3}}}"#); +} From c190ee825103ace62264f423537768285ed69af8 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 01:29:54 +0000 Subject: [PATCH 02/18] bench: widen encoded miss corpus for stable buckets --- examples/encoded_merge_bench.rs | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/examples/encoded_merge_bench.rs b/examples/encoded_merge_bench.rs index 9a2bb31..672cd4a 100644 --- a/examples/encoded_merge_bench.rs +++ b/examples/encoded_merge_bench.rs @@ -13,7 +13,7 @@ use serde_json::json; const DEFAULT_ITERATIONS: usize = 8_192; const WARMUP: usize = 1_024; -const ALPHABET: &[u8] = b"abde"; +const ALPHABET: &[u8] = b"abcdefghijklmnop"; fn fixture() -> Bpe { serde_json::from_value(json!({ @@ -23,7 +23,18 @@ fn fixture() -> Bpe { "c": 2, "d": 3, "e": 4, - "ab": 5 + "f": 5, + "g": 6, + "h": 7, + "i": 8, + "j": 9, + "k": 10, + "l": 11, + "m": 12, + "n": 13, + "o": 14, + "p": 15, + "ab": 16 }, "merges": ["a b"] })) From 6fbd0231cee0b216eed4d44c76a0a7adc81311ee Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 01:38:16 +0000 Subject: [PATCH 03/18] bench: use merge-dense encoded miss buckets --- examples/encoded_merge_bench.rs | 60 ++++++++++++++++----------------- 1 file changed, 29 insertions(+), 31 deletions(-) diff --git a/examples/encoded_merge_bench.rs b/examples/encoded_merge_bench.rs index 672cd4a..4ae65c7 100644 --- a/examples/encoded_merge_bench.rs +++ b/examples/encoded_merge_bench.rs @@ -2,43 +2,41 @@ //! //! Each input is a fresh ASCII pretoken-shaped string, so it is already in the //! representation consumed by `Bpe::tokenize` and cannot hit either BPE cache. -//! The leading `c` also keeps the whole input from matching the `ab` vocabulary -//! token, forcing the encoded merge path. Use `--symbols` to select the exact +//! The leading `z` also keeps the whole input from matching a pair token, +//! forcing the encoded merge path. Use `--symbols` to select the exact //! initial-symbol bucket, including the 32/33 crossover. use std::{collections::HashSet, env, hint::black_box, time::Instant}; use fastokens::models::bpe::Bpe; -use serde_json::json; +use serde_json::{Map, Value, json}; const DEFAULT_ITERATIONS: usize = 8_192; const WARMUP: usize = 1_024; const ALPHABET: &[u8] = b"abcdefghijklmnop"; fn fixture() -> Bpe { - serde_json::from_value(json!({ - "vocab": { - "a": 0, - "b": 1, - "c": 2, - "d": 3, - "e": 4, - "f": 5, - "g": 6, - "h": 7, - "i": 8, - "j": 9, - "k": 10, - "l": 11, - "m": 12, - "n": 13, - "o": 14, - "p": 15, - "ab": 16 - }, - "merges": ["a b"] - })) - .expect("benchmark BPE fixture must deserialize") + let mut vocab = Map::new(); + for (id, &byte) in ALPHABET.iter().enumerate() { + vocab.insert((byte as char).to_string(), Value::from(id as u32)); + } + vocab.insert("z".into(), Value::from(ALPHABET.len() as u32)); + + // Every pair of body symbols is mergeable. The leading `z` is deliberately + // not part of this table, so it prevents the whole input from being a + // vocabulary match while leaving the measured body merge-heavy. + let mut merges = Vec::with_capacity(ALPHABET.len() * ALPHABET.len()); + for &left in ALPHABET { + for &right in ALPHABET { + let merged = format!("{}{}", left as char, right as char); + let id = vocab.len() as u32; + vocab.insert(merged, Value::from(id)); + merges.push(Value::String(format!("{} {}", left as char, right as char))); + } + } + + serde_json::from_value(json!({"vocab": vocab, "merges": merges})) + .expect("benchmark BPE fixture must deserialize") } fn inputs(symbols: usize, count: usize, state: &mut u64) -> Vec { @@ -51,7 +49,7 @@ fn inputs(symbols: usize, count: usize, state: &mut u64) -> Vec { while result.len() < count { let mut value = *state; let mut input = String::with_capacity(symbols); - input.push('c'); + input.push('z'); for _ in 1..symbols { // A deterministic stream gives every invocation the same workload, // while the set makes each call a cache miss in the BPE caches. @@ -101,16 +99,16 @@ fn main() { let bpe = fixture(); let mut state = 0x243f_6a88_85a3_08d3u64 ^ symbols as u64; - let warmup = inputs(symbols, WARMUP, &mut state); - let measured = inputs(symbols, iterations, &mut state); + let all_inputs = inputs(symbols, WARMUP + iterations, &mut state); + let (warmup, measured) = all_inputs.split_at(WARMUP); - for input in &warmup { + for input in warmup { black_box(bpe.tokenize(input).expect("benchmark input must tokenize")); } let start = Instant::now(); let mut checksum = 0u64; - for input in &measured { + for input in measured { let ids = bpe.tokenize(input).expect("benchmark input must tokenize"); for id in ids { checksum = checksum.rotate_left(7) ^ u64::from(id); From 848a71ffca07f6c7ce5610299a3d36fbbce93d8a Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 01:39:39 +0000 Subject: [PATCH 04/18] bench: document encoded merger scope --- examples/encoded_merge_bench.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/encoded_merge_bench.rs b/examples/encoded_merge_bench.rs index 4ae65c7..f71701e 100644 --- a/examples/encoded_merge_bench.rs +++ b/examples/encoded_merge_bench.rs @@ -2,6 +2,8 @@ //! //! Each input is a fresh ASCII pretoken-shaped string, so it is already in the //! representation consumed by `Bpe::tokenize` and cannot hit either BPE cache. +//! The generic non-fused tokenizer calls this model entry for each encoded +//! split; using it directly keeps pre-tokenization outside the timed operation. //! The leading `z` also keeps the whole input from matching a pair token, //! forcing the encoded merge path. Use `--symbols` to select the exact //! initial-symbol bucket, including the 32/33 crossover. From f6da42af4ca3205195118ef1e4ec412241098cce Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 19:11:46 +0000 Subject: [PATCH 05/18] bench: add process CPU-time measurement --- Cargo.toml | 1 + examples/encoded_merge_bench.rs | 50 +++++++++++++++++++++++++++++++-- 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index e15fe02..9431efb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -49,6 +49,7 @@ libc = "0.2" [dev-dependencies] anyhow = "1" clap = { version = "4.5.58", features = ["derive"] } +libc = "0.2" csv = "1.4.0" indicatif = "0.18.4" tokenizers = { version = "0.22.2", features = ["http"] } diff --git a/examples/encoded_merge_bench.rs b/examples/encoded_merge_bench.rs index f71701e..bfba3a4 100644 --- a/examples/encoded_merge_bench.rs +++ b/examples/encoded_merge_bench.rs @@ -6,9 +6,16 @@ //! split; using it directly keeps pre-tokenization outside the timed operation. //! The leading `z` also keeps the whole input from matching a pair token, //! forcing the encoded merge path. Use `--symbols` to select the exact -//! initial-symbol bucket, including the 32/33 crossover. +//! initial-symbol bucket, including the 32/33 crossover. Pass `--cpu-time` to +//! also report process CPU time for the measured loop on Unix. -use std::{collections::HashSet, env, hint::black_box, time::Instant}; +use std::{ + collections::HashSet, + env, + hint::black_box, + mem::MaybeUninit, + time::{Duration, Instant}, +}; use fastokens::models::bpe::Bpe; use serde_json::{Map, Value, json}; @@ -17,6 +24,33 @@ const DEFAULT_ITERATIONS: usize = 8_192; const WARMUP: usize = 1_024; const ALPHABET: &[u8] = b"abcdefghijklmnop"; +#[cfg(unix)] +fn process_cpu_time() -> Duration { + let mut usage = MaybeUninit::::uninit(); + let result = unsafe { libc::getrusage(libc::RUSAGE_SELF, usage.as_mut_ptr()) }; + assert_eq!(result, 0, "getrusage failed with status {result}"); + + let usage = unsafe { usage.assume_init() }; + let user = timeval_duration(usage.ru_utime); + let system = timeval_duration(usage.ru_stime); + user.checked_add(system) + .expect("process CPU time overflowed") +} + +#[cfg(unix)] +fn timeval_duration(timeval: libc::timeval) -> Duration { + let seconds = u64::try_from(timeval.tv_sec).expect("negative CPU time seconds"); + let micros = u64::try_from(timeval.tv_usec).expect("negative CPU time microseconds"); + Duration::from_secs(seconds) + .checked_add(Duration::from_micros(micros)) + .expect("CPU time overflowed") +} + +#[cfg(not(unix))] +fn process_cpu_time() -> Duration { + panic!("--cpu-time requires a Unix process CPU clock") +} + fn fixture() -> Bpe { let mut vocab = Map::new(); for (id, &byte) in ALPHABET.iter().enumerate() { @@ -71,6 +105,7 @@ fn inputs(symbols: usize, count: usize, state: &mut u64) -> Vec { fn main() { let mut symbols = None; let mut iterations = DEFAULT_ITERATIONS; + let mut measure_cpu_time = false; let mut args = env::args().skip(1); while let Some(arg) = args.next() { match arg.as_str() { @@ -89,6 +124,7 @@ fn main() { .parse() .expect("invalid iteration count") } + "--cpu-time" => measure_cpu_time = true, other => panic!("unknown argument {other:?}"), } } @@ -108,6 +144,7 @@ fn main() { black_box(bpe.tokenize(input).expect("benchmark input must tokenize")); } + let cpu_start = measure_cpu_time.then(process_cpu_time); let start = Instant::now(); let mut checksum = 0u64; for input in measured { @@ -117,6 +154,11 @@ fn main() { } } let elapsed = start.elapsed(); + let cpu_elapsed = cpu_start.map(|start| { + process_cpu_time() + .checked_sub(start) + .expect("process CPU clock moved backwards") + }); black_box(checksum); let ns_per_op = elapsed.as_secs_f64() * 1e9 / iterations as f64; @@ -124,4 +166,8 @@ fn main() { "encoded generic cache misses: symbols={symbols}, iterations={iterations}, checksum={checksum}" ); println!(r#"{{"metric":"ns/op","value":{ns_per_op:.3}}}"#); + if let Some(cpu_elapsed) = cpu_elapsed { + let cpu_ns_per_op = cpu_elapsed.as_secs_f64() * 1e9 / iterations as f64; + println!(r#"{{"metric":"cpu-ns/op","value":{cpu_ns_per_op:.3}}}"#); + } } From 161569e0bd05b78b3d4d6ad06ae48b0b8c34e78c Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 01:39:53 +0000 Subject: [PATCH 06/18] perf: merge short encoded BPE pretokens on the stack --- src/models/bpe.rs | 345 ++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 334 insertions(+), 11 deletions(-) diff --git a/src/models/bpe.rs b/src/models/bpe.rs index ea92829..2f73b22 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -631,6 +631,48 @@ impl PretokenCache { } } +#[inline(always)] +fn append_encoded_symbol( + id: u32, + small_ids: &mut [u32; SMALL_MERGE_MAX], + small_n: &mut usize, + heap_mode: &mut bool, + scratch: &mut MergeScratch, +) { + if !*heap_mode { + if *small_n < SMALL_MERGE_MAX { + small_ids[*small_n] = id; + *small_n += 1; + return; + } + + // The 33rd symbol crosses the stack branch's bound. Transfer the + // already validated prefix once, then keep collecting in the existing + // linked-list representation used by the heap merger. + for (i, &c) in small_ids[..*small_n].iter().enumerate() { + scratch.symbols.push(MergeSymbol { + c, + prev: if i == 0 { -1 } else { (i - 1) as i32 }, + next: -1, + }); + if i > 0 { + scratch.symbols[i - 1].next = i as i32; + } + } + *heap_mode = true; + } + + let i = scratch.symbols.len(); + scratch.symbols.push(MergeSymbol { + c: id, + prev: if i == 0 { -1 } else { (i - 1) as i32 }, + next: -1, + }); + if i > 0 { + scratch.symbols[i - 1].next = i as i32; + } +} + thread_local! { static TL_BPE_CACHE: RefCell = RefCell::new(FlatCache::new()); static TL_FUSED_CACHE: RefCell = RefCell::new(PretokenCache::new()); @@ -1382,11 +1424,97 @@ impl Bpe { Ok(()) } - /// Priority-queue BPE merge on already-encoded (ByteLevel) text. + /// BPE merge on already-encoded (ByteLevel) text. Short inputs collect + /// their validated initial symbols on the stack; longer inputs retain the + /// existing priority-queue merger. fn merge_all_encoded_into(&self, input: &str, out: &mut Vec) -> Result<()> { if input.is_empty() { return Ok(()); } + let long_input = if input.len() <= SMALL_MERGE_MAX { + false + } else if input.is_ascii() { + true + } else { + input.chars().nth(SMALL_MERGE_MAX).is_some() + }; + if long_input { + return self.merge_all_encoded_heap_into(input, out); + } + + TL_MERGE_SCRATCH.with(|s| { + let mut scratch = s.borrow_mut(); + scratch.symbols.clear(); + scratch.heap.clear(); + + let mut small_ids = [0u32; SMALL_MERGE_MAX]; + let mut small_n = 0usize; + let mut heap_mode = false; + + for ch in input.chars() { + let mut buf = [0u8; 4]; + let encoded = ch.encode_utf8(&mut buf); + let found = if ch.is_ascii() { + let id = self.single_char_token[ch as usize]; + (id != INVALID_TOKEN).then_some(id) + } else { + self.token_to_id.get(encoded).copied() + }; + if let Some(id) = found { + append_encoded_symbol( + id, + &mut small_ids, + &mut small_n, + &mut heap_mode, + &mut scratch, + ); + continue; + } + + if !self.byte_fallback { + return Err(format!("character {ch:?} not in vocabulary")); + } + + for &byte in encoded.as_bytes() { + let id = self.byte_fallback_token_ids[byte as usize]; + if id == INVALID_TOKEN { + return Err(format!( + "byte fallback token <0x{byte:02X}> not in vocabulary" + )); + } + append_encoded_symbol( + id, + &mut small_ids, + &mut small_n, + &mut heap_mode, + &mut scratch, + ); + } + } + + if !heap_mode { + if small_n == 1 { + out.push(small_ids[0]); + } else { + self.merge_small_encoded(&mut small_ids, small_n, out); + } + return Ok(()); + } + + let n = scratch.symbols.len(); + self.init_merge_heap(&mut scratch, n); + self.run_merge_loop(&mut scratch, out); + Ok(()) + }) + } + + /// Reference priority-queue BPE merge on already-encoded (ByteLevel) text. + /// It remains the fallback for long inputs and the correctness oracle for + /// the stack merger's tests. + fn merge_all_encoded_heap_into(&self, input: &str, out: &mut Vec) -> Result<()> { + if input.is_empty() { + return Ok(()); + } TL_MERGE_SCRATCH.with(|s| { let mut scratch = s.borrow_mut(); @@ -1396,12 +1524,12 @@ impl Bpe { let mut n = 0usize; for ch in input.chars() { let mut buf = [0u8; 4]; - let s = ch.encode_utf8(&mut buf); + let encoded = ch.encode_utf8(&mut buf); let found = if ch.is_ascii() { let id = self.single_char_token[ch as usize]; (id != INVALID_TOKEN).then_some(id) } else { - self.token_to_id.get(s).copied() + self.token_to_id.get(encoded).copied() }; if let Some(id) = found { scratch.symbols.push(MergeSymbol { @@ -1420,7 +1548,7 @@ impl Bpe { return Err(format!("character {ch:?} not in vocabulary")); } - for &byte in s.as_bytes() { + for &byte in encoded.as_bytes() { let id = self.byte_fallback_token_ids[byte as usize]; if id == INVALID_TOKEN { return Err(format!( @@ -1450,19 +1578,83 @@ impl Bpe { }) } - /// Linear-scan BPE merge for short pretokens (`n <= SMALL_MERGE_MAX` - /// initial symbols). Avoids the `BinaryHeap` entirely: a stack-resident - /// doubly-linked list plus a per-position rank array, find-min by a short - /// scan over stack `u32`s, merge (O(1) pointer update), then refresh only - /// the two neighbor pairs. At these sizes this beats the heap's - /// sift/stale-entry traffic and does zero heap allocation. + /// Linear-scan BPE merge for short encoded pretokens (`n <= + /// SMALL_MERGE_MAX` initial symbols). Avoids the `BinaryHeap` entirely: a + /// stack-resident doubly-linked list plus a per-position rank array, + /// find-min by a short scan over stack `u32`s, merge (O(1) pointer update), + /// then refresh only the two neighbor pairs. /// /// Produces the identical token sequence as [`Self::run_merge_loop`]: both /// process the globally lowest-`(rank, pos)` active pair each step (the /// heap's `MergeEntry` key is `(rank << 32) | pos`; the scan's strict `<` /// keeps the leftmost/lowest-`pos` position on ties). Enforced by the /// `merge_small_matches_heap` differential test. `ids[..n]` are the - /// per-byte initial token ids. + /// initial encoded-symbol token ids. + fn merge_small_encoded(&self, ids: &mut [u32; SMALL_MERGE_MAX], n: usize, out: &mut Vec) { + let mut next = [0u8; SMALL_MERGE_MAX]; + let mut prev = [0u8; SMALL_MERGE_MAX]; + let mut ranks = [u32::MAX; SMALL_MERGE_MAX]; + let mut new_ids = [0u32; SMALL_MERGE_MAX]; + for i in 0..n { + next[i] = (i + 1) as u8; + prev[i] = (i as u8).wrapping_sub(1); // prev[0] = 255 (>= n): sentinel + } + for i in 0..n - 1 { + if let Some((rank, new_id)) = self.merge_adj.get(ids[i], ids[i + 1]) { + ranks[i] = rank; + new_ids[i] = new_id; + } + } + loop { + let mut best = u32::MAX; + let mut best_i = 0usize; + for (i, &rank) in ranks[..n - 1].iter().enumerate() { + if rank < best { + best = rank; + best_i = i; + } + } + if best == u32::MAX { + break; + } + let i = best_i; + ids[i] = new_ids[i]; + let dead = next[i] as usize; + let new_right = next[dead] as usize; + next[i] = new_right as u8; + ranks[dead] = u32::MAX; + if new_right < n { + prev[new_right] = i as u8; + match self.merge_adj.get(ids[i], ids[new_right]) { + Some((rank, new_id)) => { + ranks[i] = rank; + new_ids[i] = new_id; + } + None => ranks[i] = u32::MAX, + } + } else { + ranks[i] = u32::MAX; + } + let left = prev[i] as usize; + if left < n { + match self.merge_adj.get(ids[left], ids[i]) { + Some((rank, new_id)) => { + ranks[left] = rank; + new_ids[left] = new_id; + } + None => ranks[left] = u32::MAX, + } + } + } + let mut i = 0usize; + while i < n { + out.push(ids[i]); + i = next[i] as usize; + } + } + + /// Linear-scan BPE merge for short raw pretokens (`n <= SMALL_MERGE_MAX` + /// bytes). `ids[..n]` are the per-byte initial token ids. fn merge_small_raw( &self, bytes: &[u8], @@ -2135,6 +2327,137 @@ mod tests { } } + fn assert_encoded_matches_heap(bpe: &Bpe, input: &str) { + let mut optimized = vec![0xdead_beefu32, 0xcafe_babe]; + let mut heap = optimized.clone(); + let optimized_result = bpe.merge_all_encoded_into(input, &mut optimized); + let heap_result = bpe.merge_all_encoded_heap_into(input, &mut heap); + assert_eq!( + optimized_result, heap_result, + "result mismatch for {input:?}" + ); + assert_eq!(optimized, heap, "output mismatch for {input:?}"); + } + + #[test] + fn encoded_small_matches_heap_for_all_symbol_counts() { + let bpe = test_bpe(); + let alphabet = [b'a', b'b', b'c', b'd']; + let mut state = 0x9e37_79b9_7f4a_7c15u64; + + for len in 1..=SMALL_MERGE_MAX + 1 { + for _ in 0..128 { + let mut input = String::with_capacity(len); + for _ in 0..len { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + input.push(alphabet[state as usize & 3] as char); + } + assert_encoded_matches_heap(&bpe, &input); + } + } + } + + #[test] + fn encoded_small_preserves_leftmost_equal_rank_ties() { + let vocab: Vocab = [("a", 0), ("b", 1), ("c", 2), ("ab", 3), ("bc", 4)] + .into_iter() + .map(|(s, id)| (s.to_string(), id)) + .collect(); + let merge_map = [((0, 1), (7, 3)), ((1, 2), (7, 4))].into_iter().collect(); + let bpe = Bpe::new(&vocab, merge_map).unwrap(); + + assert_encoded_matches_heap(&bpe, "abc"); + let mut out = Vec::new(); + bpe.merge_all_encoded_into("abc", &mut out).unwrap(); + assert_eq!(out, vec![3, 2]); + } + + #[test] + fn encoded_small_handles_bytelevel_chars_and_combining_marks() { + let left = BYTE_TO_CHAR[0]; + let right = BYTE_TO_CHAR[1]; + let left_right = format!("{left}{right}"); + let vocab: Vocab = [ + (left.to_string(), 0), + (right.to_string(), 1), + (left_right, 2), + ] + .into_iter() + .map(|(s, id)| (s, id)) + .collect(); + let merge_map = [((0, 1), (0, 2))].into_iter().collect(); + let bytelevel_bpe = Bpe::new(&vocab, merge_map).unwrap(); + assert_encoded_matches_heap(&bytelevel_bpe, &format!("{left}{right}")); + assert_encoded_matches_heap(&bytelevel_bpe, &left.to_string().repeat(32)); + assert_encoded_matches_heap(&bytelevel_bpe, &left.to_string().repeat(33)); + + let combining = '\u{301}'; + let combined = format!("e{combining}"); + let vocab: Vocab = [("e".into(), 0), (combining.to_string(), 1), (combined, 2)] + .into_iter() + .collect(); + let merge_map = [((0, 1), (0, 2))].into_iter().collect(); + let combining_bpe = Bpe::new(&vocab, merge_map).unwrap(); + assert_encoded_matches_heap(&combining_bpe, &format!("e{combining}")); + } + + fn fallback_bpe(include_second_byte: bool) -> Bpe { + let mut vocab: Vocab = [("a".into(), 0), ("<0xC3>".into(), 1)] + .into_iter() + .collect(); + if include_second_byte { + vocab.insert("<0xA9>".into(), 2); + vocab.insert("<0xC3><0xA9>".into(), 3); + } + let merge_map = if include_second_byte { + [((1, 2), (0, 3))].into_iter().collect() + } else { + ParsedMergeMap::new() + }; + let mut bpe = Bpe::new(&vocab, merge_map).unwrap(); + bpe.byte_fallback = true; + bpe + } + + #[test] + fn encoded_small_handles_byte_fallback_expansion_and_boundary() { + let bpe = fallback_bpe(true); + assert_encoded_matches_heap(&bpe, "é"); + assert_encoded_matches_heap(&bpe, &format!("{}é", "a".repeat(30))); + assert_encoded_matches_heap(&bpe, &format!("{}é", "a".repeat(31))); + } + + #[test] + fn encoded_errors_are_equal_and_leave_prefilled_output_untouched() { + let bpe = fallback_bpe(false); + let input = "aé"; + let mut optimized = vec![17, 19]; + let mut heap = optimized.clone(); + let optimized_result = bpe.merge_all_encoded_into(input, &mut optimized); + let heap_result = bpe.merge_all_encoded_heap_into(input, &mut heap); + assert_eq!(optimized_result, heap_result); + assert_eq!( + optimized_result.unwrap_err(), + "byte fallback token <0xA9> not in vocabulary" + ); + assert_eq!(optimized, vec![17, 19]); + assert_eq!(heap, optimized); + + let mut optimized = vec![23]; + let mut heap = optimized.clone(); + let optimized_result = bpe.merge_all_encoded_into("?", &mut optimized); + let heap_result = bpe.merge_all_encoded_heap_into("?", &mut heap); + assert_eq!(optimized_result, heap_result); + assert_eq!( + optimized_result.unwrap_err(), + "byte fallback token <0x3F> not in vocabulary" + ); + assert_eq!(optimized, vec![23]); + assert_eq!(heap, optimized); + } + #[test] fn deserialize_from_json() { let json = serde_json::json!({ From 6bc077c3ccfdc6d72a94c9f1bc55fbabdef1aa49 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 01:41:27 +0000 Subject: [PATCH 07/18] perf: keep long encoded fallback on the heap path --- src/models/bpe.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/models/bpe.rs b/src/models/bpe.rs index 2f73b22..67a1335 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -1431,13 +1431,12 @@ impl Bpe { if input.is_empty() { return Ok(()); } - let long_input = if input.len() <= SMALL_MERGE_MAX { - false - } else if input.is_ascii() { - true - } else { - input.chars().nth(SMALL_MERGE_MAX).is_some() - }; + // A 33rd ASCII byte takes the conservative heap branch without a + // second scan. When that byte is part of a multi-byte character, count + // chars so a short ByteLevel pretoken can still use the stack branch. + let long_input = input.len() > SMALL_MERGE_MAX + && (input.as_bytes()[SMALL_MERGE_MAX].is_ascii() + || input.chars().nth(SMALL_MERGE_MAX).is_some()); if long_input { return self.merge_all_encoded_heap_into(input, out); } From 91a997ac12151cf587b217948cc52ad478e88bee Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 02:26:12 +0000 Subject: [PATCH 08/18] perf: isolate short encoded merger dispatch --- src/models/bpe.rs | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/src/models/bpe.rs b/src/models/bpe.rs index 67a1335..174232a 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -1424,23 +1424,32 @@ impl Bpe { Ok(()) } - /// BPE merge on already-encoded (ByteLevel) text. Short inputs collect - /// their validated initial symbols on the stack; longer inputs retain the - /// existing priority-queue merger. + #[inline(always)] + fn encoded_small_candidate(input: &str) -> bool { + input.len() <= SMALL_MERGE_MAX + || (!input.as_bytes()[SMALL_MERGE_MAX].is_ascii() + && input.chars().nth(SMALL_MERGE_MAX).is_none()) + } + + /// BPE merge on already-encoded (ByteLevel) text. Dispatches short + /// candidates to the stack merger and retains the priority-queue merger for + /// the other inputs. + #[inline(always)] fn merge_all_encoded_into(&self, input: &str, out: &mut Vec) -> Result<()> { if input.is_empty() { return Ok(()); } - // A 33rd ASCII byte takes the conservative heap branch without a - // second scan. When that byte is part of a multi-byte character, count - // chars so a short ByteLevel pretoken can still use the stack branch. - let long_input = input.len() > SMALL_MERGE_MAX - && (input.as_bytes()[SMALL_MERGE_MAX].is_ascii() - || input.chars().nth(SMALL_MERGE_MAX).is_some()); - if long_input { - return self.merge_all_encoded_heap_into(input, out); + if Self::encoded_small_candidate(input) { + self.merge_all_encoded_small_into(input, out) + } else { + self.merge_all_encoded_heap_into(input, out) } + } + /// Stack merger for encoded inputs whose initial-symbol count can fit in + /// the bounded representation. The collector promotes byte-fallback + /// expansions that cross the bound to the existing heap representation. + fn merge_all_encoded_small_into(&self, input: &str, out: &mut Vec) -> Result<()> { TL_MERGE_SCRATCH.with(|s| { let mut scratch = s.borrow_mut(); scratch.symbols.clear(); From 6eaeedb8a959e91a32e19d9eae52fa98a1768bd8 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 02:44:25 +0000 Subject: [PATCH 09/18] perf: use byte length for short encoded dispatch --- src/models/bpe.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/models/bpe.rs b/src/models/bpe.rs index 174232a..a3ecd99 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -1427,8 +1427,6 @@ impl Bpe { #[inline(always)] fn encoded_small_candidate(input: &str) -> bool { input.len() <= SMALL_MERGE_MAX - || (!input.as_bytes()[SMALL_MERGE_MAX].is_ascii() - && input.chars().nth(SMALL_MERGE_MAX).is_none()) } /// BPE merge on already-encoded (ByteLevel) text. Dispatches short From 6a9c87849a54fea881ad6e23631a31f28e78996f Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 11:55:32 +0000 Subject: [PATCH 10/18] bench: make encoded miss corpus finite --- examples/encoded_merge_bench.rs | 99 ++++++++++++++++++++++++--------- 1 file changed, 74 insertions(+), 25 deletions(-) diff --git a/examples/encoded_merge_bench.rs b/examples/encoded_merge_bench.rs index bfba3a4..fd08552 100644 --- a/examples/encoded_merge_bench.rs +++ b/examples/encoded_merge_bench.rs @@ -1,6 +1,6 @@ //! Benchmark the generic encoded-BPE cache-miss merger. //! -//! Each input is a fresh ASCII pretoken-shaped string, so it is already in the +//! Each input is a fresh encoded pretoken-shaped string, so it is already in the //! representation consumed by `Bpe::tokenize` and cannot hit either BPE cache. //! The generic non-fused tokenizer calls this model entry for each encoded //! split; using it directly keeps pre-tokenization outside the timed operation. @@ -10,7 +10,6 @@ //! also report process CPU time for the measured loop on Unix. use std::{ - collections::HashSet, env, hint::black_box, mem::MaybeUninit, @@ -23,6 +22,12 @@ use serde_json::{Map, Value, json}; const DEFAULT_ITERATIONS: usize = 8_192; const WARMUP: usize = 1_024; const ALPHABET: &[u8] = b"abcdefghijklmnop"; +const ALIAS_BASE: u32 = 0x1000; +const ALIAS_COUNT: usize = 16_384; + +fn alias_char(index: usize) -> char { + char::from_u32(ALIAS_BASE + index as u32).expect("benchmark alias must be a Unicode scalar") +} #[cfg(unix)] fn process_cpu_time() -> Duration { @@ -54,13 +59,13 @@ fn process_cpu_time() -> Duration { fn fixture() -> Bpe { let mut vocab = Map::new(); for (id, &byte) in ALPHABET.iter().enumerate() { - vocab.insert((byte as char).to_string(), Value::from(id as u32)); + vocab.insert((byte as char).to_string(), Value::from((id + 1) as u32)); } - vocab.insert("z".into(), Value::from(ALPHABET.len() as u32)); + vocab.insert("z".into(), Value::from(0u32)); - // Every pair of body symbols is mergeable. The leading `z` is deliberately - // not part of this table, so it prevents the whole input from being a - // vocabulary match while leaving the measured body merge-heavy. + // Every pair of body symbols is mergeable. The leading `z` is not part of + // this table, so it prevents the whole input from matching a vocabulary + // token while leaving the measured body merge-heavy. let mut merges = Vec::with_capacity(ALPHABET.len() * ALPHABET.len()); for &left in ALPHABET { for &right in ALPHABET { @@ -71,33 +76,51 @@ fn fixture() -> Bpe { } } + // The aliases give the cache-miss corpus a large finite key space without + // changing the 16-symbol merge topology. Each alias maps to one of the + // dense body IDs, but no alias is itself a representative vocabulary token, + // so a generated input cannot take the whole-token fast path. + for index in 0..ALIAS_COUNT { + let id = 1 + (index % ALPHABET.len()) as u32; + vocab.insert(alias_char(index).to_string(), Value::from(id)); + } + serde_json::from_value(json!({"vocab": vocab, "merges": merges})) .expect("benchmark BPE fixture must deserialize") } -fn inputs(symbols: usize, count: usize, state: &mut u64) -> Vec { +fn inputs(symbols: usize, count: usize) -> Vec { assert!( symbols >= 2, "a measured input must not be a one-token match" ); - let mut seen = HashSet::with_capacity(count); + let body_len = symbols - 1; + let capacity = (0..body_len) + .try_fold(1usize, |capacity, _| capacity.checked_mul(ALIAS_COUNT)) + .unwrap_or(usize::MAX); + assert!( + count <= capacity, + "requested {count} inputs but --symbols {symbols} has capacity {capacity}" + ); + + // Decode each ordinal in base ALIAS_COUNT. The per-position affine map is + // bijective because ALIAS_COUNT is a power of two and 5 is odd, so this + // produces exactly `count` distinct strings without retrying a HashSet. let mut result = Vec::with_capacity(count); - while result.len() < count { - let mut value = *state; - let mut input = String::with_capacity(symbols); + for ordinal in 0..count { + let mut value = ordinal; + let mut input = String::with_capacity(1 + body_len * 3); input.push('z'); - for _ in 1..symbols { - // A deterministic stream gives every invocation the same workload, - // while the set makes each call a cache miss in the BPE caches. - value ^= value << 13; - value ^= value >> 7; - value ^= value << 17; - input.push(ALPHABET[value as usize & (ALPHABET.len() - 1)] as char); - } - *state = value.wrapping_add(0x9e37_79b9_7f4a_7c15); - if seen.insert(input.clone()) { - result.push(input); + for position in 0..body_len { + let digit = value % ALIAS_COUNT; + value /= ALIAS_COUNT; + let alias = digit + .wrapping_mul(5) + .wrapping_add(position.wrapping_mul(257)) + & (ALIAS_COUNT - 1); + input.push(alias_char(alias)); } + result.push(input); } result } @@ -136,8 +159,10 @@ fn main() { assert!(iterations > 0, "iterations must be positive"); let bpe = fixture(); - let mut state = 0x243f_6a88_85a3_08d3u64 ^ symbols as u64; - let all_inputs = inputs(symbols, WARMUP + iterations, &mut state); + let total_inputs = WARMUP + .checked_add(iterations) + .expect("warmup plus iterations overflowed"); + let all_inputs = inputs(symbols, total_inputs); let (warmup, measured) = all_inputs.split_at(WARMUP); for input in warmup { @@ -171,3 +196,27 @@ fn main() { println!(r#"{{"metric":"cpu-ns/op","value":{cpu_ns_per_op:.3}}}"#); } } + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use super::*; + + #[test] + fn corpus_is_finite_and_unique_at_small_and_boundary_sizes() { + for symbols in [2, 3, 4, 16, 32, 33] { + let inputs = inputs(symbols, 256); + assert_eq!(inputs.len(), 256); + assert!(inputs.iter().all(|input| input.chars().count() == symbols)); + let unique: HashSet<_> = inputs.iter().collect(); + assert_eq!(unique.len(), inputs.len()); + } + } + + #[test] + #[should_panic(expected = "has capacity 16384")] + fn corpus_rejects_requests_larger_than_the_smallest_domain() { + let _ = inputs(2, ALIAS_COUNT + 1); + } +} From a82768ce9609bf510d3ed51ca045c7f4f3981017 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 11:55:56 +0000 Subject: [PATCH 11/18] perf: repair encoded short merger dispatch --- src/models/bpe.rs | 357 ++++++++++++++++++++++++++-------------------- 1 file changed, 202 insertions(+), 155 deletions(-) diff --git a/src/models/bpe.rs b/src/models/bpe.rs index a3ecd99..9223a77 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -1424,99 +1424,107 @@ impl Bpe { Ok(()) } - #[inline(always)] - fn encoded_small_candidate(input: &str) -> bool { - input.len() <= SMALL_MERGE_MAX - } - - /// BPE merge on already-encoded (ByteLevel) text. Dispatches short - /// candidates to the stack merger and retains the priority-queue merger for - /// the other inputs. + /// BPE merge on already-encoded (ByteLevel) text. The collector counts + /// emitted initial symbols, not UTF-8 bytes: it uses the stack merger while + /// the count fits and promotes to the existing priority-queue merger on the + /// first symbol above the bound. #[inline(always)] fn merge_all_encoded_into(&self, input: &str, out: &mut Vec) -> Result<()> { if input.is_empty() { return Ok(()); } - if Self::encoded_small_candidate(input) { - self.merge_all_encoded_small_into(input, out) - } else { - self.merge_all_encoded_heap_into(input, out) - } - } - /// Stack merger for encoded inputs whose initial-symbol count can fit in - /// the bounded representation. The collector promotes byte-fallback - /// expansions that cross the bound to the existing heap representation. - fn merge_all_encoded_small_into(&self, input: &str, out: &mut Vec) -> Result<()> { TL_MERGE_SCRATCH.with(|s| { let mut scratch = s.borrow_mut(); scratch.symbols.clear(); scratch.heap.clear(); let mut small_ids = [0u32; SMALL_MERGE_MAX]; - let mut small_n = 0usize; - let mut heap_mode = false; - - for ch in input.chars() { - let mut buf = [0u8; 4]; - let encoded = ch.encode_utf8(&mut buf); - let found = if ch.is_ascii() { - let id = self.single_char_token[ch as usize]; - (id != INVALID_TOKEN).then_some(id) - } else { - self.token_to_id.get(encoded).copied() - }; - if let Some(id) = found { - append_encoded_symbol( - id, - &mut small_ids, - &mut small_n, - &mut heap_mode, - &mut scratch, - ); - continue; + match self.collect_encoded_symbols(input, &mut small_ids, &mut scratch)? { + Some(n) => { + if n == 1 { + out.push(small_ids[0]); + } else { + self.merge_small_encoded(&mut small_ids, n, out); + } } - - if !self.byte_fallback { - return Err(format!("character {ch:?} not in vocabulary")); + None => { + let n = scratch.symbols.len(); + self.init_merge_heap(&mut scratch, n); + self.run_merge_loop(&mut scratch, out); } + } + Ok(()) + }) + } - for &byte in encoded.as_bytes() { - let id = self.byte_fallback_token_ids[byte as usize]; - if id == INVALID_TOKEN { - return Err(format!( - "byte fallback token <0x{byte:02X}> not in vocabulary" - )); - } - append_encoded_symbol( - id, - &mut small_ids, - &mut small_n, - &mut heap_mode, - &mut scratch, - ); - } + /// Collect encoded characters into the bounded stack representation. A + /// `Some` result is the exact initial-symbol count that fits on the stack; + /// `None` means the 33rd emitted symbol promoted the validated prefix to + /// the heap representation. + fn collect_encoded_symbols( + &self, + input: &str, + small_ids: &mut [u32; SMALL_MERGE_MAX], + scratch: &mut MergeScratch, + ) -> Result> { + let mut small_n = 0usize; + let mut heap_mode = false; + + for ch in input.chars() { + let mut buf = [0u8; 4]; + let encoded = ch.encode_utf8(&mut buf); + let found = if ch.is_ascii() { + let id = self.single_char_token[ch as usize]; + (id != INVALID_TOKEN).then_some(id) + } else { + self.token_to_id.get(encoded).copied() + }; + if let Some(id) = found { + append_encoded_symbol(id, small_ids, &mut small_n, &mut heap_mode, scratch); + continue; } - if !heap_mode { - if small_n == 1 { - out.push(small_ids[0]); - } else { - self.merge_small_encoded(&mut small_ids, small_n, out); + if !self.byte_fallback { + return Err(format!("character {ch:?} not in vocabulary")); + } + + for &byte in encoded.as_bytes() { + let id = self.byte_fallback_token_ids[byte as usize]; + if id == INVALID_TOKEN { + return Err(format!( + "byte fallback token <0x{byte:02X}> not in vocabulary" + )); } - return Ok(()); + append_encoded_symbol(id, small_ids, &mut small_n, &mut heap_mode, scratch); } + } - let n = scratch.symbols.len(); - self.init_merge_heap(&mut scratch, n); - self.run_merge_loop(&mut scratch, out); - Ok(()) + if heap_mode { + Ok(None) + } else { + Ok(Some(small_n)) + } + } + + #[cfg(test)] + fn encoded_collection_mode(&self, input: &str) -> Result<(usize, bool)> { + TL_MERGE_SCRATCH.with(|s| { + let mut scratch = s.borrow_mut(); + scratch.symbols.clear(); + scratch.heap.clear(); + let mut small_ids = [0u32; SMALL_MERGE_MAX]; + self.collect_encoded_symbols(input, &mut small_ids, &mut scratch) + .map(|small_n| match small_n { + Some(n) => (n, true), + None => (scratch.symbols.len(), false), + }) }) } /// Reference priority-queue BPE merge on already-encoded (ByteLevel) text. - /// It remains the fallback for long inputs and the correctness oracle for - /// the stack merger's tests. + /// It remains the correctness oracle for the stack merger's tests. + #[cfg(test)] fn merge_all_encoded_heap_into(&self, input: &str, out: &mut Vec) -> Result<()> { if input.is_empty() { return Ok(()); @@ -1585,82 +1593,26 @@ impl Bpe { } /// Linear-scan BPE merge for short encoded pretokens (`n <= - /// SMALL_MERGE_MAX` initial symbols). Avoids the `BinaryHeap` entirely: a - /// stack-resident doubly-linked list plus a per-position rank array, - /// find-min by a short scan over stack `u32`s, merge (O(1) pointer update), - /// then refresh only the two neighbor pairs. - /// - /// Produces the identical token sequence as [`Self::run_merge_loop`]: both - /// process the globally lowest-`(rank, pos)` active pair each step (the - /// heap's `MergeEntry` key is `(rank << 32) | pos`; the scan's strict `<` - /// keeps the leftmost/lowest-`pos` position on ties). Enforced by the - /// `merge_small_matches_heap` differential test. `ids[..n]` are the - /// initial encoded-symbol token ids. + /// SMALL_MERGE_MAX` initial symbols). The active bit is separate from the + /// rank so every `u32` rank, including `u32::MAX`, remains valid. fn merge_small_encoded(&self, ids: &mut [u32; SMALL_MERGE_MAX], n: usize, out: &mut Vec) { - let mut next = [0u8; SMALL_MERGE_MAX]; - let mut prev = [0u8; SMALL_MERGE_MAX]; - let mut ranks = [u32::MAX; SMALL_MERGE_MAX]; + let mut ranks = [0u32; SMALL_MERGE_MAX]; + let mut active = [false; SMALL_MERGE_MAX]; let mut new_ids = [0u32; SMALL_MERGE_MAX]; - for i in 0..n { - next[i] = (i + 1) as u8; - prev[i] = (i as u8).wrapping_sub(1); // prev[0] = 255 (>= n): sentinel - } for i in 0..n - 1 { if let Some((rank, new_id)) = self.merge_adj.get(ids[i], ids[i + 1]) { ranks[i] = rank; + active[i] = true; new_ids[i] = new_id; } } - loop { - let mut best = u32::MAX; - let mut best_i = 0usize; - for (i, &rank) in ranks[..n - 1].iter().enumerate() { - if rank < best { - best = rank; - best_i = i; - } - } - if best == u32::MAX { - break; - } - let i = best_i; - ids[i] = new_ids[i]; - let dead = next[i] as usize; - let new_right = next[dead] as usize; - next[i] = new_right as u8; - ranks[dead] = u32::MAX; - if new_right < n { - prev[new_right] = i as u8; - match self.merge_adj.get(ids[i], ids[new_right]) { - Some((rank, new_id)) => { - ranks[i] = rank; - new_ids[i] = new_id; - } - None => ranks[i] = u32::MAX, - } - } else { - ranks[i] = u32::MAX; - } - let left = prev[i] as usize; - if left < n { - match self.merge_adj.get(ids[left], ids[i]) { - Some((rank, new_id)) => { - ranks[left] = rank; - new_ids[left] = new_id; - } - None => ranks[left] = u32::MAX, - } - } - } - let mut i = 0usize; - while i < n { - out.push(ids[i]); - i = next[i] as usize; - } + self.merge_small_ids::(ids, n, &mut ranks, &mut active, &mut new_ids, out); } /// Linear-scan BPE merge for short raw pretokens (`n <= SMALL_MERGE_MAX` - /// bytes). `ids[..n]` are the per-byte initial token ids. + /// bytes). The raw initial table retains its existing `u32::MAX` absence + /// representation; the shared merge loop still keeps active state separate + /// from ranks after initialization. fn merge_small_raw( &self, bytes: &[u8], @@ -1668,14 +1620,9 @@ impl Bpe { n: usize, out: &mut Vec, ) { - let mut next = [0u8; SMALL_MERGE_MAX]; - let mut prev = [0u8; SMALL_MERGE_MAX]; - let mut ranks = [u32::MAX; SMALL_MERGE_MAX]; + let mut ranks = [0u32; SMALL_MERGE_MAX]; + let mut active = [false; SMALL_MERGE_MAX]; let mut new_ids = [0u32; SMALL_MERGE_MAX]; - for i in 0..n { - next[i] = (i + 1) as u8; - prev[i] = (i as u8).wrapping_sub(1); // prev[0] = 255 (>= n): sentinel - } // Round-1 ranks via the dense byte-pair table: one direct-indexed load // per pair instead of a CSR neighbor scan. for i in 0..n - 1 { @@ -1683,50 +1630,76 @@ impl Bpe { self.byte_pair_initial[bytes[i] as usize * 256 + bytes[i + 1] as usize]; if rank != u32::MAX { ranks[i] = rank; + active[i] = true; new_ids[i] = new_id; } } + self.merge_small_ids::(ids, n, &mut ranks, &mut active, &mut new_ids, out); + } + + /// Shared linked-list merge loop for the encoded and raw stack paths. + /// Scanning in ascending position order and replacing only on a strictly + /// lower rank preserves the heap's leftmost tie order. + fn merge_small_ids( + &self, + ids: &mut [u32; SMALL_MERGE_MAX], + n: usize, + ranks: &mut [u32; SMALL_MERGE_MAX], + active: &mut [bool; SMALL_MERGE_MAX], + new_ids: &mut [u32; SMALL_MERGE_MAX], + out: &mut Vec, + ) { + let mut next = [0u8; SMALL_MERGE_MAX]; + let mut prev = [0u8; SMALL_MERGE_MAX]; + for i in 0..n { + next[i] = (i + 1) as u8; + prev[i] = (i as u8).wrapping_sub(1); // prev[0] = 255 (>= n): sentinel + } + loop { - let mut best = u32::MAX; - let mut best_i = 0usize; - for (i, &rank) in ranks[..n - 1].iter().enumerate() { - if rank < best { - best = rank; + let mut best_i = n; + for i in 0..n - 1 { + if active[i] && (best_i == n || ranks[i] < ranks[best_i]) { best_i = i; } } - if best == u32::MAX { + if best_i == n { break; } + let i = best_i; ids[i] = new_ids[i]; let dead = next[i] as usize; let new_right = next[dead] as usize; next[i] = new_right as u8; - ranks[dead] = u32::MAX; + active[dead] = false; if new_right < n { prev[new_right] = i as u8; match self.merge_adj.get(ids[i], ids[new_right]) { - Some((rank, new_id)) => { + Some((rank, new_id)) if ALLOW_MAX_RANK || rank != u32::MAX => { ranks[i] = rank; + active[i] = true; new_ids[i] = new_id; } - None => ranks[i] = u32::MAX, + _ => active[i] = false, } } else { - ranks[i] = u32::MAX; + active[i] = false; } + let left = prev[i] as usize; if left < n { match self.merge_adj.get(ids[left], ids[i]) { - Some((rank, new_id)) => { + Some((rank, new_id)) if ALLOW_MAX_RANK || rank != u32::MAX => { ranks[left] = rank; + active[left] = true; new_ids[left] = new_id; } - None => ranks[left] = u32::MAX, + _ => active[left] = false, } } } + let mut i = 0usize; while i < n { out.push(ids[i]); @@ -2345,6 +2318,14 @@ mod tests { assert_eq!(optimized, heap, "output mismatch for {input:?}"); } + fn assert_public_tokenize_matches_heap(bpe: &Bpe, input: &str, expected: &[u32]) { + let actual = bpe.tokenize(input).unwrap(); + let mut heap = Vec::new(); + bpe.merge_all_encoded_heap_into(input, &mut heap).unwrap(); + assert_eq!(actual, expected); + assert_eq!(actual, heap, "public output mismatch for {input:?}"); + } + #[test] fn encoded_small_matches_heap_for_all_symbol_counts() { let bpe = test_bpe(); @@ -2365,6 +2346,43 @@ mod tests { } } + #[test] + fn encoded_small_public_api_preserves_absent_and_extreme_ranks() { + for rank in [ + None, + Some(0u32), + Some(7u32), + Some(u32::MAX - 1), + Some(u32::MAX), + ] { + let vocab: Vocab = [("a", 0), ("b", 1), ("ab", 2), ("z", 3)] + .into_iter() + .map(|(s, id)| (s.to_string(), id)) + .collect(); + let merge_map = rank + .map(|rank| [((0, 1), (rank, 2))].into_iter().collect()) + .unwrap_or_default(); + let bpe = Bpe::new(&vocab, merge_map).unwrap(); + let expected = if rank.is_some() { + vec![3, 2] + } else { + vec![3, 0, 1] + }; + assert_public_tokenize_matches_heap(&bpe, "zab", &expected); + } + } + + #[test] + fn encoded_small_public_api_preserves_leftmost_equal_rank_ties() { + let vocab: Vocab = [("z", 0), ("a", 1), ("b", 2), ("c", 3), ("ab", 4), ("bc", 5)] + .into_iter() + .map(|(s, id)| (s.to_string(), id)) + .collect(); + let merge_map = [((1, 2), (7, 4)), ((2, 3), (7, 5))].into_iter().collect(); + let bpe = Bpe::new(&vocab, merge_map).unwrap(); + assert_public_tokenize_matches_heap(&bpe, "zabc", &[0, 4, 3]); + } + #[test] fn encoded_small_preserves_leftmost_equal_rank_ties() { let vocab: Vocab = [("a", 0), ("b", 1), ("c", 2), ("ab", 3), ("bc", 4)] @@ -2380,6 +2398,25 @@ mod tests { assert_eq!(out, vec![3, 2]); } + #[test] + fn encoded_dispatch_uses_emitted_symbol_count() { + let left = BYTE_TO_CHAR[0]; + let bpe = { + let vocab: Vocab = [(left.to_string(), 0)].into_iter().collect(); + Bpe::new(&vocab, ParsedMergeMap::new()).unwrap() + }; + let short = left.to_string().repeat(SMALL_MERGE_MAX); + let long = left.to_string().repeat(SMALL_MERGE_MAX + 1); + assert!(short.len() > SMALL_MERGE_MAX); + assert_eq!(bpe.encoded_collection_mode(&short).unwrap(), (32, true)); + assert_eq!( + bpe.encoded_collection_mode(&long).unwrap(), + (SMALL_MERGE_MAX + 1, false) + ); + assert_encoded_matches_heap(&bpe, &short); + assert_encoded_matches_heap(&bpe, &long); + } + #[test] fn encoded_small_handles_bytelevel_chars_and_combining_marks() { let left = BYTE_TO_CHAR[0]; @@ -2431,8 +2468,18 @@ mod tests { fn encoded_small_handles_byte_fallback_expansion_and_boundary() { let bpe = fallback_bpe(true); assert_encoded_matches_heap(&bpe, "é"); - assert_encoded_matches_heap(&bpe, &format!("{}é", "a".repeat(30))); - assert_encoded_matches_heap(&bpe, &format!("{}é", "a".repeat(31))); + let at_boundary = format!("{}é", "a".repeat(30)); + let above_boundary = format!("{}é", "a".repeat(31)); + assert_eq!( + bpe.encoded_collection_mode(&at_boundary).unwrap(), + (32, true) + ); + assert_eq!( + bpe.encoded_collection_mode(&above_boundary).unwrap(), + (SMALL_MERGE_MAX + 1, false) + ); + assert_encoded_matches_heap(&bpe, &at_boundary); + assert_encoded_matches_heap(&bpe, &above_boundary); } #[test] From 3e131dbbcbdf31250e10cad5f5104f5e5acb76fb Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 12:31:19 +0000 Subject: [PATCH 12/18] test: bound encoded merge benchmark corpus --- examples/encoded_merge_bench.rs | 103 ++++++++++++++++++++++++-------- src/models/bpe.rs | 1 + 2 files changed, 78 insertions(+), 26 deletions(-) diff --git a/examples/encoded_merge_bench.rs b/examples/encoded_merge_bench.rs index fd08552..9a109b3 100644 --- a/examples/encoded_merge_bench.rs +++ b/examples/encoded_merge_bench.rs @@ -10,6 +10,7 @@ //! also report process CPU time for the measured loop on Unix. use std::{ + collections::HashSet, env, hint::black_box, mem::MaybeUninit, @@ -56,12 +57,12 @@ fn process_cpu_time() -> Duration { panic!("--cpu-time requires a Unix process CPU clock") } -fn fixture() -> Bpe { +fn fixture(symbols: usize) -> Bpe { let mut vocab = Map::new(); for (id, &byte) in ALPHABET.iter().enumerate() { - vocab.insert((byte as char).to_string(), Value::from((id + 1) as u32)); + vocab.insert((byte as char).to_string(), Value::from(id as u32)); } - vocab.insert("z".into(), Value::from(0u32)); + vocab.insert("z".into(), Value::from(ALPHABET.len() as u32)); // Every pair of body symbols is mergeable. The leading `z` is not part of // this table, so it prevents the whole input from matching a vocabulary @@ -76,48 +77,55 @@ fn fixture() -> Bpe { } } - // The aliases give the cache-miss corpus a large finite key space without - // changing the 16-symbol merge topology. Each alias maps to one of the - // dense body IDs, but no alias is itself a representative vocabulary token, - // so a generated input cannot take the whole-token fast path. - for index in 0..ALIAS_COUNT { - let id = 1 + (index % ALPHABET.len()) as u32; - vocab.insert(alias_char(index).to_string(), Value::from(id)); + if symbols <= 4 { + // Small smoke-test buckets need more distinct one- and two-symbol + // strings than the ASCII alphabet provides. These aliases share the + // existing dense body IDs and never form a representative whole token. + for index in 0..alias_domain(symbols) { + let id = (index % ALPHABET.len()) as u32; + vocab.insert(alias_char(index).to_string(), Value::from(id)); + } } serde_json::from_value(json!({"vocab": vocab, "merges": merges})) .expect("benchmark BPE fixture must deserialize") } -fn inputs(symbols: usize, count: usize) -> Vec { - assert!( - symbols >= 2, - "a measured input must not be a one-token match" - ); +fn alias_domain(symbols: usize) -> usize { + match symbols { + 2 => ALIAS_COUNT, + 3 => 128, + 4 => 32, + _ => ALPHABET.len(), + } +} + +fn alias_inputs(symbols: usize, count: usize) -> Vec { let body_len = symbols - 1; + let domain = alias_domain(symbols); let capacity = (0..body_len) - .try_fold(1usize, |capacity, _| capacity.checked_mul(ALIAS_COUNT)) + .try_fold(1usize, |capacity, _| capacity.checked_mul(domain)) .unwrap_or(usize::MAX); assert!( count <= capacity, "requested {count} inputs but --symbols {symbols} has capacity {capacity}" ); - // Decode each ordinal in base ALIAS_COUNT. The per-position affine map is - // bijective because ALIAS_COUNT is a power of two and 5 is odd, so this - // produces exactly `count` distinct strings without retrying a HashSet. + // Decode each ordinal in the selected power-of-two alias domain. The + // per-position affine map is bijective because 5 is odd, so this produces + // exactly `count` distinct strings without retrying a HashSet. let mut result = Vec::with_capacity(count); for ordinal in 0..count { let mut value = ordinal; let mut input = String::with_capacity(1 + body_len * 3); input.push('z'); for position in 0..body_len { - let digit = value % ALIAS_COUNT; - value /= ALIAS_COUNT; + let digit = value % domain; + value /= domain; let alias = digit .wrapping_mul(5) .wrapping_add(position.wrapping_mul(257)) - & (ALIAS_COUNT - 1); + & (domain - 1); input.push(alias_char(alias)); } result.push(input); @@ -125,6 +133,46 @@ fn inputs(symbols: usize, count: usize) -> Vec { result } +fn inputs(symbols: usize, count: usize, state: &mut u64) -> Vec { + assert!( + symbols >= 2, + "a measured input must not be a one-token match" + ); + if symbols <= 4 { + return alias_inputs(symbols, count); + } + + let body_len = symbols - 1; + let capacity = (0..body_len) + .try_fold(1usize, |capacity, _| capacity.checked_mul(ALPHABET.len())) + .unwrap_or(usize::MAX); + assert!( + count <= capacity, + "requested {count} inputs but --symbols {symbols} has capacity {capacity}" + ); + + let mut seen = HashSet::with_capacity(count); + let mut result = Vec::with_capacity(count); + while result.len() < count { + let mut value = *state; + let mut input = String::with_capacity(symbols); + input.push('z'); + for _ in 1..symbols { + // A deterministic stream gives every invocation the same workload, + // while the set makes each call a cache miss in the BPE caches. + value ^= value << 13; + value ^= value >> 7; + value ^= value << 17; + input.push(ALPHABET[value as usize & (ALPHABET.len() - 1)] as char); + } + *state = value.wrapping_add(0x9e37_79b9_7f4a_7c15); + if seen.insert(input.clone()) { + result.push(input); + } + } + result +} + fn main() { let mut symbols = None; let mut iterations = DEFAULT_ITERATIONS; @@ -158,11 +206,12 @@ fn main() { ); assert!(iterations > 0, "iterations must be positive"); - let bpe = fixture(); + let bpe = fixture(symbols); let total_inputs = WARMUP .checked_add(iterations) .expect("warmup plus iterations overflowed"); - let all_inputs = inputs(symbols, total_inputs); + let mut state = 0x243f_6a88_85a3_08d3u64 ^ symbols as u64; + let all_inputs = inputs(symbols, total_inputs, &mut state); let (warmup, measured) = all_inputs.split_at(WARMUP); for input in warmup { @@ -206,7 +255,8 @@ mod tests { #[test] fn corpus_is_finite_and_unique_at_small_and_boundary_sizes() { for symbols in [2, 3, 4, 16, 32, 33] { - let inputs = inputs(symbols, 256); + let mut state = 0x243f_6a88_85a3_08d3u64 ^ symbols as u64; + let inputs = inputs(symbols, 256, &mut state); assert_eq!(inputs.len(), 256); assert!(inputs.iter().all(|input| input.chars().count() == symbols)); let unique: HashSet<_> = inputs.iter().collect(); @@ -217,6 +267,7 @@ mod tests { #[test] #[should_panic(expected = "has capacity 16384")] fn corpus_rejects_requests_larger_than_the_smallest_domain() { - let _ = inputs(2, ALIAS_COUNT + 1); + let mut state = 0; + let _ = inputs(2, ALIAS_COUNT + 1, &mut state); } } diff --git a/src/models/bpe.rs b/src/models/bpe.rs index 9223a77..724db58 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -1640,6 +1640,7 @@ impl Bpe { /// Shared linked-list merge loop for the encoded and raw stack paths. /// Scanning in ascending position order and replacing only on a strictly /// lower rank preserves the heap's leftmost tie order. + #[inline(always)] fn merge_small_ids( &self, ids: &mut [u32; SMALL_MERGE_MAX], From baf2431418ce2ec2ec190caa469c76c49794a5ee Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 13:05:36 +0000 Subject: [PATCH 13/18] test: reject undersized encoded benchmark corpora --- examples/encoded_merge_bench.rs | 76 +++------------------------------ 1 file changed, 7 insertions(+), 69 deletions(-) diff --git a/examples/encoded_merge_bench.rs b/examples/encoded_merge_bench.rs index 9a109b3..c369503 100644 --- a/examples/encoded_merge_bench.rs +++ b/examples/encoded_merge_bench.rs @@ -23,12 +23,6 @@ use serde_json::{Map, Value, json}; const DEFAULT_ITERATIONS: usize = 8_192; const WARMUP: usize = 1_024; const ALPHABET: &[u8] = b"abcdefghijklmnop"; -const ALIAS_BASE: u32 = 0x1000; -const ALIAS_COUNT: usize = 16_384; - -fn alias_char(index: usize) -> char { - char::from_u32(ALIAS_BASE + index as u32).expect("benchmark alias must be a Unicode scalar") -} #[cfg(unix)] fn process_cpu_time() -> Duration { @@ -57,7 +51,7 @@ fn process_cpu_time() -> Duration { panic!("--cpu-time requires a Unix process CPU clock") } -fn fixture(symbols: usize) -> Bpe { +fn fixture() -> Bpe { let mut vocab = Map::new(); for (id, &byte) in ALPHABET.iter().enumerate() { vocab.insert((byte as char).to_string(), Value::from(id as u32)); @@ -77,71 +71,15 @@ fn fixture(symbols: usize) -> Bpe { } } - if symbols <= 4 { - // Small smoke-test buckets need more distinct one- and two-symbol - // strings than the ASCII alphabet provides. These aliases share the - // existing dense body IDs and never form a representative whole token. - for index in 0..alias_domain(symbols) { - let id = (index % ALPHABET.len()) as u32; - vocab.insert(alias_char(index).to_string(), Value::from(id)); - } - } - serde_json::from_value(json!({"vocab": vocab, "merges": merges})) .expect("benchmark BPE fixture must deserialize") } -fn alias_domain(symbols: usize) -> usize { - match symbols { - 2 => ALIAS_COUNT, - 3 => 128, - 4 => 32, - _ => ALPHABET.len(), - } -} - -fn alias_inputs(symbols: usize, count: usize) -> Vec { - let body_len = symbols - 1; - let domain = alias_domain(symbols); - let capacity = (0..body_len) - .try_fold(1usize, |capacity, _| capacity.checked_mul(domain)) - .unwrap_or(usize::MAX); - assert!( - count <= capacity, - "requested {count} inputs but --symbols {symbols} has capacity {capacity}" - ); - - // Decode each ordinal in the selected power-of-two alias domain. The - // per-position affine map is bijective because 5 is odd, so this produces - // exactly `count` distinct strings without retrying a HashSet. - let mut result = Vec::with_capacity(count); - for ordinal in 0..count { - let mut value = ordinal; - let mut input = String::with_capacity(1 + body_len * 3); - input.push('z'); - for position in 0..body_len { - let digit = value % domain; - value /= domain; - let alias = digit - .wrapping_mul(5) - .wrapping_add(position.wrapping_mul(257)) - & (domain - 1); - input.push(alias_char(alias)); - } - result.push(input); - } - result -} - fn inputs(symbols: usize, count: usize, state: &mut u64) -> Vec { assert!( symbols >= 2, "a measured input must not be a one-token match" ); - if symbols <= 4 { - return alias_inputs(symbols, count); - } - let body_len = symbols - 1; let capacity = (0..body_len) .try_fold(1usize, |capacity, _| capacity.checked_mul(ALPHABET.len())) @@ -206,7 +144,7 @@ fn main() { ); assert!(iterations > 0, "iterations must be positive"); - let bpe = fixture(symbols); + let bpe = fixture(); let total_inputs = WARMUP .checked_add(iterations) .expect("warmup plus iterations overflowed"); @@ -254,10 +192,10 @@ mod tests { #[test] fn corpus_is_finite_and_unique_at_small_and_boundary_sizes() { - for symbols in [2, 3, 4, 16, 32, 33] { + for (symbols, count) in [(2, 16), (3, 256), (4, 256), (16, 256), (32, 256), (33, 256)] { let mut state = 0x243f_6a88_85a3_08d3u64 ^ symbols as u64; - let inputs = inputs(symbols, 256, &mut state); - assert_eq!(inputs.len(), 256); + let inputs = inputs(symbols, count, &mut state); + assert_eq!(inputs.len(), count); assert!(inputs.iter().all(|input| input.chars().count() == symbols)); let unique: HashSet<_> = inputs.iter().collect(); assert_eq!(unique.len(), inputs.len()); @@ -265,9 +203,9 @@ mod tests { } #[test] - #[should_panic(expected = "has capacity 16384")] + #[should_panic(expected = "requested 17 inputs but --symbols 2 has capacity 16")] fn corpus_rejects_requests_larger_than_the_smallest_domain() { let mut state = 0; - let _ = inputs(2, ALIAS_COUNT + 1, &mut state); + let _ = inputs(2, 17, &mut state); } } From 7c8dc1e4b5936a80119bc32a19e419e966e700cb Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 14:00:26 +0000 Subject: [PATCH 14/18] perf: use a bounded heap for encoded short merges --- src/models/bpe.rs | 143 +++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 130 insertions(+), 13 deletions(-) diff --git a/src/models/bpe.rs b/src/models/bpe.rs index 724db58..9795c96 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -26,6 +26,7 @@ const INVALID_TOKEN: u32 = u32::MAX; /// stack-resident linear-scan merge instead of the heap. Sized so the stack /// arrays fit in registers/L1 and `u8` linked-list indices stay in range. const SMALL_MERGE_MAX: usize = 32; +const SMALL_MERGE_HEAP_CAP: usize = SMALL_MERGE_MAX * 2; /// Open-addressing hash table for merge lookups. #[derive(Clone, PartialEq)] @@ -771,6 +772,11 @@ impl MergeEntry { } } + #[inline(always)] + fn rank(&self) -> u32 { + (self.key >> 32) as u32 + } + #[inline(always)] fn pos(&self) -> u32 { self.key as u32 @@ -1606,7 +1612,7 @@ impl Bpe { new_ids[i] = new_id; } } - self.merge_small_ids::(ids, n, &mut ranks, &mut active, &mut new_ids, out); + self.merge_small_ids::(ids, n, &mut ranks, &mut active, &mut new_ids, out); } /// Linear-scan BPE merge for short raw pretokens (`n <= SMALL_MERGE_MAX` @@ -1634,14 +1640,15 @@ impl Bpe { new_ids[i] = new_id; } } - self.merge_small_ids::(ids, n, &mut ranks, &mut active, &mut new_ids, out); + self.merge_small_ids::(ids, n, &mut ranks, &mut active, &mut new_ids, out); } /// Shared linked-list merge loop for the encoded and raw stack paths. - /// Scanning in ascending position order and replacing only on a strictly - /// lower rank preserves the heap's leftmost tie order. + /// The encoded path uses a bounded binary heap of pair positions; the raw + /// path retains its original linear scan. Both order candidates by + /// `(rank, position)`, preserving the heap's leftmost tie order. #[inline(always)] - fn merge_small_ids( + fn merge_small_ids( &self, ids: &mut [u32; SMALL_MERGE_MAX], n: usize, @@ -1657,18 +1664,57 @@ impl Bpe { prev[i] = (i as u8).wrapping_sub(1); // prev[0] = 255 (>= n): sentinel } - loop { - let mut best_i = n; + let mut candidates = [MergeEntry { key: 0, val: 0 }; SMALL_MERGE_HEAP_CAP]; + let mut candidate_len = 0usize; + if USE_STACK_HEAP { for i in 0..n - 1 { - if active[i] && (best_i == n || ranks[i] < ranks[best_i]) { - best_i = i; + if active[i] { + let right = next[i] as usize; + Self::small_heap_push( + &mut candidates, + &mut candidate_len, + MergeEntry::new(ranks[i], i as u32, ids[i], ids[right]), + ); } } - if best_i == n { - break; - } + } + + loop { + let i = if USE_STACK_HEAP { + let mut selected = None; + while let Some(entry) = Self::small_heap_pop(&mut candidates, &mut candidate_len) { + let pos = entry.pos() as usize; + if pos >= n - 1 || !active[pos] { + continue; + } + let right = next[pos] as usize; + if right >= n + || ranks[pos] != entry.rank() + || ids[pos] != entry.left_c() + || ids[right] != entry.right_c() + { + continue; + } + selected = Some(pos); + break; + } + match selected { + Some(pos) => pos, + None => break, + } + } else { + let mut best_i = n; + for pos in 0..n - 1 { + if active[pos] && (best_i == n || ranks[pos] < ranks[best_i]) { + best_i = pos; + } + } + if best_i == n { + break; + } + best_i + }; - let i = best_i; ids[i] = new_ids[i]; let dead = next[i] as usize; let new_right = next[dead] as usize; @@ -1681,6 +1727,13 @@ impl Bpe { ranks[i] = rank; active[i] = true; new_ids[i] = new_id; + if USE_STACK_HEAP { + Self::small_heap_push( + &mut candidates, + &mut candidate_len, + MergeEntry::new(rank, i as u32, ids[i], ids[new_right]), + ); + } } _ => active[i] = false, } @@ -1695,6 +1748,13 @@ impl Bpe { ranks[left] = rank; active[left] = true; new_ids[left] = new_id; + if USE_STACK_HEAP { + Self::small_heap_push( + &mut candidates, + &mut candidate_len, + MergeEntry::new(rank, left as u32, ids[left], ids[i]), + ); + } } _ => active[left] = false, } @@ -1708,6 +1768,63 @@ impl Bpe { } } + #[inline(always)] + fn small_heap_push( + heap: &mut [MergeEntry; SMALL_MERGE_HEAP_CAP], + len: &mut usize, + entry: MergeEntry, + ) { + debug_assert!(*len < SMALL_MERGE_HEAP_CAP); + let mut index = *len; + *len += 1; + while index > 0 { + let parent = (index - 1) / 2; + if heap[parent].key <= entry.key { + break; + } + heap[index] = heap[parent]; + index = parent; + } + heap[index] = entry; + } + + #[inline(always)] + fn small_heap_pop( + heap: &mut [MergeEntry; SMALL_MERGE_HEAP_CAP], + len: &mut usize, + ) -> Option { + if *len == 0 { + return None; + } + let result = heap[0]; + *len -= 1; + if *len == 0 { + return Some(result); + } + + let replacement = heap[*len]; + let mut index = 0usize; + loop { + let left = index * 2 + 1; + if left >= *len { + break; + } + let right = left + 1; + let child = if right < *len && heap[right].key < heap[left].key { + right + } else { + left + }; + if replacement.key <= heap[child].key { + break; + } + heap[index] = heap[child]; + index = child; + } + heap[index] = replacement; + Some(result) + } + /// `ignore_merges` whole-pretoken lookup: is the ByteLevel-encoded form of /// the entire pretoken a single vocab token? Encodes into a stack buffer /// for short pretokens (`BYTE_TO_CHAR` codepoints are <= U+0143, so <= 2 From c2f0ce69f0357dda6e9781bbcd2894e8723da998 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 14:55:33 +0000 Subject: [PATCH 15/18] perf: use a tournament tree for encoded short merges --- src/models/bpe.rs | 158 +++++++++++++++------------------------------- 1 file changed, 50 insertions(+), 108 deletions(-) diff --git a/src/models/bpe.rs b/src/models/bpe.rs index 9795c96..1c7112f 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -22,11 +22,11 @@ type Vocab = HashMap; const INVALID_TOKEN: u32 = u32::MAX; -/// Pretokens with at most this many initial (per-byte) symbols use the -/// stack-resident linear-scan merge instead of the heap. Sized so the stack -/// arrays fit in registers/L1 and `u8` linked-list indices stay in range. +/// Pretokens with at most this many initial symbols use the stack-resident +/// short merger instead of the heap. Sized so the stack arrays fit in +/// registers/L1 and `u8` linked-list indices stay in range. const SMALL_MERGE_MAX: usize = 32; -const SMALL_MERGE_HEAP_CAP: usize = SMALL_MERGE_MAX * 2; +const SMALL_MERGE_TREE_SIZE: usize = SMALL_MERGE_MAX * 2; /// Open-addressing hash table for merge lookups. #[derive(Clone, PartialEq)] @@ -772,11 +772,6 @@ impl MergeEntry { } } - #[inline(always)] - fn rank(&self) -> u32 { - (self.key >> 32) as u32 - } - #[inline(always)] fn pos(&self) -> u32 { self.key as u32 @@ -1598,7 +1593,7 @@ impl Bpe { }) } - /// Linear-scan BPE merge for short encoded pretokens (`n <= + /// Tournament-tree BPE merge for short encoded pretokens (`n <= /// SMALL_MERGE_MAX` initial symbols). The active bit is separate from the /// rank so every `u32` rank, including `u32::MAX`, remains valid. fn merge_small_encoded(&self, ids: &mut [u32; SMALL_MERGE_MAX], n: usize, out: &mut Vec) { @@ -1644,11 +1639,11 @@ impl Bpe { } /// Shared linked-list merge loop for the encoded and raw stack paths. - /// The encoded path uses a bounded binary heap of pair positions; the raw + /// The encoded path uses a fixed tournament tree of pair positions; the raw /// path retains its original linear scan. Both order candidates by /// `(rank, position)`, preserving the heap's leftmost tie order. #[inline(always)] - fn merge_small_ids( + fn merge_small_ids( &self, ids: &mut [u32; SMALL_MERGE_MAX], n: usize, @@ -1664,44 +1659,27 @@ impl Bpe { prev[i] = (i as u8).wrapping_sub(1); // prev[0] = 255 (>= n): sentinel } - let mut candidates = [MergeEntry { key: 0, val: 0 }; SMALL_MERGE_HEAP_CAP]; - let mut candidate_len = 0usize; - if USE_STACK_HEAP { + // u64::MAX is an inactive leaf. A valid rank of u32::MAX still has + // room for its position in the low half of the key. + let mut candidates = [u64::MAX; SMALL_MERGE_TREE_SIZE]; + if USE_STACK_TREE { for i in 0..n - 1 { if active[i] { - let right = next[i] as usize; - Self::small_heap_push( - &mut candidates, - &mut candidate_len, - MergeEntry::new(ranks[i], i as u32, ids[i], ids[right]), - ); + candidates[SMALL_MERGE_MAX + i] = (ranks[i] as u64) << 32 | i as u64; } } + for i in (1..SMALL_MERGE_MAX).rev() { + candidates[i] = candidates[i * 2].min(candidates[i * 2 + 1]); + } } loop { - let i = if USE_STACK_HEAP { - let mut selected = None; - while let Some(entry) = Self::small_heap_pop(&mut candidates, &mut candidate_len) { - let pos = entry.pos() as usize; - if pos >= n - 1 || !active[pos] { - continue; - } - let right = next[pos] as usize; - if right >= n - || ranks[pos] != entry.rank() - || ids[pos] != entry.left_c() - || ids[right] != entry.right_c() - { - continue; - } - selected = Some(pos); + let i = if USE_STACK_TREE { + let key = candidates[1]; + if key == u64::MAX { break; } - match selected { - Some(pos) => pos, - None => break, - } + (key & u32::MAX as u64) as usize } else { let mut best_i = n; for pos in 0..n - 1 { @@ -1727,13 +1705,6 @@ impl Bpe { ranks[i] = rank; active[i] = true; new_ids[i] = new_id; - if USE_STACK_HEAP { - Self::small_heap_push( - &mut candidates, - &mut candidate_len, - MergeEntry::new(rank, i as u32, ids[i], ids[new_right]), - ); - } } _ => active[i] = false, } @@ -1748,17 +1719,34 @@ impl Bpe { ranks[left] = rank; active[left] = true; new_ids[left] = new_id; - if USE_STACK_HEAP { - Self::small_heap_push( - &mut candidates, - &mut candidate_len, - MergeEntry::new(rank, left as u32, ids[left], ids[i]), - ); - } } _ => active[left] = false, } } + + if USE_STACK_TREE { + Self::small_tree_update(&mut candidates, dead, u64::MAX); + Self::small_tree_update( + &mut candidates, + i, + if active[i] { + (ranks[i] as u64) << 32 | i as u64 + } else { + u64::MAX + }, + ); + if left < n { + Self::small_tree_update( + &mut candidates, + left, + if active[left] { + (ranks[left] as u64) << 32 | left as u64 + } else { + u64::MAX + }, + ); + } + } } let mut i = 0usize; @@ -1769,60 +1757,14 @@ impl Bpe { } #[inline(always)] - fn small_heap_push( - heap: &mut [MergeEntry; SMALL_MERGE_HEAP_CAP], - len: &mut usize, - entry: MergeEntry, - ) { - debug_assert!(*len < SMALL_MERGE_HEAP_CAP); - let mut index = *len; - *len += 1; - while index > 0 { - let parent = (index - 1) / 2; - if heap[parent].key <= entry.key { - break; - } - heap[index] = heap[parent]; - index = parent; - } - heap[index] = entry; - } - - #[inline(always)] - fn small_heap_pop( - heap: &mut [MergeEntry; SMALL_MERGE_HEAP_CAP], - len: &mut usize, - ) -> Option { - if *len == 0 { - return None; - } - let result = heap[0]; - *len -= 1; - if *len == 0 { - return Some(result); - } - - let replacement = heap[*len]; - let mut index = 0usize; - loop { - let left = index * 2 + 1; - if left >= *len { - break; - } - let right = left + 1; - let child = if right < *len && heap[right].key < heap[left].key { - right - } else { - left - }; - if replacement.key <= heap[child].key { - break; - } - heap[index] = heap[child]; - index = child; + fn small_tree_update(tree: &mut [u64; SMALL_MERGE_TREE_SIZE], pos: usize, key: u64) { + let mut index = SMALL_MERGE_MAX + pos; + tree[index] = key; + index >>= 1; + while index != 0 { + tree[index] = tree[index * 2].min(tree[index * 2 + 1]); + index >>= 1; } - heap[index] = replacement; - Some(result) } /// `ignore_merges` whole-pretoken lookup: is the ByteLevel-encoded form of From eea93088e439ab364e2c3cd7a0414fd2a8dac911 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 15:25:32 +0000 Subject: [PATCH 16/18] perf: keep raw fused short merge unchanged --- src/models/bpe.rs | 148 +++++++++++++++++++++++++++------------------- 1 file changed, 88 insertions(+), 60 deletions(-) diff --git a/src/models/bpe.rs b/src/models/bpe.rs index 1c7112f..1e2e00d 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -1607,13 +1607,11 @@ impl Bpe { new_ids[i] = new_id; } } - self.merge_small_ids::(ids, n, &mut ranks, &mut active, &mut new_ids, out); + self.merge_small_encoded_ids(ids, n, &mut ranks, &mut active, &mut new_ids, out); } /// Linear-scan BPE merge for short raw pretokens (`n <= SMALL_MERGE_MAX` - /// bytes). The raw initial table retains its existing `u32::MAX` absence - /// representation; the shared merge loop still keeps active state separate - /// from ranks after initialization. + /// bytes). `ids[..n]` are the per-byte initial token ids. fn merge_small_raw( &self, bytes: &[u8], @@ -1621,9 +1619,14 @@ impl Bpe { n: usize, out: &mut Vec, ) { - let mut ranks = [0u32; SMALL_MERGE_MAX]; - let mut active = [false; SMALL_MERGE_MAX]; + let mut next = [0u8; SMALL_MERGE_MAX]; + let mut prev = [0u8; SMALL_MERGE_MAX]; + let mut ranks = [u32::MAX; SMALL_MERGE_MAX]; let mut new_ids = [0u32; SMALL_MERGE_MAX]; + for i in 0..n { + next[i] = (i + 1) as u8; + prev[i] = (i as u8).wrapping_sub(1); // prev[0] = 255 (>= n): sentinel + } // Round-1 ranks via the dense byte-pair table: one direct-indexed load // per pair instead of a CSR neighbor scan. for i in 0..n - 1 { @@ -1631,19 +1634,62 @@ impl Bpe { self.byte_pair_initial[bytes[i] as usize * 256 + bytes[i + 1] as usize]; if rank != u32::MAX { ranks[i] = rank; - active[i] = true; new_ids[i] = new_id; } } - self.merge_small_ids::(ids, n, &mut ranks, &mut active, &mut new_ids, out); + loop { + let mut best = u32::MAX; + let mut best_i = 0usize; + for (i, &rank) in ranks[..n - 1].iter().enumerate() { + if rank < best { + best = rank; + best_i = i; + } + } + if best == u32::MAX { + break; + } + let i = best_i; + ids[i] = new_ids[i]; + let dead = next[i] as usize; + let new_right = next[dead] as usize; + next[i] = new_right as u8; + ranks[dead] = u32::MAX; + if new_right < n { + prev[new_right] = i as u8; + match self.merge_adj.get(ids[i], ids[new_right]) { + Some((rank, new_id)) => { + ranks[i] = rank; + new_ids[i] = new_id; + } + None => ranks[i] = u32::MAX, + } + } else { + ranks[i] = u32::MAX; + } + let left = prev[i] as usize; + if left < n { + match self.merge_adj.get(ids[left], ids[i]) { + Some((rank, new_id)) => { + ranks[left] = rank; + new_ids[left] = new_id; + } + None => ranks[left] = u32::MAX, + } + } + } + let mut i = 0usize; + while i < n { + out.push(ids[i]); + i = next[i] as usize; + } } - /// Shared linked-list merge loop for the encoded and raw stack paths. - /// The encoded path uses a fixed tournament tree of pair positions; the raw - /// path retains its original linear scan. Both order candidates by - /// `(rank, position)`, preserving the heap's leftmost tie order. + /// Tournament-tree merge for a short encoded pretoken. The tree stores + /// `(rank, original position)` keys, so its root is the same greedy, + /// leftmost choice as the reference heap without stale entries. #[inline(always)] - fn merge_small_ids( + fn merge_small_encoded_ids( &self, ids: &mut [u32; SMALL_MERGE_MAX], n: usize, @@ -1662,37 +1708,21 @@ impl Bpe { // u64::MAX is an inactive leaf. A valid rank of u32::MAX still has // room for its position in the low half of the key. let mut candidates = [u64::MAX; SMALL_MERGE_TREE_SIZE]; - if USE_STACK_TREE { - for i in 0..n - 1 { - if active[i] { - candidates[SMALL_MERGE_MAX + i] = (ranks[i] as u64) << 32 | i as u64; - } - } - for i in (1..SMALL_MERGE_MAX).rev() { - candidates[i] = candidates[i * 2].min(candidates[i * 2 + 1]); + for i in 0..n - 1 { + if active[i] { + candidates[SMALL_MERGE_MAX + i] = (ranks[i] as u64) << 32 | i as u64; } } + for i in (1..SMALL_MERGE_MAX).rev() { + candidates[i] = candidates[i * 2].min(candidates[i * 2 + 1]); + } loop { - let i = if USE_STACK_TREE { - let key = candidates[1]; - if key == u64::MAX { - break; - } - (key & u32::MAX as u64) as usize - } else { - let mut best_i = n; - for pos in 0..n - 1 { - if active[pos] && (best_i == n || ranks[pos] < ranks[best_i]) { - best_i = pos; - } - } - if best_i == n { - break; - } - best_i - }; - + let key = candidates[1]; + if key == u64::MAX { + break; + } + let i = (key & u32::MAX as u64) as usize; ids[i] = new_ids[i]; let dead = next[i] as usize; let new_right = next[dead] as usize; @@ -1701,12 +1731,12 @@ impl Bpe { if new_right < n { prev[new_right] = i as u8; match self.merge_adj.get(ids[i], ids[new_right]) { - Some((rank, new_id)) if ALLOW_MAX_RANK || rank != u32::MAX => { + Some((rank, new_id)) => { ranks[i] = rank; active[i] = true; new_ids[i] = new_id; } - _ => active[i] = false, + None => active[i] = false, } } else { active[i] = false; @@ -1715,37 +1745,35 @@ impl Bpe { let left = prev[i] as usize; if left < n { match self.merge_adj.get(ids[left], ids[i]) { - Some((rank, new_id)) if ALLOW_MAX_RANK || rank != u32::MAX => { + Some((rank, new_id)) => { ranks[left] = rank; active[left] = true; new_ids[left] = new_id; } - _ => active[left] = false, + None => active[left] = false, } } - if USE_STACK_TREE { - Self::small_tree_update(&mut candidates, dead, u64::MAX); + Self::small_tree_update(&mut candidates, dead, u64::MAX); + Self::small_tree_update( + &mut candidates, + i, + if active[i] { + (ranks[i] as u64) << 32 | i as u64 + } else { + u64::MAX + }, + ); + if left < n { Self::small_tree_update( &mut candidates, - i, - if active[i] { - (ranks[i] as u64) << 32 | i as u64 + left, + if active[left] { + (ranks[left] as u64) << 32 | left as u64 } else { u64::MAX }, ); - if left < n { - Self::small_tree_update( - &mut candidates, - left, - if active[left] { - (ranks[left] as u64) << 32 | left as u64 - } else { - u64::MAX - }, - ); - } } } From 756870a4672266bef9f9839ae9bfdaabe21b7d28 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 8 Sep 2026 15:51:53 +0000 Subject: [PATCH 17/18] perf: keep long ASCII encoded misses on the heap path --- src/models/bpe.rs | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/models/bpe.rs b/src/models/bpe.rs index 1e2e00d..2579296 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -1435,6 +1435,14 @@ impl Bpe { return Ok(()); } + // For ASCII, each valid character contributes exactly one emitted + // symbol, including byte-fallback characters. Skip the short-buffer + // collector once that count is provably above the bound, keeping the + // larger-input heap route's original single-pass setup. + if input.is_ascii() && input.len() > SMALL_MERGE_MAX { + return self.merge_all_encoded_heap_into(input, out); + } + TL_MERGE_SCRATCH.with(|s| { let mut scratch = s.borrow_mut(); scratch.symbols.clear(); @@ -1524,8 +1532,9 @@ impl Bpe { } /// Reference priority-queue BPE merge on already-encoded (ByteLevel) text. - /// It remains the correctness oracle for the stack merger's tests. - #[cfg(test)] + /// It remains the fallback for long inputs and the correctness oracle for + /// the stack merger's tests. + #[inline] fn merge_all_encoded_heap_into(&self, input: &str, out: &mut Vec) -> Result<()> { if input.is_empty() { return Ok(()); From e1338b9268208b5fee3e417945daeed1576a3bbe Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Wed, 9 Sep 2026 13:51:13 +0000 Subject: [PATCH 18/18] perf: simplify encoded short merger --- src/models/bpe.rs | 59 +++++++++++++++-------------------------------- 1 file changed, 18 insertions(+), 41 deletions(-) diff --git a/src/models/bpe.rs b/src/models/bpe.rs index 2579296..6e17976 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -26,7 +26,6 @@ const INVALID_TOKEN: u32 = u32::MAX; /// short merger instead of the heap. Sized so the stack arrays fit in /// registers/L1 and `u8` linked-list indices stay in range. const SMALL_MERGE_MAX: usize = 32; -const SMALL_MERGE_TREE_SIZE: usize = SMALL_MERGE_MAX * 2; /// Open-addressing hash table for merge lookups. #[derive(Clone, PartialEq)] @@ -1602,7 +1601,7 @@ impl Bpe { }) } - /// Tournament-tree BPE merge for short encoded pretokens (`n <= + /// Linear-scan BPE merge for short encoded pretokens (`n <= /// SMALL_MERGE_MAX` initial symbols). The active bit is separate from the /// rank so every `u32` rank, including `u32::MAX`, remains valid. fn merge_small_encoded(&self, ids: &mut [u32; SMALL_MERGE_MAX], n: usize, out: &mut Vec) { @@ -1694,8 +1693,8 @@ impl Bpe { } } - /// Tournament-tree merge for a short encoded pretoken. The tree stores - /// `(rank, original position)` keys, so its root is the same greedy, + /// Linear-scan merge for a short encoded pretoken. Candidate slots store + /// `(rank, original position)` keys, so the minimum is the same greedy, /// leftmost choice as the reference heap without stale entries. #[inline(always)] fn merge_small_encoded_ids( @@ -1714,20 +1713,17 @@ impl Bpe { prev[i] = (i as u8).wrapping_sub(1); // prev[0] = 255 (>= n): sentinel } - // u64::MAX is an inactive leaf. A valid rank of u32::MAX still has - // room for its position in the low half of the key. - let mut candidates = [u64::MAX; SMALL_MERGE_TREE_SIZE]; + // u64::MAX is an inactive candidate. A valid rank of u32::MAX still + // has room for its position in the low half of the key. + let mut candidates = [u64::MAX; SMALL_MERGE_MAX]; for i in 0..n - 1 { if active[i] { - candidates[SMALL_MERGE_MAX + i] = (ranks[i] as u64) << 32 | i as u64; + candidates[i] = (ranks[i] as u64) << 32 | i as u64; } } - for i in (1..SMALL_MERGE_MAX).rev() { - candidates[i] = candidates[i * 2].min(candidates[i * 2 + 1]); - } loop { - let key = candidates[1]; + let key = *candidates[..n - 1].iter().min().unwrap(); if key == u64::MAX { break; } @@ -1763,26 +1759,18 @@ impl Bpe { } } - Self::small_tree_update(&mut candidates, dead, u64::MAX); - Self::small_tree_update( - &mut candidates, - i, - if active[i] { - (ranks[i] as u64) << 32 | i as u64 + candidates[dead] = u64::MAX; + candidates[i] = if active[i] { + (ranks[i] as u64) << 32 | i as u64 + } else { + u64::MAX + }; + if left < n { + candidates[left] = if active[left] { + (ranks[left] as u64) << 32 | left as u64 } else { u64::MAX - }, - ); - if left < n { - Self::small_tree_update( - &mut candidates, - left, - if active[left] { - (ranks[left] as u64) << 32 | left as u64 - } else { - u64::MAX - }, - ); + }; } } @@ -1793,17 +1781,6 @@ impl Bpe { } } - #[inline(always)] - fn small_tree_update(tree: &mut [u64; SMALL_MERGE_TREE_SIZE], pos: usize, key: u64) { - let mut index = SMALL_MERGE_MAX + pos; - tree[index] = key; - index >>= 1; - while index != 0 { - tree[index] = tree[index * 2].min(tree[index * 2 + 1]); - index >>= 1; - } - } - /// `ignore_merges` whole-pretoken lookup: is the ByteLevel-encoded form of /// the entire pretoken a single vocab token? Encodes into a stack buffer /// for short pretokens (`BYTE_TO_CHAR` codepoints are <= U+0143, so <= 2