From abf1536e229f24a840b165fbc31f5cdcf94935bb Mon Sep 17 00:00:00 2001 From: James Burton Date: Mon, 8 Jun 2026 18:09:03 +0100 Subject: [PATCH 1/2] tokenizers: TryDecode(Span) zero-allocation overload + IncrementalDetokenizer adoption (#121) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Item #12 follow-up from Gemini's review on PR #122. Adds a span-writing decode overload to the tokenizer surface and adopts it in `IncrementalDetokenizer` to eliminate the two per-step `string` allocations that remained after PR #122 amortized the O(n^2) detokenization cost. Surface - `ITokenizer.TryDecode(ReadOnlySpan, bool stripBosSpace, Span, out int charsWritten)` default interface method that forwards to `Decode(...)` + copy (external impls unchanged). - `IBpeEncoding.TryDecode(...)` mirrors the same default forward. - `BpeCore.TryFlushByteBuffer` — span variant of the SPM byte-run flush used by the SPM zero-alloc path. Encoding overrides - `Gpt2TiktokenEncoding`: materializes the GPT-2 byte stream once, then `Encoding.UTF8.GetCharCount` pre-checks the destination size for an atomic return — destination is byte-for-byte untouched on overflow. Existing `Decode(...)` factored over the same helper to keep both paths in sync. - `SentencePieceEncoding`: cursor write into the destination span with `MemoryExtensions.Replace` for the inline `_` -> space substitution and an in-place left-shift for the `stripBosSpace` leading-space strip. Non-atomic on overflow (per the documented contract). Adoption - `IncrementalDetokenizer`: replaces `string _windowText` with an ArrayPool `char[] _windowBuf` + `int _windowLen` (grow-and-re-rent on overflow), a matching `_tailBuf` for `TryEvictOldest`, and an `IDisposable` implementation that returns the buffers. `TextGenerator` declares the detokenizer outside its `try` block at both call sites and disposes in the `finally`. Tests + benchmarks - `TryDecodeZeroAllocTests` (15 tests): parity vs `Decode` for SPM + tiktoken across curated and random sequences with both `stripBosSpace` values; buffer- too-small contract (tiktoken atomic, SPM bounded-write); `GC.GetAllocated- BytesForCurrentThread` delta == 0 over 1000 warm calls for both encodings; baseline sanity check that the allocating `Decode` path still allocates. - `IncrementalDetokenizerTests` (10 tests): existing 9 regression-gate the buffer rewrite; new `Append_SteadyState_AllocationFloorIsBoundedAndFarBelow- PreviousPath` asserts the adopted path's per-200-call delta stays well below the pre-PR baseline (it is now 0 bytes in practice; ceiling is conservative because the committed StringBuilder chunk-grows amortized). - `DetokenizerAllocBenchmark` — BenchmarkDotNet MemoryDiagnoser baseline for the primitive `Decode` vs `TryDecode` allocation delta on both encodings. Co-Authored-By: Claude Opus 4.7 --- .../Tokenizers/DetokenizerAllocBenchmark.cs | 107 ++++++ src/DotLLM.Engine/IncrementalDetokenizer.cs | 112 ++++-- src/DotLLM.Engine/TextGenerator.cs | 20 +- src/DotLLM.Tokenizers/Bpe/BpeCore.cs | 38 ++ src/DotLLM.Tokenizers/Bpe/BpeTokenizer.cs | 11 + .../Bpe/Gpt2TiktokenEncoding.cs | 60 +++- .../Bpe/SentencePieceEncoding.cs | 77 ++++ src/DotLLM.Tokenizers/ITokenizer.cs | 40 +++ .../Engine/IncrementalDetokenizerTests.cs | 36 ++ .../Tokenizers/TryDecodeZeroAllocTests.cs | 334 ++++++++++++++++++ 10 files changed, 791 insertions(+), 44 deletions(-) create mode 100644 benchmarks/DotLLM.Benchmarks/Tokenizers/DetokenizerAllocBenchmark.cs create mode 100644 tests/DotLLM.Tests.Unit/Tokenizers/TryDecodeZeroAllocTests.cs diff --git a/benchmarks/DotLLM.Benchmarks/Tokenizers/DetokenizerAllocBenchmark.cs b/benchmarks/DotLLM.Benchmarks/Tokenizers/DetokenizerAllocBenchmark.cs new file mode 100644 index 00000000..f309c242 --- /dev/null +++ b/benchmarks/DotLLM.Benchmarks/Tokenizers/DetokenizerAllocBenchmark.cs @@ -0,0 +1,107 @@ +using BenchmarkDotNet.Attributes; +using DotLLM.Tokenizers; +using DotLLM.Tokenizers.Bpe; + +namespace DotLLM.Benchmarks.Tokenizers; + +/// +/// Allocation-focused benchmark comparing the existing allocating +/// path against the +/// zero-allocation overload introduced for the +/// IncrementalDetokenizer hot path. The interesting column in the BenchmarkDotNet +/// report is Allocated — the bytes-per-op delta is what motivates the new surface. +/// +[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!; + + /// Number of decode calls in the benchmark loop. + [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]; + } + + [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 = + [ + "", "▁", "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 = + [ + "", "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); + } +} diff --git a/src/DotLLM.Engine/IncrementalDetokenizer.cs b/src/DotLLM.Engine/IncrementalDetokenizer.cs index 9748b176..12fa01f9 100644 --- a/src/DotLLM.Engine/IncrementalDetokenizer.cs +++ b/src/DotLLM.Engine/IncrementalDetokenizer.cs @@ -1,3 +1,4 @@ +using System.Buffers; using System.Runtime.InteropServices; using System.Text; using DotLLM.Tokenizers; @@ -23,16 +24,25 @@ namespace DotLLM.Engine; /// A hard cap prevents unbounded growth in pathological cases by force-committing the current /// window text. /// +/// +/// Zero-allocation hot path: window decode and tail decode both go through +/// into -rented char buffers +/// that grow-and-re-rent on overflow. returns the buffers; callers should +/// hold the instance in a using or otherwise invoke deterministically. +/// /// -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 _window; - private string _windowText; + private char[] _windowBuf; + private int _windowLen; + private char[] _tailBuf; private int _deltaBaseline; public IncrementalDetokenizer(ITokenizer tokenizer, int initialCapacity = 1024) @@ -40,17 +50,19 @@ public IncrementalDetokenizer(ITokenizer tokenizer, int initialCapacity = 1024) _tokenizer = tokenizer; _committed = new StringBuilder(initialCapacity); _window = new List(HardWindowLimit + 1); - _windowText = string.Empty; + _windowBuf = ArrayPool.Shared.Rent(InitialWindowBufSize); + _tailBuf = ArrayPool.Shared.Rent(InitialWindowBufSize); + _windowLen = 0; } /// Total number of decoded characters (committed + window). - public int Length => _committed.Length + _windowText.Length; + public int Length => _committed.Length + _windowLen; /// Adds a token and advances the decoded state. Amortized O(1) per call. 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.Shared.Rent(buffer.Length * 2); + ArrayPool.Shared.Return(buffer); + 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.Shared.Rent(_tailBuf.Length * 2); + ArrayPool.Shared.Return(_tailBuf); + _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 windowSpan = _windowBuf.AsSpan(0, _windowLen); + ReadOnlySpan 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; } /// /// Returns a tail view over the last characters of the decoded text. - /// The view aliases when possible (zero allocation), otherwise writes + /// The view aliases the window buffer when possible (zero allocation), otherwise writes /// into . /// /// Maximum number of trailing characters to expose. @@ -104,17 +148,17 @@ public ReadOnlySpan GetTailView(int maxChars, Span 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,35 @@ 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..]); }); } + /// Returns pooled buffers to . Idempotent. + public void Dispose() + { + if (_windowBuf is { Length: > 0 } wb) + { + ArrayPool.Shared.Return(wb); + _windowBuf = []; + } + if (_tailBuf is { Length: > 0 } tb) + { + ArrayPool.Shared.Return(tb); + _tailBuf = []; + } + _windowLen = 0; + } + private string SliceRange(int start, int endExclusive) { int length = endExclusive - start; @@ -158,7 +218,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 +230,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..]); }); } } diff --git a/src/DotLLM.Engine/TextGenerator.cs b/src/DotLLM.Engine/TextGenerator.cs index 2b9fa141..3902ba37 100644 --- a/src/DotLLM.Engine/TextGenerator.cs +++ b/src/DotLLM.Engine/TextGenerator.cs @@ -139,6 +139,11 @@ public InferenceResponse Generate(string prompt, InferenceOptions? options = nul int stopTailSize = ComputeStopTailSize(stopConditions); char[] stopScratch = ArrayPool.Shared.Rent(stopTailSize); + // Incremental detokenizer keeps stop-check cost O(1) amortized per token + // instead of decoding the entire generated sequence each step (O(n²)). + // Declared outside the try so the finally can deterministically return its pooled buffers. + var detok = new IncrementalDetokenizer(_tokenizer, initialCapacity: Math.Max(64, maxTokens * 4)); + try { var generatedIds = new List(maxTokens); @@ -148,10 +153,6 @@ public InferenceResponse Generate(string prompt, InferenceOptions? options = nul long samplerTicks = 0; int cacheSize = kvCache.MaxLength; - // Incremental detokenizer keeps stop-check cost O(1) amortized per token - // instead of decoding the entire generated sequence each step (O(n²)). - var detok = new IncrementalDetokenizer(_tokenizer, initialCapacity: Math.Max(64, maxTokens * 4)); - // Local helper: snapshot log-softmax before sampling (which modifies logits in-place), // sample a token, then build logprob info. (int tokenId, TokenLogprobInfo? logprob) SampleWithLogprobs(Span logitSpan) @@ -399,6 +400,7 @@ public InferenceResponse Generate(string prompt, InferenceOptions? options = nul finally { ArrayPool.Shared.Return(stopScratch); + detok.Dispose(); if (ownsKvCache) kvCache.Dispose(); } @@ -481,6 +483,11 @@ public async IAsyncEnumerable GenerateStreamingTokensAsync( int stopTailSize = ComputeStopTailSize(stopConditions); char[] stopScratch = ArrayPool.Shared.Rent(stopTailSize); + // Incremental detokenizer: O(1) amortized per token for stop-check + streaming delta, + // instead of decoding the full generated sequence at every step. Lifted out of the try + // so the finally can return its pooled buffers even on cancellation. + var detok = new IncrementalDetokenizer(_tokenizer, initialCapacity: Math.Max(64, maxTokens * 4)); + try { var generatedIds = new List(maxTokens); @@ -489,10 +496,6 @@ public async IAsyncEnumerable GenerateStreamingTokensAsync( long samplerTicks = 0; int cacheSize = kvCache.MaxLength; - // Incremental detokenizer: O(1) amortized per token for stop-check + streaming delta, - // instead of decoding the full generated sequence at every step. - var detok = new IncrementalDetokenizer(_tokenizer, initialCapacity: Math.Max(64, maxTokens * 4)); - // Local helper: snapshot log-softmax before sampling (which modifies logits in-place), // sample a token, then build logprob info. (int tokenId, TokenLogprobInfo? logprob) SampleWithLogprobs(Span logitSpan) @@ -796,6 +799,7 @@ public async IAsyncEnumerable GenerateStreamingTokensAsync( finally { ArrayPool.Shared.Return(stopScratch); + detok.Dispose(); if (ownsKvCache) kvCache.Dispose(); } diff --git a/src/DotLLM.Tokenizers/Bpe/BpeCore.cs b/src/DotLLM.Tokenizers/Bpe/BpeCore.cs index fadd4b52..1150b1e0 100644 --- a/src/DotLLM.Tokenizers/Bpe/BpeCore.cs +++ b/src/DotLLM.Tokenizers/Bpe/BpeCore.cs @@ -19,6 +19,27 @@ internal interface IBpeEncoding int[] EncodeSegment(string text) => Encode(text); string Decode(ReadOnlySpan tokenIds); string Decode(ReadOnlySpan tokenIds, bool stripBosSpace) => Decode(tokenIds); + + /// + /// Zero-allocation decode into a caller-provided . Returns + /// on success with the character count in , + /// or when the destination was too small (contents of + /// unspecified on failure). The default forwards to the + /// allocating + copy. + /// + bool TryDecode(ReadOnlySpan tokenIds, bool stripBosSpace, Span destination, out int charsWritten) + { + string decoded = Decode(tokenIds, stripBosSpace); + if (decoded.Length > destination.Length) + { + charsWritten = 0; + return false; + } + decoded.AsSpan().CopyTo(destination); + charsWritten = decoded.Length; + return true; + } + string DecodeToken(int tokenId); } @@ -109,6 +130,23 @@ internal static void FlushByteBuffer(StringBuilder sb, byte[]? buffer, ref int c count = 0; } + /// + /// Writes buffered bytes as UTF-8 chars into starting at + /// , advances , and resets + /// . Returns when the destination is + /// too small (cursor and count are unchanged in that case). + /// + internal static bool TryFlushByteBuffer(Span destination, ref int cursor, byte[]? buffer, ref int count) + { + if (count == 0) return true; + int needed = Encoding.UTF8.GetCharCount(buffer!, 0, count); + if (cursor + needed > destination.Length) return false; + int written = Encoding.UTF8.GetChars(buffer.AsSpan(0, count), destination[cursor..]); + cursor += written; + count = 0; + return true; + } + /// /// Builds a 256-entry byte→token-ID mapping from <0xNN> vocab entries. /// Entries with no matching byte token are set to -1. diff --git a/src/DotLLM.Tokenizers/Bpe/BpeTokenizer.cs b/src/DotLLM.Tokenizers/Bpe/BpeTokenizer.cs index afd1376f..58f2b7c7 100644 --- a/src/DotLLM.Tokenizers/Bpe/BpeTokenizer.cs +++ b/src/DotLLM.Tokenizers/Bpe/BpeTokenizer.cs @@ -116,6 +116,17 @@ public string Decode(ReadOnlySpan tokenIds) => public string Decode(ReadOnlySpan tokenIds, bool stripBosSpace) => tokenIds.IsEmpty ? string.Empty : _encoding.Decode(tokenIds, stripBosSpace); + /// + public bool TryDecode(ReadOnlySpan tokenIds, bool stripBosSpace, Span destination, out int charsWritten) + { + if (tokenIds.IsEmpty) + { + charsWritten = 0; + return true; + } + return _encoding.TryDecode(tokenIds, stripBosSpace, destination, out charsWritten); + } + /// public string DecodeToken(int tokenId) => _encoding.DecodeToken(tokenId); diff --git a/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs b/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs index 6786b811..e7c769f1 100644 --- a/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs +++ b/src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs @@ -212,12 +212,56 @@ private void EncodeSegmentInto(ReadOnlySpan segment, List dest) public string Decode(ReadOnlySpan tokenIds) { - // GPT-2 decode: every char in a token string is a GPT-2-encoded byte. - // Map each char back to its byte, then UTF-8 decode the combined byte stream. - int maxBytes = tokenIds.Length * 8; - byte[] buf = ArrayPool.Shared.Rent(maxBytes); - int count = 0; + // Materializes the decoded byte stream once into a pooled buffer, then UTF-8 decodes. + // The shared helper is used both here and by TryDecode to avoid duplicating the + // GPT-2 char-to-byte unmapping logic. + byte[] buf = MaterializeBytes(tokenIds, out int count); + string result = Encoding.UTF8.GetString(buf, 0, count); + ArrayPool.Shared.Return(buf); + return result; + } + /// + /// Zero-allocation decode into . Atomic on overflow — + /// pre-computes the UTF-8 char count against the materialized byte stream and returns + /// without touching when the + /// result would not fit. + /// + /// + /// is accepted for interface compatibility but is a no-op + /// for the GPT-2 / tiktoken encoding — it has no SentencePiece-style BOS space marker. + /// + public bool TryDecode(ReadOnlySpan tokenIds, bool stripBosSpace, Span destination, out int charsWritten) + { + _ = stripBosSpace; // no-op for tiktoken; here only to satisfy ITokenizer surface + byte[] buf = MaterializeBytes(tokenIds, out int count); + try + { + int needed = Encoding.UTF8.GetCharCount(buf, 0, count); + if (needed > destination.Length) + { + charsWritten = 0; + return false; + } + charsWritten = Encoding.UTF8.GetChars(buf.AsSpan(0, count), destination); + return true; + } + finally + { + ArrayPool.Shared.Return(buf); + } + } + + /// + /// Maps each GPT-2-encoded token char back to its raw byte value, materializing the + /// concatenated byte stream into a pooled buffer. Caller owns the + /// returned buffer and must return it to . + /// + private byte[] MaterializeBytes(ReadOnlySpan tokenIds, out int count) + { + int maxBytes = Math.Max(16, tokenIds.Length * 8); + byte[] buf = ArrayPool.Shared.Rent(maxBytes); + count = 0; foreach (int id in tokenIds) { if ((uint)id >= (uint)_idToToken.Length) continue; @@ -231,7 +275,6 @@ public string Decode(ReadOnlySpan tokenIds) ArrayPool.Shared.Return(buf); buf = larger; } - // Look up the byte value for this GPT-2 Unicode char. int idx = (int)c; if ((uint)idx < (uint)Gpt2UnicodeToByteTable.Length) { @@ -240,10 +283,7 @@ public string Decode(ReadOnlySpan tokenIds) } } } - - string result = Encoding.UTF8.GetString(buf, 0, count); - ArrayPool.Shared.Return(buf); - return result; + return buf; } public string DecodeToken(int tokenId) diff --git a/src/DotLLM.Tokenizers/Bpe/SentencePieceEncoding.cs b/src/DotLLM.Tokenizers/Bpe/SentencePieceEncoding.cs index 5a03ea0f..8cd70eb9 100644 --- a/src/DotLLM.Tokenizers/Bpe/SentencePieceEncoding.cs +++ b/src/DotLLM.Tokenizers/Bpe/SentencePieceEncoding.cs @@ -128,6 +128,83 @@ public string Decode(ReadOnlySpan tokenIds, bool stripBosSpace) return stripLeading ? sb.ToString(1, sb.Length - 1) : sb.ToString(); } + /// + /// Zero-allocation decode into . Writes tokens with a cursor + /// (▁ → space in-place) and flushes accumulated byte-fallback runs via + /// . Returns when the + /// destination is too small — the contents of are + /// unspecified on failure (the cursor may have already written some characters). + /// + public bool TryDecode(ReadOnlySpan tokenIds, bool stripBosSpace, Span destination, out int charsWritten) + { + byte[]? byteBuffer = null; + int byteCount = 0; + int cursor = 0; + + try + { + foreach (int id in tokenIds) + { + if ((uint)id >= (uint)_idToToken.Length) continue; + string token = _idToToken[id]; + if (BpeCore.IsByteToken(token, out byte b)) + { + byteBuffer ??= ArrayPool.Shared.Rent(16); + if (byteCount >= byteBuffer.Length) + { + byte[] larger = ArrayPool.Shared.Rent(byteBuffer.Length * 2); + byteBuffer.AsSpan(0, byteCount).CopyTo(larger); + ArrayPool.Shared.Return(byteBuffer); + byteBuffer = larger; + } + byteBuffer[byteCount++] = b; + } + else + { + if (!BpeCore.TryFlushByteBuffer(destination, ref cursor, byteBuffer, ref byteCount)) + { + charsWritten = 0; + return false; + } + + if (cursor + token.Length > destination.Length) + { + charsWritten = 0; + return false; + } + + Span slot = destination.Slice(cursor, token.Length); + token.AsSpan().CopyTo(slot); + // ▁ → space in place. MemoryExtensions.Replace is a vectorized scan. + MemoryExtensions.Replace(slot, SpaceMarker, ' '); + cursor += token.Length; + } + } + + if (!BpeCore.TryFlushByteBuffer(destination, ref cursor, byteBuffer, ref byteCount)) + { + charsWritten = 0; + return false; + } + + // Strip the single leading space introduced by ▁ prepending. In-place left-shift + // — the only mutation required to honour stripBosSpace semantics on a span cursor. + if (stripBosSpace && _addBosSpace && cursor > 0 && destination[0] == ' ') + { + destination.Slice(1, cursor - 1).CopyTo(destination); + cursor--; + } + + charsWritten = cursor; + return true; + } + finally + { + if (byteBuffer is not null) + ArrayPool.Shared.Return(byteBuffer); + } + } + public string DecodeToken(int tokenId) { if ((uint)tokenId >= (uint)_idToToken.Length) return string.Empty; diff --git a/src/DotLLM.Tokenizers/ITokenizer.cs b/src/DotLLM.Tokenizers/ITokenizer.cs index a8aea26a..60b8bca4 100644 --- a/src/DotLLM.Tokenizers/ITokenizer.cs +++ b/src/DotLLM.Tokenizers/ITokenizer.cs @@ -26,6 +26,46 @@ public interface ITokenizer /// Decoded text. string Decode(ReadOnlySpan tokenIds, bool stripBosSpace) => Decode(tokenIds); + /// + /// Attempts to decode a sequence of token IDs into the caller-provided destination span + /// without allocating an intermediate . Returns + /// when the decoded text fit in (with the character count + /// written to ), or when the + /// destination was too small. + /// + /// + /// + /// The default interface method forwards to + /// and copies into — providing the surface for every + /// implementation, but only the implementations that override + /// this method gain the zero-allocation benefit. + /// overrides it for both the tiktoken (Llama-3, GPT-4, Qwen-2) and SentencePiece (Llama-1/2, + /// Mistral, TinyLlama) encodings. + /// + /// + /// On a return, the contents of + /// are unspecified — callers must either retry with a larger buffer or fall back to + /// the allocating overload. + /// + /// + /// Token IDs to decode. + /// When , strips the leading space introduced by SentencePiece BOS ▁ prepending. + /// Destination character buffer. May be empty when is empty. + /// Number of characters written to on success; 0 on failure. + /// if the decoded text fit; otherwise . + bool TryDecode(ReadOnlySpan tokenIds, bool stripBosSpace, Span destination, out int charsWritten) + { + string decoded = Decode(tokenIds, stripBosSpace); + if (decoded.Length > destination.Length) + { + charsWritten = 0; + return false; + } + decoded.AsSpan().CopyTo(destination); + charsWritten = decoded.Length; + return true; + } + /// Decodes a single token ID to its string representation. /// Token ID to decode. /// String representation of the token. diff --git a/tests/DotLLM.Tests.Unit/Engine/IncrementalDetokenizerTests.cs b/tests/DotLLM.Tests.Unit/Engine/IncrementalDetokenizerTests.cs index fa5f17a0..d0e6de6d 100644 --- a/tests/DotLLM.Tests.Unit/Engine/IncrementalDetokenizerTests.cs +++ b/tests/DotLLM.Tests.Unit/Engine/IncrementalDetokenizerTests.cs @@ -183,4 +183,40 @@ public void ByteFallbackTokens_LongRun_MatchesBulkDecode() Assert.Equal(bulk, detok.ToString()); } + + /// + /// Steady-state Append() must not allocate the per-token Decode() that + /// the pre-PR implementation produced. The committed + /// still grows in chunks and TakeDelta()/ToString() materialize strings on demand, so we + /// don't assert zero bytes — but a small, bounded ceiling proves the per-Append decode + /// allocations are gone. The pre-PR implementation allocated ~2 small strings per Append + /// (one for `_windowText`, one for the eviction-tail decode); at 200 calls × small-window-size + /// that comfortably exceeded 6 KB of Gen-0 traffic. + /// + [Fact] + public void Append_SteadyState_AllocationFloorIsBoundedAndFarBelowPreviousPath() + { + var tok = BuildSpaceMarkerVocab(); + // Cycle through h/e/l/l/o tokens — purely scalar single-char tokens so the window + // resolves cleanly each step and no byte-fallback run is held back. + int[] cycle = [2, 3, 4, 4, 5]; + + using var detok = new IncrementalDetokenizer(tok); + // Warm-up: pay any one-off ArrayPool rent / first-call costs outside the window. + for (int i = 0; i < 32; i++) + detok.Append(cycle[i % cycle.Length]); + + long before = GC.GetAllocatedBytesForCurrentThread(); + for (int i = 0; i < 200; i++) + detok.Append(cycle[i % cycle.Length]); + long after = GC.GetAllocatedBytesForCurrentThread(); + + long delta = after - before; + // 200 Append calls under the previous path = ~400 small string allocs ~= O(few KB). + // The new path's only steady-state allocation source is the committed StringBuilder + // chunk growth (amortized — typically zero across 200 calls once warm). Assert a + // conservative ceiling well below the previous baseline. + Assert.True(delta < 2_048, + $"Append() allocated {delta} bytes across 200 calls — expected far less than the pre-PR ~4 KB+ baseline"); + } } diff --git a/tests/DotLLM.Tests.Unit/Tokenizers/TryDecodeZeroAllocTests.cs b/tests/DotLLM.Tests.Unit/Tokenizers/TryDecodeZeroAllocTests.cs new file mode 100644 index 00000000..31701f98 --- /dev/null +++ b/tests/DotLLM.Tests.Unit/Tokenizers/TryDecodeZeroAllocTests.cs @@ -0,0 +1,334 @@ +using DotLLM.Tokenizers; +using DotLLM.Tokenizers.Bpe; +using Xunit; + +namespace DotLLM.Tests.Unit.Tokenizers; + +/// +/// Tests for the zero-allocation overload. +/// Verifies bit-exact parity with the allocating +/// path across the SentencePiece and tiktoken (GPT-2) BPE encodings, and confirms that the +/// hot decode path produces zero managed allocation per call when the destination buffer is +/// pre-sized adequately. +/// +public sealed class TryDecodeZeroAllocTests +{ + // ------------------------------------------------------------------------- + // Vocabulary factories + // ------------------------------------------------------------------------- + + private static BpeTokenizer BuildSpaceMarkerVocab(bool addBosSpace) + { + string[] tokens = + [ + "", // 0 + "▁", // 1 ▁ + "h", // 2 + "e", // 3 + "l", // 4 + "o", // 5 + "▁h", // 6 ▁h + "▁he", // 7 ▁he + "▁hel", // 8 ▁hel + "▁hell", // 9 ▁hell + "▁hello", // 10 ▁hello + "▁world", // 11 ▁world + "▁w", // 12 ▁w + "w", // 13 + "r", // 14 + "d", // 15 + ]; + float[] scores = new float[tokens.Length]; + return BpeTokenizer.CreateSentencePiece(tokens, scores, tokenTypes: null, + bosId: 0, eosId: 0, addBosSpace: addBosSpace); + } + + private static BpeTokenizer BuildByteFallbackVocab() + { + var tokens = new List { "", "a" }; + for (int i = 0; i < 256; i++) + tokens.Add($"<0x{i:X2}>"); + float[] scores = new float[tokens.Count]; + return BpeTokenizer.CreateSentencePiece(tokens.ToArray(), scores, tokenTypes: null, + bosId: 0, eosId: 0, addBosSpace: false); + } + + private static BpeTokenizer BuildTiktokenVocab() + { + // Minimal GPT-2 / tiktoken vocab: ASCII chars + a couple of merges. + // Each char in a token string is interpreted as one byte via the GPT-2 byte-to-unicode mapping. + string[] tokens = + [ + "", // 0 + "h", // 1 + "e", // 2 + "l", // 3 + "o", // 4 + " ", // 5 (raw space is one byte) + "w", // 6 + "r", // 7 + "d", // 8 + "he", // 9 + "lo", // 10 + "hello", // 11 + " w", // 12 + "wo", // 13 + "rl", // 14 + "world", // 15 + ]; + 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); + } + + // ------------------------------------------------------------------------- + // Parity tests — TryDecode output must equal Decode output + // ------------------------------------------------------------------------- + + [Theory] + [InlineData(true)] + [InlineData(false)] + public void Parity_SentencePiece_ShortSequence(bool stripBosSpace) + { + BpeTokenizer tok = BuildSpaceMarkerVocab(addBosSpace: true); + int[] ids = [10, 1, 11]; // "▁hello" + "▁" + "▁world" + + string expected = tok.Decode(ids, stripBosSpace); + + Span buffer = stackalloc char[64]; + bool ok = tok.TryDecode(ids, stripBosSpace, buffer, out int written); + + Assert.True(ok); + Assert.Equal(expected, buffer[..written].ToString()); + } + + [Theory] + [InlineData(true)] + [InlineData(false)] + public void Parity_SentencePiece_NoBosSpace(bool stripBosSpace) + { + // addBosSpace=false means the encoding never prepends ▁ — stripBosSpace should be a no-op. + BpeTokenizer tok = BuildSpaceMarkerVocab(addBosSpace: false); + int[] ids = [2, 3, 4, 4, 5, 1, 13, 5, 14, 4, 15]; // "hello" + "▁" + "world" + + string expected = tok.Decode(ids, stripBosSpace); + + Span buffer = stackalloc char[64]; + bool ok = tok.TryDecode(ids, stripBosSpace, buffer, out int written); + + Assert.True(ok); + Assert.Equal(expected, buffer[..written].ToString()); + } + + [Fact] + public void Parity_SentencePiece_ByteFallbackRun() + { + BpeTokenizer tok = BuildByteFallbackVocab(); + const int aId = 1; + int byteC3 = 2 + 0xC3; + int byteA9 = 2 + 0xA9; + int[] ids = [aId, byteC3, byteA9, aId]; // "aéa" + + string expected = tok.Decode(ids, stripBosSpace: false); + + Span buffer = stackalloc char[32]; + bool ok = tok.TryDecode(ids, stripBosSpace: false, buffer, out int written); + + Assert.True(ok); + Assert.Equal(expected, buffer[..written].ToString()); + } + + [Fact] + public void Parity_SentencePiece_RandomSequences() + { + BpeTokenizer tok = BuildSpaceMarkerVocab(addBosSpace: true); + var rng = new Random(42); + Span buffer = stackalloc char[256]; + + for (int trial = 0; trial < 50; trial++) + { + int len = rng.Next(1, 24); + int[] ids = new int[len]; + for (int i = 0; i < len; i++) ids[i] = rng.Next(1, 16); + + foreach (bool stripBosSpace in new[] { true, false }) + { + string expected = tok.Decode(ids, stripBosSpace); + bool ok = tok.TryDecode(ids, stripBosSpace, buffer, out int written); + Assert.True(ok, $"TryDecode unexpectedly returned false on trial {trial} (stripBosSpace={stripBosSpace})"); + Assert.Equal(expected, buffer[..written].ToString()); + } + } + } + + [Theory] + [InlineData(true)] + [InlineData(false)] + public void Parity_Tiktoken_ShortSequence(bool stripBosSpace) + { + BpeTokenizer tok = BuildTiktokenVocab(); + int[] ids = [11, 5, 15]; // "hello" + " " + "world" + + string expected = tok.Decode(ids, stripBosSpace); + + Span buffer = stackalloc char[64]; + bool ok = tok.TryDecode(ids, stripBosSpace, buffer, out int written); + + Assert.True(ok); + Assert.Equal(expected, buffer[..written].ToString()); + } + + [Fact] + public void Parity_Tiktoken_RandomSequences() + { + BpeTokenizer tok = BuildTiktokenVocab(); + var rng = new Random(123); + Span buffer = stackalloc char[256]; + + for (int trial = 0; trial < 50; trial++) + { + int len = rng.Next(1, 20); + int[] ids = new int[len]; + for (int i = 0; i < len; i++) ids[i] = rng.Next(1, 16); + + foreach (bool stripBosSpace in new[] { true, false }) + { + string expected = tok.Decode(ids, stripBosSpace); + bool ok = tok.TryDecode(ids, stripBosSpace, buffer, out int written); + Assert.True(ok, $"TryDecode unexpectedly returned false on trial {trial}"); + Assert.Equal(expected, buffer[..written].ToString()); + } + } + } + + // ------------------------------------------------------------------------- + // Buffer-too-small contract + // ------------------------------------------------------------------------- + + [Fact] + public void BufferTooSmall_SentencePiece_ReturnsFalse_WrittenIsZero() + { + BpeTokenizer tok = BuildSpaceMarkerVocab(addBosSpace: false); + int[] ids = [10, 1, 11]; // expands to >2 chars + + Span tiny = stackalloc char[2]; + bool ok = tok.TryDecode(ids, stripBosSpace: false, tiny, out int written); + + Assert.False(ok); + Assert.Equal(0, written); + } + + [Fact] + public void BufferTooSmall_Tiktoken_ReturnsFalse_WrittenIsZero_AndAtomic() + { + BpeTokenizer tok = BuildTiktokenVocab(); + int[] ids = [11, 5, 15]; // "hello world" + + Span tiny = stackalloc char[3]; + // Pre-fill the buffer with a sentinel to verify atomicity (tiktoken pre-computes char count). + tiny.Fill('X'); + bool ok = tok.TryDecode(ids, stripBosSpace: false, tiny, out int written); + + Assert.False(ok); + Assert.Equal(0, written); + Assert.Equal("XXX", tiny.ToString()); + } + + [Fact] + public void EmptyInput_ReturnsTrue_WrittenIsZero() + { + BpeTokenizer tok = BuildSpaceMarkerVocab(addBosSpace: true); + + Span buffer = stackalloc char[8]; + bool ok = tok.TryDecode([], stripBosSpace: true, buffer, out int written); + + Assert.True(ok); + Assert.Equal(0, written); + } + + // ------------------------------------------------------------------------- + // Zero-allocation assertions (steady-state) + // ------------------------------------------------------------------------- + + [Fact] + public void TryDecode_SentencePiece_ZeroAllocationPerCall() + { + BpeTokenizer tok = BuildSpaceMarkerVocab(addBosSpace: false); + int[] ids = [10, 1, 11]; + Span buffer = stackalloc char[128]; + + // Warm-up: pay JIT + first-call ArrayPool rent costs outside the measurement window. + for (int i = 0; i < 32; i++) + { + bool ok = tok.TryDecode(ids, stripBosSpace: false, buffer, out _); + Assert.True(ok); + } + + long before = GC.GetAllocatedBytesForCurrentThread(); + for (int i = 0; i < 1000; i++) + { + bool ok = tok.TryDecode(ids, stripBosSpace: false, buffer, out _); + Assert.True(ok); + } + long after = GC.GetAllocatedBytesForCurrentThread(); + + long delta = after - before; + // Expect bit-exact zero managed allocation on the steady-state hot path. + Assert.True(delta == 0, + $"Expected 0 managed bytes allocated across 1000 TryDecode calls, got {delta} bytes"); + } + + [Fact] + public void TryDecode_Tiktoken_ZeroAllocationPerCall() + { + BpeTokenizer tok = BuildTiktokenVocab(); + int[] ids = [11, 5, 15]; + Span buffer = stackalloc char[128]; + + for (int i = 0; i < 32; i++) + { + bool ok = tok.TryDecode(ids, stripBosSpace: false, buffer, out _); + Assert.True(ok); + } + + long before = GC.GetAllocatedBytesForCurrentThread(); + for (int i = 0; i < 1000; i++) + { + bool ok = tok.TryDecode(ids, stripBosSpace: false, buffer, out _); + Assert.True(ok); + } + long after = GC.GetAllocatedBytesForCurrentThread(); + + long delta = after - before; + Assert.True(delta == 0, + $"Expected 0 managed bytes allocated across 1000 TryDecode calls, got {delta} bytes"); + } + + [Fact] + public void Decode_SentencePiece_AllocatesAtLeastOneStringPerCall_Baseline() + { + // Sanity check: the existing allocating Decode path should allocate per call. + // This is the baseline the TryDecode improvement is measured against — without it, + // a passing zero-alloc test for TryDecode would be meaningless. + BpeTokenizer tok = BuildSpaceMarkerVocab(addBosSpace: false); + int[] ids = [10, 1, 11]; + + for (int i = 0; i < 32; i++) _ = tok.Decode(ids, stripBosSpace: false); + + long before = GC.GetAllocatedBytesForCurrentThread(); + for (int i = 0; i < 1000; i++) _ = tok.Decode(ids, stripBosSpace: false); + long after = GC.GetAllocatedBytesForCurrentThread(); + + long delta = after - before; + Assert.True(delta > 0, + $"Sanity check failed: allocating Decode reported {delta} bytes — expected > 0"); + } +} From 3a5d34e168543e8566e90e261dacca4c388b9d9a Mon Sep 17 00:00:00 2001 From: James Burton Date: Fri, 31 Jul 2026 11:10:24 +0100 Subject: [PATCH 2/2] tokenizers: address Copilot review on TryDecode zero-alloc PR (#121) Four of six review comments applied, two answered in place: - TextGenerator (sync + streaming): construct IncrementalDetokenizer inside the try with a nullable local, so a constructor failure still runs the finally that returns the rented stopScratch buffer. - IncrementalDetokenizer.Dispose + both grow-and-re-rent paths, and the TextGenerator stopScratch returns, now use Return(clearArray: true) so decoded model output is not left readable by an unrelated renter of ArrayPool.Shared. Off the per-token path; buffers are small. - DetokenizerAllocBenchmark: validate in GlobalSetup that the scratch buffer fits both decodes, so a vocab change fails fast instead of silently benchmarking the buffer-too-small failure path. Measured loops unchanged. - TryDecodeZeroAllocTests: replace the bit-exact `delta == 0` assertion with a sub-one-byte-per-call budget (1000 bytes / 1000 calls). The minimum managed object is 24 bytes, so this still rules out any per-call allocation while not flaking on one-off runtime bookkeeping. - IncrementalDetokenizerTests: pre-size the committed StringBuilder (4096) so StringBuilder chunk growth cannot perturb the measurement; the existing 2 KB ceiling is kept rather than loosened. --- .../Tokenizers/DetokenizerAllocBenchmark.cs | 15 ++++++++++ src/DotLLM.Engine/IncrementalDetokenizer.cs | 18 ++++++++---- src/DotLLM.Engine/TextGenerator.cs | 28 +++++++++++++------ .../Engine/IncrementalDetokenizerTests.cs | 12 +++++--- .../Tokenizers/TryDecodeZeroAllocTests.cs | 23 +++++++++++---- 5 files changed, 73 insertions(+), 23 deletions(-) diff --git a/benchmarks/DotLLM.Benchmarks/Tokenizers/DetokenizerAllocBenchmark.cs b/benchmarks/DotLLM.Benchmarks/Tokenizers/DetokenizerAllocBenchmark.cs index f309c242..0c15ee18 100644 --- a/benchmarks/DotLLM.Benchmarks/Tokenizers/DetokenizerAllocBenchmark.cs +++ b/benchmarks/DotLLM.Benchmarks/Tokenizers/DetokenizerAllocBenchmark.cs @@ -33,6 +33,21 @@ public void Setup() _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)")] diff --git a/src/DotLLM.Engine/IncrementalDetokenizer.cs b/src/DotLLM.Engine/IncrementalDetokenizer.cs index 12fa01f9..4538b4a4 100644 --- a/src/DotLLM.Engine/IncrementalDetokenizer.cs +++ b/src/DotLLM.Engine/IncrementalDetokenizer.cs @@ -91,7 +91,7 @@ private void DecodeWindowInto(ref char[] buffer, out int written) // 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.Shared.Rent(buffer.Length * 2); - ArrayPool.Shared.Return(buffer); + ArrayPool.Shared.Return(buffer, clearArray: true); buffer = larger; } } @@ -104,7 +104,7 @@ private void DecodeTailInto(out int written) if (_tokenizer.TryDecode(ids, stripBosSpace: false, _tailBuf, out written)) return; char[] larger = ArrayPool.Shared.Rent(_tailBuf.Length * 2); - ArrayPool.Shared.Return(_tailBuf); + ArrayPool.Shared.Return(_tailBuf, clearArray: true); _tailBuf = larger; } } @@ -194,17 +194,25 @@ public override string ToString() }); } - /// Returns pooled buffers to . Idempotent. + /// + /// Returns pooled buffers to . Idempotent. + /// + /// + /// Buffers are returned with clearArray: true 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. + /// public void Dispose() { if (_windowBuf is { Length: > 0 } wb) { - ArrayPool.Shared.Return(wb); + ArrayPool.Shared.Return(wb, clearArray: true); _windowBuf = []; } if (_tailBuf is { Length: > 0 } tb) { - ArrayPool.Shared.Return(tb); + ArrayPool.Shared.Return(tb, clearArray: true); _tailBuf = []; } _windowLen = 0; diff --git a/src/DotLLM.Engine/TextGenerator.cs b/src/DotLLM.Engine/TextGenerator.cs index 3902ba37..83927bac 100644 --- a/src/DotLLM.Engine/TextGenerator.cs +++ b/src/DotLLM.Engine/TextGenerator.cs @@ -141,11 +141,14 @@ public InferenceResponse Generate(string prompt, InferenceOptions? options = nul // Incremental detokenizer keeps stop-check cost O(1) amortized per token // instead of decoding the entire generated sequence each step (O(n²)). - // Declared outside the try so the finally can deterministically return its pooled buffers. - var detok = new IncrementalDetokenizer(_tokenizer, initialCapacity: Math.Max(64, maxTokens * 4)); + // Declared outside the try but constructed inside it, so a failure in the constructor + // (which itself rents from the pool) still runs the finally that returns stopScratch. + IncrementalDetokenizer? detok = null; try { + detok = new IncrementalDetokenizer(_tokenizer, initialCapacity: Math.Max(64, maxTokens * 4)); + var generatedIds = new List(maxTokens); var finishReason = FinishReason.Length; long prefillTicks = 0; @@ -399,8 +402,10 @@ public InferenceResponse Generate(string prompt, InferenceOptions? options = nul } finally { - ArrayPool.Shared.Return(stopScratch); - detok.Dispose(); + // clearArray: the scratch holds decoded model output; don't leave it readable + // by an unrelated renter of the shared pool. + ArrayPool.Shared.Return(stopScratch, clearArray: true); + detok?.Dispose(); if (ownsKvCache) kvCache.Dispose(); } @@ -484,12 +489,15 @@ public async IAsyncEnumerable GenerateStreamingTokensAsync( char[] stopScratch = ArrayPool.Shared.Rent(stopTailSize); // Incremental detokenizer: O(1) amortized per token for stop-check + streaming delta, - // instead of decoding the full generated sequence at every step. Lifted out of the try - // so the finally can return its pooled buffers even on cancellation. - var detok = new IncrementalDetokenizer(_tokenizer, initialCapacity: Math.Max(64, maxTokens * 4)); + // instead of decoding the full generated sequence at every step. Declared outside the try + // so the finally can return its pooled buffers even on cancellation, but constructed + // inside it so a constructor failure still returns stopScratch. + IncrementalDetokenizer? detok = null; try { + detok = new IncrementalDetokenizer(_tokenizer, initialCapacity: Math.Max(64, maxTokens * 4)); + var generatedIds = new List(maxTokens); long prefillTicks = 0; long decodeTicks = 0; @@ -798,8 +806,10 @@ public async IAsyncEnumerable GenerateStreamingTokensAsync( } finally { - ArrayPool.Shared.Return(stopScratch); - detok.Dispose(); + // clearArray: the scratch holds decoded model output; don't leave it readable + // by an unrelated renter of the shared pool. + ArrayPool.Shared.Return(stopScratch, clearArray: true); + detok?.Dispose(); if (ownsKvCache) kvCache.Dispose(); } diff --git a/tests/DotLLM.Tests.Unit/Engine/IncrementalDetokenizerTests.cs b/tests/DotLLM.Tests.Unit/Engine/IncrementalDetokenizerTests.cs index d0e6de6d..7fb06eb9 100644 --- a/tests/DotLLM.Tests.Unit/Engine/IncrementalDetokenizerTests.cs +++ b/tests/DotLLM.Tests.Unit/Engine/IncrementalDetokenizerTests.cs @@ -201,7 +201,10 @@ public void Append_SteadyState_AllocationFloorIsBoundedAndFarBelowPreviousPath() // resolves cleanly each step and no byte-fallback run is held back. int[] cycle = [2, 3, 4, 4, 5]; - using var detok = new IncrementalDetokenizer(tok); + // Pre-size the committed StringBuilder well beyond the ~232 chars this test produces, + // so incidental chunk growth cannot contribute to the measured delta and the assertion + // reflects only per-Append allocation behaviour. + using var detok = new IncrementalDetokenizer(tok, initialCapacity: 4096); // Warm-up: pay any one-off ArrayPool rent / first-call costs outside the window. for (int i = 0; i < 32; i++) detok.Append(cycle[i % cycle.Length]); @@ -213,9 +216,10 @@ public void Append_SteadyState_AllocationFloorIsBoundedAndFarBelowPreviousPath() long delta = after - before; // 200 Append calls under the previous path = ~400 small string allocs ~= O(few KB). - // The new path's only steady-state allocation source is the committed StringBuilder - // chunk growth (amortized — typically zero across 200 calls once warm). Assert a - // conservative ceiling well below the previous baseline. + // With the buffer pre-sized above, the new path has no steady-state allocation source at + // all; the ceiling stays deliberately generous (10 bytes/call vs the previous ~20+) so the + // test cannot flake on runtime bookkeeping while still failing loudly if per-Append + // string decoding ever returns. Assert.True(delta < 2_048, $"Append() allocated {delta} bytes across 200 calls — expected far less than the pre-PR ~4 KB+ baseline"); } diff --git a/tests/DotLLM.Tests.Unit/Tokenizers/TryDecodeZeroAllocTests.cs b/tests/DotLLM.Tests.Unit/Tokenizers/TryDecodeZeroAllocTests.cs index 31701f98..3fcf0068 100644 --- a/tests/DotLLM.Tests.Unit/Tokenizers/TryDecodeZeroAllocTests.cs +++ b/tests/DotLLM.Tests.Unit/Tokenizers/TryDecodeZeroAllocTests.cs @@ -258,6 +258,13 @@ public void EmptyInput_ReturnsTrue_WrittenIsZero() // Zero-allocation assertions (steady-state) // ------------------------------------------------------------------------- + /// + /// Total managed bytes tolerated across a warmed 1000-call loop. Below one byte per call: + /// since the minimum managed object size is 24 bytes, staying under this bound proves the + /// measured path allocates nothing per call, without being brittle to one-off runtime noise. + /// + private const long AllocBudgetPer1000Calls = 1000; + [Fact] public void TryDecode_SentencePiece_ZeroAllocationPerCall() { @@ -281,9 +288,14 @@ public void TryDecode_SentencePiece_ZeroAllocationPerCall() long after = GC.GetAllocatedBytesForCurrentThread(); long delta = after - before; - // Expect bit-exact zero managed allocation on the steady-state hot path. - Assert.True(delta == 0, - $"Expected 0 managed bytes allocated across 1000 TryDecode calls, got {delta} bytes"); + // Budget rather than bit-exact zero: the smallest possible managed allocation is 24 bytes, + // so a total under 1000 bytes across 1000 calls mathematically rules out *any* per-call + // allocation, while absorbing one-off runtime bookkeeping (tiered re-JIT, diagnostics) + // that can otherwise make a strict `== 0` assertion flaky across TFMs/configurations. + // For reference, the allocating Decode path costs tens of bytes *per call* here. + Assert.True(delta < AllocBudgetPer1000Calls, + $"Expected no per-call managed allocation across 1000 TryDecode calls, got {delta} bytes " + + $"({delta / 1000.0:F2} bytes/call)"); } [Fact] @@ -308,8 +320,9 @@ public void TryDecode_Tiktoken_ZeroAllocationPerCall() long after = GC.GetAllocatedBytesForCurrentThread(); long delta = after - before; - Assert.True(delta == 0, - $"Expected 0 managed bytes allocated across 1000 TryDecode calls, got {delta} bytes"); + Assert.True(delta < AllocBudgetPer1000Calls, + $"Expected no per-call managed allocation across 1000 TryDecode calls, got {delta} bytes " + + $"({delta / 1000.0:F2} bytes/call)"); } [Fact]