A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 20:15 UTC
fix: update kimi k2 tool parser logic (#31207)
Signed-off-by: wangln19 <wanglinian@dev.wanglinian.msh-dev.svc.cluster.local> Signed-off-by: Wang Linian <wanglinian@stu.pku.edu.cn> Co-authored-by: wangln19 <wanglinian@dev.wanglinian.msh-dev.svc.cluster.local> Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
wangln19 committed
Dec 30, 2025 at 10:01 UTC
358bfd315cad6956f0921a4e8a9c4baa053abafc
2 files changed
+211
-202
tests/tool_parsers/test_kimi_k2_tool_parser.py
+192
-191
@@ -44,6 +44,33 @@ def assert_tool_calls(
44
)
45
46
47
+def run_streaming_sequence(parser, deltas):
48
+ """Helper to simulate a streaming sequence and return results."""
49
+ previous_text = ""
50
+ previous_token_ids: list[int] = []
51
+ results = []
52
+
53
+ for delta_text, delta_token_ids in deltas:
54
+ current_text = previous_text + delta_text
55
+ current_token_ids = previous_token_ids + delta_token_ids
56
+
57
+ result = parser.extract_tool_calls_streaming(
58
+ previous_text=previous_text,
59
+ current_text=current_text,
60
+ delta_text=delta_text,
61
+ previous_token_ids=previous_token_ids,
62
+ current_token_ids=current_token_ids,
63
+ delta_token_ids=delta_token_ids,
64
+ request=None,
65
+ )
66
+ results.append(result)
67
+
68
+ previous_text = current_text
69
+ previous_token_ids = current_token_ids
70
+
71
+ return results
72
+
73
+
74
def test_extract_tool_calls_no_tools(kimi_k2_tool_parser):
75
model_output = "This is a test"
76
extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(
@@ -346,61 +373,32 @@ def test_token_leak_between_section_and_tool_begin(kimi_k2_tool_parser):
373
tool_call_begin_token_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>")
374
375
# Simulate streaming sequence:
376
+ deltas = [
377
+ ("I'll help you with that. ", [1, 2, 3]),
378
+ ("<|tool_calls_section_begin|>", [section_begin_token_id]),
379
+ (" spurious text ", [4, 5]),
380
+ ("<|tool_call_begin|>", [tool_call_begin_token_id]),
381
+ ]
382
+
383
+ results = run_streaming_sequence(kimi_k2_tool_parser, deltas)
384
+
385
# Delta 1: "I'll help you with that. "
350
- result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(
351
- previous_text="",
352
- current_text="I'll help you with that. ",
353
- delta_text="I'll help you with that. ",
354
- previous_token_ids=[],
355
- current_token_ids=[1, 2, 3], # Regular tokens
356
- delta_token_ids=[1, 2, 3],
357
- request=None,
358
- )
359
- assert result1 is not None
360
- assert result1.content == "I'll help you with that. "
386
+ assert results[0] is not None
387
+ assert results[0].content == "I'll help you with that. "
388
389
# Delta 2: "<|tool_calls_section_begin|>"
363
- prev_ids = [1, 2, 3]
364
- curr_ids = prev_ids + [section_begin_token_id]
365
- result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
366
- previous_text="I'll help you with that. ",
367
- current_text="I'll help you with that. <|tool_calls_section_begin|>",
368
- delta_text="<|tool_calls_section_begin|>",
369
- previous_token_ids=prev_ids,
370
- current_token_ids=curr_ids,
371
- delta_token_ids=[section_begin_token_id],
372
- request=None,
373
- )
390
# Section marker should be stripped and suppressed
375
- assert result2 is None or (result2.content is None or result2.content == "")
391
+ assert results[1] is None or (
392
+ results[1].content is None or results[1].content == ""
393
+ )
394
395
# Delta 3: " spurious text or tokens " (THE LEAK SCENARIO)
378
- prev_ids = curr_ids
379
- curr_ids = curr_ids + [4, 5]
380
- result3 = kimi_k2_tool_parser.extract_tool_calls_streaming(
381
- previous_text="I'll help you with that. <|tool_calls_section_begin|>",
382
- current_text="I'll help you with that. <|tool_calls_section_begin|> spurious text ",
383
- delta_text=" spurious text ",
384
- previous_token_ids=prev_ids,
385
- current_token_ids=curr_ids,
386
- delta_token_ids=[4, 5],
387
- request=None,
388
- )
396
# CRITICAL: This text should be suppressed, NOT returned as reasoning_delta
390
- assert result3 is None or (result3.content is None or result3.content == "")
397
+ assert results[2] is None or (
398
+ results[2].content is None or results[2].content == ""
399
+ )
400
401
# Delta 4: "<|tool_call_begin|>..."
393
- prev_ids = curr_ids
394
- curr_ids = curr_ids + [tool_call_begin_token_id]
395
- _result4 = kimi_k2_tool_parser.extract_tool_calls_streaming(
396
- previous_text="I'll help you with that. <|tool_calls_section_begin|> spurious text ",
397
- current_text="I'll help you with that. <|tool_calls_section_begin|> spurious text <|tool_call_begin|>",
398
- delta_text="<|tool_call_begin|>",
399
- previous_token_ids=prev_ids,
400
- current_token_ids=curr_ids,
401
- delta_token_ids=[tool_call_begin_token_id],
402
- request=None,
403
- )
402
# Now we're in tool call mode, result depends on internal state
403
# The key is that the spurious text from Delta 3 was not leaked
404
@@ -416,31 +414,15 @@ def test_split_markers_across_deltas(kimi_k2_tool_parser):
414
"<|tool_calls_section_begin|>"
415
)
416
419
- # Delta 1: "...reasoning<|tool_calls_sec"
420
- _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(
421
- previous_text="Some reasoning",
422
- current_text="Some reasoning<|tool_calls_sec",
423
- delta_text="<|tool_calls_sec",
424
- previous_token_ids=[1, 2],
425
- current_token_ids=[1, 2, 3], # Partial token
426
- delta_token_ids=[3],
427
- request=None,
428
- )
429
- # Partial token not recognized yet, might be buffered
430
- # Should return as content or None (depends on implementation)
417
+ # Delta 1: partial token, Delta 2: complete marker
418
+ deltas = [
419
+ ("<|tool_calls_sec", [3]),
420
+ ("tion_begin|> ", [section_begin_token_id, 4]),
421
+ ]
422
+
423
+ _results = run_streaming_sequence(kimi_k2_tool_parser, deltas)
424
432
- # Delta 2: "tion_begin|> " (completes the marker)
433
- _result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
434
- previous_text="Some reasoning<|tool_calls_sec",
435
- current_text="Some reasoning<|tool_calls_section_begin|> ",
436
- delta_text="tion_begin|> ",
437
- previous_token_ids=[1, 2, 3],
438
- current_token_ids=[1, 2, section_begin_token_id, 4],
439
- delta_token_ids=[section_begin_token_id, 4],
440
- request=None,
441
- )
425
# Now the complete marker should be detected via buffer
443
- # The parser should enter tool section mode
426
assert kimi_k2_tool_parser.in_tool_section is True
427
428
@@ -475,42 +457,17 @@ def test_reentry_to_reasoning_after_tool_section(kimi_k2_tool_parser):
457
section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")
458
section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")
459
478
- # Enter tool section
479
- _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(
480
- previous_text="",
481
- current_text="<|tool_calls_section_begin|>",
482
- delta_text="<|tool_calls_section_begin|>",
483
- previous_token_ids=[],
484
- current_token_ids=[section_begin_id],
485
- delta_token_ids=[section_begin_id],
486
- request=None,
487
- )
488
- assert kimi_k2_tool_parser.in_tool_section is True
460
+ deltas = [
461
+ ("<|tool_calls_section_begin|>", [section_begin_id]),
462
+ ("<|tool_calls_section_end|>", [section_end_id]),
463
+ (" More reasoning", [10, 11]),
464
+ ]
465
490
- # Exit tool section
491
- _result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
492
- previous_text="<|tool_calls_section_begin|>",
493
- current_text="<|tool_calls_section_begin|><|tool_calls_section_end|>",
494
- delta_text="<|tool_calls_section_end|>",
495
- previous_token_ids=[section_begin_id],
496
- current_token_ids=[section_begin_id, section_end_id],
497
- delta_token_ids=[section_end_id],
498
- request=None,
499
- )
500
- assert kimi_k2_tool_parser.in_tool_section is False
466
+ results = run_streaming_sequence(kimi_k2_tool_parser, deltas)
467
502
- # Subsequent reasoning text should be returned normally
503
- result3 = kimi_k2_tool_parser.extract_tool_calls_streaming(
504
- previous_text="<|tool_calls_section_begin|><|tool_calls_section_end|>",
505
- current_text="<|tool_calls_section_begin|><|tool_calls_section_end|> More reasoning",
506
- delta_text=" More reasoning",
507
- previous_token_ids=[section_begin_id, section_end_id],
508
- current_token_ids=[section_begin_id, section_end_id, 10, 11],
509
- delta_token_ids=[10, 11],
510
- request=None,
511
- )
512
- assert result3 is not None
513
- assert result3.content == " More reasoning"
468
+ assert kimi_k2_tool_parser.in_tool_section is False
469
+ assert results[2] is not None
470
+ assert results[2].content == " More reasoning"
471
472
473
def test_empty_tool_section(kimi_k2_tool_parser):
@@ -819,106 +776,150 @@ def test_tool_call_end_and_section_end_same_chunk(kimi_k2_tool_parser):
776
tool_end_id = kimi_k2_tool_parser.vocab.get("<|tool_call_end|>")
777
778
# Simulate a streaming sequence for a SHORT tool call (all in one chunk):
822
- # 1. Reasoning text
823
- result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(
824
- previous_text="",
825
- current_text="Let me help. ",
826
- delta_text="Let me help. ",
827
- previous_token_ids=[],
828
- current_token_ids=[1, 2],
829
- delta_token_ids=[1, 2],
830
- request=None,
831
- )
832
- assert result1 is not None
833
- assert result1.content == "Let me help. "
834
-
835
- # 2. Section begin
836
- _result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
837
- previous_text="Let me help. ",
838
- current_text="Let me help. <|tool_calls_section_begin|>",
839
- delta_text="<|tool_calls_section_begin|>",
840
- previous_token_ids=[1, 2],
841
- current_token_ids=[1, 2, section_begin_id],
842
- delta_token_ids=[section_begin_id],
843
- request=None,
844
- )
845
- assert kimi_k2_tool_parser.in_tool_section is True
846
-
847
- # 3. Tool call begin + full content + tool_end + section_end ALL IN ONE CHUNK
848
- # This is the critical scenario for short tool calls
779
combined = (
780
'<|tool_call_begin|>get_weather:0 <|tool_call_argument_begin|> {"city": "Paris"} '
781
"<|tool_call_end|><|tool_calls_section_end|>"
782
)
783
854
- # Build up the previous text gradually to simulate realistic streaming
855
- prev_text = "Let me help. <|tool_calls_section_begin|>"
856
- curr_text = prev_text + combined
784
+ deltas = [
785
+ ("Let me help. ", [1, 2]),
786
+ ("<|tool_calls_section_begin|>", [section_begin_id]),
787
+ (combined, [tool_begin_id, 10, 11, 12, tool_end_id, section_end_id]),
788
+ (" Done", [20]),
789
+ ]
790
858
- result3 = kimi_k2_tool_parser.extract_tool_calls_streaming(
859
- previous_text=prev_text,
860
- current_text=curr_text,
861
- delta_text=combined,
862
- previous_token_ids=[1, 2, section_begin_id],
863
- current_token_ids=[
864
- 1,
865
- 2,
866
- section_begin_id,
867
- tool_begin_id,
868
- 10,
869
- 11,
870
- 12,
871
- tool_end_id,
872
- section_end_id,
873
- ],
874
- delta_token_ids=[tool_begin_id, 10, 11, 12, tool_end_id, section_end_id],
875
- request=None,
876
- )
791
+ results = run_streaming_sequence(kimi_k2_tool_parser, deltas)
792
793
# CRITICAL: Parser should have exited section AFTER processing tool
794
assert kimi_k2_tool_parser.in_tool_section is False
795
796
# Tool call should have been emitted (not dropped)
882
- # The result might be the tool name or None depending on state, but
883
- # importantly, it shouldn't be returning the literal tokens as content
884
-
885
- if result3 is not None and result3.content is not None:
797
+ if results[2] is not None and results[2].content is not None:
798
# Verify no special tokens leaked into content
887
- assert "<|tool_call_end|>" not in result3.content
888
- assert "<|tool_calls_section_end|>" not in result3.content
799
+ assert "<|tool_call_end|>" not in results[2].content
800
+ assert "<|tool_calls_section_end|>" not in results[2].content
801
890
- # 4. Verify subsequent content streams normally
891
- result4 = kimi_k2_tool_parser.extract_tool_calls_streaming(
892
- previous_text=curr_text,
893
- current_text=curr_text + " Done",
894
- delta_text=" Done",
895
- previous_token_ids=[
896
- 1,
897
- 2,
898
- section_begin_id,
899
- tool_begin_id,
900
- 10,
901
- 11,
902
- 12,
903
- tool_end_id,
904
- section_end_id,
905
- ],
906
- current_token_ids=[
907
- 1,
908
- 2,
909
- section_begin_id,
910
- tool_begin_id,
911
- 10,
912
- 11,
913
- 12,
914
- tool_end_id,
915
- section_end_id,
916
- 20,
917
- ],
918
- delta_token_ids=[20],
919
- request=None,
802
+ # Content after tool section should stream normally
803
+ assert results[3] is not None
804
+ assert results[3].content == " Done"
805
+
806
+
807
+def test_streaming_tool_call_markers_not_leaked(kimi_k2_tool_parser):
808
+ """
809
+ CRITICAL TEST: Verify that tool call markers (<|tool_call_begin|>,
810
+ <|tool_call_end|>, <|tool_call_argument_begin|>) are NOT leaked
811
+ into the content field during streaming.
812
+
813
+ This reproduces the AWS Bedrock bug where tool call markers appeared
814
+ in the 'text' field of responses.
815
+ """
816
+ kimi_k2_tool_parser.reset_streaming_state()
817
+
818
+ section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")
819
+ section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")
820
+ tool_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>")
821
+ tool_end_id = kimi_k2_tool_parser.vocab.get("<|tool_call_end|>")
822
+
823
+ # List of markers that should NEVER appear in content
824
+ forbidden_markers = [
825
+ "<|tool_call_begin|>",
826
+ "<|tool_call_end|>",
827
+ "<|tool_call_argument_begin|>",
828
+ "<|tool_calls_section_begin|>",
829
+ "<|tool_calls_section_end|>",
830
+ ]
831
+
832
+ all_content = []
833
+
834
+ # Steps: reasoning, section begin, tool call, section end, more reasoning
835
+ tool_chunk = (
836
+ "<|tool_call_begin|> functions.get_weather:0 "
837
+ '<|tool_call_argument_begin|> {"city": "Tokyo"} <|tool_call_end|>'
838
)
839
+ deltas = [
840
+ ("I'll check the weather. ", [1, 2, 3]),
841
+ ("<|tool_calls_section_begin|>", [section_begin_id]),
842
+ (tool_chunk, [tool_begin_id, 10, 11, tool_end_id]),
843
+ ("<|tool_calls_section_end|>", [section_end_id]),
844
+ (" Here's the result.", [20, 21]),
845
+ ]
846
+
847
+ results = run_streaming_sequence(kimi_k2_tool_parser, deltas)
848
+
849
+ for res in results:
850
+ if res and res.content:
851
+ all_content.append(res.content)
852
+
853
+ # CRITICAL ASSERTIONS: No forbidden markers in any content
854
+ full_content = "".join(all_content)
855
+ for marker in forbidden_markers:
856
+ assert marker not in full_content, (
857
+ f"MARKER LEAK DETECTED: '{marker}' found in content. "
858
+ f"Full content: {repr(full_content)}"
859
+ )
860
922
- # Content after tool section should stream normally
923
- assert result4 is not None
924
- assert result4.content == " Done"
861
+ # Also check that tool call content (function name, arguments) is not leaked
862
+ assert "get_weather" not in full_content, (
863
+ f"TOOL CALL CONTENT LEAKED: 'get_weather' found in content. "
864
+ f"Full content: {repr(full_content)}"
865
+ )
866
+ assert "Tokyo" not in full_content, (
867
+ f"TOOL CALL CONTENT LEAKED: 'Tokyo' found in content. "
868
+ f"Full content: {repr(full_content)}"
869
+ )
870
+
871
+ # Verify that legitimate content was preserved
872
+ assert "I'll check the weather." in full_content or len(all_content) > 0
873
+
874
+
875
+def test_streaming_multiple_tool_calls_not_leaked(kimi_k2_tool_parser):
876
+ """
877
+ Test that MULTIPLE tool calls in streaming mode do not leak into content.
878
+ This reproduces the AWS Bedrock scenario: "Compare weather in Tokyo and NYC".
879
+ """
880
+ kimi_k2_tool_parser.reset_streaming_state()
881
+
882
+ section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")
883
+ section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")
884
+ tool_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>")
885
+ tool_end_id = kimi_k2_tool_parser.vocab.get("<|tool_call_end|>")
886
+
887
+ all_content = []
888
+
889
+ tool1 = '<|tool_call_begin|> get_weather:0 <|tool_call_argument_begin|> {"city": "Tokyo"} <|tool_call_end|>'
890
+ tool2 = ' <|tool_call_begin|> get_weather:1 <|tool_call_argument_begin|> {"city": "New York"} <|tool_call_end|>'
891
+
892
+ deltas = [
893
+ ("I'll compare the weather. ", [1, 2, 3]),
894
+ ("<|tool_calls_section_begin|>", [section_begin_id]),
895
+ (tool1, [tool_begin_id, 10, tool_end_id]),
896
+ (tool2, [tool_begin_id, 20, tool_end_id]),
897
+ ("<|tool_calls_section_end|>", [section_end_id]),
898
+ (" Here's the comparison.", [30]),
899
+ ]
900
+
901
+ results = run_streaming_sequence(kimi_k2_tool_parser, deltas)
902
+
903
+ for res in results:
904
+ if res and res.content:
905
+ all_content.append(res.content)
906
+
907
+ # Assertions
908
+ full_content = "".join(all_content)
909
+
910
+ # Check no markers leaked
911
+ forbidden = ["<|tool_call", "<|tool_calls_section"]
912
+ for marker in forbidden:
913
+ assert marker not in full_content, (
914
+ f"MARKER LEAKED: {marker} in {repr(full_content)}"
915
+ )
916
+
917
+ # Check no tool call content leaked (both tools)
918
+ assert "get_weather" not in full_content, f"TOOL NAME LEAKED: {repr(full_content)}"
919
+ assert "Tokyo" not in full_content, f"TOOL ARG LEAKED (Tokyo): {repr(full_content)}"
920
+ assert "New York" not in full_content, (
921
+ f"TOOL ARG LEAKED (NYC): {repr(full_content)}"
922
+ )
923
+
924
+ # Legitimate content preserved
925
+ assert "compare" in full_content.lower() or len(all_content) > 0
vllm/tool_parsers/kimi_k2_tool_parser.py
+19
-11
@@ -122,7 +122,6 @@ class KimiK2ToolParser(ToolParser):
122
if variant in cleaned:
123
cleaned = cleaned.replace(variant, "")
124
found_end = True
125
-
125
return cleaned, found_begin, found_end
126
127
def _reset_section_state(self) -> None:
@@ -238,6 +237,7 @@ class KimiK2ToolParser(ToolParser):
237
self.in_tool_section = True
238
self.token_buffer = buffered_text # Use cleaned buffer
239
self.section_char_count = 0 # Reset counter for new section
240
+
241
if found_section_end and self.in_tool_section:
242
logger.debug("Detected section end marker")
243
# CRITICAL: Don't exit early if tool_call_end is in this chunk.
@@ -252,13 +252,18 @@ class KimiK2ToolParser(ToolParser):
252
else:
253
# No tool call ending, safe to exit immediately
254
logger.debug("Exiting tool section")
255
- remaining = buffered_text
255
self._reset_section_state()
257
- # Return remaining text as reasoning content if non-empty
258
- if remaining.strip():
259
- return DeltaMessage(content=remaining)
260
- # Return empty delta to maintain function contract
261
- # (always returns DeltaMessage)
256
+ # Extract any content AFTER the section end marker in delta_text
257
+ # (don't use buffered_text as it contains tool call data)
258
+ post_section_content = ""
259
+ for variant in self.tool_calls_end_token_variants:
260
+ if variant in delta_text:
261
+ parts = delta_text.split(variant, 1)
262
+ if len(parts) > 1:
263
+ post_section_content = parts[1]
264
+ break
265
+ if post_section_content.strip():
266
+ return DeltaMessage(content=post_section_content)
267
return DeltaMessage(content="")
268
else:
269
self.token_buffer = buffered_text
@@ -316,12 +321,12 @@ class KimiK2ToolParser(ToolParser):
321
and prev_tool_end_count == cur_tool_end_count
322
and self.tool_call_end_token not in delta_text
323
):
319
- # CRITICAL FIX: Suppress content if in tool section but
320
- # no tool calls started
324
+ # Suppress content between section begin and first tool begin
325
+ # (header noise). Don't suppress content between tools to avoid
326
+ # breaking potential delimiter characters.
327
if self.in_tool_section and cur_tool_start_count == 0:
328
logger.debug(
323
- "In tool section but no tool calls started yet. "
324
- "Suppressing: %s",
329
+ "In tool section before first tool, suppressing: %s",
330
delta_text,
331
)
332
# Return empty delta to maintain iterator contract
@@ -488,6 +493,9 @@ class KimiK2ToolParser(ToolParser):
493
if tool_call_portion is None:
494
# if there's text but not tool calls, send that -
495
# otherwise None to skip chunk
496
+ # CRITICAL: Never return content if we're in a tool section
497
+ if self.in_tool_section:
498
+ return None
499
delta = (
500
DeltaMessage(content=delta_text)
501
if text_portion is not None