diff --git a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md index 7d1da946f34d..6c770d76c0f2 100644 --- a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md @@ -213,7 +213,7 @@ These options are set within the YAML file passed to `trtllm-serve` via the `--c When the served model is Kimi K3, `trtllm-serve` applies Kimi/Moonshot API semantics on `/v1/chat/completions` (all of these are exercised by Moonshot's [Kimi Vendor Verifier](https://www.kimi.com/blog/kimi-vendor-verifier.html)): -* **Request extensions:** the `thinking` object (`{"type": "enabled"|"disabled", "keep": "all", "effort": "low"|"high"|"max"}`), `reasoning_effort` values `"low"`, `"high"`, `"max"`, and `"none"` (an explicit `thinking` object takes precedence), `tool_choice: "required"`, message-level (dynamic) tools declared on system messages, and `response_format` `json_object`/`json_schema` (the `json_schema` wrapper must carry a non-empty `name` and a `schema` object). These map onto the checkpoint chat template's native control messages; explicit `chat_template_kwargs` always win. +* **Request extensions:** the `thinking` object (`{"type": "enabled"|"disabled", "keep": "all", "effort": "low"|"high"|"max"}`), `reasoning_effort` values `"low"`, `"high"`, `"max"`, and `"none"` (an explicit `thinking` object takes precedence), `tool_choice: "required"`, message-level (dynamic) tools declared on system messages, and `response_format` `json_object`/`json_schema` (the `json_schema` wrapper must carry a non-empty `name` and a `schema` object). These map onto the checkpoint chat template's native control messages. For thinking and its effort, which the template reads from the native `thinking` and `thinking_effort` chat-template kwargs, each setting comes from the first of these sources that sets it: (1) `thinking` / `thinking_effort` in `chat_template_kwargs`; (2) the request-level `thinking` object; (3) the request-level `reasoning_effort`; (4) the generic `enable_thinking` (boolean) and `reasoning_effort` (`"none"`, `"low"`, `"high"`, `"max"`) keys in `chat_template_kwargs`, which other chat templates read and `/v1/messages` fills from `thinking` and `output_config.effort`; (5) the template default, thinking on with effort `max`. An effort of `"none"` turns thinking off, and no effort is derived while thinking is off; a `thinking_effort` passed in `chat_template_kwargs` is kept, but the template renders no effort while thinking is off, so it has no effect then. Other effort values (for example `"medium"`) have no K3 equivalent and leave the default. For the other controls, explicit `chat_template_kwargs` win over the values derived from request fields. * **Streaming usage:** `usage` is reported in the final streaming chunk even when the client does not send `stream_options` (Kimi API parity). * **Prompt-token accounting:** reported `usage.prompt_tokens` excludes the trailing 3-token generation channel opener, matching Kimi's reference accounting; the model still consumes the full rendered prompt. * **`TRTLLM_KIMI_PARAM_POLICY`** (default `0`, off): when set to `1`, enforces Kimi's immutable sampling parameters — `top_p` pinned to 0.95 (unset or the OpenAI default `1.0` are coerced to 0.95; other values are rejected with HTTP 400), `presence_penalty`/`frequency_penalty` 0, `n` 1, and `temperature` bounded to [0, 1]. Off by default so existing deployments keep accepting the requests they accept today; a Kimi-Vendor-Verifier certification run must set it to `1` (the KVV params suite requires the rejections). diff --git a/tensorrt_llm/serve/extensions/kimi_k3.py b/tensorrt_llm/serve/extensions/kimi_k3.py index db24c04ea604..2b2b8964b2c6 100644 --- a/tensorrt_llm/serve/extensions/kimi_k3.py +++ b/tensorrt_llm/serve/extensions/kimi_k3.py @@ -73,6 +73,23 @@ def enforce_kimi_param_policy(request: ChatCompletionRequest) -> None: raise ValueError(f"n is fixed at 1 for this model; got {request.n}.") +def _derive_from_reasoning_effort(kwargs: dict[str, Any], effort: object) -> None: + """Derive the K3 ``thinking`` / ``thinking_effort`` kwargs from a reasoning effort. + + A no-op when ``kwargs`` turns thinking off. Otherwise ``"none"`` turns + thinking off unless ``kwargs`` sets it, ``"low"``, ``"high"`` and ``"max"`` + set the effort unless ``kwargs`` sets one, and other values (e.g. harmony's + ``"medium"``) have no K3 equivalent and leave the template default. + """ + if not isinstance(effort, str) or kwargs.get("thinking") is False: + return + effort = effort.lower() + if effort == "none": + kwargs.setdefault("thinking", False) + elif effort in ("low", "high", "max"): + kwargs.setdefault("thinking_effort", effort) + + def dynamic_tool_dicts(messages: Optional[List[ChatCompletionMessageParam]]) -> list[dict]: """Collect message-level (dynamic) tool declarations from system messages.""" tools: list[dict] = [] @@ -124,6 +141,8 @@ def validate_kimi_dynamic_tools(request: ChatCompletionRequest) -> None: class KimiK3ServingExtension(ServingExtension): """Kimi/Moonshot API semantics for kimi_k3 chat requests.""" + render_version = 2 + # Trailing tokens of the native K3 generation prompt (the channel opener # ``<|open|>think|response<|sep|>``) that are not reported as prompt usage. _GENERATION_CHANNEL_OPENER_TOKENS = 3 @@ -170,7 +189,10 @@ def apply_chat_extensions(self, request: ChatCompletionRequest) -> None: effort, tool_choice, and response_format, but only reads them from chat-template kwargs. Derive those kwargs from the request-level fields so the OpenAI-style API surface drives the template; explicit - client-supplied ``chat_template_kwargs`` win over derived values. The + client-supplied ``chat_template_kwargs`` win over derived values. + The template reads ``thinking`` / ``thinking_effort``, not the + ``enable_thinking`` / ``reasoning_effort`` kwargs most other chat + templates read, so those generic kwargs are translated as well. The merged kwargs also steer the kimi_k3 reasoning parser's initial channel, the guided-decoding structural tag, and the thinking-budget logits processor downstream. @@ -183,32 +205,28 @@ def apply_chat_extensions(self, request: ChatCompletionRequest) -> None: if request.stream and request.stream_options is None: # StreamOptions defaults: include_usage=True, continuous off. request.stream_options = StreamOptions() + # Sources of the ``thinking`` / ``thinking_effort`` kwargs, highest + # precedence first; each one only sets what no source above it set: + # 1. those two keys in the client's ``chat_template_kwargs``; + # 2. the request-level ``thinking`` object; + # 3. the request-level ``reasoning_effort``; + # 4. the generic ``enable_thinking`` / ``reasoning_effort`` keys in + # ``chat_template_kwargs`` (the Anthropic adapter maps ``thinking`` + # and ``output_config.effort`` onto them); + # 5. the template default: thinking on, effort ``max``. + # An effort of ``"none"`` turns thinking off. No effort is derived for a + # request whose thinking is off; a ``thinking_effort`` the client sets is + # passed through, and the template ignores it while thinking is off. derived: dict[str, Any] = {} if request.thinking is not None: enabled = request.thinking.type != "disabled" derived["thinking"] = enabled if enabled and request.thinking.effort is not None: derived["thinking_effort"] = request.thinking.effort - if ( - "reasoning_effort" in request.model_fields_set - and request.reasoning_effort is not None - and "thinking_effort" not in derived - and (request.thinking is None or request.thinking.type != "disabled") - ): - # An explicit thinking.effort wins, and an explicit thinking - # object also wins the on/off axis: reasoning_effort only supplies - # the effort when thinking.effort is absent, and - # reasoning_effort="none" only disables thinking when no thinking - # object was sent. No effort is ever derived for an explicitly - # disabled request. - effort = getattr(request.reasoning_effort, "value", request.reasoning_effort).lower() - if effort == "none": - if request.thinking is None: - derived["thinking"] = False - elif effort in ("low", "high", "max"): - derived["thinking_effort"] = effort - # Other efforts (e.g. harmony's "medium") have no K3 equivalent; - # leave the template default. + if "reasoning_effort" in request.model_fields_set and request.reasoning_effort is not None: + _derive_from_reasoning_effort( + derived, getattr(request.reasoning_effort, "value", request.reasoning_effort) + ) if ( "tool_choice" in request.model_fields_set and (request.tools or dynamic_tool_dicts(request.messages)) @@ -234,11 +252,33 @@ def apply_chat_extensions(self, request: ChatCompletionRequest) -> None: if "strict" in json_schema and not isinstance(json_schema["strict"], bool): raise ValueError("response_format.json_schema.strict must be a boolean.") derived["response_schema"] = json_schema["schema"] - if derived: - request.chat_template_kwargs = { - **derived, - **(request.chat_template_kwargs or {}), - } + chat_template_kwargs = {**derived, **(request.chat_template_kwargs or {})} + unset = {"thinking", "thinking_effort"} - chat_template_kwargs.keys() + enable_thinking = chat_template_kwargs.get("enable_thinking") + if isinstance(enable_thinking, bool): + chat_template_kwargs.setdefault("thinking", enable_thinking) + _derive_from_reasoning_effort( + chat_template_kwargs, chat_template_kwargs.get("reasoning_effort") + ) + translated = unset & chat_template_kwargs.keys() + exempted = set(translated) + generic_effort = chat_template_kwargs.get("reasoning_effort") + if isinstance(generic_effort, str) and generic_effort.lower() in ( + "none", + "low", + "high", + "max", + ): + # The extension consumes this control even if a higher-precedence + # K3 setting wins. Keep it available to custom templates as well. + exempted.add("reasoning_effort") + if exempted: + # Derived and consumed controls need not be read again by a Jinja template. + request.injected_chat_template_kwargs = sorted( + {*(request.injected_chat_template_kwargs or ()), *exempted} + ) + if derived or translated: + request.chat_template_kwargs = chat_template_kwargs def structured_output_format( self, content: Dict[str, Any], chat_template_kwargs: Optional[Dict[str, Any]] diff --git a/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py b/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py index b39bdf437771..f2cc94d81a50 100644 --- a/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py +++ b/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py @@ -388,6 +388,154 @@ def test_response_format_json_schema_wrapper_validation(self, wrapper: dict, msg self.apply(response_format={"type": "json_schema", "json_schema": wrapper}) +def k3_thinking_kwargs(req: ChatCompletionRequest) -> dict: + """The K3 template's own thinking controls in the request's merged kwargs.""" + kwargs = req.chat_template_kwargs or {} + return {key: kwargs[key] for key in ("thinking", "thinking_effort") if key in kwargs} + + +class TestKimiGenericThinkingKwargs: + """``enable_thinking`` / ``reasoning_effort``: the kwargs other chat templates read.""" + + def apply(self, **kwargs) -> ChatCompletionRequest: + req = make_request(**kwargs) + apply_model_chat_extensions(req, "kimi_k3") + return req + + @pytest.mark.parametrize( + "generic, expected", + [ + ({"enable_thinking": False}, {"thinking": False}), + ({"enable_thinking": True}, {"thinking": True}), + ({"reasoning_effort": "low"}, {"thinking_effort": "low"}), + ({"reasoning_effort": "high"}, {"thinking_effort": "high"}), + ({"reasoning_effort": "max"}, {"thinking_effort": "max"}), + ({"reasoning_effort": "High"}, {"thinking_effort": "high"}), + ({"reasoning_effort": "none"}, {"thinking": False}), + # No K3 equivalent: the template default applies. + ({"reasoning_effort": "medium"}, {}), + # No effort for a request whose thinking is off ... + ({"enable_thinking": False, "reasoning_effort": "high"}, {"thinking": False}), + # ... and "none" does not override an explicit thinking setting. + ({"enable_thinking": True, "reasoning_effort": "none"}, {"thinking": True}), + # Only a bool is a thinking setting. + ({"enable_thinking": "false"}, {}), + # "none" turns thinking off even beside a native effort, which is kept; the template + # ignores it while thinking is off. + ( + {"thinking_effort": "high", "reasoning_effort": "none"}, + {"thinking": False, "thinking_effort": "high"}, + ), + # A native effort wins over the generic one. + ({"thinking_effort": "low", "reasoning_effort": "max"}, {"thinking_effort": "low"}), + ], + ids=[ + "thinking_off", + "thinking_on", + "effort_low", + "effort_high", + "effort_max", + "effort_any_case", + "effort_none", + "effort_medium", + "no_effort_when_off", + "none_keeps_thinking_on", + "non_bool_thinking", + "none_beside_native_effort", + "native_effort_wins", + ], + ) + def test_generic_kwargs_map_to_the_k3_kwargs(self, generic: dict, expected: dict) -> None: + req = self.apply(chat_template_kwargs=generic) + assert k3_thinking_kwargs(req) == expected + + def test_k3_kwargs_win_over_the_generic_ones(self) -> None: + req = self.apply( + chat_template_kwargs={ + "thinking": True, + "thinking_effort": "high", + "enable_thinking": False, + "reasoning_effort": "low", + } + ) + assert k3_thinking_kwargs(req) == {"thinking": True, "thinking_effort": "high"} + assert req.injected_chat_template_kwargs == ["reasoning_effort"] + + @pytest.mark.parametrize( + "fields, expected", + [ + ( + {"thinking": {"type": "enabled", "effort": "high"}}, + {"thinking": True, "thinking_effort": "high"}, + ), + ({"thinking": {"type": "disabled"}}, {"thinking": False}), + ({"reasoning_effort": "none"}, {"thinking": False}), + ({"reasoning_effort": "high"}, {"thinking": True, "thinking_effort": "high"}), + ], + ids=["thinking_effort", "thinking_disabled", "reasoning_effort_none", "reasoning_effort"], + ) + def test_request_fields_win_over_the_generic_kwargs(self, fields: dict, expected: dict) -> None: + req = self.apply( + chat_template_kwargs={"enable_thinking": True, "reasoning_effort": "low"}, **fields + ) + assert k3_thinking_kwargs(req) == expected + + def test_translated_kwargs_are_marked_as_server_injected(self) -> None: + # The unused-kwargs guard exempts injected keys when a Jinja template renders. + req = self.apply( + chat_template_kwargs={"enable_thinking": True, "reasoning_effort": "low"}, + injected_chat_template_kwargs=["enable_thinking"], + ) + assert req.injected_chat_template_kwargs == [ + "enable_thinking", + "reasoning_effort", + "thinking", + "thinking_effort", + ] + + def test_thinking_off_is_parsed_as_an_answer_from_the_start(self) -> None: + from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory + + req = self.apply(chat_template_kwargs={"enable_thinking": False}) + parser = ReasoningParserFactory.create_reasoning_parser("kimi_k3", req.chat_template_kwargs) + result = parser.parse("plain answer<|close|>response<|sep|><|end_of_msg|>") + assert (result.content, result.reasoning_content) == ("plain answer", "") + + +class TestKimiAnthropicThinkingControls: + """Anthropic ``thinking`` / ``output_config.effort`` reach the K3 template kwargs.""" + + def convert(self, **fields) -> ChatCompletionRequest: + from tensorrt_llm.serve.anthropic_adapter import convert_anthropic_request + from tensorrt_llm.serve.anthropic_protocol import AnthropicMessagesRequest + + req = convert_anthropic_request( + AnthropicMessagesRequest(model="m", max_tokens=4096, messages=[USER_MSG], **fields) + ) + apply_model_chat_extensions(req, "kimi_k3") + return req + + def test_thinking_disabled_turns_k3_thinking_off(self) -> None: + req = self.convert(thinking={"type": "disabled"}, output_config={"effort": "high"}) + assert k3_thinking_kwargs(req) == {"thinking": False} + assert "thinking" in req.injected_chat_template_kwargs + + @pytest.mark.parametrize("effort", ["low", "high", "max"]) + def test_effort_sets_the_k3_effort(self, effort: str) -> None: + req = self.convert( + thinking={"type": "enabled", "budget_tokens": 1024}, output_config={"effort": effort} + ) + assert k3_thinking_kwargs(req) == {"thinking": True, "thinking_effort": effort} + assert "thinking_effort" in req.injected_chat_template_kwargs + + def test_applying_the_extension_again_changes_nothing(self) -> None: + # A router renders a request and forwards it; the worker applies the extension again. + req = self.convert(output_config={"effort": "low"}) + kwargs, injected = dict(req.chat_template_kwargs), list(req.injected_chat_template_kwargs) + apply_model_chat_extensions(req, "kimi_k3") + assert (req.chat_template_kwargs, req.injected_chat_template_kwargs) == (kwargs, injected) + + class TestKimiParamPolicy: @pytest.fixture(autouse=True) def _enable_policy(self, monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/unittest/llmapi/apps/test_render_core.py b/tests/unittest/llmapi/apps/test_render_core.py index 3221d6087636..348ac2b9ec21 100644 --- a/tests/unittest/llmapi/apps/test_render_core.py +++ b/tests/unittest/llmapi/apps/test_render_core.py @@ -561,6 +561,83 @@ def test_a_gpt_oss_model_without_the_harmony_switch_uses_the_template( != render_chat(chat_request(), harmony_resources()).token_ids ) + @pytest.mark.parametrize( + ("kwargs", "template", "expected"), + [ + ({"reasoning_effort": "LOW"}, "{{ thinking_effort }}", "low"), + ({"reasoning_effort": "none"}, "{{ thinking }}", "False"), + ({"enable_thinking": False}, "{{ thinking }}", "False"), + ( + {"reasoning_effort": "low", "thinking_effort": "high"}, + "{{ thinking_effort }}", + "high", + ), + ( + {"reasoning_effort": "low"}, + "{{ reasoning_effort }}:{{ thinking_effort }}", + "low:low", + ), + ], + ) + def test_kimi_generic_effort_is_consumed_before_jinja_validation( + self, tokenizer, monkeypatch, kwargs, template, expected + ) -> None: + monkeypatch.delenv("TRTLLM_ALLOW_UNUSED_CHAT_TEMPLATE_KWARGS", raising=False) + request = chat_request(chat_template=template, chat_template_kwargs=kwargs) + res = RenderResources.from_tokenizer(tokenizer, model_type="kimi_k3") + + assert render_chat(request, res, tokenize=False).text == expected + + def test_kimi_unrecognized_generic_effort_still_requires_a_template_consumer( + self, tokenizer, monkeypatch + ) -> None: + from tensorrt_llm.inputs.chat_template_guard import UnusedChatTemplateKwargsError + + monkeypatch.delenv("TRTLLM_ALLOW_UNUSED_CHAT_TEMPLATE_KWARGS", raising=False) + request = chat_request( + chat_template="{{ thinking_effort | default('max') }}", + chat_template_kwargs={"reasoning_effort": "medium"}, + ) + res = RenderResources.from_tokenizer(tokenizer, model_type="kimi_k3") + + with pytest.raises(UnusedChatTemplateKwargsError, match="reasoning_effort"): + render_chat(request, res, tokenize=False) + + def test_the_generic_thinking_kwargs_reach_the_kimi_k3_prompt(self) -> None: + # K3's template reads `thinking` / `thinking_effort`; the extension derives them from the + # `enable_thinking` / `reasoning_effort` that other templates read and the Anthropic + # adapter sends. + res = RenderResources.from_tokenizer(make_tokenizer(name=KIMI_K3), model_type="kimi_k3") + + def render(**chat_template_kwargs) -> str: + request = chat_request( + messages=[{"role": "user", "content": "hi"}], + chat_template_kwargs=chat_template_kwargs or None, + ) + return render_chat(request, res, tokenize=False).text + + default = render() + assert default.startswith('<|open|>message role="system" type="thinking-effort"<|sep|>') + assert "`thinking_effort=max`" in default + assert default.endswith("<|open|>think<|sep|>") + + # Thinking off: the prompt opens with the user turn and ends in the answer channel. + off = render(enable_thinking=False) + assert off.startswith('<|open|>message role="user"<|sep|>hi') + assert off.endswith("<|open|>response<|sep|>") + + low = render(reasoning_effort="low") + assert low == default.replace("`thinking_effort=max`", "`thinking_effort=low`") + + # A native effort is passed through while thinking is off, and the template ignores it. + native = chat_request( + messages=[{"role": "user", "content": "hi"}], + chat_template_kwargs={"thinking_effort": "high", "reasoning_effort": "none"}, + ) + assert render_chat(native, res, tokenize=False).text == off + assert native.chat_template_kwargs["thinking"] is False + assert native.chat_template_kwargs["thinking_effort"] == "high" + class TestRenderConversation: def test_matches_render_chat_for_a_parsed_conversation(self, tokenizer) -> None: