Allow tokenizers to provide streaming detokenizers

[aleroot]

Proposed changes

Generation currently constructs NaiveStreamingDetokenizer directly, preventing tokenizer adapters from supplying an incremental implementation. This can require repeated decoding of a growing text segment, especially for long responses without newlines.

Add Tokenizer.makeStreamingDetokenizer() with a default implementation returning the existing naive decoder. Tokenizer adapters can opt into their own decoding implementation without adding a tokenizer dependency to MLX or changing the existing swift-transformers adapter.

This is an extension point: existing adapters retain the default decoder. Performance improvements require an adapter that supplies an optimized implementation.

Example: swift-tokenizers

An application using aleroot/swift-tokenizers can implement the factory inside its existing MLXLMCommon.Tokenizer bridge, where upstream is its wrapped Tokenizers.Tokenizer:

func makeStreamingDetokenizer() -> any MLXLMCommon.StreamingDetokenizer {
    SwiftTokenizersDetokenizer(tokenizer: upstream)
}

The following adapter connects that library's independent decoding session to MLX's append/next/finish lifecycle:

import MLXLMCommon
import Tokenizers

struct SwiftTokenizersDetokenizer: MLXLMCommon.StreamingDetokenizer {
    private var decoder: Tokenizers.TokenStreamDecoder
    private var pending = ""

    init(tokenizer: any Tokenizers.Tokenizer) {
        decoder = tokenizer.makeStreamDecoder(skipSpecialTokens: false)
    }

    mutating func append(token: Int) {
        do {
            pending += try decoder.append(token: token)
        } catch {
            preconditionFailure("Create a new detokenizer after finish(): \(error)")
        }
    }

    mutating func next() -> String? {
        guard !pending.isEmpty else { return nil }
        defer { pending = "" }
        return pending
    }

    mutating func finish() -> String? {
        pending += decoder.finish()
        return next()
    }
}

Special tokens remain enabled because MLX's tool and reasoning parsers consume them. Each factory call creates independent state. finish() returns any unread output and the decoder's final text. Appending after finish() is a programmer error under this lifecycle.

Note

In my own benchmark, coupled with swift-tokenizers, incremental decoding substantially reduced tokenizer CPU time compared with repeatedly decoding the full generated prefix. For 4,096 tokens of newline-free text, decoding took approximately 0.19 ms instead of 55.6 ms for Qwen, and 1.64 ms instead of 1.16 seconds for Llama 2.

These are decoder-only measurements, not end-to-end generation speedups. MLX’s default detokenizer resets its decoding segment at newlines, so its actual improvement will depend on the output and tokenizer. Long uninterrupted responses, JSON, and code are likely to benefit most.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

AI usage

• I have read this PR description in full and approve it as my own, and it
  accurately describes the code changes.
• AI usage disclosure:

[davidkoski]

This is a great idea!

[aleroot]

This is a great idea!

Thanks, see the performances notes 😀.

[davidkoski]

Thanks, see the performances notes 😀.

Wow, that is huge and a big win especially for apps that want to show the streaming output.

[ronaldmannak]

Edit: Apologies for the original version of this comment. It was posted by an agent without my consent or review. Let me rewrite it in my own words.