@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 36229e1337

[Rust Frontend] Token-aware marker parsing for Kimi K3 (#58358)

Signed-off-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Bugen Zhao committed Oct 1, 2026 at 12:21 UTC 36229e1337ae0e44e60075a11ebfde210c11d706

8 files changed +1135 -133

rust/Cargo.lock

@@ -5964,8 +5964,10 @@ dependencies = [

5964 name = "vllm-parser"

5965 version = "0.1.0"

5966 dependencies = [

5967 + "auto_impl",

5968 "criterion",

5969 "easy-ext",

5970 + "educe",

5971 "expect-test",

5972 "futures",

5973 "openai-protocol",

rust/Cargo.toml

@@ -31,6 +31,7 @@ async-openai = { version = "0.33.1", default-features = false, features = ["nati

31 async-trait = "0.1.89"

32 asynk-strim-attr = "0.1.0"

33 auto_enums = { version = "0.8.9", features = ["tokio1"] }

34 +auto_impl = "1.3.0"

35 axum = "0.8.8"

36 base64 = "0.22.1"

37 bytemuck = { version = "1.25.0", features = ["extern_crate_alloc"] }

rust/src/chat/src/parser/unified.rs

+2 -1

@@ -141,7 +141,8 @@ mod tests {

141 let tokenizer = TestTokenizer::new()

142 .with_regular_token("<|open|>", 1001)

143 .with_regular_token("<|close|>", 1002)

144 - .with_regular_token("<|sep|>", 1003);

144 + .with_regular_token("<|sep|>", 1003)

145 + .with_regular_token("<|end_of_msg|>", 1004);

146

147 assert!(factory.contains(names::KIMI_K3));

148 assert_eq!(

rust/src/parser/Cargo.toml

@@ -8,7 +8,9 @@ license.workspace = true

8 test-util = []

10 [dependencies]

11 +auto_impl.workspace = true

12 easy-ext.workspace = true

13 +educe.workspace = true

14 serde.workspace = true

15 serde_json.workspace = true

16 serde_with.workspace = true

rust/src/parser/src/unified/kimi_k3.rs

+340 -54

@@ -26,6 +26,14 @@

26 //! starts *inside* that channel without re-emitting the open tag;

27 //! [`UnifiedParser::initialize`] detects this from the prompt token IDs.

28 //!

29 +//! Channel markers are structural only when their `<|open|>` / `<|close|>` /

30 +//! `<|sep|>` bytes come from the dedicated special tokens. The model can spell

31 +//! the same text with ordinary BPE tokens (for example while quoting a

32 +//! configuration snippet in its reasoning); such text stays content in the

33 +//! current channel. The parser therefore expects token-attributed input from

34 +//! the incremental detokenizer: fed unattributed text, it treats every marker

35 +//! as content. Channel names such as `think` are ordinary text either way.

36 +//!

37 //! Argument decoding mirrors the renderer's type tagging (inverse encoding):

38 //! `type="string"` values pass the raw text through, other types are

39 //! JSON-decoded, and a raw `json` block is passed through unmodified. Attribute

@@ -34,29 +42,33 @@

42 //! Known limitation (shared with the Python parser): string argument and

43 //! response bodies are emitted raw, so a value that literally contains

44 //! `<|close|>argument<|sep|>` or `<|close|>response<|sep|>` is

37 -//! indistinguishable from a real closing marker.

45 +//! indistinguishable from a real closing marker. With token attribution this

46 +//! now only applies to `argument` and `json` blocks inside a call body, which

47 +//! are still parsed from text.

49 mod structural_tag;

51 pub use structural_tag::KimiK3StructuralTagBuilder;

53 use serde_json::{Map, Value};

44 -use vllm_tokenizer::{DecodedText, DynTokenizer};

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

55 use winnow::ascii::{multispace0 as ws0, multispace1 as ws1};

56 use winnow::combinator::{alt, delimited, eof, preceded, repeat, seq, terminated};

57 use winnow::error::{ContextError, ErrMode, ModalResult, StrContext};

58 use winnow::prelude::*;

49 -use winnow::stream::Partial;

59 use winnow::token::{literal, rest, take_till, take_until, take_while};

61 use self::structural_tag::KIMI_K3_STRUCTURAL_TAG_BUILDER;

53 -use super::{Result, UnifiedParser, UnifiedParserOutput, token_id};

62 +use super::{Result, UnifiedParser, UnifiedParserOutput, special_token};

63 use crate::output_grammar::{

64 self, BuiltOutputGrammar, OutputGrammarContext, visible_format_from_builder,

65 };

66 use crate::tool::{Tool, ToolCallDelta};

67 use crate::unified::parsing_failed;

59 -use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker};

68 +use crate::utils::{

69 + Attributed, AttributedExt as _, Marker, MarkerScanState, SpecialToken,

70 + parse_buffered_event_attributed, safe_text_len_mul, take_until_marker,

71 +};

73 const OPEN: &str = "<|open|>";

74 const CLOSE: &str = "<|close|>";

@@ -70,31 +82,82 @@ const RESPONSE_CLOSE: &str = "<|close|>response<|sep|>";

82 const TOOLS_OPEN: &str = "<|open|>tools<|sep|>";

83 const TOOLS_CLOSE: &str = "<|close|>tools<|sep|>";

84 const MESSAGE_CLOSE: &str = "<|close|>message<|sep|>";

73 -const CALL_OPEN: &str = "<|open|>call";

85 const CALL_CLOSE: &str = "<|close|>call<|sep|>";

86 const ARG_OPEN: &str = "<|open|>argument";

87 const ARG_CLOSE: &str = "<|close|>argument<|sep|>";

88 const JSON_OPEN: &str = "<|open|>json";

89 const JSON_CLOSE: &str = "<|close|>json<|sep|>";

80 -const IDLE_MARKERS: &[&str] = &[

81 - THINK_OPEN,

82 - RESPONSE_OPEN,

83 - TOOLS_OPEN,

84 - MESSAGE_CLOSE,

85 - END_OF_MSG,

86 -];

87 -const REASONING_MARKERS: &[&str] = &[THINK_CLOSE, END_OF_MSG];

88 -const RESPONSE_MARKERS: &[&str] = &[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG];

89 -const EPILOGUE_MARKERS: &[&str] = &[TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG];

90 -const TOOLS_MARKERS: &[&str] = &[CALL_OPEN, TOOLS_CLOSE, MESSAGE_CLOSE, END_OF_MSG];

91 -

91 /// Channel tags are a couple of text tokens; longer `<|open|>…<|sep|>` spans in

92 /// the prompt tail (attribute-bearing message opens, message bodies) never name

93 /// a generation channel.

94 const MAX_PREFILL_TAG_TOKENS: usize = 8;

97 -type KimiK3Input<'i> = Partial<&'i str>;

96 +type KimiK3Input<'i> = Attributed<'i, Markers>;

97 +

98 +/// The K3 structural special tokens, resolved once from the active tokenizer.

99 +struct SpecialTokens {

100 + open: SpecialToken,

101 + close: SpecialToken,

102 + sep: SpecialToken,

103 + end_of_msg: SpecialToken,

104 +}

105 +

106 +impl SpecialTokens {

107 + fn new(tokenizer: &dyn Tokenizer) -> Result<Self> {

108 + Ok(Self {

109 + open: special_token(tokenizer, OPEN)?,

110 + close: special_token(tokenizer, CLOSE)?,

111 + sep: special_token(tokenizer, SEP)?,

112 + end_of_msg: special_token(tokenizer, END_OF_MSG)?,

113 + })

114 + }

115 +}

116 +

117 +/// The K3 channel markers, with their special-token segments guarded by the

118 +/// active tokenizer's IDs.

119 +///

120 +/// One definition serves both the `alt` branches that consume a marker and the

121 +/// safe-text scans that stop in front of one, so the two can never disagree on

122 +/// whether a spelling is structural.

123 +#[derive(Debug)]

124 +struct Markers {

125 + think_open: Marker,

126 + think_close: Marker,

127 + response_open: Marker,

128 + response_close: Marker,

129 + tools_open: Marker,

130 + tools_close: Marker,

131 + message_close: Marker,

132 + end_of_msg: Marker,

133 + /// `<|open|>call`; the tag attributes and the terminating `<|sep|>` follow.

134 + call_open: Marker,

135 + call_close: Marker,

136 + sep: Marker,

137 +}

138 +

139 +impl Markers {

140 + fn new(tokens: &SpecialTokens) -> Self {

141 + let open =

142 + |tag: &str| Marker::special(&tokens.open).then_text(tag).then_special(&tokens.sep);

143 + let close =

144 + |tag: &str| Marker::special(&tokens.close).then_text(tag).then_special(&tokens.sep);

145 +

146 + Self {

147 + think_open: open("think"),

148 + think_close: close("think"),

149 + response_open: open("response"),

150 + response_close: close("response"),

151 + tools_open: open("tools"),

152 + tools_close: close("tools"),

153 + message_close: close("message"),

154 + end_of_msg: Marker::special(&tokens.end_of_msg),

155 + call_open: Marker::special(&tokens.open).then_text("call"),

156 + call_close: close("call"),

157 + sep: Marker::special(&tokens.sep),

158 + }

159 + }

160 +}

161

162 #[derive(Debug, Clone, PartialEq, Eq)]

163 enum KimiK3Event {

@@ -148,26 +211,25 @@ enum KimiK3Mode {

211 pub struct KimiK3UnifiedParser {

212 buffer: DecodedText,

213 mode: KimiK3Mode,

214 + markers: Markers,

215 /// Number of calls emitted in the current response.

216 emitted_call_count: usize,

217 tokenizer: DynTokenizer,

154 - open_token_id: u32,

155 - sep_token_id: u32,

218 + special_tokens: SpecialTokens,

219 }

220

221 impl KimiK3UnifiedParser {

159 - /// Create a Kimi K3 parser.

222 + /// Create a Kimi K3 parser expecting token-attributed input.

223 pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result<Self> {

161 - let open_token_id = token_id(tokenizer.as_ref(), OPEN)?;

162 - let sep_token_id = token_id(tokenizer.as_ref(), SEP)?;

224 + let special_tokens = SpecialTokens::new(tokenizer.as_ref())?;

225

226 Ok(Self {

227 buffer: DecodedText::default(),

228 mode: KimiK3Mode::default(),

229 + markers: Markers::new(&special_tokens),

230 emitted_call_count: 0,

231 tokenizer,

169 - open_token_id,

170 - sep_token_id,

232 + special_tokens,

233 })

234 }

235

@@ -180,11 +242,14 @@ impl KimiK3UnifiedParser {

242 fn initialize_mode(&mut self, prompt_token_ids: &[u32]) {

243 self.mode = KimiK3Mode::Idle;

244

183 - let Some(sep_pos) = prompt_token_ids.iter().rposition(|&id| id == self.sep_token_id) else {

245 + let Some(sep_pos) =

246 + prompt_token_ids.iter().rposition(|&id| id == self.special_tokens.sep.id)

247 + else {

248 return;

249 };

186 - let Some(open_pos) =

187 - prompt_token_ids[..sep_pos].iter().rposition(|&id| id == self.open_token_id)

250 + let Some(open_pos) = prompt_token_ids[..sep_pos]

251 + .iter()

252 + .rposition(|&id| id == self.special_tokens.open.id)

253 else {

254 return;

255 };

@@ -291,9 +356,11 @@ impl UnifiedParser for KimiK3UnifiedParser {

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

357 self.buffer.append(delta);

358

294 - while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer.text, |input| {

295 - parse_next_kimi_k3_event(input, &mut self.mode)

296 - })? {

359 + while let Some((event, consumed_len)) =

360 + parse_buffered_event_attributed(&self.buffer, &self.markers, |input| {

361 + parse_next_kimi_k3_event(input, &mut self.mode)

362 + })?

363 + {

364 let piece = self.buffer.drain_prefix(consumed_len);

365 self.apply_event(event, piece, output)?;

366 }

@@ -347,10 +414,11 @@ fn parse_next_kimi_k3_event(

414

415 /// Parse an event while waiting for the next channel open.

416 fn parse_idle_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

417 + let markers = input.markers();

418 alt((

351 - literal(THINK_OPEN).value(KimiK3Event::ThinkOpen),

352 - literal(RESPONSE_OPEN).value(KimiK3Event::ResponseOpen),

353 - literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen),

419 + markers.think_open.value(KimiK3Event::ThinkOpen),

420 + markers.response_open.value(KimiK3Event::ResponseOpen),

421 + markers.tools_open.value(KimiK3Event::ToolsOpen),

422 message_end_event,

423 safe_idle_text_event,

424 ))

@@ -359,11 +427,12 @@ fn parse_idle_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

427

428 /// Parse an event inside the `think` channel.

429 fn parse_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

430 + let markers = input.markers();

431 alt((

363 - literal(THINK_CLOSE).value(KimiK3Event::ThinkClose),

432 + markers.think_close.value(KimiK3Event::ThinkClose),

433 // `<|end_of_msg|>` can reach the parser under `ignore_eos` or

434 // `include_stop_str_in_output`; never leak it into reasoning.

366 - literal(END_OF_MSG).value(KimiK3Event::MessageEnd),

435 + markers.end_of_msg.value(KimiK3Event::MessageEnd),

436 safe_reasoning_event,

437 ))

438 .parse_next(input)

@@ -371,10 +440,11 @@ fn parse_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event

440

441 /// Parse an event inside the `response` channel.

442 fn parse_response_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

443 + let markers = input.markers();

444 alt((

375 - literal(RESPONSE_CLOSE).value(KimiK3Event::ResponseClose),

445 + markers.response_close.value(KimiK3Event::ResponseClose),

446 // The response body also implicitly ends at a `tools` channel.

377 - literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen),

447 + markers.tools_open.value(KimiK3Event::ToolsOpen),

448 message_end_event,

449 safe_response_text_event,

450 ))

@@ -383,8 +453,9 @@ fn parse_response_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event>

453

454 /// Parse an event after the response channel closed.

455 fn parse_epilogue_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

456 + let markers = input.markers();

457 alt((

387 - literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen),

458 + markers.tools_open.value(KimiK3Event::ToolsOpen),

459 message_end_event,

460 skip_epilogue_noise_event,

461 ))

@@ -393,9 +464,10 @@ fn parse_epilogue_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event>

464

465 /// Parse an event inside the `tools` channel, between `call` blocks.

466 fn parse_tools_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

467 + let markers = input.markers();

468 alt((

469 call_open_event,

398 - literal(TOOLS_CLOSE).value(KimiK3Event::ToolsClose),

470 + markers.tools_close.value(KimiK3Event::ToolsClose),

471 // Defensive: an unterminated tools channel still ends with the message.

472 message_end_event,

473 skip_tools_noise_event,

@@ -405,7 +477,8 @@ fn parse_tools_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

477

478 /// Parse a message close or end-of-message marker.

479 fn message_end_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

408 - alt((literal(MESSAGE_CLOSE), literal(END_OF_MSG)))

480 + let markers = input.markers();

481 + alt((&markers.message_close, &markers.end_of_msg))

482 .value(KimiK3Event::MessageEnd)

483 .parse_next(input)

484 }

@@ -417,35 +490,82 @@ fn parse_done_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

490

491 /// Parse safe text while waiting for the next channel marker.

492 fn safe_idle_text_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

420 - safe_text_len_mul(input, IDLE_MARKERS).map(|_| KimiK3Event::Text)

493 + let markers = input.markers();

494 + safe_text_len_mul(

495 + input,

496 + &[

497 + &markers.think_open,

498 + &markers.response_open,

499 + &markers.tools_open,

500 + &markers.message_close,

501 + &markers.end_of_msg,

502 + ],

503 + )

504 + .map(|_| KimiK3Event::Text)

505 }

506

507 /// Parse safe reasoning before the think close marker.

508 fn safe_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

425 - safe_text_len_mul(input, REASONING_MARKERS).map(|_| KimiK3Event::Reasoning)

509 + let markers = input.markers();

510 + safe_text_len_mul(input, &[&markers.think_close, &markers.end_of_msg])

511 + .map(|_| KimiK3Event::Reasoning)

512 }

513

514 /// Parse safe response text before the next channel marker.

515 fn safe_response_text_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

430 - safe_text_len_mul(input, RESPONSE_MARKERS).map(|_| KimiK3Event::Text)

516 + let markers = input.markers();

517 + safe_text_len_mul(

518 + input,

519 + &[

520 + &markers.response_close,

521 + &markers.tools_open,

522 + &markers.message_close,

523 + &markers.end_of_msg,

524 + ],

525 + )

526 + .map(|_| KimiK3Event::Text)

527 }

528

529 /// Skip non-content noise after the response channel closed.

530 fn skip_epilogue_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

435 - safe_text_len_mul(input, EPILOGUE_MARKERS).map(|_| KimiK3Event::Skip)

531 + let markers = input.markers();

532 + safe_text_len_mul(

533 + input,

534 + &[

535 + &markers.tools_open,

536 + &markers.message_close,

537 + &markers.end_of_msg,

538 + ],

539 + )

540 + .map(|_| KimiK3Event::Skip)

541 }

542

543 /// Skip non-content noise between `call` blocks.

544 fn skip_tools_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

440 - safe_text_len_mul(input, TOOLS_MARKERS).map(|_| KimiK3Event::Skip)

545 + let markers = input.markers();

546 + safe_text_len_mul(

547 + input,

548 + &[

549 + &markers.call_open,

550 + &markers.tools_close,

551 + &markers.message_close,

552 + &markers.end_of_msg,

553 + ],

554 + )

555 + .map(|_| KimiK3Event::Skip)

556 }

557

558 /// Parse a `call` open tag into its tool name and one-based index.

559 +///

560 +/// The attributes end at the structural `<|sep|>`; an ordinary spelling of

561 +/// `<|sep|>` inside an attribute value stays attribute text.

562 fn call_open_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {

563 + let markers = input.markers();

564 + let mut scan = MarkerScanState::default();

565 let (attrs,) = seq!(

446 - _: literal(CALL_OPEN),

447 - take_until(0.., SEP),

448 - _: literal(SEP),

566 + _: &markers.call_open,

567 + take_until_marker(&markers.sep, &mut scan),

568 + _: &markers.sep,

569 )

570 .parse_next(input)?;

571 let attrs = parse_tag_attrs(attrs)?;

@@ -464,9 +584,10 @@ fn call_body_event(

584 input: &mut KimiK3Input<'_>,

585 scan: &mut MarkerScanState,

586 ) -> ModalResult<KimiK3Event> {

587 + let markers = input.markers();

588 let (body,) = seq!(

468 - take_until_marker(CALL_CLOSE, scan),

469 - _: literal(CALL_CLOSE),

589 + take_until_marker(&markers.call_close, scan),

590 + _: &markers.call_close,

591 )

592 .parse_next(input)?;

593 let arguments = parse_call_arguments(body)?;

@@ -614,9 +735,11 @@ mod tests {

735 }

736

737 impl UnifiedParserTestExt for KimiK3UnifiedParser {

738 + /// Parse `chunk` as the detokenizer would deliver it: markers spelled in

739 + /// full are structural special tokens, everything else ordinary text.

740 fn parse_chunk(&mut self, chunk: &str) -> super::Result<UnifiedParserOutput> {

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

619 - self.parse_into(DecodedText::unattributed(chunk), &mut output)?;

742 + self.parse_into(decode_ids(&structural(chunk)), &mut output)?;

743 Ok(output)

744 }

745

@@ -665,14 +788,41 @@ mod tests {

788 }

789 }

790

791 + /// The parser as constructed in production: markers must be spelled by

792 + /// their special tokens.

793 fn test_parser() -> KimiK3UnifiedParser {

794 KimiK3UnifiedParser::new(&[], Arc::new(tokenizer())).unwrap()

795 }

796

797 + /// Run token IDs through the real incremental decoder to obtain attributed

798 + /// input, exactly as the frontend does.

799 + fn decode_ids(ids: &[u32]) -> DecodedText {

800 + let tokenizer = tokenizer();

801 + let mut decoder = tokenizer.create_decode_stream(&[], false, 0);

802 + for &id in ids {

803 + decoder.push_token(id).unwrap();

804 + }

805 + decoder.flush(None).unwrap().1

806 + }

807 +

808 + /// Encode with the special tokens, so markers are structural.

809 + fn structural(text: &str) -> Vec<u32> {

810 + tokenizer().encode(text, false).unwrap()

811 + }

812 +

813 + /// Encode as ordinary bytes only, so marker spellings are content.

814 + fn ordinary(text: &str) -> Vec<u32> {

815 + tokenizer().encode_ordinary(text).unwrap()

816 + }

817 +

818 + /// Stream the structural encoding of the concatenated `chunks`, split at the

819 + /// chunk boundaries, so a marker split across chunks keeps its token anchor.

820 fn collect_stream(parser: &mut KimiK3UnifiedParser, chunks: &[&str]) -> UnifiedParserOutput {

821 + let mut rest = decode_ids(&structural(&chunks.concat()));

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

823 for chunk in chunks {

675 - output.append(parser.parse_chunk(chunk).unwrap());

824 + let piece = rest.drain_prefix(chunk.len());

825 + parser.parse_into(piece, &mut output).unwrap();

826 }

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

828 output

@@ -1136,4 +1286,140 @@ mod tests {

1286

1287 assert_eq!(raw, "<|close|>resp");

1288 }

1289 +

1290 + // --- Ordinary marker spellings ------------------------------------------

1291 +

1292 + fn start_in_reasoning(parser: &mut KimiK3UnifiedParser) {

1293 + parser.initialize(&structural(THINK_OPEN)).unwrap();

1294 + }

1295 +

1296 + fn parse_decoded(

1297 + parser: &mut KimiK3UnifiedParser,

1298 + decoded: DecodedText,

1299 + ) -> UnifiedParserOutput {

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

1301 + parser.parse_into(decoded, &mut output).unwrap();

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

1303 + output

1304 + }

1305 +

1306 + #[test]

1307 + fn kimi_k3_markers_spell_the_documented_strings() {

1308 + let special_tokens = super::SpecialTokens::new(&tokenizer()).unwrap();

1309 + let markers = super::Markers::new(&special_tokens);

1310 +

1311 + assert_eq!(markers.think_open.as_str(), THINK_OPEN);

1312 + assert_eq!(markers.think_close.as_str(), THINK_CLOSE);

1313 + assert_eq!(markers.response_open.as_str(), RESPONSE_OPEN);

1314 + assert_eq!(markers.response_close.as_str(), RESPONSE_CLOSE);

1315 + assert_eq!(markers.tools_open.as_str(), TOOLS_OPEN);

1316 + assert_eq!(markers.tools_close.as_str(), TOOLS_CLOSE);

1317 + assert_eq!(markers.message_close.as_str(), super::MESSAGE_CLOSE);

1318 + assert_eq!(markers.call_close.as_str(), super::CALL_CLOSE);

1319 + assert_eq!(markers.end_of_msg.as_str(), END_OF_MSG);

1320 + }

1321 +

1322 + #[test]

1323 + fn kimi_k3_ordinary_close_marker_text_stays_in_reasoning() {

1324 + // The production incident: the model spells `<|close|>think<|sep|>` with

1325 + // ordinary tokens while reasoning, then really closes the channel later.

1326 + let mut ids = ordinary(THINK_CLOSE);

1327 + ids.extend(structural(&format!(

1328 + " after{THINK_CLOSE}{RESPONSE_OPEN}answer{RESPONSE_CLOSE}<|close|>message{SEP}"

1329 + )));

1330 + let decoded = decode_ids(&ids);

1331 + assert!(

1332 + decoded.text.starts_with(THINK_CLOSE),

1333 + "same bytes as the structural marker"

1334 + );

1335 +

1336 + let mut parser = test_parser();

1337 + start_in_reasoning(&mut parser);

1338 + let output = parse_decoded(&mut parser, decoded);

1339 +

1340 + assert_eq!(output.reasoning_text(), format!("{THINK_CLOSE} after"));

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

1342 + }

1343 +

1344 + #[test]

1345 + fn kimi_k3_ordinary_open_marker_text_stays_in_reasoning() {

1346 + let mut ids = structural("config: ");

1347 + ids.extend(ordinary(THINK_OPEN));

1348 + ids.extend(structural(&format!(

1349 + "{THINK_CLOSE}{RESPONSE_OPEN}ok{RESPONSE_CLOSE}"

1350 + )));

1351 +

1352 + let mut parser = test_parser();

1353 + start_in_reasoning(&mut parser);

1354 + let output = parse_decoded(&mut parser, decode_ids(&ids));

1355 +

1356 + assert_eq!(output.reasoning_text(), format!("config: {THINK_OPEN}"));

1357 + assert_eq!(output.normal_text(), "ok");

1358 + }

1359 +

1360 + #[test]

1361 + fn kimi_k3_attributed_streaming_splits_at_every_char_boundary() {

1362 + // A lookalike close inside reasoning, then genuine markers and a call.

1363 + let mut ids = structural(THINK_OPEN);

1364 + ids.extend(ordinary(THINK_CLOSE));

1365 + ids.extend(structural(&format!(

1366 + " step{THINK_CLOSE}{RESPONSE_OPEN}the answer{RESPONSE_CLOSE}{TOOLS_OPEN}{}{TOOLS_CLOSE}<|close|>message{SEP}",

1367 + call("tool=\"calc\" index=\"1\"", &arg("x", "number", "42")),

1368 + )));

1369 + let full = decode_ids(&ids);

1370 +

1371 + for split in (0..=full.text.len()).filter(|&index| full.text.is_char_boundary(index)) {

1372 + let mut rest = full.clone();

1373 + let head = rest.drain_prefix(split);

1374 + let mut parser = test_parser();

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

1376 + parser.parse_into(head, &mut output).unwrap();

1377 + parser.parse_into(rest, &mut output).unwrap();

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

1379 +

1380 + assert_eq!(

1381 + output.reasoning_text(),

1382 + format!("{THINK_CLOSE} step"),

1383 + "split {split}"

1384 + );

1385 + assert_eq!(output.normal_text(), "the answer", "split {split}");

1386 + assert_eq!(

1387 + first_call(&output).arguments,

1388 + r#"{"x":42}"#,

1389 + "split {split}"

1390 + );

1391 + }

1392 + }

1393 +

1394 + #[test]

1395 + fn kimi_k3_incomplete_utf8_before_structural_close_still_closes() {

1396 + // A stray incomplete byte right before the close marker: the decoder

1397 + // anchors the special token at its own spelling, so the boundary holds.

1398 + let mut ids = vec![0xe4];

1399 + ids.extend(structural(&format!(

1400 + "{THINK_CLOSE}{RESPONSE_OPEN}ok{RESPONSE_CLOSE}"

1401 + )));

1402 +

1403 + let mut parser = test_parser();

1404 + start_in_reasoning(&mut parser);

1405 + let output = parse_decoded(&mut parser, decode_ids(&ids));

1406 +

1407 + assert_eq!(output.reasoning_text(), "\u{FFFD}");

1408 + assert_eq!(output.normal_text(), "ok");

1409 + }

1410 +

1411 + #[test]

1412 + fn kimi_k3_ordinary_sep_text_inside_call_attributes_stays_attribute_text() {

1413 + let mut ids = structural(&format!("{TOOLS_OPEN}{OPEN}call tool=\"a"));

1414 + ids.extend(ordinary(SEP));

1415 + ids.extend(structural(&format!(

1416 + "b\" index=\"1\"{SEP}{}<|close|>call{SEP}{TOOLS_CLOSE}",

1417 + arg("x", "number", "1")

1418 + )));

1419 +

1420 + let output = parse_decoded(&mut test_parser(), decode_ids(&ids));

1421 +

1422 + assert_eq!(first_call(&output).name, Some(format!("a{SEP}b")));

1423 + assert_eq!(first_call(&output).arguments, r#"{"x":1}"#);

1424 + }

1425 }

rust/src/parser/src/unified/mod.rs

@@ -21,6 +21,7 @@ use vllm_tokenizer::{DecodedText, DynTokenizer};

21 use crate::output_grammar::{self, BuiltOutputGrammar, OutputGrammarContext};

22 use crate::reasoning::ReasoningError;

23 use crate::tool::{Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput};

24 +use crate::utils::SpecialToken;

26 /// Result alias for unified parser operations.

27 pub type Result<T> = std::result::Result<T, UnifiedParserError>;

@@ -288,3 +289,11 @@ fn token_id(tokenizer: &dyn vllm_tokenizer::Tokenizer, token: &str) -> Result<u3

289 token: token.to_string(),

290 })

291 }

292 +

293 +/// Resolves `token` to a [`SpecialToken`], or an error if it's not found.

294 +fn special_token(tokenizer: &dyn vllm_tokenizer::Tokenizer, token: &str) -> Result<SpecialToken> {

295 + Ok(SpecialToken {

296 + id: token_id(tokenizer, token)?,

297 + text: token.to_string(),

298 + })

299 +}

rust/src/parser/src/utils.rs

+119 -78

@@ -3,14 +3,23 @@

4 //! Shared helpers for streaming parsers.

6 +use std::fmt::Debug;

7 +

8 +use vllm_tokenizer::DecodedText;

9 use winnow::Parser;

10 use winnow::error::{ContextError, ErrMode, ModalResult, Needed, StrContext, StrContextValue};

8 -use winnow::stream::{FindSlice, Offset, Partial, Stream};

11 +use winnow::stream::{Partial, Stream};

13 use crate::tool::{Result, ToolParserError};

15 +pub(crate) mod marker;

16 pub(crate) mod recursion;

18 +pub use marker::{

19 + Attributed, AttributedExt, Marker, MarkerLike, MarkerStream, SpecialToken,

20 + attributed_with_markers,

21 +};

22 +

23 /// Return the byte length of the longest proper prefix of `token` that is also

24 /// a suffix of `buffer`.

25 ///

@@ -48,48 +57,81 @@ pub fn partial_prefix_len(buffer: &str, token: &str) -> usize {

57 0

58 }

60 +/// Where the next marker may be, in the coordinate system of

61 +/// [`MarkerStream::offset`].

62 +enum Scan {

63 + /// A complete marker starts at this offset.

64 + Found(usize),

65 + /// A marker may start at this offset but needs more input to decide; text

66 + /// from here on must be held back.

67 + Pending(usize),

68 + /// Every candidate before this offset (the end of the input) was rejected.

69 + Clear(usize),

70 +}

71 +

72 +/// Scan for the earliest of `markers` at or after `from`.

73 +///

74 +/// Candidates come from the stream ([`MarkerStream::next_candidate`]): text

75 +/// occurrences for plain markers, visible anchors for markers guarded by a

76 +/// special token. Each candidate is decided by probing the marker's own parser

77 +/// on a cloned cursor, so the scanner and the consuming `alt` branch can never

78 +/// disagree; a rejected candidate (an ordinary lookalike) is skipped and the

79 +/// scan continues.

80 +fn scan_markers<'i, I: MarkerStream<'i>, M: MarkerLike>(

81 + input: &I,

82 + markers: &[M],

83 + from: usize,

84 +) -> ModalResult<Scan> {

85 + let start = input.offset();

86 + let mut from = from.max(start);

87 + while let Some(at) = input.next_candidate(markers, from) {

88 + for marker in markers {

89 + let mut probe = input.clone();

90 + probe.next_slice(at - start);

91 + match marker.as_marker().parse_next(&mut probe) {

92 + Ok(_) => return Ok(Scan::Found(at)),

93 + Err(ErrMode::Backtrack(_)) => {}

94 + Err(ErrMode::Incomplete(_)) => return Ok(Scan::Pending(at)),

95 + Err(error @ ErrMode::Cut(_)) => return Err(error),

96 + }

97 + }

98 + from = at + 1;

99 + }

100 + Ok(Scan::Clear(start + input.eof_offset()))

101 +}

102 +

103 /// Parse a safe text run before the next marker.

104 /// This is the single-marker variant of [`safe_text_len_mul`].

105 ///

106 /// Returns the text length in bytes, and advances the input.

55 -pub fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalResult<usize> {

56 - let text = **input;

57 - if text.is_empty() {

58 - return incomplete();

59 - }

60 -

61 - if let Some(start_idx) = text.find(marker) {

62 - input.next_slice(start_idx);

63 - return Ok(start_idx);

64 - }

65 -

66 - let keep_len = partial_prefix_len(text, marker);

67 - let emit_len = text.len().saturating_sub(keep_len);

68 - if emit_len == 0 {

69 - return incomplete();

70 - }

71 -

72 - input.next_slice(emit_len);

73 - Ok(emit_len)

107 +pub fn safe_text_len<'i, I: MarkerStream<'i>, M: MarkerLike>(

108 + input: &mut I,

109 + marker: M,

110 +) -> ModalResult<usize> {

111 + safe_text_len_mul(input, &[marker])

112 }

113

114 /// Parse a safe text run before the earliest next marker.

115 /// This is the multi-marker variant of [`safe_text_len`].

116 ///

117 +/// Text that may still grow into a marker is held back: a partial spelling at

118 +/// the end of the input for plain markers, and only a genuine special token's

119 +/// incomplete marker for guarded ones.

120 +///

121 /// Returns the text length in bytes, and advances the input.

80 -pub fn safe_text_len_mul(input: &mut Partial<&str>, markers: &[&str]) -> ModalResult<usize> {

81 - let text = **input;

82 - if text.is_empty() {

122 +pub fn safe_text_len_mul<'i, I: MarkerStream<'i>, M: MarkerLike>(

123 + input: &mut I,

124 + markers: &[M],

125 +) -> ModalResult<usize> {

126 + if input.eof_offset() == 0 {

127 return incomplete();

128 }

129

86 - if let Some(start_idx) = find_slice_mul(text, markers) {

87 - input.next_slice(start_idx);

88 - return Ok(start_idx);

89 - }

90 -

91 - let keep_len = markers.iter().map(|marker| partial_prefix_len(text, marker)).max().unwrap_or(0);

92 - let emit_len = text.len().saturating_sub(keep_len);

130 + let start = input.offset();

131 + let stop = match scan_markers(input, markers, start)? {

132 + Scan::Found(at) | Scan::Pending(at) | Scan::Clear(at) => at,

133 + };

134 + let emit_len = stop - start;

135 if emit_len == 0 {

136 return incomplete();

137 }

@@ -98,23 +140,11 @@ pub fn safe_text_len_mul(input: &mut Partial<&str>, markers: &[&str]) -> ModalRe

140 Ok(emit_len)

141 }

142

101 -#[inline(always)]

102 -fn find_slice_mul(text: &str, markers: &[&str]) -> Option<usize> {

103 - let range = match markers {

104 - // Use the fast specialized `winnow::stream::FindSlice` impl for 1-3 markers.

105 - [first] => text.find_slice(*first),

106 - [first, second] => text.find_slice((*first, *second)),

107 - [first, second, third] => text.find_slice((*first, *second, *third)),

108 - // Fall back to a linear scan for 4+ markers.

109 - _ => return markers.iter().filter_map(|marker| text.find(marker)).min(),

110 - };

111 - range.map(|range| range.start)

112 -}

113 -

143 /// Streaming scan state for a buffered marker search [`take_until_marker`],

144 /// so that we don't have to rescan the whole buffered prefix when resuming.

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

146 pub struct MarkerScanState {

147 + /// Offset relative to the cursor at call time from which the scan resumes.

148 scan_start: usize,

149 }

150

@@ -128,55 +158,47 @@ impl MarkerScanState {

158 ///

159 /// This is the streaming-buffered variant of `winnow::token::take_until(0..,

160 /// marker)`: it returns the slice before `marker` and leaves `marker` for the

131 -/// caller to consume. On incomplete input, it stores the earliest byte offset

132 -/// that can still match `marker` and returns `Incomplete` without consuming

133 -/// input, so the next parse can avoid rescanning the whole buffered prefix.

161 +/// caller to consume. On incomplete input, it stores the earliest offset that

162 +/// can still start `marker` and returns `Incomplete` without consuming input,

163 +/// so the next parse can avoid rescanning the whole buffered prefix.

164 ///

165 /// Use this for outer parser states that keep the full buffered input across

166 /// chunks while waiting for a closing marker. Plain `take_until` is still a

167 /// better fit for one-shot parsers over a complete body, and for `1..` cases

168 /// where an empty slice before the marker should be rejected.

139 -pub fn take_until_marker<'i, 'a>(

140 - marker: &'a str,

169 +pub fn take_until_marker<'i, 'a, I: MarkerStream<'i>, M: MarkerLike + 'a>(

170 + marker: M,

171 state: &'a mut MarkerScanState,

142 -) -> impl Parser<Partial<&'i str>, &'i str, ErrMode<ContextError>> + 'a {

143 - move |input: &mut Partial<&'i str>| take_until_marker_(input, marker, state)

172 +) -> impl Parser<I, &'i str, ErrMode<ContextError>> + 'a {

173 + move |input: &mut I| take_until_marker_(input, &marker, state)

174 }

175

146 -fn take_until_marker_<'i>(

147 - input: &mut Partial<&'i str>,

148 - marker: &str,

176 +fn take_until_marker_<'i, I: MarkerStream<'i>, M: MarkerLike>(

177 + input: &mut I,

178 + marker: &M,

179 state: &mut MarkerScanState,

180 ) -> ModalResult<&'i str> {

151 - debug_assert!(!marker.is_empty());

181 + debug_assert!(!marker.as_marker().as_str().is_empty());

182

153 - let text = **input;

154 - if text.is_empty() {

183 + if input.eof_offset() == 0 {

184 return incomplete();

185 }

186

158 - // Normal updates store a char boundary; this keeps stale or misused state from panicking.

159 - let scan_start = floor_char_boundary(text, state.scan_start);

160 -

161 - if let Some(offset) = text[scan_start..].find(marker) {

162 - let marker_start = scan_start + offset;

163 - let body = &text[..marker_start];

164 - input.next_slice(marker_start);

165 - state.reset();

166 - return Ok(body);

167 - }

168 -

169 - let keep_len = partial_prefix_len(text, marker);

170 - state.scan_start = text.len() - keep_len;

171 - incomplete()

172 -}

173 -

174 -fn floor_char_boundary(text: &str, index: usize) -> usize {

175 - let mut index = index.min(text.len());

176 - while !text.is_char_boundary(index) {

177 - index -= 1;

187 + let start = input.offset();

188 + match scan_markers(

189 + input,

190 + std::slice::from_ref(marker),

191 + start + state.scan_start,

192 + )? {

193 + Scan::Found(at) => {

194 + state.reset();

195 + Ok(input.next_slice(at - start))

196 + }

197 + Scan::Pending(at) | Scan::Clear(at) => {

198 + state.scan_start = at - start;

199 + incomplete()

200 + }

201 }

179 - index

202 }

203

204 /// Streaming lexical state for a top-level JSON object.

@@ -394,7 +416,26 @@ pub fn parse_buffered_event<E>(

416 buffer: &str,

417 parse: impl FnOnce(&mut Partial<&str>) -> ModalResult<E>,

418 ) -> Result<Option<(E, usize)>> {

397 - let mut input = Partial::new(buffer);

419 + parse_buffered_stream_event(Partial::new(buffer), parse)

420 +}

421 +

422 +/// Parse one event from a buffered [`DecodedText`], for token-aware parsers.

423 +///

424 +/// Same contract as [`parse_buffered_event`]; the input also carries the

425 +/// buffer's token anchors and the parser's `markers` ([`Attributed`]).

426 +pub fn parse_buffered_event_attributed<M: Debug, E>(

427 + buffer: &DecodedText,

428 + markers: &M,

429 + parse: impl FnOnce(&mut Attributed<'_, M>) -> ModalResult<E>,

430 +) -> Result<Option<(E, usize)>> {

431 + parse_buffered_stream_event(attributed_with_markers(buffer, markers), parse)

432 +}

433 +

434 +fn parse_buffered_stream_event<'i, I: MarkerStream<'i>, E>(

435 + mut input: I,

436 + parse: impl FnOnce(&mut I) -> ModalResult<E>,

437 +) -> Result<Option<(E, usize)>> {

438 + let buffer = input.remaining();

439 let checkpoint = input.checkpoint();

440 let event = match parse(&mut input) {

441 Ok(event) => event,

rust/src/parser/src/utils/marker.rs new

+660

@@ -0,0 +1,660 @@

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

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

3 +

4 +//! Token-aware marker matching for streaming parsers.

5 +//!

6 +//! Model protocols delimit channels with dedicated special tokens such as

7 +//! `<|close|>` and `<|sep|>`. Ordinary BPE tokens can decode to the very same

8 +//! spelling, so matching a marker by text alone mistakes model-written content

9 +//! for structure. The helpers here keep text as the parse axis and consult the

10 +//! generated-token attribution carried by [`DecodedText`] only where a marker's

11 +//! bytes must have come from a special token.

12 +//!

13 +//! - [`Attributed`] is a winnow input over one `DecodedText` buffer: the text

14 +//! plus its buffer-local token anchors, composed from winnow's own stream

15 +//! wrappers so every existing combinator keeps working.

16 +//! - [`Marker`] is a fixed marker string whose special-token segments are

17 +//! *guarded*: they match only when the expected special token produced them.

18 +//! A marker without guards matches by spelling alone, which is exactly what a

19 +//! plain `&str` marker is; [`MarkerRef`] is the borrowed view shared by both.

20 +//! One marker definition is both the parser that consumes the marker and the

21 +//! definition the safe-text scanner in the parent module stops in front of,

22 +//! so the two never disagree.

23 +//! - [`MarkerStream`] abstracts over `Partial<&str>` (text-only parsers) and

24 +//! [`Attributed`], letting one generic scanner serve both.

25 +

26 +use std::fmt::Debug;

27 +use std::ops::Range;

28 +

29 +use auto_impl::auto_impl;

30 +use educe::Educe;

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

32 +use winnow::Parser;

33 +use winnow::error::{ContextError, ErrMode, ModalResult};

34 +use winnow::stream::{

35 + Compare, FindSlice, LocatingSlice, Location, Partial, Stateful, Stream, StreamIsPartial,

36 +};

37 +use winnow::token::literal;

38 +

39 +use super::partial_prefix_len;

40 +

41 +/// Streaming winnow input over one [`DecodedText`] buffer.

42 +///

43 +/// `LocatingSlice` reports the cursor's byte offset from the buffer start, which

44 +/// is the coordinate system of the anchors; `Partial` keeps the incomplete-input

45 +/// semantics streaming parsers rely on; `Stateful` carries the immutable

46 +/// [`AttributedState`]. winnow forwards `Compare`, `FindSlice`, `Location`, and

47 +/// `UpdateSlice` through all three, so `literal`, `alt`, `seq!`, `take_until`,

48 +/// and `rest` work unchanged.

49 +pub type Attributed<'i, M = ()> = Stateful<Partial<LocatingSlice<&'i str>>, AttributedState<'i, M>>;

50 +

51 +/// Immutable state of an [`Attributed`] input.

52 +#[derive(Debug, Educe)]

53 +#[educe(Clone, Copy)]

54 +pub struct AttributedState<'i, M = ()> {

55 + /// Token anchors of the buffer, sorted by byte offset.

56 + anchors: &'i [TokenAttribution],

57 + /// The parser's markers; see [`AttributedExt::markers`].

58 + markers: &'i M,

59 +}

60 +

61 +#[easy_ext::ext(AttributedExt)]

62 +pub impl<'i, M> Attributed<'i, M> {

63 + /// The parser's markers, reachable from every grammar function through the

64 + /// input instead of an extra parameter.

65 + fn markers(&self) -> &'i M {

66 + self.state.markers

67 + }

68 +}

69 +

70 +/// Create an [`Attributed`] input over `buffer`, carrying `markers`.

71 +pub fn attributed_with_markers<'i, M>(

72 + buffer: &'i DecodedText,

73 + markers: &'i M,

74 +) -> Attributed<'i, M> {

75 + Stateful {

76 + input: Partial::new(LocatingSlice::new(buffer.text.as_str())),

77 + state: AttributedState {

78 + anchors: &buffer.attributions,

79 + markers,

80 + },

81 + }

82 +}

83 +

84 +/// A `&str`-sliced streaming input that can say where markers may begin.

85 +///

86 +/// Implemented for `Partial<&str>` (text-only parsers, no anchors) and

87 +/// [`Attributed`].

88 +pub trait MarkerStream<'i>:

89 + Stream<Slice = &'i str> + StreamIsPartial + Clone + for<'m> Compare<&'m str>

90 +{

91 + /// Remaining input from the cursor.

92 + fn remaining(&self) -> &'i str;

93 +

94 + /// Byte offset of the cursor in the coordinate system of

95 + /// [`Self::anchors`] and of resumable scan state.

96 + ///

97 + /// `Partial<&str>` tracks no origin and reports 0, so its offsets are

98 + /// relative to the cursor at the time of the call; [`Attributed`] reports

99 + /// the offset from the buffer start.

100 + fn offset(&self) -> usize;

101 +

102 + /// Token anchors of the buffer, sorted by byte offset; empty for text-only

103 + /// input.

104 + fn anchors(&self) -> &'i [TokenAttribution];

105 +

106 + /// Earliest offset at or after `from` where one of `markers` may begin.

107 + ///

108 + /// A guarded marker can only begin at a visible anchor of its first special

109 + /// token, so text-only input never yields one. Any other marker begins at a

110 + /// complete occurrence of its spelling, or at a proper prefix of it that

111 + /// ends the input.

112 + fn next_candidate<M: MarkerLike>(&self, markers: &[M], from: usize) -> Option<usize> {

113 + let start = self.offset();

114 + let from = from.max(start);

115 + let markers = || markers.iter().map(MarkerLike::as_marker);

116 +

117 + let anchored = markers().any(|marker| marker.first_token_id().is_some()).then(|| {

118 + visible_in(self.anchors(), from..usize::MAX).find_map(|(at, token_id)| {

119 + markers().any(|marker| marker.first_token_id() == Some(token_id)).then_some(at)

120 + })

121 + });

122 + let by_text = text_candidate(

123 + self.remaining(),

124 + markers()

125 + .filter(|marker| marker.first_token_id().is_none())

126 + .map(MarkerRef::as_str),

127 + from - start,

128 + )

129 + .map(|at| start + at);

130 +

131 + match (anchored.flatten(), by_text) {

132 + (Some(anchored), Some(by_text)) => Some(anchored.min(by_text)),

133 + (anchored, by_text) => anchored.or(by_text),

134 + }

135 + }

136 +}

137 +

138 +impl<'i> MarkerStream<'i> for Partial<&'i str> {

139 + fn remaining(&self) -> &'i str {

140 + **self

141 + }

142 +

143 + fn offset(&self) -> usize {

144 + 0

145 + }

146 +

147 + fn anchors(&self) -> &'i [TokenAttribution] {

148 + &[]

149 + }

150 +}

151 +

152 +impl<'i, M: Debug> MarkerStream<'i> for Attributed<'i, M> {

153 + fn remaining(&self) -> &'i str {

154 + ****self

155 + }

156 +

157 + fn offset(&self) -> usize {

158 + self.current_token_start()

159 + }

160 +

161 + fn anchors(&self) -> &'i [TokenAttribution] {

162 + self.state.anchors

163 + }

164 +}

165 +

166 +fn anchor_offset(anchor: TokenAnchor) -> usize {

167 + match anchor {

168 + TokenAnchor::Visible { byte_offset } | TokenAnchor::ZeroWidth { byte_offset } => {

169 + byte_offset as usize

170 + }

171 + }

172 +}

173 +

174 +/// Visible anchors within `range`, as `(byte_offset, token_id)`, in generation

175 +/// order.

176 +fn visible_in(

177 + anchors: &[TokenAttribution],

178 + range: Range<usize>,

179 +) -> impl Iterator<Item = (usize, u32)> + '_ {

180 + let first =

181 + anchors.partition_point(|attribution| anchor_offset(attribution.anchor) < range.start);

182 + anchors[first..]

183 + .iter()

184 + .take_while(move |attribution| anchor_offset(attribution.anchor) < range.end)

185 + .filter_map(|attribution| match attribution.anchor {

186 + TokenAnchor::Visible { byte_offset } => {

187 + Some((byte_offset as usize, attribution.token_id))

188 + }

189 + TokenAnchor::ZeroWidth { .. } => None,

190 + })

191 +}

192 +

193 +/// Return whether `range` is exactly the visible span of the single token

194 +/// `token_id`: one visible anchor at `range.start` with that ID and no other

195 +/// visible anchor inside. Zero-width records are ignored.

196 +///

197 +/// Special tokens never merge with neighbouring text, so this is the identity

198 +/// test for a special token's spelling. Anchors beyond the available input do

199 +/// not exist yet, so a partially arrived genuine spelling passes and is left to

200 +/// `literal` to report as incomplete.

201 +fn is_single_token(anchors: &[TokenAttribution], range: Range<usize>, token_id: u32) -> bool {

202 + let mut visible = visible_in(anchors, range.clone());

203 + visible.next() == Some((range.start, token_id)) && visible.next().is_none()

204 +}

205 +

206 +/// Earliest text position at or after `from` where one of `spellings` occurs

207 +/// in full, or where a proper prefix of one of them ends the input.

208 +fn text_candidate<'m>(

209 + text: &str,

210 + spellings: impl Iterator<Item = &'m str> + Clone,

211 + from: usize,

212 +) -> Option<usize> {

213 + // Round up: a marker starts on a char boundary at or after `from`, and a

214 + // rejected candidate inside a multi-byte char must not be found again.

215 + let from = text.ceil_char_boundary(from);

216 + let tail = &text[from..];

217 + if tail.is_empty() {

218 + return None;

219 + }

220 + if let Some(at) = find_earliest(tail, spellings.clone()) {

221 + return Some(from + at);

222 + }

223 + let keep_len = spellings.map(|spelling| partial_prefix_len(tail, spelling)).max().unwrap_or(0);

224 + (keep_len > 0).then(|| text.len() - keep_len)

225 +}

226 +

227 +/// Earliest occurrence of any spelling in `text`.

228 +#[inline(always)]

229 +fn find_earliest<'m>(text: &str, mut spellings: impl Iterator<Item = &'m str>) -> Option<usize> {

230 + let first = spellings.next()?;

231 + let second = spellings.next();

232 + let third = second.and_then(|_| spellings.next());

233 + let fourth = third.and_then(|_| spellings.next());

234 + // Use the fast specialized `FindSlice` impls for 1-3 spellings, and fall

235 + // back to a linear scan for 4+.

236 + let range = match (second, third, fourth) {

237 + (None, _, _) => text.find_slice(first),

238 + (Some(second), None, _) => text.find_slice((first, second)),

239 + (Some(second), Some(third), None) => text.find_slice((first, second, third)),

240 + (Some(second), Some(third), Some(fourth)) => {

241 + return [first, second, third, fourth]

242 + .into_iter()

243 + .chain(spellings)

244 + .filter_map(|spelling| text.find(spelling))

245 + .min();

246 + }

247 + };

248 + range.map(|range| range.start)

249 +}

250 +

251 +/// A special token's spelling and its ID in the active tokenizer, resolved

252 +/// together so a marker cannot pair one token's spelling with another's ID.

253 +#[derive(Clone, Debug, PartialEq, Eq)]

254 +pub struct SpecialToken {

255 + pub text: String,

256 + pub id: u32,

257 +}

258 +

259 +/// A fixed marker string whose special-token segments are guarded by token

260 +/// identity.

261 +///

262 +/// Build it once from tokenizer IDs, then use it both as the parser that

263 +/// consumes the marker (it implements [`Parser`] by reference) and as the

264 +/// definition the safe-text scanner stops in front of.

265 +///

266 +/// ```ignore

267 +/// let think_close = Marker::special(&close).then_text("think").then_special(&sep);

268 +/// ```

269 +#[derive(Clone, Debug, PartialEq, Eq)]

270 +pub struct Marker {

271 + text: String,

272 + guards: Vec<Guard>,

273 +}

274 +

275 +/// One special-token segment of a marker.

276 +#[derive(Clone, Debug, PartialEq, Eq)]

277 +struct Guard {

278 + /// Byte range of the spelling within the marker text.

279 + range: Range<usize>,

280 + token_id: u32,

281 +}

282 +

283 +impl Marker {

284 + /// A marker matched by spelling alone.

285 + pub fn text(text: impl Into<String>) -> Self {

286 + Self {

287 + text: text.into(),

288 + guards: Vec::new(),

289 + }

290 + }

291 +

292 + /// A marker beginning with the special token `token`.

293 + pub fn special(token: &SpecialToken) -> Self {

294 + Self::text(String::new()).then_special(token)

295 + }

296 +

297 + /// Append ordinary text, matched by spelling.

298 + #[must_use]

299 + pub fn then_text(mut self, text: &str) -> Self {

300 + self.text.push_str(text);

301 + self

302 + }

303 +

304 + /// Append the special token `token`.

305 + #[must_use]

306 + pub fn then_special(mut self, token: &SpecialToken) -> Self {

307 + let start = self.text.len();

308 + self.text.push_str(&token.text);

309 + self.guards.push(Guard {

310 + range: start..self.text.len(),

311 + token_id: token.id,

312 + });

313 + self

314 + }

315 +

316 + /// Drop the token-identity guards, so the marker matches by spelling alone,

317 + /// for callers whose input carries no token attribution.

318 + #[must_use]

319 + pub fn without_guards(mut self) -> Self {

320 + self.guards.clear();

321 + self

322 + }

323 +

324 + /// The full spelling of the marker.

325 + pub fn as_str(&self) -> &str {

326 + &self.text

327 + }

328 +}

329 +

330 +/// Anything the scanning helpers accept as a marker, viewed through

331 +/// [`MarkerRef`].

332 +///

333 +/// A plain `str` is a marker without guards, so text-only parsers keep passing

334 +/// `&[&str]` (or `&String`); token-aware parsers pass `&Marker`.

335 +#[auto_impl(&)]

336 +pub trait MarkerLike {

337 + /// Borrow the marker's spelling and guards.

338 + fn as_marker(&self) -> MarkerRef<'_>;

339 +}

340 +

341 +impl MarkerLike for str {

342 + fn as_marker(&self) -> MarkerRef<'_> {

343 + MarkerRef {

344 + text: self,

345 + guards: &[],

346 + }

347 + }

348 +}

349 +

350 +impl MarkerLike for String {

351 + fn as_marker(&self) -> MarkerRef<'_> {

352 + self.as_str().as_marker()

353 + }

354 +}

