@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 21a2211040

[Rust Frontend] Use token-attributed text in reasoning and unified parsers (#54884)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Sep 4, 2026 at 04:05 UTC 21a22110404d98dcab5f51cd90e9acb0acc15e47

22 files changed +1012 -350

rust/src/chat/src/output/default/unified.rs

+13 -10

@@ -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, SampledDelta};

16 +use vllm_text::output::{DecodedText, DecodedTextEvent, SampledDelta};

18 use crate::Result;

19 use crate::error::Error;

@@ -63,14 +63,17 @@ impl UnifiedParserState {

63 }

65 /// Convert one decoded text delta into zero or more parsed assistant events.

66 - fn process_delta(&mut self, delta: String) -> Result<Vec<AssistantEvent>> {

66 + fn process_delta(&mut self, delta: DecodedText) -> Result<Vec<AssistantEvent>> {

67 if self.parser_failed {

68 self.open_call_index = None;

69 - return Ok(text_event(AssistantBlockKind::Text, delta).into_iter().collect());

69 + return Ok(text_event(AssistantBlockKind::Text, delta.text).into_iter().collect());

70 }

72 let mut output = UnifiedParserOutput::default();

73 - match self.parser.parse_into(&delta, &mut output) {

73 + // The parser consumes the delta; keep its text for the fallback path

74 + // that re-emits it as plain text.

75 + let fallback_text = delta.text.clone();

76 + match self.parser.parse_into(delta, &mut output) {

77 Ok(()) => {

78 let mut events = Vec::new();

79 self.process_parser_output(output, &mut events)?;

@@ -89,7 +92,7 @@ impl UnifiedParserState {

93 let recovered = self.parser.reset();

94 if recovered.is_empty() && events.is_empty() {

92 - push_text_delta(&mut events, AssistantBlockKind::Text, delta);

95 + push_text_delta(&mut events, AssistantBlockKind::Text, fallback_text);

96 } else {

97 push_text_delta(&mut events, AssistantBlockKind::Text, recovered);

98 }

@@ -137,7 +140,7 @@ impl UnifiedParserState {

140 }

141 UnifiedParserEvent::Reasoning(delta) => {

142 self.open_call_index = None;

140 - push_text_delta(events, AssistantBlockKind::Reasoning, delta);

143 + push_text_delta(events, AssistantBlockKind::Reasoning, delta.text);

144 }

145 UnifiedParserEvent::ToolCall(item) => {

146 self.process_tool_item(item, events)?;

@@ -245,7 +248,7 @@ pub(crate) async fn unified_event_stream(

248 },

249 finished,

250 } => {

248 - for next in state.process_delta(decoded.text)? {

251 + for next in state.process_delta(decoded)? {

252 y.yield_ok(next).await;

253 }

254 if logprobs.is_some() || !token_ids.is_empty() {

@@ -284,7 +287,7 @@ mod tests {

287 use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParserError, UnifiedParserOutput};

288 use vllm_tokenizer::test_utils::TestTokenizer;

289

287 - use super::unified_event_stream;

290 + use super::{DecodedText, unified_event_stream};

291 use crate::event::AssistantBlockKind;

292 use crate::output::AssistantEvent;

293

@@ -332,7 +335,7 @@ mod tests {

335

336 fn parse_into(

337 &mut self,

335 - _delta: &str,

338 + _delta: DecodedText,

339 output: &mut UnifiedParserOutput,

340 ) -> vllm_parser::unified::Result<()> {

341 match self.steps.pop_front().expect("unexpected parser call") {

@@ -418,7 +421,7 @@ mod tests {

421

422 fn reasoning(delta: &str) -> UnifiedParserOutput {

423 let mut output = UnifiedParserOutput::default();

421 - output.push_reasoning(delta.to_string());

424 + output.push_reasoning(DecodedText::unattributed(delta));

425 output

426 }

427

rust/src/parser/benches/utils/adapter.rs

+5 -2

@@ -9,7 +9,7 @@ use vllm_parser::tool::{

9 use vllm_parser::unified::{

10 UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,

11 };

12 -use vllm_tokenizer::Tokenizer;

12 +use vllm_tokenizer::{DecodedText, Tokenizer};

14 /// Tokenizer stub used by unified-parser benchmarks.

15 struct BenchTokenizer;

@@ -99,7 +99,10 @@ impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {

100 fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

101 let mut unified_output = UnifiedParserOutput::default();

102 - let result = self.inner.parse_into(chunk, &mut unified_output).map_err(map_unified_error);

102 + let result = self

103 + .inner

104 + .parse_into(DecodedText::unattributed(chunk), &mut unified_output)

105 + .map_err(map_unified_error);

106 append_unified_output(unified_output, output)?;

107 result

108 }

rust/src/parser/src/reasoning/cohere_cmd.rs

+2 -2

@@ -1,7 +1,7 @@

1 // SPDX-License-Identifier: Apache-2.0

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -use vllm_tokenizer::DynTokenizer;

4 +use vllm_tokenizer::{DecodedText, DynTokenizer};

6 use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};

@@ -38,7 +38,7 @@ impl ReasoningParser for CohereCmdReasoningParser {

38 Ok(())

39 }

41 - fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {

41 + fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {

42 Ok(self.inner.push(delta))

43 }

rust/src/parser/src/reasoning/deepseek_r1.rs

+2 -2

@@ -1,7 +1,7 @@

1 // SPDX-License-Identifier: Apache-2.0

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -use vllm_tokenizer::DynTokenizer;

4 +use vllm_tokenizer::{DecodedText, DynTokenizer};

6 use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};

@@ -37,7 +37,7 @@ impl ReasoningParser for DeepSeekR1ReasoningParser {

37 Ok(())

38 }

40 - fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {

40 + fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {

41 Ok(self.inner.push(delta))

42 }

rust/src/parser/src/reasoning/delimited.rs

+36 -22

@@ -1,7 +1,7 @@

1 // SPDX-License-Identifier: Apache-2.0

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -use vllm_tokenizer::{DynTokenizer, Tokenizer};

4 +use vllm_tokenizer::{DecodedText, DynTokenizer, Tokenizer};

6 use super::{ReasoningDelta, ReasoningError, Result};

@@ -19,7 +19,7 @@ use super::{ReasoningDelta, ReasoningError, Result};

19 pub(crate) struct DelimitedReasoningParser {

20 tokenizer: DynTokenizer,

21 current_in_reasoning: bool,

22 - buffer: String,

22 + buffer: DecodedText,

23 start_token: String,

24 end_token: String,

25 start_token_id: u32,

@@ -55,7 +55,7 @@ impl DelimitedReasoningParser {

55 Ok(Self {

56 tokenizer,

57 current_in_reasoning: default_in_reasoning,

58 - buffer: String::new(),

58 + buffer: DecodedText::default(),

59 start_token,

60 end_token,

61 start_token_id,

@@ -81,44 +81,58 @@ impl DelimitedReasoningParser {

81 }

83 /// Parse one decoded text delta and return its reasoning/content split.

84 - pub(crate) fn push(&mut self, delta: &str) -> ReasoningDelta {

85 - self.buffer.push_str(delta);

84 + pub(crate) fn push(&mut self, delta: DecodedText) -> ReasoningDelta {

85 + self.buffer.append(delta);

87 - let partial_suffix_len = self.partial_suffix_len(&self.buffer);

88 - let stable_len = self.buffer.len() - partial_suffix_len;

89 - let pending_suffix = self.buffer.split_off(stable_len);

90 - let stable_text = std::mem::replace(&mut self.buffer, pending_suffix);

87 + let partial_suffix_len = self.partial_suffix_len(&self.buffer.text);

88 + let stable_len = self.buffer.text.len() - partial_suffix_len;

89 + let stable = self.buffer.drain_prefix(stable_len);

92 - self.parse_stable_text(&stable_text)

91 + self.parse_stable_text(stable)

92 }

94 /// Flush any buffered partial delimiter suffix at end of stream.

95 pub(crate) fn finish(&mut self) -> ReasoningDelta {

97 - let stable_text = std::mem::take(&mut self.buffer);

98 - self.parse_stable_text(&stable_text)

96 + // `drain_prefix(text.len())` takes trailing zero-width tokens too.

97 + let stable = self.buffer.drain_prefix(self.buffer.text.len());

98 + self.parse_stable_text(stable)

99 }

100

101 /// Parse text that is known not to end with a partial delimiter suffix.

102 - fn parse_stable_text(&mut self, mut stable: &str) -> ReasoningDelta {

102 + ///

103 + /// Reasoning and content pieces keep the attributions of the tokens that

104 + /// produced them; delimiter marker spans are drained and dropped, keeping

105 + /// marker tokens out of any count.

106 + fn parse_stable_text(&mut self, mut stable: DecodedText) -> ReasoningDelta {

107 let mut delta = ReasoningDelta::default();

108

105 - while !stable.is_empty() {

109 + while !stable.text.is_empty() {

110 if self.current_in_reasoning {

107 - if let Some(end_idx) = stable.find(&self.end_token) {

108 - delta.push_reasoning(&stable[..end_idx]);

109 - stable = &stable[end_idx + self.end_token.len()..];

111 + if let Some(end_idx) = stable.text.find(&self.end_token) {

112 + delta.push_reasoning(stable.drain_prefix(end_idx));

113 + let _ = stable.drain_prefix(self.end_token.len());

114 self.current_in_reasoning = false;

115 } else {

116 delta.push_reasoning(stable);

113 - break;

117 + return delta;

118 }

115 - } else if let Some(start_idx) = stable.find(&self.start_token) {

116 - delta.push_content(&stable[..start_idx]);

117 - stable = &stable[start_idx + self.start_token.len()..];

119 + } else if let Some(start_idx) = stable.text.find(&self.start_token) {

120 + delta.push_content(stable.drain_prefix(start_idx));

121 + let _ = stable.drain_prefix(self.start_token.len());

122 self.current_in_reasoning = true;

123 } else {

124 delta.push_content(stable);

121 - break;

125 + return delta;

126 + }

127 + }

128 +

129 + // A remainder with empty text may still carry zero-width tokens;

130 + // attribute them to the current state.

131 + if !stable.attributions.is_empty() {

132 + if self.current_in_reasoning {

133 + delta.push_reasoning(stable);

134 + } else {

135 + delta.push_content(stable);

136 }

137 }

138

rust/src/parser/src/reasoning/hy.rs

+2 -2

@@ -1,7 +1,7 @@

1 // SPDX-License-Identifier: Apache-2.0

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -use vllm_tokenizer::DynTokenizer;

4 +use vllm_tokenizer::{DecodedText, DynTokenizer};

6 use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningError, ReasoningParser, Result};

@@ -41,7 +41,7 @@ impl ReasoningParser for HyReasoningParser {

41 Ok(())

42 }

44 - fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {

44 + fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {

45 Ok(self.inner.push(delta))

46 }

rust/src/parser/src/reasoning/kimi.rs

+2 -2

@@ -1,7 +1,7 @@

1 // SPDX-License-Identifier: Apache-2.0

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -use vllm_tokenizer::DynTokenizer;

4 +use vllm_tokenizer::{DecodedText, DynTokenizer};

6 use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};

@@ -32,7 +32,7 @@ impl ReasoningParser for KimiReasoningParser {

32 Ok(())

33 }

35 - fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {

35 + fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {

36 Ok(self.inner.push(delta))

37 }

rust/src/parser/src/reasoning/minimax_m3.rs

+17 -18

@@ -1,7 +1,7 @@

1 // SPDX-License-Identifier: Apache-2.0

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -use vllm_tokenizer::DynTokenizer;

4 +use vllm_tokenizer::{DecodedText, DynTokenizer};

6 use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};

@@ -20,7 +20,7 @@ pub struct MiniMaxM3ReasoningParser {

20 at_response_start: bool,

21 /// Holds an initial suffix like `</mm` while it may still complete into the

22 /// leading closer on a later chunk.

23 - leading_end_buffer: String,

23 + leading_end_buffer: DecodedText,

24 }

26 impl MiniMaxM3ReasoningParser {

@@ -29,31 +29,30 @@ impl MiniMaxM3ReasoningParser {

29 Ok(Self {

30 inner: DelimitedReasoningParser::new(tokenizer, M3_THINK_START, M3_THINK_END, false)?,

31 at_response_start: true,

32 - leading_end_buffer: String::new(),

32 + leading_end_buffer: DecodedText::default(),

33 })

34 }

36 /// Drop a response-leading `</mm:think>` while preserving later unmatched

37 /// closers as ordinary content.

38 - fn push_inner(&mut self, delta: &str) -> ReasoningDelta {

38 + fn push_inner(&mut self, delta: DecodedText) -> ReasoningDelta {

39 if self.at_response_start && !self.inner.in_reasoning() {

40 - self.leading_end_buffer.push_str(delta);

41 - let buffered = std::mem::take(&mut self.leading_end_buffer);

40 + self.leading_end_buffer.append(delta);

41 + let mut buffered = self.leading_end_buffer.take();

43 - if buffered.is_empty() {

44 - return ReasoningDelta::default();

45 - }

46 - if let Some(rest) = buffered.strip_prefix(M3_THINK_END) {

43 + if buffered.text.starts_with(M3_THINK_END) {

44 self.at_response_start = false;

48 - return self.inner.push(rest);

45 + // The dropped marker span keeps its tokens out of any count.

46 + let _ = buffered.drain_prefix(M3_THINK_END.len());

47 + return self.inner.push(buffered);

48 }

50 - if M3_THINK_END.starts_with(buffered.as_str()) {

49 + if M3_THINK_END.starts_with(buffered.text.as_str()) {

50 self.leading_end_buffer = buffered;

51 return ReasoningDelta::default();

52 }

54 self.at_response_start = false;

56 - return self.inner.push(&buffered);

55 + return self.inner.push(buffered);

56 }

58 self.inner.push(delta)

@@ -62,10 +61,10 @@ impl MiniMaxM3ReasoningParser {

62 fn append_delta(target: &mut ReasoningDelta, delta: ReasoningDelta) {

63 if let Some(reasoning) = delta.reasoning {

65 - target.push_reasoning(&reasoning);

64 + target.push_reasoning(reasoning);

65 }

66 if let Some(content) = delta.content {

68 - target.push_content(&content);

67 + target.push_content(content);

68 }

69 }

@@ -84,16 +83,16 @@ impl ReasoningParser for MiniMaxM3ReasoningParser {

83 Ok(())

84 }

87 - fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {

86 + fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {

87 Ok(self.push_inner(delta))

88 }

90 fn finish(&mut self) -> Result<ReasoningDelta> {

91 let mut delta = ReasoningDelta::default();

92 if !self.leading_end_buffer.is_empty() {

94 - let pending = std::mem::take(&mut self.leading_end_buffer);

93 + let pending = self.leading_end_buffer.take();

94 self.at_response_start = false;

96 - append_delta(&mut delta, self.inner.push(&pending));

95 + append_delta(&mut delta, self.inner.push(pending));

96 }

97 append_delta(&mut delta, self.inner.finish());

98 Ok(delta)

rust/src/parser/src/reasoning/mod.rs

+22 -12

@@ -28,7 +28,7 @@ mod seed_oss;

28 mod step3p5;

30 use thiserror::Error;

31 -use vllm_tokenizer::DynTokenizer;

31 +use vllm_tokenizer::{DecodedText, DynTokenizer};

33 pub use self::cohere_cmd::CohereCmdReasoningParser;

34 pub use self::deepseek_r1::DeepSeekR1ReasoningParser;

@@ -61,10 +61,14 @@ pub type Step3ReasoningParser = Qwen3ReasoningParser;

61 pub type Result<T> = std::result::Result<T, ReasoningError>;

63 /// One parsed streaming delta split into reasoning and visible content.

64 +///

65 +/// Each portion carries the attributions of the generated tokens that produced

66 +/// it, so downstream consumers can count reasoning tokens exactly. Marker

67 +/// spans are dropped by the parsers, keeping marker tokens out of any count.

68 #[derive(Debug, Default, Clone, PartialEq, Eq)]

69 pub struct ReasoningDelta {

66 - pub reasoning: Option<String>,

67 - pub content: Option<String>,

70 + pub reasoning: Option<DecodedText>,

71 + pub content: Option<DecodedText>,

72 }

74 impl ReasoningDelta {

@@ -74,24 +78,30 @@ impl ReasoningDelta {

78 }

80 /// Append text to the reasoning portion, creating it on first use.

77 - pub(crate) fn push_reasoning(&mut self, text: &str) {

78 - if text.is_empty() {

81 + ///

82 + /// A piece with empty text but non-empty attributions (zero-width tokens

83 + /// only) is still kept: the tokens are attributed to the current state.

84 + pub(crate) fn push_reasoning(&mut self, piece: DecodedText) {

85 + if piece.is_empty() {

86 return;

87 }

88 match &mut self.reasoning {

82 - Some(existing) => existing.push_str(text),

83 - None => self.reasoning = Some(text.to_string()),

89 + Some(existing) => existing.append(piece),

90 + None => self.reasoning = Some(piece),

91 }

92 }

94 /// Append text to the visible content portion, creating it on first use.

88 - pub(crate) fn push_content(&mut self, text: &str) {

89 - if text.is_empty() {

95 + ///

96 + /// A piece with empty text but non-empty attributions (zero-width tokens

97 + /// only) is still kept: the tokens are attributed to the current state.

98 + pub(crate) fn push_content(&mut self, piece: DecodedText) {

99 + if piece.is_empty() {

100 return;

101 }

102 match &mut self.content {

93 - Some(existing) => existing.push_str(text),

94 - None => self.content = Some(text.to_string()),

103 + Some(existing) => existing.append(piece),

104 + None => self.content = Some(piece),

105 }

106 }

107 }

@@ -119,7 +129,7 @@ pub trait ReasoningParser: Send {

129 }

130

131 /// Feed one decoded text delta into the parser.

122 - fn push(&mut self, delta: &str) -> Result<ReasoningDelta>;

132 + fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta>;

133

134 /// Flush any buffered partial delimiter state at end of stream.

135 fn finish(&mut self) -> Result<ReasoningDelta> {

rust/src/parser/src/reasoning/qwen3.rs

+2 -2

@@ -1,7 +1,7 @@

1 // SPDX-License-Identifier: Apache-2.0

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -use vllm_tokenizer::DynTokenizer;

4 +use vllm_tokenizer::{DecodedText, DynTokenizer};

6 use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};

@@ -36,7 +36,7 @@ impl ReasoningParser for Qwen3ReasoningParser {

36 Ok(())

37 }

39 - fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {

39 + fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {

40 Ok(self.inner.push(delta))

41 }

rust/src/parser/src/reasoning/seed_oss.rs

+23 -20

@@ -1,7 +1,7 @@

1 // SPDX-License-Identifier: Apache-2.0

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -use vllm_tokenizer::DynTokenizer;

4 +use vllm_tokenizer::{DecodedText, DynTokenizer};

6 use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};

@@ -38,7 +38,7 @@ impl ReasoningParser for SeedOssReasoningParser {

38 Ok(())

39 }

41 - fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {

41 + fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {

42 Ok(self.inner.push(delta))

43 }

@@ -53,17 +53,20 @@ mod tests {

54 use super::SeedOssReasoningParser;

55 use crate::reasoning::ReasoningParser;

56 - use crate::reasoning::tests::{SEED_THINK_END_ID, SEED_THINK_START_ID, fake_tokenizer};

56 + use crate::reasoning::tests::{

57 + SEED_THINK_END_ID, SEED_THINK_START_ID, content_str, fake_tokenizer, push_str,

58 + reasoning_str,

59 + };

61 #[test]

62 fn without_prompt_markers_expects_start_token() {

63 let tokenizer = Arc::new(fake_tokenizer());

64 let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();

63 - let delta = parser.push("implicit reasoning</seed:think>answer").unwrap();

66 + let delta = push_str(&mut parser, "implicit reasoning</seed:think>answer");

67 assert_eq!(delta.reasoning, None);

68 assert_eq!(

66 - delta.content.as_deref(),

69 + content_str(&delta),

70 Some("implicit reasoning</seed:think>answer")

71 );

72 }

@@ -75,9 +78,9 @@ mod tests {

78 // Prompt prefills `<seed:think>`, opening reasoning before the stream.

79 parser.initialize(&[SEED_THINK_START_ID]).unwrap();

78 - let delta = parser.push("reason</seed:think>answer").unwrap();

79 - assert_eq!(delta.reasoning.as_deref(), Some("reason"));

80 - assert_eq!(delta.content.as_deref(), Some("answer"));

81 + let delta = push_str(&mut parser, "reason</seed:think>answer");

82 + assert_eq!(reasoning_str(&delta), Some("reason"));

83 + assert_eq!(content_str(&delta), Some("answer"));

84 }

86 #[test]

@@ -87,9 +90,9 @@ mod tests {

90 // Prompt already closed reasoning with `</seed:think>`.

91 parser.initialize(&[SEED_THINK_END_ID]).unwrap();

90 - let delta = parser.push("answer").unwrap();

93 + let delta = push_str(&mut parser, "answer");

94 assert_eq!(delta.reasoning, None);

92 - assert_eq!(delta.content.as_deref(), Some("answer"));

95 + assert_eq!(content_str(&delta), Some("answer"));

96 }

98 #[test]

@@ -98,9 +101,9 @@ mod tests {

101 let tokenizer = Arc::new(fake_tokenizer());

102 let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();

103

101 - let delta = parser.push("<seed:think>reason</seed:think>answer").unwrap();

102 - assert_eq!(delta.reasoning.as_deref(), Some("reason"));

103 - assert_eq!(delta.content.as_deref(), Some("answer"));

104 + let delta = push_str(&mut parser, "<seed:think>reason</seed:think>answer");

105 + assert_eq!(reasoning_str(&delta), Some("reason"));

106 + assert_eq!(content_str(&delta), Some("answer"));

107 }

108

109 #[test]

@@ -121,12 +124,12 @@ mod tests {

124 "Final ",

125 "answer",

126 ] {

124 - let delta = parser.push(delta_str).unwrap();

127 + let delta = push_str(&mut parser, delta_str);

128 if let Some(r) = delta.reasoning {

126 - reasoning.push_str(&r);

129 + reasoning.push_str(&r.text);

130 }

131 if let Some(c) = delta.content {

129 - content.push_str(&c);

132 + content.push_str(&c.text);

133 }

134 }

135 assert_eq!(reasoning, "Some reasoning content");

@@ -140,12 +143,12 @@ mod tests {

143 parser.initialize(&[SEED_THINK_START_ID]).unwrap();

144

145 // Closing delimiter `</seed:think>` arrives in two halves.

143 - let first = parser.push("reason</seed:").unwrap();

144 - assert_eq!(first.reasoning.as_deref(), Some("reason"));

146 + let first = push_str(&mut parser, "reason</seed:");

147 + assert_eq!(reasoning_str(&first), Some("reason"));

148 assert_eq!(first.content, None);

149

147 - let second = parser.push("think>answer").unwrap();

150 + let second = push_str(&mut parser, "think>answer");

151 assert_eq!(second.reasoning, None);

149 - assert_eq!(second.content.as_deref(), Some("answer"));

152 + assert_eq!(content_str(&second), Some("answer"));

153 }

154 }

rust/src/parser/src/reasoning/step3p5.rs

+69 -65

@@ -1,7 +1,7 @@

1 // SPDX-License-Identifier: Apache-2.0

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -use vllm_tokenizer::DynTokenizer;

4 +use vllm_tokenizer::{DecodedText, DynTokenizer};

6 use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};

@@ -16,7 +16,7 @@ use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};

16 pub struct Step3p5ReasoningParser {

17 inner: DelimitedReasoningParser,

18 /// `\n` at end of last reasoning delta, held in case `</think>` follows.

19 - pending_reasoning_newline: bool,

19 + pending_reasoning_newline: Option<DecodedText>,

20 /// Last push ended on `</think>` without emitting content; the next

21 /// content delta's leading `\n` should be dropped.

22 just_ended_reasoning: bool,

@@ -27,7 +27,7 @@ impl Step3p5ReasoningParser {

27 pub fn new(tokenizer: DynTokenizer) -> Result<Self> {

28 Ok(Self {

29 inner: DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false)?,

30 - pending_reasoning_newline: false,

30 + pending_reasoning_newline: None,

31 just_ended_reasoning: false,

32 })

33 }

@@ -45,40 +45,42 @@ impl Step3p5ReasoningParser {

45 !now_in_reasoning && (was_in_reasoning || inner_delta.reasoning.is_some());

47 // Replay or drop a previously-held trailing reasoning newline.

48 - if self.pending_reasoning_newline {

48 + if self.pending_reasoning_newline.is_some() {

49 if let Some(reasoning) = inner_delta.reasoning.as_mut() {

50 - reasoning.insert(0, '\n');

51 - self.pending_reasoning_newline = false;

50 + let mut held = self.pending_reasoning_newline.take().unwrap();

51 + held.append(std::mem::take(reasoning));

52 + *reasoning = held;

53 } else if transitioned {

54 // The held `\n` was the one right before `</think>`: drop it.

54 - self.pending_reasoning_newline = false;

55 + self.pending_reasoning_newline = None;

56 }

57 }

59 // Hold back a trailing reasoning `\n` until we know if `</think>` follows.

60 if let Some(reasoning) = inner_delta.reasoning.as_mut()

60 - && reasoning.ends_with('\n')

61 + && reasoning.text.ends_with('\n')

62 {

62 - reasoning.pop();

63 + let kept = reasoning.drain_prefix(reasoning.text.len() - 1);

64 + let newline = std::mem::replace(reasoning, kept);

65 if !transitioned {

64 - self.pending_reasoning_newline = true;

66 + self.pending_reasoning_newline = Some(newline);

67 }

68 }

70 // Drop a leading `\n` of content emitted right after `</think>`.

71 if let Some(content) = inner_delta.content.as_mut()

72 && (transitioned || self.just_ended_reasoning)

71 - && content.starts_with('\n')

73 + && content.text.starts_with('\n')

74 {

73 - content.remove(0);

75 + let _ = content.drain_prefix(1);

76 }

78 self.just_ended_reasoning = transitioned && inner_delta.content.is_none();

78 - if inner_delta.reasoning.as_deref() == Some("") {

80 + if inner_delta.reasoning.as_ref().is_some_and(DecodedText::is_empty) {

81 inner_delta.reasoning = None;

82 }

81 - if inner_delta.content.as_deref() == Some("") {

83 + if inner_delta.content.as_ref().is_some_and(DecodedText::is_empty) {

84 inner_delta.content = None;

85 }

@@ -99,7 +101,7 @@ impl ReasoningParser for Step3p5ReasoningParser {

101 Ok(())

102 }

103

102 - fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {

104 + fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {

105 let was = self.inner.in_reasoning();

106 let inner_delta = self.inner.push(delta);

107 let now = self.inner.in_reasoning();

@@ -113,12 +115,11 @@ impl ReasoningParser for Step3p5ReasoningParser {

115 let mut delta = self.process(inner_delta, was, now);

116

117 // Emit a still-held newline rather than silently dropping it.

116 - if self.pending_reasoning_newline {

118 + if let Some(held) = self.pending_reasoning_newline.take() {

119 match delta.reasoning.as_mut() {

118 - Some(existing) => existing.push('\n'),

119 - None => delta.reasoning = Some("\n".to_string()),

120 + Some(existing) => existing.append(held),

121 + None => delta.reasoning = Some(held),

122 }

121 - self.pending_reasoning_newline = false;

123 }

124

125 Ok(delta)

@@ -131,7 +132,9 @@ mod tests {

132

133 use super::Step3p5ReasoningParser;

134 use crate::reasoning::ReasoningParser;

134 - use crate::reasoning::tests::{THINK_START_ID, fake_tokenizer};

135 + use crate::reasoning::tests::{

136 + THINK_START_ID, content_str, fake_tokenizer, push_str, reasoning_str,

137 + };

138

139 #[test]

140 fn picks_up_prompt_start_boundary() {

@@ -140,12 +143,12 @@ mod tests {

143 // Prompt prefills `<think>`, opening reasoning before the stream.

144 parser.initialize(&[THINK_START_ID]).unwrap();

145

143 - let delta = parser.push("This is a reasoning section</think>This is the rest").unwrap();

144 - assert_eq!(

145 - delta.reasoning.as_deref(),

146 - Some("This is a reasoning section")

146 + let delta = push_str(

147 + &mut parser,

148 + "This is a reasoning section</think>This is the rest",

149 );

148 - assert_eq!(delta.content.as_deref(), Some("This is the rest"));

150 + assert_eq!(reasoning_str(&delta), Some("This is a reasoning section"));

151 + assert_eq!(content_str(&delta), Some("This is the rest"));

152 }

153

154 #[test]

@@ -153,8 +156,8 @@ mod tests {

156 let tokenizer = Arc::new(fake_tokenizer());

157 let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();

158

156 - let pushed = parser.push("<think>reason without end").unwrap();

157 - assert_eq!(pushed.reasoning.as_deref(), Some("reason without end"));

159 + let pushed = push_str(&mut parser, "<think>reason without end");

160 + assert_eq!(reasoning_str(&pushed), Some("reason without end"));

161 assert_eq!(pushed.content, None);

162

163 let flushed = parser.finish().unwrap();

@@ -166,7 +169,7 @@ mod tests {

169 let tokenizer = Arc::new(fake_tokenizer());

170 let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();

171

169 - let pushed = parser.push("").unwrap();

172 + let pushed = push_str(&mut parser, "");

173 assert!(pushed.is_empty());

174 let flushed = parser.finish().unwrap();

175 assert!(flushed.is_empty());

@@ -180,14 +183,15 @@ mod tests {

183 let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();

184 parser.initialize(&[THINK_START_ID]).unwrap();

185

183 - let delta = parser

184 - .push("\n This is a \n reasoning section\n\n\n</think>\n\nThis is the rest")

185 - .unwrap();

186 + let delta = push_str(

187 + &mut parser,

188 + "\n This is a \n reasoning section\n\n\n</think>\n\nThis is the rest",

189 + );

190 assert_eq!(

187 - delta.reasoning.as_deref(),

191 + reasoning_str(&delta),

192 Some("\n This is a \n reasoning section\n\n")

193 );

190 - assert_eq!(delta.content.as_deref(), Some("\nThis is the rest"));

194 + assert_eq!(content_str(&delta), Some("\nThis is the rest"));

195 }

196

197 #[test]

@@ -195,9 +199,9 @@ mod tests {

199 let tokenizer = Arc::new(fake_tokenizer());

200 let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();

201

198 - let delta = parser.push("<think>reason\n</think>\nanswer").unwrap();

199 - assert_eq!(delta.reasoning.as_deref(), Some("reason"));

200 - assert_eq!(delta.content.as_deref(), Some("answer"));

202 + let delta = push_str(&mut parser, "<think>reason\n</think>\nanswer");

203 + assert_eq!(reasoning_str(&delta), Some("reason"));

204 + assert_eq!(content_str(&delta), Some("answer"));

205 }

206

207 #[test]

@@ -207,18 +211,18 @@ mod tests {

211

212 // The trailing `\n` from the first push is held until we know whether

213 // `</think>` follows.

210 - let first = parser.push("<think>reason\n").unwrap();

211 - assert_eq!(first.reasoning.as_deref(), Some("reason"));

214 + let first = push_str(&mut parser, "<think>reason\n");

215 + assert_eq!(reasoning_str(&first), Some("reason"));

216 assert_eq!(first.content, None);

217

218 // `</think>` arrives standalone; the held newline should be dropped.

215 - let second = parser.push("</think>").unwrap();

219 + let second = push_str(&mut parser, "</think>");

220 assert!(second.is_empty());

221

222 // The leading newline of the first content delta is dropped.

219 - let third = parser.push("\nanswer").unwrap();

223 + let third = push_str(&mut parser, "\nanswer");

224 assert_eq!(third.reasoning, None);

221 - assert_eq!(third.content.as_deref(), Some("answer"));

225 + assert_eq!(content_str(&third), Some("answer"));

226 }

227

228 #[test]

@@ -226,11 +230,11 @@ mod tests {

230 let tokenizer = Arc::new(fake_tokenizer());

231 let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();

232

229 - let first = parser.push("<think>reason\n").unwrap();

230 - assert_eq!(first.reasoning.as_deref(), Some("reason"));

233 + let first = push_str(&mut parser, "<think>reason\n");

234 + assert_eq!(reasoning_str(&first), Some("reason"));

235

232 - let second = parser.push("more reason").unwrap();

233 - assert_eq!(second.reasoning.as_deref(), Some("\nmore reason"));

236 + let second = push_str(&mut parser, "more reason");

237 + assert_eq!(reasoning_str(&second), Some("\nmore reason"));

238 assert_eq!(second.content, None);

239 }

240

@@ -239,11 +243,11 @@ mod tests {

243 let tokenizer = Arc::new(fake_tokenizer());

244 let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();

245

242 - let first = parser.push("<think>reason\n").unwrap();

243 - assert_eq!(first.reasoning.as_deref(), Some("reason"));

246 + let first = push_str(&mut parser, "<think>reason\n");

247 + assert_eq!(reasoning_str(&first), Some("reason"));

248

249 let flushed = parser.finish().unwrap();

246 - assert_eq!(flushed.reasoning.as_deref(), Some("\n"));

250 + assert_eq!(reasoning_str(&flushed), Some("\n"));

251 assert_eq!(flushed.content, None);

252 }

253

@@ -252,9 +256,9 @@ mod tests {

256 let tokenizer = Arc::new(fake_tokenizer());

257 let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();

258

255 - let delta = parser.push("<think>line1\nline2</think>tail").unwrap();

256 - assert_eq!(delta.reasoning.as_deref(), Some("line1\nline2"));

257 - assert_eq!(delta.content.as_deref(), Some("tail"));

259 + let delta = push_str(&mut parser, "<think>line1\nline2</think>tail");

260 + assert_eq!(reasoning_str(&delta), Some("line1\nline2"));

261 + assert_eq!(content_str(&delta), Some("tail"));

262 }

263

264 #[test]

@@ -264,9 +268,9 @@ mod tests {

268 let tokenizer = Arc::new(fake_tokenizer());

269 let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();

270

267 - let delta = parser.push("<think>reason\n\n</think>answer").unwrap();

268 - assert_eq!(delta.reasoning.as_deref(), Some("reason\n"));

269 - assert_eq!(delta.content.as_deref(), Some("answer"));

271 + let delta = push_str(&mut parser, "<think>reason\n\n</think>answer");

272 + assert_eq!(reasoning_str(&delta), Some("reason\n"));

273 + assert_eq!(content_str(&delta), Some("answer"));

274 }

275

276 #[test]

@@ -276,18 +280,18 @@ mod tests {

280 let tokenizer = Arc::new(fake_tokenizer());

281 let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();

282

279 - let first = parser.push("<think>reason</think>").unwrap();

280 - assert_eq!(first.reasoning.as_deref(), Some("reason"));

283 + let first = push_str(&mut parser, "<think>reason</think>");

284 + assert_eq!(reasoning_str(&first), Some("reason"));

285 assert_eq!(first.content, None);

286

283 - let second = parser.push("\nfirst").unwrap();

287 + let second = push_str(&mut parser, "\nfirst");

288 assert_eq!(second.reasoning, None);

285 - assert_eq!(second.content.as_deref(), Some("first"));

289 + assert_eq!(content_str(&second), Some("first"));

290

291 // A `\n` arriving in a later content delta must NOT be dropped.

288 - let third = parser.push("\nsecond").unwrap();

292 + let third = push_str(&mut parser, "\nsecond");

293 assert_eq!(third.reasoning, None);

290 - assert_eq!(third.content.as_deref(), Some("\nsecond"));

294 + assert_eq!(content_str(&third), Some("\nsecond"));

295 }

296

297 #[test]

@@ -295,9 +299,9 @@ mod tests {

299 let tokenizer = Arc::new(fake_tokenizer());

300 let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();

301

298 - let delta = parser.push("<think>reason</think>tail").unwrap();

299 - assert_eq!(delta.reasoning.as_deref(), Some("reason"));

300 - assert_eq!(delta.content.as_deref(), Some("tail"));

302 + let delta = push_str(&mut parser, "<think>reason</think>tail");

303 + assert_eq!(reasoning_str(&delta), Some("reason"));

304 + assert_eq!(content_str(&delta), Some("tail"));

305 }

306

307 #[test]

@@ -305,8 +309,8 @@ mod tests {

309 let tokenizer = Arc::new(fake_tokenizer());

310 let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();

311

308 - let delta = parser.push("<think></think>answer").unwrap();

312 + let delta = push_str(&mut parser, "<think></think>answer");

313 assert_eq!(delta.reasoning, None);

310 - assert_eq!(delta.content.as_deref(), Some("answer"));

314 + assert_eq!(content_str(&delta), Some("answer"));

315 }

316 }

rust/src/parser/src/reasoning/tests.rs

+254 -49

@@ -4,10 +4,12 @@

4 use std::sync::Arc;

6 use vllm_tokenizer::test_utils::TestTokenizer;

7 +use vllm_tokenizer::{DecodedText, DynTokenizer, TokenAnchor, TokenAttribution};

9 use super::{

9 - DeepSeekR1ReasoningParser, DelimitedReasoningParser, MiniMaxM3ReasoningParser,

10 - Qwen3ReasoningParser, ReasoningParser,

10 + CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DelimitedReasoningParser,

11 + KimiReasoningParser, MiniMaxM3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta,

12 + ReasoningParser, Result, SeedOssReasoningParser, Step3p5ReasoningParser,

13 };

15 pub(crate) const THINK_START_ID: u32 = 256;

@@ -37,16 +39,29 @@ pub(crate) fn fake_tokenizer() -> TestTokenizer {

39 .with_regular_token("</seed:think>", SEED_THINK_END_ID)

40 }

42 +/// Feed an unattributed text delta into a reasoning parser.

43 +pub(crate) fn push_str<P: ReasoningParser + ?Sized>(parser: &mut P, delta: &str) -> ReasoningDelta {

44 + parser.push(DecodedText::unattributed(delta)).unwrap()

45 +}

46 +

47 +/// Return the reasoning text of a delta, if any.

48 +pub(crate) fn reasoning_str(delta: &ReasoningDelta) -> Option<&str> {

49 + delta.reasoning.as_ref().map(|piece| piece.text.as_str())

50 +}

51 +

52 +/// Return the content text of a delta, if any.

53 +pub(crate) fn content_str(delta: &ReasoningDelta) -> Option<&str> {

54 + delta.content.as_ref().map(|piece| piece.text.as_str())

55 +}

56 +

57 #[test]

58 fn delimited_content_only_stream() {

59 let tokenizer = Arc::new(fake_tokenizer());

60 let mut parser =

61 DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();

46 - assert_eq!(

47 - parser.push("plain content").content.as_deref(),

48 - Some("plain content")

49 - );

63 + let delta = parser.push(DecodedText::unattributed("plain content"));

64 + assert_eq!(content_str(&delta), Some("plain content"));

65 }

67 #[test]

@@ -55,9 +70,9 @@ fn delimited_single_chunk_with_reasoning_and_content() {

70 let mut parser =

71 DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();

58 - let delta = parser.push("<think>reason</think>answer");

59 - assert_eq!(delta.reasoning.as_deref(), Some("reason"));

60 - assert_eq!(delta.content.as_deref(), Some("answer"));

73 + let delta = parser.push(DecodedText::unattributed("<think>reason</think>answer"));

74 + assert_eq!(reasoning_str(&delta), Some("reason"));

75 + assert_eq!(content_str(&delta), Some("answer"));

76 }

78 #[test]

@@ -66,10 +81,10 @@ fn delimited_partial_tokens_across_chunks() {

81 let mut parser =

82 DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();

69 - assert!(parser.push("<thi").is_empty());

70 - let delta = parser.push("nk>reason</think>answer");

71 - assert_eq!(delta.reasoning.as_deref(), Some("reason"));

72 - assert_eq!(delta.content.as_deref(), Some("answer"));

84 + assert!(parser.push(DecodedText::unattributed("<thi")).is_empty());

85 + let delta = parser.push(DecodedText::unattributed("nk>reason</think>answer"));

86 + assert_eq!(reasoning_str(&delta), Some("reason"));

87 + assert_eq!(content_str(&delta), Some("answer"));

88 }

90 #[test]

@@ -79,10 +94,10 @@ fn delimited_finish_flushes_buffer() {

94 DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();

95 parser.initialize(&[THINK_START_ID]);

82 - let delta = parser.push("unfinished</thi");

83 - assert_eq!(delta.reasoning.as_deref(), Some("unfinished"));

97 + let delta = parser.push(DecodedText::unattributed("unfinished</thi"));

98 + assert_eq!(reasoning_str(&delta), Some("unfinished"));

99 let final_delta = parser.finish();

85 - assert_eq!(final_delta.reasoning.as_deref(), Some("</thi"));

100 + assert_eq!(reasoning_str(&final_delta), Some("</thi"));

101 }

102

103 #[test]

@@ -90,9 +105,9 @@ fn qwen3_without_prompt_markers_expects_start_token() {

105 let tokenizer = Arc::new(fake_tokenizer());

106 let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();

107

93 - let delta = parser.push("reason</think>answer").unwrap();

108 + let delta = push_str(&mut parser, "reason</think>answer");

109 assert_eq!(delta.reasoning, None);

95 - assert_eq!(delta.content.as_deref(), Some("reason</think>answer"));

110 + assert_eq!(content_str(&delta), Some("reason</think>answer"));

111 }

112

113 #[test]

@@ -101,9 +116,9 @@ fn qwen3_prompt_end_marker_starts_in_content() {

116 let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();

117 parser.initialize(&[THINK_END_ID]).unwrap();

118

104 - let delta = parser.push("answer").unwrap();

119 + let delta = push_str(&mut parser, "answer");

120 assert_eq!(delta.reasoning, None);

106 - assert_eq!(delta.content.as_deref(), Some("answer"));

121 + assert_eq!(content_str(&delta), Some("answer"));

122 }

123

124 #[test]

@@ -111,15 +126,15 @@ fn qwen3_tolerates_old_and_new_formats() {

126 let tokenizer = Arc::new(fake_tokenizer());

127

128 let mut old_parser = Qwen3ReasoningParser::new(tokenizer.clone()).unwrap();

114 - let old = old_parser.push("<think>reason</think>answer").unwrap();

115 - assert_eq!(old.reasoning.as_deref(), Some("reason"));

116 - assert_eq!(old.content.as_deref(), Some("answer"));

129 + let old = push_str(&mut old_parser, "<think>reason</think>answer");

130 + assert_eq!(reasoning_str(&old), Some("reason"));

131 + assert_eq!(content_str(&old), Some("answer"));

132

133 let mut new_parser = Qwen3ReasoningParser::new(tokenizer).unwrap();

134 new_parser.initialize(&[THINK_START_ID]).unwrap();

120 - let new = new_parser.push("reason</think>answer").unwrap();

121 - assert_eq!(new.reasoning.as_deref(), Some("reason"));

122 - assert_eq!(new.content.as_deref(), Some("answer"));

135 + let new = push_str(&mut new_parser, "reason</think>answer");

136 + assert_eq!(reasoning_str(&new), Some("reason"));

137 + assert_eq!(content_str(&new), Some("answer"));

138 }

139

140 #[test]

@@ -129,9 +144,9 @@ fn qwen3_stops_scanning_at_last_special_token() {

144

145 parser.initialize(&[THINK_START_ID, SPECIAL_BOUNDARY_ID]).unwrap();

146

132 - let delta = parser.push("answer").unwrap();

147 + let delta = push_str(&mut parser, "answer");

148 assert_eq!(delta.reasoning, None);

134 - assert_eq!(delta.content.as_deref(), Some("answer"));

149 + assert_eq!(content_str(&delta), Some("answer"));

150 }

151

152 #[test]

@@ -139,9 +154,9 @@ fn deepseek_r1_defaults_to_reasoning_without_prompt_boundary() {

154 let tokenizer = Arc::new(fake_tokenizer());

155 let mut parser = DeepSeekR1ReasoningParser::new(tokenizer).unwrap();

156

142 - let delta = parser.push("reason</think>answer").unwrap();

143 - assert_eq!(delta.reasoning.as_deref(), Some("reason"));

144 - assert_eq!(delta.content.as_deref(), Some("answer"));

157 + let delta = push_str(&mut parser, "reason</think>answer");

158 + assert_eq!(reasoning_str(&delta), Some("reason"));

159 + assert_eq!(content_str(&delta), Some("answer"));

160 }

161

162 #[test]

@@ -151,9 +166,9 @@ fn deepseek_r1_stops_scanning_at_last_special_token() {

166

167 parser.initialize(&[THINK_END_ID, SPECIAL_BOUNDARY_ID]).unwrap();

168

154 - let delta = parser.push("reason</think>answer").unwrap();

155 - assert_eq!(delta.reasoning.as_deref(), Some("reason"));

156 - assert_eq!(delta.content.as_deref(), Some("answer"));

169 + let delta = push_str(&mut parser, "reason</think>answer");

170 + assert_eq!(reasoning_str(&delta), Some("reason"));

171 + assert_eq!(content_str(&delta), Some("answer"));

172 }

173

174 #[test]

@@ -161,9 +176,9 @@ fn minimax_m3_handles_explicit_think_delimiters() {

176 let tokenizer = Arc::new(fake_tokenizer());

177 let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();

178

164 - let delta = parser.push("<mm:think>reason</mm:think>answer").unwrap();

165 - assert_eq!(delta.reasoning.as_deref(), Some("reason"));

166 - assert_eq!(delta.content.as_deref(), Some("answer"));

179 + let delta = push_str(&mut parser, "<mm:think>reason</mm:think>answer");

180 + assert_eq!(reasoning_str(&delta), Some("reason"));

181 + assert_eq!(content_str(&delta), Some("answer"));

182 }

183

184 #[test]

@@ -171,9 +186,9 @@ fn minimax_m3_drops_leading_end_marker() {

186 let tokenizer = Arc::new(fake_tokenizer());

187 let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();

188

174 - let delta = parser.push("</mm:think>answer").unwrap();

189 + let delta = push_str(&mut parser, "</mm:think>answer");

190 assert_eq!(delta.reasoning, None);

176 - assert_eq!(delta.content.as_deref(), Some("answer"));

191 + assert_eq!(content_str(&delta), Some("answer"));

192 }

193

194 #[test]

@@ -181,9 +196,9 @@ fn minimax_m3_preserves_non_leading_end_marker() {

196 let tokenizer = Arc::new(fake_tokenizer());

197 let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();

198

184 - let delta = parser.push("XXX</mm:think>YYY").unwrap();

199 + let delta = push_str(&mut parser, "XXX</mm:think>YYY");

200 assert_eq!(delta.reasoning, None);

186 - assert_eq!(delta.content.as_deref(), Some("XXX</mm:think>YYY"));

201 + assert_eq!(content_str(&delta), Some("XXX</mm:think>YYY"));

202 }

203

204 #[test]

@@ -191,10 +206,10 @@ fn minimax_m3_drops_split_leading_end_marker() {

206 let tokenizer = Arc::new(fake_tokenizer());

207 let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();

208

194 - assert!(parser.push("</mm").unwrap().is_empty());

195 - let delta = parser.push(":think>answer").unwrap();

209 + assert!(push_str(&mut parser, "</mm").is_empty());

210 + let delta = push_str(&mut parser, ":think>answer");

211 assert_eq!(delta.reasoning, None);

197 - assert_eq!(delta.content.as_deref(), Some("answer"));

212 + assert_eq!(content_str(&delta), Some("answer"));

213 }

214

215 #[test]

@@ -203,9 +218,9 @@ fn minimax_m3_uses_prompt_prefilled_start_marker() {

218 let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();

219 parser.initialize(&[MM_THINK_START_ID]).unwrap();

220

206 - let delta = parser.push("reason</mm:think>answer").unwrap();

207 - assert_eq!(delta.reasoning.as_deref(), Some("reason"));

208 - assert_eq!(delta.content.as_deref(), Some("answer"));

221 + let delta = push_str(&mut parser, "reason</mm:think>answer");

222 + assert_eq!(reasoning_str(&delta), Some("reason"));

223 + assert_eq!(content_str(&delta), Some("answer"));

224 }

225

226 #[test]

@@ -214,7 +229,197 @@ fn minimax_m3_uses_prompt_prefilled_end_marker() {

229 let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();

230 parser.initialize(&[MM_THINK_END_ID]).unwrap();

231

217 - let delta = parser.push("answer").unwrap();

232 + let delta = push_str(&mut parser, "answer");

233 assert_eq!(delta.reasoning, None);

219 - assert_eq!(delta.content.as_deref(), Some("answer"));

234 + assert_eq!(content_str(&delta), Some("answer"));

235 +}

236 +

237 +#[test]

238 +fn delimited_zero_width_only_piece_is_attributed_to_current_state() {

239 + let tokenizer = Arc::new(fake_tokenizer());

240 + let mut parser =

241 + DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();

242 + parser.initialize(&[THINK_START_ID]);

243 +

244 + // A filtered special token produces a zero-width attribution with no text;

245 + // it must survive as a reasoning piece rather than being dropped.

246 + let delta = parser.push(DecodedText {

247 + text: String::new(),

248 + attributions: [TokenAttribution {

249 + token_id: 42,

250 + anchor: TokenAnchor::ZeroWidth { byte_offset: 0 },

251 + }]

252 + .into_iter()

253 + .collect(),

254 + });

255 + assert_eq!(reasoning_str(&delta), Some(""));

256 + assert_eq!(

257 + delta

258 + .reasoning

259 + .as_ref()

260 + .unwrap()

261 + .attributions

262 + .iter()

263 + .map(|attr| attr.token_id)

264 + .collect::<Vec<_>>(),

265 + [42]

266 + );

267 +}

268 +

269 +#[test]

270 +fn delimited_marker_tokens_are_dropped_from_attributions() {

271 + let tokenizer = Arc::new(fake_tokenizer());

272 + let mut parser =

273 + DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();

274 +

275 + let mut collected = CollectedAttributions::default();

276 + for chunk in attributed_chunks(&[

277 + Some("<think>"),

278 + Some("reason"),

279 + Some("</think>"),

280 + Some("answer"),

281 + ]) {

282 + collected.record(parser.push(chunk));

283 + }

284 + collected.record(parser.finish());

285 +

286 + assert_eq!(collected.reasoning_text, "reason");

287 + assert_eq!(collected.content_text, "answer");

288 + // The marker tokens (1 and 3) are dropped with their spans; only the

289 + // reasoning and content tokens remain attributed.

290 + assert_eq!(collected.reasoning_ids, [2]);

291 + assert_eq!(collected.content_ids, [4]);

292 +}

293 +

294 +#[test]

295 +fn reasoning_parsers_conserve_token_attributions() {

296 + type ParserCase = (

297 + fn(DynTokenizer) -> Result<Box<dyn ReasoningParser>>,

298 + u32,

299 + &'static str,

300 + );

301 +

302 + let parsers: Vec<ParserCase> = vec![

303 + (Qwen3ReasoningParser::create, THINK_START_ID, "qwen3"),

304 + (

305 + DeepSeekR1ReasoningParser::create,

306 + THINK_START_ID,

307 + "deepseek_r1",

308 + ),

309 + (KimiReasoningParser::create, MINIMAX_THINK_START_ID, "kimi"),

310 + (

311 + CohereCmdReasoningParser::create,

312 + START_THINKING_ID,

313 + "cohere_cmd",

314 + ),

315 + (

316 + SeedOssReasoningParser::create,

317 + SEED_THINK_START_ID,

318 + "seed_oss",

319 + ),

320 + (

321 + MiniMaxM3ReasoningParser::create,

322 + MM_THINK_START_ID,

323 + "minimax_m3",

324 + ),

325 + (Step3p5ReasoningParser::create, THINK_START_ID, "step3p5"),

326 + ];

327 +

328 + for (create, think_start_id, name) in parsers {

329 + let mut parser = create(Arc::new(fake_tokenizer())).unwrap();

330 + // Prompt prefills the start marker, opening reasoning before the stream.

331 + parser.initialize(&[think_start_id]).unwrap();

332 +

333 + // Script: reasoning text (1), a zero-width token inside reasoning (2),

334 + // the end marker (3), then visible content (4). The zero-width token

335 + // is attributed to reasoning, matching Python's current-state rule;

336 + // the marker token is dropped with its span.

337 + let end_marker = match name {

338 + "kimi" => "◁/think▷",

339 + "cohere_cmd" => "<|END_THINKING|>",

340 + "seed_oss" => "</seed:think>",

341 + "minimax_m3" => "</mm:think>",

342 + _ => "</think>",

343 + };

344 + let chunks = attributed_chunks(&[Some("reason"), None, Some(end_marker), Some("answer")]);

345 +

346 + let collected = collect_attributed(parser.as_mut(), chunks);

347 +

348 + assert_eq!(collected.reasoning_text, "reason", "{name}");

349 + assert_eq!(collected.content_text, "answer", "{name}");

350 + assert_eq!(collected.reasoning_ids, [1, 2], "{name}");

351 + assert_eq!(collected.content_ids, [4], "{name}");

352 + }

353 +}

354 +

355 +/// Build one chunk per piece, each carrying a single token attributed to it.

356 +///

357 +/// A `Some(text)` piece produces one token visibly anchored at the piece start;

358 +/// a `None` piece produces one zero-width token. Token IDs are 1-based

359 +/// positions in the script.

360 +pub(crate) fn attributed_chunks(pieces: &[Option<&str>]) -> Vec<DecodedText> {

361 + pieces

362 + .iter()

363 + .enumerate()

364 + .map(|(index, piece)| {

365 + let token_id = index as u32 + 1;

366 + match piece {

367 + Some(text) => DecodedText {

368 + text: (*text).to_string(),

369 + attributions: [TokenAttribution {

370 + token_id,

371 + anchor: TokenAnchor::Visible { byte_offset: 0 },

372 + }]

373 + .into_iter()

374 + .collect(),

375 + },

376 + None => DecodedText {

377 + text: String::new(),

378 + attributions: [TokenAttribution {

379 + token_id,

380 + anchor: TokenAnchor::ZeroWidth { byte_offset: 0 },

381 + }]

382 + .into_iter()

383 + .collect(),

384 + },

385 + }

386 + })

387 + .collect()

388 +}

389 +

390 +/// Token attributions and text collected from a stream of reasoning deltas.

391 +#[derive(Default)]

392 +pub(crate) struct CollectedAttributions {

393 + pub reasoning_ids: Vec<u32>,

394 + pub content_ids: Vec<u32>,

395 + pub reasoning_text: String,

396 + pub content_text: String,

397 +}

398 +

399 +impl CollectedAttributions {

400 + pub fn record(&mut self, delta: ReasoningDelta) {

401 + if let Some(reasoning) = delta.reasoning {

402 + self.reasoning_ids

403 + .extend(reasoning.attributions.iter().map(|attr| attr.token_id));

404 + self.reasoning_text.push_str(&reasoning.text);

405 + }

406 + if let Some(content) = delta.content {

407 + self.content_ids.extend(content.attributions.iter().map(|attr| attr.token_id));

408 + self.content_text.push_str(&content.text);

409 + }

410 + }

411 +}

412 +

413 +/// Push attributed chunks through a reasoning parser and collect what is

414 +/// attributed to reasoning and content pieces.

415 +pub(crate) fn collect_attributed(

416 + parser: &mut dyn ReasoningParser,

417 + chunks: Vec<DecodedText>,

418 +) -> CollectedAttributions {

419 + let mut collected = CollectedAttributions::default();

420 + for chunk in chunks {

421 + collected.record(parser.push(chunk).unwrap());

422 + }

423 + collected.record(parser.finish().unwrap());

424 + collected

425 }

rust/src/parser/src/unified/combined.rs

+74 -11

@@ -3,7 +3,7 @@

4 //! Adapter that combines reasoning and tool parsers.

6 -use vllm_tokenizer::DynTokenizer;

6 +use vllm_tokenizer::{DecodedText, DynTokenizer};

8 use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};

9 use crate::reasoning::ReasoningParser;

@@ -87,9 +87,9 @@ impl UnifiedParser for CombinedParser {

87 self.tool.as_ref().and_then(|parser| parser.tool_call_id(tool_index))

88 }

90 - fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()> {

90 + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> {

91 let Some(reasoning) = self.reasoning.as_mut() else {

92 - return self.parse_tool(delta, output);

92 + return self.parse_tool(&delta.text, output);

93 };

95 let reasoning_delta = reasoning.push(delta)?;

@@ -97,7 +97,11 @@ impl UnifiedParser for CombinedParser {

97 output.push_reasoning(reasoning);

98 }

99 if let Some(content) = reasoning_delta.content {

100 - self.parse_tool(&content, output)?;

100 + // Content attributions stop at this boundary: the tool parser trait

101 + // consumes plain text.

102 + if !content.text.is_empty() {

103 + self.parse_tool(&content.text, output)?;

104 + }

105 }

106 Ok(())

107 }

@@ -109,8 +113,10 @@ impl UnifiedParser for CombinedParser {

113 if let Some(reasoning) = reasoning_delta.reasoning {

114 output.push_reasoning(reasoning);

115 }

112 - if let Some(content) = reasoning_delta.content {

113 - self.parse_tool(&content, &mut output)?;

116 + if let Some(content) = reasoning_delta.content

117 + && !content.text.is_empty()

118 + {

119 + self.parse_tool(&content.text, &mut output)?;

120 }

121 }

122 output.append(self.flush_tool()?);

@@ -127,6 +133,7 @@ mod tests {

133 use std::sync::Arc;

134

135 use vllm_tokenizer::test_utils::TestTokenizer;

136 + use vllm_tokenizer::{DecodedText, TokenAnchor, TokenAttribution};

137

138 use super::CombinedParser;

139 use crate::reasoning::{Qwen3ReasoningParser, ReasoningDelta, ReasoningParser};

@@ -156,7 +163,7 @@ mod tests {

163 fn collect(parser: &mut dyn UnifiedParser, chunks: &[&str]) -> UnifiedParserOutput {

164 let mut output = UnifiedParserOutput::default();

165 for chunk in chunks {

159 - parser.parse_into(chunk, &mut output).unwrap();

166 + parser.parse_into(DecodedText::unattributed(*chunk), &mut output).unwrap();

167 }

168 output.append(parser.finish().unwrap());

169 output

@@ -178,10 +185,10 @@ mod tests {

185 true

186 }

187

181 - fn push(&mut self, delta: &str) -> crate::reasoning::Result<ReasoningDelta> {

188 + fn push(&mut self, delta: DecodedText) -> crate::reasoning::Result<ReasoningDelta> {

189 Ok(ReasoningDelta {

190 reasoning: None,

184 - content: Some(delta.to_string()),

191 + content: Some(delta),

192 })

193 }

194 }

@@ -259,7 +266,63 @@ mod tests {

266 assert_eq!(

267 output.events,

268 vec![

262 - UnifiedParserEvent::Reasoning("work".to_string()),

269 + UnifiedParserEvent::Reasoning(DecodedText::unattributed("work")),

270 + UnifiedParserEvent::Text("answer".to_string()),

271 + ]

272 + );

273 + }

274 +

275 + #[test]

276 + fn combined_parser_reasoning_events_carry_token_attributions() {

277 + let tokenizer = Arc::new(tokenizer());

278 + let reasoning = Qwen3ReasoningParser::create(tokenizer).unwrap();

279 + let mut parser = CombinedParser::new(Some(reasoning), None);

280 +

281 + let chunk = |token_id: u32, text: &str| DecodedText {

282 + text: text.to_string(),

283 + attributions: [TokenAttribution {

284 + token_id,

285 + anchor: TokenAnchor::Visible { byte_offset: 0 },

286 + }]

287 + .into_iter()

288 + .collect(),

289 + };

290 +

291 + let mut output = UnifiedParserOutput::default();

292 + for (token_id, text) in [

293 + (1, "<think>"),

294 + (2, "reason"),

295 + (3, "</think>"),

296 + (4, "answer"),

297 + ] {

298 + parser.parse_into(chunk(token_id, text), &mut output).unwrap();

299 + }

300 + output.append(parser.finish().unwrap());

301 +

302 + // The reasoning tokens keep their attributions through the combined

303 + // parser; marker tokens (1 and 3) are dropped with their spans.

304 + let reasoning_ids: Vec<u32> = output

305 + .events

306 + .iter()

307 + .filter_map(|event| match event {

308 + UnifiedParserEvent::Reasoning(piece) => Some(piece),

309 + _ => None,

310 + })

311 + .flat_map(|piece| piece.attributions.iter().map(|attr| attr.token_id))

312 + .collect();

313 + assert_eq!(reasoning_ids, [2]);

314 + assert_eq!(

315 + output.events,

316 + vec![

317 + UnifiedParserEvent::Reasoning(DecodedText {

318 + text: "reason".to_string(),

319 + attributions: [TokenAttribution {

320 + token_id: 2,

321 + anchor: TokenAnchor::Visible { byte_offset: 0 },

322 + }]

323 + .into_iter()

324 + .collect(),

325 + }),

326 UnifiedParserEvent::Text("answer".to_string()),

327 ]

328 );

@@ -300,7 +363,7 @@ mod tests {

363 let mut parser = CombinedParser::new(None, Some(Box::new(PartialThenErrorToolParser)));

364 let mut output = UnifiedParserOutput::default();

365

303 - let error = parser.parse_into("bad", &mut output).unwrap_err();

366 + let error = parser.parse_into(DecodedText::unattributed("bad"), &mut output).unwrap_err();

367

368 assert!(matches!(error, crate::unified::UnifiedParserError::Tool(_)));

369 assert_eq!(

rust/src/parser/src/unified/gemma4.rs

+89 -35

@@ -2,7 +2,7 @@

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 use serde_json::{Map, Number, Value};

5 -use vllm_tokenizer::DynTokenizer;

5 +use vllm_tokenizer::{DecodedText, DynTokenizer};

6 use winnow::ascii::multispace0 as ws0;

7 use winnow::combinator::{alt, delimited, eof, opt, separated, seq, terminated};

8 use winnow::error::{ContextError, ErrMode, ModalResult};

@@ -29,8 +29,8 @@ type Gemma4Input<'i> = Partial<&'i str>;

30 #[derive(Debug, Clone, PartialEq)]

31 enum Gemma4Event {

32 - Text { len: usize },

33 - Reasoning { len: usize },

32 + Text,

33 + Reasoning,

34 ReasoningStart,

35 ReasoningEnd,

36 ToolCallStart,

@@ -69,7 +69,7 @@ enum Gemma4Mode {

69 ///

70 /// Arguments are emitted only after a full Gemma4 tool call is parsed.

71 pub struct Gemma4UnifiedParser {

72 - buffer: String,

72 + buffer: DecodedText,

73 mode: Gemma4Mode,

74 emitted_tool_count: usize,

75 tokenizer: DynTokenizer,

@@ -84,7 +84,7 @@ impl Gemma4UnifiedParser {

84 let channel_end_token_id = token_id(tokenizer.as_ref(), CHANNEL_END)?;

86 Ok(Self {

87 - buffer: String::new(),

87 + buffer: DecodedText::default(),

88 mode: Gemma4Mode::default(),

89 emitted_tool_count: 0,

90 channel_start_token_id,

@@ -93,14 +93,16 @@ impl Gemma4UnifiedParser {

93 })

94 }

96 - fn apply_event(&mut self, event: Gemma4Event, output: &mut UnifiedParserOutput) -> Result<()> {

96 + fn apply_event(

97 + &mut self,

98 + event: Gemma4Event,

99 + piece: DecodedText,

100 + output: &mut UnifiedParserOutput,

101 + ) -> Result<()> {

102 match event {

98 - Gemma4Event::Text { len: consumed_len } => {

99 - output.push_text(self.buffer[..consumed_len].to_string());

100 - }

101 - Gemma4Event::Reasoning { len: consumed_len } => {

102 - output.push_reasoning(self.buffer[..consumed_len].to_string());

103 - }

103 + Gemma4Event::Text => output.push_text(piece.text),

104 + Gemma4Event::Reasoning => output.push_reasoning(piece),

105 + // Marker and tool-syntax spans are drained and dropped with their tokens.

106 Gemma4Event::ReasoningStart => self.mode = Gemma4Mode::Reasoning,

107 Gemma4Event::ReasoningEnd => self.mode = Gemma4Mode::Text,

108 Gemma4Event::ToolCallStart => self.mode = Gemma4Mode::Header,

@@ -144,22 +146,17 @@ impl Gemma4UnifiedParser {

146 }

147

148 fn reset(&mut self) -> String {

149 + let buffer = self.buffer.take().text;

150 let raw = match std::mem::replace(&mut self.mode, Gemma4Mode::Text) {

148 - Gemma4Mode::Text => std::mem::take(&mut self.buffer),

151 + Gemma4Mode::Text => buffer,

152 Gemma4Mode::Reasoning => {

150 - format!("{}{}", REASONING_START, std::mem::take(&mut self.buffer))

153 + format!("{REASONING_START}{buffer}")

154 }

155 Gemma4Mode::Header => {

153 - format!("{}{}", TOOL_CALL_START, std::mem::take(&mut self.buffer))

156 + format!("{TOOL_CALL_START}{buffer}")

157 }

158 Gemma4Mode::ToolCall { name, .. } => {

156 - format!(

157 - "{}{}{}{{{}",

158 - TOOL_CALL_START,

159 - CALL_PREFIX,

160 - name,

161 - std::mem::take(&mut self.buffer)

162 - )

159 + format!("{TOOL_CALL_START}{CALL_PREFIX}{name}{{{buffer}")

160 }

161 };

162 self.mode = Gemma4Mode::Text;

@@ -187,16 +184,16 @@ impl UnifiedParser for Gemma4UnifiedParser {

184 true

185 }

186

190 - fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> {

191 - self.buffer.push_str(chunk);

187 + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> {

188 + self.buffer.append(delta);

189

190 while let Some((event, consumed_len)) = {

194 - parse_buffered_event(&self.buffer, |input| {

191 + parse_buffered_event(&self.buffer.text, |input| {

192 parse_next_gemma4_event(input, &mut self.mode)

193 })?

194 } {

198 - self.apply_event(event, output)?;

199 - self.buffer.drain(..consumed_len);

195 + let piece = self.buffer.drain_prefix(consumed_len);

196 + self.apply_event(event, piece, output)?;

197 }

198

199 Ok(())

@@ -206,8 +203,8 @@ impl UnifiedParser for Gemma4UnifiedParser {

203 let mut output = UnifiedParserOutput::default();

204

205 match &self.mode {

209 - Gemma4Mode::Text => output.push_text(std::mem::take(&mut self.buffer)),

210 - Gemma4Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)),

206 + Gemma4Mode::Text => output.push_text(self.buffer.take().text),

207 + Gemma4Mode::Reasoning => output.push_reasoning(self.buffer.take()),

208 Gemma4Mode::Header | Gemma4Mode::ToolCall { .. } => {

209 return Err(parsing_failed!("incomplete Gemma4 tool call"));

210 }

@@ -306,14 +303,12 @@ fn gemma4_tool_name(input: &mut Gemma4Input<'_>) -> ModalResult<String> {

303

304 /// Parse a safe text run before the next Gemma4 marker.

305 fn safe_text_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {

309 - safe_text_len_mul(input, &[REASONING_START, TOOL_CALL_START])

310 - .map(|len| Gemma4Event::Text { len })

306 + safe_text_len_mul(input, &[REASONING_START, TOOL_CALL_START]).map(|_| Gemma4Event::Text)

307 }

308

309 /// Parse a safe reasoning run before the next Gemma4 marker.

310 fn safe_reasoning_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {

315 - safe_text_len_mul(input, &[CHANNEL_END, TOOL_CALL_START])

316 - .map(|len| Gemma4Event::Reasoning { len })

311 + safe_text_len_mul(input, &[CHANNEL_END, TOOL_CALL_START]).map(|_| Gemma4Event::Reasoning)

312 }

313

314 /// Parse raw Gemma4 arguments through the first end marker outside a Gemma string.

@@ -514,6 +509,7 @@ mod tests {

509 use serde_json::{Value, json};

510 use thiserror_ext::AsReport;

511 use vllm_tokenizer::test_utils::TestTokenizer;

512 + use vllm_tokenizer::{DecodedText, TokenAnchor, TokenAttribution};

513 use winnow::combinator::{eof, terminated};

514 use winnow::error::ErrMode;

515 use winnow::prelude::*;

@@ -545,7 +541,7 @@ mod tests {

541 impl UnifiedParserTestExt for Gemma4UnifiedParser {

542 fn parse_chunk(&mut self, chunk: &str) -> super::Result<UnifiedParserOutput> {

543 let mut output = UnifiedParserOutput::default();

548 - self.parse_into(chunk, &mut output)?;

544 + self.parse_into(DecodedText::unattributed(chunk), &mut output)?;

545 Ok(output)

546 }

547

@@ -578,7 +574,7 @@ mod tests {

574 self.events

575 .iter()

576 .filter_map(|event| match event {

581 - UnifiedParserEvent::Reasoning(text) => Some(text.as_str()),

577 + UnifiedParserEvent::Reasoning(text) => Some(text.text.as_str()),

578 UnifiedParserEvent::Text(_) | UnifiedParserEvent::ToolCall(_) => None,

579 })

580 .collect()

@@ -1116,4 +1112,62 @@ mod tests {

1112

1113 assert_eq!(raw, input);

1114 }

1115 +

1116 + #[test]

1117 + fn gemma4_reasoning_events_conserve_token_attributions() {

1118 + let mut parser = test_parser();

1119 +

1120 + let chunk = |token_id: u32, text: &str| DecodedText {

1121 + text: text.to_string(),

1122 + attributions: [TokenAttribution {

1123 + token_id,

1124 + anchor: TokenAnchor::Visible { byte_offset: 0 },

1125 + }]

1126 + .into_iter()

1127 + .collect(),

1128 + };

1129 + // A zero-width token inside the reasoning channel counts as reasoning,

1130 + // matching Python's current-state rule.

1131 + let zero_width = |token_id: u32| DecodedText {

1132 + text: String::new(),

1133 + attributions: [TokenAttribution {

1134 + token_id,

1135 + anchor: TokenAnchor::ZeroWidth { byte_offset: 0 },

1136 + }]

1137 + .into_iter()

1138 + .collect(),

1139 + };

1140 +

1141 + let mut output = UnifiedParserOutput::default();

1142 + for delta in [

1143 + chunk(1, "<|channel>thought\n"),

1144 + zero_width(2),

1145 + chunk(3, "reason"),

1146 + chunk(4, "<channel|>"),

1147 + chunk(5, "answer"),

1148 + chunk(6, "<|tool_call>"),

1149 + chunk(7, "call:get_weather{location:<|\"|>Paris<|\"|>}"),

1150 + chunk(8, "<tool_call|>"),

1151 + ] {

1152 + parser.parse_into(delta, &mut output).unwrap();

1153 + }

1154 + output.append(parser.finish().unwrap());

1155 +

1156 + // The reasoning tokens keep their attributions; marker and tool-syntax

1157 + // tokens (1, 4, 6, 7, 8) are dropped with their spans, and text/tool

1158 + // pieces carry no attributions.

1159 + let reasoning_ids: Vec<u32> = output

1160 + .events

1161 + .iter()

1162 + .filter_map(|event| match event {

1163 + UnifiedParserEvent::Reasoning(piece) => Some(piece),

1164 + _ => None,

1165 + })

1166 + .flat_map(|piece| piece.attributions.iter().map(|attr| attr.token_id))

1167 + .collect();

1168 + assert_eq!(reasoning_ids, [2, 3]);

1169 + assert_eq!(output.reasoning_text(), "reason");

1170 + assert_eq!(output.normal_text(), "answer");

1171 + assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));

1172 + }

1173 }

rust/src/parser/src/unified/hy/hy_v3.rs

+5 -5

@@ -1,7 +1,7 @@

1 // SPDX-License-Identifier: Apache-2.0

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -use vllm_tokenizer::DynTokenizer;

4 +use vllm_tokenizer::{DecodedText, DynTokenizer};

6 use super::detect_hy_token_suffix;

7 use crate::reasoning::HyReasoningParser;

@@ -54,7 +54,7 @@ impl UnifiedParser for HyV3UnifiedParser {

54 self.inner.tool_call_id(tool_index)

55 }

57 - fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()> {

57 + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> {

58 self.inner.parse_into(delta, output)

59 }

@@ -72,7 +72,7 @@ mod tests {

72 use std::sync::Arc;

74 use serde_json::json;

75 - use vllm_tokenizer::{Tokenizer, test_utils::TestTokenizer};

75 + use vllm_tokenizer::{DecodedText, Tokenizer, test_utils::TestTokenizer};

76 use xgrammar_structural_tag::builders::StructuralTagOptions;

77 use xgrammar_structural_tag::{

78 FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag,

@@ -131,14 +131,14 @@ mod tests {

131 ];

132 let mut output = UnifiedParserOutput::default();

133 for chunk in chunks {

134 - parser.parse_into(chunk, &mut output).unwrap();

134 + parser.parse_into(DecodedText::unattributed(chunk), &mut output).unwrap();

135 }

136 output.append(parser.finish().unwrap());

137

138 assert_eq!(

139 output.events,

140 vec![

141 - UnifiedParserEvent::Reasoning("reasoning".to_string()),

141 + UnifiedParserEvent::Reasoning(DecodedText::unattributed("reasoning")),

142 UnifiedParserEvent::Text("answer".to_string()),

143 UnifiedParserEvent::ToolCall(crate::tool::ToolCallDelta {

144 tool_index: 0,

rust/src/parser/src/unified/hy/hy_v4.rs

+5 -5

@@ -1,7 +1,7 @@

1 // SPDX-License-Identifier: Apache-2.0

2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -use vllm_tokenizer::DynTokenizer;

4 +use vllm_tokenizer::{DecodedText, DynTokenizer};

6 use super::detect_hy_token_suffix;

7 use crate::reasoning::HyReasoningParser;

@@ -54,7 +54,7 @@ impl UnifiedParser for HyV4UnifiedParser {

54 self.inner.tool_call_id(tool_index)

55 }

57 - fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()> {

57 + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> {

58 self.inner.parse_into(delta, output)

59 }

@@ -72,7 +72,7 @@ mod tests {

72 use std::sync::Arc;

74 use serde_json::json;

75 - use vllm_tokenizer::{Tokenizer, test_utils::TestTokenizer};

75 + use vllm_tokenizer::{DecodedText, Tokenizer, test_utils::TestTokenizer};

76 use xgrammar_structural_tag::builders::StructuralTagOptions;

77 use xgrammar_structural_tag::{

78 FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag,

@@ -139,14 +139,14 @@ mod tests {

139 ];

140 let mut output = UnifiedParserOutput::default();

141 for chunk in chunks {

142 - parser.parse_into(chunk, &mut output).unwrap();

142 + parser.parse_into(DecodedText::unattributed(chunk), &mut output).unwrap();

143 }

144 output.append(parser.finish().unwrap());

145

146 assert_eq!(

147 output.events,

148 vec![

149 - UnifiedParserEvent::Reasoning("reasoning".to_string()),

149 + UnifiedParserEvent::Reasoning(DecodedText::unattributed("reasoning")),

150 UnifiedParserEvent::Text("answer".to_string()),

151 UnifiedParserEvent::ToolCall(crate::tool::ToolCallDelta {

152 tool_index: 0,

rust/src/parser/src/unified/inkling.rs

+87 -44

@@ -8,7 +8,7 @@ use winnow::prelude::*;

8 use winnow::stream::Partial;

9 use winnow::token::literal;

11 -use vllm_tokenizer::DynTokenizer;

11 +use vllm_tokenizer::{DecodedText, DynTokenizer};

13 use super::{Result, UnifiedParser, UnifiedParserOutput, token_id};

14 use crate::tool::json::{

@@ -56,15 +56,15 @@ type InklingInput<'i> = Partial<&'i str>;

57 #[derive(Debug, Clone, PartialEq, Eq)]

58 enum InklingEvent {

59 - Text { len: usize },

60 - Reasoning { len: usize },

59 + Text,

60 + Reasoning,

61 TextStart,

62 ReasoningStart,

63 MessageStart,

64 - Header { len: usize },

64 + Header,

65 ToolJsonStart,

66 ToolJsonHeader { name: String },

67 - ToolJsonArgs { len: usize, complete: bool },

67 + ToolJsonArgs { complete: bool },

68 BlockEnd,

69 }

@@ -73,7 +73,7 @@ enum InklingMode {

73 #[default]

74 Idle,

75 MessageHeader {

76 - pending: String,

76 + pending: DecodedText,

77 },

78 Text,

79 Reasoning,

@@ -86,7 +86,7 @@ enum InklingMode {

87 /// Unified parser for Inkling typed content blocks.

88 pub struct InklingUnifiedParser {

89 - buffer: String,

89 + buffer: DecodedText,

90 mode: InklingMode,

91 emitted_tool_count: usize,

92 active_tool_index: Option<usize>,

@@ -104,7 +104,7 @@ impl InklingUnifiedParser {

104 let content_thinking_token_id = token_id(tokenizer.as_ref(), CONTENT_THINKING)?;

105

106 Ok(Self {

107 - buffer: String::new(),

107 + buffer: DecodedText::default(),

108 mode: InklingMode::Idle,

109 emitted_tool_count: 0,

110 active_tool_index: None,

@@ -120,7 +120,7 @@ impl InklingUnifiedParser {

120 for token_id in prompt_token_ids.iter().rev().copied() {

121 if token_id == self.message_model_token_id {

122 self.mode = InklingMode::MessageHeader {

123 - pending: String::new(),

123 + pending: DecodedText::default(),

124 };

125 return;

126 }

@@ -138,24 +138,27 @@ impl InklingUnifiedParser {

138 }

139 }

140

141 - fn apply_event(&mut self, event: InklingEvent, output: &mut UnifiedParserOutput) -> Result<()> {

141 + fn apply_event(

142 + &mut self,

143 + event: InklingEvent,

144 + piece: DecodedText,

145 + output: &mut UnifiedParserOutput,

146 + ) -> Result<()> {

147 match event {

143 - InklingEvent::Text { len } => output.push_text(self.buffer[..len].to_string()),

144 - InklingEvent::Reasoning { len } => {

145 - output.push_reasoning(self.buffer[..len].to_string());

146 - }

148 + InklingEvent::Text => output.push_text(piece.text),

149 + InklingEvent::Reasoning => output.push_reasoning(piece),

150 InklingEvent::MessageStart => {

151 self.mode = InklingMode::MessageHeader {

149 - pending: String::new(),

152 + pending: DecodedText::default(),

153 };

154 }

152 - InklingEvent::Header { len } => {

155 + InklingEvent::Header => {

156 let InklingMode::MessageHeader { pending } = &mut self.mode else {

157 return Err(parsing_failed!(

158 "Inkling header text outside a message header"

159 ));

160 };

158 - pending.push_str(&self.buffer[..len]);

161 + pending.append(piece);

162 }

163 InklingEvent::TextStart => self.mode = InklingMode::Text,

164 InklingEvent::ReasoningStart => self.mode = InklingMode::Reasoning,

@@ -173,7 +176,7 @@ impl InklingUnifiedParser {

176 arguments: String::new(),

177 });

178 }

176 - InklingEvent::ToolJsonArgs { len, complete } => {

179 + InklingEvent::ToolJsonArgs { complete } => {

180 let Some(tool_index) = self.active_tool_index else {

181 return Err(parsing_failed!(

182 "Inkling arguments without an active tool call"

@@ -182,7 +185,7 @@ impl InklingUnifiedParser {

185 output.push_call(ToolCallDelta {

186 tool_index,

187 name: None,

185 - arguments: self.buffer[..len].to_string(),

188 + arguments: piece.text,

189 });

190 if complete {

191 self.mode = InklingMode::ToolJsonClose;

@@ -191,7 +194,7 @@ impl InklingUnifiedParser {

194 InklingEvent::BlockEnd => {

195 let mode = std::mem::take(&mut self.mode);

196 if let InklingMode::MessageHeader { pending } = mode {

194 - output.push_text(pending);

197 + output.push_text(pending.text);

198 }

199 self.active_tool_index = None;

200 }

@@ -201,12 +204,12 @@ impl InklingUnifiedParser {

204

205 fn reset(&mut self) -> String {

206 let mut uncommitted = match std::mem::take(&mut self.mode) {

204 - InklingMode::MessageHeader { pending } => pending,

207 + InklingMode::MessageHeader { pending } => pending.text,

208 _ => String::new(),

209 };

210 self.active_tool_index = None;

211 self.emitted_tool_count = 0;

209 - uncommitted.push_str(&std::mem::take(&mut self.buffer));

212 + uncommitted.push_str(&self.buffer.take().text);

213 uncommitted

214 }

215 }

@@ -231,14 +234,14 @@ impl UnifiedParser for InklingUnifiedParser {

234 true

235 }

236

234 - fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> {

235 - self.buffer.push_str(chunk);

237 + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> {

238 + self.buffer.append(delta);

239

237 - while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {

240 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer.text, |input| {

241 parse_next_inkling_event(input, &mut self.mode)

242 })? {

240 - self.apply_event(event, output)?;

241 - self.buffer.drain(..consumed_len);

243 + let piece = self.buffer.drain_prefix(consumed_len);

244 + self.apply_event(event, piece, output)?;

245 }

246

247 Ok(())

@@ -248,15 +251,12 @@ impl UnifiedParser for InklingUnifiedParser {

251 let mut output = UnifiedParserOutput::default();

252

253 match &mut self.mode {

251 - InklingMode::Idle | InklingMode::Text => {

252 - output.push_text(std::mem::take(&mut self.buffer))

253 - }

254 + InklingMode::Idle | InklingMode::Text => output.push_text(self.buffer.take().text),

255 InklingMode::MessageHeader { pending } => {

255 - pending.push_str(&self.buffer);

256 - self.buffer.clear();

257 - output.push_text(std::mem::take(pending));

256 + pending.append(self.buffer.take());

257 + output.push_text(std::mem::take(pending).text);

258 }

259 - InklingMode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)),

259 + InklingMode::Reasoning => output.push_reasoning(self.buffer.take()),

260 InklingMode::ToolJsonHeader

261 | InklingMode::ToolJsonArgs { .. }

262 | InklingMode::ToolJsonClose => {

@@ -367,22 +367,22 @@ fn block_end_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {

367

368 /// Parse safe text while waiting for the next Inkling marker.

369 fn safe_idle_text_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {

370 - safe_text_len_mul(input, IDLE_MARKERS).map(|len| InklingEvent::Text { len })

370 + safe_text_len_mul(input, IDLE_MARKERS).map(|_| InklingEvent::Text)

371 }

372

373 /// Parse safe header text before the next Inkling marker.

374 fn safe_header_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {

375 - safe_text_len_mul(input, IDLE_MARKERS).map(|len| InklingEvent::Header { len })

375 + safe_text_len_mul(input, IDLE_MARKERS).map(|_| InklingEvent::Header)

376 }

377

378 /// Parse safe text before the end of a Inkling text block.

379 fn safe_text_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {

380 - safe_text_len_mul(input, BLOCK_END_MARKERS).map(|len| InklingEvent::Text { len })

380 + safe_text_len_mul(input, BLOCK_END_MARKERS).map(|_| InklingEvent::Text)

381 }

382

383 /// Parse safe reasoning before the end of a Inkling reasoning block.

384 fn safe_reasoning_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {

385 - safe_text_len_mul(input, BLOCK_END_MARKERS).map(|len| InklingEvent::Reasoning { len })

385 + safe_text_len_mul(input, BLOCK_END_MARKERS).map(|_| InklingEvent::Reasoning)

386 }

387

388 /// Parse a Inkling JSON tool-call header.

@@ -400,9 +400,8 @@ fn parse_tool_json_args_event(

400 input: &mut JsonToolInput<'_>,

401 json_scan: &mut JsonObjectScanState,

402 ) -> ModalResult<InklingEvent> {

403 - let len = take_json_object(input, json_scan)?;

403 + take_json_object(input, json_scan)?;

404 Ok(InklingEvent::ToolJsonArgs {

405 - len,

405 complete: json_scan.complete(),

406 })

407 }

@@ -423,11 +422,11 @@ fn parse_tool_json_close_event(input: &mut InklingInput<'_>) -> ModalResult<Inkl

422 mod tests {

423 use std::sync::Arc;

424

426 - use super::{CONTENT_TEXT, CONTENT_THINKING, InklingUnifiedParser, MESSAGE_MODEL};

425 + use super::{CONTENT_TEXT, CONTENT_THINKING, END_MESSAGE, InklingUnifiedParser, MESSAGE_MODEL};

426 use crate::tool::Tool;

427 use crate::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};

428 use thiserror_ext::AsReport;

430 - use vllm_tokenizer::Tokenizer;

429 + use vllm_tokenizer::{DecodedText, TokenAnchor, TokenAttribution, Tokenizer};

430

431 struct FakeTokenizer;

432

@@ -497,7 +496,7 @@ mod tests {

496 impl<T: UnifiedParser + ?Sized> UnifiedParserTestExt for T {

497 fn parse_chunk(&mut self, chunk: &str) -> super::Result<UnifiedParserOutput> {

498 let mut output = UnifiedParserOutput::default();

500 - self.parse_into(chunk, &mut output)?;

499 + self.parse_into(DecodedText::unattributed(chunk), &mut output)?;

500 Ok(output)

501 }

502

@@ -529,7 +528,7 @@ mod tests {

528 self.events

529 .iter()

530 .filter_map(|event| match event {

532 - UnifiedParserEvent::Reasoning(text) => Some(text.as_str()),

531 + UnifiedParserEvent::Reasoning(text) => Some(text.text.as_str()),

532 _ => None,

533 })

534 .collect()

@@ -587,6 +586,50 @@ mod tests {

586 assert!(output.calls().is_empty());

587 }

588

589 + #[test]

590 + fn inkling_reasoning_events_conserve_token_attributions() {

591 + let chunk = |token_id: u32, text: &str| DecodedText {

592 + text: text.to_string(),

593 + attributions: [TokenAttribution {

594 + token_id,

595 + anchor: TokenAnchor::Visible { byte_offset: 0 },

596 + }]

597 + .into_iter()

598 + .collect(),

599 + };

600 +

601 + let mut parser = test_parser();

602 + let mut output = UnifiedParserOutput::default();

603 + for (token_id, text) in [

604 + (1, CONTENT_THINKING),

605 + (2, "reason"),

606 + (3, END_MESSAGE),

607 + (4, MESSAGE_MODEL),

608 + (5, CONTENT_TEXT),

609 + (6, "answer"),

610 + (7, END_MESSAGE),

611 + ] {

612 + parser.parse_into(chunk(token_id, text), &mut output).unwrap();

613 + }

614 + output.append(parser.finish().unwrap());

615 +

616 + // The reasoning tokens keep their attributions; marker tokens

617 + // (1, 3, 4, 5, 7) are dropped with their spans, and text pieces

618 + // carry no attributions.

619 + let reasoning_ids: Vec<u32> = output

620 + .events

621 + .iter()

622 + .filter_map(|event| match event {

623 + UnifiedParserEvent::Reasoning(piece) => Some(piece),

624 + _ => None,

625 + })

626 + .flat_map(|piece| piece.attributions.iter().map(|attr| attr.token_id))

627 + .collect();

628 + assert_eq!(reasoning_ids, [2]);

629 + assert_eq!(output.reasoning_text(), "reason");

630 + assert_eq!(output.normal_text(), "answer");

631 + }

632 +

633 #[test]

634 fn inkling_streaming_holds_split_markers() {

635 let output = collect_stream(&[

rust/src/parser/src/unified/kimi_k3.rs

+31 -28

@@ -41,7 +41,7 @@ mod structural_tag;

41 pub use structural_tag::KimiK3StructuralTagBuilder;

43 use serde_json::{Map, Value};

44 -use vllm_tokenizer::DynTokenizer;

44 +use vllm_tokenizer::{DecodedText, DynTokenizer};

45 use winnow::ascii::{multispace0 as ws0, multispace1 as ws1};

46 use winnow::combinator::{alt, delimited, eof, preceded, repeat, seq, terminated};

47 use winnow::error::{ContextError, ErrMode, ModalResult, StrContext};

@@ -95,12 +95,8 @@ type KimiK3Input<'i> = Partial<&'i str>;

96 #[derive(Debug, Clone, PartialEq, Eq)]

97 enum KimiK3Event {

98 - Text {

99 - len: usize,

100 - },

101 - Reasoning {

102 - len: usize,

103 - },

98 + Text,

99 + Reasoning,

100 /// Structural noise consumed without emitting anything.

101 Skip,

102 ThinkOpen,

@@ -147,7 +143,7 @@ enum KimiK3Mode {

143

144 /// Unified parser for Kimi K3 XTML think / response / tools channels.

145 pub struct KimiK3UnifiedParser {

150 - buffer: String,

146 + buffer: DecodedText,

147 mode: KimiK3Mode,

148 /// Number of calls emitted in the current response.

149 emitted_call_count: usize,

@@ -163,7 +159,7 @@ impl KimiK3UnifiedParser {

159 let sep_token_id = token_id(tokenizer.as_ref(), SEP)?;

160

161 Ok(Self {

166 - buffer: String::new(),

162 + buffer: DecodedText::default(),

163 mode: KimiK3Mode::default(),

164 emitted_call_count: 0,

165 tokenizer,

@@ -203,12 +199,16 @@ impl KimiK3UnifiedParser {

199 };

200 }

201

206 - fn apply_event(&mut self, event: KimiK3Event, output: &mut UnifiedParserOutput) -> Result<()> {

202 + fn apply_event(

203 + &mut self,

204 + event: KimiK3Event,

205 + piece: DecodedText,

206 + output: &mut UnifiedParserOutput,

207 + ) -> Result<()> {

208 match event {

208 - KimiK3Event::Text { len } => output.push_text(self.buffer[..len].to_string()),

209 - KimiK3Event::Reasoning { len } => {

210 - output.push_reasoning(self.buffer[..len].to_string());

211 - }

209 + KimiK3Event::Text => output.push_text(piece.text),

210 + KimiK3Event::Reasoning => output.push_reasoning(piece),

211 + // Marker and noise spans are drained and dropped with their tokens.

212 KimiK3Event::Skip => {}

213 KimiK3Event::ThinkOpen => self.mode = KimiK3Mode::Reasoning,

214 KimiK3Event::ThinkClose => self.mode = KimiK3Mode::Idle,

@@ -252,7 +252,7 @@ impl KimiK3UnifiedParser {

252 fn reset_state(&mut self) -> String {

253 self.mode = KimiK3Mode::Idle;

254 self.emitted_call_count = 0;

255 - std::mem::take(&mut self.buffer)

255 + self.buffer.take().text

256 }

257 }

258

@@ -279,14 +279,14 @@ impl UnifiedParser for KimiK3UnifiedParser {

279 Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER)

280 }

281

282 - fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> {

283 - self.buffer.push_str(chunk);

282 + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> {

283 + self.buffer.append(delta);

284

285 - while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {

285 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer.text, |input| {

286 parse_next_kimi_k3_event(input, &mut self.mode)

287 })? {

288 - self.apply_event(event, output)?;

289 - self.buffer.drain(..consumed_len);

288 + let piece = self.buffer.drain_prefix(consumed_len);

289 + self.apply_event(event, piece, output)?;

290 }

291

292 Ok(())

@@ -297,13 +297,15 @@ impl UnifiedParser for KimiK3UnifiedParser {

297

298 match &self.mode {

299 KimiK3Mode::Idle | KimiK3Mode::Response => {

300 - output.push_text(std::mem::take(&mut self.buffer));

300 + output.push_text(self.buffer.take().text);

301 + }

302 + KimiK3Mode::Reasoning => {

303 + output.push_reasoning(self.buffer.take());

304 }

302 - KimiK3Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)),

305 KimiK3Mode::Epilogue | KimiK3Mode::Done => self.buffer.clear(),

306 // A tools channel truncated between complete calls loses only its

307 // closing markers; keep the calls already emitted.

306 - KimiK3Mode::Tools if self.buffer.is_empty() => {}

308 + KimiK3Mode::Tools if self.buffer.text.is_empty() => {}

309 KimiK3Mode::Tools | KimiK3Mode::Call { .. } => {

310 return Err(parsing_failed!("incomplete Kimi K3 tool call"));

311 }

@@ -406,17 +408,17 @@ fn parse_done_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

408

409 /// Parse safe text while waiting for the next channel marker.

410 fn safe_idle_text_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

409 - safe_text_len_mul(input, IDLE_MARKERS).map(|len| KimiK3Event::Text { len })

411 + safe_text_len_mul(input, IDLE_MARKERS).map(|_| KimiK3Event::Text)

412 }

413

414 /// Parse safe reasoning before the think close marker.

415 fn safe_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

414 - safe_text_len_mul(input, REASONING_MARKERS).map(|len| KimiK3Event::Reasoning { len })

416 + safe_text_len_mul(input, REASONING_MARKERS).map(|_| KimiK3Event::Reasoning)

417 }

418

419 /// Parse safe response text before the next channel marker.

420 fn safe_response_text_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

419 - safe_text_len_mul(input, RESPONSE_MARKERS).map(|len| KimiK3Event::Text { len })

421 + safe_text_len_mul(input, RESPONSE_MARKERS).map(|_| KimiK3Event::Text)

422 }

423

424 /// Skip non-content noise after the response channel closed.

@@ -571,6 +573,7 @@ mod tests {

573

574 use serde_json::{Value, json};

575 use thiserror_ext::AsReport;

576 + use vllm_tokenizer::DecodedText;

577 use vllm_tokenizer::Tokenizer as _;

578 use vllm_tokenizer::test_utils::TestTokenizer;

579

@@ -604,7 +607,7 @@ mod tests {

607 impl UnifiedParserTestExt for KimiK3UnifiedParser {

608 fn parse_chunk(&mut self, chunk: &str) -> super::Result<UnifiedParserOutput> {

609 let mut output = UnifiedParserOutput::default();

607 - self.parse_into(chunk, &mut output)?;

610 + self.parse_into(DecodedText::unattributed(chunk), &mut output)?;

611 Ok(output)

612 }

613

@@ -636,7 +639,7 @@ mod tests {

639 self.events

640 .iter()

641 .filter_map(|event| match event {

639 - UnifiedParserEvent::Reasoning(text) => Some(text.as_str()),

642 + UnifiedParserEvent::Reasoning(text) => Some(text.text.as_str()),

643 _ => None,

644 })

645 .collect()

rust/src/parser/src/unified/mod.rs

+78 -13

@@ -16,7 +16,7 @@ pub use inkling::InklingUnifiedParser;

16 pub use kimi_k3::{KimiK3StructuralTagBuilder, KimiK3UnifiedParser};

17 use thiserror::Error;

18 use thiserror_ext::Macro;

19 -use vllm_tokenizer::DynTokenizer;

19 +use vllm_tokenizer::{DecodedText, DynTokenizer};

21 use crate::reasoning::ReasoningError;

22 use crate::tool::{

@@ -32,7 +32,11 @@ pub enum UnifiedParserEvent {

32 /// Normal assistant-visible text.

33 Text(String),

34 /// Reasoning text hidden from the normal content stream.

35 - Reasoning(String),

35 + ///

36 + /// Carries the attributions of the generated tokens that produced it, so

37 + /// downstream consumers can count reasoning tokens exactly. Marker spans

38 + /// are dropped by the parsers, keeping marker tokens out of any count.

39 + Reasoning(DecodedText),

40 /// A tool-call update extracted from visible assistant text.

41 ToolCall(ToolCallDelta),

42 }

@@ -58,15 +62,18 @@ impl UnifiedParserOutput {

62 }

64 /// Append one reasoning text event if `delta` is non-empty.

61 - pub fn push_reasoning(&mut self, delta: impl AsRef<str> + Into<String>) {

62 - if delta.as_ref().is_empty() {

65 + ///

66 + /// A piece with empty text but non-empty attributions (zero-width tokens

67 + /// only) is still kept: the tokens count as reasoning.

68 + pub fn push_reasoning(&mut self, delta: DecodedText) {

69 + if delta.is_empty() {

70 return;

71 }

72 if let Some(UnifiedParserEvent::Reasoning(last_text)) = self.events.last_mut() {

66 - last_text.push_str(delta.as_ref());

73 + last_text.append(delta);

74 return;

75 }

69 - self.events.push(UnifiedParserEvent::Reasoning(delta.into()));

76 + self.events.push(UnifiedParserEvent::Reasoning(delta));

77 }

79 /// Append one tool-call event.

@@ -98,6 +105,8 @@ impl UnifiedParserOutput {

105

106 #[cfg(test)]

107 mod tests {

108 + use vllm_tokenizer::{DecodedText, TokenAnchor, TokenAttribution};

109 +

110 use super::{UnifiedParserEvent, UnifiedParserOutput};

111 use crate::tool::ToolCallDelta;

112

@@ -107,8 +116,8 @@ mod tests {

116 output.push_text("hello");

117 output.push_text(" ");

118 output.push_text("world");

110 - output.push_reasoning("think");

111 - output.push_reasoning("ing");

119 + output.push_reasoning(DecodedText::unattributed("think"));

120 + output.push_reasoning(DecodedText::unattributed("ing"));

121 output.push_call(ToolCallDelta {

122 tool_index: 0,

123 name: Some("lookup".to_string()),

@@ -120,7 +129,7 @@ mod tests {

129 output.events,

130 vec![

131 UnifiedParserEvent::Text("hello world".to_string()),

123 - UnifiedParserEvent::Reasoning("thinking".to_string()),

132 + UnifiedParserEvent::Reasoning(DecodedText::unattributed("thinking")),

133 UnifiedParserEvent::ToolCall(ToolCallDelta {

134 tool_index: 0,

135 name: Some("lookup".to_string()),

@@ -139,11 +148,11 @@ mod tests {

148 let mut other = UnifiedParserOutput::default();

149 other.push_text(" ");

150 other.push_text("world");

142 - other.push_reasoning("think");

151 + other.push_reasoning(DecodedText::unattributed("think"));

152 output.append(other);

153

154 let mut after_reasoning = UnifiedParserOutput::default();

146 - after_reasoning.push_reasoning("ing");

155 + after_reasoning.push_reasoning(DecodedText::unattributed("ing"));

156 after_reasoning.push_text("!");

157 output.append(after_reasoning);

158

@@ -151,11 +160,67 @@ mod tests {

160 output.events,

161 vec![

162 UnifiedParserEvent::Text("hello world".to_string()),

154 - UnifiedParserEvent::Reasoning("thinking".to_string()),

163 + UnifiedParserEvent::Reasoning(DecodedText::unattributed("thinking")),

164 UnifiedParserEvent::Text("!".to_string()),

165 ]

166 );

167 }

168 +

169 + #[test]

170 + fn unified_parser_output_reasoning_coalescing_rebases_attributions() {

171 + let mut output = UnifiedParserOutput::default();

172 + output.push_reasoning(DecodedText {

173 + text: "think".to_string(),

174 + attributions: [TokenAttribution {

175 + token_id: 1,

176 + anchor: TokenAnchor::Visible { byte_offset: 0 },

177 + }]

178 + .into_iter()

179 + .collect(),

180 + });

181 + output.push_reasoning(DecodedText {

182 + text: "ing".to_string(),

183 + attributions: [TokenAttribution {

184 + token_id: 2,

185 + anchor: TokenAnchor::Visible { byte_offset: 0 },

186 + }]

187 + .into_iter()

188 + .collect(),

189 + });

190 + // A zero-width-token-only piece is kept: the tokens count as reasoning.

191 + output.push_reasoning(DecodedText {

192 + text: String::new(),

193 + attributions: [TokenAttribution {

194 + token_id: 3,

195 + anchor: TokenAnchor::ZeroWidth { byte_offset: 0 },

196 + }]

197 + .into_iter()

198 + .collect(),

199 + });

200 +

201 + assert_eq!(

202 + output.events,

203 + vec![UnifiedParserEvent::Reasoning(DecodedText {

204 + text: "thinking".to_string(),

205 + attributions: [

206 + TokenAttribution {

207 + token_id: 1,

208 + anchor: TokenAnchor::Visible { byte_offset: 0 },

209 + },

210 + TokenAttribution {

211 + token_id: 2,

212 + anchor: TokenAnchor::Visible { byte_offset: 5 },

213 + },

214 + TokenAttribution {

215 + token_id: 3,

216 + anchor: TokenAnchor::ZeroWidth { byte_offset: 8 },

217 + },

218 + ]

219 + .into_iter()

220 + .collect(),

221 + })]

222 + );

223 + }

224 }

225

226 /// Incremental parser that extracts reasoning and tool-call events from assistant output.

@@ -186,7 +251,7 @@ pub trait UnifiedParser: Send {

251 }

252

253 /// Feed one decoded text delta into the parser, appending committed output into `output`.

189 - fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()>;

254 + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()>;

255

256 /// Flush any buffered parser state at end of stream.

257 fn finish(&mut self) -> Result<UnifiedParserOutput> {

rust/src/tokenizer/src/incremental/attribution.rs

+55 -1

@@ -67,9 +67,63 @@ impl DecodedText {

67 }

68 }

70 + /// Return true when both the text and the attributions are empty.

71 + pub fn is_empty(&self) -> bool {

72 + self.text.is_empty() && self.attributions.is_empty()

73 + }

74 +

75 + /// Take the decoded text and attributions, leaving `self` empty.

76 + pub fn take(&mut self) -> Self {

77 + take(self)

78 + }

79 +

80 + /// Clear the text and attributions, keeping the allocations.

81 + pub fn clear(&mut self) {

82 + self.text.clear();

83 + self.attributions.clear();

84 + }

85 +

86 + /// Split off and return the prefix `[0, len)`; the suffix stays in `self`,

87 + /// with its token anchors rebased by `-len`.

88 + ///

89 + /// Anchors follow the first-byte ownership rule: a `Visible` token moves to

90 + /// the prefix when its first byte lies inside the prefix, so a token

91 + /// straddling the split belongs to the span holding its first byte. A

92 + /// `ZeroWidth` token moves to the prefix when its offset is `<= len`,

93 + /// matching the `take_ready` cutoff rule and letting

94 + /// `drain_prefix(text.len())` take trailing zero-width tokens. Relative

95 + /// order across anchor kinds is preserved.

96 + ///

97 + /// Panics if `len` is not a char boundary, via `String::drain`.

98 + #[must_use]

99 + pub fn drain_prefix(&mut self, len: usize) -> DecodedText {

100 + let text: String = self.text.drain(..len).collect();

101 + let len = offset_as_u32(len);

102 +

103 + let mut prefix = DecodedText {

104 + text,

105 + attributions: SmallVec::new(),

106 + };

107 + for attribution in take(&mut self.attributions) {

108 + let in_prefix = match attribution.anchor {

109 + TokenAnchor::Visible { byte_offset } => byte_offset < len,

110 + TokenAnchor::ZeroWidth { byte_offset } => byte_offset <= len,

111 + };

112 + if in_prefix {

113 + prefix.attributions.push(attribution);

114 + } else {

115 + self.attributions.push(TokenAttribution {

116 + token_id: attribution.token_id,

117 + anchor: attribution.anchor.offset_by(-i64::from(len)),

118 + });

119 + }

120 + }

121 + prefix

122 + }

123 +

124 /// Append another decoded fragment, rebasing its token anchors.

125 pub fn append(&mut self, other: Self) {

72 - if self.text.is_empty() && self.attributions.is_empty() {

126 + if self.is_empty() {

127 *self = other;

128 return;

129 }

rust/src/tokenizer/src/incremental/attribution/tests.rs

+139

@@ -169,6 +169,145 @@ fn decoded_text_append_preserves_zero_width_only_receiver() {

169 );

170 }

171

172 +#[test]

173 +fn decoded_text_drain_prefix_moves_anchors_before_boundary() {

174 + let mut dt = decoded("ab.cd", &[1, 2, 3], &[visible(0), visible(1), visible(3)]);

175 +

176 + let prefix = dt.drain_prefix(2);

177 +

178 + assert_eq!(prefix, decoded("ab", &[1, 2], &[visible(0), visible(1)]));

179 + assert_eq!(dt, decoded(".cd", &[3], &[visible(1)]));

180 +}

181 +

182 +#[test]

183 +fn decoded_text_drain_prefix_keeps_straddling_token_in_prefix() {

184 + // One token decodes to "ab"; a split inside it keeps the token in the

185 + // prefix, the span holding its first byte.

186 + let mut dt = decoded("ab", &[1], &[visible(0)]);

187 +

188 + let prefix = dt.drain_prefix(1);

189 +

190 + assert_eq!(prefix, decoded("a", &[1], &[visible(0)]));

191 + assert_eq!(dt, decoded("b", &[], &[]));

192 +}

193 +

194 +#[test]

195 +fn decoded_text_drain_prefix_visible_anchor_at_boundary_goes_to_suffix() {

196 + let mut dt = decoded("ab", &[1, 2], &[visible(0), visible(1)]);

197 +

198 + let prefix = dt.drain_prefix(1);

199 +

200 + assert_eq!(prefix, decoded("a", &[1], &[visible(0)]));

201 + assert_eq!(dt, decoded("b", &[2], &[visible(0)]));

202 +}

203 +

204 +#[test]

205 +fn decoded_text_drain_prefix_zero_width_at_boundary_goes_to_prefix() {

206 + let mut dt = decoded("ab", &[1, 2], &[zero_width(1), visible(1)]);

207 +

208 + let prefix = dt.drain_prefix(1);

209 +

210 + assert_eq!(prefix, decoded("a", &[1], &[zero_width(1)]));

211 + assert_eq!(dt, decoded("b", &[2], &[visible(0)]));

212 +}

213 +

214 +#[test]

215 +fn decoded_text_drain_prefix_shared_anchor_group_moves_together() {

216 + // A byte-fallback group jointly decoding one character shares one anchor;

217 + // the whole group stays on whichever side holds that first byte.

218 + let mut dt = decoded(

219 + "a你",

220 + &[1, 2, 3, 4],

221 + &[visible(0), visible(1), visible(1), visible(1)],

222 + );

223 +

224 + let prefix = dt.drain_prefix(1);

225 +

226 + assert_eq!(prefix, decoded("a", &[1], &[visible(0)]));

227 + assert_eq!(

228 + dt,

229 + decoded("你", &[2, 3, 4], &[visible(0), visible(0), visible(0)])

230 + );

231 +

232 + let prefix = dt.drain_prefix(3);

233 +

234 + assert_eq!(

235 + prefix,

236 + decoded("你", &[2, 3, 4], &[visible(0), visible(0), visible(0)])

237 + );

238 + assert_eq!(dt, decoded("", &[], &[]));

239 +}

240 +

241 +#[test]

242 +fn decoded_text_drain_prefix_zero_drains_leading_zero_width() {

243 + let mut dt = decoded("ab", &[1, 2], &[zero_width(0), visible(0)]);

244 +

245 + let prefix = dt.drain_prefix(0);

246 +

247 + assert_eq!(prefix, decoded("", &[1], &[zero_width(0)]));

248 + assert_eq!(dt, decoded("ab", &[2], &[visible(0)]));

249 +}

250 +

251 +#[test]

252 +fn decoded_text_drain_prefix_full_length_takes_trailing_zero_width() {

253 + let mut dt = decoded("ab", &[1, 2, 3], &[visible(0), visible(1), zero_width(2)]);

254 +

255 + let prefix = dt.drain_prefix(2);

256 +

257 + assert_eq!(

258 + prefix,

259 + decoded("ab", &[1, 2, 3], &[visible(0), visible(1), zero_width(2)])

260 + );

261 + assert_eq!(dt, decoded("", &[], &[]));

262 +}

263 +

264 +#[test]

265 +fn decoded_text_drain_prefix_preserves_order_across_anchor_kinds() {

266 + let mut dt = decoded(

267 + "ab",

268 + &[1, 2, 3, 4],

269 + &[visible(0), zero_width(1), visible(1), zero_width(2)],

270 + );

271 +

272 + let prefix = dt.drain_prefix(1);

273 +

274 + assert_eq!(prefix, decoded("a", &[1, 2], &[visible(0), zero_width(1)]));

275 + assert_eq!(dt, decoded("b", &[3, 4], &[visible(0), zero_width(1)]));

276 +}

277 +

278 +#[test]

279 +fn decoded_text_drain_prefix_then_append_reconstructs_original() {

280 + // ASCII, a byte-fallback group, and mid and trailing zero-width tokens.

281 + let original = decoded(

282 + "a你b",

283 + &[1, 2, 3, 4, 5, 6, 7],

284 + &[

285 + visible(0),

286 + visible(1),

287 + visible(1),

288 + visible(1),

289 + zero_width(4),

290 + visible(4),

291 + zero_width(5),

292 + ],

293 + );

294 +

295 + for len in [0, 1, 4, 5] {

296 + let mut rest = original.clone();

297 + let mut prefix = rest.drain_prefix(len);

298 + prefix.append(rest);

299 + assert_eq!(prefix, original, "split at byte {len}");

300 + }

301 +}

302 +

303 +#[test]

304 +#[should_panic(expected = "is_char_boundary")]

305 +fn decoded_text_drain_prefix_panics_off_char_boundary() {

306 + let mut dt = decoded("a你", &[1, 2], &[visible(0), visible(1)]);

307 +

308 + let _ = dt.drain_prefix(2);

309 +}

310 +

311 struct AttributionCase<'a> {

312 tokenizer: &'a dyn Tokenizer,

313 prompt_token_ids: Vec<u32>,