perf(server): incremental detokenization and streaming overhead cuts

[inureyes]

Summary

Reduces three host-side per-token costs on the serving path (issue #633): O(N)-per-token detokenization, three per-token mutex acquisitions in the chat streaming callback, and prompt tokenization on the scheduler thread. The dominant win is the incremental detokenizer, which turns the quadratic-with-output-length detok cost into O(window) per token while staying byte-exact.

What changed

1. Incremental detokenizer (src/server/model_worker.rs)

StreamingDecodeState previously called tokenizer.decode(&all_ids, false) over the full token history on every generated token (regular emit path and byte-fallback re-sync), which is O(N) per token and O(N^2) over a full generation, on the single generation thread. It now re-decodes only a bounded suffix window: two decodes anchored at the same prefix_offset (the already-emitted window prefix, and that prefix plus the newest tokens), emitting the delta past the shared prefix. Since both decodes share the left context, the shared prefix cancels out and the delta is byte-identical to a whole-history decode. read_offset only advances to a UTF-8-complete boundary, so an incomplete multibyte sequence (byte-fallback <0xXX> runs or split byte-level pieces) is held until it completes and is never streamed as U+FFFD. Mirrors the mlx-lm NaiveStreamingDetokenizer and the llama.cpp/vLLM incremental detokenizers.

A MAX_HELD_TOKENS cap bounds the window on degenerate never-completing output, and a strip_prefix guard keeps the delta slice panic-safe when a decoder retroactively rewrites already-emitted text (the tokenizers ByteFallback decoder renders an entire byte run as U+FFFD once it holds an incomplete byte). The public StreamingDecodeState surface is unchanged, so the batched, single, speculative, diffusion, XLA, and CLI decode call sites are untouched. The now-redundant token_piece byte-fallback buffer, parse_byte_fallback_token, and safe_emit_boundary helpers are removed; their behavior (never emit incomplete UTF-8) is preserved by the window guard and covered end-to-end by the streaming tests.

2. Streaming-path diet (src/server/routes/chat.rs)

The chat streaming callback locked three separate mutexes per token (tool-call accumulator, logprobs buffer, stream filter). They are collapsed into one StreamCallbackState behind a single Mutex, so the hot callback locks once per token; emitted chunks are collected and sent after the lock is released. The lock is required only for the Send bound on the spawn_blocking closure (an Arc<RefCell> would not be Send) and is never contended. The SSE wire format is byte-identical.

3. Tokenize off the generation thread (src/server/model_provider.rs, src/server/batch/scheduler.rs)

Prompt tokenization moved out of BatchScheduler::enqueue_request (scheduler thread, which stalled concurrent decode ticks for a long prompt) to the request-dispatch path. ModelProvider holds an optional pre-tokenizer and encodes the prompt in send_generate_request_with_cancellation (HTTP-side thread), passing the ids through ModelRequest::Generate. The scheduler uses them when present and falls back to tokenizing itself otherwise. tokenize_prompt_for_generation is the single source of truth for the add_special convention (suppressed on a literal <bos>/<s> prefix), so both paths are byte-identical. Legacy, XLA, diffusion, and test paths keep the None fallback.

Test plan

Hardware: Apple M1 Ultra (macOS), features metal,accelerate. Model: models/qwen2.5-0.5b-bf16 (bf16, ByteLevel tokenizer, clean_up_tokenization_spaces=false).

Measured here:

• Detok is O(window), not O(N). Single greedy 4096-token stream, tok/s bucketed by output position (custom SSE arrival timer):

  	first 512 tok/s	last 512 tok/s	end/start ratio	overall decode tok/s
  BEFORE (main b73be08)	196.9	160.4	0.815	163.0
  AFTER (this branch)	191.8	179.6	0.936	181.2

  The length-dependent degradation drops from 18.5% to 6.4% (the residual is GPU warmup/thermal, not detok). Tail throughput +12% (160.4 -> 179.6), overall decode +11% (163.0 -> 181.2). scripts/bench_serving_concurrency.py --concurrency 1 --max-tokens 4096 corroborates the overall figure (162.8 -> 180.9 aggregate tok/s).

• Byte-exact streaming. New unit tests beside the code assert the accumulated streamed text equals a one-shot tokenizer.decode of the final ids across Korean, mixed CJK/emoji/ASCII, regular-piece + multibyte, a non-empty prompt seam, mid-multibyte truncation, and a 300-case fuzz. All existing byte-fallback streaming regressions and the server::routes::chat / server::streaming tests pass.

• cargo test --release --features metal,accelerate --lib -- server::model_provider::model_worker::tests -> 24 passed, 0 failed.

• Combined server::model_provider::model_worker::tests server::routes::chat::tests server::streaming::tests -> all pass.

• cargo check --lib --tests --features metal,accelerate, cargo clippy --lib --tests --features metal,accelerate -- -D warnings, cargo fmt --all -- --check -> clean.

Pending (not cheaply reproducible on this hardware; recorded per acceptance criteria):

• perf(bench): long-prompt prefill benchmark coverage and serving TTFT/decode-rate telemetry #624 long-prompt admission (8k-prompt request does not stall a concurrent short request's TTFT). Item 3 moves tokenization off the scheduler thread by construction, but the mixed-workload scenario needs a client that fires one 8k-prompt request alongside short requests; scripts/bench_serving_concurrency.py fires N identical requests and does not model the mixed case. Suggested command for a follow-up run: python3 scripts/bench_serving_concurrency.py --concurrency 8 --prompt-tokens 8192 --max-tokens 64 while separately timing a short request's TTFT.

Review follow-up: stream the held detokenizer tail

The incremental detokenizer holds back the whole window while its tail is an incomplete UTF-8 sequence, and a single byte-level token (Qwen-style) can carry complete text plus a trailing incomplete byte. If generation stopped on such a token, StreamingDecodeState::flush appended that complete text to generated_text (seen by the non-streaming result.text) but returned nothing, so streaming clients dropped it. flush now returns the flushed delta (Option<String>, #[must_use]) and every streaming finish site forwards it as one final GenerateEvent::Token before Done: both BatchScheduler finish paths, the speculative-burst finalize, the diffusion worker, and the XLA worker's Finished branch (after the stop-matcher tail). The XLA stop-string finalize still drops its post-stop tail and the CLI print paths already recover the tail via a separate decode(generated_ids) + strip_prefix, so both are unchanged.

New/updated tests: a ByteLevel stub whose single token decodes to abc + the first byte of a 3-byte char proves flush returns the held abc tail on stop and stays byte-exact when the char completes mid-stream; the byte-fallback stream helper now folds the flush return into the streamed chunks and the incomplete-at-end regression asserts the U+FFFD tail is streamed exactly once and equals result.text. cargo test --release --features metal,accelerate --lib -- server::model_provider::model_worker::tests server::routes::chat::tests server::streaming::tests -> 66 passed, 0 failed. cargo check --lib --tests --bins, cargo clippy --lib --tests -- -D warnings, cargo fmt --all -- --check clean.

Closes #633

[inureyes]

Orchestrator E2E verification on the branch build (Apple M1 Ultra, models/qwen2.5-0.5b-bf16, real mlxcel-server): accumulated SSE stream deltas vs the non-streamed response content for five cases — Korean, emoji, mixed CJK/Korean/English, byte-fallback (café), and a 2554-byte long generation.

Result: korean / cjk-mixed / byte-fallback / long are byte-identical between stream and non-stream. The emoji case differed, and both differences trace to pre-existing behavior, not this PR:

• Leading space: the non-stream chat path passes the text through clean_structural_tokens → strip_thinking/clean_content_markers, which .trim()s (src/server/tool_calls/parser.rs, untouched by this PR), while the stream path emits raw deltas. Reproduced 3/3 deterministically on both paths; same asymmetry class as the router-vs-single-node leading space noted in PR fix(router): emit usage on the disaggregated /v1/chat/completions responses #707's E2E. Pre-existing main behavior.
• Tail divergence: repeating the identical non-stream request three times produced two different generations (md5 1864c0232d49 vs 01a3b170b22d) — greedy run-to-run jitter on this bf16 checkpoint (Metal reduction-order nondeterminism, the known f16-jitter class), so comparing two separate requests cannot be token-stable for this model.

Net: no streaming byte-exactness regression attributable to this PR in E2E; the strong invariant (accumulated stream text == one-shot decode of final ids) is pinned by the new unit/fuzz tests.

[inureyes]

Implementation Review Summary

Intent

Cut three host-side per-token serving costs (issue #633): O(N)-per-token detokenization, three per-token chat-stream mutexes, and prompt tokenization on the scheduler thread. The core is a windowed incremental detokenizer mirroring vLLM/mlx-lm, kept byte-exact.

Findings

HIGH: streaming loses trailing complete text when the last content token ends mid-multibyte.
on_token holds the entire window whenever new_text ends in U+FFFD (src/server/model_worker.rs:2161). When a single token carries new complete text plus an incomplete trailing byte (common for ByteLevel/byte-fallback tokenizers) and generation stops right after (max_tokens, EOS mid-character, or a stop-string), that complete text is never emitted as a GenerateEvent::Token. It is only appended to generated_text by flush() (:2183, returns nothing) and surfaces solely in GenerateEvent::Done(result.text). The streaming drain consumes only Token/TokenWithLogprobs and ignores Done.text (src/server/model_provider.rs:1214-1225), so streaming clients miss the final characters while non-streaming clients keep them. This is a regression from the old safe_emit_boundary, which emitted the complete prefix eagerly and held only the incomplete tail. The new tests do not catch it because the byte-fallback stub emits one byte per token, so complete text can never be stranded together with an incomplete tail in the same held window; only a multi-byte token that spans a UTF-8 boundary triggers it.

Recommended fix: have flush() return the flushed delta and emit it as a final GenerateEvent::Token before Done at every streaming finish site (scheduler.rs:4047 and :4970, speculative_burst.rs:1776, xla_worker.rs:357, diffusion_worker.rs:371); mirror it into the CLI print paths. Validate with the cohort/handoff parity suites and a real checkpoint whose max_tokens lands mid-multibyte. Not auto-fixed here because a correct fix spans hot serving paths outside the narrow test scope this review could execute.

LOW (fixed): stale // Used by: discovery comment on MlxcelTokenizer::token_piece still named StreamingDecodeState, which no longer calls it. Corrected in d4fb044.

Verified clean

• Removed helpers (parse_byte_fallback_token, safe_emit_boundary, byte-fallback buffer) have no remaining non-test callers; all StreamingDecodeState call sites (batched, single, speculative, diffusion, XLA, CLI) compile against the unchanged public surface (cargo check --lib --tests and --bins clean).
• Window seam is byte-exact: both decodes share prefix_offset, so leading-space/SentencePiece context cancels; legitimate U+FFFD stays byte-exact (delayed up to MAX_HELD_TOKENS).
• Chat mutex collapse is per-request, holds one std mutex only within the synchronous spawn_blocking callback (no lock across .await), releases before sending, and preserves reasoning -> content -> placeholder chunk order and logprob wire format.
• Logprob token strings still come from build_single_token_chat_logprobs, independent of the removed token_piece path.
• Pre-tokenization funnels through a single ModelRequest::Generate construction with a shared tokenize_prompt_for_generation (add_special identical to the old scheduler call); None fallback intact; no double-tokenization; Moondream text-only re-tokenization unaffected.
• Fuzz test is deterministically seeded (xorshift64), not flaky.

Checks

• cargo fmt --all -- --check, cargo check --lib --tests, cargo clippy --lib --tests -- -D warnings: clean.
• server::model_provider::model_worker::tests + server::routes::chat::tests + server::streaming::tests: 64 passed, 0 failed.