355 +

356 +impl MarkerLike for Marker {

357 + fn as_marker(&self) -> MarkerRef<'_> {

358 + MarkerRef {

359 + text: &self.text,

360 + guards: &self.guards,

361 + }

362 + }

363 +}

364 +

365 +/// Borrowed view of a marker: its spelling and guarded segments.

366 +///

367 +/// A plain `&str` is a marker without guards; see [`MarkerLike`].

368 +#[derive(Clone, Copy, Debug)]

369 +pub struct MarkerRef<'m> {

370 + text: &'m str,

371 + guards: &'m [Guard],

372 +}

373 +

374 +impl<'m> MarkerRef<'m> {

375 + /// The full spelling of the marker.

376 + pub fn as_str(self) -> &'m str {

377 + self.text

378 + }

379 +

380 + /// The special token the marker begins with, if it begins with one.

381 + pub fn first_token_id(self) -> Option<u32> {

382 + self.guards

383 + .first()

384 + .filter(|guard| guard.range.start == 0)

385 + .map(|guard| guard.token_id)

386 + }

387 +

388 + /// Parse the marker at the cursor, returning its spelling.

389 + ///

390 + /// Every guard whose first byte is available is checked before the spelling:

391 + /// when those bytes do not come from the expected special token the parser

392 + /// backtracks at once instead of asking for more input, so an ordinary

393 + /// partial lookalike streams as content immediately. A genuine special token

394 + /// split by output holdback still reports incomplete, because its anchor

395 + /// arrived with its first byte. Nothing is consumed on failure.

396 + pub fn parse_next<'i, I: MarkerStream<'i>>(self, input: &mut I) -> ModalResult<&'i str> {

397 + let start = input.offset();

398 + let available = input.eof_offset();

399 + let anchors = input.anchors();

400 + for guard in self.guards.iter().take_while(|guard| guard.range.start < available) {

401 + let range = start + guard.range.start..start + guard.range.end;

402 + if !is_single_token(anchors, range, guard.token_id) {

403 + return Err(ErrMode::Backtrack(ContextError::new()));

404 + }

405 + }

406 + literal(self.text).parse_next(input)

407 + }

408 +}

409 +

410 +impl<'i, I: MarkerStream<'i>> Parser<I, &'i str, ErrMode<ContextError>> for &Marker {

411 + fn parse_next(&mut self, input: &mut I) -> ModalResult<&'i str> {

412 + self.as_marker().parse_next(input)

413 + }

414 +}

415 +

416 +#[cfg(test)]

417 +mod tests {

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

419 + use winnow::Parser;

420 + use winnow::error::ErrMode;

421 + use winnow::stream::{Partial, Stream};

422 +

423 + use super::{

424 + Attributed, Marker, MarkerLike, MarkerStream, SpecialToken, attributed_with_markers,

425 + };

426 + use crate::utils::safe_text_len_mul;

427 +

428 + const CLOSE_ID: u32 = 256;

429 + const SEP_ID: u32 = 258;

430 +

431 + /// Build attributed text from `(piece, token_id)` pairs, one visible anchor

432 + /// at the first byte of each piece.

433 + fn pieces(pieces: &[(&str, u32)]) -> DecodedText {

434 + let mut decoded = DecodedText::default();

435 + for &(piece, token_id) in pieces {

436 + decoded.attributions.push(TokenAttribution {

437 + token_id,

438 + anchor: TokenAnchor::Visible {

439 + byte_offset: decoded.text.len() as u32,

440 + },

441 + });

442 + decoded.text.push_str(piece);

443 + }

444 + decoded

445 + }

446 +

447 + fn token(text: &str, id: u32) -> SpecialToken {

448 + SpecialToken {

449 + text: text.to_string(),

450 + id,

451 + }

452 + }

453 +

454 + /// An input over `buffer` without parser markers.

455 + fn attributed(buffer: &DecodedText) -> Attributed<'_> {

456 + attributed_with_markers(buffer, &())

457 + }

458 +

459 + fn close() -> Marker {

460 + Marker::special(&token("<|close|>", CLOSE_ID))

461 + }

462 +

463 + fn think_close() -> Marker {

464 + close().then_text("think").then_special(&token("<|sep|>", SEP_ID))

465 + }

466 +

467 + /// `<|close|>` spelled by ordinary tokens.

468 + fn ordinary_close() -> Vec<(&'static str, u32)> {

469 + vec![("<", 60), ("|", 124), ("close", 900), ("|", 124), (">", 62)]

470 + }

471 +

472 + #[test]

473 + fn special_marker_accepts_single_token_with_expected_id() {

474 + let buffer = pieces(&[("<|close|>", CLOSE_ID), ("x", 1)]);

475 + let mut input = attributed(&buffer);

476 +

477 + let matched = (&close()).parse_next(&mut input).unwrap();

478 +

479 + assert_eq!(matched, "<|close|>");

480 + assert_eq!(input.remaining(), "x");

481 + }

482 +

483 + #[test]

484 + fn special_marker_rejects_ordinary_tokens_with_same_spelling() {

485 + let mut fixture = ordinary_close();

486 + fixture.push(("x", 1));

487 + let buffer = pieces(&fixture);

488 + let mut input = attributed(&buffer);

489 +

490 + let error = (&close()).parse_next(&mut input).unwrap_err();

491 +

492 + assert!(matches!(error, ErrMode::Backtrack(_)));

493 + assert_eq!(input.remaining(), "<|close|>x", "nothing consumed");

494 + }

