tokenizers: TryDecode(Span<char>) zero-allocation overload + IncrementalDetokenizer adoption

benchmarks/DotLLM.Benchmarks/Tokenizers/DetokenizerAllocBenchmark.cs

@@ -0,0 +1,122 @@
+using BenchmarkDotNet.Attributes;
+using DotLLM.Tokenizers;
+using DotLLM.Tokenizers.Bpe;
+
+namespace DotLLM.Benchmarks.Tokenizers;
+
+/// <summary>
+/// Allocation-focused benchmark comparing the existing allocating
+/// <see cref="ITokenizer.Decode(System.ReadOnlySpan{int}, bool)"/> path against the
+/// zero-allocation <see cref="ITokenizer.TryDecode"/> overload introduced for the
+/// <c>IncrementalDetokenizer</c> hot path. The interesting column in the BenchmarkDotNet
+/// report is <c>Allocated</c> — the bytes-per-op delta is what motivates the new surface.
+/// </summary>
+[MemoryDiagnoser]
+[SimpleJob(warmupCount: 3, iterationCount: 10)]
+public class DetokenizerAllocBenchmark
+{
+    private BpeTokenizer _spm = null!;
+    private BpeTokenizer _tiktoken = null!;
+    private int[] _spmIds = null!;
+    private int[] _tiktokenIds = null!;
+    private char[] _scratch = null!;
+
+    /// <summary>Number of decode calls in the benchmark loop.</summary>
+    [Params(1000)]
+    public int CallsPerInvoke { get; set; }
+
+    [GlobalSetup]
+    public void Setup()
+    {
+        _spm = BuildSpmVocab();
+        _tiktoken = BuildTiktokenVocab();
+        _spmIds = [10, 1, 11]; // "▁hello ▁world"-ish
+        _tiktokenIds = [11, 5, 15]; // "hello world"
+        _scratch = new char[256];
+
+        // Fail fast if _scratch can no longer hold either decode. The measured loops
+        // deliberately ignore the TryDecode result to keep the benchmark body minimal;
+        // without this guard a vocab change would silently make them time the
+        // buffer-too-small failure path (written == 0) and report meaningless numbers.
+        EnsureFits(_spm, _spmIds, nameof(_spmIds));
+        EnsureFits(_tiktoken, _tiktokenIds, nameof(_tiktokenIds));
+
+        void EnsureFits(ITokenizer tokenizer, int[] ids, string name)
+        {
+            if (!tokenizer.TryDecode(ids, stripBosSpace: false, _scratch, out int written) || written == 0)
+                throw new InvalidOperationException(
+                    $"Scratch buffer ({_scratch.Length} chars) is too small to decode {name}; " +
+                    "the benchmark would measure the failure path.");
+        }
+    }
+
+    [Benchmark(Baseline = true, Description = "SPM.Decode (allocating)")]
+    public int SpmDecode()
+    {
+        int total = 0;
+        for (int i = 0; i < CallsPerInvoke; i++)
+            total += _spm.Decode(_spmIds, stripBosSpace: false).Length;
+        return total;
+    }
+
+    [Benchmark(Description = "SPM.TryDecode (zero-alloc)")]
+    public int SpmTryDecode()
+    {
+        int total = 0;
+        for (int i = 0; i < CallsPerInvoke; i++)
+        {
+            _spm.TryDecode(_spmIds, stripBosSpace: false, _scratch, out int written);
+            total += written;
+        }
+        return total;
+    }
+
+    [Benchmark(Description = "Tiktoken.Decode (allocating)")]
+    public int TiktokenDecode()
+    {
+        int total = 0;
+        for (int i = 0; i < CallsPerInvoke; i++)
+            total += _tiktoken.Decode(_tiktokenIds, stripBosSpace: false).Length;
+        return total;
+    }
+
+    [Benchmark(Description = "Tiktoken.TryDecode (zero-alloc)")]
+    public int TiktokenTryDecode()
+    {
+        int total = 0;
+        for (int i = 0; i < CallsPerInvoke; i++)
+        {
+            _tiktoken.TryDecode(_tiktokenIds, stripBosSpace: false, _scratch, out int written);
+            total += written;
+        }
+        return total;
+    }
+
+    private static BpeTokenizer BuildSpmVocab()
+    {
+        string[] tokens =
+        [
+            "<unk>", "▁", "h", "e", "l", "o",
+            "▁h", "▁he", "▁hel", "▁hell", "▁hello",
+            "▁world", "▁w", "w", "r", "d",
+        ];
+        float[] scores = new float[tokens.Length];
+        return BpeTokenizer.CreateSentencePiece(tokens, scores, tokenTypes: null,
+            bosId: 0, eosId: 0, addBosSpace: false);
+    }
+
+    private static BpeTokenizer BuildTiktokenVocab()
+    {
+        string[] tokens =
+        [
+            "<unk>", "h", "e", "l", "o", " ", "w", "r", "d",
+            "he", "lo", "hello", " w", "wo", "rl", "world",
+        ];
+        string[] merges =
+        [
+            "h e", "l o", "he llo", "w o", "r l", "wo rld",
+        ];
+        return BpeTokenizer.CreateTiktoken(tokens, merges, tokenTypes: null,
+            bosId: 0, eosId: 0, preTokenizerType: null);
+    }
+}

