fix: Kimi K2.5 - tiktoken incomplete multi-byte sequence handling with regression tests

[biswapanda]

Overview:

The nvidia/Kimi-K2.5-NVFP4 model's worker thread panicked during streaming inference, crashing the server.

thread 'tokio-runtime-worker' panicked at lib/llm/src/backend.rs:185:83:
called Result::unwrap() on an Err value: Error decoding tiktoken tokens:
Unable to decode into a valid UTF-8 string: incomplete utf-8 byte sequence from index 6

Use _decode_native_and_split() with from_utf8_lossy() instead of CoreBPE::decode() so partial multi-byte sequences produce U+FFFD instead of panicking. Replace .unwrap() in backend.rs with error logging for defense in depth.

Includes 8 regression tests: 6 reproduce the original panic on incomplete byte sequences, 2 verify complete sequences are unaffected.

Also terminate stream gracefully on token decode errors instead of silently continuing the stream when process_token_ids fails, stop upstream generation, mark the stream as finished, and set FinishReason::Error so downstream consumers know decoding failed.

Details:

Why It Happened

The model's vocabulary is the root cause. Kimi K2.5 uses a tiktoken-based tokenizer (.model file with the Kimi BPE pattern). Unlike HuggingFace tokenizers that handle byte-fallback internally, tiktoken vocabularies can include raw byte tokens — individual tokens whose content is a single byte like 0xE4, 0xBD, or 0xA0.

These byte tokens exist to handle text that doesn't match any multi-byte BPE merge. For example, the CJK character "你" is encoded as 3 bytes in UTF-8: [0xE4, 0xBD, 0xA0]. If the BPE algorithm doesn't have a merged token for "你", it falls back to emitting three separate byte tokens, one per byte.

The problem surfaces during incremental (streaming) detokenization. When the model generates tokens one at a time, each token is decoded in the context of a sliding window of surrounding tokens (see DecodeStream::step() and Sequence::append_token_id()). These methods call tokenizer.decode() on subsequences of the token stream. When a subsequence contains only 1 or 2 of the 3 bytes needed for "你", the resulting byte array is not valid UTF-8.

The call chain that panicked:

1. backend.rs:185 — calls state.decoder.process_token_ids(&data.token_ids).unwrap()
2. This eventually calls TikTokenTokenizer::decode() on a token subsequence
3. tiktoken.rs:103-105 (old code) — calls self.bpe.decode(ids) which internally does String::from_utf8(bytes)
4. String::from_utf8() returns Err because [0xE4] alone is not valid UTF-8
5. The error propagates up, .unwrap() panics, and the tokio worker thread crashes

Why HuggingFace tokenizers don't have this problem: The tokenizers crate's decode path handles byte-fallback tokens specially, reconstructing multi-byte sequences before converting to UTF-8. Tiktoken's CoreBPE::decode()
does not — it concatenates raw token bytes and calls String::from_utf8() directly.

Why this doesn't happen in batch (non-streaming) decoding: In batch mode, all tokens for a complete response are decoded together. The full sequence [0xE4, 0xBD, 0xA0] is present, forming valid UTF-8. The bug only manifests during incremental detokenization where tokens are decoded in partial windows.

How It Was Fixed

Fix 1 — lib/llm/src/tokenizers/tiktoken.rs (core fix):

  // OLD: panics on incomplete UTF-8
  self.bpe.decode(ids)
      .map_err(|err| Error::msg(format!("Error decoding tiktoken tokens: {err}")))

  // NEW: gracefully handles incomplete UTF-8
  let bytes: Vec<u8> = self.bpe._decode_native_and_split(ids).flatten().collect();
  Ok(String::from_utf8_lossy(&bytes).into_owned())

Instead of CoreBPE::decode() (which requires valid UTF-8), the fix uses _decode_native_and_split() to get the raw bytes per token, concatenates them, and applies String::from_utf8_lossy(). Incomplete UTF-8 sequences become
the Unicode replacement character U+FFFD (�).

This works hand-in-hand with the existing DecodeStream::step() logic in tokenizers.rs:226:

  if new_text.len() > prefix_text.len() && !new_text.ends_with("�") {
      // emit text
  } else {
      Ok(None)  // buffer — incomplete character, wait for more tokens
  }

When decode() returns "�" for a partial byte sequence, DecodeStream::step() sees the replacement character, returns Ok(None), and buffers the token. On the next call, the window expands to include the next byte token. Once
all bytes of the character are in the window, from_utf8_lossy() produces the correct character (e.g., "你"), the ends_with("�") check passes, and the complete text is emitted. The same buffering logic exists in
Sequence::append_token_id() at line 337.

Fix 2 — lib/llm/src/backend.rs (defense in depth):

  // OLD: crashes the worker thread
  let result = state.decoder.process_token_ids(&data.token_ids).unwrap();

  // NEW: logs error and continues
  let result = match state.decoder.process_token_ids(&data.token_ids) {
      Ok(result) => result,
      Err(e) => {
          tracing::error!("Failed to process token_ids: {e}");
          return Some((output, state));
      }
  };

Replaced .unwrap() with error handling that logs the failure and passes through the original output. This prevents any future decode edge case from taking down the entire worker thread and crashing the server.

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• closes DIS-1584