495 +

496 + #[test]

497 + fn special_marker_rejects_single_ordinary_token_with_wrong_id() {

498 + let buffer = pieces(&[("<|close|>", 900)]);

499 + let mut input = attributed(&buffer);

500 +

501 + let error = (&close()).parse_next(&mut input).unwrap_err();

502 +

503 + assert!(matches!(error, ErrMode::Backtrack(_)));

504 + }

505 +

506 + #[test]

507 + fn special_marker_reports_incomplete_for_split_genuine_spelling() {

508 + // Output holdback delivered only the first bytes; the anchor came with them.

509 + let buffer = pieces(&[("<|clo", CLOSE_ID)]);

510 + let mut input = attributed(&buffer);

511 +

512 + let error = (&close()).parse_next(&mut input).unwrap_err();

513 +

514 + assert!(matches!(error, ErrMode::Incomplete(_)));

515 + }

516 +

517 + #[test]

518 + fn special_marker_rejects_partial_lookalike_without_waiting() {

519 + let buffer = pieces(&[("<", 60), ("|", 124), ("clo", 901)]);

520 + let mut input = attributed(&buffer);

521 +

522 + let error = (&close()).parse_next(&mut input).unwrap_err();

523 +

524 + assert!(matches!(error, ErrMode::Backtrack(_)));

525 + }

