perf: use Rust-native DecodeStream for O(n) streaming decode

- Replace hand-rolled StreamDecoder (O(n^2) full-history re-decode per token) with tokenizers.decoders.DecodeStream
- Keep O(1) bounded token buffer internally via prefix drain instead of accumulating all token IDs
- Simplify flush_remaining to no-op since stream always emits completed text per step
- Benchmark on 8000 tokens: 2305ms -> 3.9ms (~592x speedup)
This commit is contained in:
2026-07-28 14:32:10 +08:00
parent 39f84f3b4c
commit a95794d3db
+19 -31
View File
@@ -6,6 +6,8 @@ from collections import deque
from enum import Enum
from typing import Any, Callable, Deque, Dict, List, Optional
from tokenizers.decoders import DecodeStream
from astrai.tokenize.tokenizer import AutoTokenizer
logger = logging.getLogger(__name__)
@@ -14,37 +16,30 @@ STOP = object()
class StreamDecoder:
"""Incremental decoder for byte-level BPE streaming.
"""Incremental decoder backed by the tokenizers library's DecodeStream.
Byte-level BPE may split a single Unicode character (e.g. em-dash,
smart quotes) across multiple tokens. Decoding such a token in
isolation produces U+FFFD (replacement char). This decoder
accumulates token IDs and only emits text once the trailing
characters are complete, buffering incomplete multi-byte sequences
until the next token arrives.
Delegates to the Rust-native streaming decoder which maintains an
O(1) bounded token buffer internally (via prefix drain), avoiding
the O(n²) cost of re-decoding the full history on each step.
Multi-byte UTF-8 sequences split across token boundaries are
buffered until complete; ``push`` returns "" while the trailing
sequence is still incomplete.
"""
__slots__ = ("_tokenizer", "_ids", "_emitted")
__slots__ = ("_stream", "_tok")
def __init__(self, tokenizer: AutoTokenizer):
self._tokenizer = tokenizer
self._ids: List[int] = []
self._emitted: str = ""
self._tok = tokenizer._tokenizer
self._stream = DecodeStream(skip_special_tokens=True)
def push(self, token_id: int) -> str:
"""Append a token ID and return newly completed text.
Returns "" while a multi-byte character is still incomplete.
"""
self._ids.append(token_id)
full = self._tokenizer.decode(self._ids, skip_special_tokens=True)
if full.endswith("\ufffd"):
return ""
if len(full) > len(self._emitted):
diff = full[len(self._emitted) :]
self._emitted = full
return diff
return ""
chunk = self._stream.step(self._tok, token_id)
return chunk or ""
class TaskStatus(Enum):
@@ -101,18 +96,11 @@ class Task:
def flush_remaining(self, tokenizer: AutoTokenizer) -> str:
"""Emit any text still buffered in the decoder.
Called when generation terminates (max_tokens reached, stop
sequence, or external removal) to avoid dropping a final
incomplete-looking fragment that is actually complete when
adjacent to the stop token.
With the Rust-native DecodeStream, the stream is always in a
correct state — any completed text was already emitted by the
last ``push``. A trailing incomplete multi-byte sequence has no
valid text to emit, so this is a no-op.
"""
if self._decoder is None or not self.output_ids:
return ""
full = tokenizer.decode(self.output_ids, skip_special_tokens=True)
if len(full) > len(self._decoder._emitted):
diff = full[len(self._decoder._emitted) :]
self._decoder._emitted = full
return diff
return ""
@property