Summary by CodeRabbit

• Bug Fixes

  • Improved error handling to prevent crashes during token processing failures.

• Tests

  • Added comprehensive tests for multi-byte character support, including emoji, CJK text, and various language encoding roundtrips.

[coderabbitai]

Walkthrough

This PR enhances error handling in token processing and improves UTF-8 decoding robustness in the TikToken tokenizer by replacing CoreBPE-based UTF-8 decoding with lossy UTF-8 conversion that gracefully handles incomplete multi-byte sequences. Comprehensive roundtrip tests validate multi-byte UTF-8 handling across encoding, decoding, and streaming paths.

Changes

Cohort / File(s)	Summary
Error Handling Enhancement lib/llm/src/backend.rs	Replaces unwrap on token processing with match-based error handling. On failure, logs error message and returns early instead of panicking, enabling graceful termination.
UTF-8 Decoding Logic lib/llm/src/tokenizers/tiktoken.rs	Refactors decode path from CoreBPE-based to lossy UTF-8 conversion via from_utf8_lossy. Adds extensive test utilities and inline tests for byte-token handling, multi-byte sequences (CJK, emoji), incomplete sequences, and stream-based incremental decoding.
Roundtrip Test Coverage lib/llm/tests/tokenizers.rs	Adds three new integration tests validating encode-decode roundtrips for multi-byte UTF-8 content across direct paths, DecodeStream incremental decoding, and Sequence token reconstruction with mixed ASCII and multi-byte inputs.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~55 minutes

Poem

🐰 With error handling now so graceful and kind,
UTF-8 sequences no longer make us unwind,
Lossy paths decode what's broken and bent,
While tests guard the emoji and scripts heaven-sent! ✨

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Title check	✅ Passed	The title accurately summarizes the main changes: fixing tiktoken incomplete multi-byte sequence handling and adding regression tests, directly matching the core technical fix and test additions in the changeset.
Description check	✅ Passed	The description is comprehensive and follows the template structure with Overview, Details, and Related Issues sections, providing thorough context on the root cause, fix implementation, and test coverage.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

lib/llm/src/backend.rs (1)

185-191: Consider signaling stream termination on decode error.

The error handling prevents panics, which is good. However, returning Some((output, state)) with the original unprocessed output continues the stream as if nothing happened. The caller receives incomplete data without knowing decoding failed.

Consider either:

1. Setting state.finished = true to terminate the stream gracefully after logging
2. Setting a finish_reason (e.g., FinishReason::Error) on the output to signal the issue to downstream consumers

Current behavior may silently drop decoded text for that iteration, which could cause subtle data loss.

♻️ Proposed option: terminate stream on error

                 let result = match state.decoder.process_token_ids(&data.token_ids) {
                     Ok(result) => result,
                     Err(e) => {
                         tracing::error!("Failed to process token_ids: {e}");
+                        state.finished = true;
                         return Some((output, state));
                     }
                 };

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/backend.rs` around lines 185 - 191, The decoder error handling
currently logs and returns Some((output, state)) which silently continues the
stream; update the Err(e) branch inside the match on
state.decoder.process_token_ids(&data.token_ids) to both log the error and
signal termination by setting state.finished = true and setting a finish reason
on output (e.g., output.finish_reason = Some(FinishReason::Error) or
equivalent), so downstream consumers know decoding failed and the stream ends
gracefully; ensure you reference the same output and state variables used in
this function and preserve any necessary cleanup before returning.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@lib/llm/src/backend.rs`:
- Around line 185-191: The decoder error handling currently logs and returns
Some((output, state)) which silently continues the stream; update the Err(e)
branch inside the match on state.decoder.process_token_ids(&data.token_ids) to
both log the error and signal termination by setting state.finished = true and
setting a finish reason on output (e.g., output.finish_reason =
Some(FinishReason::Error) or equivalent), so downstream consumers know decoding
failed and the stream ends gracefully; ensure you reference the same output and
state variables used in this function and preserve any necessary cleanup before
returning.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 7974f525-03cd-4a96-9e15-0dea10aa5dfa

📥 Commits

Reviewing files that changed from the base of the PR and between e02321f and 389e59d.

📒 Files selected for processing (3)

• lib/llm/src/backend.rs
• lib/llm/src/tokenizers/tiktoken.rs
• lib/llm/tests/tokenizers.rs

[ryanolson]

Nice catch.

Can we put parameters on the test using rstest? parametrize by tokenizer. Then for each type of tokenizer we support, run the same tests with an instance of that tokenizer?

Basically we want all tokenizers to pass the same set of scenarios.

This will be an easier policy to enforce when adding a tokenizer. pass their Arc to the parasitized scenario tests and if all scenarios pass we are good-to-got.

[biswapanda]

Nice catch.

Can we put parameters on the test using rstest? parametrize by tokenizer. Then for each type of tokenizer we support, run the same tests with an instance of that tokenizer?

Basically we want all tokenizers to pass the same set of scenarios.

This will be an easier policy to enforce when adding a tokenizer. pass their Arc to the parasitized scenario tests and if all scenarios pass we are good-to-got.

Good point @ryanolson. I added the rtest for tiktoken + hf. Please take another look.