526 +

527 + #[test]

528 + fn special_marker_ignores_zero_width_records_inside_span() {

529 + let mut buffer = pieces(&[("<|close|>", CLOSE_ID)]);

530 + buffer.attributions.push(TokenAttribution {

531 + token_id: 7,

532 + anchor: TokenAnchor::ZeroWidth { byte_offset: 3 },

533 + });

534 + let mut input = attributed(&buffer);

535 +

536 + (&close()).parse_next(&mut input).unwrap();

537 + }

538 +

539 + #[test]

540 + fn marker_without_guards_matches_spelling_only() {

541 + let buffer = pieces(&ordinary_close());

542 + let marker = close().without_guards();

543 +

544 + (&marker).parse_next(&mut attributed(&buffer)).unwrap();

545 + (&marker).parse_next(&mut Partial::new("<|close|>")).unwrap();

546 + }

547 +

548 + #[test]

549 + fn guarded_marker_never_matches_text_only_input() {

550 + let error = (&close()).parse_next(&mut Partial::new("<|close|>")).unwrap_err();

551 +

552 + assert!(matches!(error, ErrMode::Backtrack(_)));

553 + }

554 +

555 + #[test]

556 + fn marker_requires_every_special_segment() {

557 + let marker = think_close();

558 + let genuine = pieces(&[("<|close|>", CLOSE_ID), ("think", 5), ("<|sep|>", SEP_ID)]);

559 + assert_eq!(

560 + (&marker).parse_next(&mut attributed(&genuine)).unwrap(),

561 + "<|close|>think<|sep|>"

562 + );

563 +

564 + // Real close, ordinary separator: the strict rule keeps it as content.

565 + let mixed = pieces(&[("<|close|>", CLOSE_ID), ("think", 5), ("<|sep|>", 902)]);

566 + let mut input = attributed(&mixed);

567 + assert!(matches!(

568 + (&marker).parse_next(&mut input).unwrap_err(),

569 + ErrMode::Backtrack(_)

570 + ));

571 + assert_eq!(

572 + input.remaining(),

573 + "<|close|>think<|sep|>",

574 + "nothing consumed"

575 + );

576 + }