src/DotLLM.Engine/IncrementalDetokenizer.cs

@@ -1,3 +1,4 @@
+using System.Buffers;
 using System.Runtime.InteropServices;
 using System.Text;
 using DotLLM.Tokenizers;
@@ -23,34 +24,45 @@ namespace DotLLM.Engine;
 /// A hard cap prevents unbounded growth in pathological cases by force-committing the current
 /// window text.
 /// </para>
+/// <para>
+/// Zero-allocation hot path: window decode and tail decode both go through
+/// <see cref="ITokenizer.TryDecode"/> into <see cref="ArrayPool{T}.Shared"/>-rented char buffers
+/// that grow-and-re-rent on overflow. <see cref="Dispose"/> returns the buffers; callers should
+/// hold the instance in a <c>using</c> or otherwise invoke <see cref="Dispose"/> deterministically.
+/// </para>
 /// </remarks>
-internal sealed class IncrementalDetokenizer
+internal sealed class IncrementalDetokenizer : IDisposable
 {
     private const int SoftWindowLimit = 4;
     private const int HardWindowLimit = 32;
+    private const int InitialWindowBufSize = 64;
 
     private readonly ITokenizer _tokenizer;
     private readonly StringBuilder _committed;
     private readonly List<int> _window;
-    private string _windowText;
+    private char[] _windowBuf;
+    private int _windowLen;
+    private char[] _tailBuf;
     private int _deltaBaseline;
 
     public IncrementalDetokenizer(ITokenizer tokenizer, int initialCapacity = 1024)
     {
         _tokenizer = tokenizer;
         _committed = new StringBuilder(initialCapacity);
         _window = new List<int>(HardWindowLimit + 1);
-        _windowText = string.Empty;
+        _windowBuf = ArrayPool<char>.Shared.Rent(InitialWindowBufSize);
+        _tailBuf = ArrayPool<char>.Shared.Rent(InitialWindowBufSize);
+        _windowLen = 0;
     }
 
     /// <summary>Total number of decoded characters (committed + window).</summary>
-    public int Length => _committed.Length + _windowText.Length;
+    public int Length => _committed.Length + _windowLen;
 
     /// <summary>Adds a token and advances the decoded state. Amortized O(1) per call.</summary>
     public void Append(int tokenId)
     {
         _window.Add(tokenId);
-        _windowText = _tokenizer.Decode(CollectionsMarshal.AsSpan(_window), stripBosSpace: false);
+        DecodeWindowInto(ref _windowBuf, out _windowLen);
 
         while (_window.Count > SoftWindowLimit)
         {
@@ -61,37 +73,69 @@ public void Append(int tokenId)
             {
                 // Pathological case: leading byte-token run that never resolves cleanly.
                 // Force-commit to prevent unbounded memory growth. Rare by construction.
-                _committed.Append(_windowText);
+                _committed.Append(_windowBuf, 0, _windowLen);
                 _window.Clear();
-                _windowText = string.Empty;
+                _windowLen = 0;
             }
             break;
         }
     }
 
+    private void DecodeWindowInto(ref char[] buffer, out int written)
+    {
+        var ids = CollectionsMarshal.AsSpan(_window);
+        while (true)
+        {
+            if (_tokenizer.TryDecode(ids, stripBosSpace: false, buffer, out written))
+                return;
+            // Overflow: grow and retry. ArrayPool.Rent rounds up to a power of two,
+            // so growth converges quickly even on multi-byte glyph runs.
+            char[] larger = ArrayPool<char>.Shared.Rent(buffer.Length * 2);
+            ArrayPool<char>.Shared.Return(buffer, clearArray: true);
+            buffer = larger;
+        }
+    }
+
+    private void DecodeTailInto(out int written)
+    {
+        var ids = CollectionsMarshal.AsSpan(_window).Slice(1);
+        while (true)
+        {
+            if (_tokenizer.TryDecode(ids, stripBosSpace: false, _tailBuf, out written))
+                return;
+            char[] larger = ArrayPool<char>.Shared.Rent(_tailBuf.Length * 2);
+            ArrayPool<char>.Shared.Return(_tailBuf, clearArray: true);
+            _tailBuf = larger;
+        }
+    }
+
     private bool TryEvictOldest()
     {
-        var tail = _tokenizer.Decode(
-            CollectionsMarshal.AsSpan(_window).Slice(1), stripBosSpace: false);
+        DecodeTailInto(out int tailLen);
 
-        if (tail.Length > _windowText.Length)
+        if (tailLen > _windowLen)
             return false;
 
-        var windowSpan = _windowText.AsSpan();
-        if (!windowSpan[(windowSpan.Length - tail.Length)..].SequenceEqual(tail))
+        ReadOnlySpan<char> windowSpan = _windowBuf.AsSpan(0, _windowLen);
+        ReadOnlySpan<char> tailSpan = _tailBuf.AsSpan(0, tailLen);
+        if (!windowSpan[(windowSpan.Length - tailLen)..].SequenceEqual(tailSpan))
             return false;
 
-        int commitLen = _windowText.Length - tail.Length;
+        int commitLen = _windowLen - tailLen;
         if (commitLen > 0)
-            _committed.Append(_windowText, 0, commitLen);
+            _committed.Append(_windowBuf, 0, commitLen);
         _window.RemoveAt(0);
-        _windowText = tail;
+        // Shift the surviving tail to the front of _windowBuf so subsequent appends see
+        // [0, _windowLen) as the current decoded window.
+        if (tailLen > 0)
+            _tailBuf.AsSpan(0, tailLen).CopyTo(_windowBuf);
+        _windowLen = tailLen;
         return true;
     }
 
     /// <summary>
     /// Returns a tail view over the last <paramref name="maxChars"/> characters of the decoded text.
-    /// The view aliases <see cref="_windowText"/> when possible (zero allocation), otherwise writes
+    /// The view aliases the window buffer when possible (zero allocation), otherwise writes
     /// into <paramref name="scratch"/>.
     /// </summary>
     /// <param name="maxChars">Maximum number of trailing characters to expose.</param>
@@ -104,17 +148,17 @@ public ReadOnlySpan<char> GetTailView(int maxChars, Span<char> scratch)
         if (take == 0)
             return default;
 
-        if (take <= _windowText.Length)
-            return _windowText.AsSpan(_windowText.Length - take);
+        if (take <= _windowLen)
+            return _windowBuf.AsSpan(_windowLen - take, take);
 
         if (scratch.Length < take)
             throw new ArgumentException("Scratch buffer too small for requested tail.", nameof(scratch));
 
-        int fromWindow = _windowText.Length;
+        int fromWindow = _windowLen;
         int fromCommitted = take - fromWindow;
         int committedStart = _committed.Length - fromCommitted;
         _committed.CopyTo(committedStart, scratch[..fromCommitted], fromCommitted);
-        _windowText.AsSpan().CopyTo(scratch.Slice(fromCommitted, fromWindow));
+        _windowBuf.AsSpan(0, fromWindow).CopyTo(scratch.Slice(fromCommitted, fromWindow));
         return scratch[..take];
     }
 
