Wave 8: manual review fixes — speculative correctness, server semantics, cache lifecycle, hot-path allocations

[kkokosa]

Scope

Pre-public-release audit response. A manual review surfaced correctness, semantics, and hot-path issues in the orchestration/serving layers (not in the core transformer forward path). Items here were triaged as either future work or as ones that can ship behind a smaller guard in the preceding PR. The greedy-only guard and doc-demotion PR (referenced below) ships first; the items in this issue are the actual implementations.

1. Non-greedy speculative decoding is not distributionally correct

File: src/DotLLM.Engine/SpeculativeDecoder.cs:101,103,199

q (draft probability) is taken from TensorPrimitives.SoftMax on constraint-masked raw logits at line 103, and p (target probability) from raw softmax at line 199. The actual draft token at line 105 is sampled via the full SamplerPipeline (temperature / top-k / top-p / min-p / repetition penalty). That means q/p do not describe the distribution the pipeline actually samples from, so modified rejection sampling does not reproduce the target distribution. Only greedy matches because argmax is invariant to monotonic pruning — and only when repetition penalty is 1.0.

Fix: Compute q and p from the post-transform distribution the SamplerPipeline actually samples from (apply the same chain of transforms to both draft and target logits before the softmax). Add non-greedy unit tests comparing empirical output distribution against the target model across seeds. Once non-greedy is correct, remove the constructor guard and tightened call-site predicate added in the preceding PR.

2. Prompt-cache undersize-but-fits-prompt branch caps generation silently

File: src/DotLLM.Engine/TextGenerator.cs:824

if (entry.KvCache.MaxLength >= requiredSize || entry.KvCache.MaxLength >= promptLen)
    return (entry.KvCache, matchedTokens, false);

The second clause accepts any cache that fits the prompt. If a later request has larger max_tokens, decode runs out of cache capacity mid-generation and the response ends early as if the model simply stopped.

Fix: Delete the || entry.KvCache.MaxLength >= promptLen clause. Undersized caches must fall through to reallocation. Regression test: prime the prefix cache with a small max_tokens, then issue a larger-max_tokens request with the same prompt and assert the response reaches max_tokens (or a genuine stop) rather than the cached length.

3. tool_choice is parsed and ignored

File: src/DotLLM.Server/Endpoints/ChatCompletionEndpoint.cs:58

RequestConverter.ParseToolChoice runs, but the parsed value is never used. A client sending tool_choice: "required" or a specific function gets "auto" behavior silently.

Fix (short-term): Until constraint-driven enforcement lands, reject tool_choice values other than "auto" (or unset) with HTTP 400 and an OpenAI-shaped {"error": {"type": "invalid_request_error", ...}} body. Track "wire required / specific-function enforcement via constrained decoding" as a follow-up sub-item of this bullet.

4. Streaming tool-call text leaks as delta.content

File: src/DotLLM.Server/Endpoints/ChatCompletionEndpoint.cs:213-249

Every generated token is streamed as normal delta.content, and tool-call parsing runs only after the stream completes (line 246). Clients see raw <tool_call>… markup mid-stream and only get the reclassified finish_reason: "tool_calls" at the end. Non-streaming path is correct.

Fix: Maintain an incremental parser with a small rolling buffer. Suppress delta.content for tokens inside a detected tool-call region, emit delta.tool_calls fragments per the OpenAI SSE contract as arguments are parsed.

5. Empty-text tokens under-count completion_tokens in SSE

File: src/DotLLM.Server/Endpoints/ChatCompletionEndpoint.cs:213

completionTokens is only incremented when token.Text.Length > 0, but GenerationToken (src/DotLLM.Engine/GenerationToken.cs:7) explicitly allows empty Text for incomplete UTF-8 continuation bytes. Usage is undercounted relative to the non-streaming path.

Fix: Count every yielded GenerationToken in completionTokens. Only gate the SSE delta.content emission on Text.Length > 0.

6. Prefix-cache ownership is wrong on cache miss — leaks on throw

File: src/DotLLM.Engine/TextGenerator.cs:834

A freshly allocated cache is returned with ownsKvCache = false before StoreInPrefixCache (line 874) actually transfers ownership. If generation throws between allocation and store (exception, cancellation), the finally blocks at TextGenerator.cs:387,773 skip disposal and the cache leaks.

Fix: Return ownsKvCache = true from the miss path. Flip to false only after _prefixCache.Store successfully runs (the assignment already exists at line 875). Regression test: cancel decode between allocation and store, assert the cache is disposed.

7. Prefix cache stores non-reusable cache types

Files: src/DotLLM.Engine/TextGenerator.cs:809,874

ResolveKvCache reuses only SimpleKvCache and PagedKvCache; the default branch at line 817 falls through to reallocation. But StoreInPrefixCache at line 874 stores any cache instance — quantized, GPU, custom factories. Those entries will never be reused and still hold RAM/VRAM.

Fix: Gate StoreInPrefixCache on the same type check ResolveKvCache uses (or lift the check to a shared helper / an IKvCache.SupportsPrefixReuse property).

8. StopStringCondition drops the whole last token instead of trimming to the matched suffix

