[Frontend] Prevent Kimi K3 reserved markers in response text (#52889)
Signed-off-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Bugen Zhao committed
Aug 22, 2026 at 02:06 UTC
08b73d07cc467ea29f5f13f2eec067231c7e7322
4 files changed
+31
-7
rust/src/parser/src/unified/kimi_k3.rs
+1
@@ -56,6 +56,7 @@ use crate::unified::parsing_failed;
56
use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker};
57
58
const OPEN: &str = "<|open|>";
59
+const CLOSE: &str = "<|close|>";
60
const SEP: &str = "<|sep|>";
61
const END_OF_MSG: &str = "<|end_of_msg|>";
62
rust/src/parser/src/unified/kimi_k3/structural_tag.rs
+6
-4
@@ -12,8 +12,8 @@ use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, T
12
use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters};
13
14
use super::{
15
- ARG_CLOSE, CALL_CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, RESPONSE_CLOSE,
16
- RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN,
15
+ ARG_CLOSE, CALL_CLOSE, CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN,
16
+ RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN,
17
};
18
19
pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder =
@@ -64,7 +64,9 @@ fn response_prefix(reasoning: bool) -> Vec<Format> {
64
}
65
elements.push(Format::tag(
66
"",
67
- Format::any_text_excluding(&[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]),
67
+ // Keep marker-looking prefixes out of response text so the grammar
68
+ // must resolve them through a valid channel boundary.
69
+ Format::any_text_excluding(&[OPEN, CLOSE, END_OF_MSG]),
70
RESPONSE_CLOSE,
71
));
72
elements
@@ -384,7 +386,7 @@ mod tests {
386
)
387
.unwrap();
388
387
- expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|close|>response<|sep|>","<|open|>tools<|sep|>","<|close|>message<|sep|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap());
389
+ expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|open|>","<|close|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap());
390
}
391
392
#[test]
tests/tool_parsers/test_structural_tag_registry.py
+15
@@ -177,6 +177,7 @@ _K3_TOOLS_CLOSE = "<|close|>tools<|sep|>"
177
_K3_CALL_CLOSE = "<|close|>call<|sep|>"
178
_K3_ARG_CLOSE = "<|close|>argument<|sep|>"
179
_K3_MESSAGE_CLOSE = "<|close|>message<|sep|>"
180
+_K3_END_OF_MSG = "<|end_of_msg|>"
181
182
183
def _k3_tools_by_name() -> list[ChatCompletionToolsParam]:
@@ -357,6 +358,20 @@ def test_kimi_k3_auto_strict_allows_response_only(sample_tools_strict):
358
assert _is_grammar_accept_string(grammar, _k3_response("Just answering."))
359
360
361
+@pytest.mark.parametrize(
362
+ "body",
363
+ [
364
+ _K3_RESPONSE_OPEN + "answer<|close|>message",
365
+ _K3_RESPONSE_OPEN + "answer<|open|>tools",
366
+ _K3_RESPONSE_OPEN + "answer" + _K3_END_OF_MSG,
367
+ ],
368
+)
369
+def test_kimi_k3_response_body_rejects_reserved_marker_prefixes(body: str):
370
+ assert not _is_grammar_accept_string(
371
+ _k3_grammar("required"), body, require_termination=False
372
+ )
373
+
374
+
375
@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS))
376
def test_get_model_structural_tag_supports_named_tool_choice(
377
model: str,
vllm/tool_parsers/structural_tag_registry.py
+9
-3
@@ -376,6 +376,7 @@ _K3_ARG_CLOSE = f"{_K3_CLOSE}argument{_K3_SEP}"
376
# the end-of-message token. It is generated (not part of the prompt prefix), so
377
# the tag must permit it or the FSM would mask the model's natural terminator.
378
_K3_MESSAGE_CLOSE = f"{_K3_CLOSE}message{_K3_SEP}"
379
+_K3_END_OF_MSG = "<|end_of_msg|>"
380
381
# JSON-schema type -> K3 XTML ``type=`` attribute value. Mirrors
382
# ``encoding_k3._xtml_type`` (integer collapses onto number).
@@ -556,12 +557,17 @@ def _k3_response_prefix() -> list[Any]:
557
558
``response`` is generated in thinking mode (prefix ends at
559
``<|open|>think<|sep|>``) but is part of the generation prefix in
559
- non-thinking mode, so its open marker is optional. The body is bounded by
560
- the response close marker.
560
+ non-thinking mode, so its open marker is optional. Excluding the reserved
561
+ marker heads commits ``<|close|>`` to the response-close branch immediately
562
+ and prevents a nested channel from starting inside response text.
563
"""
564
return [
565
OptionalFormat(content=ConstStringFormat(value=_K3_RESPONSE_OPEN)),
564
- TagFormat(begin="", content=AnyTextFormat(), end=_K3_RESPONSE_CLOSE),
566
+ TagFormat(
567
+ begin="",
568
+ content=AnyTextFormat(excludes=[_K3_OPEN, _K3_CLOSE, _K3_END_OF_MSG]),
569
+ end=_K3_RESPONSE_CLOSE,
570
+ ),
571
]
572
573