@@ -137,19 +181,43 @@ public string TakeDelta()
     public override string ToString()
     {
         if (_committed.Length == 0)
-            return _windowText;
-        if (_windowText.Length == 0)
+            return _windowLen == 0 ? string.Empty : new string(_windowBuf, 0, _windowLen);
+        if (_windowLen == 0)
             return _committed.ToString();
         return string.Create(
-            _committed.Length + _windowText.Length,
+            _committed.Length + _windowLen,
             this,
             static (span, self) =>
             {
                 self._committed.CopyTo(0, span[..self._committed.Length], self._committed.Length);
-                self._windowText.AsSpan().CopyTo(span[self._committed.Length..]);
+                self._windowBuf.AsSpan(0, self._windowLen).CopyTo(span[self._committed.Length..]);
             });
     }
 
+    /// <summary>
+    /// Returns pooled buffers to <see cref="ArrayPool{T}.Shared"/>. Idempotent.
+    /// </summary>
+    /// <remarks>
+    /// Buffers are returned with <c>clearArray: true</c> so decoded model output (which may
+    /// contain user prompt text or other sensitive content) is not left readable by an unrelated
+    /// renter of the shared pool. The buffers are small (tens to a few hundred chars) and this
+    /// runs once per generation, not per token, so it is off the hot path.
+    /// </remarks>
+    public void Dispose()
+    {
+        if (_windowBuf is { Length: > 0 } wb)
+        {
+            ArrayPool<char>.Shared.Return(wb, clearArray: true);
+            _windowBuf = [];
+        }
+        if (_tailBuf is { Length: > 0 } tb)
+        {
+            ArrayPool<char>.Shared.Return(tb, clearArray: true);
+            _tailBuf = [];
+        }
+        _windowLen = 0;
+    }
+
     private string SliceRange(int start, int endExclusive)
     {
         int length = endExclusive - start;
@@ -158,7 +226,7 @@ private string SliceRange(int start, int endExclusive)
         int committedLen = _committed.Length;
 
         if (start >= committedLen)
-            return _windowText.Substring(start - committedLen, length);
+            return new string(_windowBuf, start - committedLen, length);
 
         if (endExclusive <= committedLen)
             return _committed.ToString(start, length);
@@ -170,7 +238,7 @@ private string SliceRange(int start, int endExclusive)
             static (span, s) =>
             {
                 s.self._committed.CopyTo(s.start, span[..s.fromCommitted], s.fromCommitted);
-                s.self._windowText.AsSpan(0, s.fromWindow).CopyTo(span[s.fromCommitted..]);
+                s.self._windowBuf.AsSpan(0, s.fromWindow).CopyTo(span[s.fromCommitted..]);
             });
     }
 }