@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b911fe85c4

[Rust Frontend] Attribute decoded text to tokens (#52910)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Sep 2, 2026 at 08:25 UTC b911fe85c42ae984f27d183a0b1de309adab9bf2
17 files changed +1182 -320
rust/Cargo.lock
+1
@@ -6043,6 +6043,7 @@ dependencies = [
6043 "rustc-hash 1.1.0",
6044 "serde",
6045 "serde_json",
6046 + "smallvec",
6047 "tekken",
6048 "tempfile",
6049 "thiserror 2.0.18",
rust/Cargo.toml
+1
@@ -99,6 +99,7 @@ serde_tuple = "1.1.3"
99 serde_with = "3.18.0"
100 serial_test = { version = "3.2.0", features = ["file_locks"] }
101 sha2 = "0.10.9"
102 +smallvec = "1.15.1"
103 socket2 = "0.6.3"
104 strum = { version = "0.27.2", features = ["derive"] }
105 subenum = "1.1.3"
rust/src/chat/src/output/default/unified.rs
+14 -13
@@ -13,7 +13,7 @@ use futures::{StreamExt as _, pin_mut};
13 use thiserror_ext::AsReport;
14 use tracing::warn;
15 use vllm_parser::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};
16 -use vllm_text::output::DecodedTextEvent;
16 +use vllm_text::output::{DecodedTextEvent, SampledDelta};
17
18 use crate::Result;
19 use crate::error::Error;
@@ -237,12 +237,15 @@ pub(crate) async fn unified_event_stream(
237 .await;
238 }
239 DecodedTextEvent::TextDelta {
240 - delta,
241 - token_ids,
242 - logprobs,
240 + decoded,
241 + sampled:
242 + SampledDelta {
243 + token_ids,
244 + logprobs,
245 + },
246 finished,
247 } => {
245 - for next in state.process_delta(delta)? {
248 + for next in state.process_delta(decoded.text)? {
249 y.yield_ok(next).await;
250 }
251 if logprobs.is_some() || !token_ids.is_empty() {
@@ -375,24 +378,22 @@ mod tests {
378
379 fn decoded_delta(delta: &str) -> vllm_text::output::DecodedTextEvent {
380 vllm_text::output::DecodedTextEvent::TextDelta {
378 - delta: delta.to_string(),
379 - token_ids: Vec::new(),
380 - logprobs: None,
381 + decoded: vllm_text::DecodedText::unattributed(delta),
382 + sampled: vllm_text::SampledDelta::default(),
383 finished: None,
384 }
385 }
386
387 fn finished_delta(delta: &str) -> vllm_text::output::DecodedTextEvent {
388 vllm_text::output::DecodedTextEvent::TextDelta {
387 - delta: delta.to_string(),
388 - token_ids: Vec::new(),
389 - logprobs: None,
390 - finished: Some(vllm_text::output::Finished {
389 + decoded: vllm_text::DecodedText::unattributed(delta),
390 + sampled: vllm_text::SampledDelta::default(),
391 + finished: Some(Box::new(vllm_text::output::Finished {
392 usage: vllm_llm::TokenUsage::default(),
393 finish_reason: crate::FinishReason::Stop(None),
394 kv_transfer_params: None,
395 ec_transfer_params: None,
395 - }),
396 + })),
397 }
398 }
399
rust/src/chat/src/output/harmony/mod.rs
+7 -4
@@ -11,7 +11,7 @@ use asynk_strim_attr::{TryYielder, try_stream};
11 use futures::StreamExt as _;
12 use openai_harmony::chat::{Content as HarmonyContent, Message as HarmonyMessage, Role};
13 use openai_harmony::{HarmonyEncoding, StreamableParser};
14 -use vllm_text::output::DecodedTextEvent;
14 +use vllm_text::output::{DecodedTextEvent, SampledDelta};
15
16 use crate::Result as ChatResult;
17 use crate::error::{Error, Result};
@@ -345,9 +345,12 @@ async fn harmony_assistant_event_stream(
345 .await;
346 }
347 DecodedTextEvent::TextDelta {
348 - delta: _, // harmony takes raw token IDs as input, so we ignore text deltas here
349 - token_ids,
350 - logprobs,
348 + decoded: _, // harmony takes raw token IDs as input, so we ignore decoded text here
349 + sampled:
350 + SampledDelta {
351 + token_ids,
352 + logprobs,
353 + },
354 finished,
355 } => {
356 for event in state.process_token_ids(&token_ids)? {
rust/src/chat/src/output/harmony/tests.rs
+30 -49
@@ -6,7 +6,9 @@ use std::sync::Arc;
6 use futures::executor::block_on;
7 use futures::{TryStreamExt as _, stream};
8 use openai_harmony::chat::{Message, Role};
9 -use vllm_text::output::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, Finished};
9 +use vllm_text::output::{
10 + DecodedLogprobs, DecodedPositionLogprobs, DecodedText, DecodedTextEvent, Finished, SampledDelta,
11 +};
12
13 use super::*;
14 use crate::output::ChatOutputProcessor;
@@ -59,6 +61,21 @@ fn finished() -> Finished {
61 }
62 }
63
64 +fn decoded_tokens(
65 + token_ids: Vec<u32>,
66 + logprobs: Option<DecodedLogprobs>,
67 + finished: Option<Finished>,
68 +) -> DecodedTextEvent {
69 + DecodedTextEvent::TextDelta {
70 + decoded: DecodedText::default(),
71 + sampled: SampledDelta {
72 + token_ids,
73 + logprobs,
74 + },
75 + finished: finished.map(Box::new),
76 + }
77 +}
78 +
79 async fn collect_events(
80 processor: HarmonyChatOutputProcessor,
81 events: Vec<DecodedTextEvent>,
@@ -96,12 +113,7 @@ fn interrupted_final_message_is_preserved() {
113 HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
114 vec![
115 decoded_start(),
99 - DecodedTextEvent::TextDelta {
100 - delta: String::new(),
101 - token_ids: tokens[..tokens.len() - 1].to_vec(),
102 - logprobs: None,
103 - finished: Some(finished()),
104 - },
116 + decoded_tokens(tokens[..tokens.len() - 1].to_vec(), None, Some(finished())),
117 ],
118 ));
119
@@ -135,12 +147,7 @@ fn eos_flush_preserves_trailing_replacement_text() {
147 HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
148 vec![
149 decoded_start(),
138 - DecodedTextEvent::TextDelta {
139 - delta: String::new(),
140 - token_ids: tokens,
141 - logprobs: None,
142 - finished: Some(finished()),
143 - },
150 + decoded_tokens(tokens, None, Some(finished())),
151 ],
152 ));
153
@@ -157,12 +164,7 @@ fn interrupted_analysis_message_is_preserved() {
164 HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
165 vec![
166 decoded_start(),
160 - DecodedTextEvent::TextDelta {
161 - delta: String::new(),
162 - token_ids: tokens[..tokens.len() - 1].to_vec(),
163 - logprobs: None,
164 - finished: Some(finished()),
165 - },
167 + decoded_tokens(tokens[..tokens.len() - 1].to_vec(), None, Some(finished())),
168 ],
169 ));
170
@@ -196,12 +198,7 @@ fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() {
198 HarmonyChatOutputProcessor::new(&request_with_tools()).unwrap(),
199 vec![
200 decoded_start(),
199 - DecodedTextEvent::TextDelta {
200 - delta: String::new(),
201 - token_ids: tokens,
202 - logprobs: None,
203 - finished: Some(finished()),
204 - },
201 + decoded_tokens(tokens, None, Some(finished())),
202 ],
203 ));
204
@@ -225,12 +222,7 @@ fn multiple_messages_get_newline_separators() {
222 HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
223 vec![
224 decoded_start(),
228 - DecodedTextEvent::TextDelta {
229 - delta: String::new(),
230 - token_ids: tokens,
231 - logprobs: None,
232 - finished: Some(finished()),
233 - },
225 + decoded_tokens(tokens, None, Some(finished())),
226 ],
227 ));
228
@@ -256,18 +248,8 @@ fn tool_calls_stream_arguments_and_finish_with_local_id_shape() {
248 HarmonyChatOutputProcessor::new(&request_with_tools()).unwrap(),
249 vec![
250 decoded_start(),
259 - DecodedTextEvent::TextDelta {
260 - delta: String::new(),
261 - token_ids: tokens[..midpoint].to_vec(),
262 - logprobs: None,
263 - finished: None,
264 - },
265 - DecodedTextEvent::TextDelta {
266 - delta: String::new(),
267 - token_ids: tokens[midpoint..].to_vec(),
268 - logprobs: None,
269 - finished: Some(finished()),
270 - },
251 + decoded_tokens(tokens[..midpoint].to_vec(), None, None),
252 + decoded_tokens(tokens[midpoint..].to_vec(), None, Some(finished())),
253 ],
254 ));
255
@@ -303,14 +285,13 @@ fn semantic_events_precede_same_update_logprobs() {
285 HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
286 vec![
287 decoded_start(),
306 - DecodedTextEvent::TextDelta {
307 - delta: String::new(),
308 - token_ids: tokens,
309 - logprobs: Some(DecodedLogprobs {
288 + decoded_tokens(
289 + tokens,
290 + Some(DecodedLogprobs {
291 positions: vec![DecodedPositionLogprobs { entries: vec![] }],
292 }),
312 - finished: Some(finished()),
313 - },
293 + Some(finished()),
294 + ),
295 ],
296 ));
297
rust/src/chat/tests/roundtrip.rs
+13 -12
@@ -784,15 +784,14 @@ fn decoded_completion_stream(
784 if chunks.is_empty() {
785 events.push({
786 DecodedTextEvent::TextDelta {
787 - delta: String::new(),
788 - token_ids: Vec::new(),
789 - logprobs: None,
790 - finished: Some(Finished {
787 + decoded: vllm_text::DecodedText::default(),
788 + sampled: vllm_text::SampledDelta::default(),
789 + finished: Some(Box::new(Finished {
790 usage: Default::default(),
791 finish_reason: FinishReason::stop_eos(),
792 kv_transfer_params: None,
793 ec_transfer_params: None,
795 - }),
794 + })),
795 }
796 });
797 } else {
@@ -805,10 +804,12 @@ fn decoded_completion_stream(
804 ec_transfer_params: None,
805 });
806 events.push(DecodedTextEvent::TextDelta {
808 - delta: chunk.delta,
809 - token_ids: chunk.token_ids,
810 - logprobs: None,
811 - finished,
807 + decoded: vllm_text::DecodedText::unattributed(chunk.delta),
808 + sampled: vllm_text::SampledDelta {
809 + token_ids: chunk.token_ids,
810 + logprobs: None,
811 + },
812 + finished: finished.map(Box::new),
813 });
814 }
815 }
@@ -864,7 +865,7 @@ fn incremental_decode_chunks(
865 for token_id in chunk_token_ids.iter().copied() {
866 decoder.push_token(token_id)?;
867 while let Some(chunk) = decoder.next_chunk() {
867 - delta.push_str(&chunk);
868 + delta.push_str(&chunk.text);
869 }
870 }
871 chunks.push(DecodedCompletionChunk {
@@ -876,10 +877,10 @@ fn incremental_decode_chunks(
877 let (last_chunk, _) = decoder.flush(None)?;
878 if let Some(last_chunk) = last_chunk {
879 if let Some(delta) = chunks.last_mut() {
879 - delta.delta.push_str(&last_chunk);
880 + delta.delta.push_str(&last_chunk.text);
881 } else {
882 chunks.push(DecodedCompletionChunk {
882 - delta: last_chunk,
883 + delta: last_chunk.text,
884 token_ids: Vec::new(),
885 });
886 }
rust/src/server/src/grpc/inference.rs
+9 -6
@@ -14,7 +14,7 @@ use tracing::{Span, info, info_span, warn};
14 use tracing_futures::Instrument as _;
15 use uuid::Uuid;
16 use vllm_llm::current_unix_timestamp_secs;
17 -use vllm_text::{DecodedTextEvent, Prompt, TextOutputStreamExt as _, TextRequest};
17 +use vllm_text::{DecodedTextEvent, Prompt, SampledDelta, TextOutputStreamExt as _, TextRequest};
18
19 use super::convert::{self, ResponseOpts};
20 use super::{InferenceServer, pb};
@@ -264,17 +264,20 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
264 })
265 }
266 Ok(DecodedTextEvent::TextDelta {
267 - delta,
268 - token_ids,
269 - logprobs,
267 + decoded,
268 + sampled:
269 + SampledDelta {
270 + token_ids,
271 + logprobs,
272 + },
273 finished,
274 }) => Ok(pb::GenerateResponse {
275 prompt_info: None,
276 outputs: Some(convert::to_sequence_output(
274 - &delta,
277 + &decoded.text,
278 &token_ids,
279 logprobs.as_ref(),
277 - finished.as_ref(),
280 + finished.as_deref(),
281 &response_opts,
282 )),
283 }),
rust/src/server/src/routes/openai/completions.rs
+68 -48
@@ -23,7 +23,7 @@ use tracing_futures::Instrument as _;
23 use vllm_engine_core_client::protocol::output::StopReason;
24 use vllm_text::tokenizer::Tokenizer;
25 use vllm_text::{
26 - DecodedPromptLogprobs, DecodedTextEvent, FinishReason, TextOutputStream,
26 + DecodedPromptLogprobs, DecodedTextEvent, FinishReason, SampledDelta, TextOutputStream,
27 TextOutputStreamExt as _, TextRequest,
28 };
29
@@ -317,11 +317,15 @@ async fn completion_chunk_stream(
317 }
318 }
319 Ok(DecodedTextEvent::TextDelta {
320 - delta,
321 - token_ids,
322 - logprobs,
320 + decoded,
321 + sampled:
322 + SampledDelta {
323 + token_ids,
324 + logprobs,
325 + },
326 finished,
327 }) => {
328 + let delta = decoded.text;
329 // Prompt-only streaming already emitted the echoed prompt in the Start chunk.
330 // The one generated token is only used to drive the engine to a finished event,
331 // so hide its delta and forward only the terminal finish/usage metadata.
@@ -561,8 +565,8 @@ mod tests {
565 use itertools::Itertools as _;
566 use vllm_engine_core_client::protocol::output::StopReason;
567 use vllm_text::{
564 - DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, DecodedTextEvent,
565 - DecodedTokenLogprob, FinishReason, Finished,
568 + DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, DecodedText,
569 + DecodedTextEvent, DecodedTokenLogprob, FinishReason, Finished, SampledDelta,
570 };
571
572 use super::{
@@ -570,6 +574,22 @@ mod tests {
574 StreamResponseEnvelope, completion_chunk_stream, final_chunk,
575 };
576
577 + fn decoded_delta(
578 + text: impl Into<String>,
579 + token_ids: Vec<u32>,
580 + logprobs: Option<DecodedLogprobs>,
581 + finished: Option<Finished>,
582 + ) -> DecodedTextEvent {
583 + DecodedTextEvent::TextDelta {
584 + decoded: DecodedText::unattributed(text),
585 + sampled: SampledDelta {
586 + token_ids,
587 + logprobs,
588 + },
589 + finished: finished.map(Box::new),
590 + }
591 + }
592 +
593 fn stream_envelope() -> Arc<StreamResponseEnvelope> {
594 Arc::new(StreamResponseEnvelope::new(
595 "cmpl-1".to_string(),
@@ -619,10 +639,10 @@ mod tests {
639 prompt_token_ids: vec![1, 2, 3, 4, 5].into(),
640 prompt_logprobs: None,
641 }),
622 - Ok(DecodedTextEvent::TextDelta {
623 - delta: "h".to_string(),
624 - token_ids: vec![b'h' as u32],
625 - logprobs: Some(DecodedLogprobs {
642 + Ok(decoded_delta(
643 + "h",
644 + vec![b'h' as u32],
645 + Some(DecodedLogprobs {
646 positions: vec![DecodedPositionLogprobs {
647 entries: vec![
648 DecodedTokenLogprob {
@@ -640,12 +660,12 @@ mod tests {
660 ],
661 }],
662 }),
643 - finished: None,
644 - }),
645 - Ok(DecodedTextEvent::TextDelta {
646 - delta: String::new(),
647 - token_ids: vec![b'!' as u32],
648 - logprobs: Some(DecodedLogprobs {
663 + None,
664 + )),
665 + Ok(decoded_delta(
666 + "",
667 + vec![b'!' as u32],
668 + Some(DecodedLogprobs {
669 positions: vec![DecodedPositionLogprobs {
670 entries: vec![
671 DecodedTokenLogprob {
@@ -663,7 +683,7 @@ mod tests {
683 ],
684 }],
685 }),
666 - finished: Some(Finished {
686 + Some(Finished {
687 usage: vllm_llm::TokenUsage {
688 prompt_token_count: 5,
689 output_token_count: 2,
@@ -675,7 +695,7 @@ mod tests {
695 kv_transfer_params: None,
696 ec_transfer_params: None,
697 }),
678 - }),
698 + )),
699 ]);
700
701 let chunks = completion_chunk_stream(
@@ -773,11 +793,11 @@ mod tests {
793 prompt_token_ids: vec![1, 2].into(),
794 prompt_logprobs: None,
795 }),
776 - Ok(DecodedTextEvent::TextDelta {
777 - delta: " leaked".to_string(),
778 - token_ids: vec![3],
779 - logprobs: None,
780 - finished: Some(Finished {
796 + Ok(decoded_delta(
797 + " leaked",
798 + vec![3],
799 + None,
800 + Some(Finished {
801 usage: vllm_llm::TokenUsage {
802 prompt_token_count: 2,
803 output_token_count: 1,
@@ -787,7 +807,7 @@ mod tests {
807 kv_transfer_params: None,
808 ec_transfer_params: None,
809 }),
790 - }),
810 + )),
811 ]);
812
813 let response = super::collect_completion(
@@ -825,11 +845,11 @@ mod tests {
845 prompt_token_ids: vec![9707].into(),
846 prompt_logprobs: None,
847 }),
828 - Ok(DecodedTextEvent::TextDelta {
829 - delta: " leaked".to_string(),
830 - token_ids: vec![3],
831 - logprobs: None,
832 - finished: Some(Finished {
848 + Ok(decoded_delta(
849 + " leaked",
850 + vec![3],
851 + None,
852 + Some(Finished {
853 usage: vllm_llm::TokenUsage {
854 prompt_token_count: 1,
855 output_token_count: 1,
@@ -839,7 +859,7 @@ mod tests {
859 kv_transfer_params: None,
860 ec_transfer_params: None,
861 }),
842 - }),
862 + )),
863 ]);
864
865 let response = super::collect_completion(
@@ -880,11 +900,11 @@ mod tests {
900 prompt_token_ids: vec![1, 2].into(),
901 prompt_logprobs: None,
902 }),
883 - Ok(DecodedTextEvent::TextDelta {
884 - delta: " leaked".to_string(),
885 - token_ids: vec![3],
886 - logprobs: None,
887 - finished: Some(Finished {
903 + Ok(decoded_delta(
904 + " leaked",
905 + vec![3],
906 + None,
907 + Some(Finished {
908 usage: vllm_llm::TokenUsage {
909 prompt_token_count: 2,
910 output_token_count: 1,
@@ -894,7 +914,7 @@ mod tests {
914 kv_transfer_params: None,
915 ec_transfer_params: None,
916 }),
897 - }),
917 + )),
918 ]);
919
920 let chunks = completion_chunk_stream(
@@ -952,11 +972,11 @@ mod tests {
972 prompt_token_ids: vec![9707].into(),
973 prompt_logprobs: None,
974 }),
955 - Ok(DecodedTextEvent::TextDelta {
956 - delta: " leaked".to_string(),
957 - token_ids: vec![3],
958 - logprobs: None,
959 - finished: Some(Finished {
975 + Ok(decoded_delta(
976 + " leaked",
977 + vec![3],
978 + None,
979 + Some(Finished {
980 usage: vllm_llm::TokenUsage {
981 prompt_token_count: 1,
982 output_token_count: 1,
@@ -966,7 +986,7 @@ mod tests {
986 kv_transfer_params: None,
987 ec_transfer_params: None,
988 }),
969 - }),
989 + )),
990 ]);
991
992 let chunks = completion_chunk_stream(
@@ -1026,11 +1046,11 @@ mod tests {
1046 }],
1047 }),
1048 }),
1029 - Ok(DecodedTextEvent::TextDelta {
1030 - delta: " leaked".to_string(),
1031 - token_ids: vec![3],
1032 - logprobs: None,
1033 - finished: Some(Finished {
1049 + Ok(decoded_delta(
1050 + " leaked",
1051 + vec![3],
1052 + None,
1053 + Some(Finished {
1054 usage: vllm_llm::TokenUsage {
1055 prompt_token_count: 2,
1056 output_token_count: 1,
@@ -1040,7 +1060,7 @@ mod tests {
1060 kv_transfer_params: None,
1061 ec_transfer_params: None,
1062 }),
1043 - }),
1063 + )),
1064 ]);
1065
1066 let chunks = completion_chunk_stream(
rust/src/text/src/lib.rs
+2 -1
@@ -19,7 +19,8 @@ pub use lower::{
19 };
20 pub use output::{
21 CollectedTextOutput, DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs,
22 - DecodedTextEvent, DecodedTokenLogprob, Finished, TextDecodeOptions, TextOutputStreamExt,
22 + DecodedText, DecodedTextEvent, DecodedTokenLogprob, Finished, SampledDelta, TextDecodeOptions,
23 + TextOutputStreamExt, TokenAnchor, TokenAttribution,
24 };
25 pub use request::{Prompt, SamplingParams, TextRequest, normalize_top_k};
26 use trait_set::trait_set;
rust/src/text/src/output/decoded.rs
+182 -48
@@ -10,7 +10,7 @@ use tracing::{Level, debug, trace};
10 use vllm_engine_core_client::AbortCause;
11 use vllm_engine_core_client::protocol::output::StopReason;
12 use vllm_llm::{FinishReason, GenerateOutput, TokenUsage};
13 -use vllm_tokenizer::{DynTokenizer, IncrementalDecoder};
13 +use vllm_tokenizer::{DecodedText, DynTokenizer, IncrementalDecoder};
14
15 use super::logprobs::{
16 DecodedLogprobs, DecodedPromptLogprobs, decode_logprobs, decode_prompt_logprobs,
@@ -52,6 +52,19 @@ pub struct Finished {
52 pub ec_transfer_params: Option<serde_json::Value>,
53 }
54
55 +/// Sample metadata emitted by one engine output update.
56 +///
57 +/// Token IDs and logprobs share the engine-update clock and align by generated
58 +/// token position. They may arrive before the corresponding decoded text is
59 +/// visible or alongside decoded text attributed to earlier updates.
60 +#[derive(Debug, Clone, Default, PartialEq)]
61 +pub struct SampledDelta {
62 + /// Token IDs produced by this engine update, after stop-string truncation.
63 + pub token_ids: Vec<u32>,
64 + /// Per-position logprobs aligned with `token_ids`, when requested.
65 + pub logprobs: Option<DecodedLogprobs>,
66 +}
67 +
68 /// Internal decoded-text event emitted before higher-level assistant
69 /// adaptation.
70 #[derive(Debug, Clone, PartialEq)]
@@ -68,25 +81,25 @@ pub enum DecodedTextEvent {
81 /// remaining prompt positions.
82 prompt_logprobs: Option<DecodedPromptLogprobs>,
83 },
71 - /// A delta of text has been decoded, optionally alongside token-position
72 - /// logprobs.
84 + /// Decoded text and sampled metadata observed during one output update.
85 ///
74 - /// `delta` is the newly visible decoded text fragment for this update.
86 + /// `decoded` follows the decoder clock. It may be empty while tokens remain
87 + /// undecodable, and may later contain token attributions from earlier
88 + /// engine updates. Tokens suppressed before decoding are absent.
89 ///
76 - /// `logprobs` covers the newly generated token positions from the same
77 - /// update, but is not guaranteed to align with `delta` by character
78 - /// span. One update may carry token logprobs but no newly visible text
79 - /// yet, and one visible text fragment may reflect multiple token
80 - /// positions becoming decodable together.
90 + /// `sampled` follows the engine-update clock. Its token IDs and logprobs
91 + /// align with each other. `decoded.attributions` has an independent decoder
92 + /// clock and positional sequence.
93 ///
82 - /// Upper-level may further parse `delta` as reasoning or tool calls.
94 + /// Upper-level may further parse `decoded.text` as reasoning or tool calls.
95 ///
96 /// When `finished` is `Some`, this is the terminal event for the request.
97 TextDelta {
86 - delta: String,
87 - token_ids: Vec<u32>,
88 - logprobs: Option<DecodedLogprobs>,
89 - finished: Option<Finished>,
98 + decoded: DecodedText,
99 + sampled: SampledDelta,
100 + /// Heap-backed because terminal metadata is large and appears once per
101 + /// request, while every streaming event carries this enum variant.
102 + finished: Option<Box<Finished>>,
103 },
104 }
105
@@ -104,9 +117,9 @@ pub async fn decoded_text_event_stream(
117 let mut decoder: Option<Box<dyn IncrementalDecoder>> = None;
118 let mut prompt_token_count = 0_usize;
119 let mut cached_token_count = 0_usize;
107 - let mut token_ids = Vec::new();
120 + let mut sampled_token_ids = Vec::new();
121 let mut output_token_count: usize = 0;
109 - let mut logprobs: Option<DecodedLogprobs> = None;
122 + let mut sampled_logprobs: Option<DecodedLogprobs> = None;
123
124 while let Some(next) = raw_stream.next().await {
125 let output = next?;
@@ -168,7 +181,7 @@ pub async fn decoded_text_event_stream(
181 &output.token_ids
182 };
183
171 - let mut delta: Option<String> = None;
184 + let mut decoded = DecodedText::default();
185 let mut truncate_output_to = None;
186 let mut truncate_tokens_to = None;
187 for (tok_idx, &token_id) in decodable_token_ids.iter().enumerate() {
@@ -189,12 +202,9 @@ pub async fn decoded_text_event_stream(
202 break;
203 }
204
205 + // TODO: avoid generating a chunk every time a token is pushed
206 if intermediate && let Some(chunk) = decoder.next_chunk() {
193 - if let Some(delta_str) = delta.as_mut() {
194 - delta_str.push_str(&chunk);
195 - } else {
196 - delta = Some(chunk);
197 - }
207 + decoded.append(chunk);
208 }
209 }
210
@@ -223,9 +233,9 @@ pub async fn decoded_text_event_stream(
233 .transpose()?;
234
235 if !intermediate {
226 - token_ids.extend(&new_token_ids);
236 + sampled_token_ids.extend(&new_token_ids);
237 if let Some(dlp) = decoded_logprobs.as_ref() {
228 - logprobs
238 + sampled_logprobs
239 .get_or_insert_with(|| DecodedLogprobs { positions: vec![] })
240 .positions
241 .extend_from_slice(&dlp.positions);
@@ -234,22 +244,19 @@ pub async fn decoded_text_event_stream(
244
245 if let Some(reason) = finish_reason {
246 // Flush any remaining buffered text.
237 - let (last_chunk, mut text) = decoder.flush(truncate_output_to)?;
238 - let text_len = text.len();
239 - let full_text = tracing::enabled!(Level::TRACE).then(|| text.clone());
240 -
241 - if intermediate {
247 + let (last_chunk, full_decoded) = decoder.flush(truncate_output_to)?;
248 + let text_len = full_decoded.text.len();
249 + let full_text = tracing::enabled!(Level::TRACE).then(|| full_decoded.text.clone());
250 + let decoded = if intermediate {
251 if let Some(chunk) = last_chunk {
243 - if let Some(delta_str) = delta.as_mut() {
244 - delta_str.push_str(&chunk);
245 - } else {
246 - delta = Some(chunk);
247 - }
252 + decoded.append(chunk);
253 }
249 - token_ids = new_token_ids;
250 - logprobs = decoded_logprobs;
251 - text = delta.unwrap_or_default();
252 - }
254 + sampled_token_ids = new_token_ids;
255 + sampled_logprobs = decoded_logprobs;
256 + decoded
257 + } else {
258 + full_decoded
259 + };
260
261 debug!(
262 finish_reason = ?reason,
@@ -269,10 +276,12 @@ pub async fn decoded_text_event_stream(
276 }
277
278 y.yield_ok(DecodedTextEvent::TextDelta {
272 - delta: text,
273 - token_ids,
274 - logprobs,
275 - finished: Some(Finished {
279 + decoded,
280 + sampled: SampledDelta {
281 + token_ids: sampled_token_ids,
282 + logprobs: sampled_logprobs,
283 + },
284 + finished: Some(Box::new(Finished {
285 usage: TokenUsage {
286 prompt_token_count,
287 output_token_count,
@@ -281,7 +290,7 @@ pub async fn decoded_text_event_stream(
290 finish_reason: reason,
291 kv_transfer_params,
292 ec_transfer_params,
284 - }),
293 + })),
294 })
295 .await;
296 return Ok(());
@@ -289,9 +298,11 @@ pub async fn decoded_text_event_stream(
298
299 if intermediate {
300 y.yield_ok(DecodedTextEvent::TextDelta {
292 - delta: delta.unwrap_or_default(),
293 - token_ids: new_token_ids,
294 - logprobs: decoded_logprobs,
301 + decoded,
302 + sampled: SampledDelta {
303 + token_ids: new_token_ids,
304 + logprobs: decoded_logprobs,
305 + },
306 finished: None,
307 })
308 .await;
@@ -338,10 +349,10 @@ mod tests {
349 use std::sync::{Arc, Mutex};
350 use std::task::{Context, Poll};
351
341 - use futures::{Stream, stream};
352 + use futures::{Stream, StreamExt as _, stream};
353 use vllm_engine_core_client::AbortCause;
354 use vllm_llm::GenerateOutput;
344 - use vllm_tokenizer::test_utils::TestTokenizer;
355 + use vllm_tokenizer::{TokenAnchor, TokenAttribution, test_utils::TestTokenizer};
356
357 use super::*;
358 use crate::output::TextOutputStreamExt as _;
@@ -378,6 +389,129 @@ mod tests {
389 }
390 }
391
392 + fn visible(token_id: u32, byte_offset: u32) -> TokenAttribution {
393 + TokenAttribution {
394 + token_id,
395 + anchor: TokenAnchor::Visible { byte_offset },
396 + }
397 + }
398 +
399 + fn text_delta(
400 + text: impl Into<String>,
401 + attributions: impl IntoIterator<Item = TokenAttribution>,
402 + sampled_token_ids: Vec<u32>,
403 + finished: Option<Finished>,
404 + ) -> DecodedTextEvent {
405 + DecodedTextEvent::TextDelta {
406 + decoded: DecodedText {
407 + text: text.into(),
408 + attributions: attributions.into_iter().collect(),
409 + },
410 + sampled: SampledDelta {
411 + token_ids: sampled_token_ids,
412 + logprobs: None,
413 + },
414 + finished: finished.map(Box::new),
415 + }
416 + }
417 +
418 + async fn run_intermediate(
419 + outputs: Vec<GenerateOutput>,
420 + decode_options: TextDecodeOptions,
421 + ) -> Vec<DecodedTextEvent> {
422 + let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
423 + decoded_text_event_stream(
424 + "test".into(),
425 + tokenizer,
426 + stream::iter(outputs.into_iter().map(Ok)),
427 + decode_options,
428 + true,
429 + )
430 + .collect::<Vec<_>>()
431 + .await
432 + .into_iter()
433 + .collect::<crate::Result<Vec<_>>>()
434 + .unwrap()
435 + }
436 +
437 + #[tokio::test]
438 + async fn decoded_and_sampled_clocks_are_independent() {
439 + struct Case {
440 + name: &'static str,
441 + outputs: Vec<GenerateOutput>,
442 + expected: Vec<DecodedTextEvent>,
443 + }
444 +
445 + let prompt: Arc<[u32]> = Arc::from([]);
446 + let cases = [
447 + Case {
448 + name: "utf8 byte fallback resolves across engine updates",
449 + outputs: vec![
450 + GenerateOutput::for_test(Some(Arc::clone(&prompt)), vec![0xe4], None),
451 + GenerateOutput::for_test(None, vec![0xbd, 0xa0], Some(FinishReason::Length)),
452 + ],
453 + expected: vec![
454 + DecodedTextEvent::Start {
455 + prompt_token_ids: Arc::clone(&prompt),
456 + prompt_logprobs: None,
457 + },
458 + text_delta("", [], vec![0xe4], None),
459 + text_delta(
460 + "你",
461 + [visible(0xe4, 0), visible(0xbd, 0), visible(0xa0, 0)],
462 + vec![0xbd, 0xa0],
463 + Some(Finished {
464 + usage: TokenUsage {
465 + prompt_token_count: 0,
466 + output_token_count: 3,
467 + cached_token_count: 0,
468 + },
469 + finish_reason: FinishReason::Length,
470 + kv_transfer_params: None,
471 + ec_transfer_params: None,
472 + }),
473 + ),
474 + ],
475 + },
476 + Case {
477 + name: "terminal stop token remains sampled and bypasses decoder",
478 + outputs: vec![GenerateOutput::for_test(
479 + Some(Arc::clone(&prompt)),
480 + ascii_tokens("a!"),
481 + Some(FinishReason::Stop(Some(StopReason::TokenId(b'!' as u32)))),
482 + )],
483 + expected: vec![
484 + DecodedTextEvent::Start {
485 + prompt_token_ids: Arc::clone(&prompt),
486 + prompt_logprobs: None,
487 + },
488 + text_delta(
489 + "a",
490 + [visible(b'a' as u32, 0)],
491 + ascii_tokens("a!"),
492 + Some(Finished {
493 + usage: TokenUsage {
494 + prompt_token_count: 0,
495 + output_token_count: 2,
496 + cached_token_count: 0,
497 + },
498 + finish_reason: FinishReason::Stop(Some(StopReason::TokenId(
499 + b'!' as u32,
500 + ))),
501 + kv_transfer_params: None,
502 + ec_transfer_params: None,
503 + }),
504 + ),
505 + ],
506 + },
507 + ];
508 +
509 + for case in cases {
510 + let actual = run_intermediate(case.outputs, TextDecodeOptions::default()).await;
511 + assert_eq!(actual, case.expected, "{}", case.name);
512 + }
513 + }
514 +
515 struct DropRecordingStream {
516 next: Option<vllm_llm::Result<GenerateOutput>>,
517 dropped_cause: Arc<Mutex<Option<AbortCause>>>,
rust/src/text/src/output/mod.rs
+96 -84
@@ -3,10 +3,13 @@
3
4 //! Output processing helpers shared by text and chat layers.
5
6 -pub use decoded::{DecodedTextEvent, Finished, TextDecodeOptions, decoded_text_event_stream};
6 +pub use decoded::{
7 + DecodedTextEvent, Finished, SampledDelta, TextDecodeOptions, decoded_text_event_stream,
8 +};
9 pub use logprobs::{
10 DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, DecodedTokenLogprob,
11 };
12 +pub use vllm_tokenizer::{DecodedText, TokenAnchor, TokenAttribution};
13
14 mod decoded;
15 mod logprobs;
@@ -58,13 +61,16 @@ impl<T: TextOutputStream> T {
61 prompt_token_ids = start_prompt_token_ids;
62 }
63 DecodedTextEvent::TextDelta {
61 - delta,
62 - token_ids: delta_token_ids,
63 - logprobs: mut delta_logprobs,
64 + decoded,
65 + sampled:
66 + SampledDelta {
67 + token_ids: delta_token_ids,
68 + logprobs: mut delta_logprobs,
69 + },
70 finished,
71 } => {
72 if let Some(c) = collected.as_mut() {
67 - c.text.push_str(&delta);
73 + c.text.push_str(&decoded.text);
74 c.token_ids.extend(delta_token_ids);
75 if let Some(dlp) = delta_logprobs.as_mut() {
76 if let Some(lp) = c.logprobs.as_mut() {
@@ -75,7 +81,7 @@ impl<T: TextOutputStream> T {
81 }
82 } else {
83 collected = Some(CollectedTextOutput {
78 - text: delta,
84 + text: decoded.text,
85 prompt_token_ids: Arc::clone(&prompt_token_ids),
86 prompt_logprobs: prompt_logprobs.take(),
87 logprobs: delta_logprobs,
@@ -134,29 +140,31 @@ mod tests {
140 }),
141 }),
142 Ok(DecodedTextEvent::TextDelta {
137 - delta: "bc".to_string(),
138 - token_ids: vec![1, 2],
139 - logprobs: Some(DecodedLogprobs {
140 - positions: vec![
141 - DecodedPositionLogprobs {
142 - entries: vec![DecodedTokenLogprob {
143 - token_id: 0,
144 - token: "a".to_string(),
145 - logprob: -0.2,
146 - rank: 1,
147 - }],
148 - },
149 - DecodedPositionLogprobs {
150 - entries: vec![DecodedTokenLogprob {
151 - token_id: 0,
152 - token: "bc".to_string(),
153 - logprob: -0.3,
154 - rank: 1,
155 - }],
156 - },
157 - ],
158 - }),
159 - finished: Some(Finished {
143 + decoded: DecodedText::unattributed("bc"),
144 + sampled: SampledDelta {
145 + token_ids: vec![1, 2],
146 + logprobs: Some(DecodedLogprobs {
147 + positions: vec![
148 + DecodedPositionLogprobs {
149 + entries: vec![DecodedTokenLogprob {
150 + token_id: 0,
151 + token: "a".to_string(),
152 + logprob: -0.2,
153 + rank: 1,
154 + }],
155 + },
156 + DecodedPositionLogprobs {
157 + entries: vec![DecodedTokenLogprob {
158 + token_id: 0,
159 + token: "bc".to_string(),
160 + logprob: -0.3,
161 + rank: 1,
162 + }],
163 + },
164 + ],
165 + }),
166 + },
167 + finished: Some(Box::new(Finished {
168 usage: vllm_llm::TokenUsage {
169 prompt_token_count: 2,
170 output_token_count: 2,
@@ -165,7 +173,7 @@ mod tests {
173 finish_reason: FinishReason::stop_eos(),
174 kv_transfer_params: None,
175 ec_transfer_params: None,
168 - }),
176 + })),
177 }),
178 ]);
179
@@ -219,62 +227,66 @@ mod tests {
227 prompt_logprobs: None,
228 }),
229 Ok(DecodedTextEvent::TextDelta {
222 - delta: "he".to_string(),
223 - token_ids: vec![1, 2],
224 - logprobs: Some(DecodedLogprobs {
225 - positions: vec![
226 - DecodedPositionLogprobs {
227 - entries: vec![DecodedTokenLogprob {
228 - token_id: 0,
229 - token: "h".to_string(),
230 - logprob: -0.1,
231 - rank: 1,
232 - }],
233 - },
234 - DecodedPositionLogprobs {
235 - entries: vec![DecodedTokenLogprob {
236 - token_id: 0,
237 - token: "e".to_string(),
238 - logprob: -0.2,
239 - rank: 1,
240 - }],
241 - },
242 - ],
243 - }),
230 + decoded: DecodedText::unattributed("he"),
231 + sampled: SampledDelta {
232 + token_ids: vec![1, 2],
233 + logprobs: Some(DecodedLogprobs {
234 + positions: vec![
235 + DecodedPositionLogprobs {
236 + entries: vec![DecodedTokenLogprob {
237 + token_id: 0,
238 + token: "h".to_string(),
239 + logprob: -0.1,
240 + rank: 1,
241 + }],
242 + },
243 + DecodedPositionLogprobs {
244 + entries: vec![DecodedTokenLogprob {
245 + token_id: 0,
246 + token: "e".to_string(),
247 + logprob: -0.2,
248 + rank: 1,
249 + }],
250 + },
251 + ],
252 + }),
253 + },
254 finished: None,
255 }),
256 Ok(DecodedTextEvent::TextDelta {
247 - delta: "llo".to_string(),
248 - token_ids: vec![3, 4, 5],
249 - logprobs: Some(DecodedLogprobs {
250 - positions: vec![
251 - DecodedPositionLogprobs {
252 - entries: vec![DecodedTokenLogprob {
253 - token_id: 0,
254 - token: "l".to_string(),
255 - logprob: -0.3,
256 - rank: 1,
257 - }],
258 - },
259 - DecodedPositionLogprobs {
260 - entries: vec![DecodedTokenLogprob {
261 - token_id: 0,
262 - token: "l".to_string(),
263 - logprob: -0.4,
264 - rank: 1,
265 - }],
266 - },
267 - DecodedPositionLogprobs {
268 - entries: vec![DecodedTokenLogprob {
269 - token_id: 0,
270 - token: "o".to_string(),
271 - logprob: -0.5,
272 - rank: 1,
273 - }],
274 - },
275 - ],
276 - }),
277 - finished: Some(Finished {
257 + decoded: DecodedText::unattributed("llo"),
258 + sampled: SampledDelta {
259 + token_ids: vec![3, 4, 5],
260 + logprobs: Some(DecodedLogprobs {
261 + positions: vec![
262 + DecodedPositionLogprobs {
263 + entries: vec![DecodedTokenLogprob {
264 + token_id: 0,
265 + token: "l".to_string(),
266 + logprob: -0.3,
267 + rank: 1,
268 + }],
269 + },
270 + DecodedPositionLogprobs {
271 + entries: vec![DecodedTokenLogprob {
272 + token_id: 0,
273 + token: "l".to_string(),
274 + logprob: -0.4,
275 + rank: 1,
276 + }],
277 + },
278 + DecodedPositionLogprobs {
279 + entries: vec![DecodedTokenLogprob {
280 + token_id: 0,
281 + token: "o".to_string(),
282 + logprob: -0.5,
283 + rank: 1,
284 + }],
285 + },
286 + ],
287 + }),
288 + },
289 + finished: Some(Box::new(Finished {
290 usage: vllm_llm::TokenUsage {
291 prompt_token_count: 2,
292 output_token_count: 5,
@@ -283,7 +295,7 @@ mod tests {
295 finish_reason: FinishReason::stop_eos(),
296 kv_transfer_params: None,
297 ec_transfer_params: None,
286 - }),
298 + })),
299 }),
300 ]);
301
rust/src/tokenizer/Cargo.toml
+1
@@ -14,6 +14,7 @@ riptoken.workspace = true
14 rustc-hash.workspace = true
15 serde.workspace = true
16 serde_json.workspace = true
17 +smallvec.workspace = true
18 tekken.workspace = true
19 thiserror.workspace = true
20 thiserror-ext.workspace = true
rust/src/tokenizer/src/incremental.rs
+63 -46
@@ -1,10 +1,13 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -use std::mem::take;
5 -
4 use crate::{Result, Tokenizer};
5
6 +mod attribution;
7 +
8 +use attribution::AttributedTextBuffer;
9 +pub use attribution::{DecodedText, TokenAnchor, TokenAttribution};
10 +
11 /// Stateful incremental decoder that emits text chunks one token at a time.
12 pub trait IncrementalDecoder: Send {
13 /// Push one generated token and return how many new string bytes were
@@ -12,13 +15,16 @@ pub trait IncrementalDecoder: Send {
15 fn push_token(&mut self, token_id: u32) -> Result<usize>;
16
17 /// Consume any text which is currently ready.
15 - fn next_chunk(&mut self) -> Option<String>;
18 + fn next_chunk(&mut self) -> Option<DecodedText>;
19
20 /// Flush any remaining buffered text that has not yet been emitted.
21 ///
22 /// Called after the final generated token to force out buffered/incomplete
23 /// fragments.
21 - fn flush(&mut self, truncate_output_to: Option<usize>) -> Result<(Option<String>, String)>;
24 + fn flush(
25 + &mut self,
26 + truncate_output_to: Option<usize>,
27 + ) -> Result<(Option<DecodedText>, DecodedText)>;
28
29 /// Return cumulative decoded text so far.
30 fn output(&self) -> &str;
@@ -36,8 +42,7 @@ pub(crate) struct DecodeStream<'a, T: Tokenizer + ?Sized> {
42 prefix: String,
43 prefix_index: usize,
44 prefix_seeded: bool,
39 - cumulative_output: String,
40 - output_index: usize,
45 + decoded: AttributedTextBuffer,
46 }
47
48 impl<'a, T: Tokenizer + ?Sized> DecodeStream<'a, T> {
@@ -55,8 +60,7 @@ impl<'a, T: Tokenizer + ?Sized> DecodeStream<'a, T> {
60 prefix: String::new(),
61 prefix_index: 0,
62 prefix_seeded: prompt_token_ids.is_empty(),
58 - cumulative_output: String::new(),
59 - output_index: 0,
63 + decoded: AttributedTextBuffer::default(),
64 }
65 }
66 }
@@ -68,6 +72,14 @@ const SAFE_SUFFIX_MIN: usize = 4;
72 const SAFE_SUFFIX_MAX: usize = 6;
73
74 impl<T: Tokenizer + ?Sized> DecodeStream<'_, T> {
75 + fn record_token(&mut self, token_id: u32) {
76 + if self.skip_special_tokens && self.tokenizer.is_special_id(token_id) {
77 + self.decoded.record_zero_width_token(token_id);
78 + } else {
79 + self.decoded.record_pending_token(token_id);
80 + }
81 + }
82 +
83 /// Decode prompt-only context for prefix seeding.
84 ///
85 /// Prompt ids may come from the model vocabulary rather than the local
@@ -132,6 +144,7 @@ impl<T: Tokenizer + ?Sized> IncrementalDecoder for DecodeStream<'_, T> {
144 }
145
146 self.ids.push(token_id);
147 + self.record_token(token_id);
148 let string = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
149 let prefix_len = self.prefix.len();
150 if string.len() <= prefix_len || string.ends_with('\u{FFFD}') {
@@ -139,49 +152,48 @@ impl<T: Tokenizer + ?Sized> IncrementalDecoder for DecodeStream<'_, T> {
152 }
153 // Ensure we split at a utf-8 char boundary.
154 let new_chunk = &string[string.floor_char_boundary(prefix_len)..];
142 - self.cumulative_output.push_str(new_chunk);
155 + self.decoded.append_visible_text(new_chunk);
156 self.ids.drain(..self.prefix_index);
157 self.prefix = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
158 self.prefix_index = self.ids.len();
159 Ok(new_chunk.len())
160 }
161
149 - fn next_chunk(&mut self) -> Option<String> {
150 - let cutoff = self.cumulative_output.len().saturating_sub(self.min_bytes_to_buffer);
162 + fn next_chunk(&mut self) -> Option<DecodedText> {
163 + let cutoff = self.decoded.len().saturating_sub(self.min_bytes_to_buffer);
164 // Ensure we split at a utf-8 char boundary.
152 - let cutoff = self.cumulative_output.floor_char_boundary(cutoff);
153 - (cutoff > self.output_index).then(|| {
154 - let chunk = self.cumulative_output[self.output_index..cutoff].to_string();
155 - self.output_index = cutoff;
156 - chunk
157 - })
165 + let cutoff = self.decoded.text().floor_char_boundary(cutoff);
166 + self.decoded.take_ready(cutoff)
167 }
168
160 - fn flush(&mut self, truncate_output_to: Option<usize>) -> Result<(Option<String>, String)> {
169 + fn flush(
170 + &mut self,
171 + truncate_output_to: Option<usize>,
172 + ) -> Result<(Option<DecodedText>, DecodedText)> {
173 // If the prefix was never seeded (no push_token was called), `ids`
174 // holds only prompt context — decoding it would re-emit prompt text.
175 if self.prefix_seeded && !self.ids.is_empty() {
176 let string = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
177 let prefix_len = self.prefix.len();
178 // Ensure we split at a utf-8 char boundary.
167 - self.cumulative_output
168 - .push_str(&string[string.floor_char_boundary(prefix_len)..]);
179 + let new_chunk = &string[string.floor_char_boundary(prefix_len)..];
180 + if !new_chunk.is_empty() {
181 + self.decoded.append_visible_text(new_chunk);
182 + }
183 }
184 + self.decoded.resolve_pending_zero_width();
185 self.ids.clear();
186 self.prefix.clear();
187 self.prefix_index = 0;
188 self.prefix_seeded = true;
189 if let Some(truncate_output_to) = truncate_output_to {
175 - self.cumulative_output.truncate(truncate_output_to);
190 + self.decoded.truncate(truncate_output_to);
191 }
177 - let last_chunk = (self.output_index < self.cumulative_output.len())
178 - .then(|| self.cumulative_output[self.output_index..].to_string());
179 - self.output_index = 0;
180 - Ok((last_chunk, take(&mut self.cumulative_output)))
192 + Ok(self.decoded.finish())
193 }
194
195 fn output(&self) -> &str {
184 - &self.cumulative_output
196 + self.decoded.text()
197 }
198 }
199
@@ -245,13 +257,13 @@ mod tests {
257 let mut decoder = backend.create_decode_stream(&[], false, 0);
258
259 assert_eq!(decoder.push_token(b'o' as u32).unwrap(), 1);
248 - assert_eq!(decoder.next_chunk().as_deref(), Some("o"));
260 + assert_eq!(decoder.next_chunk().unwrap().text, "o");
261 assert_eq!(decoder.push_token(b'k' as u32).unwrap(), 1);
250 - assert_eq!(decoder.next_chunk().as_deref(), Some("k"));
262 + assert_eq!(decoder.next_chunk().unwrap().text, "k");
263 // All text already consumed via next_chunk
264 let (last_chunk, full_text) = decoder.flush(None).unwrap();
265 assert_eq!(last_chunk, None);
254 - assert_eq!(full_text, "ok");
266 + assert_eq!(full_text.text, "ok");
267 }
268
269 #[test]
@@ -288,6 +300,7 @@ mod tests {
300 0 if !skip_special_tokens => text.push_str("<special>"),
301 0 => {}
302 1 => text.push('a'),
303 + 2 => text.push('b'),
304 _ => {}
305 }
306 }
@@ -301,6 +314,10 @@ mod tests {
314 fn id_to_token(&self, _id: u32) -> Option<String> {
315 unreachable!()
316 }
317 +
318 + fn is_special_id(&self, token_id: u32) -> bool {
319 + token_id == 0
320 + }
321 }
322
323 #[test]
@@ -393,17 +410,17 @@ mod tests {
410 let mut decoder = backend.create_decode_stream(&[], false, 0);
411
412 let input = b"Hello, world!";
396 - let mut full = String::new();
413 + let mut out = String::new();
414 for &byte in input {
415 decoder.push_token(byte as u32).unwrap();
416 if let Some(chunk) = decoder.next_chunk() {
400 - full.push_str(&chunk);
417 + out.push_str(&chunk.text);
418 }
419 }
403 - let (last_chunk, full_text) = decoder.flush(None).unwrap();
420 + let (last_chunk, full) = decoder.flush(None).unwrap();
421 assert_eq!(last_chunk, None); // all consumed via next_chunk
405 - assert_eq!(full, "Hello, world!");
406 - assert_eq!(full_text, "Hello, world!");
422 + assert_eq!(out, "Hello, world!");
423 + assert_eq!(full.text, "Hello, world!");
424 }
425
426 /// Backend simulating non-monotonic decode where adding a token changes how
@@ -470,15 +487,15 @@ mod tests {
487 for &byte in input {
488 decoder.push_token(byte as u32).unwrap();
489 if let Some(chunk) = decoder.next_chunk() {
473 - chunks.push_str(&chunk);
490 + chunks.push_str(&chunk.text);
491 }
492 }
493 // With hold_back_bytes=3, last 3 bytes ("lo!") are held back
494 assert_eq!(chunks, "Hel");
495 // Flush returns the rest
479 - let (last_chunk, full_text) = decoder.flush(None).unwrap();
480 - assert_eq!(last_chunk.as_deref(), Some("lo!"));
481 - assert_eq!(full_text, "Hello!");
496 + let (last_chunk, full) = decoder.flush(None).unwrap();
497 + assert_eq!(last_chunk.unwrap().text, "lo!");
498 + assert_eq!(full.text, "Hello!");
499 }
500
501 #[test]
@@ -493,14 +510,14 @@ mod tests {
510 for byte in "你好A".bytes() {
511 decoder.push_token(u32::from(byte)).unwrap();
512 if let Some(chunk) = decoder.next_chunk() {
496 - out.push_str(&chunk);
513 + out.push_str(&chunk.text);
514 }
515 }
499 - let (last_chunk, full_text) = decoder.flush(None).unwrap();
516 + let (last_chunk, full) = decoder.flush(None).unwrap();
517 if let Some(chunk) = last_chunk {
501 - out.push_str(&chunk);
518 + out.push_str(&chunk.text);
519 }
503 - assert_eq!(full_text, "你好A");
520 + assert_eq!(full.text, "你好A");
521 assert_eq!(out, "你好A");
522 }
523
@@ -515,9 +532,9 @@ mod tests {
532 .collect();
533 let mut decoder = backend.create_decode_stream(&prompt, false, 0);
534
518 - let (last_chunk, full_text) = decoder.flush(None).unwrap();
535 + let (last_chunk, full) = decoder.flush(None).unwrap();
536 assert_eq!(last_chunk, None);
520 - assert_eq!(full_text, "");
537 + assert_eq!(full.text, "");
538 }
539
540 #[test]
@@ -526,8 +543,8 @@ mod tests {
543 let prompt = vec![0xe4, 0xbd];
544 let mut decoder = backend.create_decode_stream(&prompt, false, 0);
545
529 - let (last_chunk, full_text) = decoder.flush(None).unwrap();
546 + let (last_chunk, full) = decoder.flush(None).unwrap();
547 assert_eq!(last_chunk, None);
531 - assert_eq!(full_text, "");
548 + assert_eq!(full.text, "");
549 }
550 }
rust/src/tokenizer/src/incremental/attribution.rs new
+261
@@ -0,0 +1,261 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::mem::take;
5 +
6 +use smallvec::SmallVec;
7 +
8 +/// Position of one generated token in decoded text.
9 +#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10 +pub enum TokenAnchor {
11 + /// The token contributed to visible decoded text.
12 + Visible { byte_offset: u32 },
13 + /// The token produced no visible decoded bytes.
14 + ZeroWidth { byte_offset: u32 },
15 +}
16 +
17 +impl TokenAnchor {
18 + fn byte_offset(self) -> u32 {
19 + match self {
20 + Self::Visible { byte_offset } | Self::ZeroWidth { byte_offset } => byte_offset,
21 + }
22 + }
23 +
24 + /// Return a new anchor offset by `delta` bytes.
25 + #[must_use]
26 + pub(super) fn offset_by(self, delta: i64) -> Self {
27 + let byte_offset = i64::from(self.byte_offset())
28 + .checked_add(delta)
29 + .and_then(|offset| u32::try_from(offset).ok())
30 + .expect("token anchor byte offset out of range");
31 + match self {
32 + Self::Visible { .. } => Self::Visible { byte_offset },
33 + Self::ZeroWidth { .. } => Self::ZeroWidth { byte_offset },
34 + }
35 + }
36 +}
37 +
38 +/// One generated token and its position in decoded text.
39 +#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40 +pub struct TokenAttribution {
41 + /// Token ID passed to [`super::IncrementalDecoder::push_token`].
42 + pub token_id: u32,
43 + /// Position of the token in the decoded text.
44 + pub anchor: TokenAnchor,
45 +}
46 +
47 +/// Decoded text and the generated tokens attributed to it.
48 +#[derive(Clone, Debug, Default, PartialEq, Eq)]
49 +pub struct DecodedText {
50 + /// Decoded UTF-8 text.
51 + pub text: String,
52 + /// One entry per generated token, in generation order.
53 + ///
54 + /// Offsets are local to `text`. Visible offsets may repeat when multiple
55 + /// tokens jointly decode into one character. A zero-width offset may equal
56 + /// `text.len()`.
57 + /// Four inline records cover one complete UTF-8 byte-fallback sequence.
58 + pub attributions: SmallVec<[TokenAttribution; 4]>,
59 +}
60 +
61 +impl DecodedText {
62 + /// Create decoded text without token attribution records. Mostly for testing purposes.
63 + pub fn unattributed(text: impl Into<String>) -> Self {
64 + Self {
65 + text: text.into(),
66 + attributions: SmallVec::new(),
67 + }
68 + }
69 +
70 + /// Append another decoded fragment, rebasing its token anchors.
71 + pub fn append(&mut self, other: Self) {
72 + if self.text.is_empty() && self.attributions.is_empty() {
73 + *self = other;
74 + return;
75 + }
76 +
77 + let byte_offset = offset_as_u32(self.text.len());
78 + let other_len = offset_as_u32(other.text.len());
79 + let _combined_len = byte_offset.checked_add(other_len).expect("decoded text exceeds 4 GiB");
80 +
81 + self.text.push_str(&other.text);
82 + self.attributions.extend(other.attributions.into_iter().map(|attribution| {
83 + TokenAttribution {
84 + token_id: attribution.token_id,
85 + anchor: attribution.anchor.offset_by(i64::from(byte_offset)),
86 + }
87 + }));
88 + }
89 +}
90 +
91 +#[derive(Clone, Copy, Debug)]
92 +enum PendingAnchor {
93 + Unresolved,
94 + Resolved(TokenAnchor),
95 +}
96 +
97 +#[derive(Clone, Copy, Debug)]
98 +struct PendingAttribution {
99 + token_id: u32,
100 + anchor: PendingAnchor,
101 +}
102 +
103 +/// Cumulative decoded text plus the attribution state needed to emit local
104 +/// chunks.
105 +#[derive(Default)]
106 +pub(super) struct AttributedTextBuffer {
107 + text: String,
108 + emitted_byte_offset: usize,
109 + attributions: Vec<PendingAttribution>,
110 + emitted_attribution_count: usize,
111 + pending_start: Option<usize>,
112 +}
113 +
114 +impl AttributedTextBuffer {
115 + /// Return all decoded text accumulated so far.
116 + pub(super) fn text(&self) -> &str {
117 + &self.text
118 + }
119 +
120 + /// Return the decoded text length in bytes.
121 + pub(super) fn len(&self) -> usize {
122 + self.text.len()
123 + }
124 +
125 + /// Record a token awaiting visible decoded output.
126 + pub(super) fn record_pending_token(&mut self, token_id: u32) {
127 + self.pending_start.get_or_insert(self.attributions.len());
128 + self.attributions.push(PendingAttribution {
129 + token_id,
130 + anchor: PendingAnchor::Unresolved,
131 + });
132 + }
133 +
134 + /// Record a token that produced no visible decoded output.
135 + pub(super) fn record_zero_width_token(&mut self, token_id: u32) {
136 + self.attributions.push(PendingAttribution {
137 + token_id,
138 + anchor: PendingAnchor::Resolved(TokenAnchor::ZeroWidth {
139 + byte_offset: offset_as_u32(self.text.len()),
140 + }),
141 + });
142 + }
143 +
144 + /// Append visible text and anchor pending tokens at its start.
145 + pub(super) fn append_visible_text(&mut self, text: &str) {
146 + let byte_offset = offset_as_u32(self.text.len());
147 + if let Some(pending_start) = self.pending_start.take() {
148 + for attribution in &mut self.attributions[pending_start..] {
149 + if matches!(attribution.anchor, PendingAnchor::Unresolved) {
150 + attribution.anchor =
151 + PendingAnchor::Resolved(TokenAnchor::Visible { byte_offset });
152 + }
153 + }
154 + }
155 + self.text.push_str(text);
156 + }
157 +
158 + /// Resolve remaining pending tokens at the current text boundary.
159 + pub(super) fn resolve_pending_zero_width(&mut self) {
160 + let Some(pending_start) = self.pending_start.take() else {
161 + return;
162 + };
163 + let byte_offset = offset_as_u32(self.text.len());
164 + for attribution in &mut self.attributions[pending_start..] {
165 + if matches!(attribution.anchor, PendingAnchor::Unresolved) {
166 + attribution.anchor =
167 + PendingAnchor::Resolved(TokenAnchor::ZeroWidth { byte_offset });
168 + }
169 + }
170 + }
171 +
172 + /// Truncate text and collapse removed anchors to the new boundary.
173 + pub(super) fn truncate(&mut self, byte_offset: usize) {
174 + self.text.truncate(byte_offset);
175 + let byte_offset = offset_as_u32(byte_offset);
176 + for attribution in &mut self.attributions {
177 + let PendingAnchor::Resolved(anchor) = &mut attribution.anchor else {
178 + continue;
179 + };
180 + if anchor.byte_offset() >= byte_offset {
181 + *anchor = TokenAnchor::ZeroWidth { byte_offset };
182 + }
183 + }
184 + }
185 +
186 + /// Emit decoded text and attributions ready through `cutoff`.
187 + pub(super) fn take_ready(&mut self, cutoff: usize) -> Option<DecodedText> {
188 + let chunk_start = self.emitted_byte_offset;
189 + let cutoff_u32 = offset_as_u32(cutoff);
190 + let mut attribution_end = self.emitted_attribution_count;
191 +
192 + for attribution in &self.attributions[self.emitted_attribution_count..] {
193 + let PendingAnchor::Resolved(anchor) = attribution.anchor else {
194 + break;
195 + };
196 + let ready = match anchor {
197 + TokenAnchor::Visible { byte_offset } => byte_offset < cutoff_u32,
198 + TokenAnchor::ZeroWidth { byte_offset } => byte_offset <= cutoff_u32,
199 + };
200 + if !ready {
201 + break;
202 + }
203 + attribution_end += 1;
204 + }
205 +
206 + if cutoff == chunk_start && attribution_end == self.emitted_attribution_count {
207 + return None;
208 + }
209 +
210 + let chunk_start_u32 = offset_as_u32(chunk_start);
211 + let attributions = self.attributions[self.emitted_attribution_count..attribution_end]
212 + .iter()
213 + .map(|attribution| {
214 + let PendingAnchor::Resolved(anchor) = attribution.anchor else {
215 + unreachable!("ready attributions must be resolved")
216 + };
217 + TokenAttribution {
218 + token_id: attribution.token_id,
219 + anchor: anchor.offset_by(-i64::from(chunk_start_u32)),
220 + }
221 + })
222 + .collect();
223 + let text = self.text[chunk_start..cutoff].to_string();
224 +
225 + self.emitted_byte_offset = cutoff;
226 + self.emitted_attribution_count = attribution_end;
227 + Some(DecodedText { text, attributions })
228 + }
229 +
230 + /// Emit the final chunk and take the complete decoded output.
231 + pub(super) fn finish(&mut self) -> (Option<DecodedText>, DecodedText) {
232 + let last_chunk = self.take_ready(self.text.len());
233 + let attributions = take(&mut self.attributions)
234 + .into_iter()
235 + .map(|attribution| {
236 + let PendingAnchor::Resolved(anchor) = attribution.anchor else {
237 + unreachable!("finish requires every attribution to be resolved")
238 + };
239 + TokenAttribution {
240 + token_id: attribution.token_id,
241 + anchor,
242 + }
243 + })
244 + .collect();
245 + let full = DecodedText {
246 + text: take(&mut self.text),
247 + attributions,
248 + };
249 + self.emitted_byte_offset = 0;
250 + self.emitted_attribution_count = 0;
251 + self.pending_start = None;
252 + (last_chunk, full)
253 + }
254 +}
255 +
256 +fn offset_as_u32(offset: usize) -> u32 {
257 + u32::try_from(offset).expect("decoded text exceeds 4 GiB")
258 +}
259 +
260 +#[cfg(test)]
261 +mod tests;
rust/src/tokenizer/src/incremental/attribution/tests.rs new
+425
@@ -0,0 +1,425 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use super::*;
5 +use crate::{Result, Tokenizer, test_utils::TestTokenizer};
6 +
7 +const SPECIAL_TOKEN_ID: u32 = 0x100;
8 +
9 +/// Backend with tokens that decode to whole text pieces.
10 +#[derive(Debug)]
11 +struct PieceBackend;
12 +
13 +impl Tokenizer for PieceBackend {
14 + fn encode(&self, _text: &str, _add_special_tokens: bool) -> Result<Vec<u32>> {
15 + unreachable!()
16 + }
17 +
18 + fn encode_ordinary(&self, _text: &str) -> Result<Vec<u32>> {
19 + unreachable!()
20 + }
21 +
22 + fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result<String> {
23 + Ok(token_ids
24 + .iter()
25 + .map(|token_id| match token_id {
26 + 1 => "abcd",
27 + 2 => "ab",
28 + 3 => "<stop>",
29 + 4 => "",
30 + _ => unreachable!("unexpected token id: {token_id}"),
31 + })
32 + .collect())
33 + }
34 +
35 + fn token_to_id(&self, _token: &str) -> Option<u32> {
36 + unreachable!()
37 + }
38 +
39 + fn id_to_token(&self, id: u32) -> Option<String> {
40 + Some(id.to_string())
41 + }
42 +}
43 +
44 +/// Raw-byte backend with one filterable special token.
45 +#[derive(Debug)]
46 +struct Utf8SpecialBackend;
47 +
48 +impl Tokenizer for Utf8SpecialBackend {
49 + fn encode(&self, _text: &str, _add_special_tokens: bool) -> Result<Vec<u32>> {
50 + unreachable!()
51 + }
52 +
53 + fn encode_ordinary(&self, _text: &str) -> Result<Vec<u32>> {
54 + unreachable!()
55 + }
56 +
57 + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
58 + let mut bytes = Vec::new();
59 + for &token_id in token_ids {
60 + if token_id == SPECIAL_TOKEN_ID {
61 + if !skip_special_tokens {
62 + bytes.extend_from_slice(b"<special>");
63 + }
64 + } else {
65 + bytes.push(token_id as u8);
66 + }
67 + }
68 + Ok(String::from_utf8_lossy(&bytes).into_owned())
69 + }
70 +
71 + fn token_to_id(&self, _token: &str) -> Option<u32> {
72 + unreachable!()
73 + }
74 +
75 + fn id_to_token(&self, id: u32) -> Option<String> {
76 + Some(id.to_string())
77 + }
78 +
79 + fn is_special_id(&self, token_id: u32) -> bool {
80 + token_id == SPECIAL_TOKEN_ID
81 + }
82 +}
83 +
84 +fn tokenizer_with_special_token() -> TestTokenizer {
85 + TestTokenizer::new().with_special_token("<special>", SPECIAL_TOKEN_ID)
86 +}
87 +
88 +fn visible(byte_offset: u32) -> TokenAnchor {
89 + TokenAnchor::Visible { byte_offset }
90 +}
91 +
92 +fn zero_width(byte_offset: u32) -> TokenAnchor {
93 + TokenAnchor::ZeroWidth { byte_offset }
94 +}
95 +
96 +fn decoded(text: &str, token_ids: &[u32], anchors: &[TokenAnchor]) -> DecodedText {
97 + assert_eq!(token_ids.len(), anchors.len());
98 + DecodedText {
99 + text: text.to_string(),
100 + attributions: token_ids
101 + .iter()
102 + .copied()
103 + .zip(anchors.iter().copied())
104 + .map(|(token_id, anchor)| TokenAttribution { token_id, anchor })
105 + .collect(),
106 + }
107 +}
108 +
109 +#[test]
110 +fn token_anchor_offset_by_preserves_anchor_kind() {
111 + for (anchor, delta, expected) in [
112 + (visible(3), -2, visible(1)),
113 + (zero_width(1), 4, zero_width(5)),
114 + ] {
115 + assert_eq!(anchor.offset_by(delta), expected);
116 + }
117 +}
118 +
119 +#[test]
120 +fn decoded_text_append_rebases_attributions() {
121 + let mut combined = decoded("a", &[1], &[visible(0)]);
122 + combined.append(decoded(
123 + "你",
124 + &[2, 3, 4, 5],
125 + &[visible(0), visible(0), visible(0), zero_width(3)],
126 + ));
127 +
128 + assert_eq!(
129 + combined,
130 + decoded(
131 + "a你",
132 + &[1, 2, 3, 4, 5],
133 + &[
134 + visible(0),
135 + visible(1),
136 + visible(1),
137 + visible(1),
138 + zero_width(4),
139 + ],
140 + )
141 + );
142 +}
143 +
144 +#[test]
145 +fn decoded_text_append_moves_allocations_into_empty_receiver() {
146 + let source = decoded(
147 + "hello",
148 + &[1, 2, 3, 4, 5],
149 + &[visible(0), visible(1), visible(2), visible(3), visible(4)],
150 + );
151 + let text_ptr = source.text.as_ptr();
152 + let attributions_ptr = source.attributions.as_ptr();
153 + let mut combined = DecodedText::default();
154 +
155 + combined.append(source);
156 +
157 + assert_eq!(combined.text.as_ptr(), text_ptr);
158 + assert_eq!(combined.attributions.as_ptr(), attributions_ptr);
159 +}
160 +
161 +#[test]
162 +fn decoded_text_append_preserves_zero_width_only_receiver() {
163 + let mut combined = decoded("", &[1], &[zero_width(0)]);
164 + combined.append(decoded("a", &[2], &[visible(0)]));
165 +
166 + assert_eq!(
167 + combined,
168 + decoded("a", &[1, 2], &[zero_width(0), visible(0)])
169 + );
170 +}
171 +
172 +struct AttributionCase<'a> {
173 + tokenizer: &'a dyn Tokenizer,
174 + prompt_token_ids: Vec<u32>,
175 + token_ids: Vec<u32>,
176 + skip_special_tokens: bool,
177 + min_bytes_to_buffer: usize,
178 + truncate_output_to: Option<usize>,
179 + expected_chunks: Vec<DecodedText>,
180 + expected_full: DecodedText,
181 +}
182 +
183 +fn run_attribution_case(case: AttributionCase<'_>) {
184 + let mut decoder = case.tokenizer.create_decode_stream(
185 + &case.prompt_token_ids,
186 + case.skip_special_tokens,
187 + case.min_bytes_to_buffer,
188 + );
189 + let mut chunks = Vec::new();
190 +
191 + for token_id in &case.token_ids {
192 + decoder.push_token(*token_id).unwrap();
193 + while let Some(chunk) = decoder.next_chunk() {
194 + chunks.push(chunk);
195 + }
196 + }
197 + let (remaining, full) = decoder.flush(case.truncate_output_to).unwrap();
198 + if let Some(remaining) = remaining {
199 + chunks.push(remaining);
200 + }
201 +
202 + assert_eq!(chunks, case.expected_chunks, "chunks");
203 + assert_eq!(full, case.expected_full, "full output");
204 + assert_eq!(
205 + full.attributions.len(),
206 + case.token_ids.len(),
207 + "one attribution per pushed token"
208 + );
209 + assert_eq!(
210 + full.attributions
211 + .iter()
212 + .map(|attribution| attribution.token_id)
213 + .collect::<Vec<_>>(),
214 + case.token_ids,
215 + "attribution preserves pushed token IDs"
216 + );
217 +
218 + let mut reconstructed_text = String::new();
219 + let mut reconstructed_attributions = Vec::new();
220 + for chunk in &chunks {
221 + let chunk_start = reconstructed_text.len() as u32;
222 + reconstructed_text.push_str(&chunk.text);
223 + reconstructed_attributions.extend(chunk.attributions.iter().map(|attribution| {
224 + TokenAttribution {
225 + token_id: attribution.token_id,
226 + anchor: attribution.anchor.offset_by(i64::from(chunk_start)),
227 + }
228 + }));
229 + }
230 + assert_eq!(
231 + reconstructed_text, full.text,
232 + "chunk text reconstructs full text"
233 + );
234 + assert_eq!(
235 + reconstructed_attributions.as_slice(),
236 + full.attributions.as_slice(),
237 + "chunk attributions reconstruct full attributions"
238 + );
239 +}
240 +
241 +#[test]
242 +fn attribution_ascii_tokens_have_one_visible_anchor_each() {
243 + let tokenizer = TestTokenizer::new();
244 + run_attribution_case(AttributionCase {
245 + tokenizer: &tokenizer,
246 + prompt_token_ids: vec![],
247 + token_ids: vec![b'o' as u32, b'k' as u32],
248 + skip_special_tokens: false,
249 + min_bytes_to_buffer: 0,
250 + truncate_output_to: None,
251 + expected_chunks: vec![
252 + decoded("o", &[b'o' as u32], &[visible(0)]),
253 + decoded("k", &[b'k' as u32], &[visible(0)]),
254 + ],
255 + expected_full: decoded("ok", &[b'o' as u32, b'k' as u32], &[visible(0), visible(1)]),
256 + });
257 +}
258 +
259 +#[test]
260 +fn attribution_byte_fallback_tokens_share_one_visible_anchor() {
261 + let tokenizer = TestTokenizer::new();
262 + run_attribution_case(AttributionCase {
263 + tokenizer: &tokenizer,
264 + prompt_token_ids: vec![],
265 + token_ids: vec![0xe4, 0xbd, 0xa0],
266 + skip_special_tokens: false,
267 + min_bytes_to_buffer: 0,
268 + truncate_output_to: None,
269 + expected_chunks: vec![decoded(
270 + "你",
271 + &[0xe4, 0xbd, 0xa0],
272 + &[visible(0), visible(0), visible(0)],
273 + )],
274 + expected_full: decoded(
275 + "你",
276 + &[0xe4, 0xbd, 0xa0],
277 + &[visible(0), visible(0), visible(0)],
278 + ),
279 + });
280 +}
281 +
282 +#[test]
283 +fn attribution_holdback_splits_text_without_repeating_token() {
284 + run_attribution_case(AttributionCase {
285 + tokenizer: &PieceBackend,
286 + prompt_token_ids: vec![],
287 + token_ids: vec![1],
288 + skip_special_tokens: false,
289 + min_bytes_to_buffer: 2,
290 + truncate_output_to: None,
291 + expected_chunks: vec![decoded("ab", &[1], &[visible(0)]), decoded("cd", &[], &[])],
292 + expected_full: decoded("abcd", &[1], &[visible(0)]),
293 + });
294 +}
295 +
296 +#[test]
297 +fn attribution_filtered_special_token_is_zero_width_at_its_byte_boundary() {
298 + let tokenizer = tokenizer_with_special_token();
299 + run_attribution_case(AttributionCase {
300 + tokenizer: &tokenizer,
301 + prompt_token_ids: vec![],
302 + token_ids: vec![b'a' as u32, SPECIAL_TOKEN_ID, b'b' as u32],
303 + skip_special_tokens: true,
304 + min_bytes_to_buffer: 0,
305 + truncate_output_to: None,
306 + expected_chunks: vec![
307 + decoded("a", &[b'a' as u32], &[visible(0)]),
308 + decoded("", &[SPECIAL_TOKEN_ID], &[zero_width(0)]),
309 + decoded("b", &[b'b' as u32], &[visible(0)]),
310 + ],
311 + expected_full: decoded(
312 + "ab",
313 + &[b'a' as u32, SPECIAL_TOKEN_ID, b'b' as u32],
314 + &[visible(0), zero_width(1), visible(1)],
315 + ),
316 + });
317 +}
318 +
319 +#[test]
320 +fn attribution_zero_width_preserves_order_inside_byte_fallback_group() {
321 + run_attribution_case(AttributionCase {
322 + tokenizer: &Utf8SpecialBackend,
323 + prompt_token_ids: vec![],
324 + token_ids: vec![0xe4, SPECIAL_TOKEN_ID, 0xbd, 0xa0],
325 + skip_special_tokens: true,
326 + min_bytes_to_buffer: 0,
327 + truncate_output_to: None,
328 + expected_chunks: vec![decoded(
329 + "你",
330 + &[0xe4, SPECIAL_TOKEN_ID, 0xbd, 0xa0],
331 + &[visible(0), zero_width(0), visible(0), visible(0)],
332 + )],
333 + expected_full: decoded(
334 + "你",
335 + &[0xe4, SPECIAL_TOKEN_ID, 0xbd, 0xa0],
336 + &[visible(0), zero_width(0), visible(0), visible(0)],
337 + ),
338 + });
339 +}
340 +
341 +#[test]
342 +fn attribution_retained_special_token_has_visible_anchor() {
343 + run_attribution_case(AttributionCase {
344 + tokenizer: &Utf8SpecialBackend,
345 + prompt_token_ids: vec![],
346 + token_ids: vec![SPECIAL_TOKEN_ID],
347 + skip_special_tokens: false,
348 + min_bytes_to_buffer: 0,
349 + truncate_output_to: None,
350 + expected_chunks: vec![decoded("<special>", &[SPECIAL_TOKEN_ID], &[visible(0)])],
351 + expected_full: decoded("<special>", &[SPECIAL_TOKEN_ID], &[visible(0)]),
352 + });
353 +}
354 +
355 +#[test]
356 +fn attribution_empty_decode_resolves_to_zero_width_on_flush() {
357 + run_attribution_case(AttributionCase {
358 + tokenizer: &PieceBackend,
359 + prompt_token_ids: vec![],
360 + token_ids: vec![4],
361 + skip_special_tokens: false,
362 + min_bytes_to_buffer: 0,
363 + truncate_output_to: None,
364 + expected_chunks: vec![decoded("", &[4], &[zero_width(0)])],
365 + expected_full: decoded("", &[4], &[zero_width(0)]),
366 + });
367 +}
368 +
369 +#[test]
370 +fn attribution_incomplete_utf8_tokens_share_replacement_anchor_on_flush() {
371 + let tokenizer = TestTokenizer::new();
372 + run_attribution_case(AttributionCase {
373 + tokenizer: &tokenizer,
374 + prompt_token_ids: vec![],
375 + token_ids: vec![0xe4, 0xbd],
376 + skip_special_tokens: false,
377 + min_bytes_to_buffer: 0,
378 + truncate_output_to: None,
379 + expected_chunks: vec![decoded("�", &[0xe4, 0xbd], &[visible(0), visible(0)])],
380 + expected_full: decoded("�", &[0xe4, 0xbd], &[visible(0), visible(0)]),
381 + });
382 +}
383 +
384 +#[test]
385 +fn attribution_truncation_converts_fully_removed_token_to_zero_width() {
386 + run_attribution_case(AttributionCase {
387 + tokenizer: &PieceBackend,
388 + prompt_token_ids: vec![],
389 + token_ids: vec![2, 3],
390 + skip_special_tokens: false,
391 + min_bytes_to_buffer: 32,
392 + truncate_output_to: Some(2),
393 + expected_chunks: vec![decoded("ab", &[2, 3], &[visible(0), zero_width(2)])],
394 + expected_full: decoded("ab", &[2, 3], &[visible(0), zero_width(2)]),
395 + });
396 +}
397 +
398 +#[test]
399 +fn attribution_truncation_inside_token_retains_first_byte_anchor() {
400 + run_attribution_case(AttributionCase {
401 + tokenizer: &PieceBackend,
402 + prompt_token_ids: vec![],
403 + token_ids: vec![1],
404 + skip_special_tokens: false,
405 + min_bytes_to_buffer: 32,
406 + truncate_output_to: Some(2),
407 + expected_chunks: vec![decoded("ab", &[1], &[visible(0)])],
408 + expected_full: decoded("ab", &[1], &[visible(0)]),
409 + });
410 +}
411 +
412 +#[test]
413 +fn attribution_excludes_prompt_context() {
414 + let tokenizer = TestTokenizer::new();
415 + run_attribution_case(AttributionCase {
416 + tokenizer: &tokenizer,
417 + prompt_token_ids: vec![b'H' as u32, b'i' as u32],
418 + token_ids: vec![b'!' as u32],
419 + skip_special_tokens: false,
420 + min_bytes_to_buffer: 0,
421 + truncate_output_to: None,
422 + expected_chunks: vec![decoded("!", &[b'!' as u32], &[visible(0)])],
423 + expected_full: decoded("!", &[b'!' as u32], &[visible(0)]),
424 + });
425 +}
rust/src/tokenizer/src/lib.rs
+1 -1
@@ -17,7 +17,7 @@ mod tiktoken;
17
18 pub use error::{Result, TokenizerError};
19 pub use hf::HuggingFaceTokenizer;
20 -pub use incremental::IncrementalDecoder;
20 +pub use incremental::{DecodedText, IncrementalDecoder, TokenAnchor, TokenAttribution};
21 pub use tekken::TekkenTokenizer;
22 pub use tiktoken::TiktokenTokenizer;
23
rust/src/tokenizer/src/tiktoken.rs
+8 -8
@@ -964,16 +964,16 @@ mod tests {
964 for &id in &ids {
965 decoder.push_token(id).unwrap();
966 if let Some(chunk) = decoder.next_chunk() {
967 - output.push_str(&chunk);
967 + output.push_str(&chunk.text);
968 }
969 }
970 - let (last_chunk, full_text) = decoder.flush(None).unwrap();
970 + let (last_chunk, full) = decoder.flush(None).unwrap();
971 if let Some(chunk) = last_chunk {
972 - output.push_str(&chunk);
972 + output.push_str(&chunk.text);
973 }
974
975 assert_eq!(output, text);
976 - assert_eq!(full_text, text);
976 + assert_eq!(full.text, text);
977 }
978 }
979
@@ -991,16 +991,16 @@ mod tests {
991 for &id in &ids {
992 decoder.push_token(id).unwrap();
993 if let Some(chunk) = decoder.next_chunk() {
994 - output.push_str(&chunk);
994 + output.push_str(&chunk.text);
995 }
996 }
997 - let (last_chunk, full_text) = decoder.flush(None).unwrap();
997 + let (last_chunk, full) = decoder.flush(None).unwrap();
998 if let Some(chunk) = last_chunk {
999 - output.push_str(&chunk);
999 + output.push_str(&chunk.text);
1000 }
1001
1002 assert_eq!(output, text);
1003 - assert_eq!(full_text, text);
1003 + assert_eq!(full.text, text);
1004 }
1005 }
1006