| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Reasoning parser for the Kimi K3 (XTML) chat format. |
| 4 | |
| 5 | This strips the ``think`` channel out of generated text and hands the remainder |
| 6 | (``response`` + ``tools`` channels) downstream. Kimi K3 wraps the thinking |
| 7 | channel as an XTML element built from special tokens:: |
| 8 | |
| 9 | <|open|>think<|sep|> <reasoning> <|close|>think<|sep|> |
| 10 | |
| 11 | Two subtleties drive the implementation: |
| 12 | * Unlike Kimi-K2 (a single ``<think>`` token), each K3 marker is a 3-token |
| 13 | sequence, so the token-id helpers search for the marker *subsequence* |
| 14 | rather than a single id. |
| 15 | * In thinking mode the serving layer may feed ``<|open|>think<|sep|>`` as the |
| 16 | generation prefix, so the model's output can begin *inside* the think |
| 17 | channel with no open marker. The text paths therefore treat a missing open |
| 18 | marker as "reasoning starts at offset 0". |
| 19 | |
| 20 | When thinking is disabled (``chat_template_kwargs={"thinking": False}`` or |
| 21 | ``{"enable_thinking": False}``, i.e. instruct mode) the parser returns every |
| 22 | delta as normal content; there is simply no think channel to extract. |
| 23 | """ |
| 24 | |
| 25 | from collections.abc import Iterable, Sequence |
| 26 | from typing import TYPE_CHECKING |
| 27 | |
| 28 | import regex as re |
| 29 | from transformers import PreTrainedTokenizerBase |
| 30 | |
| 31 | from vllm.entrypoints.generate.base.protocol import DeltaMessage |
| 32 | from vllm.reasoning import ReasoningParser |
| 33 | |
| 34 | if TYPE_CHECKING: |
| 35 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 36 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 37 | |
| 38 | |
| 39 | def _match_at(haystack: Sequence[int], i: int, needle: Sequence[int]) -> bool: |
| 40 | """Whether *needle* occurs in *haystack* starting at *i*. |
| 41 | |
| 42 | Compares element by element rather than slicing: ``haystack`` is usually a |
| 43 | ``ConstantList``, whose slices cost a Python ``__getitem__`` plus a list |
| 44 | allocation on every probe. |
| 45 | """ |
| 46 | return all(haystack[i + k] == needle[k] for k in range(len(needle))) |
| 47 | |
| 48 | |
| 49 | def _subseq_index(haystack: Sequence[int], needle: Sequence[int]) -> int: |
| 50 | """Return start index of the last occurrence of needle in haystack, or -1.""" |
| 51 | n = len(needle) |
| 52 | if n == 0: |
| 53 | return -1 |
| 54 | first = needle[0] |
| 55 | for i in range(len(haystack) - n, -1, -1): |
| 56 | if haystack[i] == first and _match_at(haystack, i, needle): |
| 57 | return i |
| 58 | return -1 |
| 59 | |
| 60 | |
| 61 | def _newest_marker(haystack: Sequence[int], a: Sequence[int], b: Sequence[int]) -> int: |
| 62 | """Report which of *a* / *b* occurs last in *haystack*. |
| 63 | |
| 64 | Returns 0 if *a* is the newest marker, 1 if *b* is, -1 if neither occurs. |
| 65 | Equivalent to comparing two ``_subseq_index`` results, but a single backward |
| 66 | pass that stops at the first hit instead of walking to index 0 twice. |
| 67 | """ |
| 68 | if not a or not b: |
| 69 | return -1 |
| 70 | a0, b0 = a[0], b[0] |
| 71 | for i in range(len(haystack) - min(len(a), len(b)), -1, -1): |
| 72 | head = haystack[i] |
| 73 | if head == a0 and i + len(a) <= len(haystack) and _match_at(haystack, i, a): |
| 74 | return 0 |
| 75 | if head == b0 and i + len(b) <= len(haystack) and _match_at(haystack, i, b): |
| 76 | return 1 |
| 77 | return -1 |
| 78 | |
| 79 | |
| 80 | class KimiK3ReasoningParser(ReasoningParser): |
| 81 | """Reasoning parser for the Kimi K3 (XTML) think channel.""" |
| 82 | |
| 83 | def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs): |
| 84 | super().__init__(tokenizer) |
| 85 | |
| 86 | if not self.model_tokenizer: |
| 87 | raise ValueError( |
| 88 | "The model tokenizer must be passed to the ReasoningParser " |
| 89 | "constructor during construction." |
| 90 | ) |
| 91 | |
| 92 | # thinking can be disabled via chat_template_kwargs -> identity fallthrough |
| 93 | chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} |
| 94 | thinking = chat_kwargs.get("thinking", None) |
| 95 | if thinking is None: |
| 96 | thinking = chat_kwargs.get("enable_thinking", True) |
| 97 | self._thinking_enabled = bool(thinking) |
| 98 | |
| 99 | # XTML markers as literal strings (skip_special_tokens=False at serve time) |
| 100 | self._think_open = "<|open|>think<|sep|>" |
| 101 | self._think_close = "<|close|>think<|sep|>" |
| 102 | |
| 103 | # Content-channel markers that must be stripped from the final output. |
| 104 | # The tool parser handles these when active, but when tool_choice is |
| 105 | # "none" (the default for requests without tools) the tool parser is |
| 106 | # bypassed and the reasoning parser must strip them itself. |
| 107 | self._response_open = "<|open|>response<|sep|>" |
| 108 | self._response_close = "<|close|>response<|sep|>" |
| 109 | self._message_close = "<|close|>message<|sep|>" |
| 110 | |
| 111 | # Tolerant matchers: defense-in-depth against vLLM's added-token spacing |
| 112 | # ("<|open|> think <|sep|>"). The `\s*` is a no-op on clean input, so the |
| 113 | # normal path stays byte-exact; it only helps if some serving path leaves |
| 114 | # spaces_between_special_tokens on. See adjust_request for the real fix. |
| 115 | open_marker = r"<\|open\|>" |
| 116 | close_marker = r"<\|close\|>" |
| 117 | sep_marker = r"<\|sep\|>" |
| 118 | self._think_open_re = re.compile(open_marker + r"\s*think\s*" + sep_marker) |
| 119 | self._think_close_re = re.compile(close_marker + r"\s*think\s*" + sep_marker) |
| 120 | self._response_open_re = re.compile( |
| 121 | open_marker + r"\s*response\s*" + sep_marker |
| 122 | ) |
| 123 | self._response_close_re = re.compile( |
| 124 | close_marker + r"\s*response\s*" + sep_marker |
| 125 | ) |
| 126 | self._message_close_re = re.compile( |
| 127 | close_marker + r"\s*message\s*" + sep_marker |
| 128 | ) |
| 129 | |
| 130 | # marker token-id subsequences (each marker is 3 tokens) |
| 131 | self._think_open_ids = tokenizer.encode( |
| 132 | self._think_open, add_special_tokens=False |
| 133 | ) |
| 134 | self._think_close_ids = tokenizer.encode( |
| 135 | self._think_close, add_special_tokens=False |
| 136 | ) |
| 137 | self._response_open_ids = tokenizer.encode( |
| 138 | self._response_open, add_special_tokens=False |
| 139 | ) |
| 140 | self._last_streaming_delta_token_ids: tuple[int, ...] | None = None |
| 141 | self._last_streaming_content_token_ids: list[int] | None = None |
| 142 | |
| 143 | @property |
| 144 | def reasoning_start_str(self) -> str | None: |
| 145 | return self._think_open |
| 146 | |
| 147 | @property |
| 148 | def reasoning_end_str(self) -> str | None: |
| 149 | return self._think_close |
| 150 | |
| 151 | def adjust_request( |
| 152 | self, |
| 153 | request: "ChatCompletionRequest | ResponsesRequest", |
| 154 | ) -> "ChatCompletionRequest | ResponsesRequest": |
| 155 | request.skip_special_tokens = False |
| 156 | if hasattr(request, "spaces_between_special_tokens"): |
| 157 | request.spaces_between_special_tokens = False |
| 158 | return request |
| 159 | |
| 160 | def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: |
| 161 | if not self._thinking_enabled: |
| 162 | return True |
| 163 | # Reasoning has ended only if the *most recent* think block is closed: |
| 164 | # the last close marker must come after the last open marker. A plain |
| 165 | # "a close marker exists anywhere" check false-positives in multi-turn / |
| 166 | # agent continuations, where the chat template keeps a prior turn's |
| 167 | # think channel (with its <|close|>think<|sep|>) in the prompt while the |
| 168 | # current turn is still reasoning (its <|open|>think<|sep|> is the newest |
| 169 | # marker). A missing open marker (e.g. it was consumed as the generation |
| 170 | # prefix) means a close marker alone ends reasoning, which is what |
| 171 | # "close is the newest marker" already encodes. |
| 172 | return ( |
| 173 | _newest_marker(input_ids, self._think_close_ids, self._think_open_ids) == 0 |
| 174 | ) |
| 175 | |
| 176 | def is_reasoning_end_streaming( |
| 177 | self, input_ids: Sequence[int], delta_ids: Iterable[int] |
| 178 | ) -> bool: |
| 179 | """Reasoning-end check for a single decode step. |
| 180 | |
| 181 | The engine calls this once per structured-output request per decode step |
| 182 | while the request is still inside the think channel, so it only has to |
| 183 | look at the tokens generated this step, plus ``len(marker) - 1`` tokens |
| 184 | of context in case a marker straddles the step boundary. |
| 185 | |
| 186 | The inherited default re-runs the full-sequence ``is_reasoning_end``, |
| 187 | which makes the scheduler O(context) per request per step and starves the |
| 188 | GPU on long agentic contexts. |
| 189 | """ |
| 190 | if not self._thinking_enabled: |
| 191 | return True |
| 192 | delta = list(delta_ids) |
| 193 | if not delta: |
| 194 | return False |
| 195 | carry = max(len(self._think_close_ids), len(self._think_open_ids)) - 1 |
| 196 | head = len(input_ids) - len(delta) |
| 197 | window = list(input_ids[max(0, head - carry) : head]) + delta |
| 198 | return _newest_marker(window, self._think_close_ids, self._think_open_ids) == 0 |
| 199 | |
| 200 | def _extract_content_ids(self, input_ids: list[int]) -> list[int]: |
| 201 | if not self._thinking_enabled: |
| 202 | return input_ids |
| 203 | idx = _subseq_index(input_ids, self._think_close_ids) |
| 204 | if idx == -1: |
| 205 | return [] # still reasoning |
| 206 | return input_ids[idx + len(self._think_close_ids) :] |
| 207 | |
| 208 | def extract_content_ids(self, input_ids: list[int]) -> list[int]: |
| 209 | cached_delta_ids = self._last_streaming_delta_token_ids |
| 210 | cached_content_ids = self._last_streaming_content_token_ids |
| 211 | self._last_streaming_delta_token_ids = None |
| 212 | self._last_streaming_content_token_ids = None |
| 213 | if cached_delta_ids == tuple(input_ids) and cached_content_ids is not None: |
| 214 | return cached_content_ids |
| 215 | return self._extract_content_ids(input_ids) |
| 216 | |
| 217 | def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: |
| 218 | if not self._thinking_enabled: |
| 219 | return 0 |
| 220 | think_open = self._think_open_ids |
| 221 | think_close = self._think_close_ids |
| 222 | response_open = self._response_open_ids |
| 223 | first_ids = {m[0] for m in (think_open, think_close, response_open) if m} |
| 224 | n = len(token_ids) |
| 225 | count = 0 |
| 226 | in_reasoning = True |
| 227 | seen_open = False |
| 228 | i = 0 |
| 229 | while i < n: |
| 230 | head = token_ids[i] |
| 231 | if head in first_ids: |
| 232 | if i + len(think_open) <= n and _match_at(token_ids, i, think_open): |
| 233 | in_reasoning = True |
| 234 | seen_open = True |
| 235 | i += len(think_open) |
| 236 | continue |
| 237 | if i + len(think_close) <= n and _match_at(token_ids, i, think_close): |
| 238 | in_reasoning = False |
| 239 | i += len(think_close) |
| 240 | continue |
| 241 | if ( |
| 242 | not seen_open |
| 243 | and i + len(response_open) <= n |
| 244 | and _match_at(token_ids, i, response_open) |
| 245 | ): |
| 246 | in_reasoning = False |
| 247 | i += len(response_open) |
| 248 | continue |
| 249 | if in_reasoning: |
| 250 | count += 1 |
| 251 | i += 1 |
| 252 | return count |
| 253 | |
| 254 | def _strip_content_wrapper(self, text: str) -> str: |
| 255 | """Strip ``<|open|>response<|sep|>…<|close|>response<|sep|>`` wrapper and |
| 256 | ``<|close|>message<|sep|>`` from *text*. |
| 257 | |
| 258 | When Kimi K3 tool parsing is active it needs the raw XTML |
| 259 | ``response`` + ``tools`` channels. Otherwise the reasoning parser cleans |
| 260 | up the response wrapper itself so API users do not see XTML markers. |
| 261 | """ |
| 262 | # Try structured unwrap first: extract body of response channel |
| 263 | m_ro = self._response_open_re.search(text) |
| 264 | m_rc = self._response_close_re.search(text, m_ro.end() if m_ro else 0) |
| 265 | if m_ro is not None and m_rc is not None: |
| 266 | text = text[m_ro.end() : m_rc.start()] |
| 267 | elif m_ro is not None: |
| 268 | # response opened but not closed (truncated) |
| 269 | text = text[m_ro.end() :] |
| 270 | else: |
| 271 | # No response wrapper — strip stray markers if present |
| 272 | text = self._response_open_re.sub("", text) |
| 273 | text = self._response_close_re.sub("", text) |
| 274 | text = self._message_close_re.sub("", text) |
| 275 | return text |
| 276 | |
| 277 | @staticmethod |
| 278 | def _should_preserve_tool_channels( |
| 279 | request: "ChatCompletionRequest | ResponsesRequest", |
| 280 | ) -> bool: |
| 281 | return bool(getattr(request, "tools", None)) and ( |
| 282 | getattr(request, "tool_choice", None) != "none" |
| 283 | ) |
| 284 | |
| 285 | def _content_after_reasoning( |
| 286 | self, |
| 287 | text: str, |
| 288 | request: "ChatCompletionRequest | ResponsesRequest", |
| 289 | ) -> str | None: |
| 290 | if self._should_preserve_tool_channels(request): |
| 291 | return text or None |
| 292 | return self._strip_content_wrapper(text) or None |
| 293 | |
| 294 | def extract_reasoning( |
| 295 | self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest" |
| 296 | ) -> tuple[str | None, str | None]: |
| 297 | """Split full text into ``(reasoning, rest)`` for the non-streaming path. |
| 298 | |
| 299 | Handles four shapes: |
| 300 | * response opener without think markers -> no think channel; all content |
| 301 | * open marker present -> reasoning starts after ``<|open|>think<|sep|>`` |
| 302 | * open marker absent but a close marker exists (gen-prefix consumed |
| 303 | the open) -> reasoning starts at offset 0 |
| 304 | * neither marker present -> truncated reasoning after a consumed |
| 305 | generation prefix |
| 306 | ``rest`` is whatever follows the close marker, fed on to the tool parser. |
| 307 | """ |
| 308 | if not self._thinking_enabled: |
| 309 | return None, self._content_after_reasoning(model_output, request) |
| 310 | |
| 311 | m_open = self._think_open_re.search(model_output) |
| 312 | # reasoning content begins right after think-open (or at start if the |
| 313 | # open marker was already consumed as a generation prefix) |
| 314 | content_start = m_open.end() if m_open is not None else 0 |
| 315 | |
| 316 | m_close = self._think_close_re.search(model_output, content_start) |
| 317 | if ( |
| 318 | m_open is None |
| 319 | and m_close is None |
| 320 | and self._response_open_re.search(model_output) is not None |
| 321 | ): |
| 322 | return None, self._content_after_reasoning(model_output, request) |
| 323 | if m_close is not None: |
| 324 | reasoning = model_output[content_start : m_close.start()] |
| 325 | rest = model_output[m_close.end() :] |
| 326 | return (reasoning or None, self._content_after_reasoning(rest, request)) |
| 327 | # think not closed -> still reasoning, no content yet |
| 328 | return (model_output[content_start:] or None, None) |
| 329 | |
| 330 | def _reasoning_text_ready_to_emit(self, text: str) -> str: |
| 331 | """Return the reasoning prefix that is safe to stream now. |
| 332 | |
| 333 | Work from accumulated text, not from the current delta alone. That turns |
| 334 | split-open and split-close handling into the same prefix-diff problem. |
| 335 | |
| 336 | Example 1: split think-open marker. |
| 337 | chunks: ``<|open|>`` / ``think`` / ``<|sep|>reasoning`` |
| 338 | current text after chunk 1: ``<|open|>`` -> emit ``""`` |
| 339 | current text after chunk 2: ``<|open|>think`` -> emit ``""`` |
| 340 | current text after chunk 3: ``<|open|>think<|sep|>reasoning`` |
| 341 | -> emit ``reasoning`` |
| 342 | |
| 343 | Example 2: split think-close marker. |
| 344 | chunks: ``reasoning`` / ``<|close|>`` / ``think<|sep|>...`` |
| 345 | after ``<|close|>``, the suffix is only a partial close marker, so the |
| 346 | sendable reasoning is still ``reasoning`` and the delta is empty. Once |
| 347 | ``think<|sep|>`` arrives, the close branch hands the following response |
| 348 | or tools text to the downstream parser. |
| 349 | """ |
| 350 | m_open = self._think_open_re.search(text) |
| 351 | if m_open is not None: |
| 352 | text = text[m_open.end() :] |
| 353 | overlap = 0 |
| 354 | for marker in (self._think_open, self._think_close): |
| 355 | max_check = min(len(marker) - 1, len(text)) |
| 356 | for n in range(max_check, 0, -1): |
| 357 | if text.endswith(marker[:n]): |
| 358 | overlap = max(overlap, n) |
| 359 | break |
| 360 | return text[:-overlap] if overlap else text |
| 361 | |
| 362 | def _content_ready_to_emit(self, text: str) -> str: |
| 363 | """Return the content prefix that is safe to stream now. |
| 364 | |
| 365 | Mirrors ``_reasoning_text_ready_to_emit`` but for the post-reasoning |
| 366 | content phase. Strips the ``<|open|>response<|sep|>`` prefix, holds back |
| 367 | any partial marker suffix, and removes complete |
| 368 | ``<|close|>response<|sep|>`` / ``<|close|>message<|sep|>`` markers. |
| 369 | """ |
| 370 | # Strip response-open prefix |
| 371 | m_open = self._response_open_re.search(text) |
| 372 | if m_open is not None: |
| 373 | text = text[m_open.end() :] |
| 374 | |
| 375 | # Remove complete close/message markers |
| 376 | text = self._response_close_re.sub("", text) |
| 377 | text = self._message_close_re.sub("", text) |
| 378 | |
| 379 | # Hold back partial markers at the end |
| 380 | overlap = 0 |
| 381 | for marker in ( |
| 382 | self._response_open, |
| 383 | self._response_close, |
| 384 | self._message_close, |
| 385 | ): |
| 386 | max_check = min(len(marker) - 1, len(text)) |
| 387 | for n in range(max_check, 0, -1): |
| 388 | if text.endswith(marker[:n]): |
| 389 | overlap = max(overlap, n) |
| 390 | break |
| 391 | return text[:-overlap] if overlap else text |
| 392 | |
| 393 | def strip_content_streaming( |
| 394 | self, |
| 395 | previous_text: str, |
| 396 | current_text: str, |
| 397 | ) -> DeltaMessage | None: |
| 398 | """Strip XTML content wrappers from streaming deltas after reasoning. |
| 399 | |
| 400 | Called by KimiK3Parser when no tool parser is configured, so the |
| 401 | reasoning parser handles ``<|open|>response<|sep|>`` / |
| 402 | ``<|close|>response<|sep|>`` / ``<|close|>message<|sep|>`` stripping itself. |
| 403 | |
| 404 | Works from accumulated text (``previous_text`` / ``current_text`` |
| 405 | already contain only post-reasoning content). |
| 406 | """ |
| 407 | current_safe = self._content_ready_to_emit(current_text) |
| 408 | previous_safe = self._content_ready_to_emit(previous_text) |
| 409 | if current_safe.startswith(previous_safe): |
| 410 | delta = current_safe[len(previous_safe) :] |
| 411 | else: |
| 412 | delta = current_safe |
| 413 | return DeltaMessage(content=delta) if delta else None |
| 414 | |
| 415 | def extract_reasoning_streaming( |
| 416 | self, |
| 417 | previous_text: str, |
| 418 | current_text: str, |
| 419 | delta_text: str, |
| 420 | previous_token_ids: Sequence[int], |
| 421 | current_token_ids: Sequence[int], |
| 422 | delta_token_ids: Sequence[int], |
| 423 | ) -> DeltaMessage | None: |
| 424 | self._last_streaming_delta_token_ids = None |
| 425 | self._last_streaming_content_token_ids = None |
| 426 | if not self._thinking_enabled: |
| 427 | return DeltaMessage(content=delta_text) |
| 428 | |
| 429 | # reasoning already ended -> downstream content |
| 430 | if self._think_close_re.search(previous_text): |
| 431 | return DeltaMessage(content=delta_text) |
| 432 | |
| 433 | # the close marker completes within this delta's accumulated text: |
| 434 | # split the buffer at the close marker into reasoning vs trailing content. |
| 435 | m_close = self._think_close_re.search(current_text) |
| 436 | if m_close is not None: |
| 437 | self._last_streaming_delta_token_ids = tuple(delta_token_ids) |
| 438 | self._last_streaming_content_token_ids = self._extract_content_ids( |
| 439 | list(current_token_ids) |
| 440 | ) |
| 441 | m_open = self._think_open_re.search(current_text) |
| 442 | r_start = m_open.end() if m_open is not None else 0 |
| 443 | reasoning = current_text[r_start : m_close.start()] |
| 444 | already_sent = self._reasoning_text_ready_to_emit(previous_text) |
| 445 | if reasoning.startswith(already_sent): |
| 446 | reasoning_delta = reasoning[len(already_sent) :] |
| 447 | else: |
| 448 | reasoning_delta = reasoning |
| 449 | content = current_text[m_close.end() :] |
| 450 | return DeltaMessage( |
| 451 | reasoning=reasoning_delta or None, |
| 452 | content=content or None, |
| 453 | ) |
| 454 | |
| 455 | current_reasoning = self._reasoning_text_ready_to_emit(current_text) |
| 456 | previous_reasoning = self._reasoning_text_ready_to_emit(previous_text) |
| 457 | if current_reasoning.startswith(previous_reasoning): |
| 458 | reasoning_delta = current_reasoning[len(previous_reasoning) :] |
| 459 | else: |
| 460 | reasoning_delta = current_reasoning |
| 461 | if not reasoning_delta: |
| 462 | return None |
| 463 | return DeltaMessage(reasoning=reasoning_delta) |
| 464 | |
| 465 | # Backward-compatible aliases for existing unit tests and downstream users |
| 466 | # that still call the pre-split method names. |
| 467 | extract_reasoning_content = extract_reasoning |
| 468 | extract_reasoning_content_streaming = extract_reasoning_streaming |