Implementation:Mlc ai Mlc llm Text Streamer

Overview

The Text Streamer implementation at cpp/tokenizers/streamer.cc provides two key streaming components for MLC LLM token-by-token text generation: the TextStreamer (which converts token IDs into valid UTF-8 text incrementally) and the StopStrHandler (which detects stop strings in the token stream and truncates output accordingly). Both classes are registered as TVM FFI objects for cross-language accessibility.

Purpose

During LLM inference, tokens are generated one at a time. A single token may not decode into a complete UTF-8 character -- it could represent only part of a multi-byte sequence. The TextStreamer solves this by buffering tokens until they form valid UTF-8 output. The StopStrHandler complements this by monitoring the output stream for configurable stop strings and halting generation when one is found.

TextStreamer Implementation

Construction

TextStreamerObj::TextStreamerObj(Tokenizer tokenizer) : tokenizer_(std::move(tokenizer)) {}

TextStreamer::TextStreamer(Tokenizer tokenizer) {
  data_ = tvm::ffi::make_object<TextStreamerObj>(std::move(tokenizer));
}

The streamer is constructed with a Tokenizer instance used for decoding token sequences into strings.

Put Method

std::string TextStreamerObj::Put(const std::vector<int32_t>& delta_tokens);

The core incremental decoding method. For each delta token, it:

1. Appends the token to pending_tokens_.
2. Concatenates prefix_tokens_ and pending_tokens_ to form all_tokens.
3. Decodes both prefix_tokens_ and all_tokens using the tokenizer.
4. Computes a validated delta string using one of two strategies:

Case 1: The decoded prefix string is a prefix of the full decoded string. The delta text is the suffix beyond the prefix. Tokens are popped from the back of pending_tokens_ if the suffix ends with the UTF-8 replacement character (U+FFFD, encoded as \xef\xbf\xbd), which indicates an incomplete multi-byte sequence. At most 3 tokens are popped since a valid UTF-8 character is at most 4 bytes.

Case 2: The decoded prefix string is not a prefix of the full decoded string (a rare case caused by tokenizer behavior). Tokens are popped from pending_tokens_ until the prefix relationship is restored. If fewer than 3 pending tokens exist, the loop continues to the next delta token without output.

After processing, the non-popped pending tokens become the new prefix_tokens_, and the popped tokens (reversed) become the new pending_tokens_.

Finish Method

std::string TextStreamerObj::Finish();

Flushes all remaining tokens. Concatenates prefix_tokens_ and pending_tokens_, decodes them, and returns the portion beyond the prefix string. Sets the finished_ flag to prevent further calls to Put.

StopStrHandler Implementation

KMP Partial Match Table

inline std::vector<int> CreatePartialMatchTable(const String& str);

Creates the failure function (partial match table) for the Knuth-Morris-Pratt (KMP) string matching algorithm. This table enables efficient character-by-character matching against stop strings without backtracking through the input.

Construction

StopStrHandlerObj::StopStrHandlerObj(Array<String> stop_strs,
                                     const std::vector<std::string>& token_table);

Initializes the handler with an array of stop strings and a token table for looking up token string representations. For each stop string, a KMP partial match table is precomputed. Empty stop strings are rejected with a CHECK assertion.

Put Method

void StopStrHandlerObj::Put(int32_t token_id, std::vector<int64_t>* return_token_ids);

Processes one token at a time. For each character in the token's string representation, the method:

1. Runs one step of the KMP algorithm against every stop string simultaneously.
2. Tracks the earliest possible start position of any stop string match.
3. Computes a safe cutoff length -- the number of characters that can be emitted without risk of being part of a stop string.
4. Converts the cutoff length back into whole token IDs that can be safely returned.

If a stop string is fully matched (cur_match_length == stop_str.length()), the stop_triggered_ flag is set, the cutoff is adjusted to exclude the stop string, and remaining tokens are discarded.

Token Buffering

The handler maintains:

• pending_token_ids_ -- Token IDs not yet returned, potentially part of a stop string.
• pending_token_lengths_ -- The string length of each pending token.
• pending_string_len_ -- The total character length of all pending, unmatched content.

Tokens are only returned when it is certain they do not overlap with any stop string prefix.

TVM FFI Registration

Both classes are registered as TVM FFI functions:

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("mlc.tokenizers.TextStreamer", ...)
      .def("mlc.tokenizers.TextStreamerPut", ...)
      .def_method("mlc.tokenizers.TextStreamerFinish", &TextStreamerObj::Finish);
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("mlc.tokenizers.StopStrHandler", ...)
      .def("mlc.tokenizers.StopStrHandlerPut", ...)
      .def("mlc.tokenizers.StopStringHandlerFinish", ...)
      .def_method("mlc.tokenizers.StopStrHandlerStopTriggered", &StopStrHandlerObj::StopTriggered);
}

This allows both the text streamer and stop string handler to be instantiated and invoked from Python and other TVM-supported languages.

Dependencies

• streamer.h -- Header declarations for TextStreamerObj and StopStrHandlerObj.
• tokenizers.h -- Tokenizer interface.
• tvm/ffi/function.h and tvm/ffi/reflection/registry.h -- TVM FFI registration.
• tvm/runtime/int_tuple.h -- IntTuple for returning token ID arrays.
• <algorithm> -- For std::reverse and std::min/std::max.

File Location

• Source file: cpp/tokenizers/streamer.cc
• Namespace: mlc::llm