577 +

578 + #[test]

579 + fn marker_first_token_id_requires_leading_special() {

580 + assert_eq!(think_close().as_marker().first_token_id(), Some(CLOSE_ID));

581 + assert_eq!(Marker::text("</think>").as_marker().first_token_id(), None);

582 + assert_eq!(

583 + Marker::text("x")

584 + .then_special(&token("<|sep|>", SEP_ID))

585 + .as_marker()

586 + .first_token_id(),

587 + None

588 + );

589 + assert_eq!("</think>".as_marker().first_token_id(), None);

590 + }

591 +

592 + #[test]

593 + fn attributed_next_candidate_uses_anchors_for_guarded_markers() {

594 + let mut fixture = vec![("say ", 1)];

595 + fixture.extend(ordinary_close()); // lookalike at 4

596 + fixture.extend([("think", 5), ("<|sep|>", 903), (" then ", 2)]);

597 + let genuine_at = fixture.iter().map(|(piece, _)| piece.len()).sum::<usize>();

598 + fixture.extend([("<|close|>", CLOSE_ID), ("think", 5), ("<|sep|>", SEP_ID)]);

599 + let buffer = pieces(&fixture);

600 + let input = attributed(&buffer);

601 +

602 + let marker = think_close();

603 + assert_eq!(input.next_candidate(&[&marker], 0), Some(genuine_at));

604 + assert_eq!(input.next_candidate(&[&marker], genuine_at + 1), None);

605 + }

