@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 629a938a92

[Frontend] Preserve bare Inkling text in Python and Rust parsers (#50403)

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

Bugen Zhao committed Jul 31, 2026 at 03:25 UTC 629a938a928476a167efc6428bf2fbd43b1dc888

4 files changed +144 -18

rust/src/parser/src/unified/inkling.rs

+85 -16

@@ -61,7 +61,7 @@ enum InklingEvent {

61 TextStart,

62 ReasoningStart,

63 MessageStart,

64 - Header,

64 + Header { len: usize },

65 ToolJsonStart,

66 ToolJsonHeader { name: String },

67 ToolJsonArgs { len: usize, complete: bool },

@@ -72,7 +72,9 @@ enum InklingEvent {

72 enum InklingMode {

73 #[default]

74 Idle,

75 - MessageHeader,

75 + MessageHeader {

76 + pending: String,

77 + },

78 Text,

79 Reasoning,

80 ToolJsonHeader,

@@ -117,7 +119,9 @@ impl InklingUnifiedParser {

119 self.mode = InklingMode::Idle;

120 for token_id in prompt_token_ids.iter().rev().copied() {

121 if token_id == self.message_model_token_id {

120 - self.mode = InklingMode::MessageHeader;

122 + self.mode = InklingMode::MessageHeader {

123 + pending: String::new(),

124 + };

125 return;

126 }

127 if token_id == self.content_thinking_token_id {

@@ -140,8 +144,19 @@ impl InklingUnifiedParser {

144 InklingEvent::Reasoning { len } => {

145 output.push_reasoning(self.buffer[..len].to_string());

146 }

143 - InklingEvent::MessageStart => self.mode = InklingMode::MessageHeader,

144 - InklingEvent::Header => {}

147 + InklingEvent::MessageStart => {

148 + self.mode = InklingMode::MessageHeader {

149 + pending: String::new(),

150 + };

151 + }

152 + InklingEvent::Header { len } => {

153 + let InklingMode::MessageHeader { pending } = &mut self.mode else {

154 + return Err(parsing_failed!(

155 + "Inkling header text outside a message header"

156 + ));

157 + };

158 + pending.push_str(&self.buffer[..len]);

159 + }

160 InklingEvent::TextStart => self.mode = InklingMode::Text,

161 InklingEvent::ReasoningStart => self.mode = InklingMode::Reasoning,

162 InklingEvent::ToolJsonStart => self.mode = InklingMode::ToolJsonHeader,

@@ -174,7 +189,10 @@ impl InklingUnifiedParser {

189 }

190 }

191 InklingEvent::BlockEnd => {

177 - self.mode = InklingMode::Idle;

192 + let mode = std::mem::take(&mut self.mode);

193 + if let InklingMode::MessageHeader { pending } = mode {

194 + output.push_text(pending);

195 + }

196 self.active_tool_index = None;

197 }

198 }

@@ -182,10 +200,14 @@ impl InklingUnifiedParser {

200 }

201

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

185 - self.mode = InklingMode::Idle;

203 + let mut uncommitted = match std::mem::take(&mut self.mode) {

204 + InklingMode::MessageHeader { pending } => pending,

205 + _ => String::new(),

206 + };

207 self.active_tool_index = None;

208 self.emitted_tool_count = 0;

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

209 + uncommitted.push_str(&std::mem::take(&mut self.buffer));

210 + uncommitted

211 }

212 }

213

@@ -225,11 +247,15 @@ impl UnifiedParser for InklingUnifiedParser {

247 fn finish(&mut self) -> Result<UnifiedParserOutput> {

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

249

228 - match &self.mode {

250 + match &mut self.mode {

251 InklingMode::Idle | InklingMode::Text => {

252 output.push_text(std::mem::take(&mut self.buffer))

253 }

232 - InklingMode::MessageHeader => self.buffer.clear(),

254 + InklingMode::MessageHeader { pending } => {

255 + pending.push_str(&self.buffer);

256 + self.buffer.clear();

257 + output.push_text(std::mem::take(pending));

258 + }

259 InklingMode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)),

260 InklingMode::ToolJsonHeader

261 | InklingMode::ToolJsonArgs { .. }

@@ -254,7 +280,7 @@ fn parse_next_inkling_event(

280 ) -> ModalResult<InklingEvent> {

281 match mode {

282 InklingMode::Idle => parse_idle_event(input),

257 - InklingMode::MessageHeader => parse_message_header_event(input),

283 + InklingMode::MessageHeader { .. } => parse_message_header_event(input),

284 InklingMode::Text => parse_text_event(input),

285 InklingMode::Reasoning => parse_reasoning_event(input),

286 InklingMode::ToolJsonHeader => parse_tool_json_header_event(input),

@@ -346,7 +372,7 @@ fn safe_idle_text_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEven

372

373 /// Parse safe header text before the next Inkling marker.

374 fn safe_header_event(input: &mut InklingInput<'_>) -> ModalResult<InklingEvent> {

349 - safe_text_len_mul(input, IDLE_MARKERS).map(|_| InklingEvent::Header)

375 + safe_text_len_mul(input, IDLE_MARKERS).map(|len| InklingEvent::Header { len })

376 }

377

378 /// Parse safe text before the end of a Inkling text block.

@@ -672,15 +698,58 @@ mod tests {

698 let mut parser = test_parser();

699 parser.initialize(&[200001]).unwrap();

700

701 + let mut output = parser.parse_chunk("get_").unwrap();

702 + output.append(

703 + parser

704 + .parse_complete(concat!(

705 + "weather<|content_invoke_tool_json|>",

706 + "{\"name\":\"get_weather\",\"args\":{}}<|end_message|>"

707 + ))

708 + .unwrap(),

709 + );

710 +

711 + assert!(output.normal_text().is_empty());

712 + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));

713 + }

714 +

715 + #[test]

716 + fn inkling_initialize_model_opener_flushes_bare_text_at_finish() {

717 + let mut parser = test_parser();

718 + parser.initialize(&[200001]).unwrap();

719 +

720 + let mut output = parser.parse_chunk("plain ").unwrap();

721 + output.append(parser.parse_chunk("answer").unwrap());

722 + assert!(output.normal_text().is_empty());

723 +

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

725 +

726 + assert_eq!(output.normal_text(), "plain answer");

727 + }

728 +

729 + #[test]

730 + fn inkling_initialize_model_opener_flushes_bare_text_before_end_marker() {

731 + let mut parser = test_parser();

732 + parser.initialize(&[200001]).unwrap();

733 +

734 let output = parser

735 .parse_complete(concat!(

677 - "get_weather<|content_invoke_tool_json|>",

678 - "{\"name\":\"get_weather\",\"args\":{}}<|end_message|>"

736 + "plain answer",

737 + "<|end_message|>",

738 + "<|content_model_end_sampling|>"

739 ))

740 .unwrap();

741

682 - assert!(output.normal_text().is_empty());

683 - assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));

742 + assert_eq!(output.normal_text(), "plain answer");

743 + }

744 +

745 + #[test]

746 + fn inkling_reset_returns_pending_message_header() {

747 + let mut parser = test_parser();

748 + parser.initialize(&[200001]).unwrap();

749 + parser.parse_chunk("plain ").unwrap();

750 + parser.parse_chunk("answer").unwrap();

751 +

752 + assert_eq!(parser.reset(), "plain answer");

753 }

754

755 #[test]

tests/parser/engine/test_inkling.py

+29

@@ -178,6 +178,13 @@ class TestArgConverter:

178

179

180 class TestNonStreaming:

181 + @pytest.mark.parametrize("suffix", ["", END_MESSAGE, END_SAMPLING])

182 + def test_bare_text_after_model_opener(self, parser, mock_request, suffix):

183 + reasoning, content, tools = parser.parse(f"hello world{suffix}", mock_request)

184 + assert reasoning is None

185 + assert content == "hello world"

186 + assert tools is None

187 +

188 def test_plain_text(self, parser, mock_request):

189 reasoning, content, tools = parser.parse(

190 f"{TEXT_START}hello world{END_MESSAGE}", mock_request

@@ -426,6 +433,28 @@ class TestPromptSeededState:

433 assert delta.content is None

434 assert delta.tool_calls[0].function.name == "get_weather"

435

436 + def test_generation_prompt_header_flushes_bare_text_at_finish(

437 + self, parser, mock_request

438 + ):

439 + prompt_token_ids = [_TML_VOCAB[END_MESSAGE], _TML_VOCAB[MSG_MODEL]]

440 + first = parser.parse_delta(

441 + "plain ",

442 + [ord(char) for char in "plain "],

443 + mock_request,

444 + prompt_token_ids=prompt_token_ids,

445 + finished=False,

446 + )

447 + assert first is None

448 +

449 + second = parser.parse_delta(

450 + "answer",

451 + [ord(char) for char in "answer"],

452 + mock_request,

453 + finished=True,

454 + )

455 + assert second is not None

456 + assert second.content == "plain answer"

457 +

458

459 class TestToolCallFiltering:

460 """Inkling equivalents of the generic tool-call-filtering replay tests

vllm/parser/engine/streaming_parser_engine.py

+29 -1

@@ -185,6 +185,7 @@ class StreamingParserEngine:

185 # implicit-reasoning-end (content returns None).

186 self._scanner.reset()

187 self._lexer.reset()

188 + self._message_header_buffer = ""

189 self._reset_args_state()

190

191 def feed(

@@ -270,6 +271,15 @@ class StreamingParserEngine:

271 )

272 self.state = ParserState.CONTENT

273 elif self.state == ParserState.MESSAGE_HEADER:

274 + if self._message_header_buffer:

275 + events.append(

276 + SemanticEvent(

277 + EventType.TEXT_CHUNK,

278 + value=self._message_header_buffer,

279 + tool_index=self.tool_index,

280 + )

281 + )

282 + self._message_header_buffer = ""

283 self.state = ParserState.CONTENT

284

285 return events

@@ -311,6 +321,7 @@ class StreamingParserEngine:

321 if self.skip_tool_parsing and terminal in self._tool_terminals:

322 if self.state == ParserState.MESSAGE_HEADER:

323 self.state = ParserState.CONTENT

324 + self._message_header_buffer = ""

325 return [

326 SemanticEvent(

327 EventType.TEXT_CHUNK,

@@ -345,6 +356,9 @@ class StreamingParserEngine:

356 return self._apply_transition(transition, value)

357

358 def _emit_for_state(self, text: str) -> list[SemanticEvent]:

359 + if self.state == ParserState.MESSAGE_HEADER:

360 + self._message_header_buffer += text

361 + return []

362 if self.state == ParserState.TOOL_ARGS:

363 if self.config.tool_args_json:

364 return self._feed_args_text(text)

@@ -371,6 +385,8 @@ class StreamingParserEngine:

385 value: str,

386 ) -> list[SemanticEvent]:

387 events: list[SemanticEvent] = []

388 + previous_state = self.state

389 + message_header = ""

390

391 if (

392 self.state == ParserState.TOOL_ARGS

@@ -386,15 +402,27 @@ class StreamingParserEngine:

402 )

403 self._args_buffer = ""

404

405 + if previous_state == ParserState.MESSAGE_HEADER:

406 + message_header = self._message_header_buffer

407 + self._message_header_buffer = ""

408 +

409 self.state = transition.next_state

410

411 for event_type in transition.events:

412 if event_type == EventType.TOOL_CALL_START:

413 self.tool_index += 1

414 + event_value = (

415 + message_header

416 + if previous_state == ParserState.MESSAGE_HEADER

417 + and event_type == EventType.TEXT_CHUNK

418 + else value

419 + )

420 + if event_type == EventType.TEXT_CHUNK and not event_value:

421 + continue

422 events.append(

423 SemanticEvent(

424 event_type,

397 - value=value,

425 + value=event_value,

426 tool_index=self.tool_index,

427 )

428 )

vllm/parser/inkling.py

+1 -1

@@ -268,7 +268,7 @@ def inkling_config() -> ParserEngineConfig:

268 )

269 transitions[(ParserState.MESSAGE_HEADER, end)] = Transition(

270 ParserState.CONTENT,

271 - (),

271 + (EventType.TEXT_CHUNK,),

272 )

273

274 return ParserEngineConfig(