fix: avoid O(N²) per-chunk reprocessing in long-running StreamingSession - #39
fix: avoid O(N²) per-chunk reprocessing in long-running StreamingSession#39rubdttcom wants to merge 1 commit into
Conversation
…sion process_emitted() re-detokenized all non_special_ and regroup_words() re-grouped all word_tokens_ on every chunk, and neither is trimmed across utterances -- so a long stream degrades to O(N^2) CPU with climbing RSS. Detokenize is prefix-stable, so append only the new tokens' text; re-group only the still-open word tail and keep finalized words. Output is identical (text()/tokens()/drain_words()/drain_events() unchanged), now O(1) amortized per chunk. Adds tests/test_streaming_longrun_bounded.cpp: asserts per-chunk reprocessing stays bounded and that the incremental text equals a full detokenize. Signed-off-by: Rubén Fernández <129730697+rubdttcom@users.noreply.github.com> Assisted-by: Claude:claude-opus-4-8 [Claude Code]
localai-bot
left a comment
There was a problem hiding this comment.
Confirmed the perf motivation is real: a standalone micro-benchmark (no model needed, detokenize/group_words are pure functions of tokens) driving the old full-rebuild path vs this PR's incremental path over identical synthetic token streams shows the old path scaling roughly quadratically and the new path roughly linearly:
| tokens | OLD (full rebuild/chunk) | NEW (incremental) | speedup |
|---|---|---|---|
| 2,500 | 46 ms | 2 ms | 19x |
| 5,000 | 251 ms | 14 ms | 18x |
| 10,000 | 767 ms | 45 ms | 17x |
| 20,000 | 3,291 ms | 144 ms | 23x |
But I found a real bug in the incremental word-grouping (regroup_words) that drops words from the output.
The cursor-advance logic in regroup_words picks the split point using piece_starts_with_meta alone (does this piece's raw bytes start with ▁?). group_words's actual word-start rule is (piece != decoded_text) && !is_punctuation. The shipped tokenizer (tdt_ctc-110m-q4_k.gguf, piece id 582) has a real vocab entry "▁'", a token whose piece starts with ▁ but whose decoded text is just ', which is punctuation. This is a realistic pattern: it's what the model would emit for a space followed by an opening single-quote, e.g. the word 'quote' means....
I reproduced this directly against the real pk::group_words function using real ids from that vocab ("▁the", "▁and", "▁'", "in"), split across two chunks the way StreamingSession would see them:
std::vector<std::string> pieces = {
"\xe2\x96\x81" "the", "\xe2\x96\x81" "and", "\xe2\x96\x81" "'", "in", "'"
};
std::vector<TokenInfo> full = {
{0 /* ▁the */, 0, 1.0f, 1},
{1 /* ▁and */, 1, 1.0f, 1},
{2 /* ▁' */, 2, 1.0f, 1},
{3 /* in */, 3, 1.0f, 1},
};- Ground truth (
group_words(full)):["the", "and'in"] - Incremental (this PR's algorithm, chunk 1 =
[▁the], chunk 2 =[▁and, ▁', in], then finalize):["the", "'in"]
The word "and" is silently dropped from the output.
What happens: when ▁' is the last ▁-prefixed token in a tail, the backward scan for the open-word start lands on ▁' instead of the real word-start token ▁and (which sits earlier in the tail). ▁and never gets pushed into final_words_, and wt_cursor_ advances past it, so it's excluded from every future tail too — its text is gone for good.
This affects --timestamps, the --json word array, and drain_words() in streaming. It does not affect the plain running transcript text — text()'s incremental append (detok_fragment) is a separate code path and I didn't find an issue there.
None of the existing tests catch this: tests/fixtures/speech.wav doesn't contain quoted text, and the new test_streaming_longrun_bounded regression test only asserts plain-text parity (sess.text() == full_text), not word-level parity against group_words on the full session.
Suggested fix direction: the cursor should pick the start of the open word using the same combined rule group_words uses (piece != decoded_text AND not punctuation), not piece_starts_with_meta alone. A word-level regression test comparing incremental drain_words()/finalize output against group_words on the full accumulated word_tokens_ (using a synthetic pieces vocab like the one above, no real model needed) would catch this and similar edge cases going forward.
|
@mudler This long-streaming optimization still needs contributor action before review can progress: GitHub reports the branch as conflicting with current |
Problem
StreamingSessionrebuilds the running transcript and the word grouping fromthe entire session history on every chunk:
process_emitted()runsdetokenize()over all non-special tokens so far.regroup_words()runsgroup_words()over all accumulated word tokens.Neither buffer is trimmed across utterances (the
<EOU>/<EOB>reset onlyclears the decoder LSTM state). So for a long-lived stream the per-chunk cost
grows with the whole session — O(N²) total CPU — and the repeated growing
allocations make RSS climb. A dictation session left open for hours pegs a core;
a multi-minute continuous stream already shows the per-chunk rate collapsing.
Fix
Make both steps incremental, with identical output:
detokenizeis prefix-stable (per-token piece concat +▁→space,and the single leading-space strip only touches byte 0). So append only the
newly emitted tokens' text to
text_instead of rebuilding it each chunk.word_tokens_from the last▁-word-start) and keep already-finalized words.group_wordssplits at▁-word-starts and its only cross-token effects (one-token forward lookahead,backward punctuation attach) never cross a word-start boundary, so
group_words([0,cursor)) ++ group_words([cursor,end)) == group_words([0,end)).No public contract changes:
text(),tokens(),take_new_text(),drain_words(),drain_events()return exactly what they did before — now O(1)amortized per chunk instead of O(N).
Testing
tests/test_streaming_longrun_bounded.cppdrives a long multi-utterancestream and asserts (a) the worst single-chunk reprocessing stays bounded (does
not scale with the session) via a lightweight high-water counter, and (b) the
incremental
text()is byte-for-byte identical to a fulldetokenize()of allnon-special tokens (public API). Measured on the streaming nemotron model:
max_chunk_reprocess3202 → 20 over a 1601-token session.(
test_streaming_decode,test_streaming_eou_reset,test_streaming_encoder,test_streaming_nemotron,test_capi_stream_json) — these cover the wordgrouping output.
Skips cleanly (exit 77) without
PARAKEET_TEST_GGUF_STREAM, matching the othermodel-dependent tests.