[Rust Frontend] Anchor tokens after pending UTF-8 bytes at their own first byte (#58357)
Signed-off-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Bugen Zhao <bugen@inferact.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Bugen Zhao committed
Oct 1, 2026 at 12:21 UTC
ec5ef92c87a66cb9e69c349d82d5e0980d494ab5
2 files changed
+172
-4
rust/src/tokenizer/src/incremental.rs
+72
-4
@@ -42,6 +42,10 @@ pub(crate) struct DecodeStream<'a, T: Tokenizer + ?Sized> {
42
prefix: String,
43
prefix_index: usize,
44
prefix_seeded: bool,
45
+ /// Joint decode from the last `push_token` that left tokens pending, i.e.
46
+ /// `decode(context ++ pending)`. Used to split the pending group from the
47
+ /// next token when that token provably does not merge with it.
48
+ pending_string: Option<String>,
49
decoded: AttributedTextBuffer,
50
}
51
@@ -60,6 +64,7 @@ impl<'a, T: Tokenizer + ?Sized> DecodeStream<'a, T> {
64
prefix: String::new(),
65
prefix_index: 0,
66
prefix_seeded: prompt_token_ids.is_empty(),
67
+ pending_string: None,
68
decoded: AttributedTextBuffer::default(),
69
}
70
}
@@ -72,6 +77,44 @@ const SAFE_SUFFIX_MIN: usize = 4;
77
const SAFE_SUFFIX_MAX: usize = 6;
78
79
impl<T: Tokenizer + ?Sized> DecodeStream<'_, T> {
80
+ /// Decide whether `token_id` can be anchored at its own first byte inside
81
+ /// `new_chunk`, separately from the older tokens still pending.
82
+ ///
83
+ /// Pending tokens (an incomplete UTF-8 sequence, or text still empty in
84
+ /// context) normally share one anchor with the token that completes them,
85
+ /// because the tokenizer only defines their joint decode. That sharing is
86
+ /// wrong when the new token does *not* merge with them, e.g. a stray
87
+ /// incomplete byte followed by a special token or plain ASCII: the joint
88
+ /// decode is then exactly `decode(pending) ++ decode([token_id])`, and the
89
+ /// new token's first byte is known. Returns that split offset, or `None`
90
+ /// when the tokens merge (byte fallback completing a character), when
91
+ /// nothing is pending, or when the isolated decode is not provably
92
+ /// consistent with the joint one. `pending` is the joint decode kept from the
93
+ /// push that left tokens pending. Decoded text is never changed by this;
94
+ /// only anchor placement is.
95
+ fn pending_group_split(
96
+ &self,
97
+ pending: Option<&str>,
98
+ token_id: u32,
99
+ new_chunk: &str,
100
+ ) -> Result<Option<usize>> {
101
+ let Some(pending) = pending else {
102
+ return Ok(None);
103
+ };
104
+ let prefix_len = self.prefix.len();
105
+ let pending_chunk = &pending[pending.floor_char_boundary(prefix_len)..];
106
+ // Empty pending text: today's shared anchor is already this token's first
107
+ // byte. Not a prefix: the new token completed the pending character.
108
+ if pending_chunk.is_empty() || !new_chunk.starts_with(pending_chunk) {
109
+ return Ok(None);
110
+ }
111
+ let own = self.tokenizer.decode(&[token_id], self.skip_special_tokens)?;
112
+ let independent = !own.is_empty()
113
+ && new_chunk.len() == pending_chunk.len() + own.len()
114
+ && new_chunk.ends_with(&own);
115
+ Ok(independent.then_some(pending_chunk.len()))
116
+ }
117
+
118
/// Decode prompt-only context for prefix seeding.
119
///
120
/// Prompt ids may come from the model vocabulary rather than the local
@@ -145,16 +188,40 @@ impl<T: Tokenizer + ?Sized> IncrementalDecoder for DecodeStream<'_, T> {
188
|| (!produced_text && self.tokenizer.id_to_token(token_id).is_none());
189
if zero_width {
190
self.decoded.record_zero_width_token(token_id);
148
- } else {
149
- self.decoded.record_pending_token(token_id);
191
}
151
-
192
if !produced_text {
193
+ if !zero_width {
194
+ self.decoded.record_pending_token(token_id);
195
+ self.pending_string = Some(string);
196
+ }
197
return Ok(0);
198
}
199
+
200
// Ensure we split at a utf-8 char boundary.
201
let new_chunk = &string[string.floor_char_boundary(prefix_len)..];
157
- self.decoded.append_visible_text(new_chunk);
202
+ // Text is released, so the pending group resolves in this push either
203
+ // way; never carry its joint decode into a later push.
204
+ let pending = self.pending_string.take();
205
+ let split = if zero_width {
206
+ None
207
+ } else {
208
+ self.pending_group_split(pending.as_deref(), token_id, new_chunk)?
209
+ };
210
+ match split {
211
+ Some(split) => {
212
+ // Older pending tokens own the prefix (including any U+FFFD from
213
+ // an incomplete character); this token owns exactly its own text.
214
+ self.decoded.append_visible_text(&new_chunk[..split]);
215
+ self.decoded.record_pending_token(token_id);
216
+ self.decoded.append_visible_text(&new_chunk[split..]);
217
+ }
218
+ None => {
219
+ if !zero_width {
220
+ self.decoded.record_pending_token(token_id);
221
+ }
222
+ self.decoded.append_visible_text(new_chunk);
223
+ }
224
+ }
225
self.ids.drain(..self.prefix_index);
226
self.prefix = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
227
self.prefix_index = self.ids.len();
@@ -184,6 +251,7 @@ impl<T: Tokenizer + ?Sized> IncrementalDecoder for DecodeStream<'_, T> {
251
}
252
}
253
self.decoded.resolve_pending_zero_width();
254
+ self.pending_string = None;
255
self.ids.clear();
256
self.prefix.clear();
257
self.prefix_index = 0;
rust/src/tokenizer/src/incremental/attribution/tests.rs
+100
@@ -520,6 +520,106 @@ fn attribution_incomplete_utf8_tokens_share_replacement_anchor_on_flush() {
520
});
521
}
522
523
+#[test]
524
+fn attribution_special_token_after_incomplete_utf8_anchors_at_its_own_spelling() {
525
+ // A stray incomplete byte followed by a special token: the joint decode is
526
+ // exactly `decode([0xe4]) ++ decode([special])`, so the special token owns
527
+ // its own first byte instead of sharing the group's anchor at 0.
528
+ let tokenizer = tokenizer_with_special_token();
529
+ run_attribution_case(AttributionCase {
530
+ tokenizer: &tokenizer,
531
+ prompt_token_ids: vec![],
532
+ token_ids: vec![0xe4, SPECIAL_TOKEN_ID],
533
+ skip_special_tokens: false,
534
+ min_bytes_to_buffer: 0,
535
+ truncate_output_to: None,
536
+ expected_chunks: vec![decoded(
537
+ "�<special>",
538
+ &[0xe4, SPECIAL_TOKEN_ID],
539
+ &[visible(0), visible(3)],
540
+ )],
541
+ expected_full: decoded(
542
+ "�<special>",
543
+ &[0xe4, SPECIAL_TOKEN_ID],
544
+ &[visible(0), visible(3)],
545
+ ),
546
+ });
547
+}
548
+
549
+#[test]
550
+fn attribution_ascii_token_after_incomplete_utf8_anchors_at_its_own_byte() {
551
+ // Same rule for an ordinary token: the split is decided by the non-merge
552
+ // equality, not by the token being special.
553
+ let tokenizer = TestTokenizer::new();
554
+ run_attribution_case(AttributionCase {
555
+ tokenizer: &tokenizer,
556
+ prompt_token_ids: vec![],
557
+ token_ids: vec![0xe4, b'a' as u32],
558
+ skip_special_tokens: false,
559
+ min_bytes_to_buffer: 0,
560
+ truncate_output_to: None,
561
+ expected_chunks: vec![decoded(
562
+ "�a",
563
+ &[0xe4, b'a' as u32],
564
+ &[visible(0), visible(3)],
565
+ )],
566
+ expected_full: decoded("�a", &[0xe4, b'a' as u32], &[visible(0), visible(3)]),
567
+ });
568
+}
569
+
570
+#[test]
571
+fn attribution_byte_fallback_completion_keeps_shared_anchor_before_special_token() {
572
+ // The completing byte merges with the pending group (no split); the special
573
+ // token that follows a complete character is an ordinary independent token.
574
+ let tokenizer = tokenizer_with_special_token();
575
+ run_attribution_case(AttributionCase {
576
+ tokenizer: &tokenizer,
577
+ prompt_token_ids: vec![],
578
+ token_ids: vec![0xe4, 0xbd, 0xa0, SPECIAL_TOKEN_ID],
579
+ skip_special_tokens: false,
580
+ min_bytes_to_buffer: 0,
581
+ truncate_output_to: None,
582
+ expected_chunks: vec![
583
+ decoded(
584
+ "你",
585
+ &[0xe4, 0xbd, 0xa0],
586
+ &[visible(0), visible(0), visible(0)],
587
+ ),
588
+ decoded("<special>", &[SPECIAL_TOKEN_ID], &[visible(0)]),
589
+ ],
590
+ expected_full: decoded(
591
+ "你<special>",
592
+ &[0xe4, 0xbd, 0xa0, SPECIAL_TOKEN_ID],
593
+ &[visible(0), visible(0), visible(0), visible(3)],
594
+ ),
595
+ });
596
+}
597
+
598
+#[test]
599
+fn attribution_filtered_special_between_incomplete_utf8_and_independent_token() {
600
+ // A skipped special token sits between the pending byte and the independent
601
+ // token; it stays zero-width in generation order and the split still applies.
602
+ let tokenizer = tokenizer_with_special_token();
603
+ run_attribution_case(AttributionCase {
604
+ tokenizer: &tokenizer,
605
+ prompt_token_ids: vec![],
606
+ token_ids: vec![0xe4, SPECIAL_TOKEN_ID, b'a' as u32],
607
+ skip_special_tokens: true,
608
+ min_bytes_to_buffer: 0,
609
+ truncate_output_to: None,
610
+ expected_chunks: vec![decoded(
611
+ "�a",
612
+ &[0xe4, SPECIAL_TOKEN_ID, b'a' as u32],
613
+ &[visible(0), zero_width(0), visible(3)],
614
+ )],
615
+ expected_full: decoded(
616
+ "�a",
617
+ &[0xe4, SPECIAL_TOKEN_ID, b'a' as u32],
618
+ &[visible(0), zero_width(0), visible(3)],
619
+ ),
620
+ });
621
+}
622
+
623
#[test]
624
fn attribution_truncation_converts_fully_removed_token_to_zero_width() {
625
run_attribution_case(AttributionCase {