Files: src/DotLLM.Engine/Samplers/StopConditions/StopStringCondition.cs:23 and src/DotLLM.Engine/TextGenerator.cs:129

Stop-string detection happens at token boundaries. When a match is found, the entire last token is removed from the generated output, not just the matched suffix. For BPE tokenizers that regularly merge non-stop-preceding text into the same token as the stop sequence, valid content can be lost. The code comment at line 129 acknowledges this.

Fix: Track token-boundary offsets in the decoded string during generation. On stop-match, compute the byte/char offset of the matched suffix start within the last token and trim exactly to that offset, not the whole token.

9. Full-sequence detokenization on every step — O(n²) — DONE in PR #122

Resolved by IncrementalDetokenizer (sliding window + committed StringBuilder with suffix-based eviction) + IStopCondition.ShouldStop(ReadOnlySpan<char>) overload. Six per-step Decode(generatedIds) sites in TextGenerator replaced; streaming delta is TakeDelta().

10. Generate / GenerateStreamingTokensAsync duplicate ~all decode orchestration

Files: src/DotLLM.Engine/TextGenerator.cs:68 (Generate), src/DotLLM.Engine/TextGenerator.cs:400 (GenerateStreamingTokensAsync)

Both methods re-implement prefill, stop handling, speculative decode, cache storage, and timings independently. Behavioral drift risk is already present (e.g., #5 above exists in streaming but not non-streaming at the server layer because the engine path is forked lower down).

Fix: Extract a single internal core — either a private IAsyncEnumerable<GenerationToken> GenerateCore(...) consumed by both the sync (materialize into InferenceResponse) and async (passthrough) public entry points, or a push-style core that invokes a delegate per token. Prefill, stop handling, spec decode, cache store, timings live in one place.

12. Zero-allocation detokenization via TryDecode(..., Span<char>)

Files: src/DotLLM.Tokenizers/ITokenizer.cs, src/DotLLM.Tokenizers/Bpe/BpeCore.cs:11, src/DotLLM.Tokenizers/Bpe/BpeTokenizer.cs:112, src/DotLLM.Tokenizers/Bpe/Gpt2TiktokenEncoding.cs:213, src/DotLLM.Tokenizers/Bpe/SentencePieceEncoding.cs:89, src/DotLLM.Engine/IncrementalDetokenizer.cs

Follow-up from Gemini's review on PR #122. IncrementalDetokenizer.Append and TryEvictOldest each allocate a string per call via _tokenizer.Decode(window, ...). Amortized O(1) per token today; still two small Gen-0 allocations per decode step (~4K strings on a 2K-token generation). Gate for true zero-allocation detokenization.

Fix: Add a span-writing overload to the tokenizer surface:

// ITokenizer + IBpeEncoding
int TryDecode(ReadOnlySpan<int> tokenIds, bool stripBosSpace, Span<char> destination); // returns charsWritten, or -1 if destination too small

Default interface method forwards to Decode and copies — external implementors unchanged.

Implementations:

• Gpt2TiktokenEncoding (~1–2 h): accumulate bytes into existing pooled byte buffer, then Encoding.UTF8.GetChars(ReadOnlySpan<byte>, Span<char>) into destination. Replaces the per-call GetString allocation.
• SentencePieceEncoding (~2–4 h): rewrite StringBuilder logic to a span cursor. Per-token write with ▁→space in-place replacement; flush byte-token runs via Encoding.UTF8.GetChars(bytes, dest.Slice(cursor)). BpeCore.FlushByteBuffer gains a span overload.
• BpeTokenizer: one-line forward.

Adoption in IncrementalDetokenizer (~2–3 h):

• Replace string _windowText with rented char[] _windowBuf + int _windowLen (ArrayPool-backed, grow-and-re-rent on overflow).
• Append: compute upper-bound destination size (sum of _idToToken[id].Length for window); TryDecode into buffer.
• TryEvictOldest: rent a second buffer; compare spans with MemoryExtensions.EndsWith — no string materialization.
• Add Dispose to release pooled buffers; TextGenerator's outer finally calls it.

Testing: TryDecode equivalence tests on both encodings (TryDecode output == Decode output for matching inputs); existing IncrementalDetokenizerTests exercise the adoption end-to-end.

Total: ~1 engineer-day. Eliminates the 2 string allocations per decode step that remain after PR #122.

11. Hot model swap leaks PagedFactory

File: src/DotLLM.Server/Endpoints/ModelManagementEndpoint.cs:55

ServerState expects to own and dispose PagedFactory (src/DotLLM.Server/ServerState.cs:37,101), but the hot-swap path never transfers newState.PagedFactory into the live ServerState. If the newly loaded model uses paged KV, the shared block pool detaches from lifecycle management and is not disposed on the next swap or shutdown.

Fix: In ModelManagementEndpoint.SwapAsync (around line 55), transfer newState.PagedFactory (and any other lifecycle-owned handles) into the live ServerState as part of the swap, disposing the previous one.

Acceptance Criteria

• Non-greedy speculative decoding distributionally correct; new unit tests compare empirical draws to target distribution across seeds; constructor guard + call-site predicate from the preceding PR removed
• Prompt-cache size check fixed — regression test: reuse with larger max_tokens than cached capacity produces a full-length response
• Server rejects unsupported tool_choice values with 400; follow-up sub-issue filed for required / specific-function enforcement
• Streaming chat completions: no raw tool-call text in delta.content; tool-call fragments emitted as delta.tool_calls
• Streaming completion_tokens matches non-streaming path on prompts containing multi-byte UTF-8
• Prefix-cache ownership transfer is atomic; cancellation mid-generation disposes the cache (regression test)
• StoreInPrefixCache only stores reusable cache types (shared type check with ResolveKvCache)
• StopStringCondition trims exactly the matched suffix on BPE tokenizers (regression test with a tokenizer that merges adjacent text into the stop token)
• Streaming decode path is O(n) in generated length — done in PR Restrict speculative decoding to greedy, demote overstated features (#121) #122
• Generate and GenerateStreamingTokensAsync share a single internal orchestration core
• Hot model swap disposes the previous PagedFactory and installs the new one under live ServerState ownership

References

• Preceding PR: greedy-only speculative decoding guard + README/docs honesty pass (the minimal shippable subset of the audit)
• docs/SPECULATIVE.md — distributional correctness requirements
• docs/SERVER.md, docs/TOOL_CALLING.md — OpenAI compatibility surface and tool-calling semantics
• docs/KV_CACHE.md, docs/SCHEDULING.md — prefix cache and paged allocator lifecycle
• docs/SAMPLING.md — stop-condition semantics

[jamesburton]

Starting work on this. Plan to ship in focused PRs (this comment lists what I'll attempt; I'll open PRs separately):

• Primary: item Initial scaffolding #1 — non-greedy distributional correctness in SpeculativeDecoder. Branch: issue/121-speculative-non-greedy-correctness. Approach: add ApplyTransforms/SampleFromTransformed/IsGreedy to SamplerPipeline; rewrite SpeculativeDecoder so q and p come from the post-transform distribution; rebuild repetition-penalty context in the verify pass to match the draft pass (generatedIds + draft_0…draft_{i-1}); cover bonus token, beyond-shared-vocab fallback, and SampleCorrected via the same path. Removes the greedy-only constructor guard and the IsEffectivelyGreedy gate at both TextGenerator call sites. Adds a chi-squared distributional test against the target pipeline (verified to FAIL on the unfixed implementation when only the guard is removed; passes on the fix with the same seeds). Updates docs/SPECULATIVE.md and README.

Smaller items I'll attempt next as separate focused PRs:

• Implement BPE tokenizer with GGUF vocabulary loading #6 — prefix-cache ownership transfer (ownsKvCache = true on miss path).
• Implement GQA attention with RoPE position encoding #7 — StoreInPrefixCache type gate matching ResolveKvCache.
• Add Claude Code GitHub Workflow #2 — drop the entry.KvCache.MaxLength >= promptLen clause.
• Implement GGUF model loader with memory-mapped tensor access #3 — tool_choice 400 rejection for non-auto values.
• Implement FP16 and Q8_0 dequantization kernels #5 — count empty-text GenerationTokens in SSE.
• Claude development notes #11 — hot-swap PagedFactory lifecycle transfer.

Out of scope for this issue (separate work or already deferred): #4 (incremental tool-call SSE — bigger change), #8 (StopStringCondition trim — needs char-offset tracking), #10 (Generate/streaming consolidation — large refactor), #12 (zero-alloc detok — Gemini follow-up).

[jamesburton]

Picking up item #8 (StopString suffix trim) for the streaming path on branch issue/121-stop-string-streaming (built off issue/107-stop-string-suffix). PR #107 covered non-streaming; this extends the same character-boundary trim semantics to the streaming SSE path via a rolling token-tail buffer so token-boundary-straddling stop matches never leak as delta.content. Will reuse StopSuffixTrimmer from #107.

[jamesburton]

Picking up item #4 (streaming tool-call SSE separation). Branch: jamesburton/dotLLM:issue/121-streaming-tool-call-sse.

[jamesburton]

Starting work on item #12 (zero-allocation TryDecode follow-up to PR #122).

Plan:

• Add bool TryDecode(ReadOnlySpan<int> tokenIds, bool stripBosSpace, Span<char> destination, out int charsWritten) to ITokenizer and IBpeEncoding as a default interface method that forwards to Decode + copy (external impls unchanged).
• Override on Gpt2TiktokenEncoding and SentencePieceEncoding to write directly into the destination span (UTF-8 GetCharCount pre-check for atomic overflow on the tiktoken path; cursor-write with in-place ▁→space replacement on the SPM path).
• Forward through BpeTokenizer and NormalizingTokenizer.
• Adopt in IncrementalDetokenizer by replacing string _windowText with a rented char[] _windowBuf + int _windowLen (grow-and-re-rent on overflow), implement IDisposable, and wire Dispose into TextGenerator's finally blocks.

Branch: issue/121-tokenizer-try-decode-zero-alloc on jamesburton/dotLLM. Will measure allocation delta via GC.GetAllocatedBytesForCurrentThread over a 1000-call steady-state loop.