perf(tokenizers): reuse one BPE merge queue per Encode call (#413)

src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs

@@ -137,17 +137,29 @@ public int[] Encode(string text)
                     rentedGpt2[i] = Gpt2ByteToUnicode[rentedUtf8[i]];
                 ReadOnlySpan<char> 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<BgramEntry, (int, int)>();
+
                 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 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<int>(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 +177,20 @@ public int[] Encode(string text)
     /// <summary>
     /// Encodes a single pre-tokenized segment using BPE merges.
     /// </summary>
-    private int[] EncodeSegment(ReadOnlySpan<char> segment)
+    /// <param name="segment">The segment to encode.</param>
+    /// <param name="queue">
+    /// 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.
+    /// </param>
+    private int[] EncodeSegment(ReadOnlySpan<char> segment, PriorityQueue<BgramEntry, (int, int)> queue)
     {
         Symbol[] symbols = ArrayPool<Symbol>.Shared.Rent(segment.Length * 2);
         int symbolCount;
         try
         {
             symbolCount = BuildInitialSymbols(segment, symbols);
-
-            var queue = new PriorityQueue<BgramEntry, (int, int)>(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 +204,17 @@ private int[] EncodeSegment(ReadOnlySpan<char> segment)
     /// Encodes a segment and appends token IDs directly to <paramref name="dest"/>,
     /// avoiding intermediate <c>int[]</c> allocation per segment.
     /// </summary>
-    private void EncodeSegmentInto(ReadOnlySpan<char> segment, List<int> dest)
+    /// <param name="segment">The segment to encode.</param>
+    /// <param name="dest">Destination for the segment's token ids.</param>
+    /// <param name="queue">Scratch merge queue, cleared on entry. See <see cref="EncodeSegment"/>.</param>
+    private void EncodeSegmentInto(
+        ReadOnlySpan<char> segment, List<int> dest, PriorityQueue<BgramEntry, (int, int)> queue)
     {
         Symbol[] symbols = ArrayPool<Symbol>.Shared.Rent(segment.Length * 2);
         try
         {
             int symbolCount = BuildInitialSymbols(segment, symbols);
-
-            var queue = new PriorityQueue<BgramEntry, (int, int)>(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 +224,26 @@ private void EncodeSegmentInto(ReadOnlySpan<char> segment, List<int> dest)
         }
     }
 
+    /// <summary>
+    /// Resets <paramref name="queue"/> and seeds it with every adjacent bigram.
+    /// </summary>
+    /// <remarks>
+    /// <para><see cref="PriorityQueue{TElement,TPriority}.Clear"/> keeps the backing array, which is
+    /// the point: across a call the array grows once to the largest segment and is then reused.</para>
+    /// <para>The <see cref="PriorityQueue{TElement,TPriority}.EnsureCapacity"/> 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.</para>
+    /// </remarks>
+    private void FillQueue(Symbol[] symbols, int symbolCount, PriorityQueue<BgramEntry, (int, int)> 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<int> tokenIds)
     {
         // GPT-2 decode: every char in a token string is a GPT-2-encoded byte.

tests/DotLLM.Tests.Unit/Tokenizers/BpeTokenizerTests.cs

@@ -590,4 +590,48 @@ 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);
+
+        // 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<int>();
+        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));
+    }
+
+    /// <summary>GPT-2 pre-tokenization pattern, mirroring <c>TiktokenPreTokenizer</c> (internal).</summary>
+    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);
 }