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 6, 10:45 UTC
[Bugfix] Add Kimi-K2.5 reasoning/tool parser aliases and tool_call_id support (#37438)
Signed-off-by: DorBernsohn <dor.bernsohn@gmail.com>
DorBernsohn committed
Mar 19, 2026 at 15:08 UTC
c63ca2b2e696e8dd1ae0f5ace08fd57a6a95a65f
4 files changed
+173
-18
tests/reasoning/test_kimi_k2_reasoning_parser.py
new
+155
@@ -0,0 +1,155 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
7
+from vllm.entrypoints.openai.engine.protocol import DeltaMessage
8
+from vllm.reasoning.identity_reasoning_parser import IdentityReasoningParser
9
+from vllm.reasoning.kimi_k2_reasoning_parser import KimiK2ReasoningParser
10
+from vllm.tokenizers import get_tokenizer
11
+
12
+REASONING_MODEL_NAME = "moonshotai/Kimi-K2.5"
13
+
14
+
15
+@pytest.fixture(scope="module")
16
+def kimi_k2_tokenizer():
17
+ return get_tokenizer(tokenizer_name=REASONING_MODEL_NAME, trust_remote_code=True)
18
+
19
+
20
+def test_parser_selection_thinking_enabled(kimi_k2_tokenizer):
21
+ parser = KimiK2ReasoningParser(
22
+ kimi_k2_tokenizer, chat_template_kwargs={"thinking": True}
23
+ )
24
+ assert parser._identity_parser is None
25
+
26
+
27
+def test_parser_selection_thinking_disabled(kimi_k2_tokenizer):
28
+ parser = KimiK2ReasoningParser(
29
+ kimi_k2_tokenizer, chat_template_kwargs={"thinking": False}
30
+ )
31
+ assert isinstance(parser._identity_parser, IdentityReasoningParser)
32
+
33
+
34
+def test_extract_reasoning_with_think_tags(kimi_k2_tokenizer):
35
+ parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
36
+ request = ChatCompletionRequest(model="test-model", messages=[], temperature=1.0)
37
+
38
+ reasoning, content = parser.extract_reasoning(
39
+ "<think>step by step reasoning</think>final answer", request
40
+ )
41
+ assert reasoning == "step by step reasoning"
42
+ assert content == "final answer"
43
+
44
+
45
+def test_extract_reasoning_empty_thinking(kimi_k2_tokenizer):
46
+ parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
47
+ request = ChatCompletionRequest(model="test-model", messages=[], temperature=1.0)
48
+
49
+ reasoning, content = parser.extract_reasoning(
50
+ "<think></think>final answer", request
51
+ )
52
+ assert reasoning == ""
53
+ assert content == "final answer"
54
+
55
+
56
+def test_extract_reasoning_implicit_start(kimi_k2_tokenizer):
57
+ """When there's no <think> tag, everything is treated as reasoning."""
58
+ parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
59
+ request = ChatCompletionRequest(model="test-model", messages=[], temperature=1.0)
60
+
61
+ reasoning, content = parser.extract_reasoning(
62
+ "implicit reasoning with no tags", request
63
+ )
64
+ assert reasoning == "implicit reasoning with no tags"
65
+ assert content is None
66
+
67
+
68
+def test_extract_reasoning_tool_section_ends_reasoning(kimi_k2_tokenizer):
69
+ """<|tool_calls_section_begin|> implicitly ends reasoning."""
70
+ parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
71
+ request = ChatCompletionRequest(model="test-model", messages=[], temperature=1.0)
72
+
73
+ text = "some reasoning<|tool_calls_section_begin|>tool call data"
74
+ reasoning, content = parser.extract_reasoning(text, request)
75
+ assert reasoning == "some reasoning"
76
+ assert content == "<|tool_calls_section_begin|>tool call data"
77
+
78
+
79
+def test_streaming_reasoning_then_content(kimi_k2_tokenizer):
80
+ """Token-by-token streaming: reasoning tokens then content after </think>."""
81
+ parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
82
+
83
+ think_id = parser._start_token_id
84
+ end_think_id = parser._end_token_id
85
+ # Use a real token ID from the tokenizer for regular content
86
+ regular_id = kimi_k2_tokenizer.encode("hello", add_special_tokens=False)[0]
87
+
88
+ # First token: <think> — single special token should be skipped
89
+ result = parser.extract_reasoning_streaming(
90
+ previous_text="",
91
+ current_text="<think>",
92
+ delta_text="<think>",
93
+ previous_token_ids=[],
94
+ current_token_ids=[think_id],
95
+ delta_token_ids=[think_id],
96
+ )
97
+ assert result is None
98
+
99
+ # Reasoning token
100
+ result = parser.extract_reasoning_streaming(
101
+ previous_text="<think>",
102
+ current_text="<think>step one",
103
+ delta_text="step one",
104
+ previous_token_ids=[think_id],
105
+ current_token_ids=[think_id, regular_id],
106
+ delta_token_ids=[regular_id],
107
+ )
108
+ assert isinstance(result, DeltaMessage)
109
+ assert result.reasoning == "step one"
110
+ assert result.content is None
111
+
112
+ # End token </think> as single token — should be skipped
113
+ result = parser.extract_reasoning_streaming(
114
+ previous_text="<think>step one",
115
+ current_text="<think>step one</think>",
116
+ delta_text="</think>",
117
+ previous_token_ids=[think_id, regular_id],
118
+ current_token_ids=[think_id, regular_id, end_think_id],
119
+ delta_token_ids=[end_think_id],
120
+ )
121
+ assert result is None
122
+
123
+ # Content after </think>
124
+ content_id = kimi_k2_tokenizer.encode("world", add_special_tokens=False)[0]
125
+ result = parser.extract_reasoning_streaming(
126
+ previous_text="<think>step one</think>",
127
+ current_text="<think>step one</think>answer",
128
+ delta_text="answer",
129
+ previous_token_ids=[think_id, regular_id, end_think_id],
130
+ current_token_ids=[think_id, regular_id, end_think_id, content_id],
131
+ delta_token_ids=[content_id],
132
+ )
133
+ assert isinstance(result, DeltaMessage)
134
+ assert result.content == "answer"
135
+
136
+
137
+def test_streaming_tool_section_ends_reasoning(kimi_k2_tokenizer):
138
+ """<|tool_calls_section_begin|> in delta ends reasoning during streaming."""
139
+ parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
140
+
141
+ think_id = parser._start_token_id
142
+ tool_begin_id = parser._tool_section_start_token_id
143
+ regular_id = kimi_k2_tokenizer.encode("hello", add_special_tokens=False)[0]
144
+
145
+ # Tool section token arrives — should transition from reasoning to content
146
+ result = parser.extract_reasoning_streaming(
147
+ previous_text="<think>thinking",
148
+ current_text="<think>thinking<|tool_calls_section_begin|>",
149
+ delta_text="<|tool_calls_section_begin|>",
150
+ previous_token_ids=[think_id, regular_id],
151
+ current_token_ids=[think_id, regular_id, tool_begin_id],
152
+ delta_token_ids=[tool_begin_id],
153
+ )
154
+ assert isinstance(result, DeltaMessage)
155
+ assert result.content == "<|tool_calls_section_begin|>"
vllm/entrypoints/chat_utils.py
+14
@@ -1660,6 +1660,20 @@ def get_history_tool_calls_cnt(conversation: list[ConversationMessage]):
1660
return idx
1661
1662
1663
+_KIMI_MODEL_TYPES = ("kimi_k2", "kimi_k25")
1664
+
1665
+
1666
+def get_tool_call_id_type(model_config: ModelConfig) -> str:
1667
+ """Return the tool-call ID type for a given model configuration."""
1668
+ hf_overrides = getattr(model_config, "hf_overrides", None)
1669
+ if model_config.hf_text_config.model_type in _KIMI_MODEL_TYPES or (
1670
+ isinstance(hf_overrides, dict)
1671
+ and hf_overrides.get("model_type") in _KIMI_MODEL_TYPES
1672
+ ):
1673
+ return "kimi_k2"
1674
+ return "random"
1675
+
1676
+
1677
def make_tool_call_id(id_type: str = "random", func_name=None, idx=None):
1678
if id_type == "kimi_k2":
1679
return f"functions.{func_name}:{idx}"
vllm/entrypoints/openai/chat_completion/serving.py
+2
-9
@@ -19,6 +19,7 @@ from vllm.entrypoints.chat_utils import (
19
ChatTemplateContentFormatOption,
20
ConversationMessage,
21
get_history_tool_calls_cnt,
22
+ get_tool_call_id_type,
23
make_tool_call_id,
24
)
25
from vllm.entrypoints.logger import RequestLogger
@@ -152,15 +153,7 @@ class OpenAIServingChat(OpenAIServing):
153
get_stop_tokens_for_assistant_actions()
154
)
155
155
- # Handle tool call ID type for Kimi K2 (supporting test mocking via overrides)
156
- hf_overrides = getattr(self.model_config, "hf_overrides", None)
157
- if self.model_config.hf_text_config.model_type == "kimi_k2" or (
158
- isinstance(hf_overrides, dict)
159
- and hf_overrides.get("model_type") == "kimi_k2"
160
- ):
161
- self.tool_call_id_type = "kimi_k2"
162
- else:
163
- self.tool_call_id_type = "random"
156
+ self.tool_call_id_type = get_tool_call_id_type(self.model_config)
157
158
# NOTE(woosuk): While OpenAI's chat completion API supports browsing
159
# for some models, currently vLLM doesn't support it. Please use the
vllm/entrypoints/openai/responses/serving.py
+2
-9
@@ -46,6 +46,7 @@ from vllm.engine.protocol import EngineClient
46
from vllm.entrypoints.chat_utils import (
47
ChatCompletionMessageParam,
48
ChatTemplateContentFormatOption,
49
+ get_tool_call_id_type,
50
)
51
from vllm.entrypoints.logger import RequestLogger
52
from vllm.entrypoints.mcp.tool_server import ToolServer
@@ -241,15 +242,7 @@ class OpenAIServingResponses(OpenAIServing):
242
get_stop_tokens_for_assistant_actions()
243
)
244
244
- # Handle tool call ID type for Kimi K2 (supporting test mocking via overrides)
245
- hf_overrides = getattr(self.model_config, "hf_overrides", None)
246
- if self.model_config.hf_text_config.model_type == "kimi_k2" or (
247
- isinstance(hf_overrides, dict)
248
- and hf_overrides.get("model_type") == "kimi_k2"
249
- ):
250
- self.tool_call_id_type = "kimi_k2"
251
- else:
252
- self.tool_call_id_type = "random"
245
+ self.tool_call_id_type = get_tool_call_id_type(self.model_config)
246
247
self.enable_auto_tools = enable_auto_tools
248
# HACK(woosuk): This is a hack. We should use a better store.