perf(server): incremental detokenization and per-token streaming overhead reduction #633

Part of #623. Independent of the CUDA work; measured impact is largest on fast small models (200-650 tok/s on GB10).

Context

Three host-side per-token costs on the serving path:

1. O(N) detokenization per token. StreamingDecodeState::on_token -> emit_regular_token calls tokenizer.decode(&self.all_ids, false) on the FULL token history for every generated token (src/server/model_worker.rs:2014, byte-fallback branch :1979). At 500 tok/s and 2k-token outputs this is quadratic work sitting on the single generation thread.
2. Per-token streaming overhead in the chat route. For every token the callback locks three mutexes (accumulator, logprobs buffer, stream filter), builds a ChatCompletionChunk, serializes JSON, and does two mpsc hops (src/server/routes/chat.rs:640-810, sends at :762,:781,:806; src/server/streaming.rs:133-147).
3. Prompt tokenization on the generation thread. enqueue_request tokenizes inside the scheduler thread (src/server/batch/scheduler.rs:2275), so a long prompt arriving stalls decode ticks for every active sequence.

Scope

All three items; they are small individually and share test surface (byte-exact streaming).

Implementation plan

1. Incremental detokenizer

• Replace full-history decode with an incremental scheme in StreamingDecodeState (src/server/model_worker.rs:1936-2060): keep the previously decoded string (or byte offset) plus a small tail window of token ids; on each token, decode only tail_window + new_id and emit the suffix past the previously emitted boundary. The tail window (a few tokens) handles tokenizers whose piece merging changes the tail rendering; the existing UTF-8 partial-character handling (byte-fallback branch at :1979) must be preserved: never emit an incomplete UTF-8 sequence, buffer until complete.
• Mirror the proven approach in mlx-lm/llama.cpp streaming detokenizers rather than inventing merge rules.
• The same state is used by both the batched and single decode emit paths (scheduler.rs:4513-4519, :4830-4836); no call-site changes should be needed beyond the struct internals.

2. Streaming-path diet

• Collapse the three per-token mutex acquisitions in the chat callback into one small state struct behind a single lock, or restructure so the callback owns its state (it runs on one thread; check whether the mutexes exist only for Sync bounds and can become RefCell-style or plain ownership).
• Reuse a serialization buffer for chunk JSON (serde_json to_writer into a reused Vec<u8>), and skip constructing unused optional fields on the common path.

3. Tokenize off the generation thread

• Move tokenizer.encode from enqueue_request (scheduler.rs:2275) to the HTTP handler side (src/server/routes/chat.rs, where the chat template is already applied on the tokio worker at :365) and pass token ids through the request struct. The tokenizer is already shared/thread-safe enough to construct encodings outside; verify add_special handling stays identical (scheduler.rs:2275 semantics).

Acceptance criteria

• Detokenization cost is O(tail) per token: demonstrated by a benchmark generating 4096 tokens on qwen2.5-0.5b-bf16 where server tok/s no longer degrades with output length (record tok/s for tokens 0-512 vs 3584-4096; current code degrades measurably).
• Server streaming throughput on qwen2.5-0.5b-bf16 within 10% of mlxcel-bench-decode for the same generation on GB10.
• Byte-exact stream output vs the old implementation across: Korean/CJK text, emoji, byte-fallback tokens, stop-string truncation mid-multibyte, logprobs enabled. Add a regression test comparing accumulated stream text to a one-shot tokenizer.decode of the final ids.
• Long-prompt request (8k tokens) no longer stalls concurrent decode: verified with the perf(bench): long-prompt prefill benchmark coverage and serving TTFT/decode-rate telemetry #624 concurrency script (TTFT of a concurrent short request stays bounded while an 8k-prompt request is admitted).
• cargo test --features cuda (or default features on a dev box; this issue is backend-agnostic) passes with the new streaming tests.

Validation

cargo build --release
./target/release/mlxcel-server --model ./models/qwen2.5-0.5b-bf16 &
python3 scripts/bench_serving_concurrency.py --concurrency 1 --max-tokens 4096

References

• Detok: src/server/model_worker.rs:1936-2132.
• Route callback: src/server/routes/chat.rs:640-810; SSE plumbing src/server/streaming.rs:117-165.
• Enqueue tokenize: src/server/batch/scheduler.rs:2275.