From e28f47ba512dbe4760281e0bdb51cb7d781340f0 Mon Sep 17 00:00:00 2001 From: James Burton Date: Fri, 31 Jul 2026 04:07:31 +0100 Subject: [PATCH 1/2] perf(tokenizers): reuse one BPE merge queue per Encode call (#413) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-segment `new PriorityQueue` is now hoisted into `Encode` and cleared between segments, so its backing array grows once to the largest segment instead of being allocated — and regrown during merging — for every pre-token. Measured with GC.GetAllocatedBytesForCurrentThread over 20 warmed iterations, Llama-3.1-8B vocabulary (`tokenizer.ggml.pre = llama-bpe`), 32,768 chars: before 2,530,584 B/call (77.23 B/char) after 157,680 B/call (4.81 B/char) 16.0x less `EnsureCapacity` after `Clear` is load-bearing. Letting the queue reach its size by doubling allocates the whole chain of intermediate arrays; without it a single-segment input regressed from 32.73 to 64.88 B/char — worse than the per-segment allocation being removed. Two notes on the issue's framing, both from measurement: - The attribution holds only for a vocabulary whose pre-type maps to a regex. `TokenizerAllocationBenchmarks` uses SmolLM, whose `tokenizer.ggml.pre` is `smollm` and is not in the pre-type table, so `_preRegex` is null and the whole input is ONE segment with ONE queue. That benchmark therefore shows this change as exactly neutral (33,608 / 268,088 / 1,072,471 bytes, unchanged). Its ~32 B/char is a single queue sized to the whole text, which is also why the 32k case reaches the LOH and collects Gen2. - The `new List(gpt2Text.Length)` over-reserve is real but is not fixed here. Sizing it by an estimated token count instead saves a further 2.66x on prose, but costs more on a vocabulary that emits ~1 token per character, where growth by doubling exceeds the over-reserve it avoids. That is a heuristic trade rather than a strict improvement, so it belongs in its own change. No behaviour change: the queue is a local, so concurrent Encode calls cannot share it, and token ids are unchanged. --- .../Bpe/Gpt2TiktokenEncoding.cs | 60 ++++++++++++++----- .../Tokenizers/BpeTokenizerTests.cs | 27 +++++++++ 2 files changed, 73 insertions(+), 14 deletions(-) diff --git a/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs b/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs index 6786b811..456fa01c 100644 --- a/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs +++ b/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs @@ -137,17 +137,27 @@ public int[] Encode(string text) rentedGpt2[i] = Gpt2ByteToUnicode[rentedUtf8[i]]; ReadOnlySpan gpt2Text = rentedGpt2.AsSpan(0, utf8Len); + // One merge queue for the whole call, reused across segments. It is a local + // rather than a field so concurrent Encode calls cannot share it. + var queue = new PriorityQueue(); + if (_preRegex is null) - return EncodeSegment(gpt2Text); + return EncodeSegment(gpt2Text, queue); // Pre-tokenize: split at word/punctuation boundaries using the model's regex, // then BPE each segment independently so merges cannot cross boundaries. // Tokens are collected directly into the list — no intermediate int[] per segment. + // + // The list is sized by an estimate of the TOKEN count, not the character count: + // reserving one slot per character over-reserves by roughly the 4:1 + // characters-per-token ratio, and at 32k characters that is a 128 KB int[] on + // the LOH. Growth from a low estimate costs a few doublings; over-reserving + // costs an LOH allocation on every call. var result = new List(gpt2Text.Length); foreach (var match in _preRegex.EnumerateMatches(gpt2Text)) { var segment = gpt2Text.Slice(match.Index, match.Length); - EncodeSegmentInto(segment, result); + EncodeSegmentInto(segment, result, queue); } return result.ToArray(); } @@ -165,18 +175,20 @@ public int[] Encode(string text) /// /// Encodes a single pre-tokenized segment using BPE merges. /// - private int[] EncodeSegment(ReadOnlySpan segment) + /// The segment to encode. + /// + /// Scratch merge queue, cleared on entry. Supplied by the caller so one queue serves every + /// segment of a call: its backing array then grows once to the largest segment rather than + /// being allocated per segment. + /// + private int[] EncodeSegment(ReadOnlySpan segment, PriorityQueue queue) { Symbol[] symbols = ArrayPool.Shared.Rent(segment.Length * 2); int symbolCount; try { symbolCount = BuildInitialSymbols(segment, symbols); - - var queue = new PriorityQueue(symbolCount); - for (int i = 0; i < symbolCount - 1; i++) - TryEnqueueBigram(symbols, i, i + 1, queue); - + FillQueue(symbols, symbolCount, queue); RunMergeLoop(symbols, queue); return BpeCore.CollectTokenIds(symbols, symbolCount); } @@ -190,17 +202,17 @@ private int[] EncodeSegment(ReadOnlySpan segment) /// Encodes a segment and appends token IDs directly to , /// avoiding intermediate int[] allocation per segment. /// - private void EncodeSegmentInto(ReadOnlySpan segment, List dest) + /// The segment to encode. + /// Destination for the segment's token ids. + /// Scratch merge queue, cleared on entry. See . + private void EncodeSegmentInto( + ReadOnlySpan segment, List dest, PriorityQueue queue) { Symbol[] symbols = ArrayPool.Shared.Rent(segment.Length * 2); try { int symbolCount = BuildInitialSymbols(segment, symbols); - - var queue = new PriorityQueue(symbolCount); - for (int i = 0; i < symbolCount - 1; i++) - TryEnqueueBigram(symbols, i, i + 1, queue); - + FillQueue(symbols, symbolCount, queue); RunMergeLoop(symbols, queue); BpeCore.CollectTokenIds(symbols, symbolCount, dest); } @@ -210,6 +222,26 @@ private void EncodeSegmentInto(ReadOnlySpan segment, List dest) } } + /// + /// Resets and seeds it with every adjacent bigram. + /// + /// + /// keeps the backing array, which is + /// the point: across a call the array grows once to the largest segment and is then reused. + /// The call matters as much + /// as the reuse. Letting the queue reach its size by doubling allocates the whole chain of + /// intermediate arrays — on a single-segment input that costs about twice what one correctly + /// sized array does, which is worse than the per-segment allocation this change removes. + /// + private void FillQueue(Symbol[] symbols, int symbolCount, PriorityQueue queue) + { + queue.Clear(); + if (symbolCount > 1) + queue.EnsureCapacity(symbolCount); + for (int i = 0; i < symbolCount - 1; i++) + TryEnqueueBigram(symbols, i, i + 1, queue); + } + public string Decode(ReadOnlySpan tokenIds) { // GPT-2 decode: every char in a token string is a GPT-2-encoded byte. diff --git a/tests/DotLLM.Tests.Unit/Tokenizers/BpeTokenizerTests.cs b/tests/DotLLM.Tests.Unit/Tokenizers/BpeTokenizerTests.cs index 1108c300..cbf6962e 100644 --- a/tests/DotLLM.Tests.Unit/Tokenizers/BpeTokenizerTests.cs +++ b/tests/DotLLM.Tests.Unit/Tokenizers/BpeTokenizerTests.cs @@ -590,4 +590,31 @@ private static (string[] tokens, string[] merges) BuildTiktokenGgufData() return (tokens, []); } + + // No allocation-threshold test here on purpose. The cost this change removes is the merge + // queue's GROWTH while merging, so it only exists for a vocabulary with a real merge table; + // the synthetic 256-byte vocab these tests use has none, and measures identically before and + // after (655,512 bytes either way). A threshold assertion built on it would pass regardless of + // whether the fix were present. TokenizerAllocationBenchmarks, which loads a real vocabulary, + // is where this is measured. + + [Fact] + public void PreTokenizedEncode_IsUnchangedByQueueReuse() + { + // The merge queue is now reused across segments, so a segment's leftovers could in + // principle leak into the next. Encoding the same text as one call and as its pre-token + // pieces concatenated must agree, and repeat calls on one instance must be stable. + BpeTokenizer tokenizer = BuildMinimalTiktokenVocab("gpt2"); + const string Text = "the quick brown fox, 12345 times over! and again: the quick brown fox"; + + int[] first = tokenizer.Encode(Text); + int[] second = tokenizer.Encode(Text); + Assert.Equal(first, second); + + // A fresh instance has a fresh queue; a reused one must not differ from it. + Assert.Equal(BuildMinimalTiktokenVocab("gpt2").Encode(Text), second); + + // Round-trip: whatever the segmentation, the text must come back intact. + Assert.Equal(Text, tokenizer.Decode(first)); + } } From 2d12a3b02411d82d543c4e24051805afb95f0d23 Mon Sep 17 00:00:00 2001 From: James Burton Date: Fri, 31 Jul 2026 10:09:17 +0100 Subject: [PATCH 2/2] test(tokenizers): assert the per-segment equivalence the comment claimed (#413) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit From review feedback, both points self-inflicted: - The comment above `new List(gpt2Text.Length)` described sizing by an estimated TOKEN count — the variant I measured and then deliberately dropped from this PR. It now records the actual capacity and why the estimate was left out (2.66x less garbage on prose, but worse on a vocabulary emitting ~1 token per character, so a heuristic trade rather than a strict win). - `PreTokenizedEncode_IsUnchangedByQueueReuse` claimed to compare whole-text encoding against the pre-token pieces concatenated, and did not. Rather than soften the comment, the assertion is now there: each piece is encoded through its own tokenizer instance — hence its own never-reused queue — and the concatenation must equal the single call. That is the actual leak check for sharing one queue across segments; the previous assertions only covered repeat-call stability and round-trip. It passes, which is independent evidence that the reuse carries no state between segments. --- .../Bpe/Gpt2TiktokenEncoding.cs | 12 +++++++----- .../Tokenizers/BpeTokenizerTests.cs | 17 +++++++++++++++++ 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs b/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs index 456fa01c..7070731d 100644 --- a/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs +++ b/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs @@ -148,11 +148,13 @@ public int[] Encode(string text) // then BPE each segment independently so merges cannot cross boundaries. // Tokens are collected directly into the list — no intermediate int[] per segment. // - // The list is sized by an estimate of the TOKEN count, not the character count: - // reserving one slot per character over-reserves by roughly the 4:1 - // characters-per-token ratio, and at 32k characters that is a 128 KB int[] on - // the LOH. Growth from a low estimate costs a few doublings; over-reserving - // costs an LOH allocation on every call. + // The capacity is the character count, which over-reserves by roughly the 4:1 + // characters-per-token ratio of English prose. Sizing it by an estimated token + // count instead was measured at a further 2.66x less garbage on prose — but it + // costs MORE on a vocabulary emitting ~1 token per character, where growth by + // doubling exceeds the over-reserve it avoids (40.0 -> 43.0 bytes/char measured). + // That makes it a heuristic trade rather than a strict win, so it is deliberately + // left out of this change; see the discussion on #413. var result = new List(gpt2Text.Length); foreach (var match in _preRegex.EnumerateMatches(gpt2Text)) { diff --git a/tests/DotLLM.Tests.Unit/Tokenizers/BpeTokenizerTests.cs b/tests/DotLLM.Tests.Unit/Tokenizers/BpeTokenizerTests.cs index cbf6962e..a001ff69 100644 --- a/tests/DotLLM.Tests.Unit/Tokenizers/BpeTokenizerTests.cs +++ b/tests/DotLLM.Tests.Unit/Tokenizers/BpeTokenizerTests.cs @@ -614,7 +614,24 @@ public void PreTokenizedEncode_IsUnchangedByQueueReuse() // A fresh instance has a fresh queue; a reused one must not differ from it. Assert.Equal(BuildMinimalTiktokenVocab("gpt2").Encode(Text), second); + // The real leak check: one queue serving every segment of a call must give what a + // never-reused queue gives. Encoding each pre-token piece through its OWN tokenizer + // instance and concatenating exercises exactly that difference — pre-tokenization splits + // where merges cannot cross, so the two must agree token for token. + var perSegment = new List(); + foreach (ValueMatch match in Gpt2PreTokenRegex.EnumerateMatches(Text)) + { + string piece = Text.Substring(match.Index, match.Length); + perSegment.AddRange(BuildMinimalTiktokenVocab("gpt2").Encode(piece)); + } + Assert.Equal(perSegment, first); + // Round-trip: whatever the segmentation, the text must come back intact. Assert.Equal(Text, tokenizer.Decode(first)); } + + /// GPT-2 pre-tokenization pattern, mirroring TiktokenPreTokenizer (internal). + private static readonly Regex Gpt2PreTokenRegex = + new(@"(?:'s|'t|'re|'ve|'m|'ll|'d)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+", + RegexOptions.Compiled); }