feat(inference): UTF-8-safe streaming detokenization + generate_streaming

crates/inference/src/model/qwen35/detokenize.rs

@@ -1,27 +1,140 @@
 use crate::tokenizer::bpe::BpeTokenizer;
 use std::collections::HashMap;
 
-/// Decode token IDs back to text using the BPE tokenizer's reverse lookup.
-pub(crate) fn decode_tokens(tokenizer: &BpeTokenizer, ids: &[u32]) -> String {
+/// Build the GPT-2 byte-level reverse map (unicode placeholder char → original byte).
+pub(crate) fn build_byte_decoder() -> HashMap<char, u8> {
     let byte_encoder = bytes_to_unicode();
     let mut byte_decoder: HashMap<char, u8> = HashMap::new();
     for (byte_val, &ch) in byte_encoder.iter().enumerate() {
         byte_decoder.insert(ch, byte_val as u8);
     }
+    byte_decoder
+}
 
+/// Append the raw decoded bytes for a single token id onto `out`.
+///
+/// Byte-level BPE decoding is context-free per token, so concatenating the
+/// per-token byte strings equals decoding the whole id sequence at once.
+fn append_token_bytes(
+    tokenizer: &BpeTokenizer,
+    id: u32,
+    byte_decoder: &HashMap<char, u8>,
+    out: &mut Vec<u8>,
+) {
+    if let Some(token_str) = tokenizer.token_for_id(id) {
+        for ch in token_str.chars() {
+            if let Some(&b) = byte_decoder.get(&ch) {
+                out.push(b);
+            }
+            // Skip characters not in byte_decoder (e.g., special tokens)
+        }
+    }
+}
+
+/// Decode token IDs back to text using the BPE tokenizer's reverse lookup.
+pub(crate) fn decode_tokens(tokenizer: &BpeTokenizer, ids: &[u32]) -> String {
+    let byte_decoder = build_byte_decoder();
     let mut bytes = Vec::new();
     for &id in ids {
-        if let Some(token_str) = tokenizer.token_for_id(id) {
-            for ch in token_str.chars() {
-                if let Some(&b) = byte_decoder.get(&ch) {
-                    bytes.push(b);
+        append_token_bytes(tokenizer, id, &byte_decoder, &mut bytes);
+    }
+    String::from_utf8_lossy(&bytes).to_string()
+}
+
+/// Incremental, UTF-8-boundary-safe detokenizer for streaming generation.
+///
+/// Byte-level BPE frequently splits a single multibyte codepoint (CJK, emoji,
+/// accented Latin) across two or more tokens. Decoding the running token list
+/// with `from_utf8_lossy` after every token would substitute `U+FFFD` for the
+/// incomplete trailing bytes and stream that replacement char to the caller,
+/// who cannot retract it. This type instead buffers raw bytes and emits only
+/// the longest *complete* UTF-8 prefix after each token, holding incomplete
+/// trailing bytes until a later token completes the codepoint. The concatenation
+/// of every `push` delta plus the final `finish` equals `decode_tokens` over the
+/// same ids.
+pub(crate) struct IncrementalDetokenizer {
+    byte_decoder: HashMap<char, u8>,
+    bytes: Vec<u8>,
+    flushed: usize,
+}
+
+impl IncrementalDetokenizer {
+    pub(crate) fn new() -> Self {
+        Self {
+            byte_decoder: build_byte_decoder(),
+            bytes: Vec::new(),
+            flushed: 0,
+        }
+    }
+
+    /// Feed one token id; return the complete-UTF-8 text it reveals (which may be
+    /// empty when the token only contributed continuation bytes of a codepoint
+    /// still being assembled).
+    pub(crate) fn push(&mut self, tokenizer: &BpeTokenizer, id: u32) -> String {
+        append_token_bytes(tokenizer, id, &self.byte_decoder, &mut self.bytes);
+        self.flush_complete()
+    }
+
+    fn flush_complete(&mut self) -> String {
+        let mut out = String::new();
+        loop {
+            match std::str::from_utf8(&self.bytes[self.flushed..]) {
+                // Entire remaining buffer is valid UTF-8: emit it all.
+                Ok(s) => {
+                    out.push_str(s);
+                    self.flushed = self.bytes.len();
+                    return out;
+                }
+                Err(e) => {
+                    let valid = e.valid_up_to();
+                    if valid > 0 {
+                        // [flushed .. flushed + valid] is guaranteed valid by
+                        // valid_up_to, so from_utf8_lossy substitutes nothing.
+                        out.push_str(
+                            String::from_utf8_lossy(
+                                &self.bytes[self.flushed..self.flushed + valid],
+                            )
+                            .as_ref(),
+                        );
+                        self.flushed += valid;
+                    }
+                    match e.error_len() {
+                        // None: the error is an incomplete trailing codepoint,
+                        // not a malformed one. A later token may complete it, so
+                        // keep the tail buffered and emit no replacement char.
+                        None => return out,
+                        // Some(len): a genuinely invalid sequence of `len` bytes
+                        // that no future token can repair. Emit one U+FFFD (matching
+                        // from_utf8_lossy's maximal-subpart substitution), skip past
+                        // it, and continue — without this the stream would stall on
+                        // the bad byte until finish().
+                        Some(len) => {
+                            out.push('\u{FFFD}');
+                            self.flushed += len;
+                        }
+                    }
                 }
-                // Skip characters not in byte_decoder (e.g., special tokens)
             }
         }
     }
 
-    String::from_utf8_lossy(&bytes).to_string()
+    /// Flush any trailing incomplete bytes lossily. Call once after the final
+    /// token so the concatenated stream matches `decode_tokens` exactly even when
+    /// a generation is truncated mid-codepoint.
+    pub(crate) fn finish(&mut self) -> String {
+        if self.flushed < self.bytes.len() {
+            let tail = String::from_utf8_lossy(&self.bytes[self.flushed..]).into_owned();
+            self.flushed = self.bytes.len();
+            tail
+        } else {
+            String::new()
+        }
+    }
+
+    /// The full decoded text so far (equivalent to `decode_tokens` over every fed id).
+    pub(crate) fn text(&self) -> String {
+        String::from_utf8_lossy(&self.bytes).to_string()
+    }
 }
 
 /// GPT-2 byte-to-unicode mapping (same as in bpe.rs).
@@ -48,3 +161,81 @@ pub fn bytes_to_unicode() -> Vec<char> {
     }
     table
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn incremental_flush_reconstructs_split_codepoints() {
+        // 好 = E5 A5 BD, 😀 = F0 9F 98 80 — split across simulated token boundaries
+        // so several intermediate buffers end mid-codepoint.
+        let chunks: &[&[u8]] = &[&[0xE5, 0xA5], &[0xBD], &[0xF0, 0x9F], &[0x98, 0x80], b"ok"];
+        let mut d = IncrementalDetokenizer::new();
+        let mut out = String::new();
+        for c in chunks {
+            d.bytes.extend_from_slice(c);
+            out.push_str(&d.flush_complete());
+        }
+        out.push_str(&d.finish());
+        assert_eq!(out, "好😀ok");
+        assert_eq!(out, d.text());
+    }
+
+    #[test]
+    fn incremental_flush_truncated_mid_codepoint_matches_lossy() {
+        // Generation truncated after only the first 2 of 好's 3 bytes: finish()
+        // flushes the incomplete tail lossily, matching decode_tokens' behavior.
+        let mut d = IncrementalDetokenizer::new();
+        let mut out = String::new();
+        d.bytes.extend_from_slice(&[b'h', b'i', 0xE5, 0xA5]);
+        out.push_str(&d.flush_complete());
+        out.push_str(&d.finish());
+        assert_eq!(out, String::from_utf8_lossy(&[b'h', b'i', 0xE5, 0xA5]));
+        assert_eq!(out, d.text());
+    }
+
+    #[test]
+    fn incremental_flush_invalid_byte_does_not_stall_stream() {
+        // A lone continuation byte (0x80) is a malformed sequence no later token
+        // can complete. flush_complete must emit U+FFFD for it immediately and
+        // keep going, not buffer it until finish() (which would stall the stream).
+        let mut d = IncrementalDetokenizer::new();
+        d.bytes.extend_from_slice(&[0x80]);
+        let first = d.flush_complete();
+        assert_eq!(first, "\u{FFFD}", "invalid byte must flush immediately");
+        d.bytes.extend_from_slice(b"A");
+        let second = d.flush_complete();
+        assert_eq!(second, "A", "valid byte after invalid one must flush");
+        // Concatenated stream equals from_utf8_lossy over the whole byte buffer.
+        assert_eq!(
+            format!("{first}{second}"),
+            String::from_utf8_lossy(&[0x80, b'A'])
+        );
+    }
+
+    #[test]
+    fn incremental_flush_invalid_between_valid_matches_lossy() {
+        // Valid · invalid · valid in one buffer: one U+FFFD for the bad run, both
+        // valid runs verbatim — byte-for-byte equal to from_utf8_lossy.
+        let raw = [b'h', b'i', 0xFF, b'y', b'o'];
+        let mut d = IncrementalDetokenizer::new();
+        d.bytes.extend_from_slice(&raw);
+        let out = d.flush_complete();
+        assert_eq!(out, String::from_utf8_lossy(&raw));
+        assert_eq!(out, d.text());
+    }
+
+    #[test]
+    fn incremental_flush_ascii_is_exact_per_chunk() {
+        // Pure-ASCII path: every chunk is whole codepoints, so each flush emits it all.
+        let mut d = IncrementalDetokenizer::new();
+        let mut out = String::new();
+        for c in [b"He".as_slice(), b"llo", b"!"] {
+            d.bytes.extend_from_slice(c);
+            out.push_str(&d.flush_complete());
+        }
+        out.push_str(&d.finish());
+        assert_eq!(out, "Hello!");
+    }
+}

crates/inference/src/model/qwen35/generation.rs

@@ -1,5 +1,5 @@
 use super::cache::{ForwardScratch, KvCache};
-use super::detokenize::decode_tokens;
+use super::detokenize::{IncrementalDetokenizer, decode_tokens};
 use super::model::Qwen35Model;
 use super::sampling::sample_token;
 use crate::attention::gdn::GatedDeltaNetState;
@@ -85,6 +85,143 @@ impl Qwen35Model {
             generated_tokens: generated_ids.len(),
         })
     }
+
+    /// Streaming variant of [`generate`] — identical token sequence, but invokes
+    /// `on_token` with incremental text deltas after each generated token.
+    ///
+    /// # Parity safety
+    ///
+    /// The body below is a deliberate copy of `generate` rather than a refactor of
+    /// the shared path. This ensures that no change here can silently alter the
+    /// non-streaming `generate` path, which is pinned by the e2e-parity CI gate
+    /// (greedy token match vs HF transformers). `on_token` is the only addition.
+    pub fn generate_streaming(
+        &self,
+        prompt: &str,
+        gen_cfg: &GenerateConfig,
+        mut on_token: impl FnMut(&str),
+    ) -> Result<GenerateOutput, InferenceError> {
+        let cfg = &self.config;
+
+        let mut rng_state = initial_rng_state(gen_cfg.seed);
+
+        let input = self.tokenizer.tokenize(prompt);
+        let prompt_ids: Vec<u32> = input.input_ids[..input.real_length].to_vec();
+        let prompt_len = prompt_ids.len();
+
+        if prompt_len == 0 {
+            return Err(InferenceError::Inference("empty prompt".into()));
+        }
+
+        // max_new_tokens == 0 means "generate nothing": return before sampling so
+        // we never emit a token the caller did not ask for.
+        if gen_cfg.max_new_tokens == 0 {
+            return Ok(GenerateOutput {
+                text: String::new(),
+                token_ids: vec![],
+                prompt_tokens: prompt_len,
+                generated_tokens: 0,
+            });
+        }
+
+        let num_linear = cfg.num_linear_attention_layers();
+        let num_full = cfg.num_full_attention_layers();
+        let mut gdn_states: Vec<GatedDeltaNetState> = (0..num_linear)
+            .map(|_| GatedDeltaNetState::new(cfg))
+            .collect();
+        let mut kv_cache = KvCache::new(num_full);
+        let mut scratch = ForwardScratch::new();
+
+        let mut generated_ids: Vec<u32> = Vec::with_capacity(gen_cfg.max_new_tokens);
+        let mut all_ids = prompt_ids.clone();
+
+        prefill_tokens(
+            self,
+            &prompt_ids,
+            &mut gdn_states,
+            &mut kv_cache,
+            &mut scratch,
+        );
+        kv_cache.seq_len = prompt_len;
+
+        let next_id = sample_token(
+            &scratch.logits[..cfg.vocab_size],
+            gen_cfg,
+            &all_ids,
+            &mut rng_state,
+        );
+
+        if should_stop_token(cfg, gen_cfg, next_id) {
+            return Ok(GenerateOutput {
+                text: String::new(),
+                token_ids: vec![],
+                prompt_tokens: prompt_len,
+                generated_tokens: 0,
+            });
+        }
+
+        generated_ids.push(next_id);
+        all_ids.push(next_id);
+
+        // Incremental detokenization: emit only complete-UTF-8 text deltas. A
+        // byte-level BPE codepoint can span several tokens, so we buffer raw bytes
+        // and never stream a partial codepoint (see IncrementalDetokenizer).
+        let mut detok = IncrementalDetokenizer::new();
+        let delta = detok.push(&self.tokenizer, next_id);
+        if !delta.is_empty() {
+            on_token(&delta);
+        }
+
+        // Decode loop (mirrors decode_loop free function exactly).
+        for _ in 1..gen_cfg.max_new_tokens {
+            let pos = kv_cache.seq_len;
+            let Some(&last_token) = all_ids.last() else {
+                return Err(InferenceError::Inference("empty generation state".into()));
+            };
+
+            self.forward_step(
+                last_token,
+                pos,
+                &mut gdn_states,
+                &mut kv_cache,
+                &mut scratch,
+            );
+            kv_cache.seq_len += 1;
+
+            let next_id = sample_token(
+                &scratch.logits[..cfg.vocab_size],
+                gen_cfg,
+                &all_ids,
+                &mut rng_state,
+            );
+
+            if should_stop_token(cfg, gen_cfg, next_id) {
+                break;
+            }
+
+            generated_ids.push(next_id);
+            all_ids.push(next_id);
+
+            let delta = detok.push(&self.tokenizer, next_id);
+            if !delta.is_empty() {
+                on_token(&delta);
+            }
+        }
+
+        // Flush any trailing incomplete bytes (generation truncated mid-codepoint)
+        // so the streamed deltas concatenate to exactly the returned text.
+        let tail = detok.finish();
+        if !tail.is_empty() {
+            on_token(&tail);
+        }
+
+        Ok(GenerateOutput {
+            text: detok.text(),
+            token_ids: generated_ids.clone(),
+            prompt_tokens: prompt_len,
+            generated_tokens: generated_ids.len(),
+        })
+    }
 }
 
 fn initial_rng_state(seed: Option<u64>) -> u64 {