606 +

607 + #[test]

608 + fn attributed_next_candidate_falls_back_to_text_for_unguarded_markers() {

609 + let buffer = pieces(&[("a", 1), ("</think>", 2), ("b", 3)]);

610 + let input = attributed(&buffer);

611 +

612 + assert_eq!(input.next_candidate(&["</think>"], 0), Some(1));

613 + assert_eq!(input.next_candidate(&["</thinking>"], 0), None);

614 + // Partial spelling at the end of the input is a candidate too.

615 + let buffer = pieces(&[("a", 1), ("</thi", 2)]);

616 + let input = attributed(&buffer);

617 + assert_eq!(input.next_candidate(&["</think>"], 0), Some(1));

618 + }

619 +

620 + #[test]

621 + fn text_next_candidate_matches_previous_scan_semantics() {

622 + let input = Partial::new("hello<not_marker><|tool");

623 +

624 + assert_eq!(

625 + input.next_candidate(&["<|tool_call>", "<|channel>thought\n"], 0),

626 + Some("hello<not_marker>".len())

627 + );

628 + assert_eq!(input.next_candidate(&["<|tool_call>"], 18), None);

629 + assert_eq!(

630 + Partial::new("hello<channel|><|tool_call>")

631 + .next_candidate(&["<|tool_call>", "<channel|>"], 0),

632 + Some(5)

633 + );

634 + }

635 +

636 + #[test]

637 + fn scan_skips_rejected_text_candidate_starting_with_multibyte_char() {

638 + // A marker whose leading text is found by spelling but whose guard then

639 + // fails: the scan must move past the whole char, not retry inside it.

640 + let marker = Marker::text("é").then_special(&token("<|sep|>", SEP_ID));

641 + let buffer = pieces(&[("aé", 1), ("<", 60), ("|sep|>", 904), ("b", 2)]);

642 + let mut input = attributed(&buffer);

643 +

644 + let len = safe_text_len_mul(&mut input, &[&marker]).unwrap();

645 +

646 + assert_eq!(len, buffer.text.len());

647 + }

648 +

649 + #[test]

650 + fn attributed_offsets_follow_the_cursor() {

651 + let buffer = pieces(&[("ab", 1), ("<|close|>", CLOSE_ID)]);

652 + let mut input = attributed(&buffer);

653 + input.next_slice(2);

654 +

655 + assert_eq!(input.offset(), 2);

656 + assert_eq!(input.remaining(), "<|close|>");

657 + (&close()).parse_next(&mut input).unwrap();

658 + assert_eq!(input.offset(), 11);

659 + }

660 +}